From 768ce4340ab24bee354bfb1fd4e5d2d3ca4d826d Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 27 Jul 2026 19:44:47 +0800 Subject: [PATCH 01/14] feat(sim): integrate dexsim gizmo controllers Delegate entity and robot gizmo control to dexsim, wire entity gizmos into viewer lifecycle, and exclude the default plane from manipulation. --- docs/source/features/interaction/window.md | 46 + docs/source/tutorial/gizmo.rst | 104 ++- embodichain/lab/sim/cfg.py | 4 +- embodichain/lab/sim/objects/__init__.py | 4 +- embodichain/lab/sim/objects/gizmo.py | 942 ++++++++++++--------- embodichain/lab/sim/sim_manager.py | 230 ++++- embodichain/lab/sim/utility/gizmo_utils.py | 115 +-- examples/sim/gizmo/gizmo_object.py | 42 +- examples/sim/gizmo/gizmo_robot.py | 42 +- scripts/tutorials/sim/gizmo_robot.py | 32 +- tests/sim/objects/test_gizmo.py | 189 +++++ tests/sim/test_sim_manager.py | 183 +++- 12 files changed, 1374 insertions(+), 559 deletions(-) create mode 100644 tests/sim/objects/test_gizmo.py diff --git a/docs/source/features/interaction/window.md b/docs/source/features/interaction/window.md index 10faf2110..da7279ce0 100644 --- a/docs/source/features/interaction/window.md +++ b/docs/source/features/interaction/window.md @@ -44,6 +44,52 @@ Recording hotkey registration is controlled by `SimConfig.window_record.enable_h The camera-pose hotkey is controlled by `SimulationManagerCfg.window_camera_pose.enable_hotkey` and prints look-at form by default. Set `SimulationManagerCfg.window_camera_pose.convert_to_look_at=False` to print the raw 4x4 pose matrix instead. The same output can be requested programmatically with `SimulationManager.print_window_camera_pose()`. +### Entity Gizmo Control + +Opening a non-headless `SimulationManager` window enables dexsim's world-owned +`EntityGizmoManipulator` by default: + +```python +import dexsim + +gizmo_config = dexsim.interaction.EntityGizmoConfig() +gizmo_config.max_gizmos = 0 # Unlimited simultaneous bindings. +sim.open_window(entity_gizmo_config=gizmo_config) +``` + +While enabled, left-click a render mesh, dynamic/kinematic rigid body, or +articulation link and press **G** to attach or detach its root gizmo. The +controller supports multiple simultaneous bindings and owns selection, +temporary physics-state changes, and cleanup. No `sim.update_gizmos()` call is +needed for this world-level controller. + +EmbodiChain's built-in `default_plane` is registered as an immovable target and +cannot receive an entity gizmo. Other supported scene entities remain +selectable normally. + +For a view-only window, opt out explicitly: + +```python +sim.open_window(enable_entity_gizmo=False) +``` + +Set `SimulationManagerCfg.enable_entity_gizmo_on_window_open=False` to change +the default for constructor-opened and subsequently opened windows. Headless +simulations do not create or enable the controller. + +`sim.enable_entity_gizmo(config)` can reconfigure or reactivate the controller +at any time, and `sim.disable_entity_gizmo()` cancels it without closing the +window. The last explicit configuration is restored if the window is closed +and reopened. + +Use `sim.get_entity_gizmo()` to access the native controller and +`sim.has_entity_gizmo()` to query its lifecycle state. Closing the window or +destroying the `SimulationManager` disables it automatically. + +This controller is distinct from the target-specific Robot TCP IK gizmo. When +both are active, **G** controls entity roots and **I** shows or hides the Robot +TCP IK gizmo. + ## Customizing Window Events Users can create their own custom window interaction controls by subclassing the `ObjectManipulator` class (provided by `dexsim`). This allows for the implementation of specific behaviors and responses to user inputs. diff --git a/docs/source/tutorial/gizmo.rst b/docs/source/tutorial/gizmo.rst index 6f2a5b7a0..49404a358 100644 --- a/docs/source/tutorial/gizmo.rst +++ b/docs/source/tutorial/gizmo.rst @@ -5,7 +5,7 @@ Interactive Robot Control with Gizmo .. currentmodule:: embodichain.lab.sim -This tutorial demonstrates how to use the Gizmo class for interactive robot manipulation in SimulationManager. You'll learn how to create a gizmo attached to a robot's end-effector and use it for real-time inverse kinematics (IK) control, allowing intuitive manipulation of robot poses through visual interaction. +This tutorial demonstrates how to use the Gizmo class for interactive robot manipulation in SimulationManager. Robot gizmos delegate interactive inverse kinematics (IK) to dexsim's Newton IK controller while all joint-state reads and drive-target writes continue to pass through the EmbodiChain ``Robot`` abstraction. The Code ~~~~~~~~ @@ -28,7 +28,8 @@ Similar to the previous tutorial on robot simulation, we use the :class:`Simulat -**Important:** Gizmo only supports single environment mode (`num_envs=1`). Using multiple environments will raise an exception. +**Important:** The target-specific Robot TCP, rigid-object, and Camera +``Gizmo`` wrapper supports only single-environment mode (``num_envs=1``). All gizmo creation, visibility, and destruction operations must be managed via the SimulationManager API: @@ -57,21 +58,34 @@ The :class:`objects.Gizmo` class provides a unified interface for interactive co Setting up Robot Configuration ------------------------------ -First, we configure a UR10 robot with an IK solver for end-effector control: +First, configure a UR10 robot and its controllable arm joints: .. literalinclude:: ../../../scripts/tutorials/sim/gizmo_robot.py :language: python - :start-at: # Create UR10 robot configuration + :start-at: # Create UR10 robot :end-at: robot = sim.add_robot(cfg=robot_cfg) Key components of the robot configuration: - **URDF Configuration**: Loads the robot's kinematic and visual model - **Control Parts**: Defines which joints can be controlled (``"Joint[1-6]"`` for UR10) -- **IK Solver**: :class:`solvers.PinkSolverCfg` provides inverse kinematics capabilities - **Drive Properties**: Sets stiffness and damping for joint control -The IK solver is crucial for gizmo functionality, as it enables the robot to automatically calculate joint angles needed to reach gizmo target positions. +An EmbodiChain kinematics solver is not required by the gizmo. The IK chain is declared when enabling it: + +.. code-block:: python + + gizmo_cfg = GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ) + sim.enable_gizmo( + uid="ur10_gizmo_test", + control_part="arm", + gizmo_cfg=gizmo_cfg, + ) + +For existing robot configurations, these link names and the TCP can instead be inherited from the selected control part's configured EmbodiChain solver. The solver supplies metadata only; interactive IK is still performed by dexsim's ``NewtonChainIK``. Creating and Attaching a Gizmo ------------------------------- @@ -83,7 +97,14 @@ After configuring the robot, enable the gizmo for interactive control using the .. code-block:: python # Enable gizmo for the robot's arm - sim.enable_gizmo(uid="ur10_gizmo_test", control_part="arm") + sim.enable_gizmo( + uid="ur10_gizmo_test", + control_part="arm", + gizmo_cfg=GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ), + ) if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): logger.log_error("Failed to enable gizmo!") return @@ -102,22 +123,23 @@ The Gizmo system will automatically: 1. **Detect Target Type**: Identify that the target is a robot (vs. rigid object or camera) 2. **Find End-Effector**: Locate the robot's end-effector link (``ee_link`` for UR10) -3. **Create Proxy Object**: Generate a small invisible cube at the end-effector position -4. **Set Up IK Callback**: Configure the gizmo to trigger IK solving when moved +3. **Build Newton IK Chain**: Build a reduced start-link-to-end-link model from the robot URDF +4. **Bind dexsim Controller**: Attach ``IKGizmoController`` directly to the articulation adapter How Gizmo-Robot Interaction Works ---------------------------------- -The gizmo-robot interaction follows this efficient workflow: +The gizmo-robot interaction follows this workflow: -1. **Gizmo Callback**: When the user drags the gizmo, a callback function updates the proxy object's transform -2. **Deferred IK Solving**: Instead of solving IK immediately in the callback (which causes UI lag), the target transform is stored -3. **Update Loop**: During each simulation step, ``gizmo.update()`` solves IK and applies joint commands -4. **Robot Motion**: The robot smoothly moves to follow the gizmo position +1. **Target Update**: Dragging the dexsim target gizmo updates the Newton IK target state +2. **Deferred Solve**: ``sim.update_gizmos()`` asks ``IKGizmoController`` to solve only when the target changed +3. **State Bridge**: The adapter reads the selected EmbodiChain control-part joints as the solve seed +4. **Drive Target**: The solved positions are written through ``Robot.set_qpos(..., target=True)`` so CPU and CUDA state paths stay synchronized +5. **Robot Motion**: Joint drives move the robot toward the target without teleporting its current state -This design separates UI responsiveness from computational IK solving, ensuring smooth interaction even with complex robots. +Robot gizmos no longer create or maintain an EmbodiChain proxy cube. Camera gizmos continue to use their existing proxy path, and rigid-object gizmos continue to follow the selected object directly. The Simulation Loop ------------------- @@ -163,19 +185,57 @@ Gizmo Lifecycle Management Gizmo lifecycle is managed by SimulationManager: -- Enable: `sim.enable_gizmo(...)` +- Enable a target-specific gizmo: `sim.enable_gizmo(...)` - Update: Main loop automatically calls `sim.update_gizmos()` - Destroy/disable: `sim.disable_gizmo(...)` or `sim.destroy()` (recommended) There is no need to manually create or destroy Gizmo instances. All resources are managed by SimulationManager. +World-Level Entity Gizmo +------------------------ + +For selection-based root manipulation, opening a non-headless window enables +dexsim's world-level entity gizmo by default: + +.. code-block:: python + + import dexsim + + config = dexsim.interaction.EntityGizmoConfig() + config.max_gizmos = 0 + sim.open_window(entity_gizmo_config=config) + + # Left-click an entity and press G to attach or detach a gizmo. + # Multiple entities may remain attached. + + sim.disable_entity_gizmo() + +This path supports render meshes, eligible rigid bodies, and articulation +roots. dexsim owns raycast selection, temporary body-state changes, multiple +bindings, and cleanup. It requires no ``sim.update_gizmos()`` call. + +EmbodiChain's built-in ``default_plane`` is excluded from manipulation. +Selecting it and pressing **G** does not create a gizmo. + +Use ``sim.open_window(enable_entity_gizmo=False)`` for a view-only window, or +set ``SimulationManagerCfg.enable_entity_gizmo_on_window_open=False`` to +change the default. Headless simulations do not create the controller. + +``sim.get_entity_gizmo()`` returns the native +``EntityGizmoManipulator`` and ``sim.has_entity_gizmo()`` reports whether it is +enabled. Closing the window or destroying the simulation also disables it. + +The Robot end-effector controller remains target-specific because it solves a +TCP pose rather than editing the articulation root. By default, **G** controls +the entity gizmo and **I** toggles Robot TCP IK gizmo visibility. + Available Gizmo Methods ----------------------- -If you need to access the underlying Gizmo instance (via `sim.get_gizmo`), you can use the following methods: +If you need to access the underlying Gizmo instance (via `sim.get_gizmo`), you can use the following methods. For robot targets these methods operate on dexsim's IK target gizmo: **Transform Control:** @@ -245,8 +305,8 @@ Tips and Best Practices **Robot compatibility:** -- Ensure your robot is configured with a correct IK solver -- Check the end-effector (EE) link name +- Set valid ``ik_root_link_name`` and ``ik_end_link_name`` values, or configure an EmbodiChain solver whose chain metadata can be inherited +- Set ``ik_tcp_pose`` when the desired tool center point differs from the end-link frame - Test joint limits and workspace boundaries @@ -254,7 +314,7 @@ Tips and Best Practices **Visualization customization:** - Adjust gizmo appearance via Gizmo config (e.g., ``set_line_width()``; requires access to the instance via `sim.get_gizmo`) -- Adjust gizmo scale according to robot size +- Adjust robot target size with ``GizmoCfg.ik_gizmo_scale`` - Enable collision for debugging if needed Next Steps @@ -265,6 +325,6 @@ After mastering basic gizmo usage, you can explore: - **Multi-robot Gizmos**: Attach gizmos to multiple robots simultaneously - **Custom Gizmo Callbacks**: Implement application-specific interaction logic - **Gizmo with Rigid Objects**: Use gizmos for interactive object manipulation -- **Advanced IK Configuration**: Fine-tune solver parameters for specific robots +- **Advanced IK Configuration**: Tune ``GizmoCfg.ik_iterations``, ``ik_device``, and the TCP pose -For more advanced robot control and simulation features, refer to the complete :doc:`robot` tutorial and the API documentation for :class:`objects.Gizmo` and :class:`solvers.PinkSolverCfg`. +For more advanced robot control and simulation features, refer to the complete :doc:`robot` tutorial and the API documentation for :class:`objects.Gizmo`. diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index eb0dd7f58..b30257294 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -1735,8 +1735,8 @@ class RobotCfg(ArticulationCfg): If no control part is specified, the robot will use all joints as a single control part. Note: - - if `control_parts` is specified, `solver_cfg` must be a dict with part names as - keys corresponding to the control parts name. + - `control_parts` can be used without `solver_cfg`. If `solver_cfg` is a + dictionary, its keys must correspond to control-part names. - The joint names in the control parts support regular expressions, e.g., 'joint[1-6]'. After initialization of robot, the names will be expanded to a list of full joint names. - `Robot` is a derived class of `Articulation`, with control parts support. So the `drive_pros` diff --git a/embodichain/lab/sim/objects/__init__.py b/embodichain/lab/sim/objects/__init__.py index 2254f1008..06014aef1 100644 --- a/embodichain/lab/sim/objects/__init__.py +++ b/embodichain/lab/sim/objects/__init__.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + from ..common import BatchEntity from .rigid_object import RigidObject, RigidBodyData, RigidObjectCfg from .rigid_object_group import ( @@ -26,7 +28,7 @@ from .articulation import Articulation, ArticulationData, ArticulationCfg from .robot import Robot, RobotCfg from .light import Light, LightCfg -from .gizmo import Gizmo +from .gizmo import Gizmo, GizmoCfg from .constraint import RigidConstraint diff --git a/embodichain/lab/sim/objects/gizmo.py b/embodichain/lab/sim/objects/gizmo.py index 9fb370c83..42ef89da6 100644 --- a/embodichain/lab/sim/objects/gizmo.py +++ b/embodichain/lab/sim/objects/gizmo.py @@ -13,71 +13,110 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -""" -Gizmo: A reusable controller for interactive manipulation of simulation elements (object, robot, camera, etc.) -""" +"""Interactive gizmos for simulation objects, robots, and cameras.""" -import numpy as np -import torch -import dexsim -from typing import Callable -from scipy.spatial.transform import Rotation as R +from __future__ import annotations -from embodichain.lab.sim.common import BatchEntity -from embodichain.lab.sim.objects import RigidObject, Robot -from embodichain.lab.sim.sensors import Camera -from embodichain.utils import configclass, logger +from collections.abc import Callable +from typing import TYPE_CHECKING, Any +import dexsim +import numpy as np +import torch +import warp as wp from dexsim.types import ( - AxisOption, - RotationRingsOption, AxisArrowType, AxisCornerType, + AxisOption, AxisTagType, - TransformMask, - ActorType, - RigidBodyShape, - PhysicalAttr, + InputKey, + RotationRingsOption, ) +from scipy.spatial.transform import Rotation as R +from embodichain.lab.sim.common import BatchEntity +from embodichain.lab.sim.objects.rigid_object import RigidObject +from embodichain.lab.sim.objects.robot import Robot +from embodichain.lab.sim.sensors import Camera from embodichain.lab.sim.utility.gizmo_utils import create_gizmo_callback +from embodichain.utils import configclass, logger + +if TYPE_CHECKING: + from dexsim.kit.ik import IKGizmoController, NewtonChainIK + +__all__ = ["Gizmo", "GizmoCfg"] @configclass class GizmoCfg: - """Configuration class for Gizmo parameters. + """Configure gizmo appearance and robot Newton IK behavior.""" - This class defines the visual and interaction parameters for gizmo controllers, - including axis appearance and rotation rings settings. - """ - - # Axis configuration axis_length_x: float = 0.2 - """Length of X-axis arrow.""" + """Length of the X-axis arrow.""" + axis_length_y: float = 0.2 - """Length of Y-axis arrow.""" + """Length of the Y-axis arrow.""" + axis_length_z: float = 0.2 - """Length of Z-axis arrow.""" + """Length of the Z-axis arrow.""" + axis_size: float = 0.01 - """Thickness of axis lines.""" + """Thickness of the axis lines.""" + arrow_type: AxisArrowType = AxisArrowType.CONE """Type of arrow head.""" + corner_type: AxisCornerType = AxisCornerType.SPHERE """Type of axis corner.""" + tag_type: AxisTagType = AxisTagType.PLANE """Type of axis label.""" - # Rotation rings configuration rings_radius: float = 0.15 - """Radius of rotation rings.""" + """Radius of the rotation rings.""" + rings_size: float = 0.01 - """Thickness of rotation rings.""" + """Thickness of the rotation rings.""" + + ik_root_link_name: str | None = None + """Robot IK chain root link. + + When omitted, the value is read from the selected control part's configured + EmbodiChain solver. + """ + + ik_end_link_name: str | None = None + """Robot IK chain end link. + + When omitted, the value is read from the selected control part's configured + EmbodiChain solver. + """ - def to_options_dict(self) -> dict: - """Convert configuration to options dictionary format expected by gizmo creation. + ik_tcp_pose: torch.Tensor | np.ndarray | list[list[float]] | None = None + """End-link-to-TCP transform for robot IK. + + When omitted, the configured EmbodiChain solver TCP is used if available; + otherwise the identity transform is used. + """ + + ik_iterations: int = 24 + """Number of Newton IK iterations per changed target.""" + + ik_device: str | None = None + """Warp device for the Newton IK model, or the robot device when omitted.""" + + ik_gizmo_scale: float = 1.5 + """Isotropic scale of dexsim's robot IK target gizmo.""" + + ik_toggle_key: InputKey = InputKey.SCANCODE_I + """Window key used by dexsim to toggle the robot IK gizmo.""" + + def to_options_dict(self) -> dict[str, object]: + """Convert the visual configuration to dexsim gizmo options. Returns: - Dictionary containing AxisOption and RotationRingsOption objects. + The axis and rotation-ring options used by rigid-object and camera + gizmos. """ return { "axis": AxisOption( @@ -90,19 +129,125 @@ def to_options_dict(self) -> dict: tag_type=self.tag_type, ), "rings": RotationRingsOption( - radius=self.rings_radius, size=self.rings_size + radius=self.rings_radius, + size=self.rings_size, ), } +class _RobotGizmoAdapter: + """Expose one EmbodiChain robot control part to dexsim's IK controller.""" + + def __init__(self, robot: Robot, control_part: str, env_id: int = 0) -> None: + """Create the adapter. + + Args: + robot: EmbodiChain robot whose state is synchronized. + control_part: Robot control part driven by the IK solution. + env_id: Environment instance exposed to the interactive controller. + + Raises: + ValueError: If the control part, environment, or joint selection is + invalid. + """ + if not robot.control_parts or control_part not in robot.control_parts: + raise ValueError( + f"Control part {control_part!r} is not defined. Available parts: " + f"{list(robot.control_parts or {})}." + ) + if env_id < 0 or env_id >= robot.num_instances: + raise ValueError( + f"Robot gizmo env_id={env_id} is outside [0, {robot.num_instances})." + ) + + joint_ids = robot.get_joint_ids(control_part, remove_mimic=True) + if not joint_ids: + raise ValueError( + f"Control part {control_part!r} has no non-mimic active joints." + ) + + self.robot = robot + self.control_part = control_part + self.env_id = env_id + self.joint_ids = list(joint_ids) + self.joint_names = [robot.joint_names[index] for index in self.joint_ids] + + def get_current_qpos(self) -> np.ndarray: + """Return current selected joint positions in dexsim joint-name order.""" + return self._selected_qpos(target=False) + + def get_target_qpos(self) -> np.ndarray: + """Return target selected joint positions in dexsim joint-name order.""" + return self._selected_qpos(target=True) + + def set_current_qpos(self, qpos: np.ndarray) -> None: + """Write selected current positions through the EmbodiChain abstraction.""" + self._set_qpos(qpos, target=False) + + def set_target_qpos(self, qpos: np.ndarray) -> None: + """Write selected drive targets through the EmbodiChain abstraction.""" + self._set_qpos(qpos, target=True) + + def get_actived_joint_names(self) -> list[str]: + """Return selected active joint names using dexsim's API spelling.""" + return self.joint_names.copy() + + def get_world_pose(self) -> np.ndarray: + """Return the selected robot instance's root pose as a matrix.""" + pose = self.robot.get_local_pose(to_matrix=True)[self.env_id] + return pose.detach().cpu().numpy().astype(np.float32, copy=True) + + def get_link_names(self, include_fixed: bool = True) -> list[str]: + """Return all runtime link names. + + Args: + include_fixed: Kept for compatibility with the dexsim articulation + API. EmbodiChain's link list already includes fixed links. + """ + del include_fixed + return list(self.robot.link_names) + + def get_link_pose(self, link_name: str) -> np.ndarray: + """Return one runtime link pose as a world-space matrix.""" + pose = self.robot.get_link_pose( + link_name, + env_ids=[self.env_id], + to_matrix=True, + )[0] + return pose.detach().cpu().numpy().astype(np.float32, copy=True) + + def _selected_qpos(self, target: bool) -> np.ndarray: + qpos = self.robot.get_qpos(target=target)[self.env_id, self.joint_ids] + return qpos.detach().cpu().numpy().astype(np.float32, copy=True) + + def _set_qpos(self, qpos: np.ndarray, target: bool) -> None: + values = np.asarray(qpos, dtype=np.float32) + if values.shape != (len(self.joint_ids),): + raise ValueError( + f"Expected qpos shape ({len(self.joint_ids)},), got {values.shape}." + ) + self.robot.set_qpos( + qpos=torch.as_tensor( + values, + dtype=torch.float32, + device=self.robot.device, + ).unsqueeze(0), + joint_ids=self.joint_ids, + env_ids=[self.env_id], + target=target, + ) + + class Gizmo: - """ - Generic Gizmo controller for simulation elements. - Supports RigidObject, Robot, and Camera with type-specific handling. + """Control one rigid object, robot end effector, or camera interactively. + + Robot targets use dexsim's :class:`IKGizmoController` and + :class:`NewtonChainIK`. Rigid-object and camera behavior remains on the + existing direct/proxy paths. - Note: - Gizmo can only be used in single environment mode (num_envs=1). - Will raise RuntimeError if used with multiple environments. + .. attention:: + Gizmos currently expose only one environment instance. Create them only + when ``num_envs=1``. """ def __init__( @@ -110,441 +255,426 @@ def __init__( target: BatchEntity, cfg: GizmoCfg | None = None, control_part: str | None = "arm", - ): - """ + ) -> None: + """Create and attach a gizmo. + Args: - target: The simulation element to control (RigidObject, Robot, or Camera) - cfg: Gizmo configuration parameters (optional, uses default if None) - control_part: For robots, specifies which control part to use (optional, default: "arm") + target: Simulation element to control. + cfg: Gizmo appearance and robot IK configuration. + control_part: Robot control part. When omitted, the first configured + part is selected. """ - self.target = target - self._target_type = self._detect_target_type(target) - self._control_part = control_part - self._env = dexsim.default_world().get_env() - self._windows = dexsim.default_world().get_windows() + world = dexsim.default_world() + if world is None: + raise RuntimeError("A dexsim world must exist before creating a gizmo.") - # Check if running in single environment (num_env must be 1) - num_envs = dexsim.get_world_num() - if num_envs > 1: + self.cfg = cfg if cfg is not None else GizmoCfg() + self._world = world + self._env = world.get_env() + self._control_part = control_part + self._callback: Callable[..., Any] | None = None + self._state = "active" + self._is_visible = True + self._gizmo: object | None = None + self._proxy_cube: object | None = None + self._pending_target_transform: torch.Tensor | None = None + self._ik_model: object | None = None + self._ik_solver: NewtonChainIK | None = None + self._ik_controller: IKGizmoController | None = None + self._robot_adapter: _RobotGizmoAdapter | None = None + self.target: BatchEntity | None = None + self._target_type = "" + self._attach_target(target) + + def _attach_target(self, target: BatchEntity) -> None: + num_instances = int(getattr(target, "num_instances", dexsim.get_world_num())) + if num_instances > 1: raise RuntimeError( - f"Gizmo can only be used in single environment mode (num_env=1), " - f"but current num_envs={num_envs}. Please create simulation with num_envs=1." + "Gizmo can only be used in single environment mode " + f"(num_envs=1), but target has {num_instances} instances." ) - # Use provided config or get default - if cfg is None: - cfg = self._get_default_cfg() - self.cfg = cfg + self.target = target + self._target_type = self._detect_target_type(target) + if self._target_type == "robot": + self._setup_robot_gizmo() + return + self._gizmo = self._create_gizmo(self.cfg) - self._callback = None - self._state = "active" - self._setup_gizmo_follow() + if self._target_type == "rigidobject": + self._setup_rigid_object_gizmo() + else: + self._setup_camera_gizmo() - def _detect_target_type(self, target: BatchEntity) -> str: - """Detect target type: 'rigidobject', 'robot', or 'camera' using isinstance only.""" - if Robot is not None and isinstance(target, Robot): + @staticmethod + def _detect_target_type(target: BatchEntity) -> str: + if isinstance(target, Robot): return "robot" - if Camera is not None and isinstance(target, Camera): + if isinstance(target, Camera): return "camera" - if RigidObject is not None and isinstance(target, RigidObject): + if isinstance(target, RigidObject): return "rigidobject" - raise ValueError( - f"Unsupported target type: {type(target)}. Only RigidObject, Robot, and Camera are supported." + f"Unsupported target type: {type(target)}. Only RigidObject, Robot, " + "and Camera are supported." ) - def _get_default_cfg(self) -> GizmoCfg: - """Get default gizmo configuration (same for all target types)""" - return GizmoCfg() - - def _create_gizmo(self, cfg: GizmoCfg): - """Create gizmo using configuration object""" + def _create_gizmo(self, cfg: GizmoCfg) -> object: options = cfg.to_options_dict() - axis = options["axis"] - rings = options["rings"] - return self._env.create_gizmo(axis, rings) - - def _compute_ee_pose_fk(self): - """Compute end-effector pose using forward kinematics""" - # Get current joint positions for this arm - proprioception = self.target.get_proprioception() - current_qpos_full = proprioception["qpos"] - current_joint_ids = self.target.get_joint_ids(self._robot_arm_name) - - joint_positions = current_qpos_full[:, current_joint_ids] - if joint_positions.dim() > 1: - joint_positions = joint_positions[0] - - # Compute forward kinematics - ee_pose = self.target.compute_fk( - joint_positions, name=self._control_part, to_matrix=True - ) + return self._env.create_gizmo(options["axis"], options["rings"]) - return ee_pose + def _setup_rigid_object_gizmo(self) -> None: + target = self._require_target() + target_node = target._entities[0].node + self._require_gizmo().follow(target_node) + self._require_gizmo().set_flush_localpose_callback(create_gizmo_callback()) - def _create_proxy_cube( - self, position: np.ndarray, rotation_matrix: np.ndarray, name: str - ): - """Create a proxy cube for gizmo tracking""" - # Convert rotation matrix to euler angles - euler = R.from_matrix(rotation_matrix).as_euler("xyz", degrees=False) + def _setup_robot_gizmo(self) -> None: + try: + from dexsim.kit.ik import ( + IKApplyMode, + IKGizmoController, + NewtonChainIK, + build_newton_model_from_urdf, + ) + except ImportError as error: + raise RuntimeError( + "Robot gizmo requires a dexsim build that exports " + "IKGizmoController, NewtonChainIK, and " + "build_newton_model_from_urdf." + ) from error + + target = self._require_robot() + control_parts = list(target.control_parts or {}) + if not control_parts: + raise ValueError("Robot has no control parts defined.") + if self._control_part is None: + self._control_part = control_parts[0] + if self._control_part not in control_parts: + raise ValueError( + f"Control part {self._control_part!r} is not defined. Available " + f"parts: {control_parts}." + ) - # Create small proxy cube at specified position - proxy_cube = self._env.create_cube(0.02, 0.02, 0.02) # 2cm cube - proxy_cube.set_location(position[0], position[1], position[2]) - proxy_cube.set_rotation_euler(euler[0], euler[1], euler[2]) + root_link, end_link, tcp_pose = self._resolve_robot_ik_chain(target) + if self.cfg.ik_iterations <= 0: + raise ValueError("ik_iterations must be greater than zero.") + if not np.isfinite(self.cfg.ik_gizmo_scale) or self.cfg.ik_gizmo_scale <= 0: + raise ValueError("ik_gizmo_scale must be positive and finite.") + + adapter = _RobotGizmoAdapter(target, self._control_part) + ik_device = self.cfg.ik_device or str(target.device) + with wp.ScopedDevice(ik_device): + ik_model = build_newton_model_from_urdf( + target.cfg.fpath, + hide_visuals=True, + ) + ik_solver = NewtonChainIK( + ik_model, + start_link=root_link, + end_link=end_link, + iterations=self.cfg.ik_iterations, + tcp_pose=tcp_pose, + ) - # Connect gizmo to proxy cube. - self._gizmo.follow(proxy_cube.node) + current_qpos = adapter.get_current_qpos() + ik_solver.set_qpos_from_joint_names( + adapter.get_actived_joint_names(), + current_qpos, + ) + base_pose = adapter.get_world_pose() + ik_solver.sync_target_state_from_link(adapter, base_pose) + + target_name = getattr(target.cfg, "uid", "robot") + ik_controller = IKGizmoController( + self._world, + adapter, + ik_solver, + base_state={"pose": base_pose}, + toggle_key=self.cfg.ik_toggle_key, + follow_robot_base=True, + apply_mode=IKApplyMode.DRIVE_TARGET, + gizmo_scale=self.cfg.ik_gizmo_scale, + name=f"{target_name}_{self._control_part}_ik", + ) - logger.log_info(f"{name} gizmo proxy created at position: {position}") - return proxy_cube + self._robot_adapter = adapter + self._ik_model = ik_model + self._ik_solver = ik_solver + self._ik_controller = ik_controller + self._gizmo = ik_controller.target_gizmo.gizmo + logger.log_info( + f"Robot gizmo uses dexsim Newton IK for control part " + f"{self._control_part!r} ({root_link} -> {end_link})." + ) - def _setup_camera_gizmo(self): - """Setup gizmo for Camera by creating a proxy RigidObject at camera position""" - # Get current camera pose - camera_pose = self.target.get_local_pose(to_matrix=True)[0] # Get first camera - camera_pos = camera_pose[:3, 3].cpu().numpy() - camera_rot_matrix = camera_pose[:3, :3].cpu().numpy() + def _resolve_robot_ik_chain( + self, + target: Robot, + ) -> tuple[str, str, np.ndarray]: + solver = ( + target.get_solver(self._control_part) + if target.cfg.solver_cfg is not None + else None + ) + root_link = self.cfg.ik_root_link_name or getattr( + solver, + "root_link_name", + None, + ) + end_link = self.cfg.ik_end_link_name or getattr( + solver, + "end_link_name", + None, + ) + if not root_link or not end_link: + raise ValueError( + "Robot gizmo needs an IK chain. Set GizmoCfg.ik_root_link_name " + "and ik_end_link_name, or configure a solver for the selected " + "robot control part." + ) - # Create proxy cube and set callback + tcp_pose = self.cfg.ik_tcp_pose + if tcp_pose is None and solver is not None: + tcp_pose = solver.get_tcp() + if tcp_pose is None: + tcp_pose = np.eye(4, dtype=np.float32) + if isinstance(tcp_pose, torch.Tensor): + tcp_pose = tcp_pose.detach().cpu().numpy() + return root_link, end_link, np.asarray(tcp_pose, dtype=np.float32) + + def _setup_camera_gizmo(self) -> None: + target = self._require_target() + camera_pose = target.get_local_pose(to_matrix=True)[0] + camera_pos = camera_pose[:3, 3].detach().cpu().numpy() + camera_rotation = camera_pose[:3, :3].detach().cpu().numpy() self._proxy_cube = self._create_proxy_cube( - camera_pos, camera_rot_matrix, "Camera" + camera_pos, + camera_rotation, + "Camera", ) - # New API uses set_flush_localpose_callback - self._gizmo.set_flush_localpose_callback(self._proxy_gizmo_callback) + self._require_gizmo().set_flush_localpose_callback(self._proxy_gizmo_callback) - def _proxy_gizmo_callback(self, *args): - """Generic callback for proxy-based gizmo. + def _create_proxy_cube( + self, + position: np.ndarray, + rotation_matrix: np.ndarray, + name: str, + ) -> object: + euler = R.from_matrix(rotation_matrix).as_euler("xyz", degrees=False) + proxy_cube = self._env.create_cube(0.02, 0.02, 0.02) + proxy_cube.set_location(*position) + proxy_cube.set_rotation_euler(*euler) + self._require_gizmo().follow(proxy_cube.node) + logger.log_info(f"{name} gizmo proxy created at position: {position}.") + return proxy_cube - Supports both old signature: (node, translation, rotation, flag) - and new signature: (node, local_pose, flag) where local_pose is a 4x4 matrix. - Updates the proxy cube transform and sets `_pending_target_transform`. - """ - # New API callback signature: (node, local_pose, flag) - if len(args) != 3: + def _proxy_gizmo_callback(self, *args: object) -> None: + if len(args) != 3 or self._proxy_cube is None: return node, local_pose, flag = args if node is None: return - # Check if proxy cube still exists - if not hasattr(self, "_proxy_cube") or self._proxy_cube is None: - return - - # convert to numpy 4x4 matrix if isinstance(local_pose, torch.Tensor): - lp = local_pose.cpu().numpy() + pose = local_pose.detach().cpu().numpy() else: - lp = np.asarray(local_pose) - - if lp.shape != (4, 4): + pose = np.asarray(local_pose) + if pose.shape != (4, 4): return - trans = lp[:3, 3] - rot_mat = lp[:3, :3] - euler = R.from_matrix(rot_mat).as_euler("xyz", degrees=False) - - self._proxy_cube.set_location(float(trans[0]), float(trans[1]), float(trans[2])) - self._proxy_cube.set_rotation_euler( - float(euler[0]), float(euler[1]), float(euler[2]) - ) - - # Build pending target transform (1,4,4) - target_transform = torch.eye(4, dtype=torch.float32) - target_transform[:3, 3] = torch.tensor( - [trans[0], trans[1], trans[2]], dtype=torch.float32 - ) - target_transform[:3, :3] = torch.tensor(rot_mat, dtype=torch.float32) - self._pending_target_transform = target_transform.unsqueeze(0) - - def _update_camera_pose(self, target_transform: torch.Tensor): - """Update camera pose to match target transform""" + node.set_transform(pose, flag) + position = pose[:3, 3] + euler = R.from_matrix(pose[:3, :3]).as_euler("xyz", degrees=False) + self._proxy_cube.set_location(*position) + self._proxy_cube.set_rotation_euler(*euler) + self._pending_target_transform = torch.as_tensor( + pose, + dtype=torch.float32, + ).unsqueeze(0) + + def _update_camera_pose(self, target_transform: torch.Tensor) -> bool: try: - # Set camera pose using set_local_pose method - self.target.set_local_pose(target_transform) + self._require_target().set_local_pose(target_transform) return True - except Exception as e: - logger.log_error(f"Error updating camera pose: {e}") + except Exception as error: + logger.log_error(f"Error updating camera pose: {error}") return False - def _setup_robot_gizmo(self): - """Setup gizmo for Robot by creating a proxy RigidObject at end-effector""" - # Get end-effector pose using specified control part - if self.target.cfg.solver_cfg is None: - raise ValueError( - "Robot has no solver configured for IK/FK computations for gizmo" - ) - - arm_names = list(self.target.control_parts.keys()) - if not arm_names: - raise ValueError("Robot has no control parts defined") + def attach(self, target: BatchEntity) -> None: + """Attach this controller to another supported single-instance target.""" + self._release_resources() + self._attach_target(target) - # Use specified control part or fall back to first available - if self._control_part and self._control_part in arm_names: - self._robot_arm_name = self._control_part - else: - logger.log_error(f"Control part '{self._control_part}' not found.") - - logger.log_info(f"Using control part: {self._robot_arm_name}") - - # Get end-effector pose using forward kinematics - ee_pose = self._compute_ee_pose_fk()[0] # remove batch dimension - - ee_pos = ee_pose[:3, 3].cpu().numpy() - ee_rot_matrix = ee_pose[:3, :3].cpu().numpy() - - # Create proxy cube and set callback (use new callback API) - self._proxy_cube = self._create_proxy_cube(ee_pos, ee_rot_matrix, "Robot") - self._gizmo.set_flush_localpose_callback(self._proxy_gizmo_callback) - - def _update_robot_ik(self, target_transform: torch.Tensor): - """Update robot joints using IK to reach target transform""" - try: - # Get current joint positions as seed using proprioception - proprioception = self.target.get_proprioception() - current_qpos_full = proprioception["qpos"] # Full joint positions - - # Get joint IDs for this arm - current_joint_ids = self.target.get_joint_ids(self._robot_arm_name) - - # Extract joint positions for this specific arm - if len(current_joint_ids) > 0: - joint_seed = current_qpos_full[ - :, current_joint_ids - ] # Select arm joints - if joint_seed.dim() > 1: - joint_seed = joint_seed[0] # Take first batch element - else: - logger.log_warning( - f"No joint IDs found for arm: {self._robot_arm_name}" - ) - return False - - # Solve IK - ik_success, new_qpos = self.target.compute_ik( - pose=target_transform, name=self._robot_arm_name, joint_seed=joint_seed - ) - - if ik_success: - # Ensure correct dimensions for setting qpos - # new_qpos from IK solver may be (1, N, dof) or (N, dof), flatten to (dof,) for single env - if new_qpos.dim() > 1: - new_qpos = new_qpos.squeeze() # Remove all singleton dimensions - if new_qpos.dim() == 1: - new_qpos = new_qpos.unsqueeze(0) # Make it (1, dof) for set_qpos - - # Update robot joint positions - self.target.set_qpos(qpos=new_qpos[0], joint_ids=current_joint_ids) - return True - else: - logger.log_warning("IK solution not found") - return False - - except Exception as e: - logger.log_error(f"Error in robot IK: {e}") - return False - - def _setup_gizmo_follow(self): - """Setup gizmo based on target type""" - if self._target_type == "rigidobject": - # RigidObject: direct node access through MeshObject — use follow/attach - tgt_node = self.target._entities[0].node - self._gizmo.follow(tgt_node) - # set callback (localpose-style) - self._gizmo.set_flush_localpose_callback(create_gizmo_callback()) - - elif self._target_type == "robot": - # Robot: create proxy object at end-effector position - self._setup_robot_gizmo() - elif self._target_type == "camera": - # Camera: create proxy object at camera position - self._setup_camera_gizmo() - - def attach(self, target: BatchEntity): - """Attach gizmo to a new simulation element.""" - self.target = target - self._target_type = self._detect_target_type(target) - self._setup_gizmo_follow() - - def detach(self): - """Detach gizmo from current element.""" + def detach(self) -> None: + """Detach the gizmo and release target-specific controller resources.""" + self._release_resources() self.target = None - # Detach gizmo using new API - self._gizmo.detach_parent() + self._target_type = "" - def set_transform_callback(self, callback: Callable): - """Set callback for gizmo transform events (translation/rotation).""" + def set_transform_callback(self, callback: Callable[..., Any]) -> None: + """Set an additional raw gizmo transform callback.""" self._callback = callback - self._gizmo.set_transform_flush_callback(callback) + self._require_gizmo().set_transform_flush_callback(callback) - def set_world_pose(self, pose): - """Set gizmo's world pose.""" - self._gizmo.set_world_pose(pose) + def set_world_pose(self, pose: np.ndarray) -> None: + """Set the underlying gizmo's world pose.""" + self._require_gizmo().set_world_pose(pose) - def set_local_pose(self, pose): - """Set gizmo's local pose.""" - self._gizmo.set_local_pose(pose) + def set_local_pose(self, pose: np.ndarray) -> None: + """Set the underlying gizmo's local pose.""" + self._require_gizmo().set_local_pose(pose) - def set_line_width(self, width: float): - """Set gizmo line width.""" - self._gizmo.set_line_width(width) + def set_line_width(self, width: float) -> None: + """Set the underlying gizmo line width.""" + self._require_gizmo().set_line_width(width) - def enable_collision(self, enabled: bool): + def enable_collision(self, enabled: bool) -> None: """Enable or disable gizmo collision.""" - self._gizmo.enable_collision(enabled) + self._require_gizmo().enable_collision(enabled) - def get_world_pose(self): - """Get gizmo's world pose.""" - return self._gizmo.get_world_pose() + def get_world_pose(self) -> np.ndarray: + """Return the underlying gizmo's world pose.""" + return self._require_gizmo().get_world_pose() - def get_local_pose(self): - """Get gizmo's local pose.""" - return self._gizmo.get_local_pose() + def get_local_pose(self) -> np.ndarray: + """Return the underlying gizmo's local pose.""" + return self._require_gizmo().get_local_pose() - def get_name(self): - """Get gizmo node name.""" - return self._gizmo.get_name() + def get_name(self) -> str: + """Return the underlying gizmo name.""" + return self._require_gizmo().get_name() - def get_parent(self): - """Get gizmo's parent node.""" - return self._gizmo.get_parent() + def get_parent(self) -> object: + """Return the underlying gizmo parent.""" + return self._require_gizmo().get_parent() def toggle_visibility(self) -> bool: - """ - Toggle the visibility of the gizmo. - - Returns: - bool: The new visibility state (True = visible, False = hidden) - """ - if not hasattr(self, "_is_visible"): - self._is_visible = True # Default to visible - - # Toggle the state - self._is_visible = not self._is_visible - - # Apply the visibility setting to the gizmo node - if self._gizmo: - self._gizmo.set_visible(self._is_visible) - - return self._is_visible - - def set_visible(self, visible: bool): - """ - Set the visibility of the gizmo. - - Args: - visible (bool): True to show, False to hide the gizmo - """ - self._is_visible = visible - - # Apply the visibility setting to the gizmo node - if self._gizmo: - self._gizmo.set_visible(self._is_visible) + """Toggle gizmo visibility and return the new state.""" + visible = not self.is_visible() + self.set_visible(visible) + return visible + + def set_visible(self, visible: bool) -> None: + """Set gizmo visibility.""" + self._is_visible = bool(visible) + if self._ik_controller is not None: + self._ik_controller.enabled = self._is_visible + gizmo = self._gizmo + if gizmo is not None: + gizmo.set_visible(self._is_visible) def is_visible(self) -> bool: - """ - Check if the gizmo is currently visible. - - Returns: - bool: True if visible, False if hidden - """ - return getattr(self, "_is_visible", True) + """Return whether the gizmo is visible.""" + if self._ik_controller is not None: + return bool(self._ik_controller.enabled) + return self._is_visible - def update(self): - """Synchronize gizmo with target's current transform, and handle IK solving here.""" + def update(self) -> None: + """Synchronize the gizmo and apply pending target changes.""" + if self.target is None: + return if self._target_type == "rigidobject": - tgt_node = self.target._entities[0].node - self._gizmo.follow(tgt_node) - + target_node = self.target._entities[0].node + self._require_gizmo().follow(target_node) elif self._target_type == "robot": - # If there is a pending target, solve IK and clear it - if ( - hasattr(self, "_pending_target_transform") - and self._pending_target_transform is not None - ): - self._update_robot_ik(self._pending_target_transform) - self._pending_target_transform = None + if self._ik_controller is not None: + self._ik_controller.update(iterations=self.cfg.ik_iterations) elif self._target_type == "camera": - # Update proxy cube position to match current camera pose - if hasattr(self, "_proxy_cube") and self._proxy_cube: + if self._proxy_cube is not None: camera_pose = self.target.get_local_pose(to_matrix=True)[0] - camera_pos = camera_pose[:3, 3].cpu().numpy() - self._proxy_cube.set_location( - camera_pos[0], camera_pos[1], camera_pos[2] - ) - - # If there is a pending camera target, update camera pose and clear it - if ( - hasattr(self, "_pending_target_transform") - and self._pending_target_transform is not None - ): + position = camera_pose[:3, 3].detach().cpu().numpy() + self._proxy_cube.set_location(*position) + if self._pending_target_transform is not None: self._update_camera_pose(self._pending_target_transform) self._pending_target_transform = None - def apply_transform(self, translation, rotation): - """Apply transform based on target type""" + def apply_transform( + self, + translation: np.ndarray, + rotation: np.ndarray, + ) -> None: + """Apply a direct transform where the target path supports it.""" + if self.target is None: + return if self._target_type == "rigidobject": self.target.set_location(*translation) self.target.set_rotation_euler(*rotation) - elif self._target_type == "robot": - # Robot transforms are handled by IK in the gizmo callback - if hasattr(self, "_proxy_cube") and self._proxy_cube: - self._proxy_cube.set_location(*translation) - self._proxy_cube.set_rotation_euler(*rotation) - elif self._target_type == "camera": - # Camera transforms are handled by pose update in the gizmo callback - if hasattr(self, "_proxy_cube") and self._proxy_cube: - self._proxy_cube.set_location(*translation) - self._proxy_cube.set_rotation_euler(*rotation) - else: - # Other target types - pass + elif self._target_type == "camera" and self._proxy_cube is not None: + self._proxy_cube.set_location(*translation) + self._proxy_cube.set_rotation_euler(*rotation) - def destroy(self): - """Clean up gizmo resources and release references.""" - # Clear transform callback first to avoid bad_function_call - if hasattr(self, "_gizmo") and self._gizmo and hasattr(self._gizmo, "node"): + def destroy(self) -> None: + """Release gizmo resources and target references.""" + self._release_resources() + self.target = None + self._target_type = "" + + def _release_resources(self) -> None: + gizmo = self._gizmo + if gizmo is not None: + for method_name in ( + "set_flush_localpose_callback", + "set_transform_flush_callback", + ): + method = getattr(gizmo, method_name, None) + if callable(method): + try: + method(None) + except (TypeError, RuntimeError): + pass try: - # Clear transform callback before any other cleanup - self._gizmo.node.set_flush_transform_callback(None) - logger.log_info("Cleared gizmo transform callback") - except Exception as e: - logger.log_warning(f"Failed to clear gizmo callback: {e}") - - # Remove proxy cube if exists (before detaching gizmo) - if hasattr(self, "_proxy_cube") and self._proxy_cube: + gizmo.set_visible(False) + except (AttributeError, TypeError, RuntimeError): + pass try: - # Detach gizmo from proxy cube first - if ( - hasattr(self, "_gizmo") - and self._gizmo - and hasattr(self._gizmo, "node") - ): - self._gizmo.detach_parent() - # Then remove the proxy cube - self._env.remove_actor(self._proxy_cube) - logger.log_info("Successfully removed proxy cube from environment") - except Exception as e: - logger.log_warning(f"Failed to remove proxy cube: {e}") - self._proxy_cube = None + gizmo.detach_parent() + except (AttributeError, TypeError, RuntimeError): + pass + + if self._ik_controller is not None: + target_node = self._ik_controller.target_gizmo.target_node + try: + target_node.detach_parent() + except (AttributeError, TypeError, RuntimeError): + pass - # Final gizmo cleanup - if hasattr(self, "_gizmo") and self._gizmo and hasattr(self._gizmo, "node"): + if self._proxy_cube is not None: try: - # Ensure detach_parent is called if not done above - if self._target_type in ["robot", "camera"]: - pass # Already detached above - else: - self._gizmo.node.detach_parent() - logger.log_info("Successfully cleaned up gizmo node") - except Exception as e: - logger.log_warning(f"Failed to cleanup gizmo node: {e}") - - # Clear pending transform - if hasattr(self, "_pending_target_transform"): - self._pending_target_transform = None - - # Directly release references + self._env.remove_actor(self._proxy_cube) + except (AttributeError, TypeError, RuntimeError) as error: + logger.log_warning(f"Failed to remove gizmo proxy cube: {error}") + + if gizmo is not None: + remove_gizmo = getattr(self._env, "remove_gizmo", None) + if callable(remove_gizmo): + try: + remove_gizmo(gizmo) + except (AttributeError, TypeError, RuntimeError) as error: + logger.log_warning( + f"Failed to remove gizmo from dexsim environment: {error}" + ) + + self._pending_target_transform = None + self._proxy_cube = None self._gizmo = None - self.target = None + self._ik_controller = None + self._ik_solver = None + self._ik_model = None + self._robot_adapter = None + + def _require_gizmo(self) -> object: + if self._gizmo is None: + raise RuntimeError("Gizmo is not attached.") + return self._gizmo + + def _require_target(self) -> BatchEntity: + if self.target is None: + raise RuntimeError("Gizmo has no target.") + return self.target + + def _require_robot(self) -> Robot: + target = self._require_target() + if not isinstance(target, Robot): + raise TypeError(f"Expected Robot target, got {type(target)}.") + return target diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 86be05455..65a65adaf 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -32,7 +32,7 @@ from copy import deepcopy from datetime import datetime from functools import cached_property -from typing import Callable, Dict, List, Sequence, Union +from typing import TYPE_CHECKING, Callable, Dict, List, Sequence, Union from dataclasses import dataclass, asdict, field, MISSING # Global cache directories @@ -66,7 +66,7 @@ Light, RigidConstraint, ) -from embodichain.lab.sim.objects.gizmo import Gizmo +from embodichain.lab.sim.objects.gizmo import Gizmo, GizmoCfg from embodichain.lab.sim.sensors import ( SensorCfg, BaseSensor, @@ -94,6 +94,9 @@ from embodichain.utils import configclass, logger from embodichain.utils.math import look_at_to_pose, pose_inv +if TYPE_CHECKING: + from dexsim.interaction import EntityGizmoConfig, EntityGizmoManipulator + __all__ = [ "SimulationManager", "SimulationManagerCfg", @@ -159,6 +162,9 @@ class SimulationManagerCfg: window_camera_pose: WindowCameraPoseCfg = field(default_factory=WindowCameraPoseCfg) """Interactive viewer camera-pose printing settings.""" + enable_entity_gizmo_on_window_open: bool = True + """Whether opening a viewer window automatically enables entity gizmo control.""" + @dataclass class _WindowRecordState: @@ -200,6 +206,7 @@ class SimulationManager: _instances = {} _cleanup_queue: queue.Queue = queue.Queue() + _DEFAULT_PLANE_GIZMO_TARGET_ID = (1 << 64) - 1 SUPPORTED_SENSOR_TYPES = { "Camera": Camera, @@ -249,6 +256,7 @@ def __init__( self._world: dexsim.World = dexsim.World(world_config) self._window: Windows | None = None + self._entity_gizmo_config: EntityGizmoConfig | None = None self._window_record_state: _WindowRecordState | None = None self._window_record_camera: object | None = None wr = sim_config.window_record @@ -326,6 +334,7 @@ def __init__( if sim_config.headless is False: self._window = self._world.get_windows() + self._on_window_opened() @classmethod def get_instance(cls, instance_id: int = 0) -> SimulationManager: @@ -616,11 +625,43 @@ def get_env(self, arena_index: int = -1) -> dexsim.environment.Arena: def get_world(self) -> dexsim.World: return self._world - def open_window(self) -> None: - """Open the simulation window.""" - self._world.open_window() + def open_window( + self, + *, + enable_entity_gizmo: bool | None = None, + entity_gizmo_config: EntityGizmoConfig | None = None, + ) -> None: + """Open the simulation window and initialize its interaction controls. + + Entity gizmo control is enabled by default. Set + ``enable_entity_gizmo=False`` for a view-only window. When the argument + is omitted, :attr:`SimulationManagerCfg.enable_entity_gizmo_on_window_open` + determines the behavior. + + Args: + enable_entity_gizmo: Whether to enable world-level entity gizmo + control for this window. ``None`` uses the simulation + configuration default. + entity_gizmo_config: Optional native dexsim configuration. Passing + a configuration implies entity gizmo control unless explicitly + disabled. + """ + if not self.is_window_opened or self._window is None: + self._world.open_window() self._window = self._world.get_windows() + self.is_window_opened = True + self._on_window_opened( + enable_entity_gizmo=enable_entity_gizmo, + entity_gizmo_config=entity_gizmo_config, + ) + def _on_window_opened( + self, + *, + enable_entity_gizmo: bool | None = None, + entity_gizmo_config: EntityGizmoConfig | None = None, + ) -> None: + """Initialize controls shared by constructor-opened and reopened windows.""" if ( self._window_record_hotkey_cfg is not None and self._window_record_input_control is None @@ -631,10 +672,31 @@ def open_window(self) -> None: and self._window_camera_pose_input_control is None ): self.enable_window_camera_pose_hotkey(**self._window_camera_pose_hotkey_cfg) - self.is_window_opened = True + + if enable_entity_gizmo is None: + enable_entity_gizmo = entity_gizmo_config is not None or getattr( + self.sim_config, + "enable_entity_gizmo_on_window_open", + True, + ) + + try: + if enable_entity_gizmo: + if entity_gizmo_config is not None: + self.enable_entity_gizmo(entity_gizmo_config) + elif not self.has_entity_gizmo(): + self.enable_entity_gizmo(self._entity_gizmo_config) + elif self.has_entity_gizmo(): + self.disable_entity_gizmo() + except RuntimeError as error: + logger.log_warning( + f"Entity gizmo control could not be initialized for the window: {error}" + ) def close_window(self) -> None: """Close the simulation window.""" + if self.has_entity_gizmo(): + self.disable_entity_gizmo() if self.is_window_recording(): self.stop_window_record() self._world.close_window() @@ -1620,15 +1682,140 @@ def get_robot_uid_list(self) -> List[str]: """ return list(self._robots.keys()) + def enable_entity_gizmo( + self, + config: EntityGizmoConfig | None = None, + ) -> EntityGizmoManipulator: + """Enable dexsim's world-level entity gizmo controller. + + This is a thin lifecycle wrapper around + :meth:`dexsim.World.enable_entity_gizmo`. The returned controller owns + window selection, hotkey handling, multiple gizmo bindings, temporary + physics-state changes, and rigid-body/articulation-root manipulation. + + Args: + config: Native dexsim entity-gizmo configuration. When omitted, + dexsim's defaults are used. + + Returns: + The world-owned dexsim entity gizmo manipulator. + + Raises: + RuntimeError: If the installed dexsim build does not provide entity + gizmo support or fails to create the controller. + """ + world = getattr(self, "_world", None) + enable = getattr(world, "enable_entity_gizmo", None) + if not callable(enable): + raise RuntimeError( + "The installed dexsim build does not provide " + "World.enable_entity_gizmo()." + ) + + controller = enable() if config is None else enable(config) + if controller is None: + raise RuntimeError("dexsim failed to enable the entity gizmo controller.") + self._exclude_default_plane_from_entity_gizmo(controller) + self._entity_gizmo_config = config + logger.log_info("Dexsim entity gizmo control enabled.") + return controller + + def _exclude_default_plane_from_entity_gizmo( + self, + controller: EntityGizmoManipulator, + ) -> None: + """Register the EmbodiChain ground as an immovable gizmo target. + + dexsim resolves registered external targets before its generic + render-mesh path. Registering the default plane as a static rigid body + therefore makes both raycast toggles and programmatic attachment return + ``STATIC_RIGID_BODY`` without adding physics to the visual plane. + """ + default_plane = getattr(self, "_default_plane", None) + register = getattr(controller, "register_external_target", None) + if default_plane is None: + return + if not callable(register): + logger.log_warning( + "The installed dexsim build cannot exclude the default plane " + "from entity gizmo control." + ) + return + + try: + result = register( + self._DEFAULT_PLANE_GIZMO_TARGET_ID, + dexsim.interaction.EntityGizmoTargetType.RIGID_BODY, + default_plane, + ActorType.STATIC, + ) + except (AttributeError, TypeError, RuntimeError) as error: + logger.log_warning( + "Failed to exclude the default plane from entity gizmo " + f"control: {error}." + ) + return + if result != dexsim.interaction.EntityGizmoResult.SUCCESS: + logger.log_warning( + "Failed to exclude the default plane from entity gizmo " + f"control: {result}." + ) + + def disable_entity_gizmo(self) -> bool: + """Disable dexsim's world-level entity gizmo controller. + + Returns: + ``True`` when an active controller was disabled, or ``False`` when + entity gizmo control was already disabled. + + Raises: + RuntimeError: If the installed dexsim build does not provide entity + gizmo support. + """ + world = getattr(self, "_world", None) + get_controller = getattr(world, "get_entity_gizmo", None) + disable = getattr(world, "disable_entity_gizmo", None) + if not callable(get_controller) or not callable(disable): + raise RuntimeError( + "The installed dexsim build does not provide entity gizmo " + "lifecycle APIs." + ) + if get_controller() is None: + return False + + disable() + logger.log_info("Dexsim entity gizmo control disabled.") + return True + + def get_entity_gizmo(self) -> EntityGizmoManipulator | None: + """Return dexsim's active world-level entity gizmo controller.""" + world = getattr(self, "_world", None) + get_controller = getattr(world, "get_entity_gizmo", None) + if not callable(get_controller): + raise RuntimeError( + "The installed dexsim build does not provide " + "World.get_entity_gizmo()." + ) + return get_controller() + + def has_entity_gizmo(self) -> bool: + """Return whether world-level entity gizmo control is enabled.""" + world = getattr(self, "_world", None) + get_controller = getattr(world, "get_entity_gizmo", None) + return callable(get_controller) and get_controller() is not None + def enable_gizmo( - self, uid: str, control_part: str | None = None, gizmo_cfg: object = None - ) -> Gizmo: + self, + uid: str, + control_part: str | None = None, + gizmo_cfg: GizmoCfg | None = None, + ) -> Gizmo | None: """Enable gizmo control for any simulation object (Robot, RigidObject, Camera, etc.). Args: uid (str): UID of the object to attach gizmo to (searches in robots, rigid_objects, sensors, etc.) control_part (str | None, optional): Control part name for robots. Defaults to "arm". - gizmo_cfg (object, optional): Gizmo configuration object. Defaults to None. + gizmo_cfg: Gizmo configuration. Defaults to None. """ # Create gizmo key combining uid and control_part gizmo_key = f"{uid}:{control_part}" if control_part else uid @@ -1638,7 +1825,7 @@ def enable_gizmo( logger.log_warning( f"Gizmo for '{uid}' with control_part '{control_part}' already exists." ) - return + return self._gizmos[gizmo_key] # Search for target object in different collections target = None @@ -1656,17 +1843,13 @@ def enable_gizmo( else: logger.log_error( - f"Object with uid '{uid}' not found in any collection (robots, rigid_objects, sensors, articulations)." + f"Object with uid '{uid}' not found in any supported collection " + "(robots, rigid_objects, sensors)." ) - return + return None + gizmo = None try: - gizmo = Gizmo(target, gizmo_cfg, control_part) - self._gizmos[gizmo_key] = gizmo - logger.log_info( - f"Gizmo enabled for {object_type} '{uid}' with control_part '{control_part}'" - ) - # Initialize GizmoController if not already done. if not hasattr(self, "_gizmo_controller") or self._gizmo_controller is None: window = ( @@ -1674,9 +1857,17 @@ def enable_gizmo( if hasattr(self._world, "get_windows") else None ) + if window is None: + raise RuntimeError("Gizmo requires a simulation window.") self._gizmo_controller = GizmoController() window.add_input_control(self._gizmo_controller) + gizmo = Gizmo(target, gizmo_cfg, control_part) + self._gizmos[gizmo_key] = gizmo + logger.log_info( + f"Gizmo enabled for {object_type} '{uid}' with control_part '{control_part}'" + ) + except Exception as e: logger.log_error( f"Failed to create gizmo for {object_type} '{uid}' with control_part '{control_part}': {e}" @@ -2699,6 +2890,9 @@ def destroy(self, exit_process: bool | None = None) -> None: def _deferred_destroy(self) -> None: """Destroy all simulated assets and release resources.""" + if self.has_entity_gizmo(): + self.disable_entity_gizmo() + # Clean up all gizmos before destroying the simulation for uid in list(self._gizmos.keys()): self.disable_gizmo(uid) diff --git a/embodichain/lab/sim/utility/gizmo_utils.py b/embodichain/lab/sim/utility/gizmo_utils.py index 3ff1c7de7..177612023 100644 --- a/embodichain/lab/sim/utility/gizmo_utils.py +++ b/embodichain/lab/sim/utility/gizmo_utils.py @@ -14,31 +14,33 @@ # limitations under the License. # ---------------------------------------------------------------------------- -""" -Gizmo utility functions for EmbodiSim. +"""Gizmo utility functions for EmbodiChain. This module provides utility functions for creating gizmo transform callbacks. """ -from typing import Callable -from typing import TYPE_CHECKING -from dexsim.types import TransformMask +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from embodichain.lab.sim.objects import Robot +__all__ = ["create_gizmo_callback", "run_gizmo_robot_control_loop"] + -def create_gizmo_callback() -> Callable: +def create_gizmo_callback() -> Callable[[Any, Any, Any], None]: """Create a standard gizmo transform callback function. This callback handles local pose for gizmo controls. It applies transformations directly to the node when gizmo controls are manipulated. Returns: - Callable: A callback function that can be used with gizmo.node.set_flush_transform_callback() + A callback compatible with dexsim's gizmo local-pose flush hook. """ - def gizmo_transform_callback(node, local_pose, flag): + def gizmo_transform_callback(node: Any, local_pose: Any, flag: Any) -> None: if node is not None: node.set_transform(local_pose, flag) @@ -46,8 +48,10 @@ def gizmo_transform_callback(node, local_pose, flag): def run_gizmo_robot_control_loop( - robot: object | str, control_part: str = "arm", end_link_name: str | None = None -): + robot: Robot | str, + control_part: str = "arm", + end_link_name: str | None = None, +) -> None: """Run a control loop for testing gizmo controls on a robot. This function implements a control loop that allows users to manipulate a robot @@ -75,38 +79,62 @@ def run_gizmo_robot_control_loop( np.set_printoptions(precision=5, suppress=True) from embodichain.lab.sim import SimulationManager - from embodichain.lab.sim.objects import Robot - from embodichain.lab.sim.solvers import PinkSolverCfg + from embodichain.lab.sim.objects import GizmoCfg - from embodichain.utils.logger import log_info, log_warning, log_error + from embodichain.utils.logger import log_error, log_info sim = SimulationManager.get_instance() if isinstance(robot, str): - robot = sim.get_robot(uid=robot) + robot_uid = robot + robot = sim.get_robot(uid=robot_uid) + if robot is None: + log_error(f"Robot {robot_uid!r} was not found.") + return # Enter auto-update mode. sim.set_manual_update(False) - # Replace robot's default solver with PinkSolver for gizmo control. - robot_solver = robot.get_solver(name=control_part) + # Resolve only the chain metadata. dexsim owns the Newton IK solver and + # writes its drive targets back through the EmbodiChain Robot API. + robot_solver = ( + robot.get_solver(name=control_part) + if robot.cfg.solver_cfg is not None + else None + ) control_part_link_names = robot.get_control_part_link_names(name=control_part) + if not control_part_link_names: + raise ValueError(f"Control part {control_part!r} has no links.") + root_link_name = ( + robot_solver.root_link_name + if robot_solver is not None + else control_part_link_names[0] + ) end_link_name = ( - control_part_link_names[-1] if end_link_name is None else end_link_name + ( + robot_solver.end_link_name + if robot_solver is not None + else control_part_link_names[-1] + ) + if end_link_name is None + else end_link_name ) - pink_solver_cfg = PinkSolverCfg( - urdf_path=robot.cfg.fpath, - end_link_name=end_link_name, - root_link_name=robot_solver.root_link_name, - pos_eps=1e-2, - rot_eps=5e-2, - max_iterations=300, - dt=0.1, + tcp_pose = robot_solver.get_tcp() if robot_solver is not None else None + gizmo_cfg = GizmoCfg( + ik_root_link_name=root_link_name, + ik_end_link_name=end_link_name, + ik_tcp_pose=tcp_pose, ) - robot.init_solver(cfg={control_part: pink_solver_cfg}) # Enable gizmo for the robot - gizmo = sim.enable_gizmo(uid=robot.uid, control_part=control_part) + gizmo = sim.enable_gizmo( + uid=robot.uid, + control_part=control_part, + gizmo_cfg=gizmo_cfg, + ) + if gizmo is None: + log_error(f"Failed to enable gizmo for control part {control_part!r}.") + return # Store initial robot configuration initial_qpos = robot.get_qpos(name=control_part) @@ -127,7 +155,7 @@ def run_gizmo_robot_control_loop( old_settings = termios.tcgetattr(sys.stdin) tty.setcbreak(sys.stdin.fileno()) - def get_key(): + def get_key() -> str | None: """Non-blocking keyboard input.""" if select.select([sys.stdin], [], [], 0)[0]: return sys.stdin.read(1) @@ -146,29 +174,20 @@ def get_key(): if key in ["q", "Q", "\x1b"]: # Q or ESC log_info("Exiting gizmo control loop...") sim.disable_gizmo(uid=robot.uid, control_part=control_part) - if robot_solver: - robot.init_solver( - cfg={control_part: robot_solver.cfg} - ) # Restore original solver break # Print robot state elif key in ["p", "P"]: current_qpos = robot.get_qpos(name=control_part) - eef_pose = robot.compute_fk(name=control_part, qpos=current_qpos) + eef_pose = robot.get_link_pose(end_link_name, to_matrix=True) + if tcp_pose is not None: + tcp_tensor = np.asarray(tcp_pose, dtype=np.float32) + eef_pose = eef_pose @ eef_pose.new_tensor(tcp_tensor) log_info(f"\n=== Robot State ===") log_info(f"Control part: {control_part}") log_info(f"Joint positions: {current_qpos.squeeze().tolist()}") - log_info(f"End-effector pose:\n{eef_pose.squeeze().numpy()}") - - if eef_pose is None: - log_info( - "End-effector pose unavailable: compute_fk returned None " - f"for control part '{control_part}'." - ) - else: - eef_pose_np = eef_pose.detach().cpu().numpy().squeeze() - log_info(f"End-effector pose:\n{eef_pose_np}") + eef_pose_np = eef_pose.detach().cpu().numpy().squeeze() + log_info(f"End-effector pose:\n{eef_pose_np}") elif key in ["g", "G"]: if gizmo_visible: sim.set_gizmo_visibility( @@ -189,7 +208,11 @@ def get_key(): sim.disable_gizmo(uid=robot.uid, control_part=control_part) robot.clear_dynamics() robot.set_qpos(qpos=initial_qpos, name=control_part, target=False) - sim.enable_gizmo(uid=robot.uid, control_part=control_part) + sim.enable_gizmo( + uid=robot.uid, + control_part=control_part, + gizmo_cfg=gizmo_cfg, + ) log_info("Robot reset to initial pose") # Print info @@ -206,10 +229,6 @@ def get_key(): except KeyboardInterrupt: sim.disable_gizmo(uid=robot.uid, control_part=control_part) - if robot_solver: - robot.init_solver( - cfg={control_part: robot_solver.cfg} - ) # Restore original solver log_info("\nControl loop interrupted by user (Ctrl+C)") finally: diff --git a/examples/sim/gizmo/gizmo_object.py b/examples/sim/gizmo/gizmo_object.py index b0931f241..62a9f123b 100644 --- a/examples/sim/gizmo/gizmo_object.py +++ b/examples/sim/gizmo/gizmo_object.py @@ -14,14 +14,15 @@ # limitations under the License. # ---------------------------------------------------------------------------- -""" -This script demonstrates how to create a simulation scene using SimulationManager. -It shows the basic setup of simulation context, adding objects, and sensors. -""" +"""Manipulate raycast-selected entities with dexsim's world-level gizmo.""" + +from __future__ import annotations import argparse import time +import dexsim + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg from embodichain.lab.sim.shapes import CubeCfg @@ -60,7 +61,7 @@ def main(): cfg=RigidObjectCfg( uid="cube1", shape=CubeCfg(size=[0.1, 0.1, 0.1]), - body_type="kinematic", + body_type="dynamic", attrs=RigidBodyAttributesCfg( mass=1.0, dynamic_friction=0.5, @@ -85,24 +86,20 @@ def main(): ) ) - # Enable Gizmo for both cubes using the new API (only in window mode) + # Opening a window enables the world-level controller by default. Passing + # a config here reconfigures it for unlimited simultaneous bindings. if not args.headless: - sim.enable_gizmo(uid="cube1") - sim.enable_gizmo(uid="cube2") + gizmo_config = dexsim.interaction.EntityGizmoConfig() + gizmo_config.max_gizmos = 0 + sim.open_window(entity_gizmo_config=gizmo_config) logger.log_info("Scene setup complete!") logger.log_info(f"Running simulation with 1 environment(s)") if not args.headless: - if sim.has_gizmo("cube1"): - logger.log_info("Gizmo enabled for cube1 - you can drag it around!") - if sim.has_gizmo("cube2"): - logger.log_info("Gizmo enabled for cube2 - you can drag it around!") + logger.log_info("Left-click an entity and press G to attach/detach its gizmo.") + logger.log_info("Multiple selected entities can keep gizmos simultaneously.") logger.log_info("Press Ctrl+C to stop the simulation") - # Open window when the scene has been set up - if not args.headless: - sim.open_window() - # Run the simulation run_simulation(sim) @@ -113,22 +110,19 @@ def run_simulation(sim: SimulationManager): sim.init_gpu_physics() step_count = 0 - gizmo_enabled = True + gizmo_enabled = sim.has_entity_gizmo() try: last_time = time.time() last_step = 0 while True: sim.update(step=1) - # Update all gizmos if any are enabled - sim.update_gizmos() - step_count += 1 - # Disable gizmo after 200000 steps (example) + # Demonstrate programmatic cancellation after 200000 steps. if step_count == 200000 and gizmo_enabled: - logger.log_info("Disabling gizmo at step 200000") - sim.disable_gizmo("cube") + logger.log_info("Disabling entity gizmo control at step 200000") + sim.disable_entity_gizmo() gizmo_enabled = False # Print FPS every second @@ -146,6 +140,8 @@ def run_simulation(sim: SimulationManager): except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: + if sim.has_entity_gizmo(): + sim.disable_entity_gizmo() sim.destroy() logger.log_info("Simulation terminated successfully") diff --git a/examples/sim/gizmo/gizmo_robot.py b/examples/sim/gizmo/gizmo_robot.py index 6b8b9effb..06efff88a 100644 --- a/examples/sim/gizmo/gizmo_robot.py +++ b/examples/sim/gizmo/gizmo_robot.py @@ -13,9 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -""" -Gizmo-Robot Example: Test Gizmo class on a robot (UR10) -""" +"""Control a UR10 end effector with dexsim's Newton IK gizmo.""" + +from __future__ import annotations import time import torch @@ -23,15 +23,14 @@ import argparse from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.solvers import PytorchSolverCfg from embodichain.lab.sim.cfg import ( RenderCfg, RobotCfg, URDFCfg, JointDrivePropertiesCfg, ) +from embodichain.lab.sim.objects import GizmoCfg from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.solvers import PinkSolverCfg from embodichain.data import get_data_path from embodichain.utils import logger @@ -75,19 +74,6 @@ def main(): "arm": ["JOINT[0-9]"], "hand": ["FINGER[1-2]"], }, - solver_cfg={ - "arm": PytorchSolverCfg( - end_link_name="ee_link", - root_link_name="base_link", - tcp=[ - [0.0, 1.0, 0.0, 0.0], - [-1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.12], - [0.0, 0.0, 0.0, 1.0], - ], - num_samples=30, - ) - }, drive_pros=JointDrivePropertiesCfg( stiffness={"JOINT[0-9]": 1e4, "FINGER[1-2]": 1e2}, damping={"JOINT[0-9]": 1e3, "FINGER[1-2]": 1e1}, @@ -109,8 +95,23 @@ def main(): time.sleep(0.2) # Wait for a moment to ensure everything is set up - # Enable gizmo using the new API - sim.enable_gizmo(uid="ur10_gizmo_test", control_part="arm") + # The robot needs no EmbodiChain IK solver for interactive gizmo control. + # dexsim builds and owns the Newton IK chain from this metadata. + gizmo_cfg = GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ik_tcp_pose=[ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.12], + [0.0, 0.0, 0.0, 1.0], + ], + ) + sim.enable_gizmo( + uid="ur10_gizmo_test", + control_part="arm", + gizmo_cfg=gizmo_cfg, + ) if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): logger.log_error("Failed to enable gizmo!") return @@ -119,6 +120,7 @@ def main(): logger.log_info("Gizmo-Robot example started!") logger.log_info("Use the gizmo to drag the robot end-effector (EE)") + logger.log_info("Press I to show or hide the Robot TCP IK gizmo") logger.log_info("Press Ctrl+C to stop the simulation") run_simulation(sim) diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index 6d6613f9a..7d8432079 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -13,9 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -""" -Gizmo-Robot Example: Test Gizmo class on a robot (UR10) -""" +"""Control a UR10 end effector with dexsim's Newton IK gizmo.""" + +from __future__ import annotations import time import torch @@ -30,8 +30,7 @@ URDFCfg, JointDrivePropertiesCfg, ) - -from embodichain.lab.sim.solvers import PinkSolverCfg +from embodichain.lab.sim.objects import GizmoCfg from embodichain.data import get_data_path from embodichain.utils import logger @@ -68,17 +67,6 @@ def main(): components=[{"component_type": "arm", "urdf_path": urdf_path}] ), control_parts={"arm": ["Joint[1-6]"]}, - solver_cfg={ - "arm": PinkSolverCfg( - urdf_path=urdf_path, - end_link_name="ee_link", - root_link_name="base_link", - pos_eps=1e-2, - rot_eps=5e-2, - max_iterations=300, - dt=0.1, - ) - }, drive_pros=JointDrivePropertiesCfg( stiffness={"Joint[1-6]": 1e4}, damping={"Joint[1-6]": 1e3}, @@ -97,8 +85,15 @@ def main(): time.sleep(0.2) # Wait for a moment to ensure everything is set up - # Enable gizmo using the new API - sim.enable_gizmo(uid="ur10_gizmo_test", control_part="arm") + # dexsim owns the Newton IK solver used by the interactive controller. + sim.enable_gizmo( + uid="ur10_gizmo_test", + control_part="arm", + gizmo_cfg=GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ), + ) if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): logger.log_error("Failed to enable gizmo!") return @@ -107,6 +102,7 @@ def main(): logger.log_info("Gizmo-Robot example started!") logger.log_info("Use the gizmo to drag the robot end-effector (EE)") + logger.log_info("Press I to show or hide the Robot TCP IK gizmo") logger.log_info("Press Ctrl+C to stop the simulation") run_simulation(sim) diff --git a/tests/sim/objects/test_gizmo.py b/tests/sim/objects/test_gizmo.py new file mode 100644 index 000000000..faddcc717 --- /dev/null +++ b/tests/sim/objects/test_gizmo.py @@ -0,0 +1,189 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from embodichain.lab.sim.objects.gizmo import ( + Gizmo, + GizmoCfg, + _RobotGizmoAdapter, +) + + +class _FakeRobot: + """Small Robot-compatible state holder for adapter tests.""" + + def __init__(self) -> None: + self.control_parts = {"arm": ["joint_a", "joint_mimic", "joint_b"]} + self.num_instances = 1 + self.joint_names = ["joint_a", "joint_mimic", "joint_b"] + self.link_names = ["base_link", "tool_link"] + self.device = torch.device("cpu") + self.cfg = SimpleNamespace(solver_cfg=None) + self.current_qpos = torch.tensor([[0.1, 0.2, 0.3]], dtype=torch.float32) + self.target_qpos = torch.tensor([[0.4, 0.5, 0.6]], dtype=torch.float32) + self.write_calls: list[dict[str, object]] = [] + + def get_joint_ids( + self, + name: str, + remove_mimic: bool = False, + ) -> list[int]: + assert name == "arm" + return [0, 2] if remove_mimic else [0, 1, 2] + + def get_qpos(self, target: bool = False) -> torch.Tensor: + return self.target_qpos if target else self.current_qpos + + def set_qpos(self, **kwargs: object) -> None: + self.write_calls.append(kwargs) + + def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: + assert to_matrix + return torch.eye(4, dtype=torch.float32).unsqueeze(0) + + def get_link_pose( + self, + link_name: str, + env_ids: list[int], + to_matrix: bool = False, + ) -> torch.Tensor: + assert link_name in self.link_names + assert env_ids == [0] + assert to_matrix + pose = torch.eye(4, dtype=torch.float32) + pose[2, 3] = 0.8 + return pose.unsqueeze(0) + + +def test_robot_adapter_synchronizes_selected_joint_state() -> None: + robot = _FakeRobot() + adapter = _RobotGizmoAdapter(robot, "arm") + + assert adapter.get_actived_joint_names() == ["joint_a", "joint_b"] + np.testing.assert_allclose(adapter.get_current_qpos(), [0.1, 0.3]) + np.testing.assert_allclose(adapter.get_target_qpos(), [0.4, 0.6]) + + adapter.set_target_qpos(np.array([0.7, 0.9], dtype=np.float32)) + + write = robot.write_calls[-1] + assert write["joint_ids"] == [0, 2] + assert write["env_ids"] == [0] + assert write["target"] is True + torch.testing.assert_close( + write["qpos"], + torch.tensor([[0.7, 0.9]], dtype=torch.float32), + ) + + +def test_robot_adapter_reads_root_and_link_pose_through_robot() -> None: + adapter = _RobotGizmoAdapter(_FakeRobot(), "arm") + + np.testing.assert_allclose(adapter.get_world_pose(), np.eye(4)) + link_pose = adapter.get_link_pose("tool_link") + assert link_pose[2, 3] == pytest.approx(0.8) + assert adapter.get_link_names(True) == ["base_link", "tool_link"] + + +def test_robot_adapter_rejects_wrong_qpos_shape() -> None: + adapter = _RobotGizmoAdapter(_FakeRobot(), "arm") + + with pytest.raises(ValueError, match="Expected qpos shape"): + adapter.set_target_qpos(np.zeros(3, dtype=np.float32)) + + +def test_robot_ik_chain_can_be_configured_without_embodichain_solver() -> None: + gizmo = object.__new__(Gizmo) + gizmo.cfg = GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="tool_link", + ) + gizmo._control_part = "arm" + robot = _FakeRobot() + + root_link, end_link, tcp_pose = gizmo._resolve_robot_ik_chain(robot) + + assert (root_link, end_link) == ("base_link", "tool_link") + np.testing.assert_allclose(tcp_pose, np.eye(4)) + + +def test_robot_update_delegates_to_dexsim_ik_controller() -> None: + calls: list[int] = [] + + class _Controller: + def update(self, *, iterations: int) -> None: + calls.append(iterations) + + gizmo = object.__new__(Gizmo) + gizmo.target = object() + gizmo._target_type = "robot" + gizmo._ik_controller = _Controller() + gizmo.cfg = GizmoCfg(ik_iterations=12) + + gizmo.update() + + assert calls == [12] + + +def test_destroy_removes_gizmo_from_dexsim_environment() -> None: + class _DexsimGizmo: + def __init__(self) -> None: + self.detached = False + + def set_flush_localpose_callback(self, callback: object | None) -> None: + pass + + def set_transform_flush_callback(self, callback: object | None) -> None: + pass + + def set_visible(self, visible: bool) -> None: + pass + + def detach_parent(self) -> None: + self.detached = True + + class _Environment: + def __init__(self) -> None: + self.removed: object | None = None + + def remove_gizmo(self, gizmo: object) -> None: + self.removed = gizmo + + dexsim_gizmo = _DexsimGizmo() + environment = _Environment() + gizmo = object.__new__(Gizmo) + gizmo._env = environment + gizmo._gizmo = dexsim_gizmo + gizmo._proxy_cube = None + gizmo._ik_controller = None + gizmo._ik_solver = None + gizmo._ik_model = None + gizmo._robot_adapter = None + gizmo._pending_target_transform = None + gizmo.target = object() + gizmo._target_type = "rigidobject" + + gizmo.destroy() + + assert environment.removed is dexsim_gizmo + assert dexsim_gizmo.detached is True + assert gizmo._gizmo is None diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index b78a63b24..2f1e994e0 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -19,6 +19,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock +import dexsim import numpy as np import pytest @@ -69,15 +70,61 @@ def add_loop(self, callback, time_step: float) -> str: return "loop_handle" +class FakeEntityGizmo: + """Entity-gizmo stub with external-target registration.""" + + def __init__(self) -> None: + self.active = True + self.external_targets: list[tuple[int, object, object, object]] = [] + + def register_external_target( + self, + target_id: int, + target_type: object, + target: object, + actor_type: object, + ) -> object: + self.external_targets.append((target_id, target_type, target, actor_type)) + return dexsim.interaction.EntityGizmoResult.SUCCESS + + class FakeWorld: """World stub exposing the render-thread loop API.""" def __init__(self) -> None: self.thread_runtime = FakeThreadRuntime() + self.entity_gizmo: object | None = None + self.entity_gizmo_configs: list[object | None] = [] + self.window = SimpleNamespace(add_input_control=lambda control: None) + self.window_open_count = 0 + self.window_closed = False def thread_rt(self) -> FakeThreadRuntime: return self.thread_runtime + def enable_entity_gizmo(self, config: object | None = None) -> object: + self.entity_gizmo_configs.append(config) + self.entity_gizmo = FakeEntityGizmo() + return self.entity_gizmo + + def disable_entity_gizmo(self) -> None: + if self.entity_gizmo is not None: + self.entity_gizmo.active = False + self.entity_gizmo = None + + def get_entity_gizmo(self) -> object | None: + return self.entity_gizmo + + def open_window(self) -> None: + self.window_open_count += 1 + self.window_closed = False + + def get_windows(self) -> object: + return self.window + + def close_window(self) -> None: + self.window_closed = True + class FakeEnv: """Environment stub that creates fake cameras.""" @@ -95,13 +142,24 @@ def _make_sim_manager(window: object | None = None) -> SimulationManager: """Create a minimally initialized simulation manager for recorder tests.""" sim = object.__new__(SimulationManager) sim.instance_id = 0 - sim.sim_config = SimpleNamespace(width=64, height=48) + sim.sim_config = SimpleNamespace( + width=64, + height=48, + enable_entity_gizmo_on_window_open=True, + ) sim._window = window + sim._entity_gizmo_config = None sim._window_record_state = None sim._window_record_camera = None sim._window_record_save_threads = [] + sim._window_record_hotkey_cfg = None + sim._window_camera_pose_hotkey_cfg = None + sim._window_record_input_control = None + sim._window_camera_pose_input_control = None sim._env = FakeEnv() sim._world = FakeWorld() + sim._default_plane = object() + sim.is_window_opened = window is not None return sim @@ -198,6 +256,129 @@ def fake_save_window_record_worker( assert sim._window_record_save_threads == [] +def test_entity_gizmo_lifecycle_delegates_to_dexsim_world() -> None: + sim = _make_sim_manager() + config = object() + + controller = sim.enable_entity_gizmo(config) + + assert controller is sim._world.get_entity_gizmo() + assert sim._world.entity_gizmo_configs == [config] + assert sim.get_entity_gizmo() is controller + assert sim.has_entity_gizmo() is True + assert sim.disable_entity_gizmo() is True + assert controller.active is False + assert sim.has_entity_gizmo() is False + assert sim.disable_entity_gizmo() is False + + +def test_entity_gizmo_registers_default_plane_as_static_exclusion() -> None: + sim = _make_sim_manager() + + controller = sim.enable_entity_gizmo() + + assert controller.external_targets == [ + ( + SimulationManager._DEFAULT_PLANE_GIZMO_TARGET_ID, + dexsim.interaction.EntityGizmoTargetType.RIGID_BODY, + sim._default_plane, + dexsim.types.ActorType.STATIC, + ) + ] + + +def test_open_window_enables_entity_gizmo_by_default() -> None: + sim = _make_sim_manager() + + sim.open_window() + + assert sim.is_window_opened is True + assert sim._world.window_open_count == 1 + assert sim.has_entity_gizmo() is True + assert sim._world.entity_gizmo_configs == [None] + + +def test_open_window_supports_view_only_opt_out() -> None: + sim = _make_sim_manager() + + sim.open_window(enable_entity_gizmo=False) + + assert sim.is_window_opened is True + assert sim.has_entity_gizmo() is False + assert sim._world.entity_gizmo_configs == [] + + +def test_open_window_view_only_opt_out_disables_active_controller() -> None: + sim = _make_sim_manager() + controller = sim.enable_entity_gizmo() + + sim.open_window(enable_entity_gizmo=False) + + assert controller.active is False + assert sim.has_entity_gizmo() is False + + +def test_open_window_respects_configured_entity_gizmo_default() -> None: + sim = _make_sim_manager() + sim.sim_config.enable_entity_gizmo_on_window_open = False + + sim.open_window() + + assert sim.is_window_opened is True + assert sim.has_entity_gizmo() is False + + +def test_open_window_tolerates_dexsim_without_entity_gizmo_api() -> None: + sim = _make_sim_manager() + window = object() + sim._world = SimpleNamespace( + open_window=lambda: None, + get_windows=lambda: window, + ) + + sim.open_window() + + assert sim.is_window_opened is True + assert sim._window is window + assert sim.has_entity_gizmo() is False + + +def test_open_window_preserves_active_entity_gizmo_configuration() -> None: + sim = _make_sim_manager(window=object()) + config = object() + controller = sim.enable_entity_gizmo(config) + + sim.open_window() + + assert sim.get_entity_gizmo() is controller + assert sim._world.entity_gizmo_configs == [config] + assert sim._world.window_open_count == 0 + + +def test_reopened_window_restores_last_entity_gizmo_configuration() -> None: + sim = _make_sim_manager(window=object()) + config = object() + sim.enable_entity_gizmo(config) + sim.close_window() + + sim.open_window() + + assert sim.has_entity_gizmo() is True + assert sim._world.entity_gizmo_configs == [config, config] + + +def test_close_window_disables_entity_gizmo() -> None: + sim = _make_sim_manager(window=object()) + controller = sim.enable_entity_gizmo() + + sim.close_window() + + assert controller.active is False + assert sim.has_entity_gizmo() is False + assert sim._world.window_closed is True + assert sim.is_window_opened is False + + def test_reset_objects_state_includes_soft_and_cloth_assets() -> None: sim = object.__new__(SimulationManager) sim._robots = {} From 8e523c089784313750a44bdd0da7eb856db67d68 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 1 Aug 2026 14:33:56 +0800 Subject: [PATCH 02/14] wip --- docs/source/features/interaction/window.md | 5 + docs/source/tutorial/gizmo.rst | 47 +++- embodichain/lab/sim/objects/gizmo.py | 166 +++++++------ embodichain/lab/sim/sim_manager.py | 64 +++++ embodichain/lab/visualization/__init__.py | 2 + .../lab/visualization/backends/base.py | 15 +- .../lab/visualization/backends/viser.py | 125 ++++++++++ embodichain/lab/visualization/picker.py | 216 +++++++++++++++++ embodichain/lab/visualization/protocol.py | 26 +++ embodichain/lab/visualization/runtime.py | 52 +++++ .../lab/visualization/scene_exporter.py | 20 ++ tests/sim/objects/test_gizmo.py | 89 +++---- tests/sim/test_sim_manager.py | 162 +++++++++++++ tests/visualization/test_picker.py | 218 ++++++++++++++++++ tests/visualization/test_viser_backend.py | 174 ++++++++++++++ 15 files changed, 1244 insertions(+), 137 deletions(-) create mode 100644 embodichain/lab/visualization/picker.py create mode 100644 tests/visualization/test_picker.py diff --git a/docs/source/features/interaction/window.md b/docs/source/features/interaction/window.md index 1275b95c6..6110111a3 100644 --- a/docs/source/features/interaction/window.md +++ b/docs/source/features/interaction/window.md @@ -97,6 +97,11 @@ This controller is distinct from the target-specific Robot TCP IK gizmo. When both are active, **G** controls entity roots and **I** shows or hides the Robot TCP IK gizmo. +The entity gizmo is native-window only. The Viser backend offers an analogous +**click-to-pick** flow (an *Enable click-to-pick Gizmo* checkbox instead of the +**G** hotkey, since browsers do not expose keyboard events); see +:doc:`tutorial/gizmo` for details. + ## Customizing Window Events Users can create their own custom window interaction controls by subclassing the `ObjectManipulator` class (provided by `dexsim`). This allows for the implementation of specific behaviors and responses to user inputs. diff --git a/docs/source/tutorial/gizmo.rst b/docs/source/tutorial/gizmo.rst index 1c0b1b670..41f2bb279 100644 --- a/docs/source/tutorial/gizmo.rst +++ b/docs/source/tutorial/gizmo.rst @@ -43,9 +43,9 @@ All gizmo creation, visibility, and destruction operations must be managed via t Always use the SimulationManager API to control gizmo visibility and lifecycle. Do not operate on the Gizmo instance directly. The same target behavior is available in either the DexSim window or Viser. -Native robot Gizmos use DexSim Newton IK, while headless Viser Gizmos use the -robot's configured EmbodiChain solver. The standard Viser mode includes -interactive Gizmo control: +Robot Gizmos solve IK with DexSim Newton IK in both modes; the only difference +is the input source (a native window gizmo handle vs a Viser transform control). +The standard Viser mode includes interactive Gizmo control: .. code-block:: bash @@ -54,6 +54,27 @@ interactive Gizmo control: Only expose the Viser endpoint to trusted browser clients because dragging a Gizmo mutates simulation targets. +Click-to-Pick in Viser +~~~~~~~~~~~~~~~~~~~~~~ + +Unlike the native DexSim window, the browser does not ray-cast the scene for +you, so EmbodiChain performs the click hit-test against the published scene +geometry. Enable it explicitly in the browser panel: + +1. Toggle the **Enable click-to-pick Gizmo** checkbox under the **Interaction** + folder. +2. Click a rigid object or robot link in the 3D view. A transform control is + attached to it (replacing any previously picked Gizmo); drag it to move the + target. Robot IK is solved with DexSim Newton IK, just as in the native + window. +3. Click empty space, or uncheck the checkbox, to detach the picker-owned + Gizmo. + +The picker manages at most one Gizmo at a time and never touches Gizmos you +created yourself through ``sim.enable_gizmo(...)``. Only rigid objects and +robots are pickable; articulations, soft bodies, and cameras are ignored by the +picker. + What is a Gizmo? ----------------- @@ -83,10 +104,12 @@ Key components of the robot configuration: - **IK Solver**: :class:`solvers.PinkSolverCfg` provides inverse kinematics capabilities - **Drive Properties**: Sets stiffness and damping for joint control -The configured solver drives Viser Gizmos and also provides default chain -metadata to the native controller. A native-only application may instead set -the root link, end link, and optional TCP transform directly in -:class:`objects.GizmoCfg` without configuring an EmbodiChain solver. +The configured EmbodiChain solver is optional: it only supplies default IK chain +metadata (root link, end link, and TCP transform) to the Gizmo. IK itself is +always solved by DexSim Newton IK in both native and Viser modes. A native-only +or Viser-only application may instead set the root link, end link, and optional +TCP transform directly in :class:`objects.GizmoCfg` without configuring an +EmbodiChain solver. Creating and Attaching a Gizmo ------------------------------- @@ -127,7 +150,7 @@ The Gizmo system will automatically: 1. **Detect Target Type**: Identify that the target is a robot (vs. rigid object or camera) 2. **Resolve the IK Chain**: Locate the root and end-effector links -3. **Select the Backend**: Build a native DexSim Newton controller or a headless Viser command path +3. **Select the Backend**: Build a DexSim Newton IK solver; ``enable_native`` only decides whether a native window gizmo handle is created for direct interaction or Viser commands drive the same solver 4. **Defer Simulation Writes**: Apply IK drive targets from the simulation update loop How Gizmo-Robot Interaction Works @@ -138,9 +161,9 @@ How Gizmo-Robot Interaction Works The gizmo-robot interaction follows this workflow: 1. **Target Update**: DexSim or Viser records the requested TCP transform -2. **Deferred Solve**: ``sim.update_gizmos()`` invokes the selected IK backend only when needed -3. **State Bridge**: Native DexSim IK reads and writes the selected EmbodiChain control-part joints through an adapter -4. **Drive Target**: Native solutions use ``Robot.set_qpos(..., target=True)``; Viser solutions use the configured EmbodiChain solver +2. **Deferred Solve**: ``sim.update_gizmos()`` invokes the DexSim Newton IK solver only when needed +3. **State Bridge**: Newton IK reads and writes the selected EmbodiChain control-part joints through an adapter +4. **Drive Target**: Both native and Viser solutions use ``Robot.set_qpos(..., target=True)`` to drive the joint targets 5. **Robot Motion**: Joint drives move the robot toward the target without teleporting its current state Native robot Gizmos do not create an EmbodiChain proxy cube. Camera Gizmos @@ -275,7 +298,7 @@ Tips and Best Practices **Robot compatibility:** -- Ensure your robot is configured with a correct IK solver +- Set the IK chain (root link and end-effector link) in :class:`objects.GizmoCfg`, or configure an EmbodiChain solver to supply them as defaults - Check the end-effector (EE) link name - Test joint limits and workspace boundaries diff --git a/embodichain/lab/sim/objects/gizmo.py b/embodichain/lab/sim/objects/gizmo.py index ebf41372c..43ce04528 100644 --- a/embodichain/lab/sim/objects/gizmo.py +++ b/embodichain/lab/sim/objects/gizmo.py @@ -227,9 +227,10 @@ class Gizmo: target: Simulation element controlled by this Gizmo. cfg: Appearance configuration. control_part: Robot control part used for FK and IK. - enable_native: Whether to create a native DexSim Gizmo. Native robot - Gizmos use DexSim Newton IK; headless controllers use the configured - EmbodiChain solver so Viser commands remain available. + enable_native: Whether to create a native DexSim Gizmo handle. Robot + Gizmos solve IK with DexSim Newton IK in both native and headless + (Viser) modes; ``enable_native`` only controls whether a native + window Gizmo handle is created for direct interaction. """ def __init__( @@ -271,14 +272,15 @@ def __init__( self._ik_solver: NewtonChainIK | None = None self._ik_controller: IKGizmoController | None = None self._robot_adapter: _RobotGizmoAdapter | None = None + self._native_robot_end_link: str | None = None + self._native_robot_tcp_pose: np.ndarray | None = None if self._target_type == "robot": - self._configure_robot(require_solver=not enable_native) + self._configure_robot() + self._setup_robot_ik_solver() if enable_native: self._setup_native_robot_gizmo() - self._desired_target_transform = self._read_native_robot_pose() - else: - self._desired_target_transform = self._read_target_pose() + self._desired_target_transform = self._read_native_robot_pose() else: self._desired_target_transform = self._read_target_pose() @@ -313,11 +315,9 @@ def _detect_target_type(self, target: BatchEntity) -> str: "RigidObject, Robot, or Camera." ) - def _configure_robot(self, *, require_solver: bool) -> None: + def _configure_robot(self) -> None: if self.target is None or not isinstance(self.target, Robot): raise RuntimeError("Robot Gizmo has no attached Robot.") - if require_solver and self.target.cfg.solver_cfg is None: - raise ValueError("Robot has no solver configured for Gizmo IK/FK.") arm_names = list(self.target.control_parts.keys()) if not arm_names: raise ValueError("Robot has no control parts defined.") @@ -332,20 +332,20 @@ def _configure_robot(self, *, require_solver: bool) -> None: f"available parts are {arm_names}." ) - def _setup_native_robot_gizmo(self) -> None: - """Create DexSim's Newton IK controller for a native robot Gizmo.""" + def _setup_robot_ik_solver(self) -> None: + """Build the shared DexSim Newton IK solver and robot adapter. + + The solver is shared by native and headless (Viser) robot Gizmos so both + paths solve IK with DexSim Newton IK instead of an EmbodiChain solver. + Native Gizmos additionally create an :class:`IKGizmoController` in + :meth:`_setup_native_robot_gizmo`. + """ try: - from dexsim.kit.ik import ( - IKApplyMode, - IKGizmoController, - NewtonChainIK, - build_newton_model_from_urdf, - ) + from dexsim.kit.ik import NewtonChainIK, build_newton_model_from_urdf except ImportError as error: raise RuntimeError( "Robot Gizmo requires a DexSim build that exports " - "IKGizmoController, NewtonChainIK, and " - "build_newton_model_from_urdf." + "NewtonChainIK and build_newton_model_from_urdf." ) from error if self.target is None or not isinstance(self.target, Robot): @@ -354,8 +354,6 @@ def _setup_native_robot_gizmo(self) -> None: raise RuntimeError("Robot Gizmo control part is not configured.") if self.cfg.ik_iterations <= 0: raise ValueError("ik_iterations must be greater than zero.") - if not np.isfinite(self.cfg.ik_gizmo_scale) or self.cfg.ik_gizmo_scale <= 0: - raise ValueError("ik_gizmo_scale must be positive and finite.") root_link, end_link, tcp_pose = self._resolve_robot_ik_chain(self.target) adapter = _RobotGizmoAdapter(self.target, self._robot_arm_name) @@ -380,11 +378,39 @@ def _setup_native_robot_gizmo(self) -> None: base_pose = adapter.get_world_pose() ik_solver.sync_target_state_from_link(adapter, base_pose) + self._robot_adapter = adapter + self._ik_model = ik_model + self._ik_solver = ik_solver + self._native_robot_end_link = end_link + self._native_robot_tcp_pose = tcp_pose + logger.log_info( + f"Robot Gizmo uses DexSim Newton IK for control part " + f"{self._robot_arm_name!r} ({root_link} -> {end_link})." + ) + + def _setup_native_robot_gizmo(self) -> None: + """Create DexSim's native IK controller on top of the shared solver.""" + try: + from dexsim.kit.ik import IKApplyMode, IKGizmoController + except ImportError as error: + raise RuntimeError( + "Robot Gizmo requires a DexSim build that exports " + "IKGizmoController and IKApplyMode." + ) from error + + if self._ik_solver is None or self._robot_adapter is None: + raise RuntimeError("Robot Gizmo IK solver is not configured.") + if self.target is None or not isinstance(self.target, Robot): + raise RuntimeError("Robot Gizmo has no attached Robot.") + if not np.isfinite(self.cfg.ik_gizmo_scale) or self.cfg.ik_gizmo_scale <= 0: + raise ValueError("ik_gizmo_scale must be positive and finite.") + + base_pose = self._robot_adapter.get_world_pose() target_name = getattr(self.target.cfg, "uid", "robot") ik_controller = IKGizmoController( self._world, - adapter, - ik_solver, + self._robot_adapter, + self._ik_solver, base_state={"pose": base_pose}, toggle_key=self.cfg.ik_toggle_key, follow_robot_base=True, @@ -393,17 +419,8 @@ def _setup_native_robot_gizmo(self) -> None: name=f"{target_name}_{self._robot_arm_name}_ik", ) - self._robot_adapter = adapter - self._ik_model = ik_model - self._ik_solver = ik_solver self._ik_controller = ik_controller self._gizmo = ik_controller.target_gizmo.gizmo - self._native_robot_end_link = end_link - self._native_robot_tcp_pose = tcp_pose - logger.log_info( - f"Robot Gizmo uses DexSim Newton IK for control part " - f"{self._robot_arm_name!r} ({root_link} -> {end_link})." - ) def _resolve_robot_ik_chain( self, @@ -470,29 +487,11 @@ def _as_pose_matrix(pose: object, device: torch.device) -> torch.Tensor: raise ValueError("Gizmo target pose must contain only finite values.") return matrix.detach().clone() - def _compute_ee_pose_fk(self) -> torch.Tensor: - if self.target is None or not isinstance(self.target, Robot): - raise RuntimeError("Robot Gizmo has no attached Robot.") - if self._robot_arm_name is None: - raise RuntimeError("Robot Gizmo control part is not configured.") - current_qpos = self.target.get_proprioception()["qpos"] - joint_ids = self.target.get_joint_ids(self._robot_arm_name) - joint_positions = current_qpos[:, joint_ids] - pose = self.target.compute_fk( - joint_positions, - name=self._robot_arm_name, - env_ids=[0], - to_matrix=True, - ) - if pose is None: - raise RuntimeError("Robot forward kinematics returned no pose.") - return self._as_pose_matrix(pose, self._target_device()) - def _read_target_pose(self) -> torch.Tensor: if self.target is None: raise RuntimeError("Gizmo is detached.") if self._target_type == "robot": - return self._compute_ee_pose_fk() + return self._read_native_robot_pose() pose = self.target.get_local_pose(to_matrix=True) return self._as_pose_matrix(pose[0], self._target_device()) @@ -622,39 +621,37 @@ def _update_rigid_object_pose(self, target_transform: torch.Tensor) -> bool: def _update_robot_ik(self, target_transform: torch.Tensor) -> bool: if self.target is None or not isinstance(self.target, Robot): return False - if self._robot_arm_name is None: + if self._ik_solver is None or self._robot_adapter is None: return False try: - current_qpos = self.target.get_proprioception()["qpos"] - joint_ids = self.target.get_joint_ids(self._robot_arm_name) - if len(joint_ids) == 0: - logger.log_warning( - f"No joint IDs found for control part {self._robot_arm_name!r}." - ) - return False - joint_seed = current_qpos[:, joint_ids] - result = self.target.compute_ik( - pose=target_transform, - name=self._robot_arm_name, - joint_seed=joint_seed, - env_ids=[0], + from dexsim.kit.ik.pose import ( + local_pose_from_world, + rotation_matrix_to_quat_xyzw, ) - if result is None: - return False - success, new_qpos = result - if not bool(torch.as_tensor(success).reshape(-1)[0].item()): - logger.log_warning("Gizmo IK solution not found.") - return False - new_qpos = torch.as_tensor( - new_qpos, - dtype=torch.float32, - device=self._target_device(), - ).reshape(1, -1) - self.target.set_qpos( - qpos=new_qpos, - joint_ids=joint_ids, - env_ids=[0], + + # The queued target is the TCP transform in the arena-local frame. + # Newton IK tracks a base-local target, so convert it with the same + # helper the native gizmo callback uses (inv(base_pose) @ target). + base_pose = self._robot_adapter.get_world_pose() + target_pose = target_transform[0].detach().cpu().numpy().astype(np.float32) + base_local = local_pose_from_world(base_pose, target_pose) + position = np.asarray(base_local[:3, 3], dtype=np.float32) + rotation = rotation_matrix_to_quat_xyzw(base_local[:3, :3]) + + joint_names = self._robot_adapter.get_actived_joint_names() + current_qpos = self._robot_adapter.get_current_qpos() + self._ik_solver.set_target_pose(position, rotation) + self._ik_solver.solve( + joint_names, + current_qpos, + iterations=self.cfg.ik_iterations, + ) + solved_qpos = self._ik_solver.qpos_for_joint_names( + joint_names, current_qpos ) + # Drive the joint targets (matching native IKApplyMode.DRIVE_TARGET) + # so physics moves the robot instead of snapping its current pose. + self._robot_adapter.set_target_qpos(solved_qpos) return True except Exception as error: logger.log_error(f"Error in Gizmo robot IK: {error}") @@ -700,12 +697,11 @@ def attach(self, target: BatchEntity) -> None: self._target_type = self._detect_target_type(target) self._robot_arm_name = None if self._target_type == "robot": - self._configure_robot(require_solver=not self._enable_native) + self._configure_robot() + self._setup_robot_ik_solver() if self._enable_native: self._setup_native_robot_gizmo() - desired_pose = self._read_native_robot_pose() - else: - desired_pose = self._read_target_pose() + desired_pose = self._read_native_robot_pose() else: desired_pose = self._read_target_pose() if self._enable_native: diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 6967e168d..54d315ab8 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -320,6 +320,10 @@ def __init__( # gizmo management self._gizmos: Dict[str, object] = dict() # Store active gizmos + # ``(uid, control_part)`` of the Gizmo currently owned by the Viser + # click-to-pick feature, or ``None``. Only one picker Gizmo is kept at a + # time and user-created Gizmos are never touched. + self._picker_gizmo: tuple[str, str | None] | None = None # marker management self._markers: Dict[str, MeshObject] = dict() @@ -2165,6 +2169,8 @@ def disable_gizmo(self, uid: str, control_part: str | None = None) -> None: try: gizmo = self._gizmos.pop(gizmo_key) + if self._picker_gizmo == (uid, control_part): + self._picker_gizmo = None try: if gizmo is not None: gizmo.destroy() @@ -2282,6 +2288,7 @@ def process_visualization_commands(self) -> int: def update_gizmos(self) -> None: """Apply Viser commands and update all active Gizmos.""" + self.process_pick_commands() self.process_visualization_commands() for gizmo_key, gizmo in list( getattr(self, "_gizmos", {}).items() @@ -2292,6 +2299,63 @@ def update_gizmos(self) -> None: except Exception as error: logger.log_error(f"Error updating gizmo '{gizmo_key}': {error}") + def process_pick_commands(self) -> int: + """Apply queued Viser click-pick commands on the simulation thread. + + A non-empty pick attaches a picker-owned Gizmo to the clicked node; + an empty pick (clicking empty space) clears it. Only one picker-owned + Gizmo is kept at a time and user-created Gizmos are never touched. + + Returns: + Number of pick commands drained from the visualization runtime. + """ + runtime = self._visualization_runtime + if runtime is None or not getattr( + self.sim_config.visualization, "allow_commands", False + ): + return 0 + processed = 0 + for command in runtime.drain_pick_commands(): + processed += 1 + if ( + command.run_id != runtime.exporter.run_id + or command.scene_revision != runtime.exporter.scene_revision + ): + continue + if command.node_id is None: + # Clicking empty space detaches the picker-owned Gizmo. + self._release_picker_gizmo() + continue + resolved = runtime.exporter.resolve_node_target(command.node_id) + if resolved is None: + continue + uid, kind = resolved + if kind not in {"robot", "rigid"}: + logger.log_warning( + f"Pick target kind {kind!r} (uid {uid!r}) is not gizmo-able; " + "only rigid objects and robots can be picked." + ) + self._release_picker_gizmo() + continue + # Re-clicking the already-picked target is a no-op (avoids flicker + # from recreating the same Gizmo, e.g. clicking near its handle). + if self._picker_gizmo is not None and self._picker_gizmo[0] == uid: + continue + self._release_picker_gizmo() + gizmo = self.enable_gizmo(uid=uid) + if gizmo is not None: + self._picker_gizmo = (uid, None) + return processed + + def _release_picker_gizmo(self) -> None: + """Detach the picker-owned Gizmo if one is currently attached.""" + if self._picker_gizmo is None: + return + uid, control_part = self._picker_gizmo + self._picker_gizmo = None + if self.has_gizmo(uid, control_part=control_part): + self.disable_gizmo(uid, control_part=control_part) + def toggle_gizmo_visibility( self, uid: str, control_part: str | None = None ) -> bool | None: diff --git a/embodichain/lab/visualization/__init__.py b/embodichain/lab/visualization/__init__.py index a297539e2..c61e4738a 100644 --- a/embodichain/lab/visualization/__init__.py +++ b/embodichain/lab/visualization/__init__.py @@ -33,6 +33,7 @@ GizmoSpec, GizmoState, MeshGeometry, + PickCommand, PointCloudOverlay, SceneFrame, SceneManifest, @@ -65,6 +66,7 @@ "GizmoState", "LatestFrameQueue", "MeshGeometry", + "PickCommand", "PointCloudOverlay", "RuntimeHealth", "RuntimeStats", diff --git a/embodichain/lab/visualization/backends/base.py b/embodichain/lab/visualization/backends/base.py index 36c876d0d..fc7c5f5c8 100644 --- a/embodichain/lab/visualization/backends/base.py +++ b/embodichain/lab/visualization/backends/base.py @@ -19,7 +19,13 @@ from abc import ABC, abstractmethod from collections.abc import Callable -from ..protocol import CameraImageFrame, GizmoCommand, SceneFrame, SceneManifest +from ..protocol import ( + CameraImageFrame, + GizmoCommand, + PickCommand, + SceneFrame, + SceneManifest, +) __all__ = ["VisualizationBackend"] @@ -34,6 +40,13 @@ def set_gizmo_command_sink( """Set the thread-safe sink used for browser Gizmo commands.""" self._gizmo_command_sink = sink + def set_pick_command_sink( + self, + sink: Callable[[PickCommand], None] | None, + ) -> None: + """Set the thread-safe sink used for browser click-pick commands.""" + self._pick_command_sink = sink + @property @abstractmethod def endpoint(self) -> str | None: diff --git a/embodichain/lab/visualization/backends/viser.py b/embodichain/lab/visualization/backends/viser.py index ca9c71da1..a741d70a0 100644 --- a/embodichain/lab/visualization/backends/viser.py +++ b/embodichain/lab/visualization/backends/viser.py @@ -25,6 +25,7 @@ import numpy as np from ..cfg import ViserServerCfg +from ..picker import ScenePicker from ..protocol import ( CameraImageFrame, CameraSpec, @@ -33,6 +34,7 @@ GizmoSpec, GizmoState, MeshGeometry, + PickCommand, PointCloudOverlay, SceneFrame, SceneManifest, @@ -127,6 +129,13 @@ def __init__( self._gizmo_owners: dict[str, str] = {} self._gizmo_drag_poses: dict[str, tuple[np.ndarray, np.ndarray]] = {} self._gizmo_sequence = 0 + self._picker = ScenePicker() + self._pick_enabled = False + self._node_geometry: dict[str, str] = {} + self._frame_positions: np.ndarray | None = None + self._frame_wxyz: np.ndarray | None = None + self._frame_visible: np.ndarray | None = None + self._pointer_handler: object | None = None self._world_handle: object | None = None self._ground_grid_handle: object | None = None self._camera_handles: dict[str, object] = {} @@ -208,6 +217,14 @@ def _(client: object) -> None: ) ) + if self.allow_commands and self._pointer_handler is None: + + @self._server.scene.on_pointer_event("click") + def _on_pick_click(event: object) -> None: + self._handle_pick_click(event) + + self._pointer_handler = _on_pick_click + def _register_visibility_controls(self, manifest: SceneManifest) -> None: previous_env_visibility = self._env_visibility while True: @@ -312,6 +329,22 @@ def _(event: object, selected_category: str = category) -> None: ) ) + if self.allow_commands: + with self._server.gui.add_folder("Interaction"): + pick_checkbox = self._server.gui.add_checkbox( + "Enable click-to-pick Gizmo", + initial_value=self._pick_enabled, + ) + + @pick_checkbox.on_update + def _(event: object) -> None: + self._gui_events.put( + _GuiEvent( + category="pick_enabled", + value=bool(event.target.value), + ) + ) + @staticmethod def _event_client_id(event: object) -> str | None: client_id = getattr(event, "client_id", None) @@ -322,6 +355,88 @@ def _event_client_id(event: object) -> str | None: return None return str(client_id) + def _handle_pick_click(self, event: object) -> None: + """Ray-cast a browser click and enqueue a PickCommand. + + Clicking a scene node attaches a picker-owned Gizmo to it; clicking + empty space (no hit) clears the picker-owned Gizmo. The command is + processed on the simulation thread. + """ + if not self._pick_enabled: + return + sink = getattr(self, "_pick_command_sink", None) + if sink is None or self._run_id is None: + return + ray_origin = getattr(event, "ray_origin", None) + ray_direction = getattr(event, "ray_direction", None) + if ray_origin is None or ray_direction is None: + return + client_id = self._event_client_id(event) or "unknown" + hit_node = self._picker.pick( + np.asarray(ray_origin, dtype=np.float32), + np.asarray(ray_direction, dtype=np.float32), + self._pick_instances(), + ) + sink( + PickCommand( + run_id=self._run_id, + scene_revision=self._scene_revision, + client_id=client_id, + node_id=hit_node, + ) + ) + + def _pick_instances( + self, + ) -> list[tuple[str, str, np.ndarray, np.ndarray]]: + """Build the ``(node_id, geometry_id, position, wxyz)`` pick candidates. + + Only visible, non-deformable mesh nodes are considered, since deformable + nodes update their vertices every frame and are not gizmo targets. + """ + instances: list[tuple[str, str, np.ndarray, np.ndarray]] = [] + positions = self._frame_positions + wxyz = self._frame_wxyz + visible = self._frame_visible + if positions is None or wxyz is None or visible is None: + return instances + for index, node_id in enumerate(self._frame_node_ids): + geometry_id = self._node_geometry.get(node_id) + if geometry_id is None or not bool(visible[index]): + continue + instances.append((node_id, geometry_id, positions[index], wxyz[index])) + return instances + + def _rebuild_picker( + self, + geometry_by_id: dict[str, MeshGeometry], + nodes_by_geometry: dict[str, list[SceneNode]], + ) -> None: + """Refresh cached pick geometry and the node-to-geometry map.""" + self._picker.clear() + self._node_geometry = {} + for geometry_id, nodes in nodes_by_geometry.items(): + geometry = geometry_by_id.get(geometry_id) + if geometry is None: + continue + self._picker.set_geometry(geometry_id, geometry.vertices, geometry.faces) + for node in nodes: + self._node_geometry[node.node_id] = geometry_id + + def _clear_picker_gizmo(self) -> None: + """Tell the simulation thread to release the picker-owned Gizmo.""" + sink = getattr(self, "_pick_command_sink", None) + if sink is None or self._run_id is None: + return + sink( + PickCommand( + run_id=self._run_id, + scene_revision=self._scene_revision, + client_id="picker-toggle", + node_id=None, + ) + ) + def _queue_gizmo_event( self, event: object, @@ -671,6 +786,8 @@ def publish_manifest(self, manifest: SceneManifest) -> None: else: nodes_by_geometry[node.geometry_id].append(node) + self._rebuild_picker(geometry_by_id, nodes_by_geometry) + removed_geometry_ids = set(self._mesh_batches) - set(nodes_by_geometry) for geometry_id in removed_geometry_ids: self._mesh_batches.pop(geometry_id).handle.remove() @@ -823,6 +940,10 @@ def _apply_gui_events(self) -> None: elif event.category == "overlay": category, visible = event.value self._overlay_visibility[str(category)] = bool(visible) + elif event.category == "pick_enabled": + self._pick_enabled = bool(event.value) + if not self._pick_enabled: + self._clear_picker_gizmo() elif event.category == "camera_environment": self._selected_camera_env = int(event.value) camera_uids = self._camera_uids_for_env(self._selected_camera_env) @@ -1028,6 +1149,10 @@ def publish_frame(self, frame: SceneFrame) -> bool: if not np.array_equal(batch.frame_visible, frame_visible): batch.frame_visible = frame_visible self._apply_mesh_visibility(batch) + # Retain the latest world-space poses for click-to-pick ray casting. + self._frame_positions = frame.positions + self._frame_wxyz = frame.wxyz + self._frame_visible = frame.visible for node_id, dynamic_mesh in self._dynamic_meshes.items(): index = dynamic_mesh.frame_index dynamic_mesh.frame_visible = bool(frame.visible[index]) diff --git a/embodichain/lab/visualization/picker.py b/embodichain/lab/visualization/picker.py new file mode 100644 index 000000000..a25f8d689 --- /dev/null +++ b/embodichain/lab/visualization/picker.py @@ -0,0 +1,216 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Backend-neutral ray-mesh picking for Viser click selection. + +Viser's ``on_pointer_event`` callback exposes the camera ray but not the scene +node it hits. :class:`ScenePicker` closes that gap by ray-casting the ray +against the cached scene geometry with a vectorized Möller-Trumbore test, +returning the closest hit node so the simulation can attach a Gizmo to it. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Iterable + +import numpy as np + +__all__ = ["ScenePicker"] + +_EPSILON = 1.0e-9 + + +@dataclass(frozen=True) +class _Geometry: + """Cached triangle data for one geometry, stored in local coordinates.""" + + v0: np.ndarray + edge1: np.ndarray + edge2: np.ndarray + + +def _wxyz_to_rotation(wxyz: np.ndarray) -> np.ndarray: + """Convert a normalized wxyz quaternion to a 3x3 rotation matrix.""" + w, x, y, z = np.asarray(wxyz, dtype=np.float64) + rotation = np.array( + [ + [1.0 - 2.0 * (y * y + z * z), 2.0 * (x * y - w * z), 2.0 * (x * z + w * y)], + [2.0 * (x * y + w * z), 1.0 - 2.0 * (x * x + z * z), 2.0 * (y * z - w * x)], + [2.0 * (x * z - w * y), 2.0 * (y * z + w * x), 1.0 - 2.0 * (x * x + y * y)], + ], + dtype=np.float32, + ) + return rotation + + +class ScenePicker: + """Resolve a world-space ray to the closest hit scene node. + + Geometry is cached per ``geometry_id`` in local coordinates. Each pick + transforms the ray into every instance's local frame (so cached triangle + data is reused across instances and across frames) and runs a vectorized + Möller-Trumbore test, keeping the smallest positive ray parameter. + + Args: + epsilon: Lower bound for accepted ray parameters, in world length units. + """ + + def __init__(self, epsilon: float = _EPSILON) -> None: + self._geometries: dict[str, _Geometry] = {} + self._epsilon = float(epsilon) + + def set_geometry( + self, + geometry_id: str, + vertices: np.ndarray, + faces: np.ndarray, + ) -> None: + """Cache one geometry's triangle data in local coordinates. + + Args: + geometry_id: Stable geometry identifier from the scene manifest. + vertices: Triangle mesh vertices with shape ``(V, 3)``. + faces: Triangle indices into ``vertices`` with shape ``(F, 3)``. + """ + verts = np.ascontiguousarray(np.asarray(vertices, dtype=np.float32)) + tris = np.ascontiguousarray(np.asarray(faces, dtype=np.int64)) + if verts.ndim != 2 or verts.shape[1] != 3: + raise ValueError( + f"vertices must have shape (V, 3), received {verts.shape}." + ) + if tris.ndim != 2 or tris.shape[1] != 3: + raise ValueError(f"faces must have shape (F, 3), received {tris.shape}.") + if tris.size == 0: + self._geometries.pop(geometry_id, None) + return + v0 = verts[tris[:, 0]] + v1 = verts[tris[:, 1]] + v2 = verts[tris[:, 2]] + self._geometries[geometry_id] = _Geometry( + v0=v0, + edge1=v1 - v0, + edge2=v2 - v0, + ) + + def remove_geometry(self, geometry_id: str) -> None: + """Drop one cached geometry.""" + self._geometries.pop(geometry_id, None) + + def clear(self) -> None: + """Drop all cached geometry.""" + self._geometries.clear() + + def pick( + self, + ray_origin: np.ndarray, + ray_direction: np.ndarray, + instances: Iterable[tuple[str, str, np.ndarray, np.ndarray]], + ) -> str | None: + """Return the node id of the closest instance hit by the ray. + + Each instance is a ``(node_id, geometry_id, position, wxyz)`` tuple, + where ``position`` is the world-space translation and ``wxyz`` is the + normalized ``[w, x, y, z]`` quaternion. The ray is transformed into each + instance's local frame so the cached local geometry can be reused. + + Args: + ray_origin: World-space ray origin with shape ``(3,)``. + ray_direction: World-space ray direction with shape ``(3,)``. It is + normalized internally so the returned hit distance is in world + length units. + instances: Iterable of scene instances to test. + + Returns: + The closest hit ``node_id``, or ``None`` if the ray misses every + instance. + """ + origin = np.asarray(ray_origin, dtype=np.float32) + direction = np.asarray(ray_direction, dtype=np.float32) + if origin.shape != (3,) or direction.shape != (3,): + raise ValueError("ray_origin and ray_direction must have shape (3,).") + dir_norm = float(np.linalg.norm(direction)) + if dir_norm <= self._epsilon: + return None + direction = direction / dir_norm + + best_node: str | None = None + best_t = np.inf + for node_id, geometry_id, position, wxyz in instances: + geometry = self._geometries.get(geometry_id) + if geometry is None: + continue + local_origin, local_direction = self._world_to_local_ray( + origin, direction, position, wxyz + ) + hit_t = self._ray_cast_geometry(geometry, local_origin, local_direction) + if hit_t is not None and hit_t < best_t: + best_t = hit_t + best_node = node_id + return best_node + + @staticmethod + def _world_to_local_ray( + origin: np.ndarray, + direction: np.ndarray, + position: np.ndarray, + wxyz: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray]: + """Transform a world ray into an instance's local frame. + + The direction is left unnormalized after the inverse rotation so the ray + parameter stays in world length units: the local triangle hit parameter + equals the world-space distance along the (normalized) world ray. + """ + rotation = _wxyz_to_rotation(wxyz) + inv_rotation = rotation.T + local_origin = inv_rotation @ (origin - np.asarray(position, dtype=np.float32)) + local_direction = inv_rotation @ direction + return local_origin.astype(np.float32), local_direction.astype(np.float32) + + def _ray_cast_geometry( + self, + geometry: _Geometry, + origin: np.ndarray, + direction: np.ndarray, + ) -> float | None: + """Return the smallest positive ray parameter hitting one geometry.""" + edge1 = geometry.edge1 + edge2 = geometry.edge2 + v0 = geometry.v0 + + h = np.cross(direction, edge2) # (F, 3) + a = np.einsum("fd,fd->f", edge1, h) # (F,) + parallel = np.abs(a) <= self._epsilon + # Avoid division by zero for parallel rays; mask them out later. + safe_a = np.where(parallel, 1.0, a) + f = 1.0 / safe_a + s = origin - v0 # (F, 3) + u = f * np.einsum("fd,fd->f", s, h) + q = np.cross(s, edge1) # (F, 3) + v = f * np.einsum("d,fd->f", direction, q) + t = f * np.einsum("fd,fd->f", edge2, q) + + valid = ( + (~parallel) + & (u >= 0.0) + & (u <= 1.0) + & (v >= 0.0) + & (u + v <= 1.0) + & (t > self._epsilon) + ) + if not np.any(valid): + return None + return float(np.min(t[valid])) diff --git a/embodichain/lab/visualization/protocol.py b/embodichain/lab/visualization/protocol.py index c0d786229..c2e42f1db 100644 --- a/embodichain/lab/visualization/protocol.py +++ b/embodichain/lab/visualization/protocol.py @@ -35,6 +35,7 @@ "GizmoSpec", "GizmoState", "MeshGeometry", + "PickCommand", "PointCloudOverlay", "SceneFrame", "SceneManifest", @@ -348,6 +349,31 @@ def __post_init__(self) -> None: object.__setattr__(self, "wxyz", wxyz) +@dataclass(frozen=True) +class PickCommand: + """Immutable browser click-pick command consumed on the simulation thread. + + A non-empty ``node_id`` requests a Gizmo on the clicked scene node; a + ``None`` ``node_id`` (clicking empty space) clears the picker-owned Gizmo. + """ + + run_id: str + scene_revision: int + client_id: str + node_id: str | None + schema_version: int = SCHEMA_VERSION + + def __post_init__(self) -> None: + if not self.run_id: + raise ValueError("Pick command run_id must not be empty.") + if self.scene_revision < 0: + raise ValueError("Pick command scene_revision must be non-negative.") + if not self.client_id: + raise ValueError("Pick command client_id must not be empty.") + if self.node_id is not None and not self.node_id: + raise ValueError("Pick command node_id must be None or non-empty.") + + @dataclass(frozen=True) class SceneNode: """One mesh-bearing logical node in a scene manifest.""" diff --git a/embodichain/lab/visualization/runtime.py b/embodichain/lab/visualization/runtime.py index 44bfa2d14..1137facf8 100644 --- a/embodichain/lab/visualization/runtime.py +++ b/embodichain/lab/visualization/runtime.py @@ -28,6 +28,7 @@ from .protocol import ( CameraImageFrame, GizmoCommand, + PickCommand, SceneFrame, SceneManifest, SceneOverlays, @@ -138,6 +139,44 @@ def clear(self) -> None: self._commands.clear() +class PickCommandQueue: + """Bounded queue for low-frequency browser click-pick commands. + + Only the latest pick per client is retained, so a rapid sequence of clicks + from one browser cannot pile up ahead of the simulation thread. + """ + + def __init__(self, maxsize: int = 64) -> None: + if maxsize <= 0: + raise ValueError("maxsize must be greater than zero.") + self._maxsize = maxsize + self._commands: deque[PickCommand] = deque() + self._lock = threading.Lock() + + def put(self, command: PickCommand) -> None: + """Enqueue a pick command without blocking the Viser callback thread.""" + with self._lock: + for index in range(len(self._commands) - 1, -1, -1): + if self._commands[index].client_id == command.client_id: + self._commands[index] = command + return + if len(self._commands) >= self._maxsize: + self._commands.popleft() + self._commands.append(command) + + def drain(self) -> tuple[PickCommand, ...]: + """Return and clear all queued commands in arrival order.""" + with self._lock: + commands = tuple(self._commands) + self._commands.clear() + return commands + + def clear(self) -> None: + """Discard all queued commands.""" + with self._lock: + self._commands.clear() + + @dataclass(frozen=True) class RuntimeStats: """Snapshot of scene and camera-image capture/upload telemetry.""" @@ -204,6 +243,8 @@ def __init__( self._backend = backend self._gizmo_commands = GizmoCommandQueue() self._backend.set_gizmo_command_sink(self._enqueue_gizmo_command) + self._pick_commands = PickCommandQueue() + self._backend.set_pick_command_sink(self._enqueue_pick_command) self._frames: LatestFrameQueue[SceneFrame] = LatestFrameQueue() self._camera_images: LatestFrameQueue[CameraImageFrame] = LatestFrameQueue() self._manifests: queue.Queue[SceneManifest] = queue.Queue() @@ -228,6 +269,16 @@ def drain_gizmo_commands(self) -> tuple[GizmoCommand, ...]: return () return self._gizmo_commands.drain() + def _enqueue_pick_command(self, command: PickCommand) -> None: + if self.cfg.allow_commands: + self._pick_commands.put(command) + + def drain_pick_commands(self) -> tuple[PickCommand, ...]: + """Drain browser click-pick commands for simulation-thread processing.""" + if not self.cfg.allow_commands: + return () + return self._pick_commands.drain() + @property def endpoint(self) -> str | None: """Local browser endpoint after :meth:`start` returns.""" @@ -480,6 +531,7 @@ def stop(self, timeout: float = 10.0) -> None: self._frames.clear() self._camera_images.clear() self._gizmo_commands.clear() + self._pick_commands.clear() self._raise_worker_error() def __enter__(self) -> VisualizationRuntime: diff --git a/embodichain/lab/visualization/scene_exporter.py b/embodichain/lab/visualization/scene_exporter.py index 941eb80de..9bc51899b 100644 --- a/embodichain/lab/visualization/scene_exporter.py +++ b/embodichain/lab/visualization/scene_exporter.py @@ -439,6 +439,26 @@ def _append_gizmos(self, sources: list[_GizmoSource]) -> None: ) ) + def resolve_node_target(self, node_id: str) -> tuple[str, str] | None: + """Map a published scene node id to its ``(uid, kind)``. + + Used by the simulation thread to turn a Viser click-pick result into the + asset uid that :meth:`SimulationManager.enable_gizmo` expects. + + Args: + node_id: Scene node id from the current manifest. + + Returns: + ``(uid, kind)`` where ``kind`` is the asset kind (for example + ``"rigid"``, ``"robot"``, or ``"articulation"``), or ``None`` if the + node id is not part of the current scene. + """ + for source in self._sources: + if source.node.node_id == node_id: + kind, uid = source.asset_key + return str(uid), str(kind) + return None + def _append_rigid_object_groups( self, sources: list[_NodeSource], diff --git a/tests/sim/objects/test_gizmo.py b/tests/sim/objects/test_gizmo.py index f2b444556..d6a138105 100644 --- a/tests/sim/objects/test_gizmo.py +++ b/tests/sim/objects/test_gizmo.py @@ -212,39 +212,6 @@ class _Camera(_RigidObject): pass -class _Robot: - def __init__(self) -> None: - self.device = torch.device("cpu") - self.cfg = SimpleNamespace(uid="robot", solver_cfg={"arm": object()}) - self.control_parts = {"arm": ["joint"]} - self.set_calls: list[tuple[torch.Tensor, list[int], list[int]]] = [] - - def get_proprioception(self) -> dict[str, torch.Tensor]: - return {"qpos": torch.zeros((1, 2), dtype=torch.float32)} - - def get_joint_ids(self, name: str) -> list[int]: - assert name == "arm" - return [0, 1] - - def compute_fk(self, *args: object, **kwargs: object) -> torch.Tensor: - return torch.eye(4, dtype=torch.float32).unsqueeze(0) - - def compute_ik( - self, - *args: object, - **kwargs: object, - ) -> tuple[torch.Tensor, torch.Tensor]: - return torch.tensor([True]), torch.tensor([[0.4, -0.2]]) - - def set_qpos( - self, - qpos: torch.Tensor, - joint_ids: list[int], - env_ids: list[int], - ) -> None: - self.set_calls.append((qpos.clone(), joint_ids, env_ids)) - - def _patch_headless_dexsim(monkeypatch) -> None: monkeypatch.setattr(gizmo_module.dexsim, "get_world_num", lambda: 1) monkeypatch.setattr( @@ -289,19 +256,63 @@ def test_headless_camera_gizmo_uses_shared_pose_path(monkeypatch) -> None: torch.testing.assert_close(target.pose, pose) -def test_headless_robot_gizmo_preserves_native_fk_ik_behavior(monkeypatch) -> None: - monkeypatch.setattr(gizmo_module, "Robot", _Robot) +def test_headless_robot_gizmo_uses_dexsim_newton_ik(monkeypatch) -> None: + """Headless (Viser) robot Gizmo solves IK with DexSim Newton IK. + + The queued Viser target is converted to the robot base-local frame and + driven through the Newton solver; the solved qpos is written back as a + joint drive target, mirroring the native ``IKApplyMode.DRIVE_TARGET`` path + instead of calling the EmbodiChain ``compute_ik`` solver. + """ + monkeypatch.setattr(gizmo_module, "Robot", _FakeAdapterRobot) _patch_headless_dexsim(monkeypatch) - target = _Robot() + target = _FakeAdapterRobot() + + solved_qpos = np.array([0.4, -0.2], dtype=np.float32) + + class _FakeNewtonSolver: + def __init__(self) -> None: + self.set_target_calls: list[tuple[np.ndarray, np.ndarray]] = [] + self.solve_iterations: list[int | None] = [] + + def set_target_pose(self, position, rotation) -> None: + self.set_target_calls.append( + (np.array(position, copy=True), np.array(rotation, copy=True)) + ) + + def solve(self, joint_names, current_qpos, iterations=None) -> None: + self.solve_iterations.append(iterations) + + def qpos_for_joint_names(self, joint_names, fallback_qpos): + return solved_qpos + + fake_solver = _FakeNewtonSolver() + + def _inject_solver(self) -> None: + self._robot_adapter = _RobotGizmoAdapter(target, "arm") + self._ik_solver = fake_solver + self._native_robot_end_link = "tool_link" + self._native_robot_tcp_pose = np.eye(4, dtype=np.float32) + + monkeypatch.setattr(Gizmo, "_setup_robot_ik_solver", _inject_solver) + gizmo = Gizmo(target, control_part="arm", enable_native=False) pose = torch.eye(4, dtype=torch.float32).unsqueeze(0) pose[0, 0, 3] = 0.5 + assert not gizmo.native_enabled assert gizmo.request_local_pose(pose, source_id="viser:client-a") gizmo.update() - assert target.set_calls[-1][1:] == ([0, 1], [0]) + # The Newton solver received a base-local target and was asked to solve + # with the configured iteration count. + assert fake_solver.set_target_calls + assert fake_solver.solve_iterations == [gizmo.cfg.ik_iterations] + # The solved qpos is written back as a drive target through the adapter. + write = target.write_calls[-1] + assert write["target"] is True + assert write["joint_ids"] == [0, 2] torch.testing.assert_close( - target.set_calls[-1][0], - torch.tensor([[0.4, -0.2]]), + write["qpos"], + torch.tensor([[0.4, -0.2]], dtype=torch.float32), ) diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index 0befea621..0f6da1ea4 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -32,6 +32,7 @@ ) from embodichain.lab.visualization import ( GizmoCommand, + PickCommand, PointCloudOverlay, SceneOverlays, VisualizationCfg, @@ -331,6 +332,167 @@ def test_sim_manager_routes_viser_gizmo_commands_in_local_arena_frame() -> None: ) +def _make_pick_sim_manager(pick_commands, resolve): + """Build a minimally initialized manager with stubbed gizmo lifecycle.""" + sim = object.__new__(SimulationManager) + sim._gizmos = {} + sim._picker_gizmo = None + enabled: list = [] + disabled: list = [] + + def fake_enable(uid, control_part=None, gizmo_cfg=None, *, enable_native=None): + enabled.append((uid, control_part)) + return SimpleNamespace(control_part=control_part) + + def fake_disable(uid, control_part=None): + disabled.append((uid, control_part)) + + sim.enable_gizmo = fake_enable + sim.disable_gizmo = fake_disable + sim.has_gizmo = lambda uid, control_part=None: True + sim.sim_config = SimpleNamespace( + visualization=SimpleNamespace(allow_commands=True), + ) + sim._visualization_runtime = SimpleNamespace( + exporter=SimpleNamespace( + run_id="run", + scene_revision=2, + resolve_node_target=resolve, + ), + drain_pick_commands=lambda: pick_commands, + ) + return sim, enabled, disabled + + +def test_process_pick_commands_attaches_single_picker_gizmo() -> None: + pick_commands = ( + PickCommand( + run_id="run", + scene_revision=2, + client_id="client-a", + node_id="env:0/rigid:cube", + ), + PickCommand( + run_id="run", + scene_revision=2, + client_id="client-a", + node_id="env:0/robot:ur10", + ), + PickCommand( + run_id="run", + scene_revision=2, + client_id="client-a", + node_id=None, + ), + ) + + def resolve(node_id: str): + if node_id == "env:0/rigid:cube": + return ("cube", "rigid") + if node_id == "env:0/robot:ur10": + return ("ur10", "robot") + return None + + sim, enabled, disabled = _make_pick_sim_manager(pick_commands, resolve) + + processed = sim.process_pick_commands() + + assert processed == 3 + # cube attached, then swapped to ur10 (disabling cube), then ur10 cleared. + assert enabled == [("cube", None), ("ur10", None)] + assert disabled == [("cube", None), ("ur10", None)] + assert sim._picker_gizmo is None + + +def test_process_pick_commands_skips_non_gizmo_targets() -> None: + pick_commands = ( + PickCommand( + run_id="run", + scene_revision=2, + client_id="client-a", + node_id="env:0/soft:cloth", + ), + ) + sim, enabled, disabled = _make_pick_sim_manager( + pick_commands, lambda node_id: ("cloth", "soft") + ) + + processed = sim.process_pick_commands() + + assert processed == 1 + assert enabled == [] # soft bodies are not gizmo-able + assert disabled == [] + assert sim._picker_gizmo is None + + +def test_process_pick_commands_is_noop_for_already_picked_target() -> None: + pick_commands = ( + PickCommand( + run_id="run", + scene_revision=2, + client_id="client-a", + node_id="env:0/rigid:cube", + ), + PickCommand( + run_id="run", + scene_revision=2, + client_id="client-a", + node_id="env:0/rigid:cube", # same target again + ), + ) + sim, enabled, disabled = _make_pick_sim_manager( + pick_commands, lambda node_id: ("cube", "rigid") + ) + + processed = sim.process_pick_commands() + + assert processed == 2 + # The second pick is a no-op: no flicker from disable+re-enable. + assert enabled == [("cube", None)] + assert disabled == [] + assert sim._picker_gizmo == ("cube", None) + + +def test_process_pick_commands_ignores_stale_scene_revision() -> None: + pick_commands = ( + PickCommand( + run_id="run", + scene_revision=99, # stale + client_id="client-a", + node_id="env:0/rigid:cube", + ), + ) + sim, enabled, disabled = _make_pick_sim_manager( + pick_commands, lambda node_id: ("cube", "rigid") + ) + + processed = sim.process_pick_commands() + + assert processed == 1 + assert enabled == [] + assert sim._picker_gizmo is None + + +def test_process_pick_commands_noop_without_command_permission() -> None: + sim = object.__new__(SimulationManager) + sim.sim_config = SimpleNamespace( + visualization=SimpleNamespace(allow_commands=False), + ) + sim._visualization_runtime = SimpleNamespace( + exporter=SimpleNamespace(run_id="run", scene_revision=2), + drain_pick_commands=lambda: ( + PickCommand( + run_id="run", + scene_revision=2, + client_id="client-a", + node_id="env:0/rigid:cube", + ), + ), + ) + + assert sim.process_pick_commands() == 0 + + def test_simulation_config_nests_viser_server_under_visualization() -> None: cfg = SimulationManagerCfg() diff --git a/tests/visualization/test_picker.py b/tests/visualization/test_picker.py new file mode 100644 index 000000000..e13c7356d --- /dev/null +++ b/tests/visualization/test_picker.py @@ -0,0 +1,218 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import math + +import numpy as np +import pytest + +from embodichain.lab.visualization.picker import ScenePicker + + +def _unit_cube() -> tuple[np.ndarray, np.ndarray]: + vertices = np.array( + [ + [-0.5, -0.5, -0.5], + [0.5, -0.5, -0.5], + [0.5, 0.5, -0.5], + [-0.5, 0.5, -0.5], + [-0.5, -0.5, 0.5], + [0.5, -0.5, 0.5], + [0.5, 0.5, 0.5], + [-0.5, 0.5, 0.5], + ], + dtype=np.float32, + ) + faces = np.array( + [ + [0, 1, 2], + [0, 2, 3], + [4, 6, 5], + [4, 7, 6], + [0, 4, 5], + [0, 5, 1], + [2, 6, 7], + [2, 7, 3], + [1, 5, 6], + [1, 6, 2], + [0, 3, 7], + [0, 7, 4], + ], + dtype=np.int64, + ) + return vertices, faces + + +IDENTITY_WXYZ = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32) + + +def test_pick_hits_top_face_of_unit_cube() -> None: + picker = ScenePicker() + picker.set_geometry("cube", *_unit_cube()) + + hit = picker.pick( + np.array([0.0, 0.0, 5.0], dtype=np.float32), + np.array([0.0, 0.0, -1.0], dtype=np.float32), + [("cubeA", "cube", np.zeros(3, dtype=np.float32), IDENTITY_WXYZ)], + ) + + assert hit == "cubeA" + + +def test_pick_returns_none_when_ray_misses() -> None: + picker = ScenePicker() + picker.set_geometry("cube", *_unit_cube()) + + hit = picker.pick( + np.array([5.0, 5.0, 5.0], dtype=np.float32), + np.array([0.0, 0.0, -1.0], dtype=np.float32), + [("cubeA", "cube", np.zeros(3, dtype=np.float32), IDENTITY_WXYZ)], + ) + + assert hit is None + + +def test_pick_respects_translated_instance() -> None: + picker = ScenePicker() + picker.set_geometry("cube", *_unit_cube()) + + # The cube is translated to x=3; a ray straight down at x=3 hits it. + hit = picker.pick( + np.array([3.0, 0.0, 5.0], dtype=np.float32), + np.array([0.0, 0.0, -1.0], dtype=np.float32), + [("cubeB", "cube", np.array([3.0, 0.0, 0.0], dtype=np.float32), IDENTITY_WXYZ)], + ) + + assert hit == "cubeB" + + +def test_pick_returns_closest_instance() -> None: + picker = ScenePicker() + picker.set_geometry("cube", *_unit_cube()) + + instances = [ + ("near", "cube", np.array([0.0, 0.0, 0.0], dtype=np.float32), IDENTITY_WXYZ), + ("far", "cube", np.array([0.0, 0.0, 2.0], dtype=np.float32), IDENTITY_WXYZ), + ] + # Ray from z=5 going -z hits "far" (top at z=2.5) before "near" (top at z=0.5). + hit = picker.pick( + np.array([0.0, 0.0, 5.0], dtype=np.float32), + np.array([0.0, 0.0, -1.0], dtype=np.float32), + instances, + ) + + assert hit == "far" + + +def test_pick_respects_rotated_instance() -> None: + picker = ScenePicker() + picker.set_geometry("cube", *_unit_cube()) + + # 90-degree rotation about y: the +z face now points along +x. + angle = math.pi / 2.0 + wxyz = np.array([math.cos(angle), 0.0, math.sin(angle), 0.0], dtype=np.float32) + hit = picker.pick( + np.array([5.0, 0.0, 0.0], dtype=np.float32), + np.array([-1.0, 0.0, 0.0], dtype=np.float32), + [("rotated", "cube", np.zeros(3, dtype=np.float32), wxyz)], + ) + + assert hit == "rotated" + + +def test_pick_skips_instances_with_unknown_geometry() -> None: + picker = ScenePicker() + picker.set_geometry("cube", *_unit_cube()) + + instances = [ + ("unknown", "missing", np.zeros(3, dtype=np.float32), IDENTITY_WXYZ), + ("cubeA", "cube", np.zeros(3, dtype=np.float32), IDENTITY_WXYZ), + ] + hit = picker.pick( + np.array([0.0, 0.0, 5.0], dtype=np.float32), + np.array([0.0, 0.0, -1.0], dtype=np.float32), + instances, + ) + + assert hit == "cubeA" + + +def test_pick_with_no_instances_returns_none() -> None: + picker = ScenePicker() + picker.set_geometry("cube", *_unit_cube()) + + hit = picker.pick( + np.array([0.0, 0.0, 5.0], dtype=np.float32), + np.array([0.0, 0.0, -1.0], dtype=np.float32), + [], + ) + + assert hit is None + + +def test_pick_with_empty_geometry_is_skipped() -> None: + picker = ScenePicker() + picker.set_geometry( + "empty", np.zeros((0, 3), dtype=np.float32), np.zeros((0, 3), dtype=np.int64) + ) + picker.set_geometry("cube", *_unit_cube()) + + hit = picker.pick( + np.array([0.0, 0.0, 5.0], dtype=np.float32), + np.array([0.0, 0.0, -1.0], dtype=np.float32), + [("cubeA", "cube", np.zeros(3, dtype=np.float32), IDENTITY_WXYZ)], + ) + + assert hit == "cubeA" + + +def test_set_geometry_rejects_bad_shapes() -> None: + picker = ScenePicker() + with pytest.raises(ValueError, match="vertices"): + picker.set_geometry( + "bad", np.zeros((4,), dtype=np.float32), np.zeros((1, 3), dtype=np.int64) + ) + with pytest.raises(ValueError, match="faces"): + picker.set_geometry( + "bad", np.zeros((3, 3), dtype=np.float32), np.zeros((3,), dtype=np.int64) + ) + + +def test_pick_rejects_bad_ray_shapes() -> None: + picker = ScenePicker() + picker.set_geometry("cube", *_unit_cube()) + with pytest.raises(ValueError, match="ray_origin"): + picker.pick( + np.zeros((2,), dtype=np.float32), + np.array([0.0, 0.0, -1.0], dtype=np.float32), + [], + ) + + +def test_pick_normalizes_direction_so_distance_is_in_world_units() -> None: + picker = ScenePicker() + picker.set_geometry("cube", *_unit_cube()) + + # An unnormalized direction should produce the same hit as the normalized one. + hit = picker.pick( + np.array([0.0, 0.0, 5.0], dtype=np.float32), + np.array([0.0, 0.0, -2.0], dtype=np.float32), + [("cubeA", "cube", np.zeros(3, dtype=np.float32), IDENTITY_WXYZ)], + ) + + assert hit == "cubeA" diff --git a/tests/visualization/test_viser_backend.py b/tests/visualization/test_viser_backend.py index 826ed816a..03e0a0e31 100644 --- a/tests/visualization/test_viser_backend.py +++ b/tests/visualization/test_viser_backend.py @@ -180,6 +180,17 @@ def add_transform_controls( self.transform_controls.append(handle) return handle + def on_pointer_event(self, event_type: str): + """Register a pointer callback like viser's scene API (test-only).""" + + def decorator(callback: object) -> object: + if not hasattr(self, "pointer_callbacks"): + self.pointer_callbacks = [] + self.pointer_callbacks.append((event_type, callback)) + return callback + + return decorator + class _Server: def __init__(self, **kwargs: object) -> None: @@ -563,6 +574,169 @@ def event(client_id: str, position: list[float]) -> SimpleNamespace: backend.stop() +def _unit_cube_geometry() -> MeshGeometry: + vertices = np.array( + [ + [-0.5, -0.5, -0.5], + [0.5, -0.5, -0.5], + [0.5, 0.5, -0.5], + [-0.5, 0.5, -0.5], + [-0.5, -0.5, 0.5], + [0.5, -0.5, 0.5], + [0.5, 0.5, 0.5], + [-0.5, 0.5, 0.5], + ], + dtype=np.float32, + ) + faces = np.array( + [ + [0, 1, 2], + [0, 2, 3], + [4, 6, 5], + [4, 7, 6], + [0, 4, 5], + [0, 5, 1], + [2, 6, 7], + [2, 7, 3], + [1, 5, 6], + [1, 6, 2], + [0, 3, 7], + [0, 7, 4], + ], + dtype=np.uint32, + ) + return MeshGeometry(geometry_id="cube", vertices=vertices, faces=faces) + + +def _make_pickable_scene() -> tuple[SceneManifest, SceneFrame]: + node = SceneNode( + node_id="env:0/rigid:cube", + path="/envs/0/rigid_objects/cube", + parent_id=None, + env_id=0, + kind="rigid_object", + geometry_id="cube", + ) + manifest = SceneManifest("run", 1, (node,), (_unit_cube_geometry(),)) + frame = SceneFrame( + run_id="run", + scene_revision=1, + sequence=1, + sim_step=1, + sim_time=0.01, + node_ids=("env:0/rigid:cube",), + positions=np.array([[0.0, 0.0, 0.0]], dtype=np.float32), + wxyz=np.array([[1.0, 0.0, 0.0, 0.0]], dtype=np.float32), + visible=np.array([True], dtype=np.bool_), + ) + return manifest, frame + + +def _make_pick_backend() -> tuple[object, object, list]: + server = _Server() + pick_commands: list = [] + backend = ViserBackend( + ViserServerCfg(port=8765), + server_factory=lambda **_: server, + allow_commands=True, + ) + backend.set_pick_command_sink(pick_commands.append) + return backend, server, pick_commands + + +def test_viser_backend_pick_enqueues_command_when_enabled() -> None: + backend, server, pick_commands = _make_pick_backend() + manifest, frame = _make_pickable_scene() + + backend.start() + backend.publish_manifest(manifest) + assert backend.publish_frame(frame) + + # The click handler is registered but inactive until the checkbox is on. + click_callback = server.scene.pointer_callbacks[0][1] + click_callback( + SimpleNamespace( + client_id="client-a", + ray_origin=np.array([0.0, 0.0, 5.0], dtype=np.float32), + ray_direction=np.array([0.0, 0.0, -1.0], dtype=np.float32), + ) + ) + assert pick_commands == [] + + server.gui.checkboxes["Enable click-to-pick Gizmo"].callback( + SimpleNamespace(target=SimpleNamespace(value=True)) + ) + backend.poll() + + click_callback( + SimpleNamespace( + client_id="client-a", + ray_origin=np.array([0.0, 0.0, 5.0], dtype=np.float32), + ray_direction=np.array([0.0, 0.0, -1.0], dtype=np.float32), + ) + ) + + assert len(pick_commands) == 1 + command = pick_commands[0] + assert command.node_id == "env:0/rigid:cube" + assert command.client_id == "client-a" + backend.stop() + + +def test_viser_backend_pick_miss_enqueues_empty_command() -> None: + backend, server, pick_commands = _make_pick_backend() + manifest, frame = _make_pickable_scene() + + backend.start() + backend.publish_manifest(manifest) + assert backend.publish_frame(frame) + + server.gui.checkboxes["Enable click-to-pick Gizmo"].callback( + SimpleNamespace(target=SimpleNamespace(value=True)) + ) + backend.poll() + + click_callback = server.scene.pointer_callbacks[0][1] + # Ray off to the side misses the cube. + click_callback( + SimpleNamespace( + client_id="client-a", + ray_origin=np.array([5.0, 5.0, 5.0], dtype=np.float32), + ray_direction=np.array([0.0, 0.0, -1.0], dtype=np.float32), + ) + ) + + assert len(pick_commands) == 1 + assert pick_commands[0].node_id is None + backend.stop() + + +def test_viser_backend_disabling_pick_clears_picker_gizmo() -> None: + backend, server, pick_commands = _make_pick_backend() + manifest, frame = _make_pickable_scene() + + backend.start() + backend.publish_manifest(manifest) + assert backend.publish_frame(frame) + + server.gui.checkboxes["Enable click-to-pick Gizmo"].callback( + SimpleNamespace(target=SimpleNamespace(value=True)) + ) + backend.poll() + assert backend._pick_enabled is True + + server.gui.checkboxes["Enable click-to-pick Gizmo"].callback( + SimpleNamespace(target=SimpleNamespace(value=False)) + ) + backend.poll() + + assert backend._pick_enabled is False + # Disabling emits an empty pick so the simulation releases the gizmo. + assert len(pick_commands) == 1 + assert pick_commands[0].node_id is None + backend.stop() + + def test_viser_backend_keeps_gizmos_read_only_without_command_permission() -> None: server = _Server() backend = ViserBackend( From 19e3618917a4541dd3276923abd4b38010501afd Mon Sep 17 00:00:00 2001 From: yuecideng Date: Thu, 6 Aug 2026 15:34:19 +0800 Subject: [PATCH 03/14] wip --- docs/source/guides/cli.md | 3 + docs/source/guides/run_env.md | 12 ++ docs/source/tutorial/data_generation.rst | 2 +- embodichain/lab/scripts/run_env.py | 225 +++++++++++++++++++---- tests/lab/scripts/test_run_env.py | 127 ++++++++++++- 5 files changed, 332 insertions(+), 37 deletions(-) diff --git a/docs/source/guides/cli.md b/docs/source/guides/cli.md index 3b7ca2e16..3f5412f00 100644 --- a/docs/source/guides/cli.md +++ b/docs/source/guides/cli.md @@ -222,6 +222,9 @@ configuration, remote access, and performance details, see When ``--preview`` is enabled, an interactive REPL is available: +- **``i``** — show or hide solver-backed robot IK Gizmos; in the native DexSim + window, **``I``** provides the same toggle and the Gizmos can be dragged to + operate the robot (single-environment previews only) - **``p``** — enter an IPython embed session with ``env`` in scope - **``q``** — quit diff --git a/docs/source/guides/run_env.md b/docs/source/guides/run_env.md index 5d8d8665d..0816cab2d 100644 --- a/docs/source/guides/run_env.md +++ b/docs/source/guides/run_env.md @@ -80,9 +80,21 @@ embodichain run-env \ After constructing and resetting the environment, the terminal accepts: +- `i`: show or hide the robot IK Gizmo (native single-environment preview); - `p`: enter an IPython session with `env` in scope; - `q`: close the preview. +For a native preview with one environment, `run-env` prepares an IK Gizmo for +each task-selected robot control part that has an IK solver. The controls start +hidden. Focus the DexSim window and press `I` to show or hide them, then drag an +end-effector Gizmo to operate the robot. While preview waits for terminal input, +it continues stepping the simulation so IK targets are applied immediately. +Pressing `i` in the terminal provides the same visibility toggle. + +IK Gizmos require `num_envs=1`, a native window, and solver metadata for the +selected control part. Headless and Viser previews skip this native shortcut; +use Viser's click-to-pick Gizmo interaction in the browser instead. + IPython is required only when entering the embedded session. Install it with `pip install ipython` if the `p` command reports that it is unavailable. diff --git a/docs/source/tutorial/data_generation.rst b/docs/source/tutorial/data_generation.rst index 4dbf80464..7968a856c 100644 --- a/docs/source/tutorial/data_generation.rst +++ b/docs/source/tutorial/data_generation.rst @@ -158,7 +158,7 @@ The recommended CLI entrypoint is: --headless For interactive inspection, you can use preview mode: replace ``--headless`` with ``--preview``. -When ``--preview`` is enabled, the script opens the environment in an interactive debugging mode. This mode is for inspection and does not save datasets. +When ``--preview`` is enabled, the script opens the environment in an interactive debugging mode. In a native single-environment preview, press ``I`` in the DexSim window to show or hide solver-backed robot IK Gizmos and drag them to operate the robot. This mode is for inspection and does not save datasets. For a detailed comparison of preview, structured dataset recording, debug-video recording, trajectory recording, and the three replay modes, see diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index ce9444276..f226aed76 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -22,7 +22,8 @@ import sys import time -from collections.abc import Sequence +from collections.abc import Iterator, Sequence +from contextlib import contextmanager import gymnasium import numpy as np @@ -254,6 +255,22 @@ def read_key(self, timeout: float | None = None) -> str | None: raise EOFError return value.lower() if self.single_key else value.strip().lower() + @contextmanager + def suspend_terminal(self) -> Iterator[None]: + """Restore canonical terminal input while an embedded REPL is active.""" + if self._term_attrs is None or self._fd is None: + yield + return + + import termios + import tty + + termios.tcsetattr(self._fd, termios.TCSADRAIN, self._term_attrs) + try: + yield + finally: + tty.setcbreak(self._fd) + def _read_replay_control_command( control_input: _ReplayControlInput, initial: str | None = None @@ -444,49 +461,187 @@ def main(args, env, gym_config): env.close() -def preview(env: gymnasium.Env) -> None: +def _enable_preview_ik_gizmos( + env: gymnasium.Env, +) -> tuple[tuple[str, str], ...]: + """Create hidden native IK Gizmos for the preview robot. + + Only control parts selected by the environment and backed by an IK solver + are enabled. Existing Gizmos are reused without changing their visibility. + + Args: + env: Gymnasium environment being previewed. + + Returns: + ``(robot_uid, control_part)`` pairs for the available IK Gizmos. """ - Run the following code to create a demonstration and perform env steps. + base_env = env.unwrapped + sim = getattr(base_env, "sim", None) + robot = getattr(base_env, "robot", None) + if sim is None or robot is None: + log_warning("Preview IK Gizmo is unavailable because no robot was found.") + return () + if not bool(getattr(sim, "is_window_opened", False)): + log_warning( + "Preview IK Gizmo requires a native DexSim window and is disabled " + "for headless or Viser preview." + ) + return () - ``` - # Demo version of environment rollout - for i in range(10): - qpos = env.robot.get_qpos() + num_envs = getattr(base_env, "num_envs", None) + if num_envs is None: + num_envs = getattr(robot, "num_instances", 1) + if int(num_envs) != 1: + log_warning( + "Preview IK Gizmo supports exactly one environment; " + f"received num_envs={num_envs}." + ) + return () + + robot_uid = getattr(robot, "uid", None) + control_parts = getattr(robot, "control_parts", None) or {} + get_solver = getattr(robot, "get_solver", None) + if not isinstance(robot_uid, str) or not robot_uid or not callable(get_solver): + log_warning("Preview IK Gizmo requires a named robot with IK control parts.") + return () + + configured_parts = getattr(getattr(base_env, "cfg", None), "control_parts", None) + if configured_parts: + candidate_parts = tuple( + dict.fromkeys(part for part in configured_parts if part in control_parts) + ) + else: + candidate_parts = tuple(control_parts) + ik_parts = tuple(part for part in candidate_parts if get_solver(part) is not None) + if not ik_parts: + log_warning( + f"Robot {robot_uid!r} has no active control part with an IK solver; " + "preview IK Gizmo was not enabled." + ) + return () - obs, reward, terminated, truncated, info = env.step(qpos) + gizmo_keys: list[tuple[str, str]] = [] + for control_part in ik_parts: + if sim.has_gizmo(robot_uid, control_part=control_part): + gizmo_keys.append((robot_uid, control_part)) + continue + gizmo = sim.enable_gizmo( + uid=robot_uid, + control_part=control_part, + enable_native=True, + ) + if gizmo is None: + continue + # Preview starts view-only. The native IKGizmoController owns the I + # hotkey and reveals all newly created targets on the first key press. + sim.set_gizmo_visibility( + robot_uid, + visible=False, + control_part=control_part, + ) + gizmo_keys.append((robot_uid, control_part)) - # reset the environment - env.reset() - ``` + if gizmo_keys: + part_names = ", ".join(part for _, part in gizmo_keys) + log_info( + f"Preview IK Gizmo ready for {robot_uid!r}: {part_names}. " + "Focus the DexSim window and press I to show or hide it.", + color="green", + ) + else: + log_warning(f"Failed to initialize a preview IK Gizmo for {robot_uid!r}.") + return tuple(gizmo_keys) + + +def _toggle_preview_ik_gizmos( + sim: object, + gizmo_keys: Sequence[tuple[str, str]], +) -> tuple[bool, ...]: + """Toggle preview IK Gizmos from the terminal fallback command.""" + states: list[bool] = [] + for robot_uid, control_part in gizmo_keys: + visible = sim.toggle_gizmo_visibility( + robot_uid, + control_part=control_part, + ) + if visible is not None: + states.append(bool(visible)) + return tuple(states) - Run the following code to preview the sensor observations. - ``` - env.preview_sensor_data("camera") - ``` - """ - _, _ = env.reset() +def _run_preview_loop( + env: gymnasium.Env, + control_input: _ReplayControlInput, + gizmo_keys: Sequence[tuple[str, str]], +) -> None: + """Run terminal commands while servicing interactive Gizmos.""" + sim = env.unwrapped.sim + physics_dt = float(sim.sim_config.physics_dt) + visualization = getattr(sim.sim_config, "visualization", None) + service_interactions = bool(gizmo_keys) or ( + getattr(visualization, "backend", "none") == "viser" + ) - end = False - while end is False: - print("Press `p` to enter embed mode to interact with the environment.") - print("Press `q` to quit the simulation.") - txt = input() - if txt == "p": - try: - from IPython import embed - except ImportError: - log_error( - "IPython is not installed. Preview mode requires IPython to be " - "available. Please install it with `pip install ipython` and try again." - ) - continue + print("Preview controls:") + if gizmo_keys: + print(" DexSim window: I=show/hide IK Gizmo, drag Gizmo=move robot") + print(" Terminal: i=show/hide IK Gizmo") + print(" Terminal: p=IPython embed, q=quit") - embed() - elif txt == "q": - end = True + while True: + try: + command = control_input.read_key( + timeout=physics_dt if service_interactions else None + ) + except (EOFError, KeyboardInterrupt): + break - exit(0) + if command is not None: + command = command.strip().lower() + if command in {"q", "quit"}: + break + if command == "p": + try: + from IPython import embed + except ImportError: + log_error( + "IPython is not installed. Preview embed mode requires " + "IPython. Install it with `pip install ipython`." + ) + continue + with control_input.suspend_terminal(): + embed() + elif command == "i" and gizmo_keys: + states = _toggle_preview_ik_gizmos(sim, gizmo_keys) + if states: + state = "shown" if all(states) else "hidden" + log_info(f"Preview IK Gizmo {state}.", color="green") + elif command: + print(f"Unknown preview command: {command!r}") + + if service_interactions: + # SimulationManager.update() invokes update_gizmos() before the + # physics step, allowing native or Viser controllers to apply IK + # drive targets while the terminal remains responsive. + sim.update(physics_dt, step=1) + + +def preview(env: gymnasium.Env) -> None: + """Run an interactive environment preview. + + A native single-environment preview automatically creates hidden IK + Gizmos for the robot's active solver-backed control parts. Press ``I`` in + the DexSim window to show or hide the controls, then drag an end-effector + target to operate the robot. Terminal commands remain available for the + IPython embed session and shutdown. + + Args: + env: Gymnasium environment to reset and preview. + """ + _, _ = env.reset() + gizmo_keys = _enable_preview_ik_gizmos(env) + with _ReplayControlInput() as control_input: + _run_preview_loop(env, control_input, gizmo_keys) def _create_parser() -> argparse.ArgumentParser: diff --git a/tests/lab/scripts/test_run_env.py b/tests/lab/scripts/test_run_env.py index df08e02a3..902a27f08 100644 --- a/tests/lab/scripts/test_run_env.py +++ b/tests/lab/scripts/test_run_env.py @@ -16,7 +16,8 @@ from __future__ import annotations -from unittest.mock import MagicMock +from types import SimpleNamespace +from unittest.mock import MagicMock, call from embodichain.lab.gym.utils.gym_utils import merge_args_with_gym_config from embodichain.lab.scripts import run_env @@ -28,6 +29,18 @@ ACTION_LIST_INDEX = 0 +class _PreviewInput: + """Deterministic input source for the non-blocking preview loop.""" + + def __init__(self, keys: list[str | None]) -> None: + self._keys = iter(keys) + self.timeouts: list[float | None] = [] + + def read_key(self, timeout: float | None = None) -> str | None: + self.timeouts.append(timeout) + return next(self._keys) + + def test_generate_function_displays_episode_and_action_list_indices( monkeypatch, ) -> None: @@ -87,3 +100,115 @@ def test_run_env_preserves_configured_viser_image_fps() -> None: ) assert merged["visualization"]["sensor_image_fps"] == configured_fps + + +def test_preview_enables_hidden_ik_gizmos_for_active_solver_parts() -> None: + """Preview prepares each task-selected arm that has an IK solver.""" + solvers = {"left_arm": object(), "right_arm": object()} + robot = SimpleNamespace( + uid="preview_robot", + control_parts={ + "left_arm": [], + "left_eef": [], + "right_arm": [], + }, + get_solver=MagicMock(side_effect=lambda part: solvers.get(part)), + ) + sim = MagicMock() + sim.is_window_opened = True + sim.has_gizmo.return_value = False + sim.enable_gizmo.side_effect = [object(), object()] + env = SimpleNamespace( + unwrapped=SimpleNamespace( + sim=sim, + robot=robot, + num_envs=1, + cfg=SimpleNamespace(control_parts=["left_arm", "left_eef", "right_arm"]), + ) + ) + + gizmo_keys = run_env._enable_preview_ik_gizmos(env) + + assert gizmo_keys == ( + ("preview_robot", "left_arm"), + ("preview_robot", "right_arm"), + ) + assert sim.enable_gizmo.call_args_list == [ + call(uid="preview_robot", control_part="left_arm", enable_native=True), + call(uid="preview_robot", control_part="right_arm", enable_native=True), + ] + assert sim.set_gizmo_visibility.call_args_list == [ + call("preview_robot", visible=False, control_part="left_arm"), + call("preview_robot", visible=False, control_part="right_arm"), + ] + + +def test_preview_skips_ik_gizmo_for_vectorized_environment() -> None: + """Native IK Gizmos remain limited to one simulated environment.""" + sim = MagicMock() + sim.is_window_opened = True + env = SimpleNamespace( + unwrapped=SimpleNamespace( + sim=sim, + robot=SimpleNamespace(uid="preview_robot"), + num_envs=2, + ) + ) + + gizmo_keys = run_env._enable_preview_ik_gizmos(env) + + assert gizmo_keys == () + sim.enable_gizmo.assert_not_called() + + +def test_preview_loop_services_native_ik_gizmo_while_waiting() -> None: + """Each input timeout advances Gizmo processing and one physics step.""" + physics_dt = 0.02 + sim = MagicMock() + sim.sim_config = SimpleNamespace(physics_dt=physics_dt) + env = SimpleNamespace(unwrapped=SimpleNamespace(sim=sim)) + control_input = _PreviewInput([None, "q"]) + + run_env._run_preview_loop( + env, + control_input, + (("preview_robot", "arm"),), + ) + + sim.update.assert_called_once_with(physics_dt, step=1) + assert control_input.timeouts == [physics_dt, physics_dt] + + +def test_preview_terminal_i_toggles_ik_gizmo() -> None: + """Terminal I mirrors the native-window visibility hotkey.""" + physics_dt = 0.02 + sim = MagicMock() + sim.sim_config = SimpleNamespace(physics_dt=physics_dt) + sim.toggle_gizmo_visibility.return_value = True + env = SimpleNamespace(unwrapped=SimpleNamespace(sim=sim)) + + run_env._run_preview_loop( + env, + _PreviewInput(["i", "q"]), + (("preview_robot", "arm"),), + ) + + sim.toggle_gizmo_visibility.assert_called_once_with( + "preview_robot", + control_part="arm", + ) + + +def test_preview_loop_services_viser_without_native_ik_gizmo() -> None: + """Viser preview keeps processing browser interaction commands.""" + physics_dt = 0.02 + sim = MagicMock() + sim.sim_config = SimpleNamespace( + physics_dt=physics_dt, + visualization=SimpleNamespace(backend="viser"), + ) + env = SimpleNamespace(unwrapped=SimpleNamespace(sim=sim)) + + run_env._run_preview_loop(env, _PreviewInput([None, "q"]), ()) + + sim.update.assert_called_once_with(physics_dt, step=1) From 77163dc949488ea1f197e6c86edc15f6146aa2dc Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 12 Aug 2026 15:25:57 +0800 Subject: [PATCH 04/14] wip --- .../embodichain.lab.sim.objects.rst | 3 + .../embodichain.lab.sim.utility.rst | 7 - docs/source/features/interaction/window.md | 34 +- docs/source/guides/cli.md | 3 - docs/source/guides/run_env.md | 12 - .../overview/sim/viser_visualization.md | 2 +- docs/source/tutorial/data_generation.rst | 2 +- docs/source/tutorial/gizmo.rst | 135 ++-- embodichain/lab/scripts/run_env.py | 225 +----- embodichain/lab/sim/objects/__init__.py | 2 +- embodichain/lab/sim/objects/gizmo.py | 724 ++++++------------ embodichain/lab/sim/sim_manager.py | 218 +----- embodichain/lab/sim/utility/__init__.py | 1 - embodichain/lab/sim/utility/gizmo_utils.py | 240 ------ .../lab/visualization/backends/viser.py | 7 + embodichain/lab/visualization/runtime.py | 4 +- examples/sim/gizmo/gizmo_camera.py | 43 +- examples/sim/gizmo/gizmo_object.py | 16 +- examples/sim/gizmo/gizmo_robot.py | 42 +- examples/sim/gizmo/gizmo_scene.py | 40 +- examples/sim/gizmo/gizmo_w1.py | 23 +- scripts/tutorials/sim/gizmo_robot.py | 33 +- tests/lab/scripts/test_run_env.py | 138 +--- tests/sim/objects/test_gizmo.py | 138 ++-- tests/sim/test_sim_manager.py | 181 ++--- tests/visualization/test_runtime.py | 17 + 26 files changed, 634 insertions(+), 1656 deletions(-) delete mode 100644 embodichain/lab/sim/utility/gizmo_utils.py diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst index 3af7583d5..34d8c0545 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst @@ -41,6 +41,7 @@ bodies. RobotWorkspaceCfg Gizmo GizmoCfg + create_robot_ik_gizmo_controller RigidConstraint .. currentmodule:: embodichain.lab.sim.objects @@ -186,6 +187,8 @@ Gizmo :show-inheritance: :exclude-members: __init__, copy, replace, to_dict, validate +.. autofunction:: create_robot_ik_gizmo_controller + Rigid Constraint ---------------- diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.utility.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.utility.rst index 2e45ea5db..234e8f5fe 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.utility.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.utility.rst @@ -17,7 +17,6 @@ action/solver adaptation. action_utils atom_action_utils cfg_utils - gizmo_utils import_utils io_utils keyboard_utils @@ -46,12 +45,6 @@ Configuration Utilities .. automodule:: embodichain.lab.sim.utility.cfg_utils :members: -Gizmo Utilities -~~~~~~~~~~~~~~~ - -.. automodule:: embodichain.lab.sim.utility.gizmo_utils - :members: - Import Utilities ~~~~~~~~~~~~~~~~ diff --git a/docs/source/features/interaction/window.md b/docs/source/features/interaction/window.md index 6110111a3..70404f376 100644 --- a/docs/source/features/interaction/window.md +++ b/docs/source/features/interaction/window.md @@ -53,15 +53,16 @@ The camera-pose hotkey is controlled by `SimulationManagerCfg.window_camera_pose ### Entity Gizmo Control -Opening a non-headless `SimulationManager` window enables dexsim's world-owned -`EntityGizmoManipulator` by default: +DexSim owns native entity selection and manipulation. Enable it explicitly after +opening a native window: ```python import dexsim gizmo_config = dexsim.interaction.EntityGizmoConfig() gizmo_config.max_gizmos = 0 # Unlimited simultaneous bindings. -sim.open_window(entity_gizmo_config=gizmo_config) +sim.open_window() +sim.enable_entity_gizmo(gizmo_config) ``` While enabled, left-click a render mesh, dynamic/kinematic rigid body, or @@ -74,28 +75,19 @@ EmbodiChain's built-in `default_plane` is registered as an immovable target and cannot receive an entity gizmo. Other supported scene entities remain selectable normally. -For a view-only window, opt out explicitly: +`sim.enable_entity_gizmo(config)` is a thin helper that also excludes +EmbodiChain's render-only default plane. All other lifecycle operations stay on +DexSim's world object: ```python -sim.open_window(enable_entity_gizmo=False) +world = sim.get_world() +controller = world.get_entity_gizmo() +world.disable_entity_gizmo() ``` -Set `SimulationManagerCfg.enable_entity_gizmo_on_window_open=False` to change -the default for constructor-opened and subsequently opened windows. Headless -simulations do not create or enable the controller. - -`sim.enable_entity_gizmo(config)` can reconfigure or reactivate the controller -at any time, and `sim.disable_entity_gizmo()` cancels it without closing the -window. The last explicit configuration is restored if the window is closed -and reopened. - -Use `sim.get_entity_gizmo()` to access the native controller and -`sim.has_entity_gizmo()` to query its lifecycle state. Closing the window or -destroying the `SimulationManager` disables it automatically. - -This controller is distinct from the target-specific Robot TCP IK gizmo. When -both are active, **G** controls entity roots and **I** shows or hides the Robot -TCP IK gizmo. +This controller is distinct from DexSim's target-specific Robot TCP IK +controller. When both are active, **G** controls entity roots and **I** shows or +hides the Robot TCP target. The entity gizmo is native-window only. The Viser backend offers an analogous **click-to-pick** flow (an *Enable click-to-pick Gizmo* checkbox instead of the diff --git a/docs/source/guides/cli.md b/docs/source/guides/cli.md index 815f2d132..35cdcda08 100644 --- a/docs/source/guides/cli.md +++ b/docs/source/guides/cli.md @@ -222,9 +222,6 @@ configuration, remote access, and performance details, see When ``--preview`` is enabled, an interactive REPL is available: -- **``i``** — show or hide solver-backed robot IK Gizmos; in the native DexSim - window, **``I``** provides the same toggle and the Gizmos can be dragged to - operate the robot (single-environment previews only) - **``p``** — enter an IPython embed session with ``env`` in scope - **``q``** — quit diff --git a/docs/source/guides/run_env.md b/docs/source/guides/run_env.md index a128e6548..2c79fb181 100644 --- a/docs/source/guides/run_env.md +++ b/docs/source/guides/run_env.md @@ -80,21 +80,9 @@ embodichain run-env \ After constructing and resetting the environment, the terminal accepts: -- `i`: show or hide the robot IK Gizmo (native single-environment preview); - `p`: enter an IPython session with `env` in scope; - `q`: close the preview. -For a native preview with one environment, `run-env` prepares an IK Gizmo for -each task-selected robot control part that has an IK solver. The controls start -hidden. Focus the DexSim window and press `I` to show or hide them, then drag an -end-effector Gizmo to operate the robot. While preview waits for terminal input, -it continues stepping the simulation so IK targets are applied immediately. -Pressing `i` in the terminal provides the same visibility toggle. - -IK Gizmos require `num_envs=1`, a native window, and solver metadata for the -selected control part. Headless and Viser previews skip this native shortcut; -use Viser's click-to-pick Gizmo interaction in the browser instead. - IPython is required only when entering the embedded session. Install it with `pip install ipython` if the `p` command reports that it is unavailable. diff --git a/docs/source/overview/sim/viser_visualization.md b/docs/source/overview/sim/viser_visualization.md index fe3d56ad9..ed7f01a70 100644 --- a/docs/source/overview/sim/viser_visualization.md +++ b/docs/source/overview/sim/viser_visualization.md @@ -126,7 +126,7 @@ each Gizmo through `SimulationManager.enable_gizmo`; a pure browser process can omit the DexSim handle: ```python -sim.enable_gizmo("cube", enable_native=False) +sim.enable_gizmo("cube") ``` Viser and DexSim use the same deferred target-control path: diff --git a/docs/source/tutorial/data_generation.rst b/docs/source/tutorial/data_generation.rst index 9054ef6a8..f89a288ea 100644 --- a/docs/source/tutorial/data_generation.rst +++ b/docs/source/tutorial/data_generation.rst @@ -189,7 +189,7 @@ The recommended CLI entrypoint is: --headless For interactive inspection, you can use preview mode: replace ``--headless`` with ``--preview``. -When ``--preview`` is enabled, the script opens the environment in an interactive debugging mode. In a native single-environment preview, press ``I`` in the DexSim window to show or hide solver-backed robot IK Gizmos and drag them to operate the robot. This mode is for inspection and does not save datasets. +When ``--preview`` is enabled, the script opens the environment in an interactive debugging mode. This mode is for inspection and does not save datasets. For a detailed comparison of preview, structured dataset recording, debug-video recording, trajectory recording, and the three replay modes, see diff --git a/docs/source/tutorial/gizmo.rst b/docs/source/tutorial/gizmo.rst index 41f2bb279..7fef64815 100644 --- a/docs/source/tutorial/gizmo.rst +++ b/docs/source/tutorial/gizmo.rst @@ -5,7 +5,9 @@ Interactive Robot Control with Gizmo .. currentmodule:: embodichain.lab.sim -This tutorial demonstrates how to use the Gizmo class for interactive robot manipulation in SimulationManager. You'll learn how to create a gizmo attached to a robot's end-effector and use it for real-time inverse kinematics (IK) control, allowing intuitive manipulation of robot poses through visual interaction. +This tutorial demonstrates native DexSim and browser-based Viser Gizmo control. +DexSim owns native entity and robot IK controllers; EmbodiChain keeps only the +robot control-part adapter and the Viser command path. The Code ~~~~~~~~ @@ -30,7 +32,8 @@ Similar to the previous tutorial on robot simulation, we use the :class:`Simulat **Important:** Gizmo only supports single environment mode (`num_envs=1`). Using multiple environments will raise an exception. -All gizmo creation, visibility, and destruction operations must be managed via the SimulationManager API: +Viser Gizmo creation, visibility, and destruction are managed through +SimulationManager: .. code-block:: python @@ -40,12 +43,8 @@ All gizmo creation, visibility, and destruction operations must be managed via t # Set visibility explicitly sim.set_gizmo_visibility("ur10_gizmo_test", visible=False, control_part="arm") -Always use the SimulationManager API to control gizmo visibility and lifecycle. Do not operate on the Gizmo instance directly. - -The same target behavior is available in either the DexSim window or Viser. -Robot Gizmos solve IK with DexSim Newton IK in both modes; the only difference -is the input source (a native window gizmo handle vs a Viser transform control). -The standard Viser mode includes interactive Gizmo control: +Native controls use DexSim directly. The standard Viser mode includes +interactive Gizmo control: .. code-block:: bash @@ -85,7 +84,8 @@ A Gizmo is an interactive visual tool that allows users to manipulate simulation - **Real-time Manipulation**: Provide immediate visual feedback during robot motion planning - **Debugging and Visualization**: Test robot reachability and workspace limits -The :class:`objects.Gizmo` class provides a unified interface for interactive control of different simulation elements including robots, rigid objects, and cameras. +The :class:`objects.Gizmo` class is the Viser-side target controller for robots, +rigid objects, and cameras. Native controls are DexSim controllers. Setting up Robot Configuration ------------------------------ @@ -104,54 +104,59 @@ Key components of the robot configuration: - **IK Solver**: :class:`solvers.PinkSolverCfg` provides inverse kinematics capabilities - **Drive Properties**: Sets stiffness and damping for joint control -The configured EmbodiChain solver is optional: it only supplies default IK chain -metadata (root link, end link, and TCP transform) to the Gizmo. IK itself is -always solved by DexSim Newton IK in both native and Viser modes. A native-only -or Viser-only application may instead set the root link, end link, and optional -TCP transform directly in :class:`objects.GizmoCfg` without configuring an -EmbodiChain solver. +The configured EmbodiChain solver is optional: it supplies default IK-chain +metadata (root link, end link, and TCP transform). IK itself is solved by +DexSim Newton IK. Applications may instead set this metadata directly in +:class:`objects.GizmoCfg`. Creating and Attaching a Gizmo ------------------------------- -After configuring the robot, enable the gizmo for interactive control using the SimulationManager API (supports robot, rigid object, camera; key is `uid:control_part`): +For native-window robot control, create DexSim's IK controller through the +small EmbodiChain adapter factory and retain both returned objects: .. code-block:: python - from embodichain.lab.sim.objects import GizmoCfg + from embodichain.lab.sim.objects import ( + GizmoCfg, + create_robot_ik_gizmo_controller, + ) - # Enable gizmo for the robot's arm - sim.enable_gizmo( - uid="ur10_gizmo_test", + ik_controller, input_controller = create_robot_ik_gizmo_controller( + robot, control_part="arm", - gizmo_cfg=GizmoCfg( + cfg=GizmoCfg( ik_root_link_name="base_link", ik_end_link_name="ee_link", ), - enable_native=native_window_opened, + world=sim.get_world(), ) - if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): - logger.log_error("Failed to enable gizmo!") - return - - -The Gizmo instance is managed internally by SimulationManager. If you need to access it: +Call ``ik_controller.update()`` once per frame. DexSim owns the native target, +hotkey, solve trigger, and visibility state. For Viser, use the SimulationManager +command path instead: .. code-block:: python - gizmo = sim.get_gizmo("ur10_gizmo_test", control_part="arm") + sim.enable_gizmo( + "ur10_gizmo_test", + control_part="arm", + gizmo_cfg=GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ), + ) The Gizmo system will automatically: -1. **Detect Target Type**: Identify that the target is a robot (vs. rigid object or camera) -2. **Resolve the IK Chain**: Locate the root and end-effector links -3. **Select the Backend**: Build a DexSim Newton IK solver; ``enable_native`` only decides whether a native window gizmo handle is created for direct interaction or Viser commands drive the same solver -4. **Defer Simulation Writes**: Apply IK drive targets from the simulation update loop +1. **Resolve the IK Chain**: Locate the root and end-effector links +2. **Build Newton IK**: Construct DexSim's reduced-chain solver +3. **Bridge State**: Map one EmbodiChain control part to DexSim's joint API +4. **Own the Frontend**: DexSim owns native interaction; SimulationManager owns Viser commands How Gizmo-Robot Interaction Works ---------------------------------- @@ -161,34 +166,32 @@ How Gizmo-Robot Interaction Works The gizmo-robot interaction follows this workflow: 1. **Target Update**: DexSim or Viser records the requested TCP transform -2. **Deferred Solve**: ``sim.update_gizmos()`` invokes the DexSim Newton IK solver only when needed +2. **Deferred Solve**: the native controller or ``sim.update_gizmos()`` invokes Newton IK only when needed 3. **State Bridge**: Newton IK reads and writes the selected EmbodiChain control-part joints through an adapter -4. **Drive Target**: Both native and Viser solutions use ``Robot.set_qpos(..., target=True)`` to drive the joint targets +4. **Drive Target**: Both paths use ``Robot.set_qpos(..., target=True)`` 5. **Robot Motion**: Joint drives move the robot toward the target without teleporting its current state -Native robot Gizmos do not create an EmbodiChain proxy cube. Camera Gizmos -retain their proxy path, while rigid-object Gizmos follow their selected object -directly. - The Simulation Loop ------------------- -In the main loop, simply call `sim.update_gizmos()`. There is no need to manually update any Gizmo instance. +Update the DexSim native controller explicitly, then service any Viser controls: .. code-block:: python - def run_simulation(sim: SimulationManager): + def run_simulation(sim: SimulationManager, ik_controller=None): step_count = 0 try: last_time = time.time() last_step = 0 while True: time.sleep(0.033) # 30Hz - sim.update_gizmos() # Update all gizmos + if ik_controller is not None: + ik_controller.update() + sim.update_gizmos() # Update Viser gizmos sim.capture_visualization_safely() # Publish Viser state, if enabled step_count += 1 # ...performance statistics, etc... @@ -202,8 +205,9 @@ In the main loop, simply call `sim.update_gizmos()`. There is no need to manuall Main loop highlights: -- **Gizmo update**: Only `sim.update_gizmos()` is needed, no `gizmo.update()` -- **Viser update**: Automatic-physics loops also call `sim.capture_visualization_safely()` +- **Native update**: Call DexSim's ``IKGizmoController.update()`` each frame +- **Viser command update**: Call ``sim.update_gizmos()`` +- **Viser frame update**: Automatic-physics loops also call ``sim.capture_visualization_safely()`` - **Performance monitoring**: Optional FPS statistics - **Resource cleanup**: Only `sim.destroy()` is needed, no manual Gizmo destruction - **Graceful shutdown**: Supports Ctrl+C interruption @@ -214,43 +218,18 @@ Gizmo Lifecycle Management -Gizmo lifecycle is managed by SimulationManager: +Viser Gizmo lifecycle is managed by SimulationManager: - Enable: `sim.enable_gizmo(...)` -- Update: Main loop automatically calls `sim.update_gizmos()` +- Update: Call ``sim.update_gizmos()`` from the main loop - Destroy/disable: `sim.disable_gizmo(...)` or `sim.destroy()` (recommended) -There is no need to manually create or destroy Gizmo instances. All resources are managed by SimulationManager. - -Available Gizmo Methods ------------------------ - - - - -If you need to access the underlying Gizmo instance (via `sim.get_gizmo`), you can use the following methods: - -**Transform Control:** - -- ``set_world_pose(pose)``: Set gizmo world position and orientation -- ``get_world_pose()``: Get current gizmo world transform -- ``set_local_pose(pose)``: Set gizmo local transform relative to parent -- ``get_local_pose()``: Get gizmo local transform - - - -**Visual properties (strongly recommend using SimulationManager API):** +Native controller lifecycle remains in DexSim. Viser visual properties are +available through SimulationManager: - ``sim.toggle_gizmo_visibility(uid, control_part=None)``: Toggle gizmo visibility - ``sim.set_gizmo_visibility(uid, visible, control_part=None)``: Set gizmo visibility -**Hierarchy Management:** - -- ``get_parent()``: Get gizmo's parent node in scene hierarchy -- ``get_name()``: Get gizmo node name for debugging -- ``detach()``: Disconnect gizmo from current target -- ``attach(target)``: Attach gizmo to a new simulation object - Running the Tutorial -------------------- @@ -270,7 +249,7 @@ Command-line options: Once running: -1. **Mouse Interaction**: Click and drag the gizmo (colorful axes) to move the robot +1. **Mouse Interaction**: Click and drag the gizmo to move the robot 2. **Real-time IK**: Watch the robot joints automatically adjust to follow the gizmo 3. **Workspace Limits**: Observe how the robot behaves at workspace boundaries 4. **Performance**: Monitor FPS in the console output @@ -282,7 +261,8 @@ Tips and Best Practices **Performance optimization:** -- Only call ``sim.update_gizmos()`` in the main loop, no need for ``gizmo.update()`` +- Call the native IK controller's ``update()`` once per frame; call + ``sim.update_gizmos()`` for Viser - Reduce IK solver iterations for better real-time performance if needed - Use ``set_manual_update(False)`` for smoother interaction @@ -291,7 +271,7 @@ Tips and Best Practices **Debugging tips:** - Check console output for IK solver success/failure messages -- Use ``get_world_pose()`` to check gizmo position (if needed) +- Inspect the robot TCP or Viser target pose when debugging alignment - Monitor FPS to identify performance bottlenecks @@ -306,7 +286,7 @@ Tips and Best Practices **Visualization customization:** -- Adjust gizmo appearance via Gizmo config (e.g., ``set_line_width()``; requires access to the instance via `sim.get_gizmo`) +- Adjust Viser axis lengths, ring radius, and line width through :class:`objects.GizmoCfg` - Adjust gizmo scale according to robot size - Enable collision for debugging if needed @@ -316,7 +296,6 @@ Next Steps After mastering basic gizmo usage, you can explore: - **Multi-robot Gizmos**: Attach gizmos to multiple robots simultaneously -- **Custom Gizmo Callbacks**: Implement application-specific interaction logic - **Gizmo with Rigid Objects**: Use gizmos for interactive object manipulation - **Advanced IK Configuration**: Fine-tune solver parameters for specific robots diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index af814758b..880d040fc 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -22,8 +22,7 @@ import sys import time -from collections.abc import Iterable, Iterator, Sequence -from contextlib import contextmanager +from collections.abc import Iterable, Sequence from typing import TYPE_CHECKING, Any import gymnasium @@ -415,22 +414,6 @@ def read_key(self, timeout: float | None = None) -> str | None: raise EOFError return value.lower() if self.single_key else value.strip().lower() - @contextmanager - def suspend_terminal(self) -> Iterator[None]: - """Restore canonical terminal input while an embedded REPL is active.""" - if self._term_attrs is None or self._fd is None: - yield - return - - import termios - import tty - - termios.tcsetattr(self._fd, termios.TCSADRAIN, self._term_attrs) - try: - yield - finally: - tty.setcbreak(self._fd) - def _read_replay_control_command( control_input: _ReplayControlInput, initial: str | None = None @@ -703,187 +686,49 @@ def main(args: Any, env: Any, gym_config: dict[str, Any]) -> None: ) -def _enable_preview_ik_gizmos( - env: gymnasium.Env, -) -> tuple[tuple[str, str], ...]: - """Create hidden native IK Gizmos for the preview robot. - - Only control parts selected by the environment and backed by an IK solver - are enabled. Existing Gizmos are reused without changing their visibility. - - Args: - env: Gymnasium environment being previewed. - - Returns: - ``(robot_uid, control_part)`` pairs for the available IK Gizmos. +def preview(env: gymnasium.Env) -> None: """ - base_env = env.unwrapped - sim = getattr(base_env, "sim", None) - robot = getattr(base_env, "robot", None) - if sim is None or robot is None: - log_warning("Preview IK Gizmo is unavailable because no robot was found.") - return () - if not bool(getattr(sim, "is_window_opened", False)): - log_warning( - "Preview IK Gizmo requires a native DexSim window and is disabled " - "for headless or Viser preview." - ) - return () - - num_envs = getattr(base_env, "num_envs", None) - if num_envs is None: - num_envs = getattr(robot, "num_instances", 1) - if int(num_envs) != 1: - log_warning( - "Preview IK Gizmo supports exactly one environment; " - f"received num_envs={num_envs}." - ) - return () - - robot_uid = getattr(robot, "uid", None) - control_parts = getattr(robot, "control_parts", None) or {} - get_solver = getattr(robot, "get_solver", None) - if not isinstance(robot_uid, str) or not robot_uid or not callable(get_solver): - log_warning("Preview IK Gizmo requires a named robot with IK control parts.") - return () - - configured_parts = getattr(getattr(base_env, "cfg", None), "control_parts", None) - if configured_parts: - candidate_parts = tuple( - dict.fromkeys(part for part in configured_parts if part in control_parts) - ) - else: - candidate_parts = tuple(control_parts) - ik_parts = tuple(part for part in candidate_parts if get_solver(part) is not None) - if not ik_parts: - log_warning( - f"Robot {robot_uid!r} has no active control part with an IK solver; " - "preview IK Gizmo was not enabled." - ) - return () - - gizmo_keys: list[tuple[str, str]] = [] - for control_part in ik_parts: - if sim.has_gizmo(robot_uid, control_part=control_part): - gizmo_keys.append((robot_uid, control_part)) - continue - gizmo = sim.enable_gizmo( - uid=robot_uid, - control_part=control_part, - enable_native=True, - ) - if gizmo is None: - continue - # Preview starts view-only. The native IKGizmoController owns the I - # hotkey and reveals all newly created targets on the first key press. - sim.set_gizmo_visibility( - robot_uid, - visible=False, - control_part=control_part, - ) - gizmo_keys.append((robot_uid, control_part)) + Run the following code to create a demonstration and perform env steps. - if gizmo_keys: - part_names = ", ".join(part for _, part in gizmo_keys) - log_info( - f"Preview IK Gizmo ready for {robot_uid!r}: {part_names}. " - "Focus the DexSim window and press I to show or hide it.", - color="green", - ) - else: - log_warning(f"Failed to initialize a preview IK Gizmo for {robot_uid!r}.") - return tuple(gizmo_keys) - - -def _toggle_preview_ik_gizmos( - sim: object, - gizmo_keys: Sequence[tuple[str, str]], -) -> tuple[bool, ...]: - """Toggle preview IK Gizmos from the terminal fallback command.""" - states: list[bool] = [] - for robot_uid, control_part in gizmo_keys: - visible = sim.toggle_gizmo_visibility( - robot_uid, - control_part=control_part, - ) - if visible is not None: - states.append(bool(visible)) - return tuple(states) + ``` + # Demo version of environment rollout + for i in range(10): + qpos = env.robot.get_qpos() + obs, reward, terminated, truncated, info = env.step(qpos) -def _run_preview_loop( - env: gymnasium.Env, - control_input: _ReplayControlInput, - gizmo_keys: Sequence[tuple[str, str]], -) -> None: - """Run terminal commands while servicing interactive Gizmos.""" - sim = env.unwrapped.sim - physics_dt = float(sim.sim_config.physics_dt) - visualization = getattr(sim.sim_config, "visualization", None) - service_interactions = bool(gizmo_keys) or ( - getattr(visualization, "backend", "none") == "viser" - ) + # reset the environment + env.reset() + ``` - print("Preview controls:") - if gizmo_keys: - print(" DexSim window: I=show/hide IK Gizmo, drag Gizmo=move robot") - print(" Terminal: i=show/hide IK Gizmo") - print(" Terminal: p=IPython embed, q=quit") - - while True: - try: - command = control_input.read_key( - timeout=physics_dt if service_interactions else None - ) - except (EOFError, KeyboardInterrupt): - break - - if command is not None: - command = command.strip().lower() - if command in {"q", "quit"}: - break - if command == "p": - try: - from IPython import embed - except ImportError: - log_error( - "IPython is not installed. Preview embed mode requires " - "IPython. Install it with `pip install ipython`." - ) - continue - with control_input.suspend_terminal(): - embed() - elif command == "i" and gizmo_keys: - states = _toggle_preview_ik_gizmos(sim, gizmo_keys) - if states: - state = "shown" if all(states) else "hidden" - log_info(f"Preview IK Gizmo {state}.", color="green") - elif command: - print(f"Unknown preview command: {command!r}") - - if service_interactions: - # SimulationManager.update() invokes update_gizmos() before the - # physics step, allowing native or Viser controllers to apply IK - # drive targets while the terminal remains responsive. - sim.update(physics_dt, step=1) + Run the following code to preview the sensor observations. + ``` + env.preview_sensor_data("camera") + ``` + """ + _, _ = env.reset() -def preview(env: gymnasium.Env) -> None: - """Run an interactive environment preview. + end = False + while end is False: + print("Press `p` to enter embed mode to interact with the environment.") + print("Press `q` to quit the simulation.") + txt = input() + if txt == "p": + try: + from IPython import embed + except ImportError: + log_error( + "IPython is not installed. Preview mode requires IPython to be " + "available. Please install it with `pip install ipython` and try again." + ) + continue - A native single-environment preview automatically creates hidden IK - Gizmos for the robot's active solver-backed control parts. Press ``I`` in - the DexSim window to show or hide the controls, then drag an end-effector - target to operate the robot. Terminal commands remain available for the - IPython embed session and shutdown. + embed() + elif txt == "q": + end = True - Args: - env: Gymnasium environment to reset and preview. - """ - _, _ = env.reset() - gizmo_keys = _enable_preview_ik_gizmos(env) - with _ReplayControlInput() as control_input: - _run_preview_loop(env, control_input, gizmo_keys) + return def _create_parser() -> argparse.ArgumentParser: diff --git a/embodichain/lab/sim/objects/__init__.py b/embodichain/lab/sim/objects/__init__.py index 4a5868785..838023f8f 100644 --- a/embodichain/lab/sim/objects/__init__.py +++ b/embodichain/lab/sim/objects/__init__.py @@ -33,7 +33,7 @@ from .articulation import Articulation, ArticulationData, ArticulationCfg from .robot import Robot, RobotCfg, RobotWorkspaceCfg from .light import Light, LightCfg -from .gizmo import Gizmo, GizmoCfg +from .gizmo import Gizmo, GizmoCfg, create_robot_ik_gizmo_controller from .constraint import RigidConstraint diff --git a/embodichain/lab/sim/objects/gizmo.py b/embodichain/lab/sim/objects/gizmo.py index 43ce04528..cb0352c7e 100644 --- a/embodichain/lab/sim/objects/gizmo.py +++ b/embodichain/lab/sim/objects/gizmo.py @@ -14,27 +14,18 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Backend-neutral Gizmo target control with an optional DexSim handle.""" +"""Viser Gizmo control and the EmbodiChain-to-DexSim robot IK adapter.""" from __future__ import annotations import threading -from collections.abc import Callable -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import dexsim import numpy as np import torch import warp as wp -from dexsim.types import ( - AxisArrowType, - AxisCornerType, - AxisOption, - AxisTagType, - InputKey, - RotationRingsOption, -) -from scipy.spatial.transform import Rotation +from dexsim.types import InputKey from embodichain.lab.sim.common import BatchEntity from embodichain.lab.sim.objects.rigid_object import RigidObject @@ -43,14 +34,15 @@ from embodichain.utils import configclass, logger if TYPE_CHECKING: + from dexsim.engine import GizmoController from dexsim.kit.ik import IKGizmoController, NewtonChainIK -__all__ = ["Gizmo", "GizmoCfg"] +__all__ = ["Gizmo", "GizmoCfg", "create_robot_ik_gizmo_controller"] @configclass class GizmoCfg: - """Configure Gizmo appearance and robot IK behavior.""" + """Configure Viser Gizmo appearance and robot IK behavior.""" axis_length_x: float = 0.2 """Length of the X-axis arrow.""" @@ -62,67 +54,34 @@ class GizmoCfg: """Length of the Z-axis arrow.""" axis_size: float = 0.01 - """Thickness of the native axis lines.""" - - arrow_type: AxisArrowType = AxisArrowType.CONE - """Native axis arrow-head style.""" - - corner_type: AxisCornerType = AxisCornerType.SPHERE - """Native axis corner style.""" - - tag_type: AxisTagType = AxisTagType.PLANE - """Native axis-label style.""" + """Thickness of the Viser axis lines.""" rings_radius: float = 0.15 """Radius of the rotation rings.""" rings_size: float = 0.01 - """Thickness of the native rotation rings.""" + """Thickness of the rotation rings.""" ik_root_link_name: str | None = None - """Robot IK chain root link. - - When omitted, the selected control part's EmbodiChain solver supplies it. - """ + """Robot IK chain root link, or the configured solver root when omitted.""" ik_end_link_name: str | None = None - """Robot IK chain end link. - - When omitted, the selected control part's EmbodiChain solver supplies it. - """ + """Robot IK chain end link, or the configured solver end when omitted.""" ik_tcp_pose: torch.Tensor | np.ndarray | list[list[float]] | None = None - """End-link-to-TCP transform used by native DexSim robot IK.""" + """End-link-to-TCP transform.""" ik_iterations: int = 24 - """Number of Newton IK iterations per changed native target.""" + """Number of Newton IK iterations per changed target.""" ik_device: str | None = None - """Warp device for native Newton IK, or the robot device when omitted.""" + """Warp device for Newton IK, or the robot device when omitted.""" ik_gizmo_scale: float = 1.5 - """Isotropic scale of the native robot IK target Gizmo.""" + """Isotropic scale of a native DexSim robot IK target.""" ik_toggle_key: InputKey = InputKey.SCANCODE_I - """Native-window key used to toggle the robot IK Gizmo.""" - - def to_options_dict(self) -> dict[str, AxisOption | RotationRingsOption]: - """Convert this configuration to DexSim Gizmo options.""" - return { - "axis": AxisOption( - lx=self.axis_length_x, - ly=self.axis_length_y, - lz=self.axis_length_z, - size=self.axis_size, - arrow_type=self.arrow_type, - corner_type=self.corner_type, - tag_type=self.tag_type, - ), - "rings": RotationRingsOption( - radius=self.rings_radius, - size=self.rings_size, - ), - } + """Native-window key used to toggle a DexSim robot IK target.""" class _RobotGizmoAdapter: @@ -212,82 +171,194 @@ def _set_qpos(self, qpos: np.ndarray, target: bool) -> None: ) +def _resolve_control_part(robot: Robot, control_part: str | None) -> str: + part_names = list(robot.control_parts or {}) + if not part_names: + raise ValueError("Robot has no control parts defined.") + if control_part is None: + return part_names[0] + if control_part not in part_names: + raise ValueError( + f"Control part {control_part!r} was not found; available parts are " + f"{part_names}." + ) + return control_part + + +def _resolve_robot_ik_chain( + robot: Robot, + control_part: str, + cfg: GizmoCfg, +) -> tuple[str, str, np.ndarray]: + solver = ( + robot.get_solver(control_part) if robot.cfg.solver_cfg is not None else None + ) + root_link = cfg.ik_root_link_name or getattr(solver, "root_link_name", None) + end_link = cfg.ik_end_link_name or getattr(solver, "end_link_name", None) + if not root_link or not end_link: + raise ValueError( + "Robot Gizmo needs an IK chain. Set GizmoCfg.ik_root_link_name and " + "GizmoCfg.ik_end_link_name, or configure a solver for the control part." + ) + + tcp_pose = cfg.ik_tcp_pose + if tcp_pose is None and solver is not None: + tcp_pose = solver.get_tcp() + if tcp_pose is None: + tcp_pose = np.eye(4, dtype=np.float32) + if isinstance(tcp_pose, torch.Tensor): + tcp_pose = tcp_pose.detach().cpu().numpy() + tcp_matrix = np.asarray(tcp_pose, dtype=np.float32) + if tcp_matrix.shape != (4, 4) or not np.isfinite(tcp_matrix).all(): + raise ValueError("ik_tcp_pose must be a finite 4x4 transform.") + return root_link, end_link, tcp_matrix + + +def _build_robot_ik( + robot: Robot, + control_part: str, + cfg: GizmoCfg, +) -> tuple[_RobotGizmoAdapter, NewtonChainIK, str, np.ndarray]: + from dexsim.kit.ik import NewtonChainIK, build_newton_model_from_urdf + + if robot.num_instances != 1: + raise RuntimeError( + "Robot Gizmo supports exactly one environment, " + f"but the robot has {robot.num_instances} instances." + ) + if cfg.ik_iterations <= 0: + raise ValueError("ik_iterations must be greater than zero.") + + root_link, end_link, tcp_pose = _resolve_robot_ik_chain( + robot, + control_part, + cfg, + ) + adapter = _RobotGizmoAdapter(robot, control_part) + with wp.ScopedDevice(cfg.ik_device or str(robot.device)): + model = build_newton_model_from_urdf(robot.cfg.fpath, hide_visuals=True) + solver = NewtonChainIK( + model, + start_link=root_link, + end_link=end_link, + iterations=cfg.ik_iterations, + tcp_pose=tcp_pose, + ) + + solver.set_qpos_from_joint_names( + adapter.get_actived_joint_names(), + adapter.get_current_qpos(), + ) + solver.sync_target_state_from_link(adapter, adapter.get_world_pose()) + logger.log_info( + f"Robot Gizmo uses DexSim Newton IK for control part {control_part!r} " + f"({root_link} -> {end_link})." + ) + return adapter, solver, end_link, tcp_pose + + +def create_robot_ik_gizmo_controller( + robot: Robot, + control_part: str = "arm", + cfg: GizmoCfg | None = None, + *, + world: dexsim.World | None = None, +) -> tuple[IKGizmoController, GizmoController]: + """Create DexSim's native IK controller for one EmbodiChain control part. + + The caller owns the returned objects and must call ``controller.update()`` + once per frame. Retain the input controller while the native window exists. + + Args: + robot: Single-instance EmbodiChain robot. + control_part: Robot control part driven by IK. + cfg: IK chain and target appearance settings. + world: DexSim world, or the current default world when omitted. + + Returns: + ``(ik_controller, input_controller)`` owned by the caller. + """ + from dexsim.engine import GizmoController + from dexsim.kit.ik import IKApplyMode, IKGizmoController + + cfg = cfg or GizmoCfg() + if not np.isfinite(cfg.ik_gizmo_scale) or cfg.ik_gizmo_scale <= 0: + raise ValueError("ik_gizmo_scale must be positive and finite.") + if world is None: + world = dexsim.default_world() + if world is None: + raise RuntimeError("A DexSim world must exist before creating an IK Gizmo.") + window = world.get_windows() + if window is None: + raise RuntimeError("A native DexSim window is required for an IK Gizmo.") + + control_part = _resolve_control_part(robot, control_part) + adapter, solver, _, _ = _build_robot_ik(robot, control_part, cfg) + input_controller = GizmoController() + window.add_input_control(input_controller) + controller = IKGizmoController( + world, + adapter, + solver, + base_state={"pose": adapter.get_world_pose()}, + toggle_key=cfg.ik_toggle_key, + follow_robot_base=True, + apply_mode=IKApplyMode.DRIVE_TARGET, + gizmo_scale=cfg.ik_gizmo_scale, + name=f"{robot.uid}_{control_part}_ik", + ) + return controller, input_controller + + class Gizmo: - """Control a rigid object, robot end-effector, or camera. + """Apply Viser Gizmo commands to one simulation target. - Target mutation is backend-neutral: both DexSim callbacks and Viser commands - submit a local target pose, and :meth:`update` applies it on the simulation - thread. A DexSim Gizmo handle is optional, which allows the same controller - to work in a headless Viser process. + Native-window entity manipulation is owned by DexSim. Use + :meth:`SimulationManager.enable_entity_gizmo` for entity roots and + :func:`create_robot_ik_gizmo_controller` for a native robot TCP target. .. attention:: - Gizmo control currently supports exactly one simulation environment. + Viser Gizmo control supports exactly one simulation environment. Args: - target: Simulation element controlled by this Gizmo. - cfg: Appearance configuration. + target: Rigid object, robot, or camera controlled from Viser. + cfg: Viser appearance and robot IK configuration. control_part: Robot control part used for FK and IK. - enable_native: Whether to create a native DexSim Gizmo handle. Robot - Gizmos solve IK with DexSim Newton IK in both native and headless - (Viser) modes; ``enable_native`` only controls whether a native - window Gizmo handle is created for direct interaction. """ def __init__( self, target: BatchEntity, cfg: GizmoCfg | None = None, - control_part: str | None = "arm", - *, - enable_native: bool = True, + control_part: str | None = None, ) -> None: - world = dexsim.default_world() - if world is None: - raise RuntimeError("A DexSim world must exist before creating a Gizmo.") - - num_envs = int(getattr(target, "num_instances", dexsim.get_world_num())) - if num_envs > 1: + if target.num_instances != 1: raise RuntimeError( - "Gizmo can only be used in single environment mode " - f"(num_envs=1), but target has {num_envs} instances." + "Viser Gizmo supports exactly one environment, " + f"but the target has {target.num_instances} instances." ) self.target: BatchEntity | None = target self.cfg = cfg or GizmoCfg() - self._world = world - self._control_part = control_part self._target_type = self._detect_target_type(target) - self._env = world.get_env() - self._enable_native = enable_native - self._gizmo: object | None = None - self._proxy_cube: object | None = None - self._callback: Callable[..., Any] | None = None + self._control_part = control_part self._is_visible = True self._state_lock = threading.RLock() self._interaction_owner: str | None = None self._pending_target_transform: torch.Tensor | None = None self._desired_target_transform: torch.Tensor | None = None - self._robot_arm_name: str | None = None - self._ik_model: object | None = None self._ik_solver: NewtonChainIK | None = None - self._ik_controller: IKGizmoController | None = None self._robot_adapter: _RobotGizmoAdapter | None = None - self._native_robot_end_link: str | None = None - self._native_robot_tcp_pose: np.ndarray | None = None + self._robot_end_link: str | None = None + self._robot_tcp_pose: np.ndarray | None = None if self._target_type == "robot": - self._configure_robot() + self._control_part = _resolve_control_part(target, control_part) self._setup_robot_ik_solver() - if enable_native: - self._setup_native_robot_gizmo() - self._desired_target_transform = self._read_native_robot_pose() + self._desired_target_transform = self._read_robot_pose() else: self._desired_target_transform = self._read_target_pose() - if enable_native and self._target_type != "robot": - self._gizmo = self._create_native_gizmo(self.cfg) - self._setup_native_gizmo() - @property def target_type(self) -> str: """Return ``rigid_object``, ``robot``, or ``camera``.""" @@ -298,11 +369,6 @@ def control_part(self) -> str | None: """Return the robot control part, if applicable.""" return self._control_part - @property - def native_enabled(self) -> bool: - """Whether this controller owns a DexSim Gizmo handle.""" - return self._gizmo is not None - def _detect_target_type(self, target: BatchEntity) -> str: if isinstance(target, Robot): return "robot" @@ -315,159 +381,34 @@ def _detect_target_type(self, target: BatchEntity) -> str: "RigidObject, Robot, or Camera." ) - def _configure_robot(self) -> None: - if self.target is None or not isinstance(self.target, Robot): - raise RuntimeError("Robot Gizmo has no attached Robot.") - arm_names = list(self.target.control_parts.keys()) - if not arm_names: - raise ValueError("Robot has no control parts defined.") - if self._control_part is None: - self._robot_arm_name = arm_names[0] - self._control_part = self._robot_arm_name - elif self._control_part in arm_names: - self._robot_arm_name = self._control_part - else: - raise ValueError( - f"Control part {self._control_part!r} was not found; " - f"available parts are {arm_names}." - ) - def _setup_robot_ik_solver(self) -> None: - """Build the shared DexSim Newton IK solver and robot adapter. - - The solver is shared by native and headless (Viser) robot Gizmos so both - paths solve IK with DexSim Newton IK instead of an EmbodiChain solver. - Native Gizmos additionally create an :class:`IKGizmoController` in - :meth:`_setup_native_robot_gizmo`. - """ - try: - from dexsim.kit.ik import NewtonChainIK, build_newton_model_from_urdf - except ImportError as error: - raise RuntimeError( - "Robot Gizmo requires a DexSim build that exports " - "NewtonChainIK and build_newton_model_from_urdf." - ) from error - if self.target is None or not isinstance(self.target, Robot): raise RuntimeError("Robot Gizmo has no attached Robot.") - if self._robot_arm_name is None: + if self._control_part is None: raise RuntimeError("Robot Gizmo control part is not configured.") - if self.cfg.ik_iterations <= 0: - raise ValueError("ik_iterations must be greater than zero.") - - root_link, end_link, tcp_pose = self._resolve_robot_ik_chain(self.target) - adapter = _RobotGizmoAdapter(self.target, self._robot_arm_name) - ik_device = self.cfg.ik_device or str(self.target.device) - with wp.ScopedDevice(ik_device): - ik_model = build_newton_model_from_urdf( - self.target.cfg.fpath, - hide_visuals=True, - ) - ik_solver = NewtonChainIK( - ik_model, - start_link=root_link, - end_link=end_link, - iterations=self.cfg.ik_iterations, - tcp_pose=tcp_pose, - ) - - ik_solver.set_qpos_from_joint_names( - adapter.get_actived_joint_names(), - adapter.get_current_qpos(), + adapter, solver, end_link, tcp_pose = _build_robot_ik( + self.target, + self._control_part, + self.cfg, ) - base_pose = adapter.get_world_pose() - ik_solver.sync_target_state_from_link(adapter, base_pose) - self._robot_adapter = adapter - self._ik_model = ik_model - self._ik_solver = ik_solver - self._native_robot_end_link = end_link - self._native_robot_tcp_pose = tcp_pose - logger.log_info( - f"Robot Gizmo uses DexSim Newton IK for control part " - f"{self._robot_arm_name!r} ({root_link} -> {end_link})." + self._ik_solver = solver + self._robot_end_link = end_link + self._robot_tcp_pose = tcp_pose + + def _read_robot_pose(self) -> torch.Tensor: + if ( + self._robot_adapter is None + or self._robot_end_link is None + or self._robot_tcp_pose is None + ): + raise RuntimeError("Robot Gizmo IK is not configured.") + link_pose = self._robot_adapter.get_link_pose(self._robot_end_link) + return self._as_pose_matrix( + link_pose @ self._robot_tcp_pose, + self._target_device(), ) - def _setup_native_robot_gizmo(self) -> None: - """Create DexSim's native IK controller on top of the shared solver.""" - try: - from dexsim.kit.ik import IKApplyMode, IKGizmoController - except ImportError as error: - raise RuntimeError( - "Robot Gizmo requires a DexSim build that exports " - "IKGizmoController and IKApplyMode." - ) from error - - if self._ik_solver is None or self._robot_adapter is None: - raise RuntimeError("Robot Gizmo IK solver is not configured.") - if self.target is None or not isinstance(self.target, Robot): - raise RuntimeError("Robot Gizmo has no attached Robot.") - if not np.isfinite(self.cfg.ik_gizmo_scale) or self.cfg.ik_gizmo_scale <= 0: - raise ValueError("ik_gizmo_scale must be positive and finite.") - - base_pose = self._robot_adapter.get_world_pose() - target_name = getattr(self.target.cfg, "uid", "robot") - ik_controller = IKGizmoController( - self._world, - self._robot_adapter, - self._ik_solver, - base_state={"pose": base_pose}, - toggle_key=self.cfg.ik_toggle_key, - follow_robot_base=True, - apply_mode=IKApplyMode.DRIVE_TARGET, - gizmo_scale=self.cfg.ik_gizmo_scale, - name=f"{target_name}_{self._robot_arm_name}_ik", - ) - - self._ik_controller = ik_controller - self._gizmo = ik_controller.target_gizmo.gizmo - - def _resolve_robot_ik_chain( - self, - target: Robot, - ) -> tuple[str, str, np.ndarray]: - solver = ( - target.get_solver(self._control_part) - if target.cfg.solver_cfg is not None - else None - ) - root_link = self.cfg.ik_root_link_name or getattr( - solver, - "root_link_name", - None, - ) - end_link = self.cfg.ik_end_link_name or getattr( - solver, - "end_link_name", - None, - ) - if not root_link or not end_link: - raise ValueError( - "Robot Gizmo needs an IK chain. Set GizmoCfg.ik_root_link_name " - "and ik_end_link_name, or configure a solver for the selected " - "robot control part." - ) - - tcp_pose = self.cfg.ik_tcp_pose - if tcp_pose is None and solver is not None: - tcp_pose = solver.get_tcp() - if tcp_pose is None: - tcp_pose = np.eye(4, dtype=np.float32) - if isinstance(tcp_pose, torch.Tensor): - tcp_pose = tcp_pose.detach().cpu().numpy() - tcp_matrix = np.asarray(tcp_pose, dtype=np.float32) - if tcp_matrix.shape != (4, 4) or not np.isfinite(tcp_matrix).all(): - raise ValueError("ik_tcp_pose must be a finite 4x4 transform.") - return root_link, end_link, tcp_matrix - - def _read_native_robot_pose(self) -> torch.Tensor: - """Read the native robot TCP pose without an EmbodiChain solver.""" - if self._robot_adapter is None: - raise RuntimeError("Native robot Gizmo adapter is not configured.") - link_pose = self._robot_adapter.get_link_pose(self._native_robot_end_link) - tcp_pose = link_pose @ self._native_robot_tcp_pose - return self._as_pose_matrix(tcp_pose, self._target_device()) - def _target_device(self) -> torch.device: if self.target is None: return torch.device("cpu") @@ -491,7 +432,7 @@ def _read_target_pose(self) -> torch.Tensor: if self.target is None: raise RuntimeError("Gizmo is detached.") if self._target_type == "robot": - return self._read_native_robot_pose() + return self._read_robot_pose() pose = self.target.get_local_pose(to_matrix=True) return self._as_pose_matrix(pose[0], self._target_device()) @@ -506,7 +447,7 @@ def get_control_pose(self) -> torch.Tensor: return self._read_target_pose() def begin_interaction(self, source_id: str) -> bool: - """Acquire this Gizmo for one native or Viser drag source.""" + """Acquire this Gizmo for one Viser drag source.""" if not source_id: raise ValueError("source_id must not be empty.") with self._state_lock: @@ -516,11 +457,7 @@ def begin_interaction(self, source_id: str) -> bool: return True def request_local_pose(self, pose: object, *, source_id: str) -> bool: - """Queue a local target pose for application by :meth:`update`. - - Returns: - ``False`` if another client currently owns the drag. - """ + """Queue a local target pose for application by :meth:`update`.""" matrix = self._as_pose_matrix(pose, self._target_device()) with self._state_lock: if self._interaction_owner not in {None, source_id}: @@ -530,7 +467,7 @@ def request_local_pose(self, pose: object, *, source_id: str) -> bool: return True def end_interaction(self, source_id: str) -> bool: - """Release a drag source's ownership of this Gizmo.""" + """Release a Viser drag source's ownership of this Gizmo.""" with self._state_lock: if self._interaction_owner != source_id: return False @@ -548,56 +485,6 @@ def cancel_interaction(self, source_prefix: str | None = None) -> bool: self._interaction_owner = None return True - def _create_native_gizmo(self, cfg: GizmoCfg) -> object: - options = cfg.to_options_dict() - return self._env.create_gizmo(options["axis"], options["rings"]) - - def _create_proxy_cube(self, pose: torch.Tensor, name: str) -> object: - matrix = pose[0].detach().cpu().numpy() - euler = Rotation.from_matrix(matrix[:3, :3]).as_euler("xyz", degrees=False) - proxy_cube = self._env.create_cube(0.02, 0.02, 0.02) - proxy_cube.set_location(*matrix[:3, 3].tolist()) - proxy_cube.set_rotation_euler(*euler.tolist()) - self._require_native().follow(proxy_cube.node) - logger.log_info( - f"{name} Gizmo proxy created at position: {matrix[:3, 3].tolist()}" - ) - return proxy_cube - - def _set_proxy_pose(self, pose: torch.Tensor) -> None: - if self._proxy_cube is None: - return - matrix = pose[0].detach().cpu().numpy() - euler = Rotation.from_matrix(matrix[:3, :3]).as_euler("xyz", degrees=False) - self._proxy_cube.set_location(*matrix[:3, 3].tolist()) - self._proxy_cube.set_rotation_euler(*euler.tolist()) - - def _native_pose_callback(self, *args: object) -> None: - if len(args) != 3 or args[0] is None: - return - try: - pose = self._as_pose_matrix(args[1], self._target_device()) - if self._proxy_cube is not None: - self._set_proxy_pose(pose) - if not self.request_local_pose(pose, source_id="native"): - self._set_proxy_pose(self.get_control_pose()) - except (TypeError, ValueError) as error: - logger.log_warning(f"Ignoring invalid native Gizmo pose: {error}") - - def _setup_native_gizmo(self) -> None: - native = self._require_native() - if self.target is None: - raise RuntimeError("Gizmo is detached.") - if self._target_type == "rigid_object": - native.follow(self.target._entities[0].node) - else: - label = "Robot" if self._target_type == "robot" else "Camera" - self._proxy_cube = self._create_proxy_cube( - self.get_control_pose(), - label, - ) - native.set_flush_localpose_callback(self._native_pose_callback) - def _update_camera_pose(self, target_transform: torch.Tensor) -> bool: if self.target is None or not isinstance(self.target, Camera): return False @@ -615,7 +502,7 @@ def _update_rigid_object_pose(self, target_transform: torch.Tensor) -> bool: self.target.set_local_pose(target_transform, env_ids=[0]) return True except Exception as error: - logger.log_error(f"Error updating rigid object pose: {error}") + logger.log_error(f"Error updating rigid-object pose: {error}") return False def _update_robot_ik(self, target_transform: torch.Tensor) -> bool: @@ -629,244 +516,65 @@ def _update_robot_ik(self, target_transform: torch.Tensor) -> bool: rotation_matrix_to_quat_xyzw, ) - # The queued target is the TCP transform in the arena-local frame. - # Newton IK tracks a base-local target, so convert it with the same - # helper the native gizmo callback uses (inv(base_pose) @ target). base_pose = self._robot_adapter.get_world_pose() target_pose = target_transform[0].detach().cpu().numpy().astype(np.float32) base_local = local_pose_from_world(base_pose, target_pose) - position = np.asarray(base_local[:3, 3], dtype=np.float32) - rotation = rotation_matrix_to_quat_xyzw(base_local[:3, :3]) - joint_names = self._robot_adapter.get_actived_joint_names() current_qpos = self._robot_adapter.get_current_qpos() - self._ik_solver.set_target_pose(position, rotation) + self._ik_solver.set_target_pose( + np.asarray(base_local[:3, 3], dtype=np.float32), + rotation_matrix_to_quat_xyzw(base_local[:3, :3]), + ) self._ik_solver.solve( joint_names, current_qpos, iterations=self.cfg.ik_iterations, ) solved_qpos = self._ik_solver.qpos_for_joint_names( - joint_names, current_qpos + joint_names, + current_qpos, ) - # Drive the joint targets (matching native IKApplyMode.DRIVE_TARGET) - # so physics moves the robot instead of snapping its current pose. self._robot_adapter.set_target_qpos(solved_qpos) return True except Exception as error: - logger.log_error(f"Error in Gizmo robot IK: {error}") + logger.log_error(f"Error in Viser Gizmo robot IK: {error}") return False def update(self) -> None: - """Apply the latest queued target pose on the simulation thread.""" - if self._ik_controller is not None: - self._ik_controller.update(iterations=self.cfg.ik_iterations) - return - + """Apply the latest queued Viser target pose on the simulation thread.""" with self._state_lock: pending = self._pending_target_transform self._pending_target_transform = None - if pending is not None: - if self._target_type == "rigid_object": - self._update_rigid_object_pose(pending) - elif self._target_type == "robot": - self._update_robot_ik(pending) - elif self._target_type == "camera": - self._update_camera_pose(pending) - self._set_proxy_pose(pending) - - if self._gizmo is None or self.target is None: + if pending is None: return if self._target_type == "rigid_object": - self._gizmo.follow(self.target._entities[0].node) - elif self._target_type == "camera" and pending is None: - self._set_proxy_pose(self._read_target_pose()) - - def attach(self, target: BatchEntity) -> None: - """Attach this Gizmo to a supported target.""" - num_envs = int(getattr(target, "num_instances", dexsim.get_world_num())) - if num_envs > 1: - raise RuntimeError( - "Gizmo can only be used in single environment mode " - f"(num_envs=1), but target has {num_envs} instances." - ) - - self._release_native_resources() - self.target = target - self._target_type = self._detect_target_type(target) - self._robot_arm_name = None - if self._target_type == "robot": - self._configure_robot() - self._setup_robot_ik_solver() - if self._enable_native: - self._setup_native_robot_gizmo() - desired_pose = self._read_native_robot_pose() - else: - desired_pose = self._read_target_pose() - if self._enable_native: - self._gizmo = self._create_native_gizmo(self.cfg) - self._setup_native_gizmo() - with self._state_lock: - self._interaction_owner = None - self._pending_target_transform = None - self._desired_target_transform = desired_pose - - def detach(self) -> None: - """Detach this Gizmo from its current target.""" - self._release_native_resources() - with self._state_lock: - self._interaction_owner = None - self._pending_target_transform = None - self._desired_target_transform = None - self.target = None - self._target_type = "" - - def _require_native(self) -> object: - if self._gizmo is None: - raise RuntimeError("This Gizmo was created without a native DexSim handle.") - return self._gizmo - - def set_transform_callback(self, callback: Callable[..., Any]) -> None: - """Set a callback directly on the native transform handle.""" - self._callback = callback - self._require_native().set_transform_flush_callback(callback) - - def set_world_pose(self, pose: object) -> None: - """Set the native Gizmo world pose.""" - self._require_native().set_world_pose(pose) - - def set_local_pose(self, pose: object) -> None: - """Set the native Gizmo pose or queue it for a headless controller.""" - if self._gizmo is None: - self.request_local_pose(pose, source_id="api") - else: - self._gizmo.set_local_pose(pose) - - def set_line_width(self, width: float) -> None: - """Set the native Gizmo line width.""" - self._require_native().set_line_width(width) - - def enable_collision(self, enabled: bool) -> None: - """Enable or disable native Gizmo collision.""" - self._require_native().enable_collision(enabled) - - def get_world_pose(self) -> object: - """Return the native Gizmo world pose.""" - return self._require_native().get_world_pose() - - def get_local_pose(self) -> object: - """Return the native pose, or the logical local control pose.""" - if self._gizmo is None: - return self.get_control_pose() - return self._gizmo.get_local_pose() - - def get_name(self) -> object: - """Return the native Gizmo node name.""" - return self._require_native().get_name() - - def get_parent(self) -> object: - """Return the native Gizmo parent node.""" - return self._require_native().get_parent() + self._update_rigid_object_pose(pending) + elif self._target_type == "robot": + self._update_robot_ik(pending) + elif self._target_type == "camera": + self._update_camera_pose(pending) def toggle_visibility(self) -> bool: - """Toggle visibility and return the new state.""" - self.set_visible(not self._is_visible) + """Toggle Viser visibility and return the new state.""" + self._is_visible = not self._is_visible return self._is_visible def set_visible(self, visible: bool) -> None: - """Set native and Viser Gizmo visibility.""" + """Set Viser Gizmo visibility.""" self._is_visible = bool(visible) - if self._ik_controller is not None: - self._ik_controller.enabled = self._is_visible - if self._gizmo is not None: - self._gizmo.set_visible(self._is_visible) def is_visible(self) -> bool: - """Return whether this Gizmo should be visible.""" - if self._ik_controller is not None: - return bool(self._ik_controller.enabled) + """Return whether the Viser Gizmo should be visible.""" return self._is_visible - def apply_transform( - self, - translation: object, - rotation: object, - ) -> None: - """Apply a translation and XYZ Euler rotation through the shared path.""" - matrix = np.eye(4, dtype=np.float32) - matrix[:3, 3] = np.asarray(translation, dtype=np.float32) - matrix[:3, :3] = Rotation.from_euler( - "xyz", - np.asarray(rotation, dtype=np.float32), - ).as_matrix() - self.request_local_pose(matrix, source_id="api") - - def _remove_proxy_cube(self) -> None: - if self._proxy_cube is None: - return - try: - if self._gizmo is not None: - self._gizmo.detach_parent() - self._env.remove_actor(self._proxy_cube) - except Exception as error: - logger.log_warning(f"Failed to remove Gizmo proxy: {error}") - self._proxy_cube = None - - def _release_native_resources(self) -> None: - """Release DexSim Gizmo, proxy, and native IK resources.""" - gizmo = self._gizmo - if gizmo is not None: - for method_name in ( - "set_flush_localpose_callback", - "set_transform_flush_callback", - ): - method = getattr(gizmo, method_name, None) - if callable(method): - try: - method(None) - except (TypeError, RuntimeError): - pass - try: - gizmo.set_visible(False) - except (AttributeError, TypeError, RuntimeError): - pass - try: - gizmo.detach_parent() - except (AttributeError, TypeError, RuntimeError): - pass - - if self._ik_controller is not None: - try: - self._ik_controller.target_gizmo.target_node.detach_parent() - except (AttributeError, TypeError, RuntimeError): - pass - - self._remove_proxy_cube() - - if gizmo is not None: - remove_gizmo = getattr(self._env, "remove_gizmo", None) - if callable(remove_gizmo): - try: - remove_gizmo(gizmo) - except (AttributeError, TypeError, RuntimeError) as error: - logger.log_warning( - f"Failed to remove Gizmo from DexSim environment: {error}" - ) - - self._gizmo = None - self._proxy_cube = None - self._ik_controller = None - self._ik_solver = None - self._ik_model = None - self._robot_adapter = None - def destroy(self) -> None: - """Release native resources and target references.""" - self._release_native_resources() + """Release target and IK references.""" with self._state_lock: self._interaction_owner = None self._pending_target_transform = None self._desired_target_transform = None + self._ik_solver = None + self._robot_adapter = None self.target = None self._target_type = "" diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index eb9e37894..ed661a4e0 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -54,7 +54,7 @@ from dexsim.engine import CudaArray, Material from dexsim.models import MeshObject from dexsim.render import Light as _Light, LightType, Windows -from dexsim.engine import GizmoController, ObjectManipulator +from dexsim.engine import ObjectManipulator from embodichain.lab.sim.objects import ( RigidObject, @@ -184,9 +184,6 @@ class SimulationManagerCfg: window_camera_pose: WindowCameraPoseCfg = field(default_factory=WindowCameraPoseCfg) """Interactive viewer camera-pose printing settings.""" - enable_entity_gizmo_on_window_open: bool = True - """Whether opening a native window enables world-level entity Gizmo control.""" - visualization: VisualizationCfg = field(default_factory=VisualizationCfg) """Live browser visualization settings.""" @@ -287,7 +284,6 @@ def __init__( self._world: dexsim.World = dexsim.World(world_config) self._window: Windows | None = None - self._entity_gizmo_config: EntityGizmoConfig | None = None self._window_record_state: _WindowRecordState | None = None self._window_record_camera: object | None = None wr = sim_config.window_record @@ -377,7 +373,6 @@ def __init__( if sim_config.headless is False: self._window = self._world.get_windows() - self._on_window_opened() @classmethod def get_instance(cls, instance_id: int = 0) -> SimulationManager: @@ -687,6 +682,8 @@ def stop_visualization(self) -> None: try: runtime.stop() finally: + if getattr(self, "_picker_gizmo", None) is not None: + self._release_picker_gizmo() for _, gizmo in self.get_gizmo_items(): cancel = getattr(gizmo, "cancel_interaction", None) if cancel is not None: @@ -889,25 +886,13 @@ def can_open_native_window(self) -> bool: and self._visualization_runtime is None ) - def open_window( - self, - *, - enable_entity_gizmo: bool | None = None, - entity_gizmo_config: EntityGizmoConfig | None = None, - ) -> bool: + def open_window(self) -> bool: """Open the native DexSim simulation window when allowed. Viser owns visualization while it is configured or running. In that case this method safely skips the native window so launchers do not need a separate Viser condition. - Args: - enable_entity_gizmo: Whether to enable world-level entity Gizmo - control. ``None`` uses the simulation configuration default. - entity_gizmo_config: Optional native DexSim entity-Gizmo settings. - Providing settings implies enabling the controller unless - ``enable_entity_gizmo`` is explicitly ``False``. - Returns: ``True`` when the native window is open, otherwise ``False``. """ @@ -918,28 +903,10 @@ def open_window( ) return False if self.is_window_opened: - if enable_entity_gizmo is not None or entity_gizmo_config is not None: - self._on_window_opened( - enable_entity_gizmo=enable_entity_gizmo, - entity_gizmo_config=entity_gizmo_config, - ) return True self._world.open_window() self._window = self._world.get_windows() - self.is_window_opened = True - self._on_window_opened( - enable_entity_gizmo=enable_entity_gizmo, - entity_gizmo_config=entity_gizmo_config, - ) - return True - def _on_window_opened( - self, - *, - enable_entity_gizmo: bool | None = None, - entity_gizmo_config: EntityGizmoConfig | None = None, - ) -> None: - """Initialize controls shared by constructor-opened and reopened windows.""" if ( self._window_record_hotkey_cfg is not None and self._window_record_input_control is None @@ -950,31 +917,11 @@ def _on_window_opened( and self._window_camera_pose_input_control is None ): self.enable_window_camera_pose_hotkey(**self._window_camera_pose_hotkey_cfg) - - if enable_entity_gizmo is None: - enable_entity_gizmo = entity_gizmo_config is not None or getattr( - self.sim_config, - "enable_entity_gizmo_on_window_open", - True, - ) - - try: - if enable_entity_gizmo: - if entity_gizmo_config is not None: - self.enable_entity_gizmo(entity_gizmo_config) - elif not self.has_entity_gizmo(): - self.enable_entity_gizmo(self._entity_gizmo_config) - elif self.has_entity_gizmo(): - self.disable_entity_gizmo() - except RuntimeError as error: - logger.log_warning( - f"Entity Gizmo control could not be initialized for the window: {error}" - ) + self.is_window_opened = True + return True def close_window(self) -> None: """Close the simulation window.""" - if self.has_entity_gizmo(): - self.disable_entity_gizmo() if self.is_window_recording(): self.stop_window_record() self._world.close_window() @@ -1970,133 +1917,47 @@ def enable_entity_gizmo( self, config: EntityGizmoConfig | None = None, ) -> EntityGizmoManipulator: - """Enable DexSim's world-level entity Gizmo controller. - - The world-owned controller handles selection, hotkeys, simultaneous - bindings, temporary physics-state changes, and rigid-body or - articulation-root manipulation. + """Enable DexSim entity control and exclude the EmbodiChain ground. Args: - config: Native DexSim entity-Gizmo configuration. DexSim defaults - are used when omitted. + config: Native DexSim entity-Gizmo configuration. Returns: The active world-owned entity Gizmo manipulator. - - Raises: - RuntimeError: If the installed DexSim does not provide the API or - fails to create the controller. """ - world = getattr(self, "_world", None) - enable = getattr(world, "enable_entity_gizmo", None) - if not callable(enable): - raise RuntimeError( - "The installed DexSim build does not provide " - "World.enable_entity_gizmo()." - ) - - controller = enable() if config is None else enable(config) - if controller is None: - raise RuntimeError("DexSim failed to enable the entity Gizmo controller.") - self._exclude_default_plane_from_entity_gizmo(controller) - self._entity_gizmo_config = config - logger.log_info("DexSim entity Gizmo control enabled.") - return controller - - def _exclude_default_plane_from_entity_gizmo( - self, - controller: EntityGizmoManipulator, - ) -> None: - """Register the EmbodiChain ground as an immovable Gizmo target.""" + controller = ( + self._world.enable_entity_gizmo() + if config is None + else self._world.enable_entity_gizmo(config) + ) default_plane = getattr(self, "_default_plane", None) - register = getattr(controller, "register_external_target", None) if default_plane is None: - return - if not callable(register): - logger.log_warning( - "The installed DexSim build cannot exclude the default plane " - "from entity Gizmo control." - ) - return - - try: - result = register( - self._DEFAULT_PLANE_GIZMO_TARGET_ID, - dexsim.interaction.EntityGizmoTargetType.RIGID_BODY, - default_plane, - ActorType.STATIC, - ) - except (AttributeError, TypeError, RuntimeError) as error: - logger.log_warning( - "Failed to exclude the default plane from entity Gizmo " - f"control: {error}." - ) - return + return controller + result = controller.register_external_target( + self._DEFAULT_PLANE_GIZMO_TARGET_ID, + dexsim.interaction.EntityGizmoTargetType.RIGID_BODY, + default_plane, + ActorType.STATIC, + ) if result != dexsim.interaction.EntityGizmoResult.SUCCESS: logger.log_warning( "Failed to exclude the default plane from entity Gizmo " f"control: {result}." ) - - def disable_entity_gizmo(self) -> bool: - """Disable DexSim's world-level entity Gizmo controller. - - Returns: - ``True`` when an active controller was disabled, or ``False`` when - entity Gizmo control was already disabled. - - Raises: - RuntimeError: If the installed DexSim lacks entity-Gizmo lifecycle - APIs. - """ - world = getattr(self, "_world", None) - get_controller = getattr(world, "get_entity_gizmo", None) - disable = getattr(world, "disable_entity_gizmo", None) - if not callable(get_controller) or not callable(disable): - raise RuntimeError( - "The installed DexSim build does not provide entity Gizmo " - "lifecycle APIs." - ) - if get_controller() is None: - return False - - disable() - logger.log_info("DexSim entity Gizmo control disabled.") - return True - - def get_entity_gizmo(self) -> EntityGizmoManipulator | None: - """Return DexSim's active world-level entity Gizmo controller.""" - world = getattr(self, "_world", None) - get_controller = getattr(world, "get_entity_gizmo", None) - if not callable(get_controller): - raise RuntimeError( - "The installed DexSim build does not provide " - "World.get_entity_gizmo()." - ) - return get_controller() - - def has_entity_gizmo(self) -> bool: - """Return whether world-level entity Gizmo control is enabled.""" - world = getattr(self, "_world", None) - get_controller = getattr(world, "get_entity_gizmo", None) - return callable(get_controller) and get_controller() is not None + return controller def enable_gizmo( self, uid: str, control_part: str | None = None, gizmo_cfg: GizmoCfg | None = None, - *, - enable_native: bool | None = None, ) -> Gizmo | None: - """Enable gizmo control for any simulation object (Robot, RigidObject, Camera, etc.). + """Enable Viser Gizmo control for a simulation target. Args: uid: UID of the robot, rigid object, or camera sensor. control_part: Robot control part used for IK/FK. - gizmo_cfg: Native and Viser Gizmo appearance configuration. - enable_native: Whether to create a DexSim Gizmo. By default, native - controls are created only when a native window is active. + gizmo_cfg: Viser appearance and robot IK configuration. Returns: The created Gizmo, or ``None`` if setup failed. @@ -2131,36 +1992,14 @@ def enable_gizmo( ) return None - if enable_native is None: - enable_native = self.is_window_opened or not self.sim_config.headless gizmo: Gizmo | None = None try: - gizmo = Gizmo( - target, - gizmo_cfg, - control_part, - enable_native=enable_native, - ) - if enable_native and ( - not hasattr(self, "_gizmo_controller") or self._gizmo_controller is None - ): - window = ( - self._world.get_windows() - if hasattr(self._world, "get_windows") - else None - ) - if window is None: - raise RuntimeError( - "A native window is required for the DexSim Gizmo controller." - ) - self._gizmo_controller = GizmoController() - window.add_input_control(self._gizmo_controller) + gizmo = Gizmo(target, gizmo_cfg, control_part) self._gizmos[gizmo_key] = gizmo self.notify_visualization_topology_changed() logger.log_info( - f"Gizmo enabled for {object_type} '{uid}' with control_part " - f"'{control_part}' (native={enable_native}, " - f"viser={self.sim_config.visualization.allow_commands})" + f"Viser Gizmo enabled for {object_type} '{uid}' with " + f"control_part '{control_part}'." ) except Exception as e: @@ -2362,6 +2201,8 @@ def process_pick_commands(self) -> int: if self._picker_gizmo is not None and self._picker_gizmo[0] == uid: continue self._release_picker_gizmo() + if any(key == uid or key.startswith(f"{uid}:") for key in self._gizmos): + continue gizmo = self.enable_gizmo(uid=uid) if gizmo is not None: self._picker_gizmo = (uid, None) @@ -3311,9 +3152,6 @@ def destroy(self, exit_process: bool | None = None) -> None: def _deferred_destroy(self) -> None: """Destroy all simulated assets and release resources.""" - if self.has_entity_gizmo(): - self.disable_entity_gizmo() - # Clean up all gizmos before destroying the simulation for uid in list(self._gizmos.keys()): self.disable_gizmo(uid) diff --git a/embodichain/lab/sim/utility/__init__.py b/embodichain/lab/sim/utility/__init__.py index 02f142a69..c839d646f 100644 --- a/embodichain/lab/sim/utility/__init__.py +++ b/embodichain/lab/sim/utility/__init__.py @@ -18,6 +18,5 @@ from .sim_utils import * from .mesh_utils import * -from .gizmo_utils import * from .keyboard_utils import * from .render_utils import * diff --git a/embodichain/lab/sim/utility/gizmo_utils.py b/embodichain/lab/sim/utility/gizmo_utils.py deleted file mode 100644 index 177612023..000000000 --- a/embodichain/lab/sim/utility/gizmo_utils.py +++ /dev/null @@ -1,240 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Gizmo utility functions for EmbodiChain. - -This module provides utility functions for creating gizmo transform callbacks. -""" - -from __future__ import annotations - -from collections.abc import Callable -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from embodichain.lab.sim.objects import Robot - -__all__ = ["create_gizmo_callback", "run_gizmo_robot_control_loop"] - - -def create_gizmo_callback() -> Callable[[Any, Any, Any], None]: - """Create a standard gizmo transform callback function. - - This callback handles local pose for gizmo controls. - It applies transformations directly to the node when gizmo controls are manipulated. - - Returns: - A callback compatible with dexsim's gizmo local-pose flush hook. - """ - - def gizmo_transform_callback(node: Any, local_pose: Any, flag: Any) -> None: - if node is not None: - node.set_transform(local_pose, flag) - - return gizmo_transform_callback - - -def run_gizmo_robot_control_loop( - robot: Robot | str, - control_part: str = "arm", - end_link_name: str | None = None, -) -> None: - """Run a control loop for testing gizmo controls on a robot. - - This function implements a control loop that allows users to manipulate a robot - using gizmo controls with keyboard input for additional commands. - - Args: - robot (Robot | str): The robot to control with the gizmo. - control_part (str, optional): The part of the robot to control. Defaults to "arm". - end_link_name (str | None, optional): The name of the end link for FK calculations. Defaults to None. - - Keyboard Controls: - Q/ESC: Exit the control loop - P: Print current robot state (joint positions, end-effector pose) - G: Toggle gizmo visibility - R: Reset robot to initial pose - I: Print control information - """ - import select - import sys - import tty - import termios - import time - import numpy as np - - np.set_printoptions(precision=5, suppress=True) - - from embodichain.lab.sim import SimulationManager - from embodichain.lab.sim.objects import GizmoCfg - - from embodichain.utils.logger import log_error, log_info - - sim = SimulationManager.get_instance() - - if isinstance(robot, str): - robot_uid = robot - robot = sim.get_robot(uid=robot_uid) - if robot is None: - log_error(f"Robot {robot_uid!r} was not found.") - return - - # Enter auto-update mode. - sim.set_manual_update(False) - - # Resolve only the chain metadata. dexsim owns the Newton IK solver and - # writes its drive targets back through the EmbodiChain Robot API. - robot_solver = ( - robot.get_solver(name=control_part) - if robot.cfg.solver_cfg is not None - else None - ) - control_part_link_names = robot.get_control_part_link_names(name=control_part) - if not control_part_link_names: - raise ValueError(f"Control part {control_part!r} has no links.") - root_link_name = ( - robot_solver.root_link_name - if robot_solver is not None - else control_part_link_names[0] - ) - end_link_name = ( - ( - robot_solver.end_link_name - if robot_solver is not None - else control_part_link_names[-1] - ) - if end_link_name is None - else end_link_name - ) - tcp_pose = robot_solver.get_tcp() if robot_solver is not None else None - gizmo_cfg = GizmoCfg( - ik_root_link_name=root_link_name, - ik_end_link_name=end_link_name, - ik_tcp_pose=tcp_pose, - ) - - # Enable gizmo for the robot - gizmo = sim.enable_gizmo( - uid=robot.uid, - control_part=control_part, - gizmo_cfg=gizmo_cfg, - ) - if gizmo is None: - log_error(f"Failed to enable gizmo for control part {control_part!r}.") - return - - # Store initial robot configuration - initial_qpos = robot.get_qpos(name=control_part) - - gizmo_visible = True - - log_info("\n=== Gizmo Robot Control ===") - log_info("Gizmo Controls:") - log_info(" Use the 3D gizmo to drag and manipulate the robot") - log_info("\nKeyboard Controls:") - log_info(" Q/ESC: Exit control loop") - log_info(" P: Print current robot state") - log_info(" G: Toggle gizmo visibility") - log_info(" R: Reset robot to initial pose") - log_info(" I: Print this information again") - - # Save terminal settings - old_settings = termios.tcgetattr(sys.stdin) - tty.setcbreak(sys.stdin.fileno()) - - def get_key() -> str | None: - """Non-blocking keyboard input.""" - if select.select([sys.stdin], [], [], 0)[0]: - return sys.stdin.read(1) - return None - - try: - while True: - time.sleep(0.033) # ~30Hz - sim.update_gizmos() - - # Check for keyboard input - key = get_key() - - if key: - # Exit controls - if key in ["q", "Q", "\x1b"]: # Q or ESC - log_info("Exiting gizmo control loop...") - sim.disable_gizmo(uid=robot.uid, control_part=control_part) - break - - # Print robot state - elif key in ["p", "P"]: - current_qpos = robot.get_qpos(name=control_part) - eef_pose = robot.get_link_pose(end_link_name, to_matrix=True) - if tcp_pose is not None: - tcp_tensor = np.asarray(tcp_pose, dtype=np.float32) - eef_pose = eef_pose @ eef_pose.new_tensor(tcp_tensor) - log_info(f"\n=== Robot State ===") - log_info(f"Control part: {control_part}") - log_info(f"Joint positions: {current_qpos.squeeze().tolist()}") - eef_pose_np = eef_pose.detach().cpu().numpy().squeeze() - log_info(f"End-effector pose:\n{eef_pose_np}") - elif key in ["g", "G"]: - if gizmo_visible: - sim.set_gizmo_visibility( - uid=robot.uid, control_part=control_part, visible=False - ) - log_info("Gizmo hidden") - gizmo_visible = False - else: - sim.set_gizmo_visibility( - uid=robot.uid, control_part=control_part, visible=True - ) - log_info("Gizmo shown") - gizmo_visible = True - - # Reset to initial pose - elif key in ["r", "R"]: - # TODO: Workaround for reset. Gizmo pose should be fixed in the future. - sim.disable_gizmo(uid=robot.uid, control_part=control_part) - robot.clear_dynamics() - robot.set_qpos(qpos=initial_qpos, name=control_part, target=False) - sim.enable_gizmo( - uid=robot.uid, - control_part=control_part, - gizmo_cfg=gizmo_cfg, - ) - log_info("Robot reset to initial pose") - - # Print info - elif key in ["i", "I"]: - log_info("\n=== Gizmo Robot Control ===") - log_info("Gizmo Controls:") - log_info(" Use the 3D gizmo to drag and manipulate the robot") - log_info("\nKeyboard Controls:") - log_info(" Q/ESC: Exit control loop") - log_info(" P: Print current robot state") - log_info(" G: Toggle gizmo visibility") - log_info(" R: Reset robot to initial pose") - log_info(" I: Print this information again") - - except KeyboardInterrupt: - sim.disable_gizmo(uid=robot.uid, control_part=control_part) - log_info("\nControl loop interrupted by user (Ctrl+C)") - - finally: - try: - # Restore terminal settings - termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings) - except: - pass - log_info("Gizmo control loop terminated") diff --git a/embodichain/lab/visualization/backends/viser.py b/embodichain/lab/visualization/backends/viser.py index 5a0947958..931293271 100644 --- a/embodichain/lab/visualization/backends/viser.py +++ b/embodichain/lab/visualization/backends/viser.py @@ -1681,6 +1681,13 @@ def stop(self) -> None: self._gizmo_owners.clear() self._gizmo_drag_poses.clear() self._gizmo_sequence = 0 + self._picker.clear() + self._pick_enabled = False + self._node_geometry.clear() + self._frame_positions = None + self._frame_wxyz = None + self._frame_visible = None + self._pointer_handler = None self._joint_control_handles.clear() self._joint_control_specs.clear() self._joint_control_states.clear() diff --git a/embodichain/lab/visualization/runtime.py b/embodichain/lab/visualization/runtime.py index 81c20cf31..9fa020e98 100644 --- a/embodichain/lab/visualization/runtime.py +++ b/embodichain/lab/visualization/runtime.py @@ -161,8 +161,8 @@ def put(self, command: PickCommand) -> None: with self._lock: for index in range(len(self._commands) - 1, -1, -1): if self._commands[index].client_id == command.client_id: - self._commands[index] = command - return + del self._commands[index] + break if len(self._commands) >= self._maxsize: self._commands.popleft() self._commands.append(command) diff --git a/examples/sim/gizmo/gizmo_camera.py b/examples/sim/gizmo/gizmo_camera.py index 8ff4e842a..c96352365 100644 --- a/examples/sim/gizmo/gizmo_camera.py +++ b/examples/sim/gizmo/gizmo_camera.py @@ -105,22 +105,22 @@ def main(): if not args.headless: native_window_opened = sim.open_window() - # Enable gizmo for interactive camera control using the new unified API - if native_window_opened or args.viser: + if args.viser: sim.enable_gizmo( uid="gizmo_camera", - enable_native=native_window_opened, ) if not sim.has_gizmo("gizmo_camera"): logger.log_error("Failed to enable gizmo for camera!") return - else: + elif not native_window_opened: logger.log_warning( "Gizmo interaction is disabled in headless mode without Viser." ) + else: + logger.log_warning("Camera Gizmo control is available through Viser only.") logger.log_info("Gizmo-Camera tutorial started!") - if native_window_opened or args.viser: + if args.viser: logger.log_info( "Use the gizmo to interactively control the camera position and orientation" ) @@ -145,7 +145,7 @@ def run_simulation( last_step = 0 if show_camera_window: - logger.log_info("Camera view window will open. Press Ctrl+C or 'q' to exit") + logger.log_info("Camera view window will open. Press Ctrl+C to exit") if sim.has_gizmo("gizmo_camera"): logger.log_info( "Use the gizmo in the 3D view to control camera position and orientation" @@ -177,25 +177,20 @@ def run_simulation( # Convert RGB to BGR for OpenCV bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR) - # Add text overlay - cv2.putText( - bgr_image, - "Press 'h' to toggle camera gizmo visibility", - (10, 30), - cv2.FONT_HERSHEY_SIMPLEX, - 0.6, - (0, 255, 0), - 2, - ) - - # Display the image - cv2.imshow("Gizmo Camera View", bgr_image) + # Add text overlay + cv2.putText( + bgr_image, + "Camera sensor preview", + (10, 30), + cv2.FONT_HERSHEY_SIMPLEX, + 0.6, + (0, 255, 0), + 2, + ) - # Check for key press - key = cv2.waitKey(1) & 0xFF - if key == ord("h"): - # Toggle the camera gizmo visibility using SimulationManager API - sim.toggle_gizmo_visibility("gizmo_camera") + # Display the image + cv2.imshow("Gizmo Camera View", bgr_image) + cv2.waitKey(1) # Example: Destroy gizmo after certain steps to test cleanup if step_count == 30000 and sim.has_gizmo("gizmo_camera"): diff --git a/examples/sim/gizmo/gizmo_object.py b/examples/sim/gizmo/gizmo_object.py index 5eac31b2f..b01175738 100644 --- a/examples/sim/gizmo/gizmo_object.py +++ b/examples/sim/gizmo/gizmo_object.py @@ -92,20 +92,18 @@ def main(): if not args.headless: entity_gizmo_config = dexsim.interaction.EntityGizmoConfig() entity_gizmo_config.max_gizmos = 0 - native_window_opened = sim.open_window( - entity_gizmo_config=entity_gizmo_config, - ) + native_window_opened = sim.open_window() + if native_window_opened: + sim.enable_entity_gizmo(entity_gizmo_config) - # Native windows use DexSim's raycast-selected entity controller; headless - # Viser uses one backend-neutral Gizmo per published object. + # Native windows use DexSim's entity controller; Viser publishes one + # EmbodiChain-side transform control per object. if args.viser: sim.enable_gizmo( uid="cube1", - enable_native=False, ) sim.enable_gizmo( uid="cube2", - enable_native=False, ) elif native_window_opened: logger.log_info("Left-click an entity and press G to attach/detach its Gizmo.") @@ -146,8 +144,8 @@ def run_simulation(sim: SimulationManager): # Disable Gizmo control after 200000 steps (example). if step_count == 200000 and gizmo_enabled: logger.log_info("Disabling Gizmo control at step 200000") - if sim.has_entity_gizmo(): - sim.disable_entity_gizmo() + if sim.get_world().get_entity_gizmo() is not None: + sim.get_world().disable_entity_gizmo() else: sim.disable_gizmo("cube1") sim.disable_gizmo("cube2") diff --git a/examples/sim/gizmo/gizmo_robot.py b/examples/sim/gizmo/gizmo_robot.py index 5fcde4d93..477abb924 100644 --- a/examples/sim/gizmo/gizmo_robot.py +++ b/examples/sim/gizmo/gizmo_robot.py @@ -25,7 +25,10 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.solvers import PytorchSolverCfg -from embodichain.lab.sim.objects import GizmoCfg +from embodichain.lab.sim.objects import ( + GizmoCfg, + create_robot_ik_gizmo_controller, +) from embodichain.lab.sim.cfg import ( RenderCfg, RobotCfg, @@ -116,23 +119,29 @@ def main(): if not args.headless: native_window_opened = sim.open_window() - # Enable gizmo using the new API - if native_window_opened or args.viser: - gizmo_cfg = GizmoCfg( - ik_root_link_name="base_link", - ik_end_link_name="ee_link", - ik_tcp_pose=[ - [0.0, 1.0, 0.0, 0.0], - [-1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.12], - [0.0, 0.0, 0.0, 1.0], - ], + gizmo_cfg = GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ik_tcp_pose=[ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.12], + [0.0, 0.0, 0.0, 1.0], + ], + ) + native_control = None + if native_window_opened: + native_control = create_robot_ik_gizmo_controller( + robot, + control_part="arm", + cfg=gizmo_cfg, + world=sim.get_world(), ) + elif args.viser: sim.enable_gizmo( uid="ur10_gizmo_test", control_part="arm", gizmo_cfg=gizmo_cfg, - enable_native=native_window_opened, ) if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): logger.log_error("Failed to enable gizmo!") @@ -149,17 +158,18 @@ def main(): logger.log_info("Press I to show or hide the native robot IK Gizmo") logger.log_info("Press Ctrl+C to stop the simulation") - run_simulation(sim) + run_simulation(sim, native_control) -def run_simulation(sim: SimulationManager): +def run_simulation(sim: SimulationManager, native_control=None): step_count = 0 try: last_time = time.time() last_step = 0 while True: time.sleep(0.033) # 30Hz - # Update all gizmos managed by sim + if native_control is not None: + native_control[0].update() sim.update_gizmos() sim.capture_visualization_safely() step_count += 1 diff --git a/examples/sim/gizmo/gizmo_scene.py b/examples/sim/gizmo/gizmo_scene.py index 554517a83..fec1af200 100644 --- a/examples/sim/gizmo/gizmo_scene.py +++ b/examples/sim/gizmo/gizmo_scene.py @@ -44,6 +44,7 @@ from embodichain.lab.sim.shapes import CubeCfg from embodichain.lab.sim.sensors import CameraCfg from embodichain.lab.sim.solvers import PinkSolverCfg +from embodichain.lab.sim.objects import create_robot_ik_gizmo_controller from embodichain.data import get_data_path from embodichain.utils import logger @@ -166,12 +167,21 @@ def main(): if not args.headless: native_window_opened = sim.open_window() - # Enable gizmo for all assets after all are created and initialized - if native_window_opened or args.viser: + native_controls = [] + if native_window_opened: + sim.enable_entity_gizmo() + for control_part in ("left_arm", "right_arm"): + native_controls.append( + create_robot_ik_gizmo_controller( + robot, + control_part=control_part, + world=sim.get_world(), + ) + ) + elif args.viser: sim.enable_gizmo( uid="w1_gizmo_test", control_part="left_arm", - enable_native=native_window_opened, ) if not sim.has_gizmo("w1_gizmo_test", control_part="left_arm"): logger.log_error("Failed to enable left arm gizmo!") @@ -180,7 +190,6 @@ def main(): sim.enable_gizmo( uid="w1_gizmo_test", control_part="right_arm", - enable_native=native_window_opened, ) if not sim.has_gizmo("w1_gizmo_test", control_part="right_arm"): logger.log_error("Failed to enable right arm gizmo!") @@ -188,7 +197,6 @@ def main(): sim.enable_gizmo( uid="interactive_cube", - enable_native=native_window_opened, ) if not sim.has_gizmo("interactive_cube"): logger.log_error("Failed to enable gizmo for cube!") @@ -196,7 +204,6 @@ def main(): sim.enable_gizmo( uid="scene_camera", - enable_native=native_window_opened, ) if not sim.has_gizmo("scene_camera"): logger.log_error("Failed to enable gizmo for camera!") @@ -207,7 +214,7 @@ def main(): ) logger.log_info("Gizmo Scene example started!") - if native_window_opened or args.viser: + if args.viser: logger.log_info("Four gizmos are active in the scene:") logger.log_info( "1. Left arm gizmo - Use to drag the left arm end-effector (EE)" @@ -217,14 +224,22 @@ def main(): ) logger.log_info("3. Cube gizmo - Use to drag and position the cube") logger.log_info("4. Camera gizmo - Use to drag and orient the camera") + elif native_window_opened: + logger.log_info("Press I to show or hide each robot TCP IK Gizmo.") + logger.log_info("Select a scene entity and press G to manipulate its root.") logger.log_info("Press Ctrl+C to stop the simulation") - run_simulation(sim, show_camera_window=native_window_opened) + run_simulation( + sim, + native_controls=native_controls, + show_camera_window=native_window_opened, + ) def run_simulation( sim: SimulationManager, *, + native_controls=(), show_camera_window: bool, ) -> None: step_count = 0 @@ -235,6 +250,8 @@ def run_simulation( last_step = 0 while True: time.sleep(0.033) # 30Hz + for controller, _ in native_controls: + controller.update() sim.update_gizmos() sim.capture_visualization_safely() step_count += 1 @@ -248,7 +265,7 @@ def run_simulation( bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR) cv2.putText( bgr_image, - "Press 'h' to toggle camera gizmo visibility", + "Camera sensor preview", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, @@ -256,10 +273,7 @@ def run_simulation( 2, ) cv2.imshow("Camera Sensor View", bgr_image) - key = cv2.waitKey(1) & 0xFF - if key == ord("h"): - # Toggle the camera gizmo visibility using SimulationManager API - sim.toggle_gizmo_visibility("scene_camera") + cv2.waitKey(1) if step_count % 100 == 0: current_time = time.time() diff --git a/examples/sim/gizmo/gizmo_w1.py b/examples/sim/gizmo/gizmo_w1.py index 554859c39..394585554 100644 --- a/examples/sim/gizmo/gizmo_w1.py +++ b/examples/sim/gizmo/gizmo_w1.py @@ -37,6 +37,7 @@ from embodichain.data import get_data_path from embodichain.utils import logger from embodichain.lab.sim.robots.dexforce_w1.cfg import DexforceW1Cfg +from embodichain.lab.sim.objects import create_robot_ik_gizmo_controller def main(): @@ -158,12 +159,20 @@ def main(): if not args.headless: native_window_opened = sim.open_window() - # Enable gizmo for both arms using the new API - if native_window_opened or args.viser: + native_controls = [] + if native_window_opened: + for control_part in ("left_arm", "right_arm"): + native_controls.append( + create_robot_ik_gizmo_controller( + robot, + control_part=control_part, + world=sim.get_world(), + ) + ) + elif args.viser: sim.enable_gizmo( uid="w1_gizmo_test", control_part="left_arm", - enable_native=native_window_opened, ) if not sim.has_gizmo("w1_gizmo_test", control_part="left_arm"): logger.log_error("Failed to enable left arm gizmo!") @@ -172,7 +181,6 @@ def main(): sim.enable_gizmo( uid="w1_gizmo_test", control_part="right_arm", - enable_native=native_window_opened, ) if not sim.has_gizmo("w1_gizmo_test", control_part="right_arm"): logger.log_error("Failed to enable right arm gizmo!") @@ -187,17 +195,18 @@ def main(): logger.log_info("Use the gizmos to drag both robot arms' end-effectors") logger.log_info("Press Ctrl+C to stop the simulation") - run_simulation(sim) + run_simulation(sim, native_controls) -def run_simulation(sim: SimulationManager): +def run_simulation(sim: SimulationManager, native_controls=()): step_count = 0 try: last_time = time.time() last_step = 0 while True: time.sleep(0.033) # 30Hz - # Update all gizmos managed by sim + for controller, _ in native_controls: + controller.update() sim.update_gizmos() sim.capture_visualization_safely() step_count += 1 diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index a06412f05..c14d9dd3b 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -24,7 +24,10 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.objects import GizmoCfg +from embodichain.lab.sim.objects import ( + GizmoCfg, + create_robot_ik_gizmo_controller, +) from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.cfg import ( RenderCfg, @@ -105,16 +108,23 @@ def main(): if not args.headless: native_window_opened = sim.open_window() - # Enable gizmo using the new API - if native_window_opened or args.viser: + gizmo_cfg = GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ) + native_control = None + if native_window_opened: + native_control = create_robot_ik_gizmo_controller( + robot, + control_part="arm", + cfg=gizmo_cfg, + world=sim.get_world(), + ) + elif args.viser: sim.enable_gizmo( uid="ur10_gizmo_test", control_part="arm", - gizmo_cfg=GizmoCfg( - ik_root_link_name="base_link", - ik_end_link_name="ee_link", - ), - enable_native=native_window_opened, + gizmo_cfg=gizmo_cfg, ) if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): logger.log_error("Failed to enable gizmo!") @@ -131,17 +141,18 @@ def main(): logger.log_info("Press I to show or hide the native robot IK Gizmo") logger.log_info("Press Ctrl+C to stop the simulation") - run_simulation(sim) + run_simulation(sim, native_control) -def run_simulation(sim: SimulationManager): +def run_simulation(sim: SimulationManager, native_control=None): step_count = 0 try: last_time = time.time() last_step = 0 while True: time.sleep(0.033) # 30Hz - # Update all gizmos managed by sim + if native_control is not None: + native_control[0].update() sim.update_gizmos() sim.capture_visualization_safely() step_count += 1 diff --git a/tests/lab/scripts/test_run_env.py b/tests/lab/scripts/test_run_env.py index b9bd25807..0c495a4f5 100644 --- a/tests/lab/scripts/test_run_env.py +++ b/tests/lab/scripts/test_run_env.py @@ -17,7 +17,7 @@ from __future__ import annotations from types import SimpleNamespace -from unittest.mock import MagicMock, call +from unittest.mock import MagicMock import pytest import torch @@ -40,24 +40,6 @@ VISER_POLL_INTERVAL = 0.05 -class _PreviewInput: - """Deterministic input source for the non-blocking preview loop.""" - - def __init__(self, keys: list[str | None]) -> None: - self._keys = iter(keys) - self.timeouts: list[float | None] = [] - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, traceback): - return None - - def read_key(self, timeout: float | None = None) -> str | None: - self.timeouts.append(timeout) - return next(self._keys) - - class _LegacyProgressEnv: num_envs = 1 @@ -141,118 +123,6 @@ def test_run_env_preserves_configured_viser_image_fps() -> None: assert merged["visualization"]["sensor_image_fps"] == configured_fps -def test_preview_enables_hidden_ik_gizmos_for_active_solver_parts() -> None: - """Preview prepares each task-selected arm that has an IK solver.""" - solvers = {"left_arm": object(), "right_arm": object()} - robot = SimpleNamespace( - uid="preview_robot", - control_parts={ - "left_arm": [], - "left_eef": [], - "right_arm": [], - }, - get_solver=MagicMock(side_effect=lambda part: solvers.get(part)), - ) - sim = MagicMock() - sim.is_window_opened = True - sim.has_gizmo.return_value = False - sim.enable_gizmo.side_effect = [object(), object()] - env = SimpleNamespace( - unwrapped=SimpleNamespace( - sim=sim, - robot=robot, - num_envs=1, - cfg=SimpleNamespace(control_parts=["left_arm", "left_eef", "right_arm"]), - ) - ) - - gizmo_keys = run_env._enable_preview_ik_gizmos(env) - - assert gizmo_keys == ( - ("preview_robot", "left_arm"), - ("preview_robot", "right_arm"), - ) - assert sim.enable_gizmo.call_args_list == [ - call(uid="preview_robot", control_part="left_arm", enable_native=True), - call(uid="preview_robot", control_part="right_arm", enable_native=True), - ] - assert sim.set_gizmo_visibility.call_args_list == [ - call("preview_robot", visible=False, control_part="left_arm"), - call("preview_robot", visible=False, control_part="right_arm"), - ] - - -def test_preview_skips_ik_gizmo_for_vectorized_environment() -> None: - """Native IK Gizmos remain limited to one simulated environment.""" - sim = MagicMock() - sim.is_window_opened = True - env = SimpleNamespace( - unwrapped=SimpleNamespace( - sim=sim, - robot=SimpleNamespace(uid="preview_robot"), - num_envs=2, - ) - ) - - gizmo_keys = run_env._enable_preview_ik_gizmos(env) - - assert gizmo_keys == () - sim.enable_gizmo.assert_not_called() - - -def test_preview_loop_services_native_ik_gizmo_while_waiting() -> None: - """Each input timeout advances Gizmo processing and one physics step.""" - physics_dt = 0.02 - sim = MagicMock() - sim.sim_config = SimpleNamespace(physics_dt=physics_dt) - env = SimpleNamespace(unwrapped=SimpleNamespace(sim=sim)) - control_input = _PreviewInput([None, "q"]) - - run_env._run_preview_loop( - env, - control_input, - (("preview_robot", "arm"),), - ) - - sim.update.assert_called_once_with(physics_dt, step=1) - assert control_input.timeouts == [physics_dt, physics_dt] - - -def test_preview_terminal_i_toggles_ik_gizmo() -> None: - """Terminal I mirrors the native-window visibility hotkey.""" - physics_dt = 0.02 - sim = MagicMock() - sim.sim_config = SimpleNamespace(physics_dt=physics_dt) - sim.toggle_gizmo_visibility.return_value = True - env = SimpleNamespace(unwrapped=SimpleNamespace(sim=sim)) - - run_env._run_preview_loop( - env, - _PreviewInput(["i", "q"]), - (("preview_robot", "arm"),), - ) - - sim.toggle_gizmo_visibility.assert_called_once_with( - "preview_robot", - control_part="arm", - ) - - -def test_preview_loop_services_viser_without_native_ik_gizmo() -> None: - """Viser preview keeps processing browser interaction commands.""" - physics_dt = 0.02 - sim = MagicMock() - sim.sim_config = SimpleNamespace( - physics_dt=physics_dt, - visualization=SimpleNamespace(backend="viser"), - ) - env = SimpleNamespace(unwrapped=SimpleNamespace(sim=sim)) - - run_env._run_preview_loop(env, _PreviewInput([None, "q"]), ()) - - sim.update.assert_called_once_with(physics_dt, step=1) - - def test_replay_restores_wrapper_state_without_closing_caller_env(monkeypatch) -> None: """Replay leaves the environment close to its CLI owner.""" env = MagicMock() @@ -276,12 +146,8 @@ def test_replay_restores_wrapper_state_without_closing_caller_env(monkeypatch) - def test_preview_quit_returns_without_zero_exit(monkeypatch) -> None: """Preview quit lets CLI cleanup failures determine the process status.""" env = MagicMock() - env.unwrapped = SimpleNamespace( - sim=SimpleNamespace(sim_config=SimpleNamespace(physics_dt=0.02)), - robot=None, - ) env.reset.return_value = (None, {}) - monkeypatch.setattr(run_env, "_ReplayControlInput", lambda: _PreviewInput(["q"])) + monkeypatch.setattr("builtins.input", lambda: "q") run_env.preview(env) diff --git a/tests/sim/objects/test_gizmo.py b/tests/sim/objects/test_gizmo.py index d6a138105..2cf91a13a 100644 --- a/tests/sim/objects/test_gizmo.py +++ b/tests/sim/objects/test_gizmo.py @@ -16,7 +16,6 @@ from __future__ import annotations -import threading from types import SimpleNamespace import numpy as np @@ -24,7 +23,12 @@ import torch import embodichain.lab.sim.objects.gizmo as gizmo_module -from embodichain.lab.sim.objects.gizmo import Gizmo, GizmoCfg, _RobotGizmoAdapter +from embodichain.lab.sim.objects.gizmo import ( + Gizmo, + GizmoCfg, + _RobotGizmoAdapter, + create_robot_ik_gizmo_controller, +) class _FakeAdapterRobot: @@ -36,7 +40,8 @@ def __init__(self) -> None: self.joint_names = ["joint_a", "joint_mimic", "joint_b"] self.link_names = ["base_link", "tool_link"] self.device = torch.device("cpu") - self.cfg = SimpleNamespace(solver_cfg=None) + self.uid = "robot" + self.cfg = SimpleNamespace(solver_cfg=None, fpath="robot.urdf") self.current_qpos = torch.tensor([[0.1, 0.2, 0.3]], dtype=torch.float32) self.target_qpos = torch.tensor([[0.4, 0.5, 0.6]], dtype=torch.float32) self.write_calls: list[dict[str, object]] = [] @@ -110,86 +115,67 @@ def test_robot_adapter_rejects_wrong_qpos_shape() -> None: def test_robot_native_ik_chain_can_be_configured_without_solver() -> None: - gizmo = object.__new__(Gizmo) - gizmo.cfg = GizmoCfg( + cfg = GizmoCfg( ik_root_link_name="base_link", ik_end_link_name="tool_link", ) - gizmo._control_part = "arm" - root_link, end_link, tcp_pose = gizmo._resolve_robot_ik_chain(_FakeAdapterRobot()) + root_link, end_link, tcp_pose = gizmo_module._resolve_robot_ik_chain( + _FakeAdapterRobot(), + "arm", + cfg, + ) assert (root_link, end_link) == ("base_link", "tool_link") np.testing.assert_allclose(tcp_pose, np.eye(4)) -def test_robot_update_delegates_to_dexsim_ik_controller() -> None: - calls: list[int] = [] - - class _Controller: - def update(self, *, iterations: int) -> None: - calls.append(iterations) - - gizmo = object.__new__(Gizmo) - gizmo.target = object() - gizmo._ik_controller = _Controller() - gizmo.cfg = GizmoCfg(ik_iterations=12) - - gizmo.update() - - assert calls == [12] - +def test_native_robot_factory_returns_dexsim_owned_controllers(monkeypatch) -> None: + robot = _FakeAdapterRobot() + adapter = _RobotGizmoAdapter(robot, "arm") + solver = object() + monkeypatch.setattr( + gizmo_module, + "_build_robot_ik", + lambda robot, control_part, cfg: ( + adapter, + solver, + "tool_link", + np.eye(4, dtype=np.float32), + ), + ) -def test_destroy_removes_gizmo_from_dexsim_environment() -> None: - class _DexsimGizmo: - def __init__(self) -> None: - self.detached = False + class _InputController: + pass - def set_flush_localpose_callback(self, callback: object | None) -> None: - pass + class _IKController: + def __init__(self, *args, **kwargs) -> None: + self.args = args + self.kwargs = kwargs - def set_transform_flush_callback(self, callback: object | None) -> None: - pass + import dexsim.engine + import dexsim.kit.ik - def set_visible(self, visible: bool) -> None: - pass + monkeypatch.setattr(dexsim.engine, "GizmoController", _InputController) + monkeypatch.setattr(dexsim.kit.ik, "IKGizmoController", _IKController) + window = SimpleNamespace(controls=[]) + window.add_input_control = window.controls.append + world = SimpleNamespace(get_windows=lambda: window) - def detach_parent(self) -> None: - self.detached = True + controller, input_controller = create_robot_ik_gizmo_controller( + robot, + world=world, + ) - class _Environment: - def __init__(self) -> None: - self.removed: object | None = None - - def remove_gizmo(self, gizmo: object) -> None: - self.removed = gizmo - - native_gizmo = _DexsimGizmo() - environment = _Environment() - gizmo = object.__new__(Gizmo) - gizmo._env = environment - gizmo._gizmo = native_gizmo - gizmo._proxy_cube = None - gizmo._ik_controller = None - gizmo._ik_solver = None - gizmo._ik_model = None - gizmo._robot_adapter = None - gizmo._state_lock = threading.RLock() - gizmo._interaction_owner = None - gizmo._pending_target_transform = None - gizmo._desired_target_transform = None - gizmo.target = object() - gizmo._target_type = "rigid_object" - - gizmo.destroy() - - assert environment.removed is native_gizmo - assert native_gizmo.detached is True - assert gizmo._gizmo is None + assert controller.args[:3] == (world, adapter, solver) + assert controller.kwargs["follow_robot_base"] is True + assert isinstance(input_controller, _InputController) + assert window.controls == [input_controller] class _RigidObject: def __init__(self) -> None: + self.num_instances = 1 self.device = torch.device("cpu") self.cfg = SimpleNamespace(uid="cube") self.pose = torch.eye(4, dtype=torch.float32).unsqueeze(0) @@ -212,22 +198,12 @@ class _Camera(_RigidObject): pass -def _patch_headless_dexsim(monkeypatch) -> None: - monkeypatch.setattr(gizmo_module.dexsim, "get_world_num", lambda: 1) - monkeypatch.setattr( - gizmo_module.dexsim, - "default_world", - lambda: SimpleNamespace(get_env=lambda: object()), - ) - - def test_headless_gizmo_applies_shared_pose_and_arbitrates_sources( monkeypatch, ) -> None: monkeypatch.setattr(gizmo_module, "RigidObject", _RigidObject) - _patch_headless_dexsim(monkeypatch) target = _RigidObject() - gizmo = Gizmo(target, enable_native=False) + gizmo = Gizmo(target) pose = torch.eye(4, dtype=torch.float32).unsqueeze(0) pose[0, :3, 3] = torch.tensor([0.2, 0.3, 0.4]) @@ -237,16 +213,14 @@ def test_headless_gizmo_applies_shared_pose_and_arbitrates_sources( gizmo.update() assert gizmo.end_interaction("viser:client-a") - assert not gizmo.native_enabled assert target.set_calls[-1][1] == [0] torch.testing.assert_close(target.pose, pose) def test_headless_camera_gizmo_uses_shared_pose_path(monkeypatch) -> None: monkeypatch.setattr(gizmo_module, "Camera", _Camera) - _patch_headless_dexsim(monkeypatch) target = _Camera() - gizmo = Gizmo(target, enable_native=False) + gizmo = Gizmo(target) pose = torch.eye(4, dtype=torch.float32).unsqueeze(0) pose[0, 2, 3] = 1.2 @@ -265,7 +239,6 @@ def test_headless_robot_gizmo_uses_dexsim_newton_ik(monkeypatch) -> None: instead of calling the EmbodiChain ``compute_ik`` solver. """ monkeypatch.setattr(gizmo_module, "Robot", _FakeAdapterRobot) - _patch_headless_dexsim(monkeypatch) target = _FakeAdapterRobot() solved_qpos = np.array([0.4, -0.2], dtype=np.float32) @@ -291,16 +264,15 @@ def qpos_for_joint_names(self, joint_names, fallback_qpos): def _inject_solver(self) -> None: self._robot_adapter = _RobotGizmoAdapter(target, "arm") self._ik_solver = fake_solver - self._native_robot_end_link = "tool_link" - self._native_robot_tcp_pose = np.eye(4, dtype=np.float32) + self._robot_end_link = "tool_link" + self._robot_tcp_pose = np.eye(4, dtype=np.float32) monkeypatch.setattr(Gizmo, "_setup_robot_ik_solver", _inject_solver) - gizmo = Gizmo(target, control_part="arm", enable_native=False) + gizmo = Gizmo(target, control_part="arm") pose = torch.eye(4, dtype=torch.float32).unsqueeze(0) pose[0, 0, 3] = 0.5 - assert not gizmo.native_enabled assert gizmo.request_local_pose(pose, source_id="viser:client-a") gizmo.update() diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index 292de5344..6074e1cb2 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -133,14 +133,6 @@ def enable_entity_gizmo(self, config: object | None = None) -> object: self.entity_gizmo = FakeEntityGizmo() return self.entity_gizmo - def disable_entity_gizmo(self) -> None: - if self.entity_gizmo is not None: - self.entity_gizmo.active = False - self.entity_gizmo = None - - def get_entity_gizmo(self) -> object | None: - return self.entity_gizmo - def open_window(self) -> None: self.window_open_count += 1 self.window_closed = False @@ -219,11 +211,9 @@ def _make_sim_manager(window: object | None = None) -> SimulationManager: sim.sim_config = SimpleNamespace( width=64, height=48, - enable_entity_gizmo_on_window_open=True, visualization=SimpleNamespace(backend="none"), ) sim._window = window - sim._entity_gizmo_config = None sim._window_record_state = None sim._window_record_camera = None sim._window_record_save_threads = [] @@ -386,16 +376,26 @@ def _make_pick_sim_manager(pick_commands, resolve): enabled: list = [] disabled: list = [] - def fake_enable(uid, control_part=None, gizmo_cfg=None, *, enable_native=None): + def fake_enable(uid, control_part=None, gizmo_cfg=None): enabled.append((uid, control_part)) - return SimpleNamespace(control_part=control_part) + gizmo = SimpleNamespace(control_part=control_part) + gizmo_key = f"{uid}:{control_part}" if control_part else uid + sim._gizmos[gizmo_key] = gizmo + return gizmo def fake_disable(uid, control_part=None): disabled.append((uid, control_part)) + gizmo_key = f"{uid}:{control_part}" if control_part else uid + sim._gizmos.pop(gizmo_key, None) sim.enable_gizmo = fake_enable sim.disable_gizmo = fake_disable - sim.has_gizmo = lambda uid, control_part=None: True + sim.has_gizmo = ( + lambda uid, control_part=None: ( + f"{uid}:{control_part}" if control_part else uid + ) + in sim._gizmos + ) sim.sim_config = SimpleNamespace( visualization=SimpleNamespace(allow_commands=True), ) @@ -499,6 +499,48 @@ def test_process_pick_commands_is_noop_for_already_picked_target() -> None: assert sim._picker_gizmo == ("cube", None) +@pytest.mark.parametrize( + ("node_id", "target", "gizmo_key"), + [ + ("env:0/rigid:cube", ("cube", "rigid"), "cube"), + ("env:0/robot:ur10", ("ur10", "robot"), "ur10:arm"), + ], +) +def test_process_pick_commands_preserves_user_created_gizmo( + node_id: str, + target: tuple[str, str], + gizmo_key: str, +) -> None: + pick_commands = ( + PickCommand( + run_id="run", + scene_revision=2, + client_id="client-a", + node_id=node_id, + ), + PickCommand( + run_id="run", + scene_revision=2, + client_id="client-a", + node_id=None, + ), + ) + sim, enabled, disabled = _make_pick_sim_manager( + pick_commands, + lambda node_id: target, + ) + user_gizmo = SimpleNamespace(control_part=None) + sim._gizmos[gizmo_key] = user_gizmo + + processed = sim.process_pick_commands() + + assert processed == 2 + assert enabled == [] + assert disabled == [] + assert sim._gizmos[gizmo_key] is user_gizmo + assert sim._picker_gizmo is None + + def test_process_pick_commands_ignores_stale_scene_revision() -> None: pick_commands = ( PickCommand( @@ -632,27 +674,14 @@ def test_open_window_is_idempotent() -> None: sim._world.open_window.assert_not_called() -def test_entity_gizmo_lifecycle_delegates_to_dexsim_world() -> None: +def test_entity_gizmo_delegates_to_dexsim_and_excludes_default_plane() -> None: sim = _make_sim_manager() config = object() controller = sim.enable_entity_gizmo(config) - assert controller is sim._world.get_entity_gizmo() + assert controller is sim._world.entity_gizmo assert sim._world.entity_gizmo_configs == [config] - assert sim.get_entity_gizmo() is controller - assert sim.has_entity_gizmo() is True - assert sim.disable_entity_gizmo() is True - assert controller.active is False - assert sim.has_entity_gizmo() is False - assert sim.disable_entity_gizmo() is False - - -def test_entity_gizmo_registers_default_plane_as_static_exclusion() -> None: - sim = _make_sim_manager() - - controller = sim.enable_entity_gizmo() - assert controller.external_targets == [ ( SimulationManager._DEFAULT_PLANE_GIZMO_TARGET_ID, @@ -663,94 +692,23 @@ def test_entity_gizmo_registers_default_plane_as_static_exclusion() -> None: ] -def test_open_window_enables_entity_gizmo_by_default() -> None: +def test_open_window_does_not_enable_entity_gizmo_implicitly() -> None: sim = _make_sim_manager() assert sim.open_window() assert sim.is_window_opened is True assert sim._world.window_open_count == 1 - assert sim.has_entity_gizmo() is True - assert sim._world.entity_gizmo_configs == [None] - - -def test_open_window_supports_view_only_opt_out() -> None: - sim = _make_sim_manager() - - assert sim.open_window(enable_entity_gizmo=False) - - assert sim.is_window_opened is True - assert sim.has_entity_gizmo() is False assert sim._world.entity_gizmo_configs == [] -def test_open_window_view_only_opt_out_disables_active_controller() -> None: - sim = _make_sim_manager() - controller = sim.enable_entity_gizmo() - - assert sim.open_window(enable_entity_gizmo=False) - - assert controller.active is False - assert sim.has_entity_gizmo() is False - - -def test_open_window_respects_configured_entity_gizmo_default() -> None: - sim = _make_sim_manager() - sim.sim_config.enable_entity_gizmo_on_window_open = False - - assert sim.open_window() - - assert sim.is_window_opened is True - assert sim.has_entity_gizmo() is False - - -def test_open_window_tolerates_dexsim_without_entity_gizmo_api() -> None: - sim = _make_sim_manager() - window = object() - sim._world = SimpleNamespace( - open_window=lambda: None, - get_windows=lambda: window, - ) - - assert sim.open_window() - - assert sim.is_window_opened is True - assert sim._window is window - assert sim.has_entity_gizmo() is False - - -def test_open_window_preserves_active_entity_gizmo_configuration() -> None: - sim = _make_sim_manager(window=object()) - config = object() - controller = sim.enable_entity_gizmo(config) - - assert sim.open_window() - - assert sim.get_entity_gizmo() is controller - assert sim._world.entity_gizmo_configs == [config] - assert sim._world.window_open_count == 0 - - -def test_reopened_window_restores_last_entity_gizmo_configuration() -> None: - sim = _make_sim_manager(window=object()) - config = object() - sim.enable_entity_gizmo(config) - sim.close_window() - - assert sim.open_window() - - assert sim.has_entity_gizmo() is True - assert sim._world.entity_gizmo_configs == [config, config] - - -def test_close_window_disables_entity_gizmo() -> None: +def test_close_window_leaves_entity_gizmo_lifecycle_to_dexsim() -> None: sim = _make_sim_manager(window=object()) controller = sim.enable_entity_gizmo() sim.close_window() - assert controller.active is False - assert sim.has_entity_gizmo() is False + assert controller.active is True assert sim._world.window_closed is True assert sim.is_window_opened is False @@ -853,6 +811,25 @@ def test_remove_asset_marks_visualization_topology_dirty() -> None: assert runtime.stopped +def test_stop_visualization_releases_only_picker_owned_gizmo() -> None: + sim, runtime = _make_visualization_sim_manager() + picker_gizmo = MagicMock() + user_gizmo = MagicMock() + sim._gizmos = { + "picked": picker_gizmo, + "user": user_gizmo, + } + sim._picker_gizmo = ("picked", None) + + sim.stop_visualization() + + assert runtime.stopped + assert sim._picker_gizmo is None + assert sim._gizmos == {"user": user_gizmo} + picker_gizmo.destroy.assert_called_once_with() + user_gizmo.destroy.assert_not_called() + + def test_add_stereo_camera_marks_visualization_topology_dirty() -> None: sim = object.__new__(SimulationManager) sensor = object.__new__(sim_manager_module.StereoCamera) diff --git a/tests/visualization/test_runtime.py b/tests/visualization/test_runtime.py index ee03b4c59..2505d902a 100644 --- a/tests/visualization/test_runtime.py +++ b/tests/visualization/test_runtime.py @@ -38,6 +38,8 @@ VisualizationRuntime, ) from embodichain.lab.visualization.backends.base import VisualizationBackend +from embodichain.lab.visualization.protocol import PickCommand +from embodichain.lab.visualization.runtime import PickCommandQueue REPLAY_CURRENT_STEP = 6 REPLAY_MAX_STEP = 9 @@ -124,6 +126,21 @@ def test_joint_control_queue_keeps_latest_value_per_control() -> None: ] +def test_pick_command_queue_keeps_latest_click_in_arrival_order() -> None: + commands = PickCommandQueue(maxsize=3) + + commands.put(PickCommand("run", 1, "client-a", "node-a-1")) + commands.put(PickCommand("run", 1, "client-b", "node-b")) + commands.put(PickCommand("run", 1, "client-a", "node-a-2")) + + drained = commands.drain() + + assert [(command.client_id, command.node_id) for command in drained] == [ + ("client-b", "node-b"), + ("client-a", "node-a-2"), + ] + + @dataclass class _Exporter: published: threading.Event From 5607b3b04ad038662308af55cc3e2cf8849e1ab4 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 5 Sep 2026 16:20:44 +0800 Subject: [PATCH 05/14] feat(sim): support configured IK solvers with native dexsim gizmos --- agent_context/MAP.yaml | 9 + agent_context/topics/ik-solvers/ik-solvers.md | 10 + .../sim-visualization/sim-visualization.md | 15 ++ .../simulation-system/simulation-system.md | 17 ++ .../embodichain.lab.sim.objects.rst | 11 + .../embodichain.lab.visualization.rst | 8 + docs/source/api_reference/public_api.rst | 22 ++ embodichain/lab/sim/objects/gizmo.py | 225 +++++++++++------ embodichain/lab/sim/utility/__init__.py | 2 + .../lab/visualization/backends/viser.py | 16 +- examples/sim/gizmo/gizmo_robot.py | 64 ++--- pyproject.toml | 2 +- tests/sim/objects/test_gizmo.py | 230 ++++++++++++++++++ tests/visualization/test_viser_backend.py | 49 ++++ 14 files changed, 577 insertions(+), 103 deletions(-) diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 261d596f4..3a9fe925c 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -21,6 +21,8 @@ topics: - simulation - SimulationManager - SimulationManagerCfg + - EntityGizmoManipulator + - enable_entity_gizmo - DexSim - world - arena @@ -344,11 +346,18 @@ topics: - deformable - soft body - cloth + - gizmo + - IKGizmoController + - create_robot_ik_gizmo_controller + - GizmoCfg + - ScenePicker + - PickCommand paths: - topics/sim-visualization/sim-visualization.md source_of_truth: - embodichain/lab/sim/sim_manager.py - embodichain/lab/visualization/ + - embodichain/lab/sim/objects/gizmo.py - embodichain/lab/gym/utils/gym_utils.py - embodichain/lab/gym/envs/base_env.py - embodichain/lab/scripts/preview_asset.py diff --git a/agent_context/topics/ik-solvers/ik-solvers.md b/agent_context/topics/ik-solvers/ik-solvers.md index ac0ddcbf0..2de569ecc 100644 --- a/agent_context/topics/ik-solvers/ik-solvers.md +++ b/agent_context/topics/ik-solvers/ik-solvers.md @@ -34,6 +34,16 @@ path for performance. ## Solver Hierarchy +Gizmo solver selection belongs to `embodichain/lab/sim/objects/gizmo.py`, not +to a new solver subclass. Native DexSim `IKGizmoController` remains the +controller in both modes: `GizmoCfg.ik_solver="dexsim"` selects Newton IK; +`"embodichain"` adapts `robot.get_solver(control_part)` for the native window +and Viser. It preserves configured solver convergence/limits, maps joint names +between solver and robot order, and holds current qpos on failure/non-finite +results. Gizmo root/end links must match the configured solver; TCP overrides +are converted without mutating that solver. See the robot gizmo example's +`--ik-solver pink` option for an executable Pink configuration. + ``` SolverCfg (@configclass, abstract) ├── SRSSolverCfg diff --git a/agent_context/topics/sim-visualization/sim-visualization.md b/agent_context/topics/sim-visualization/sim-visualization.md index 763e25979..86292e613 100644 --- a/agent_context/topics/sim-visualization/sim-visualization.md +++ b/agent_context/topics/sim-visualization/sim-visualization.md @@ -250,6 +250,21 @@ Soft bodies and cloth require GPU physics. Their live vertices are sampled at ## Browser Controls and Overlays +Gizmo implementation lives in `embodichain/lab/sim/objects/gizmo.py`. Native +windows delegate object picking/manipulation to DexSim's entity gizmo and robot +targets to its `IKGizmoController`. Viser transports poses to the simulation +thread. Both robot paths use Newton IK by default or, with +`GizmoCfg(ik_solver="embodichain")`, reuse the configured control-part solver +(including Pink). Chain roots follow the live robot link, including upstream +joint motion; TCP overrides adapt targets without changing the shared solver. + +Viser click picking runs on the visualization worker via the existing GUI +event queue. A manifest invalidates cached pick poses until its matching frame +arrives; stale clicks are dropped. The manager validates `PickCommand` run and +revision before attaching a gizmo and tracks picker ownership independently +from explicitly created gizmos. Clearing selection releases only the picker +gizmo. Native entity selection remains entirely DexSim-owned. + The browser GUI has: - **Environments** — visibility per exported environment; diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index 7ecd402b9..6a67718ed 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -110,6 +110,23 @@ do not belong to the simulation object; use ## Configuration Flow +### Gizmo ownership + +Native entity manipulation belongs to DexSim 0.5.0. Call +`SimulationManager.enable_entity_gizmo(config)` to obtain the world-owned +`EntityGizmoManipulator`; the manager only registers its default plane as a +static external target. Opening a window does not implicitly enable entity +gizmos. Query/disable through `sim.get_world().get_entity_gizmo()` and +`disable_entity_gizmo()`; DexSim owns window detach/reopen and controller state. + +`create_robot_ik_gizmo_controller()` in `objects/gizmo.py` returns the native +DexSim IK controller and input controller. The caller retains both and calls +the IK controller's `update()` per frame. `SimulationManager.enable_gizmo()` +creates Viser controls for robots, rigid objects, or cameras. Both robot paths +default to native Newton IK; `GizmoCfg(ik_solver="embodichain")` adapts the +control part's existing solver, such as PinkSolver. Both support one environment +and write only selected non-mimic joint drive targets through `Robot`. + `SimulationManagerCfg` owns window size, headless mode, rendering, GPU/CPU selection, arena count and spacing, physics timestep, physics and GPU-memory settings, recording, profiling, and browser visualization. diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst index 9320690b0..7c286857a 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst @@ -194,6 +194,17 @@ Gizmo .. autofunction:: create_robot_ik_gizmo_controller +The native controller defaults to DexSim Newton IK. With a ``PinkSolverCfg`` +(or another EmbodiChain solver) configured for the robot's control part, pass +``GizmoCfg(ik_solver="embodichain")`` to select that solver for either a native +controller or a Viser gizmo. Its iteration limits and convergence settings +remain owned by the configured solver; ``ik_iterations`` applies to Newton IK. +Only the selected control part's drive targets are written, and failed +EmbodiChain IK solutions preserve the current joint positions. + +The runnable example ``examples/sim/gizmo/gizmo_robot.py`` exposes +``--ik-solver dexsim|pytorch|pink`` for both the native window and ``--viser``. + Rigid Constraint ---------------- diff --git a/docs/source/api_reference/embodichain/embodichain.lab.visualization.rst b/docs/source/api_reference/embodichain/embodichain.lab.visualization.rst index 2a52d2432..fc3b01c71 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.visualization.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.visualization.rst @@ -103,6 +103,14 @@ Scene Export :members: :undoc-members: +``PickCommand`` identifies a node in a specific simulation run and scene +revision. An empty selection releases only the gizmo created by click picking; +explicitly configured gizmos retain their ownership. + +.. autoclass:: PickCommand + :members: + :undoc-members: + .. autoclass:: JointControlSpec :members: :undoc-members: diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index ea1aa59e5..1c05b3565 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -865,10 +865,15 @@ embodichain.lab.sim.objects.gizmo .. currentmodule:: embodichain.lab.sim.objects.gizmo +Native robot targets use DexSim's controller with Newton IK by default. +Set ``GizmoCfg.ik_solver="embodichain"`` to reuse the robot control part's +configured solver, including PinkSolver; Viser uses the same solver adapter. + .. autosummary:: Gizmo GizmoCfg + create_robot_ik_gizmo_controller embodichain.lab.sim.objects.rigid_object ---------------------------------------- @@ -1499,6 +1504,22 @@ embodichain.lab.visualization.cli add_viser_args_to_parser visualization_cfg_from_args +embodichain.lab.visualization.picker +------------------------------------ + +.. currentmodule:: embodichain.lab.visualization.picker + +Browser picking caches triangle geometry and returns the closest node hit by +a world-space ray. The Viser worker pairs this geometry with poses from the +same scene revision before producing a pick command. + +.. autosummary:: + + ScenePicker + +.. autoclass:: ScenePicker + :members: + embodichain.lab.visualization.protocol -------------------------------------- @@ -1520,6 +1541,7 @@ embodichain.lab.visualization.protocol JointControlSpec JointControlState MeshGeometry + PickCommand PointCloudOverlay SceneFrame SceneManifest diff --git a/embodichain/lab/sim/objects/gizmo.py b/embodichain/lab/sim/objects/gizmo.py index cb0352c7e..730cfd9e9 100644 --- a/embodichain/lab/sim/objects/gizmo.py +++ b/embodichain/lab/sim/objects/gizmo.py @@ -19,12 +19,16 @@ from __future__ import annotations import threading -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal import dexsim import numpy as np import torch -import warp as wp +from dexsim.kit.ik.pose import ( + local_pose_from_world, + pose_from_position_rotation, + rotation_matrix_to_quat_xyzw, +) from dexsim.types import InputKey from embodichain.lab.sim.common import BatchEntity @@ -36,6 +40,7 @@ if TYPE_CHECKING: from dexsim.engine import GizmoController from dexsim.kit.ik import IKGizmoController, NewtonChainIK + from embodichain.lab.sim.solvers import BaseSolver __all__ = ["Gizmo", "GizmoCfg", "create_robot_ik_gizmo_controller"] @@ -62,6 +67,14 @@ class GizmoCfg: rings_size: float = 0.01 """Thickness of the rotation rings.""" + ik_solver: Literal["dexsim", "embodichain"] = "dexsim" + """Use native Newton IK, or the control part's configured EmbodiChain solver. + + The native window always uses DexSim's ``IKGizmoController``. Selecting + ``embodichain`` reuses ``robot.get_solver(control_part)`` (for example, + PinkSolver), including that solver's convergence and joint-limit settings. + """ + ik_root_link_name: str | None = None """Robot IK chain root link, or the configured solver root when omitted.""" @@ -109,6 +122,8 @@ def __init__(self, robot: Robot, control_part: str, env_id: int = 0) -> None: self.env_id = env_id self.joint_ids = list(joint_ids) self.joint_names = [robot.joint_names[index] for index in self.joint_ids] + self.root_link_name: str | None = None + self.model_root_inverse = np.eye(4, dtype=np.float32) def get_current_qpos(self) -> np.ndarray: """Return current selected joint positions in DexSim joint-name order.""" @@ -131,7 +146,9 @@ def get_actived_joint_names(self) -> list[str]: return self.joint_names.copy() def get_world_pose(self) -> np.ndarray: - """Return the selected robot instance's root pose as a matrix.""" + """Map the IK model frame into the selected instance's arena frame.""" + if self.root_link_name is not None: + return self.get_link_pose(self.root_link_name) @ self.model_root_inverse pose = self.robot.get_local_pose(to_matrix=True)[self.env_id] return pose.detach().cpu().numpy().astype(np.float32, copy=True) @@ -214,19 +231,104 @@ def _resolve_robot_ik_chain( return root_link, end_link, tcp_matrix +class _EmbodiChainIK: + """Adapt BaseSolver to the solver contract consumed by DexSim and Viser.""" + + def __init__( + self, adapter: _RobotGizmoAdapter, solver: BaseSolver, tcp_pose: np.ndarray + ) -> None: + self.solver = solver + self.joint_names = list(solver.joint_names or adapter.joint_names) + if set(self.joint_names) != set(adapter.joint_names): + raise ValueError( + "IK solver joints must match the control part's active joints." + ) + self._tcp_inverse = np.linalg.inv(tcp_pose) + self._solution: dict[str, float] = {} + pose = ( + local_pose_from_world( + adapter.get_world_pose(), adapter.get_link_pose(solver.end_link_name) + ) + @ tcp_pose + ) + self._target_state: dict[str, np.ndarray] = {} + self.set_target_pose(pose[:3, 3], rotation_matrix_to_quat_xyzw(pose[:3, :3])) + self.reset_target_state_changed() + + def target_state(self) -> dict[str, np.ndarray]: + return self._target_state + + def set_target_pose(self, position: np.ndarray, rotation: np.ndarray) -> None: + self._target_state.update( + position=np.array(position), rotation=np.array(rotation) + ) + + def target_state_changed(self) -> bool: + return any( + not np.allclose(value, self._snapshot[key], atol=1e-5) + for key, value in self._target_state.items() + ) + + def reset_target_state_changed(self) -> None: + self._snapshot = { + key: value.copy() for key, value in self._target_state.items() + } + + def solve( + self, + joint_names: list[str], + current_qpos: np.ndarray, + *, + iterations: int | None = None, + ) -> None: + # EmbodiChain solvers retain their own iteration/convergence configuration. + del iterations + self._solution.clear() + seed = dict(zip(joint_names, current_qpos)) + target = pose_from_position_rotation(**self._target_state) + # The gizmo TCP may differ from the shared solver's TCP. Adapt the + # target without mutating the solver used by the rest of the robot. + target = target @ self._tcp_inverse @ self.solver.get_tcp() + success, qpos = self.solver.get_ik( + torch.as_tensor( + target, dtype=torch.float32, device=self.solver.device + ).unsqueeze(0), + qpos_seed=torch.tensor( + [[seed[name] for name in self.joint_names]], + dtype=torch.float32, + device=self.solver.device, + ), + return_all_solutions=False, + ) + if bool(success.all()) and bool(torch.isfinite(qpos).all()): + values = qpos.detach().cpu().numpy().reshape(len(self.joint_names)) + self._solution = dict(zip(self.joint_names, values)) + + def qpos_for_joint_names( + self, joint_names: list[str], fallback_qpos: np.ndarray + ) -> np.ndarray: + return np.array( + [ + self._solution.get(name, value) + for name, value in zip(joint_names, fallback_qpos) + ], + dtype=np.float32, + ) + + def _build_robot_ik( robot: Robot, control_part: str, cfg: GizmoCfg, -) -> tuple[_RobotGizmoAdapter, NewtonChainIK, str, np.ndarray]: - from dexsim.kit.ik import NewtonChainIK, build_newton_model_from_urdf - +) -> tuple[_RobotGizmoAdapter, NewtonChainIK | _EmbodiChainIK, str, np.ndarray]: if robot.num_instances != 1: raise RuntimeError( "Robot Gizmo supports exactly one environment, " f"but the robot has {robot.num_instances} instances." ) - if cfg.ik_iterations <= 0: + if cfg.ik_solver not in {"dexsim", "embodichain"}: + raise ValueError("ik_solver must be 'dexsim' or 'embodichain'.") + if cfg.ik_solver == "dexsim" and cfg.ik_iterations <= 0: raise ValueError("ik_iterations must be greater than zero.") root_link, end_link, tcp_pose = _resolve_robot_ik_chain( @@ -235,6 +337,24 @@ def _build_robot_ik( cfg, ) adapter = _RobotGizmoAdapter(robot, control_part) + adapter.root_link_name = root_link + if cfg.ik_solver == "embodichain": + solver = ( + robot.get_solver(control_part) if robot.cfg.solver_cfg is not None else None + ) + if solver is None: + raise ValueError( + f"Control part {control_part!r} needs a configured EmbodiChain solver." + ) + if (root_link, end_link) != (solver.root_link_name, solver.end_link_name): + raise ValueError( + "Gizmo IK links must match the configured EmbodiChain solver." + ) + return adapter, _EmbodiChainIK(adapter, solver, tcp_pose), end_link, tcp_pose + + import warp as wp + from dexsim.kit.ik import NewtonChainIK, build_newton_model_from_urdf + with wp.ScopedDevice(cfg.ik_device or str(robot.device)): model = build_newton_model_from_urdf(robot.cfg.fpath, hide_visuals=True) solver = NewtonChainIK( @@ -245,6 +365,14 @@ def _build_robot_ik( tcp_pose=tcp_pose, ) + # Newton preserves the start link's URDF rest transform in its reduced + # model. Follow the live chain root (including upstream joints) while + # retaining that model frame, rather than assuming the robot root is it. + root_index = list(solver.model.body_label).index(solver.info.start_link) + root_pose = solver.state.body_q.numpy()[root_index] + adapter.model_root_inverse = np.linalg.inv( + pose_from_position_rotation(root_pose[:3], root_pose[3:]) + ) solver.set_qpos_from_joint_names( adapter.get_actived_joint_names(), adapter.get_current_qpos(), @@ -272,7 +400,7 @@ def create_robot_ik_gizmo_controller( Args: robot: Single-instance EmbodiChain robot. control_part: Robot control part driven by IK. - cfg: IK chain and target appearance settings. + cfg: IK chain, solver selection, and target appearance settings. world: DexSim world, or the current default world when omitted. Returns: @@ -295,7 +423,6 @@ def create_robot_ik_gizmo_controller( control_part = _resolve_control_part(robot, control_part) adapter, solver, _, _ = _build_robot_ik(robot, control_part, cfg) input_controller = GizmoController() - window.add_input_control(input_controller) controller = IKGizmoController( world, adapter, @@ -307,6 +434,7 @@ def create_robot_ik_gizmo_controller( gizmo_scale=cfg.ik_gizmo_scale, name=f"{robot.uid}_{control_part}_ik", ) + window.add_input_control(input_controller) return controller, input_controller @@ -347,7 +475,7 @@ def __init__( self._interaction_owner: str | None = None self._pending_target_transform: torch.Tensor | None = None self._desired_target_transform: torch.Tensor | None = None - self._ik_solver: NewtonChainIK | None = None + self._ik_solver: NewtonChainIK | _EmbodiChainIK | None = None self._robot_adapter: _RobotGizmoAdapter | None = None self._robot_end_link: str | None = None self._robot_tcp_pose: np.ndarray | None = None @@ -485,75 +613,28 @@ def cancel_interaction(self, source_prefix: str | None = None) -> bool: self._interaction_owner = None return True - def _update_camera_pose(self, target_transform: torch.Tensor) -> bool: - if self.target is None or not isinstance(self.target, Camera): - return False - try: - self.target.set_local_pose(target_transform, env_ids=[0]) - return True - except Exception as error: - logger.log_error(f"Error updating camera pose: {error}") - return False - - def _update_rigid_object_pose(self, target_transform: torch.Tensor) -> bool: - if self.target is None or not isinstance(self.target, RigidObject): - return False - try: - self.target.set_local_pose(target_transform, env_ids=[0]) - return True - except Exception as error: - logger.log_error(f"Error updating rigid-object pose: {error}") - return False - - def _update_robot_ik(self, target_transform: torch.Tensor) -> bool: - if self.target is None or not isinstance(self.target, Robot): - return False - if self._ik_solver is None or self._robot_adapter is None: - return False - try: - from dexsim.kit.ik.pose import ( - local_pose_from_world, - rotation_matrix_to_quat_xyzw, - ) - - base_pose = self._robot_adapter.get_world_pose() - target_pose = target_transform[0].detach().cpu().numpy().astype(np.float32) - base_local = local_pose_from_world(base_pose, target_pose) - joint_names = self._robot_adapter.get_actived_joint_names() - current_qpos = self._robot_adapter.get_current_qpos() - self._ik_solver.set_target_pose( - np.asarray(base_local[:3, 3], dtype=np.float32), - rotation_matrix_to_quat_xyzw(base_local[:3, :3]), - ) - self._ik_solver.solve( - joint_names, - current_qpos, - iterations=self.cfg.ik_iterations, - ) - solved_qpos = self._ik_solver.qpos_for_joint_names( - joint_names, - current_qpos, - ) - self._robot_adapter.set_target_qpos(solved_qpos) - return True - except Exception as error: - logger.log_error(f"Error in Viser Gizmo robot IK: {error}") - return False - def update(self) -> None: """Apply the latest queued Viser target pose on the simulation thread.""" with self._state_lock: pending = self._pending_target_transform self._pending_target_transform = None - if pending is None: + if pending is None or self.target is None: return - if self._target_type == "rigid_object": - self._update_rigid_object_pose(pending) - elif self._target_type == "robot": - self._update_robot_ik(pending) - elif self._target_type == "camera": - self._update_camera_pose(pending) + if self._robot_adapter is None: + self.target.set_local_pose(pending, env_ids=[0]) + return + adapter, solver = self._robot_adapter, self._ik_solver + base_local = local_pose_from_world( + adapter.get_world_pose(), pending[0].detach().cpu().numpy() + ) + joint_names = adapter.get_actived_joint_names() + current_qpos = adapter.get_current_qpos() + solver.set_target_pose( + base_local[:3, 3], rotation_matrix_to_quat_xyzw(base_local[:3, :3]) + ) + solver.solve(joint_names, current_qpos, iterations=self.cfg.ik_iterations) + adapter.set_target_qpos(solver.qpos_for_joint_names(joint_names, current_qpos)) def toggle_visibility(self) -> bool: """Toggle Viser visibility and return the new state.""" diff --git a/embodichain/lab/sim/utility/__init__.py b/embodichain/lab/sim/utility/__init__.py index c839d646f..ae4c254cf 100644 --- a/embodichain/lab/sim/utility/__init__.py +++ b/embodichain/lab/sim/utility/__init__.py @@ -16,6 +16,8 @@ """Helper utilities for simulation state conversion, mesh/geometry handling, configuration transforms, keyboard interaction, and action/solver adaptation.""" +from __future__ import annotations + from .sim_utils import * from .mesh_utils import * from .keyboard_utils import * diff --git a/embodichain/lab/visualization/backends/viser.py b/embodichain/lab/visualization/backends/viser.py index 931293271..2fde15f7f 100644 --- a/embodichain/lab/visualization/backends/viser.py +++ b/embodichain/lab/visualization/backends/viser.py @@ -244,7 +244,9 @@ def _(client: object) -> None: @self._server.scene.on_pointer_event("click") def _on_pick_click(event: object) -> None: - self._handle_pick_click(event) + self._gui_events.put( + _GuiEvent("pick", (self._run_id, self._scene_revision, event)) + ) self._pointer_handler = _on_pick_click @@ -392,7 +394,7 @@ def _handle_pick_click(self, event: object) -> None: empty space (no hit) clears the picker-owned Gizmo. The command is processed on the simulation thread. """ - if not self._pick_enabled: + if not self._pick_enabled or self._frame_positions is None: return sink = getattr(self, "_pick_command_sink", None) if sink is None or self._run_id is None: @@ -1123,6 +1125,11 @@ def publish_manifest(self, manifest: SceneManifest) -> None: geometry_by_id = { geometry.geometry_id: geometry for geometry in manifest.geometries } + # Pick topology and poses must always come from the same revision. + # Defer clicks until a matching frame arrives after this manifest. + self._frame_positions = None + self._frame_wxyz = None + self._frame_visible = None self._frame_node_ids = tuple(node.node_id for node in manifest.nodes) node_indices = { node_id: index for index, node_id in enumerate(self._frame_node_ids) @@ -1279,6 +1286,11 @@ def _apply_gui_events(self) -> None: event = self._gui_events.get_nowait() except queue.Empty: break + if event.category == "pick": + run_id, revision, click = event.value + if (run_id, revision) == (self._run_id, self._scene_revision): + self._handle_pick_click(click) + continue if event.category == "environment": env_id, visible = event.value self._env_visibility[int(env_id)] = bool(visible) diff --git a/examples/sim/gizmo/gizmo_robot.py b/examples/sim/gizmo/gizmo_robot.py index 477abb924..75d47ed75 100644 --- a/examples/sim/gizmo/gizmo_robot.py +++ b/examples/sim/gizmo/gizmo_robot.py @@ -24,7 +24,7 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.solvers import PytorchSolverCfg +from embodichain.lab.sim.solvers import PinkSolverCfg, PytorchSolverCfg from embodichain.lab.sim.objects import ( GizmoCfg, create_robot_ik_gizmo_controller, @@ -48,6 +48,12 @@ def main(): description="Create a simulation scene with SimulationManager" ) add_env_launcher_args_to_parser(parser) + parser.add_argument( + "--ik-solver", + choices=("dexsim", "pytorch", "pink"), + default="dexsim", + help="IK solver; the native window always uses DexSim's IKGizmoController.", + ) args = parser.parse_args() # Configure the simulation @@ -68,6 +74,30 @@ def main(): ur10_urdf_path = get_data_path("UniversalRobots/UR10/UR10.urdf") gripper_urdf_path = get_data_path("DH_PGC_140_50_M/DH_PGC_140_50_M.urdf") + # Native IK needs only chain metadata. Build an EmbodiChain solver only + # when explicitly selected, and share the same TCP configuration. + gizmo_cfg = GizmoCfg( + ik_solver="dexsim" if args.ik_solver == "dexsim" else "embodichain", + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ik_tcp_pose=[ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.12], + [0.0, 0.0, 0.0, 1.0], + ], + ) + solver_cfg = None + if args.ik_solver != "dexsim": + solver_type = PinkSolverCfg if args.ik_solver == "pink" else PytorchSolverCfg + solver_cfg = { + "arm": solver_type( + root_link_name=gizmo_cfg.ik_root_link_name, + end_link_name=gizmo_cfg.ik_end_link_name, + tcp=gizmo_cfg.ik_tcp_pose, + ) + } + # Create UR10 robot robot_cfg = RobotCfg( uid="ur10_gizmo_test", @@ -78,26 +108,14 @@ def main(): ] ), control_parts={ - "arm": ["JOINT[0-9]"], + "arm": ["Joint[0-9]"], "hand": ["FINGER[1-2]"], }, - solver_cfg={ - "arm": PytorchSolverCfg( - end_link_name="ee_link", - root_link_name="base_link", - tcp=[ - [0.0, 1.0, 0.0, 0.0], - [-1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.12], - [0.0, 0.0, 0.0, 1.0], - ], - num_samples=30, - ) - }, + solver_cfg=solver_cfg, drive_pros=JointDrivePropertiesCfg( - stiffness={"JOINT[0-9]": 1e4, "FINGER[1-2]": 1e2}, - damping={"JOINT[0-9]": 1e3, "FINGER[1-2]": 1e1}, - max_effort={"JOINT[0-9]": 1e5, "FINGER[1-2]": 1e3}, + stiffness={"Joint[0-9]": 1e4, "FINGER[1-2]": 1e2}, + damping={"Joint[0-9]": 1e3, "FINGER[1-2]": 1e1}, + max_effort={"Joint[0-9]": 1e5, "FINGER[1-2]": 1e3}, drive_type="force", ), init_qpos=[0.0, -np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi / 2, 0.0, 0.0, 0.0], @@ -119,16 +137,6 @@ def main(): if not args.headless: native_window_opened = sim.open_window() - gizmo_cfg = GizmoCfg( - ik_root_link_name="base_link", - ik_end_link_name="ee_link", - ik_tcp_pose=[ - [0.0, 1.0, 0.0, 0.0], - [-1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.12], - [0.0, 0.0, 0.0, 1.0], - ], - ) native_control = None if native_window_opened: native_control = create_robot_ik_gizmo_controller( diff --git a/pyproject.toml b/pyproject.toml index faece9731..299eaf49c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ dynamic = ["version"] # Core install dependencies (kept from requirements.txt). dependencies = [ - "dexsim_engine==0.4.3", + "dexsim_engine==0.5.0", "setuptools>=78.1.1", "gymnasium>=0.29.1", "langchain", diff --git a/tests/sim/objects/test_gizmo.py b/tests/sim/objects/test_gizmo.py index 2cf91a13a..39b1f93bd 100644 --- a/tests/sim/objects/test_gizmo.py +++ b/tests/sim/objects/test_gizmo.py @@ -17,6 +17,7 @@ from __future__ import annotations from types import SimpleNamespace +from pathlib import Path import numpy as np import pytest @@ -288,3 +289,232 @@ def _inject_solver(self) -> None: write["qpos"], torch.tensor([[0.4, -0.2]], dtype=torch.float32), ) + + +class _ConfiguredRobot(_FakeAdapterRobot): + def __init__(self) -> None: + super().__init__() + self.solver = SimpleNamespace( + joint_names=["joint_b", "joint_a"], + device=torch.device("cpu"), + root_link_name="base_link", + end_link_name="tool_link", + get_tcp=lambda: np.eye(4, dtype=np.float32), + ) + self.cfg.solver_cfg = {"arm": object()} + + def get_solver(self, name: str): + assert name == "arm" + return self.solver + + +@pytest.mark.parametrize("success,finite", [(True, True), (False, True), (True, False)]) +def test_configured_solver_preserves_joint_order_tcp_and_failed_seed( + success, finite +) -> None: + robot = _ConfiguredRobot() + calls = [] + + def solve(target, qpos_seed, return_all_solutions): + calls.append((target, qpos_seed)) + assert return_all_solutions is False + return torch.tensor([success]), torch.tensor( + [[[0.9, 0.7 if finite else float("nan")]]] + ) + + robot.solver.get_ik = solve + tcp = np.eye(4, dtype=np.float32) + tcp[0, 3] = 0.2 + adapter, solver, _, _ = gizmo_module._build_robot_ik( + robot, "arm", GizmoCfg(ik_solver="embodichain", ik_tcp_pose=tcp) + ) + solver.set_target_pose(np.array([0.5, 0.2, 0.1]), np.array([0.0, 0.0, 0.0, 1.0])) + assert solver.target_state_changed() + names, seed = adapter.get_actived_joint_names(), adapter.get_current_qpos() + solver.solve(names, seed) + torch.testing.assert_close(calls[0][1], torch.tensor([[0.3, 0.1]])) + torch.testing.assert_close(calls[0][0][0, :3, 3], torch.tensor([0.3, 0.2, 0.1])) + np.testing.assert_allclose( + solver.qpos_for_joint_names(names, seed), + [0.7, 0.9] if success and finite else seed, + ) + solver.reset_target_state_changed() + assert not solver.target_state_changed() + solver.target_state()["position"][0] += 0.1 + assert solver.target_state_changed() + np.testing.assert_array_equal(robot.solver.get_tcp(), np.eye(4)) + + +@pytest.mark.parametrize( + "configuration,match", + [ + ({"ik_solver": "unknown"}, "ik_solver must"), + ({"ik_solver": "embodichain"}, "configured EmbodiChain solver"), + ({"ik_iterations": 0}, "ik_iterations"), + ], +) +def test_robot_ik_rejects_invalid_configuration(configuration, match) -> None: + cfg = GizmoCfg( + ik_root_link_name="base_link", ik_end_link_name="tool_link", **configuration + ) + with pytest.raises(ValueError, match=match): + gizmo_module._build_robot_ik(_FakeAdapterRobot(), "arm", cfg) + + +def test_configured_solver_rejects_chain_and_joint_mismatch() -> None: + robot = _ConfiguredRobot() + with pytest.raises(ValueError, match="links must match"): + gizmo_module._build_robot_ik( + robot, "arm", GizmoCfg(ik_solver="embodichain", ik_root_link_name="other") + ) + robot.solver.joint_names = ["other_joint"] + with pytest.raises(ValueError, match="joints must match"): + gizmo_module._build_robot_ik(robot, "arm", GizmoCfg(ik_solver="embodichain")) + + +_PLANAR_URDF = """ + + + + + + + + + + + + + + + + + +""" + + +def _planar_pose(angle: float, x: float = 0.0) -> np.ndarray: + pose = np.eye(4, dtype=np.float32) + c, s = np.cos(angle), np.sin(angle) + pose[:2, :2] = [[c, -s], [s, c]] + pose[0, 3] = x + return pose + + +class _PlanarRobot(_ConfiguredRobot): + def __init__(self, urdf_path: Path, backend: str) -> None: + super().__init__() + urdf_path.write_text(_PLANAR_URDF) + self.cfg.fpath = str(urdf_path) + self.cfg.solver_cfg = None + # Move the runtime chain root away from both the robot root and the + # URDF rest transform, as happens when a torso or upstream joint moves. + self.root_pose = _planar_pose(0.8, 0.7) + self.root_pose[1:3, 3] = [-0.2, 0.6] + if backend == "embodichain": + pytest.importorskip("pink") + from embodichain.lab.sim.solvers import PinkSolverCfg + + self.solver = PinkSolverCfg( + urdf_path=str(urdf_path), + joint_names=["joint_a", "joint_b"], + root_link_name="base_link", + end_link_name="tool_link", + max_iterations=150, + pos_eps=1e-5, + rot_eps=1e-5, + show_ik_warnings=False, + ).init_solver(device=torch.device("cpu")) + self.cfg.solver_cfg = {"arm": self.solver.cfg} + + def tool_pose(self, qpos: np.ndarray) -> np.ndarray: + return ( + self.root_pose + @ _planar_pose(qpos[0]) + @ _planar_pose(qpos[1], 0.5) + @ _planar_pose(0.0, 0.5) + ) + + def get_link_pose(self, link_name, env_ids, to_matrix=False): + assert env_ids == [0] and to_matrix + pose = ( + self.root_pose + if link_name == "base_link" + else self.tool_pose(self.current_qpos[0, [0, 2]].numpy()) + ) + return torch.tensor(pose).unsqueeze(0) + + +@pytest.mark.parametrize( + "backend,ik_device", + [ + ("dexsim", "cpu"), + ("embodichain", "cpu"), + pytest.param("dexsim", "cuda:0", marks=pytest.mark.gpu), + ], +) +@pytest.mark.parametrize("frontend", ["native", "viser"]) +def test_real_ik_solvers_follow_live_chain_root_and_tcp( + tmp_path, monkeypatch, backend, ik_device, frontend +) -> None: + """Run real Newton/Pink IK; replace only native window rendering and input.""" + robot = _PlanarRobot(tmp_path / "robot.urdf", backend) + tcp = np.eye(4, dtype=np.float32) + tcp[0, 3] = 0.1 + cfg = GizmoCfg( + ik_solver=backend, + ik_root_link_name="base_link", + ik_end_link_name="tool_link", + ik_tcp_pose=tcp, + ik_iterations=100, + ik_device=ik_device, + ) + if frontend == "native": + import dexsim.kit.ik.interactive as interactive + import dexsim.engine + + monkeypatch.setattr( + interactive, + "setup_target_gizmo", + lambda *args, **kwargs: SimpleNamespace( + gizmo=SimpleNamespace( + set_visible=lambda _: None, follow=lambda _: None + ), + target_node=SimpleNamespace(set_world_pose=lambda _: None), + ), + ) + monkeypatch.setattr(dexsim.engine, "GizmoController", lambda: object()) + world = SimpleNamespace( + get_windows=lambda: SimpleNamespace(add_input_control=lambda _: None), + key_state=lambda _: False, + ) + controller, _ = create_robot_ik_gizmo_controller(robot, cfg=cfg, world=world) + assert isinstance(controller, interactive.IKGizmoController) + assert not controller.update() + adapter = controller.robot + else: + monkeypatch.setattr(gizmo_module, "Robot", _PlanarRobot) + controller = Gizmo(robot, cfg, "arm") + adapter = controller._robot_adapter + + # Move again after construction to verify live root following on update. + robot.root_pose = _planar_pose(-0.2, 0.3) @ robot.root_pose + target = robot.tool_pose(np.array([0.4, -0.2])) @ tcp + if frontend == "native": + local_target = gizmo_module.local_pose_from_world( + adapter.get_world_pose(), target + ) + controller.ik.set_target_pose( + local_target[:3, 3], + gizmo_module.rotation_matrix_to_quat_xyzw(local_target[:3, :3]), + ) + assert controller.update() + else: + assert controller.request_local_pose(target, source_id="viser:test") + controller.update() + write = robot.write_calls[-1] + assert write["target"] is True + assert write["joint_ids"] == [0, 2] + assert write["env_ids"] == [0] + actual = robot.tool_pose(write["qpos"][0].numpy()) @ tcp + np.testing.assert_allclose(actual, target, atol=1e-3) diff --git a/tests/visualization/test_viser_backend.py b/tests/visualization/test_viser_backend.py index ac6eba43e..ee8140b7a 100644 --- a/tests/visualization/test_viser_backend.py +++ b/tests/visualization/test_viser_backend.py @@ -17,8 +17,10 @@ from __future__ import annotations from types import SimpleNamespace +from dataclasses import replace import numpy as np +import pytest from embodichain.lab.visualization import ( CameraImage, @@ -982,6 +984,7 @@ def test_viser_backend_pick_enqueues_command_when_enabled() -> None: ray_direction=np.array([0.0, 0.0, -1.0], dtype=np.float32), ) ) + backend.poll() assert pick_commands == [] server.gui.checkboxes["Enable click-to-pick Gizmo"].callback( @@ -997,6 +1000,7 @@ def test_viser_backend_pick_enqueues_command_when_enabled() -> None: ) ) + backend.poll() assert len(pick_commands) == 1 command = pick_commands[0] assert command.node_id == "env:0/rigid:cube" @@ -1027,11 +1031,55 @@ def test_viser_backend_pick_miss_enqueues_empty_command() -> None: ) ) + backend.poll() assert len(pick_commands) == 1 assert pick_commands[0].node_id is None backend.stop() +@pytest.mark.parametrize("node_count", [1, 2]) +def test_picker_waits_for_matching_frame_after_topology_change(node_count: int) -> None: + backend, server, commands = _make_pick_backend() + manifest, frame = _make_pickable_scene() + backend.start() + try: + backend.publish_manifest(manifest) + assert backend.publish_frame(frame) + backend._pick_enabled = True + event = SimpleNamespace( + client_id="client-a", + ray_origin=np.array([0.0, 0.0, 5.0]), + ray_direction=np.array([0.0, 0.0, -1.0]), + ) + nodes = tuple( + replace(manifest.nodes[0], node_id=f"env:0/rigid:new{index}") + for index in range(node_count) + ) + backend.publish_manifest(replace(manifest, scene_revision=2, nodes=nodes)) + server.scene.pointer_callbacks[0][1](event) + backend.poll() + # Neither select a new node using old transforms nor detach an active + # gizmo as if this temporary absence of pose data were an empty click. + assert commands == [] + assert not backend.publish_frame(frame) + backend.publish_frame( + replace( + frame, + scene_revision=2, + node_ids=tuple(node.node_id for node in nodes), + positions=np.zeros((node_count, 3), dtype=np.float32), + wxyz=np.tile([1.0, 0.0, 0.0, 0.0], (node_count, 1)), + visible=np.ones(node_count, dtype=np.bool_), + ) + ) + server.scene.pointer_callbacks[0][1](event) + backend.poll() + assert commands[-1].node_id == nodes[0].node_id + assert commands[-1].scene_revision == 2 + finally: + backend.stop() + + def test_viser_backend_disabling_pick_clears_picker_gizmo() -> None: backend, server, pick_commands = _make_pick_backend() manifest, frame = _make_pickable_scene() @@ -1053,6 +1101,7 @@ def test_viser_backend_disabling_pick_clears_picker_gizmo() -> None: assert backend._pick_enabled is False # Disabling emits an empty pick so the simulation releases the gizmo. + backend.poll() assert len(pick_commands) == 1 assert pick_commands[0].node_id is None backend.stop() From 342ead2637c35b4d42a9a2059da3e95073ee9bf6 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 5 Sep 2026 16:43:38 +0800 Subject: [PATCH 06/14] feat(sim): enable native entity gizmos by default in windows --- agent_context/MAP.yaml | 1 + .../sim-visualization/sim-visualization.md | 9 +- .../simulation-system/simulation-system.md | 23 +++- .../embodichain.lab.sim.sim_manager.rst | 7 + docs/source/features/interaction/gizmo.md | 30 +++-- docs/source/features/interaction/window.md | 31 +++-- embodichain/lab/gym/utils/gym_utils.py | 1 + embodichain/lab/sim/sim_manager.py | 36 ++++++ examples/sim/gizmo/gizmo_object.py | 2 +- examples/sim/gizmo/gizmo_scene.py | 1 - tests/gym/utils/test_gym_utils.py | 22 ++++ tests/sim/test_sim_manager.py | 122 +++++++++++++++++- 12 files changed, 247 insertions(+), 38 deletions(-) diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 3a9fe925c..acde81349 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -23,6 +23,7 @@ topics: - SimulationManagerCfg - EntityGizmoManipulator - enable_entity_gizmo + - disable_entity_gizmo - DexSim - world - arena diff --git a/agent_context/topics/sim-visualization/sim-visualization.md b/agent_context/topics/sim-visualization/sim-visualization.md index 86292e613..c55c6fc53 100644 --- a/agent_context/topics/sim-visualization/sim-visualization.md +++ b/agent_context/topics/sim-visualization/sim-visualization.md @@ -252,8 +252,13 @@ Soft bodies and cloth require GPU physics. Their live vertices are sampled at Gizmo implementation lives in `embodichain/lab/sim/objects/gizmo.py`. Native windows delegate object picking/manipulation to DexSim's entity gizmo and robot -targets to its `IKGizmoController`. Viser transports poses to the simulation -thread. Both robot paths use Newton IK by default or, with +targets to its `IKGizmoController`. Entity interaction defaults on at the first +native window open; `SimulationManagerCfg.enable_entity_gizmo=False` or +`sim.disable_entity_gizmo()` opts out and reopening preserves the choice. +Pure headless and Viser runs do not automatically create native controls, and +robot IK targets require explicit creation. Viser's `allow_commands` remains +independent of the native startup preference. Viser transports poses to the +simulation thread. Both robot paths use Newton IK by default or, with `GizmoCfg(ik_solver="embodichain")`, reuse the configured control-part solver (including Pink). Chain roots follow the live robot link, including upstream joint motion; TCP overrides adapt targets without changing the shared solver. diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index 6a67718ed..d7d008c3a 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -112,17 +112,26 @@ do not belong to the simulation object; use ### Gizmo ownership -Native entity manipulation belongs to DexSim 0.5.0. Call -`SimulationManager.enable_entity_gizmo(config)` to obtain the world-owned -`EntityGizmoManipulator`; the manager only registers its default plane as a -static external target. Opening a window does not implicitly enable entity -gizmos. Query/disable through `sim.get_world().get_entity_gizmo()` and -`disable_entity_gizmo()`; DexSim owns window detach/reopen and controller state. +Native entity manipulation belongs to DexSim 0.5.0. The first successful native +window open enables its world-owned `EntityGizmoManipulator` by default, after +the scene and default plane are ready. The manager registers the default plane +as a static external target. `SimulationManagerCfg.enable_entity_gizmo=False` +opts out; Gym deployments accept the same top-level JSON/YAML field through +`gym.utils.gym_utils.config_to_cfg()`. + +`sim.enable_entity_gizmo(config)` explicitly enables/configures the controller; +`sim.disable_entity_gizmo()` also cancels pending automatic enablement before +the first window. These explicit calls take precedence over the startup default. +Query through `sim.get_world().get_entity_gizmo()`. DexSim owns window +detach/reopen and controller state; reopening never reapplies the default or +overwrites an explicit native disable. Pure headless and Viser runs do not +automatically create a native entity controller. `create_robot_ik_gizmo_controller()` in `objects/gizmo.py` returns the native DexSim IK controller and input controller. The caller retains both and calls the IK controller's `update()` per frame. `SimulationManager.enable_gizmo()` -creates Viser controls for robots, rigid objects, or cameras. Both robot paths +creates Viser controls for robots, rigid objects, or cameras. Robot IK controls +are always created explicitly for a control part. Both robot paths default to native Newton IK; `GizmoCfg(ik_solver="embodichain")` adapts the control part's existing solver, such as PinkSolver. Both support one environment and write only selected non-mimic joint drive targets through `Robot`. diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst index 088dcb39e..48c8ddaa0 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst @@ -13,6 +13,13 @@ configuration. Downstream components (environments, planners, IK solvers, the visualization runtime) look up the active manager through its class-level instance registry instead of passing it around explicitly. +Native entity interaction defaults on when the first native window opens. +Set ``SimulationManagerCfg(enable_entity_gizmo=False)`` to opt out, or call +``sim.disable_entity_gizmo()`` at runtime. Explicit enable/disable calls and +custom DexSim controller settings survive window close/reopen. Pure headless +and Viser runs do not automatically create native gizmos; robot IK controls +are created separately for a selected control part. + .. rubric:: Classes .. autosummary:: diff --git a/docs/source/features/interaction/gizmo.md b/docs/source/features/interaction/gizmo.md index 637fb4065..cf4fa9c3f 100644 --- a/docs/source/features/interaction/gizmo.md +++ b/docs/source/features/interaction/gizmo.md @@ -3,21 +3,29 @@ ```{currentmodule} embodichain.lab.sim ``` -A Gizmo is a registered transform control for manipulating a simulation target -from either the native DexSim window or a trusted Viser browser. The simulation -owns the authoritative state: UI callbacks enqueue requested poses, and -`SimulationManager` applies them on the simulation thread. +A Gizmo is a transform control for manipulating a simulation target. Native +windows use DexSim's entity and robot IK controllers. Viser uses registered +controls whose callbacks enqueue requested poses for `SimulationManager` to +apply on the simulation thread. + +Native entity interaction is enabled by default when the first native window +opens. Select an object and press **G** to attach its gizmo. Set +`SimulationManagerCfg(enable_entity_gizmo=False)` or call +`sim.disable_entity_gizmo()` to opt out; reopening preserves that choice. +Robot IK controls require explicit creation for a selected control part. +See {doc}`native window controls ` for native entity configuration. ## Supported Targets | Target | Gizmo behavior | |---|---| -| Robot | Moves the selected control part through its configured FK/IK solver. | +| Robot | Moves the selected control part through native DexSim IK by default, or its configured EmbodiChain solver. | | Rigid object | Sets the object's local arena pose. | | Camera | Sets the camera's local pose. | -Gizmo interaction currently requires `num_envs=1`. Robot targets also require -a valid `control_part` and solver configuration. +Registered Viser controls and robot IK controls currently require `num_envs=1`. +Robot targets require a valid `control_part`; an EmbodiChain solver configuration +is needed only when selecting `GizmoCfg(ik_solver="embodichain")`. ## Quick Start @@ -33,8 +41,7 @@ Use the same target through Viser: python scripts/tutorials/sim/gizmo_robot.py --viser ``` -Register controls through the manager rather than constructing or destroying -Gizmo instances directly: +For Viser, register controls through the manager: ```python sim.enable_gizmo( @@ -60,8 +67,9 @@ Use `disable_gizmo()`, `set_gizmo_visibility()`, and ## Frontend Behavior -- The native window uses the DexSim Gizmo controller and requires an open - window. +- Native entity interaction starts with the first native window; pure headless + and Viser runs do not automatically create native controllers. Native robot + targets use `create_robot_ik_gizmo_controller()` explicitly. - Viser exports browser-native transform controls when `VisualizationCfg.allow_commands=True`; otherwise it shows read-only frames. - The standard `--viser` launcher enables registered commands for trusted diff --git a/docs/source/features/interaction/window.md b/docs/source/features/interaction/window.md index b999be81e..c45db37fe 100644 --- a/docs/source/features/interaction/window.md +++ b/docs/source/features/interaction/window.md @@ -55,8 +55,19 @@ The camera-pose hotkey is controlled by `SimulationManagerCfg.window_camera_pose ### Entity Gizmo Control -DexSim owns native entity selection and manipulation. Enable it explicitly after -opening a native window: +DexSim owns native entity selection and manipulation. EmbodiChain enables it +automatically when the first native window opens, including a window created +with `SimulationManagerCfg(headless=False)`. Pure headless and Viser runs do not +automatically create native entity gizmos. + +To start with native interaction disabled, use +`SimulationManagerCfg(enable_entity_gizmo=False)`. Gym JSON/YAML deployments +accept the top-level field `enable_entity_gizmo: false` as well. At runtime, +`sim.disable_entity_gizmo()` disables interaction and cancels automatic +enablement even before the first window opens. Closing and reopening a window +preserves the controller's current enabled state and custom configuration. + +For custom DexSim settings, enable or reconfigure the controller explicitly: ```python import dexsim @@ -78,18 +89,18 @@ cannot receive an entity gizmo. Other supported scene entities remain selectable normally. `sim.enable_entity_gizmo(config)` is a thin helper that also excludes -EmbodiChain's render-only default plane. All other lifecycle operations stay on -DexSim's world object: +EmbodiChain's render-only default plane. Query the controller through DexSim's +world object; use the manager's disable helper to preserve your preference +across future window opens: ```python -world = sim.get_world() -controller = world.get_entity_gizmo() -world.disable_entity_gizmo() +controller = sim.get_world().get_entity_gizmo() +sim.disable_entity_gizmo() ``` -This controller is distinct from DexSim's target-specific Robot TCP IK -controller. When both are active, **G** controls entity roots and **I** shows or -hides the Robot TCP target. +Robot TCP IK controllers still require explicit creation for a robot control +part. When both controllers are active, **G** controls entity roots and **I** +shows or hides the Robot TCP target. The entity gizmo is native-window only. The Viser backend offers an analogous **click-to-pick** flow (an *Enable click-to-pick Gizmo* checkbox instead of the diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 34ecc7574..775491d2a 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -620,6 +620,7 @@ class ComponentCfg: env_cfg.sim_cfg = SimulationManagerCfg( headless=config.get("headless", False), + enable_entity_gizmo=config.get("enable_entity_gizmo", True), sim_device=config.get("device", "cpu"), render_cfg=RenderCfg(**render_config), gpu_id=config.get("gpu_id", 0), diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index bb5c967ad..721770d10 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -134,6 +134,15 @@ class SimulationManagerCfg: require a native window. """ + enable_entity_gizmo: bool = True + """Enable native object interaction when the first native window opens. + + Headless and Viser simulations do not automatically create native gizmos. + Explicit runtime enable/disable calls take precedence over this default; + reopening a window preserves the current DexSim controller state. + Robot IK gizmos are created separately for a selected control part. + """ + render_cfg: RenderCfg = field(default_factory=RenderCfg) """The rendering configuration parameters.""" @@ -295,6 +304,7 @@ def __init__( self._world: dexsim.World = dexsim.World(world_config) self._window: Windows | None = None + self._auto_entity_gizmo_pending = sim_config.enable_entity_gizmo self._window_record_state: _WindowRecordState | None = None self._window_record_camera: object | None = None wr = sim_config.window_record @@ -384,6 +394,8 @@ def __init__( if sim_config.headless is False: self._window = self._world.get_windows() + self.is_window_opened = self._window is not None + self._enable_default_entity_gizmo() self._is_constructed = True @@ -993,6 +1005,9 @@ def open_window(self) -> bool: return True self._world.open_window() self._window = self._world.get_windows() + if self._window is None: + return False + self._enable_default_entity_gizmo() if ( self._window_record_hotkey_cfg is not None @@ -2006,6 +2021,11 @@ def enable_entity_gizmo( ) -> EntityGizmoManipulator: """Enable DexSim entity control and exclude the EmbodiChain ground. + Called automatically on the first native window unless disabled in + :class:`SimulationManagerCfg`. Explicit calls also work headlessly. + DexSim retains this controller and its configuration across window + close/reopen; reopening does not reapply the default configuration. + Args: config: Native DexSim entity-Gizmo configuration. @@ -2017,6 +2037,7 @@ def enable_entity_gizmo( if config is None else self._world.enable_entity_gizmo(config) ) + self._auto_entity_gizmo_pending = False default_plane = getattr(self, "_default_plane", None) if default_plane is None: return controller @@ -2033,6 +2054,21 @@ def enable_entity_gizmo( ) return controller + def disable_entity_gizmo(self) -> None: + """Disable native object interaction, including future window opens. + + Call :meth:`enable_entity_gizmo` to explicitly re-enable interaction. + Robot IK controllers and Viser command permissions are independent. + """ + self._world.disable_entity_gizmo() + self._auto_entity_gizmo_pending = False + + def _enable_default_entity_gizmo(self) -> None: + # Apply the startup default once; DexSim owns subsequent controller + # state, including explicit native disable calls and window reopens. + if self._window is not None and self._auto_entity_gizmo_pending: + self.enable_entity_gizmo() + def enable_gizmo( self, uid: str, diff --git a/examples/sim/gizmo/gizmo_object.py b/examples/sim/gizmo/gizmo_object.py index b01175738..cff7a6492 100644 --- a/examples/sim/gizmo/gizmo_object.py +++ b/examples/sim/gizmo/gizmo_object.py @@ -145,7 +145,7 @@ def run_simulation(sim: SimulationManager): if step_count == 200000 and gizmo_enabled: logger.log_info("Disabling Gizmo control at step 200000") if sim.get_world().get_entity_gizmo() is not None: - sim.get_world().disable_entity_gizmo() + sim.disable_entity_gizmo() else: sim.disable_gizmo("cube1") sim.disable_gizmo("cube2") diff --git a/examples/sim/gizmo/gizmo_scene.py b/examples/sim/gizmo/gizmo_scene.py index bd63963ee..8c3844f41 100644 --- a/examples/sim/gizmo/gizmo_scene.py +++ b/examples/sim/gizmo/gizmo_scene.py @@ -165,7 +165,6 @@ def main(): native_controls = [] if native_window_opened: - sim.enable_entity_gizmo() for control_part in ("left_arm", "right_arm"): native_controls.append( create_robot_ik_gizmo_controller( diff --git a/tests/gym/utils/test_gym_utils.py b/tests/gym/utils/test_gym_utils.py index a59bf74bd..1c00b41df 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -1385,6 +1385,28 @@ def test_task_integration_rejects_removed_version_field( source_path=tmp_path / "env.yaml", ) + @pytest.mark.parametrize("suffix", ["yaml", "json"]) + @pytest.mark.parametrize("enabled", [None, False, True]) + def test_gym_config_preserves_entity_gizmo_startup_preference( + self, tmp_path: Path, suffix: str, enabled: bool | None + ) -> None: + """Task deployments default to native interaction and can opt out.""" + config = { + "id": "EmbodiedEnv-v1", + "env": {}, + "robot": {"uid": "TestRobot"}, + } + if enabled is not None: + config["enable_entity_gizmo"] = enabled + config_path = tmp_path / f"gym_config.{suffix}" + save_config(config_path, config) + + cfg = config_to_cfg( + load_config(config_path), manager_modules=DEFAULT_MANAGER_MODULES + ) + + assert cfg.sim_cfg.enable_entity_gizmo is (enabled is not False) + def test_yaml_gym_config_parses_to_cfg(self, tmp_path): config = { "id": "EmbodiedEnv-v1", diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index 4b83cff24..8ac236b1b 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -149,9 +149,20 @@ def update(self, physics_dt: float) -> None: def enable_entity_gizmo(self, config: object | None = None) -> object: self.entity_gizmo_configs.append(config) - self.entity_gizmo = FakeEntityGizmo() + if self.entity_gizmo is None: + self.entity_gizmo = FakeEntityGizmo() + self.entity_gizmo.active = True return self.entity_gizmo + def disable_entity_gizmo(self) -> None: + if self.entity_gizmo is not None: + self.entity_gizmo.active = False + + def get_entity_gizmo(self) -> object | None: + if self.entity_gizmo is not None and self.entity_gizmo.active: + return self.entity_gizmo + return None + def open_window(self) -> None: self.window_open_count += 1 self.window_closed = False @@ -229,16 +240,20 @@ def end_interaction(self, source_id: str) -> bool: return True -def _make_sim_manager(window: object | None = None) -> SimulationManager: +def _make_sim_manager( + window: object | None = None, *, enable_entity_gizmo: bool = True +) -> SimulationManager: """Create a minimally initialized simulation manager for recorder tests.""" sim = object.__new__(SimulationManager) sim.instance_id = 0 sim.sim_config = SimpleNamespace( width=64, height=48, + enable_entity_gizmo=enable_entity_gizmo, visualization=SimpleNamespace(backend="none"), ) sim._window = window + sim._auto_entity_gizmo_pending = enable_entity_gizmo sim._window_record_state = None sim._window_record_camera = None sim._window_record_save_threads = [] @@ -703,6 +718,7 @@ def test_open_window_allows_native_backend() -> None: sim._window_record_input_control = None sim._window_camera_pose_hotkey_cfg = None sim._window_camera_pose_input_control = None + sim._auto_entity_gizmo_pending = True sim.is_window_opened = False opened = sim.open_window() @@ -745,16 +761,82 @@ def test_entity_gizmo_delegates_to_dexsim_and_excludes_default_plane() -> None: ] -def test_open_window_does_not_enable_entity_gizmo_implicitly() -> None: +def test_open_window_enables_entity_gizmo_by_default() -> None: sim = _make_sim_manager() assert sim.open_window() assert sim.is_window_opened is True assert sim._world.window_open_count == 1 + assert sim._world.entity_gizmo_configs == [None] + assert sim._world.get_entity_gizmo().external_targets[0][2] is sim._default_plane + + +def test_entity_gizmo_can_be_disabled_in_startup_configuration() -> None: + cfg = SimulationManagerCfg() + assert cfg.enable_entity_gizmo is True + cfg = SimulationManagerCfg(enable_entity_gizmo=False) + sim = _make_sim_manager(enable_entity_gizmo=cfg.enable_entity_gizmo) + + assert sim.open_window() + sim.close_window() + assert sim.open_window() + assert sim._world.get_entity_gizmo() is None assert sim._world.entity_gizmo_configs == [] +@pytest.mark.parametrize("before_first_window", [True, False]) +def test_explicit_entity_gizmo_disable_survives_window_reopen( + before_first_window: bool, +) -> None: + sim = _make_sim_manager() + if not before_first_window: + assert sim.open_window() + sim.disable_entity_gizmo() + sim.close_window() + assert sim.open_window() + assert sim._world.get_entity_gizmo() is None + assert len(sim._world.entity_gizmo_configs) == (0 if before_first_window else 1) + + controller = sim.enable_entity_gizmo() + sim.close_window() + assert sim.open_window() + assert sim._world.get_entity_gizmo() is controller + + +def test_native_entity_gizmo_disable_is_not_overridden_on_reopen() -> None: + sim = _make_sim_manager() + assert sim.open_window() + sim._world.disable_entity_gizmo() + sim.close_window() + assert sim.open_window() + assert sim._world.get_entity_gizmo() is None + assert sim._world.entity_gizmo_configs == [None] + + +def test_window_reopen_preserves_explicit_entity_gizmo_configuration() -> None: + sim = _make_sim_manager() + config = object() + controller = sim.enable_entity_gizmo(config) + assert sim.open_window() + sim.close_window() + assert sim.open_window() + assert sim._world.get_entity_gizmo() is controller + assert sim._world.entity_gizmo_configs == [config] + + +def test_failed_window_open_does_not_consume_gizmo_startup_default() -> None: + sim = _make_sim_manager() + window = sim._world.window + sim._world.window = None + assert not sim.open_window() + assert not sim.is_window_opened + assert sim._world.entity_gizmo_configs == [] + sim._world.window = window + assert sim.open_window() + assert sim._world.entity_gizmo_configs == [None] + + def test_close_window_leaves_entity_gizmo_lifecycle_to_dexsim() -> None: sim = _make_sim_manager(window=object()) controller = sim.enable_entity_gizmo() @@ -778,7 +860,22 @@ def test_start_visualization_rejects_open_native_window() -> None: sim.start_visualization() -def test_constructor_starts_visualization_after_default_scene(monkeypatch) -> None: +@pytest.mark.parametrize( + "headless,entity_gizmo,backend,expected_gizmo", + [ + (False, True, "none", True), + (False, False, "none", False), + (True, True, "none", False), + (False, True, "viser", False), + ], +) +def test_constructor_starts_visualization_after_default_scene( + monkeypatch: pytest.MonkeyPatch, + headless: bool, + entity_gizmo: bool, + backend: str, + expected_gizmo: bool, +) -> None: lifecycle: list[str] = [] world = MagicMock() world.get_physics_scene.return_value = MagicMock() @@ -837,9 +934,22 @@ def start_visualization(sim: SimulationManager) -> None: "start_visualization", start_visualization, ) + monkeypatch.setattr( + SimulationManager, + "enable_entity_gizmo", + lambda _self: lifecycle.append("entity_gizmo"), + ) sim = object.__new__(SimulationManager) - SimulationManager.__init__(sim, SimulationManagerCfg(num_envs=3)) + SimulationManager.__init__( + sim, + SimulationManagerCfg( + num_envs=3, + headless=headless, + enable_entity_gizmo=entity_gizmo, + visualization=VisualizationCfg(backend=backend), + ), + ) assert lifecycle == [ "resources", @@ -848,7 +958,7 @@ def start_visualization(sim: SimulationManager) -> None: "lighting", "arenas", "visualization:3", - ] + ] + (["entity_gizmo"] if expected_gizmo else []) def test_remove_asset_marks_visualization_topology_dirty() -> None: From eced864c3e79dd3c67fb10f5c67804853e77a728 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 5 Sep 2026 17:47:27 +0800 Subject: [PATCH 07/14] fix(tutorial): advance robot gizmos with manual physics --- .../simulation-system/simulation-system.md | 7 +++ scripts/tutorials/sim/gizmo_robot.py | 32 +++++++--- .../test_gizmo_robot_tutorial.py | 62 +++++++++++++++++++ 3 files changed, 92 insertions(+), 9 deletions(-) create mode 100644 tests/visualization/test_gizmo_robot_tutorial.py diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index d7d008c3a..845e16947 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -70,6 +70,13 @@ then advances the world for the requested number of physics steps. Each environment control step normally calls it with `sim_steps_per_control`. +`scripts/tutorials/sim/gizmo_robot.py` defaults to manual physics. It initializes +GPU physics after robot creation when needed, sets both current and target +joint positions, and advances once before creating the IK target. In its loop, +the native IK controller updates before `sim.update(step=1)`, which also handles +Viser commands/capture. The loop is paced by `physics_dt`; automatic physics +instead polls gizmos and visualization at 30 Hz without advancing the world. + `ArticulationCfg.enable_gravity` defaults to `True`. During articulation construction, `Articulation` applies this explicit runtime flag to every native entity before the first physics update, including when diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index c14d9dd3b..fa054d7a1 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -"""Control a UR10 end effector with a native DexSim or Viser Gizmo.""" +"""Control a UR10 end effector with a Gizmo and manual physics stepping.""" from __future__ import annotations @@ -63,7 +63,7 @@ def main(): ) sim = SimulationManager(sim_cfg) - sim.set_manual_update(False) + sim.set_manual_update(True) # Get UR10 URDF path urdf_path = get_data_path("UniversalRobots/UR10/UR10.urdf") @@ -92,17 +92,21 @@ def main(): ), ) robot = sim.add_robot(cfg=robot_cfg) + if sim.is_use_gpu_physics: + sim.init_gpu_physics() # Set initial joint positions initial_qpos = torch.tensor( [[0, -np.pi / 2, np.pi / 2, 0.0, np.pi / 2, 0.0]], dtype=torch.float32, - device="cpu", + device=sim.device, ) joint_ids = robot.get_joint_ids("arm") + robot.set_qpos(qpos=initial_qpos, joint_ids=joint_ids, target=False) robot.set_qpos(qpos=initial_qpos, joint_ids=joint_ids) - time.sleep(0.2) # Wait for a moment to ensure everything is set up + if sim.is_physics_manually_update: + sim.update(step=1) # Refresh link poses before creating the IK target. native_window_opened = False if not args.headless: @@ -145,20 +149,27 @@ def main(): def run_simulation(sim: SimulationManager, native_control=None): + """Update IK before each manual physics step, or poll automatic physics.""" step_count = 0 try: - last_time = time.time() + last_time = time.perf_counter() last_step = 0 while True: - time.sleep(0.033) # 30Hz + frame_start = time.perf_counter() if native_control is not None: native_control[0].update() - sim.update_gizmos() - sim.capture_visualization_safely() + if sim.is_physics_manually_update: + # update() also processes Viser commands and publishes the frame. + sim.update(step=1) + frame_dt = sim.sim_config.physics_dt + else: + sim.update_gizmos() + sim.capture_visualization_safely() + frame_dt = 1.0 / 30.0 step_count += 1 if step_count % 100 == 0: - current_time = time.time() + current_time = time.perf_counter() elapsed = current_time - last_time fps = ( sim.num_envs * (step_count - last_step) / elapsed @@ -168,6 +179,9 @@ def run_simulation(sim: SimulationManager, native_control=None): logger.log_info(f"Simulation step: {step_count}, FPS: {fps:.2f}") last_time = current_time last_step = step_count + + elapsed = time.perf_counter() - frame_start + time.sleep(max(0.0, frame_dt - elapsed)) except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: diff --git a/tests/visualization/test_gizmo_robot_tutorial.py b/tests/visualization/test_gizmo_robot_tutorial.py new file mode 100644 index 000000000..31f8e2031 --- /dev/null +++ b/tests/visualization/test_gizmo_robot_tutorial.py @@ -0,0 +1,62 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock, call + +import pytest + +from scripts.tutorials.sim import gizmo_robot + +pytestmark = pytest.mark.no_sim + + +@pytest.mark.parametrize("manual", [True, False]) +@pytest.mark.parametrize("native", [True, False]) +@pytest.mark.parametrize("work_duration", [0.002, 0.050]) +def test_gizmo_loop_advances_manual_physics_and_paces_frames( + monkeypatch: pytest.MonkeyPatch, + manual: bool, + native: bool, + work_duration: float, +) -> None: + """Apply IK before stepping, avoid duplicate updates, and clean up on exit.""" + physics_dt = 0.01 + calls = Mock() + sim = calls.sim + sim.is_physics_manually_update = manual + sim.sim_config = SimpleNamespace(physics_dt=physics_dt) + native_control = (calls.ik, object()) if native else None + clock = Mock(side_effect=[0.0, 0.0, work_duration]) + sleep = Mock(side_effect=KeyboardInterrupt) + monkeypatch.setattr(gizmo_robot.time, "perf_counter", clock) + monkeypatch.setattr(gizmo_robot.time, "sleep", sleep) + + gizmo_robot.run_simulation(sim, native_control) + + expected_calls = [call.ik.update()] if native else [] + if manual: + expected_calls.append(call.sim.update(step=1)) + else: + expected_calls.extend( + [call.sim.update_gizmos(), call.sim.capture_visualization_safely()] + ) + expected_calls.append(call.sim.destroy()) + assert calls.mock_calls == expected_calls + frame_dt = physics_dt if manual else 1.0 / 30.0 + sleep.assert_called_once_with(pytest.approx(max(0.0, frame_dt - work_duration))) From 99361c18e86f0e6f77478566ebdbd8c756a11093 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 5 Sep 2026 17:49:11 +0800 Subject: [PATCH 08/14] refactor(tutorial): keep robot gizmos in manual physics mode --- .../simulation-system/simulation-system.md | 6 +++--- scripts/tutorials/sim/gizmo_robot.py | 17 +++++------------ .../visualization/test_gizmo_robot_tutorial.py | 14 ++------------ 3 files changed, 10 insertions(+), 27 deletions(-) diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index 845e16947..88c9dfcf4 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -70,12 +70,12 @@ then advances the world for the requested number of physics steps. Each environment control step normally calls it with `sim_steps_per_control`. -`scripts/tutorials/sim/gizmo_robot.py` defaults to manual physics. It initializes +`scripts/tutorials/sim/gizmo_robot.py` supports only manual physics. It initializes GPU physics after robot creation when needed, sets both current and target joint positions, and advances once before creating the IK target. In its loop, the native IK controller updates before `sim.update(step=1)`, which also handles -Viser commands/capture. The loop is paced by `physics_dt`; automatic physics -instead polls gizmos and visualization at 30 Hz without advancing the world. +Viser commands/capture. The loop is paced by `physics_dt` and has no automatic +physics polling path. `ArticulationCfg.enable_gravity` defaults to `True`. During articulation construction, `Articulation` applies this explicit runtime flag to every native diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index fa054d7a1..6b830ac55 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -105,8 +105,7 @@ def main(): robot.set_qpos(qpos=initial_qpos, joint_ids=joint_ids, target=False) robot.set_qpos(qpos=initial_qpos, joint_ids=joint_ids) - if sim.is_physics_manually_update: - sim.update(step=1) # Refresh link poses before creating the IK target. + sim.update(step=1) # Refresh link poses before creating the IK target. native_window_opened = False if not args.headless: @@ -149,7 +148,7 @@ def main(): def run_simulation(sim: SimulationManager, native_control=None): - """Update IK before each manual physics step, or poll automatic physics.""" + """Update IK and advance one manual physics step per frame.""" step_count = 0 try: last_time = time.perf_counter() @@ -158,14 +157,8 @@ def run_simulation(sim: SimulationManager, native_control=None): frame_start = time.perf_counter() if native_control is not None: native_control[0].update() - if sim.is_physics_manually_update: - # update() also processes Viser commands and publishes the frame. - sim.update(step=1) - frame_dt = sim.sim_config.physics_dt - else: - sim.update_gizmos() - sim.capture_visualization_safely() - frame_dt = 1.0 / 30.0 + # update() also processes Viser commands and publishes the frame. + sim.update(step=1) step_count += 1 if step_count % 100 == 0: @@ -181,7 +174,7 @@ def run_simulation(sim: SimulationManager, native_control=None): last_step = step_count elapsed = time.perf_counter() - frame_start - time.sleep(max(0.0, frame_dt - elapsed)) + time.sleep(max(0.0, sim.sim_config.physics_dt - elapsed)) except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: diff --git a/tests/visualization/test_gizmo_robot_tutorial.py b/tests/visualization/test_gizmo_robot_tutorial.py index 31f8e2031..d6c661beb 100644 --- a/tests/visualization/test_gizmo_robot_tutorial.py +++ b/tests/visualization/test_gizmo_robot_tutorial.py @@ -26,12 +26,10 @@ pytestmark = pytest.mark.no_sim -@pytest.mark.parametrize("manual", [True, False]) @pytest.mark.parametrize("native", [True, False]) @pytest.mark.parametrize("work_duration", [0.002, 0.050]) def test_gizmo_loop_advances_manual_physics_and_paces_frames( monkeypatch: pytest.MonkeyPatch, - manual: bool, native: bool, work_duration: float, ) -> None: @@ -39,7 +37,6 @@ def test_gizmo_loop_advances_manual_physics_and_paces_frames( physics_dt = 0.01 calls = Mock() sim = calls.sim - sim.is_physics_manually_update = manual sim.sim_config = SimpleNamespace(physics_dt=physics_dt) native_control = (calls.ik, object()) if native else None clock = Mock(side_effect=[0.0, 0.0, work_duration]) @@ -50,13 +47,6 @@ def test_gizmo_loop_advances_manual_physics_and_paces_frames( gizmo_robot.run_simulation(sim, native_control) expected_calls = [call.ik.update()] if native else [] - if manual: - expected_calls.append(call.sim.update(step=1)) - else: - expected_calls.extend( - [call.sim.update_gizmos(), call.sim.capture_visualization_safely()] - ) - expected_calls.append(call.sim.destroy()) + expected_calls.extend([call.sim.update(step=1), call.sim.destroy()]) assert calls.mock_calls == expected_calls - frame_dt = physics_dt if manual else 1.0 / 30.0 - sleep.assert_called_once_with(pytest.approx(max(0.0, frame_dt - work_duration))) + sleep.assert_called_once_with(pytest.approx(max(0.0, physics_dt - work_duration))) From 8a13259f7dd05249412778b9a307797024badf39 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 5 Sep 2026 18:58:14 +0800 Subject: [PATCH 09/14] feat(sim): manage robot IK gizmos automatically --- agent_context/MAP.yaml | 1 + .../sim-visualization/sim-visualization.md | 11 +- .../simulation-system/simulation-system.md | 26 +- .../embodichain.lab.sim.objects.rst | 6 + .../embodichain.lab.sim.sim_manager.rst | 12 +- docs/source/features/interaction/gizmo.md | 52 +-- docs/source/features/interaction/window.md | 7 +- docs/source/tutorial/gizmo.rst | 162 ++++------ embodichain/lab/gym/utils/gym_utils.py | 1 + embodichain/lab/sim/objects/gizmo.py | 115 +++++-- embodichain/lab/sim/sim_manager.py | 112 ++++++- scripts/tutorials/sim/gizmo_robot.py | 39 +-- tests/gym/utils/test_gym_utils.py | 24 ++ tests/sim/objects/test_gizmo.py | 7 +- tests/sim/test_robot_gizmo_lifecycle.py | 304 ++++++++++++++++++ .../test_gizmo_robot_tutorial.py | 8 +- 16 files changed, 680 insertions(+), 207 deletions(-) create mode 100644 tests/sim/test_robot_gizmo_lifecycle.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index acde81349..43d19b42b 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -24,6 +24,7 @@ topics: - EntityGizmoManipulator - enable_entity_gizmo - disable_entity_gizmo + - robot_ik_gizmo - DexSim - world - arena diff --git a/agent_context/topics/sim-visualization/sim-visualization.md b/agent_context/topics/sim-visualization/sim-visualization.md index c55c6fc53..2a080a359 100644 --- a/agent_context/topics/sim-visualization/sim-visualization.md +++ b/agent_context/topics/sim-visualization/sim-visualization.md @@ -255,9 +255,14 @@ windows delegate object picking/manipulation to DexSim's entity gizmo and robot targets to its `IKGizmoController`. Entity interaction defaults on at the first native window open; `SimulationManagerCfg.enable_entity_gizmo=False` or `sim.disable_entity_gizmo()` opts out and reopening preserves the choice. -Pure headless and Viser runs do not automatically create native controls, and -robot IK targets require explicit creation. Viser's `allow_commands` remains -independent of the native startup preference. Viser transports poses to the +Pure headless and Viser runs do not automatically create native controls. +`SimulationManagerCfg.robot_ik_gizmo` automatically registers robot parts with +solver chain/TCP metadata: native IK activates on I and Viser builds IK on the +first drag. The manager updates both and preserves native controls across window +reopen. Explicit configuration and disable calls take precedence; caller-owned +native factory controllers are not duplicated. Automatic Viser registration +requires `allow_commands`, independently of native interaction preferences. +Viser transports poses to the simulation thread. Both robot paths use Newton IK by default or, with `GizmoCfg(ik_solver="embodichain")`, reuse the configured control-part solver (including Pink). Chain roots follow the live robot link, including upstream diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index 88c9dfcf4..156996f3e 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -72,9 +72,9 @@ environment control step normally calls it with `scripts/tutorials/sim/gizmo_robot.py` supports only manual physics. It initializes GPU physics after robot creation when needed, sets both current and target -joint positions, and advances once before creating the IK target. In its loop, -the native IK controller updates before `sim.update(step=1)`, which also handles -Viser commands/capture. The loop is paced by `physics_dt` and has no automatic +joint positions, and advances once before opening the window. Its loop only +calls `sim.update(step=1)`; the manager owns native IK updates and Viser +commands/capture. The loop is paced by `physics_dt` and has no automatic physics polling path. `ArticulationCfg.enable_gravity` defaults to `True`. During articulation @@ -134,11 +134,21 @@ detach/reopen and controller state; reopening never reapplies the default or overwrites an explicit native disable. Pure headless and Viser runs do not automatically create a native entity controller. -`create_robot_ik_gizmo_controller()` in `objects/gizmo.py` returns the native -DexSim IK controller and input controller. The caller retains both and calls -the IK controller's `update()` per frame. `SimulationManager.enable_gizmo()` -creates Viser controls for robots, rigid objects, or cameras. Robot IK controls -are always created explicitly for a control part. Both robot paths +`SimulationManagerCfg.robot_ik_gizmo` defaults to `GizmoCfg()`. During normal +updates the manager registers robot control parts with complete solver chain/TCP +metadata in single-environment interactive runs. Pure headless, read-only Viser, and +multi-environment runs do not register automatic controls. The first native I +press creates DexSim's `IKGizmoController`; Viser constructs IK on its first drag. +Registration never writes drive targets. `Gizmo` owns managed native input and +target-node cleanup, detaches input on window close, and reattaches the same +controller on reopen. Robot removal releases all its managed controls. + +Set `robot_ik_gizmo=None` to opt out or supply `GizmoCfg` overrides; Gym +JSON/YAML accepts the same mapping/null. `enable_gizmo()` can override one part, +and `disable_gizmo()` prevents automatic recreation (all parts when omitted). +The explicit `create_robot_ik_gizmo_controller()` factory still returns +caller-owned controllers; a weak registry prevents automatic duplicates. +Both robot paths default to native Newton IK; `GizmoCfg(ik_solver="embodichain")` adapts the control part's existing solver, such as PinkSolver. Both support one environment and write only selected non-mimic joint drive targets through `Robot`. diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst index 7c286857a..af8903319 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst @@ -194,6 +194,12 @@ Gizmo .. autofunction:: create_robot_ik_gizmo_controller +SimulationManager automatically discovers robot control parts with configured +IK chain metadata. Native controllers activate on the first I press, while +Viser constructs its solver on the first drag. ``sim.update()`` owns updates +and cleanup; ordinary applications do not need the explicit factory. +Use ``SimulationManagerCfg(robot_ik_gizmo=None)`` to disable automatic setup. + The native controller defaults to DexSim Newton IK. With a ``PinkSolverCfg`` (or another EmbodiChain solver) configured for the robot's control part, pass ``GizmoCfg(ik_solver="embodichain")`` to select that solver for either a native diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst index 48c8ddaa0..1af4b6a43 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst @@ -17,8 +17,16 @@ Native entity interaction defaults on when the first native window opens. Set ``SimulationManagerCfg(enable_entity_gizmo=False)`` to opt out, or call ``sim.disable_entity_gizmo()`` at runtime. Explicit enable/disable calls and custom DexSim controller settings survive window close/reopen. Pure headless -and Viser runs do not automatically create native gizmos; robot IK controls -are created separately for a selected control part. +and Viser runs do not automatically create native gizmos. + +``SimulationManagerCfg.robot_ik_gizmo`` defaults to ``GizmoCfg()`` and registers +robot control parts with configured IK-chain/TCP metadata during normal updates. +Native IK activates on the first **I** press; Viser constructs its solver on +the first drag and requires ``visualization.allow_commands``. Registration does +not write drive targets. Set this field to ``None`` to opt out or select +``GizmoCfg(ik_solver="embodichain")`` to reuse configured solvers. Explicit +``enable_gizmo()`` settings override automatic defaults, and ``disable_gizmo()`` +prevents automatic recreation. .. rubric:: Classes diff --git a/docs/source/features/interaction/gizmo.md b/docs/source/features/interaction/gizmo.md index cf4fa9c3f..42856e155 100644 --- a/docs/source/features/interaction/gizmo.md +++ b/docs/source/features/interaction/gizmo.md @@ -12,7 +12,8 @@ Native entity interaction is enabled by default when the first native window opens. Select an object and press **G** to attach its gizmo. Set `SimulationManagerCfg(enable_entity_gizmo=False)` or call `sim.disable_entity_gizmo()` to opt out; reopening preserves that choice. -Robot IK controls require explicit creation for a selected control part. +Robot IK controls are discovered automatically from configured control parts +and their solver's root link, end link, and TCP. See {doc}`native window controls ` for native entity configuration. ## Supported Targets @@ -24,8 +25,9 @@ See {doc}`native window controls ` for native entity configuration. | Camera | Sets the camera's local pose. | Registered Viser controls and robot IK controls currently require `num_envs=1`. -Robot targets require a valid `control_part`; an EmbodiChain solver configuration -is needed only when selecting `GizmoCfg(ik_solver="embodichain")`. +Automatic discovery requires robot solver metadata for each supported control +part. IK computation still defaults to DexSim. Robots without this metadata +can use explicit `GizmoCfg` chain settings instead; no end link or TCP is guessed. ## Quick Start @@ -41,35 +43,43 @@ Use the same target through Viser: python scripts/tutorials/sim/gizmo_robot.py --viser ``` -For Viser, register controls through the manager: +The tutorial does not create or update a controller explicitly. After adding +the robot, its ordinary manual-physics loop is sufficient: ```python -sim.enable_gizmo( - uid="robot", - control_part="arm", -) - -if not sim.has_gizmo("robot", control_part="arm"): - raise RuntimeError("Gizmo setup failed") +sim.open_window() # Safely skipped when Viser is configured. +while True: + sim.update(step=1) ``` -`SimulationManager.update()` drains pending Gizmo commands during a normal -manual-physics loop. An automatic-physics loop that does not call `update()` -must continue calling: +In the native window, the first **I** press creates the eligible robot IK +controllers; later presses toggle their visibility. Opening a window alone +does not construct an IK solver or change existing drive targets. Viser displays +the TCP handles and constructs the solver on the first drag. Both paths update +through `SimulationManager.update()`. + +The default `SimulationManagerCfg.robot_ik_gizmo` is `GizmoCfg()`. To use each +control part's configured solver (including Pink), or disable automatic setup: ```python -sim.update_gizmos() -sim.capture_visualization_safely() # Publish the authoritative pose to Viser. +SimulationManagerCfg(robot_ik_gizmo=GizmoCfg(ik_solver="embodichain")) +SimulationManagerCfg(robot_ik_gizmo=None) ``` -Use `disable_gizmo()`, `set_gizmo_visibility()`, and -`toggle_gizmo_visibility()` for lifecycle and visibility changes. +Gym JSON/YAML deployments accept the same `robot_ik_gizmo` mapping or `null`. +`sim.enable_gizmo(uid, control_part, gizmo_cfg)` provides an explicit per-part +override. `sim.disable_gizmo(uid, control_part)` prevents automatic recreation, +including across window reopen; omit the part to disable all parts of a robot. +Use `set_gizmo_visibility()` and `toggle_gizmo_visibility()` for visibility. +Activated native controls retain their IK and visibility state across window +close/reopen. Removing a robot or destroying the manager releases their input +handlers and native target nodes. ## Frontend Behavior - Native entity interaction starts with the first native window; pure headless and Viser runs do not automatically create native controllers. Native robot - targets use `create_robot_ik_gizmo_controller()` explicitly. + targets activate on **I** without a script-level controller call. - Viser exports browser-native transform controls when `VisualizationCfg.allow_commands=True`; otherwise it shows read-only frames. - The standard `--viser` launcher enables registered commands for trusted @@ -77,6 +87,10 @@ Use `disable_gizmo()`, `set_gizmo_visibility()`, and - A Viser Gizmo is owned by one browser client from drag start until drag end or disconnect. Other clients continue receiving authoritative poses. +Advanced callers may still use `create_robot_ik_gizmo_controller()` directly +and retain/update its returned objects. Automatic setup detects an existing +factory-created controller and does not create or update a duplicate. + For the complete robot setup and IK walkthrough, continue with the {doc}`Gizmo tutorial `. See {doc}`Viser browser visualization ` for diff --git a/docs/source/features/interaction/window.md b/docs/source/features/interaction/window.md index c45db37fe..d339d118e 100644 --- a/docs/source/features/interaction/window.md +++ b/docs/source/features/interaction/window.md @@ -98,9 +98,10 @@ controller = sim.get_world().get_entity_gizmo() sim.disable_entity_gizmo() ``` -Robot TCP IK controllers still require explicit creation for a robot control -part. When both controllers are active, **G** controls entity roots and **I** -shows or hides the Robot TCP target. +Robot TCP IK controls are registered automatically for parts with configured +IK chain/TCP metadata. The first **I** press activates them; later presses show +or hide their targets. **G** continues to control entity roots. Normal +`sim.update()` calls handle IK updates; no controller-specific call is needed. The entity gizmo is native-window only. The Viser backend offers an analogous **click-to-pick** flow (an *Enable click-to-pick Gizmo* checkbox instead of the diff --git a/docs/source/tutorial/gizmo.rst b/docs/source/tutorial/gizmo.rst index e9204f9de..844c65356 100644 --- a/docs/source/tutorial/gizmo.rst +++ b/docs/source/tutorial/gizmo.rst @@ -7,7 +7,7 @@ Interactive Robot Control with Gizmo This tutorial demonstrates native DexSim and browser-based Viser Gizmo control. DexSim owns native entity and robot IK controllers; EmbodiChain keeps only the -robot control-part adapter and the Viser command path. +robot control-part adapter, Viser commands, and controller lifecycle management. For the cross-frontend capability summary, supported targets, lifecycle rules, and security boundary, see :doc:`/features/interaction/gizmo`. @@ -33,9 +33,10 @@ Similar to the previous tutorial on robot simulation, we use the :class:`Simulat -**Important:** Gizmo only supports single environment mode (`num_envs=1`). Using multiple environments will raise an exception. +**Important:** Gizmo supports a single environment (``num_envs=1``). Automatic +registration is skipped for multi-environment simulations. -Viser Gizmo creation, visibility, and destruction are managed through +Robot Gizmo registration, updates, visibility, and destruction are managed by SimulationManager: .. code-block:: python @@ -46,7 +47,7 @@ SimulationManager: # Set visibility explicitly sim.set_gizmo_visibility("ur10_gizmo_test", visible=False, control_part="arm") -Native controls use DexSim directly. The standard Viser mode includes +Native interaction uses DexSim controllers. The standard Viser mode includes interactive Gizmo control: .. code-block:: bash @@ -87,8 +88,8 @@ A Gizmo is an interactive visual tool that allows users to manipulate simulation - **Real-time Manipulation**: Provide immediate visual feedback during robot motion planning - **Debugging and Visualization**: Test robot reachability and workspace limits -The :class:`objects.Gizmo` class is the Viser-side target controller for robots, -rigid objects, and cameras. Native controls are DexSim controllers. +The :class:`objects.Gizmo` class manages native robot controllers and Viser +targets. Native entity manipulation remains owned by DexSim. Setting up Robot Configuration ------------------------------ @@ -104,7 +105,7 @@ Key components of the robot configuration: - **URDF Configuration**: Loads the robot's kinematic and visual model - **Control Parts**: Defines which joints can be controlled (``"Joint[1-6]"`` for UR10) -- **IK Solver**: :class:`solvers.PinkSolverCfg` provides inverse kinematics capabilities +- **IK Solver**: :class:`solvers.PinkSolverCfg` supplies chain metadata and an optional solver override - **Drive Properties**: Sets stiffness and damping for joint control The configured EmbodiChain solver is optional: it supplies default IK-chain @@ -112,54 +113,37 @@ metadata (root link, end link, and TCP transform). IK itself is solved by DexSim Newton IK. Applications may instead set this metadata directly in :class:`objects.GizmoCfg`. -Creating and Attaching a Gizmo -------------------------------- +Automatic Robot Controls +------------------------ +With the robot configuration above, no Gizmo-specific configuration or API +call is needed. SimulationManager discovers each control part with existing +root-link and end-link metadata and uses its configured TCP transform. +- In a native window, press **I** to create and show the IK targets. Further + presses toggle their visibility using DexSim's native controller. +- In Viser, the TCP controls are available automatically when commands are + allowed. The solver is constructed on the first drag. +- Pure headless, read-only Viser, and multi-environment simulations skip + automatic registration. -For native-window robot control, create DexSim's IK controller through the -small EmbodiChain adapter factory and retain both returned objects: +Opening a window or registering a control does not initialize another IK +solver or overwrite existing joint drive targets. DexSim Newton IK is the +default. To use the robot's configured solver instead: .. code-block:: python - from embodichain.lab.sim.objects import ( - GizmoCfg, - create_robot_ik_gizmo_controller, - ) + from embodichain.lab.sim.objects import GizmoCfg - ik_controller, input_controller = create_robot_ik_gizmo_controller( - robot, - control_part="arm", - cfg=GizmoCfg( - ik_root_link_name="base_link", - ik_end_link_name="ee_link", - ), - world=sim.get_world(), + sim_cfg = SimulationManagerCfg( + robot_ik_gizmo=GizmoCfg(ik_solver="embodichain"), ) -Call ``ik_controller.update()`` once per frame. DexSim owns the native target, -hotkey, solve trigger, and visibility state. For Viser, use the SimulationManager -command path instead: - -.. code-block:: python - - sim.enable_gizmo( - "ur10_gizmo_test", - control_part="arm", - gizmo_cfg=GizmoCfg( - ik_root_link_name="base_link", - ik_end_link_name="ee_link", - ), - ) - - - -The Gizmo system will automatically: - -1. **Resolve the IK Chain**: Locate the root and end-effector links -2. **Build Newton IK**: Construct DexSim's reduced-chain solver -3. **Bridge State**: Map one EmbodiChain control part to DexSim's joint API -4. **Own the Frontend**: DexSim owns native interaction; SimulationManager owns Viser commands +Set ``robot_ik_gizmo=None`` to disable automatic setup. Robots without chain +metadata can still use ``sim.enable_gizmo(...)`` with explicit +:class:`objects.GizmoCfg` link settings. Advanced callers can use +:func:`objects.create_robot_ik_gizmo_controller` and manage its updates directly; +SimulationManager will not create a duplicate native controller for that part. How Gizmo-Robot Interaction Works ---------------------------------- @@ -177,61 +161,34 @@ The gizmo-robot interaction follows this workflow: The Simulation Loop ------------------- +The tutorial uses manual physics only. After setting initial joint positions +and drive targets, each iteration advances one physics step: +.. literalinclude:: ../../../scripts/tutorials/sim/gizmo_robot.py + :language: python + :start-at: def run_simulation( + :end-at: sim.update(step=1) -Update the DexSim native controller explicitly, then service any Viser controls: - - - -.. code-block:: python - - def run_simulation(sim: SimulationManager, ik_controller=None): - step_count = 0 - try: - last_time = time.time() - last_step = 0 - while True: - time.sleep(0.033) # 30Hz - if ik_controller is not None: - ik_controller.update() - sim.update_gizmos() # Update Viser gizmos - sim.capture_visualization_safely() # Publish Viser state, if enabled - step_count += 1 - # ...performance statistics, etc... - except KeyboardInterrupt: - logger.log_info("\nStopping simulation...") - finally: - sim.destroy() # Release all resources - logger.log_info("Simulation terminated successfully") - - - -Main loop highlights: - -- **Native update**: Call DexSim's ``IKGizmoController.update()`` each frame -- **Viser command update**: Call ``sim.update_gizmos()`` -- **Viser frame update**: Automatic-physics loops also call ``sim.capture_visualization_safely()`` -- **Performance monitoring**: Optional FPS statistics -- **Resource cleanup**: Only `sim.destroy()` is needed, no manual Gizmo destruction -- **Graceful shutdown**: Supports Ctrl+C interruption +``sim.update()`` processes native and Viser interaction, advances physics, and +publishes visualization state. No separate controller update is required. +The tutorial paces the loop using ``physics_dt`` and releases resources with +``sim.destroy()`` on Ctrl+C. Gizmo Lifecycle Management --------------------------- - - - - -Viser Gizmo lifecycle is managed by SimulationManager: +------------------------- -- Enable: `sim.enable_gizmo(...)` -- Update: Call ``sim.update_gizmos()`` from the main loop -- Destroy/disable: `sim.disable_gizmo(...)` or `sim.destroy()` (recommended) +SimulationManager handles automatic robot control registration and cleanup. +Closing a native window detaches input handlers; reopening it reuses existing +controllers and preserves their visibility without writing new drive targets. +Removing a robot also removes its managed controls. -Native controller lifecycle remains in DexSim. Viser visual properties are -available through SimulationManager: +For explicit overrides: -- ``sim.toggle_gizmo_visibility(uid, control_part=None)``: Toggle gizmo visibility -- ``sim.set_gizmo_visibility(uid, visible, control_part=None)``: Set gizmo visibility +- ``sim.enable_gizmo(uid, control_part, gizmo_cfg)`` replaces that part's settings. +- ``sim.disable_gizmo(uid, control_part)`` disables one part and prevents automatic + recreation; omitting the part disables every part of the robot. +- ``sim.toggle_gizmo_visibility(uid, control_part)`` and + ``sim.set_gizmo_visibility(uid, visible, control_part)`` control visibility. Running the Tutorial -------------------- @@ -248,14 +205,16 @@ Command-line options: - ``--device cpu|cuda``: Choose simulation device - ``--num_envs N``: Number of parallel environments - ``--headless``: Run without GUI for automated testing -- ``--renderer``: Enable ray tracing for better visuals +- ``--renderer auto|hybrid|fast-rt|rt``: Select the renderer +- ``--viser``: Use browser-based interaction Once running: -1. **Mouse Interaction**: Click and drag the gizmo to move the robot -2. **Real-time IK**: Watch the robot joints automatically adjust to follow the gizmo -3. **Workspace Limits**: Observe how the robot behaves at workspace boundaries -4. **Performance**: Monitor FPS in the console output +1. **Activate**: Press **I** in the native window, or open the Viser page +2. **Mouse Interaction**: Click and drag the gizmo to move the robot +3. **Real-time IK**: Watch the robot joints automatically adjust to follow the gizmo +4. **Workspace Limits**: Observe how the robot behaves at workspace boundaries +5. **Performance**: Monitor FPS in the console output Tips and Best Practices ------------------------ @@ -264,10 +223,9 @@ Tips and Best Practices **Performance optimization:** -- Call the native IK controller's ``update()`` once per frame; call - ``sim.update_gizmos()`` for Viser +- Use ``sim.update(step=1)`` to service interaction and advance manual physics - Reduce IK solver iterations for better real-time performance if needed -- Use ``set_manual_update(False)`` for smoother interaction +- Pace manual steps using ``physics_dt`` diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 775491d2a..15c714737 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -621,6 +621,7 @@ class ComponentCfg: env_cfg.sim_cfg = SimulationManagerCfg( headless=config.get("headless", False), enable_entity_gizmo=config.get("enable_entity_gizmo", True), + robot_ik_gizmo=config.get("robot_ik_gizmo", {}), sim_device=config.get("device", "cpu"), render_cfg=RenderCfg(**render_config), gpu_id=config.get("gpu_id", 0), diff --git a/embodichain/lab/sim/objects/gizmo.py b/embodichain/lab/sim/objects/gizmo.py index 730cfd9e9..174282df5 100644 --- a/embodichain/lab/sim/objects/gizmo.py +++ b/embodichain/lab/sim/objects/gizmo.py @@ -20,6 +20,7 @@ import threading from typing import TYPE_CHECKING, Literal +from weakref import WeakValueDictionary import dexsim import numpy as np @@ -44,10 +45,16 @@ __all__ = ["Gizmo", "GizmoCfg", "create_robot_ik_gizmo_controller"] +# Explicit factory users retain ownership; automatic controls must not create +# a second controller for the same robot/control part. +_NATIVE_IK_CONTROLLERS: WeakValueDictionary[tuple[int, str], IKGizmoController] = ( + WeakValueDictionary() +) + @configclass class GizmoCfg: - """Configure Viser Gizmo appearance and robot IK behavior.""" + """Configure Gizmo appearance and native or Viser robot IK behavior.""" axis_length_x: float = 0.2 """Length of the X-axis arrow.""" @@ -435,22 +442,23 @@ def create_robot_ik_gizmo_controller( name=f"{robot.uid}_{control_part}_ik", ) window.add_input_control(input_controller) + _NATIVE_IK_CONTROLLERS[id(robot), control_part] = controller return controller, input_controller class Gizmo: - """Apply Viser Gizmo commands to one simulation target. + """Manage native robot IK or Viser control for one simulation target. Native-window entity manipulation is owned by DexSim. Use - :meth:`SimulationManager.enable_entity_gizmo` for entity roots and - :func:`create_robot_ik_gizmo_controller` for a native robot TCP target. + :meth:`SimulationManager.enable_entity_gizmo` for entity roots. The manager + discovers robot TCP controls and owns their updates and cleanup. .. attention:: - Viser Gizmo control supports exactly one simulation environment. + Gizmo control supports exactly one simulation environment. Args: - target: Rigid object, robot, or camera controlled from Viser. - cfg: Viser appearance and robot IK configuration. + target: Robot, or a rigid object or camera controlled from Viser. + cfg: Gizmo appearance and robot IK configuration. control_part: Robot control part used for FK and IK. """ @@ -479,11 +487,19 @@ def __init__( self._robot_adapter: _RobotGizmoAdapter | None = None self._robot_end_link: str | None = None self._robot_tcp_pose: np.ndarray | None = None + self._native_controller: IKGizmoController | None = None + self._native_input_controller: GizmoController | None = None + self._native_window = None + + from dexsim.kit.ik.interactive import KeyPressTracker + + self._native_toggle = KeyPressTracker(self.cfg.ik_toggle_key) if self._target_type == "robot": self._control_part = _resolve_control_part(target, control_part) - self._setup_robot_ik_solver() - self._desired_target_transform = self._read_robot_pose() + _, self._robot_end_link, self._robot_tcp_pose = _resolve_robot_ik_chain( + target, self._control_part, self.cfg + ) else: self._desired_target_transform = self._read_target_pose() @@ -526,17 +542,62 @@ def _setup_robot_ik_solver(self) -> None: def _read_robot_pose(self) -> torch.Tensor: if ( - self._robot_adapter is None + self.target is None or self._robot_end_link is None or self._robot_tcp_pose is None ): raise RuntimeError("Robot Gizmo IK is not configured.") - link_pose = self._robot_adapter.get_link_pose(self._robot_end_link) + link_pose = self.target.get_link_pose( + self._robot_end_link, env_ids=[0], to_matrix=True + )[0] return self._as_pose_matrix( - link_pose @ self._robot_tcp_pose, + link_pose + @ torch.as_tensor( + self._robot_tcp_pose, dtype=link_pose.dtype, device=link_pose.device + ), self._target_device(), ) + def update_native(self, world: dexsim.World) -> None: + """Activate robot IK on the first toggle, then update DexSim's controller. + + Registration and window reopening never write robot drive targets. + Explicit factory-created controllers remain owned by their caller. + + Args: + world: Simulation world with an open native window. + """ + if self._target_type != "robot" or self.target is None: + return + window = world.get_windows() + if window is None: + return + if self._native_controller is None: + existing = _NATIVE_IK_CONTROLLERS.get((id(self.target), self._control_part)) + if existing is not None or not self._native_toggle.pressed(world): + return + self._native_controller, self._native_input_controller = ( + create_robot_ik_gizmo_controller( + self.target, self._control_part, self.cfg, world=world + ) + ) + self._native_window = window + # The activation press created a visible target. Consume it so the + # controller's own toggle does not immediately hide that target. + self._native_controller.toggle.pressed(world) + return + if self._native_window is not window: + self.detach_native_window() + window.add_input_control(self._native_input_controller) + self._native_window = window + self._native_controller.update() + + def detach_native_window(self) -> None: + """Detach input before window close while retaining IK state and visibility.""" + if self._native_window is not None: + self._native_window.remove_input_control(self._native_input_controller) + self._native_window = None + def _target_device(self) -> torch.device: if self.target is None: return torch.device("cpu") @@ -621,9 +682,11 @@ def update(self) -> None: if pending is None or self.target is None: return - if self._robot_adapter is None: + if self._target_type != "robot": self.target.set_local_pose(pending, env_ids=[0]) return + if self._ik_solver is None: + self._setup_robot_ik_solver() adapter, solver = self._robot_adapter, self._ik_solver base_local = local_pose_from_world( adapter.get_world_pose(), pending[0].detach().cpu().numpy() @@ -637,20 +700,36 @@ def update(self) -> None: adapter.set_target_qpos(solver.qpos_for_joint_names(joint_names, current_qpos)) def toggle_visibility(self) -> bool: - """Toggle Viser visibility and return the new state.""" - self._is_visible = not self._is_visible - return self._is_visible + """Toggle Gizmo visibility and return the new state.""" + visible = not self.is_visible() + self.set_visible(visible) + return visible def set_visible(self, visible: bool) -> None: - """Set Viser Gizmo visibility.""" + """Set Gizmo visibility, including an activated native controller.""" self._is_visible = bool(visible) + if self._native_controller is not None: + self._native_controller._set_visible(self._is_visible) def is_visible(self) -> bool: - """Return whether the Viser Gizmo should be visible.""" + """Return the active native visibility or configured Viser visibility.""" + if self._native_controller is not None: + return self._native_controller.enabled return self._is_visible def destroy(self) -> None: """Release target and IK references.""" + self.detach_native_window() + if self._native_controller is not None: + controller = self._native_controller + env = controller.world.get_env() + env.remove_gizmo(controller.target_gizmo.gizmo) + env.remove_dummy_node(controller.target_gizmo.target_node) + key = (id(self.target), self._control_part) + if _NATIVE_IK_CONTROLLERS.get(key) is controller: + _NATIVE_IK_CONTROLLERS.pop(key) + self._native_controller = None + self._native_input_controller = None with self._state_lock: self._interaction_owner = None self._pending_target_transform = None diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 721770d10..10a784d54 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -140,7 +140,17 @@ class SimulationManagerCfg: Headless and Viser simulations do not automatically create native gizmos. Explicit runtime enable/disable calls take precedence over this default; reopening a window preserves the current DexSim controller state. - Robot IK gizmos are created separately for a selected control part. + Robot IK interaction is controlled separately by ``robot_ik_gizmo``. + """ + + robot_ik_gizmo: GizmoCfg | None = field(default_factory=GizmoCfg) + """Automatically register robot parts with configured IK chain/TCP metadata. + + Native controls activate on the first I press; Viser creates its IK solver + on the first drag and requires ``visualization.allow_commands``. Pure + headless and multi-environment simulations do not register controls. + Defaults to DexSim IK. Use ``GizmoCfg(ik_solver="embodichain")`` to reuse + configured solvers such as Pink, or ``None`` to disable automatic setup. """ render_cfg: RenderCfg = field(default_factory=RenderCfg) @@ -198,6 +208,12 @@ class SimulationManagerCfg: def __post_init__(self) -> None: """Apply visualization-dependent simulation defaults.""" + if isinstance(self.robot_ik_gizmo, dict): + self.robot_ik_gizmo = GizmoCfg(**self.robot_ik_gizmo) + if self.robot_ik_gizmo is not None and not isinstance( + self.robot_ik_gizmo, GizmoCfg + ): + raise TypeError("robot_ik_gizmo must be a GizmoCfg, mapping, or None.") if self.visualization.backend == "viser": self.headless = True @@ -347,6 +363,7 @@ def __init__( # gizmo management self._gizmos: Dict[str, object] = dict() # Store active gizmos + self._disabled_robot_gizmos: set[str] = set() # ``(uid, control_part)`` of the Gizmo currently owned by the Viser # click-to-pick feature, or ``None``. Only one picker Gizmo is kept at a # time and user-created Gizmos are never touched. @@ -1026,6 +1043,8 @@ def close_window(self) -> None: """Close the simulation window.""" if self.is_window_recording(): self.stop_window_record() + for _, gizmo in self.get_gizmo_items(): + gizmo.detach_native_window() self._world.close_window() self._window = None self._window_record_input_control = None @@ -2075,12 +2094,16 @@ def enable_gizmo( control_part: str | None = None, gizmo_cfg: GizmoCfg | None = None, ) -> Gizmo | None: - """Enable Viser Gizmo control for a simulation target. + """Register Gizmo control for a simulation target. + + Robot controls support native windows and Viser; native IK activates + on the configured toggle key. Explicit configuration takes precedence + over automatically registered controls. Args: uid: UID of the robot, rigid object, or camera sensor. control_part: Robot control part used for IK/FK. - gizmo_cfg: Viser appearance and robot IK configuration. + gizmo_cfg: Gizmo appearance and robot IK configuration. Returns: The created Gizmo, or ``None`` if setup failed. @@ -2088,12 +2111,17 @@ def enable_gizmo( # Create gizmo key combining uid and control_part gizmo_key = f"{uid}:{control_part}" if control_part else uid + if hasattr(self, "_disabled_robot_gizmos"): + self._disabled_robot_gizmos.discard(gizmo_key) + # Check if gizmo already exists if gizmo_key in self._gizmos: - logger.log_warning( - f"Gizmo for '{uid}' with control_part '{control_part}' already exists." - ) - return self._gizmos[gizmo_key] + if gizmo_cfg is not None: + self.disable_gizmo(uid, control_part) + if hasattr(self, "_disabled_robot_gizmos"): + self._disabled_robot_gizmos.discard(gizmo_key) + else: + return self._gizmos[gizmo_key] # Search for target object in different collections target = None @@ -2121,7 +2149,7 @@ def enable_gizmo( self._gizmos[gizmo_key] = gizmo self.notify_visualization_topology_changed() logger.log_info( - f"Viser Gizmo enabled for {object_type} '{uid}' with " + f"Gizmo registered for {object_type} '{uid}' with " f"control_part '{control_part}'." ) @@ -2136,14 +2164,23 @@ def enable_gizmo( return gizmo def disable_gizmo(self, uid: str, control_part: str | None = None) -> None: - """Disable and remove a Gizmo. + """Disable a Gizmo and prevent its automatic recreation. Args: uid: Target asset UID. - control_part: Robot control part, if applicable. + control_part: Robot control part. Omit to disable every part of a robot. """ gizmo_key = f"{uid}:{control_part}" if control_part else uid + if hasattr(self, "_disabled_robot_gizmos"): + self._disabled_robot_gizmos.add(gizmo_key) + robot = getattr(self, "_robots", {}).get(uid) + if robot is not None and control_part is None: + for key, gizmo in self.get_gizmo_items(): + if key != gizmo_key and gizmo.target is robot: + self.disable_gizmo(key) if gizmo_key not in self._gizmos: + if robot is not None: + return logger.log_warning( f"No gizmo found for '{uid}' with control_part '{control_part}'." ) @@ -2268,8 +2305,49 @@ def process_visualization_commands(self) -> int: gizmo.end_interaction(source_id) return accepted + def _register_robot_gizmos(self) -> None: + """Discover IK-capable parts without constructing another solver.""" + cfg = getattr(self.sim_config, "robot_ik_gizmo", None) + if cfg is None or self.num_envs != 1: + return + native = self._window is not None + viser = ( + self.sim_config.visualization.backend == "viser" + and self.sim_config.visualization.allow_commands + ) + if not native and not viser: + return + for uid, robot in self._robots.items(): + solver_cfgs = robot.cfg.solver_cfg + if not isinstance(solver_cfgs, dict): + continue + for part, solver_cfg in solver_cfgs.items(): + key = f"{uid}:{part}" + if ( + uid in self._disabled_robot_gizmos + or key in self._disabled_robot_gizmos + or key in self._gizmos + or part not in (robot.control_parts or {}) + or not getattr(solver_cfg, "root_link_name", None) + or not getattr(solver_cfg, "end_link_name", None) + ): + continue + if any( + gizmo.target is robot and gizmo.control_part == part + for _, gizmo in self.get_gizmo_items() + ): + continue + try: + self.enable_gizmo(uid, part, deepcopy(cfg)) + except Exception as error: + self._disabled_robot_gizmos.add(key) + logger.log_warning( + f"Could not register robot Gizmo '{key}': {error}" + ) + def update_gizmos(self) -> None: - """Apply Viser commands and update all active Gizmos.""" + """Register robot controls, process interaction, and update active Gizmos.""" + self._register_robot_gizmos() self.process_pick_commands() self.process_visualization_commands() for gizmo_key, gizmo in list( @@ -2277,9 +2355,11 @@ def update_gizmos(self) -> None: ): # Use list() to avoid modification during iteration if gizmo is not None: try: + if getattr(self, "_window", None) is not None: + gizmo.update_native(self._world) gizmo.update() except Exception as error: - logger.log_error(f"Error updating gizmo '{gizmo_key}': {error}") + logger.log_warning(f"Error updating gizmo '{gizmo_key}': {error}") def process_pick_commands(self) -> int: """Apply queued Viser click-pick commands on the simulation thread. @@ -2457,6 +2537,14 @@ def remove_asset(self, uid: str) -> bool: if uid in self._robots: robot = self._robots.pop(uid) + for key, gizmo in self.get_gizmo_items(): + if gizmo.target is robot: + self.disable_gizmo(key) + self._disabled_robot_gizmos = { + key + for key in self._disabled_robot_gizmos + if key != uid and not key.startswith(f"{uid}:") + } robot.destroy() self.notify_visualization_topology_changed() return True diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index 6b830ac55..c041e335a 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -24,10 +24,6 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.objects import ( - GizmoCfg, - create_robot_ik_gizmo_controller, -) from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.cfg import ( RenderCfg, @@ -111,28 +107,7 @@ def main(): if not args.headless: native_window_opened = sim.open_window() - gizmo_cfg = GizmoCfg( - ik_root_link_name="base_link", - ik_end_link_name="ee_link", - ) - native_control = None - if native_window_opened: - native_control = create_robot_ik_gizmo_controller( - robot, - control_part="arm", - cfg=gizmo_cfg, - world=sim.get_world(), - ) - elif args.viser: - sim.enable_gizmo( - uid="ur10_gizmo_test", - control_part="arm", - gizmo_cfg=gizmo_cfg, - ) - if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): - logger.log_error("Failed to enable gizmo!") - return - else: + if not native_window_opened and not args.viser: logger.log_warning( "Gizmo interaction is disabled in headless mode without Viser." ) @@ -141,23 +116,21 @@ def main(): if native_window_opened or args.viser: logger.log_info("Use the gizmo to drag the robot end-effector (EE)") if native_window_opened: - logger.log_info("Press I to show or hide the native robot IK Gizmo") + logger.log_info("Press I to activate or toggle the native robot IK Gizmo") logger.log_info("Press Ctrl+C to stop the simulation") - run_simulation(sim, native_control) + run_simulation(sim) -def run_simulation(sim: SimulationManager, native_control=None): - """Update IK and advance one manual physics step per frame.""" +def run_simulation(sim: SimulationManager) -> None: + """Advance physics; the manager owns native and Viser IK interaction.""" step_count = 0 try: last_time = time.perf_counter() last_step = 0 while True: frame_start = time.perf_counter() - if native_control is not None: - native_control[0].update() - # update() also processes Viser commands and publishes the frame. + # update() owns IK interaction, physics stepping, and Viser capture. sim.update(step=1) step_count += 1 diff --git a/tests/gym/utils/test_gym_utils.py b/tests/gym/utils/test_gym_utils.py index 1c00b41df..7ea9a26cb 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -1407,6 +1407,30 @@ def test_gym_config_preserves_entity_gizmo_startup_preference( assert cfg.sim_cfg.enable_entity_gizmo is (enabled is not False) + @pytest.mark.parametrize("suffix", ["yaml", "json"]) + @pytest.mark.parametrize("settings", [None, {}, {"ik_solver": "embodichain"}]) + def test_gym_config_parses_automatic_robot_gizmo_settings( + self, tmp_path: Path, suffix: str, settings: dict | None + ) -> None: + """Deployments may disable automatic IK controls or select their solver.""" + path = tmp_path / f"gym_config.{suffix}" + save_config( + path, + { + "id": "EmbodiedEnv-v1", + "env": {}, + "robot": {"uid": "robot"}, + "robot_ik_gizmo": settings, + }, + ) + cfg = config_to_cfg(load_config(path), manager_modules=DEFAULT_MANAGER_MODULES) + if settings is None: + assert cfg.sim_cfg.robot_ik_gizmo is None + else: + assert cfg.sim_cfg.robot_ik_gizmo.ik_solver == settings.get( + "ik_solver", "dexsim" + ) + def test_yaml_gym_config_parses_to_cfg(self, tmp_path): config = { "id": "EmbodiedEnv-v1", diff --git a/tests/sim/objects/test_gizmo.py b/tests/sim/objects/test_gizmo.py index 39b1f93bd..2a525f025 100644 --- a/tests/sim/objects/test_gizmo.py +++ b/tests/sim/objects/test_gizmo.py @@ -270,7 +270,12 @@ def _inject_solver(self) -> None: monkeypatch.setattr(Gizmo, "_setup_robot_ik_solver", _inject_solver) - gizmo = Gizmo(target, control_part="arm") + gizmo = Gizmo( + target, + GizmoCfg(ik_root_link_name="base_link", ik_end_link_name="tool_link"), + control_part="arm", + ) + assert gizmo._ik_solver is None pose = torch.eye(4, dtype=torch.float32).unsqueeze(0) pose[0, 0, 3] = 0.5 diff --git a/tests/sim/test_robot_gizmo_lifecycle.py b/tests/sim/test_robot_gizmo_lifecycle.py new file mode 100644 index 000000000..2c45944dd --- /dev/null +++ b/tests/sim/test_robot_gizmo_lifecycle.py @@ -0,0 +1,304 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from dexsim.kit.ik.interactive import KeyPressTracker + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.objects import GizmoCfg +from embodichain.lab.sim.objects import gizmo as gizmo_module +from embodichain.lab.visualization import VisualizationCfg + +pytestmark = pytest.mark.no_sim + + +class _Robot: + __hash__ = None # Robot inherits dataclass equality and is not hashable. + num_instances = 1 + device = torch.device("cpu") + + def __init__(self) -> None: + self.control_parts = { + "left": ["left_joint"], + "right": ["right_joint"], + "hand": ["finger"], + } + self.cfg = SimpleNamespace( + uid="robot", + solver_cfg={ + part: SimpleNamespace( + root_link_name="base", end_link_name=f"{part}_tool" + ) + for part in ("left", "right") + }, + ) + self.pose = torch.eye(4).unsqueeze(0) + self.destroy = Mock() + + def get_solver(self, part: str) -> SimpleNamespace: + cfg = self.cfg.solver_cfg[part] + return SimpleNamespace(**vars(cfg), get_tcp=lambda: np.eye(4)) + + def get_link_pose( + self, link: str, *, env_ids: list[int], to_matrix: bool + ) -> torch.Tensor: + return self.pose.clone() + + +class _Window: + def __init__(self) -> None: + self.controls: list[object] = [] + + def add_input_control(self, control: object) -> None: + self.controls.append(control) + + def remove_input_control(self, control: object) -> None: + self.controls.remove(control) + + +class _World: + def __init__(self) -> None: + self.window = _Window() + self.key_down = False + self.env = SimpleNamespace(remove_gizmo=Mock(), remove_dummy_node=Mock()) + + def get_windows(self) -> _Window | None: + return self.window + + def get_env(self) -> SimpleNamespace: + return self.env + + def key_state(self, key: object) -> bool: + return self.key_down + + def close_window(self) -> None: + self.window = None + + def open_window(self) -> None: + self.window = _Window() + + +class _Controller: + def __init__(self, world: _World, cfg: GizmoCfg) -> None: + self.world = world + self.toggle = KeyPressTracker(cfg.ik_toggle_key) + self.enabled = True + self.target_gizmo = SimpleNamespace(gizmo=object(), target_node=object()) + + def update(self) -> None: + if self.toggle.pressed(self.world): + self.enabled = not self.enabled + + def _set_visible(self, visible: bool) -> None: + self.enabled = visible + + +@pytest.fixture +def managed_robot(monkeypatch: pytest.MonkeyPatch): + """Use the real manager/Gizmo lifecycle with rendering and IK construction stubbed.""" + monkeypatch.setattr(gizmo_module, "Robot", _Robot) + robot = _Robot() + sim = object.__new__(SimulationManager) + sim.sim_config = SimulationManagerCfg(headless=True) + sim.num_envs = 1 + sim._world = _World() + sim._window = sim._world.window + sim.is_window_opened = True + sim._window_record_state = None + sim._window_record_hotkey_cfg = None + sim._window_camera_pose_hotkey_cfg = None + sim._auto_entity_gizmo_pending = False + sim._visualization_runtime = None + sim._visualization_topology_revision = 0 + sim._gizmos = {} + sim._disabled_robot_gizmos = set() + sim._picker_gizmo = None + sim._robots = {robot.cfg.uid: robot} + for registry in ( + "_rigid_objects", + "_rigid_object_groups", + "_articulations", + "_soft_objects", + "_cloth_objects", + "_sensors", + ): + setattr(sim, registry, {}) + sim.process_pick_commands = Mock(return_value=0) + sim.process_visualization_commands = Mock(return_value=0) + + def create(robot, part, cfg, *, world): + controller = _Controller(world, cfg) + control = object() + world.get_windows().add_input_control(control) + return controller, control + + factory = Mock(side_effect=create) + monkeypatch.setattr(gizmo_module, "create_robot_ik_gizmo_controller", factory) + yield sim, robot, factory + for _, gizmo in sim.get_gizmo_items(): + gizmo.destroy() + + +@pytest.mark.parametrize( + "frontend", ["native", "viser", "read_only", "headless", "batch", "disabled"] +) +def test_automatic_registration_respects_interaction_boundaries( + managed_robot, frontend: str +) -> None: + sim, robot, factory = managed_robot + if frontend in {"viser", "read_only", "headless"}: + sim._window = None + if frontend in {"viser", "read_only"}: + sim.sim_config.visualization = VisualizationCfg( + backend="viser", allow_commands=frontend == "viser" + ) + if frontend == "batch": + sim.num_envs = 2 + if frontend == "disabled": + sim.sim_config.robot_ik_gizmo = None + sim.update_gizmos() + sim.update_gizmos() + expected = ( + {"robot:left", "robot:right"} if frontend in {"native", "viser"} else set() + ) + assert set(sim.list_gizmos()) == expected + factory.assert_not_called() + for _, gizmo in sim.get_gizmo_items(): + assert gizmo._ik_solver is None + robot.pose[0, 0, 3] = 0.5 + assert gizmo.get_control_pose()[0, 0, 3] == pytest.approx(0.5) + + +def test_native_activation_and_reopen_preserve_controller_and_visibility( + managed_robot, +) -> None: + sim, _, factory = managed_robot + sim.update_gizmos() + sim._world.key_down = True + sim.update_gizmos() + controls = [gizmo._native_controller for _, gizmo in sim.get_gizmo_items()] + assert factory.call_count == 2 + sim.update_gizmos() # Holding I does not toggle a second time. + assert all(control.enabled for control in controls) + sim._world.key_down = False + sim.update_gizmos() + sim._world.key_down = True + sim.update_gizmos() + assert not any(control.enabled for control in controls) + + old_window = sim._window + sim.close_window() + assert old_window.controls == [] + assert sim.open_window() + sim.update_gizmos() + assert factory.call_count == 2 + assert len(sim._window.controls) == 2 + assert [g._native_controller for _, g in sim.get_gizmo_items()] == controls + assert not any(control.enabled for control in controls) + + +def test_explicit_disable_and_solver_override_take_precedence(managed_robot) -> None: + sim, _, _ = managed_robot + sim.disable_gizmo("robot", "left") + sim.update_gizmos() + assert set(sim.list_gizmos()) == {"robot:right"} + sim.enable_gizmo("robot", "right", GizmoCfg(ik_solver="embodichain")) + right = sim.get_gizmo("robot", "right") + sim.update_gizmos() + assert sim.get_gizmo("robot", "right") is right + assert right.cfg.ik_solver == "embodichain" + sim.enable_gizmo("robot", "left") + assert set(sim.list_gizmos()) == {"robot:left", "robot:right"} + + +def test_whole_robot_disable_cleans_activated_controls(managed_robot) -> None: + sim, _, factory = managed_robot + sim._world.key_down = True + sim.update_gizmos() + assert not sim.toggle_gizmo_visibility("robot", "left") + assert not sim.get_gizmo("robot", "left")._native_controller.enabled + sim.set_gizmo_visibility("robot", True, "left") + assert sim.get_gizmo("robot", "left")._native_controller.enabled + sim.disable_gizmo("robot") + sim.update_gizmos() + assert not sim.list_gizmos() + assert not sim._window.controls + assert factory.call_count == 2 + + +def test_native_activation_failure_does_not_interrupt_other_gizmos( + managed_robot, +) -> None: + sim, _, factory = managed_robot + create = factory.side_effect + + def fail_left(robot, part, cfg, *, world): + if part == "left": + raise ValueError("invalid chain") + return create(robot, part, cfg, world=world) + + factory.side_effect = fail_left + sim._world.key_down = True + sim.update_gizmos() + sim.update_gizmos() + assert factory.call_count == 2 # A held key does not repeatedly retry a failure. + assert sim.get_gizmo("robot", "left")._native_controller is None + assert sim.get_gizmo("robot", "right")._native_controller is not None + + +def test_robot_gizmo_setting_rejects_boolean() -> None: + with pytest.raises(TypeError, match="robot_ik_gizmo"): + SimulationManagerCfg(robot_ik_gizmo=False) + + +def test_robot_removal_cleans_native_controls_and_allows_same_uid_again( + managed_robot, +) -> None: + sim, robot, _ = managed_robot + sim._world.key_down = True + sim.update_gizmos() + assert sim.remove_asset("robot") + assert not sim.list_gizmos() + assert not sim._window.controls + assert sim._world.env.remove_gizmo.call_count == 2 + assert sim._world.env.remove_dummy_node.call_count == 2 + robot.destroy.assert_called_once() + sim._robots["robot"] = _Robot() + sim._world.key_down = False + sim.update_gizmos() + assert set(sim.list_gizmos()) == {"robot:left", "robot:right"} + + +def test_explicit_native_factory_controller_is_not_duplicated( + managed_robot, monkeypatch: pytest.MonkeyPatch +) -> None: + sim, robot, factory = managed_robot + explicit = Mock() + monkeypatch.setitem( + gizmo_module._NATIVE_IK_CONTROLLERS, (id(robot), "left"), explicit + ) + sim._world.key_down = True + sim.update_gizmos() + assert factory.call_count == 1 + assert factory.call_args.args[1] == "right" + explicit.update.assert_not_called() diff --git a/tests/visualization/test_gizmo_robot_tutorial.py b/tests/visualization/test_gizmo_robot_tutorial.py index d6c661beb..0b30f34d4 100644 --- a/tests/visualization/test_gizmo_robot_tutorial.py +++ b/tests/visualization/test_gizmo_robot_tutorial.py @@ -26,11 +26,9 @@ pytestmark = pytest.mark.no_sim -@pytest.mark.parametrize("native", [True, False]) @pytest.mark.parametrize("work_duration", [0.002, 0.050]) def test_gizmo_loop_advances_manual_physics_and_paces_frames( monkeypatch: pytest.MonkeyPatch, - native: bool, work_duration: float, ) -> None: """Apply IK before stepping, avoid duplicate updates, and clean up on exit.""" @@ -38,15 +36,13 @@ def test_gizmo_loop_advances_manual_physics_and_paces_frames( calls = Mock() sim = calls.sim sim.sim_config = SimpleNamespace(physics_dt=physics_dt) - native_control = (calls.ik, object()) if native else None clock = Mock(side_effect=[0.0, 0.0, work_duration]) sleep = Mock(side_effect=KeyboardInterrupt) monkeypatch.setattr(gizmo_robot.time, "perf_counter", clock) monkeypatch.setattr(gizmo_robot.time, "sleep", sleep) - gizmo_robot.run_simulation(sim, native_control) + gizmo_robot.run_simulation(sim) - expected_calls = [call.ik.update()] if native else [] - expected_calls.extend([call.sim.update(step=1), call.sim.destroy()]) + expected_calls = [call.sim.update(step=1), call.sim.destroy()] assert calls.mock_calls == expected_calls sleep.assert_called_once_with(pytest.approx(max(0.0, physics_dt - work_duration))) From 5e2dafd2363af6dba9f718fdac7a4972c9e23f49 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 5 Sep 2026 19:08:36 +0800 Subject: [PATCH 10/14] feat(gizmo): allow explicit native IK activation at startup --- .../sim-visualization/sim-visualization.md | 5 ++- .../simulation-system/simulation-system.md | 9 +++- .../embodichain.lab.sim.objects.rst | 4 +- .../embodichain.lab.sim.sim_manager.rst | 5 ++- docs/source/features/interaction/gizmo.md | 11 +++-- docs/source/features/interaction/window.md | 9 +++- docs/source/tutorial/gizmo.rst | 30 +++++++++---- embodichain/lab/sim/objects/gizmo.py | 25 ++++++++--- embodichain/lab/sim/sim_manager.py | 2 + scripts/tutorials/sim/gizmo_robot.py | 6 ++- tests/gym/utils/test_gym_utils.py | 7 ++- tests/sim/test_robot_gizmo_lifecycle.py | 45 ++++++++++++++++++- 12 files changed, 128 insertions(+), 30 deletions(-) diff --git a/agent_context/topics/sim-visualization/sim-visualization.md b/agent_context/topics/sim-visualization/sim-visualization.md index 2a080a359..bf4c70623 100644 --- a/agent_context/topics/sim-visualization/sim-visualization.md +++ b/agent_context/topics/sim-visualization/sim-visualization.md @@ -257,8 +257,9 @@ native window open; `SimulationManagerCfg.enable_entity_gizmo=False` or `sim.disable_entity_gizmo()` opts out and reopening preserves the choice. Pure headless and Viser runs do not automatically create native controls. `SimulationManagerCfg.robot_ik_gizmo` automatically registers robot parts with -solver chain/TCP metadata: native IK activates on I and Viser builds IK on the -first drag. The manager updates both and preserves native controls across window +solver chain/TCP metadata: native IK activates on I by default, or on the first +update with an open window when `GizmoCfg(ik_start_enabled=True)`. The robot +tutorial uses this startup option. Viser builds IK on the first drag. The manager updates both and preserves native controls across window reopen. Explicit configuration and disable calls take precedence; caller-owned native factory controllers are not duplicated. Automatic Viser registration requires `allow_commands`, independently of native interaction preferences. diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index 156996f3e..d2dd7b620 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -72,7 +72,9 @@ environment control step normally calls it with `scripts/tutorials/sim/gizmo_robot.py` supports only manual physics. It initializes GPU physics after robot creation when needed, sets both current and target -joint positions, and advances once before opening the window. Its loop only +joint positions, and advances once before opening the window. It explicitly +sets `GizmoCfg(ik_start_enabled=True)` so the native controller activates on the +first update after opening the window. Its loop only calls `sim.update(step=1)`; the manager owns native IK updates and Viser commands/capture. The loop is paced by `physics_dt` and has no automatic physics polling path. @@ -138,7 +140,10 @@ automatically create a native entity controller. updates the manager registers robot control parts with complete solver chain/TCP metadata in single-environment interactive runs. Pure headless, read-only Viser, and multi-environment runs do not register automatic controls. The first native I -press creates DexSim's `IKGizmoController`; Viser constructs IK on its first drag. +press creates DexSim's `IKGizmoController` by default; +`GizmoCfg(ik_start_enabled=True)` opts into activation on the first update with +an open window. The startup attempt is consumed once, including on failure; +later key presses can retry. Viser constructs IK on its first drag. Registration never writes drive targets. `Gizmo` owns managed native input and target-node cleanup, detaches input on window close, and reattaches the same controller on reopen. Robot removal releases all its managed controls. diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst index af8903319..fc0a040fd 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst @@ -195,10 +195,12 @@ Gizmo .. autofunction:: create_robot_ik_gizmo_controller SimulationManager automatically discovers robot control parts with configured -IK chain metadata. Native controllers activate on the first I press, while +IK chain metadata. Native controllers activate on the first I press by default, while Viser constructs its solver on the first drag. ``sim.update()`` owns updates and cleanup; ordinary applications do not need the explicit factory. Use ``SimulationManagerCfg(robot_ik_gizmo=None)`` to disable automatic setup. +Set ``GizmoCfg(ik_start_enabled=True)`` to activate native IK on the first update +with an open window; subsequent visibility toggles and reopening are preserved. The native controller defaults to DexSim Newton IK. With a ``PinkSolverCfg`` (or another EmbodiChain solver) configured for the robot's control part, pass diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst index 1af4b6a43..1131a77f8 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst @@ -21,12 +21,13 @@ and Viser runs do not automatically create native gizmos. ``SimulationManagerCfg.robot_ik_gizmo`` defaults to ``GizmoCfg()`` and registers robot control parts with configured IK-chain/TCP metadata during normal updates. -Native IK activates on the first **I** press; Viser constructs its solver on +Native IK activates on the first **I** press by default; Viser constructs its solver on the first drag and requires ``visualization.allow_commands``. Registration does not write drive targets. Set this field to ``None`` to opt out or select ``GizmoCfg(ik_solver="embodichain")`` to reuse configured solvers. Explicit ``enable_gizmo()`` settings override automatic defaults, and ``disable_gizmo()`` -prevents automatic recreation. +prevents automatic recreation. ``GizmoCfg(ik_start_enabled=True)`` activates +native IK on the first update with an open window, as used by the robot tutorial. .. rubric:: Classes diff --git a/docs/source/features/interaction/gizmo.md b/docs/source/features/interaction/gizmo.md index 42856e155..f8992e972 100644 --- a/docs/source/features/interaction/gizmo.md +++ b/docs/source/features/interaction/gizmo.md @@ -43,8 +43,10 @@ Use the same target through Viser: python scripts/tutorials/sim/gizmo_robot.py --viser ``` -The tutorial does not create or update a controller explicitly. After adding -the robot, its ordinary manual-physics loop is sufficient: +The tutorial opts into immediate native activation with +`robot_ik_gizmo=GizmoCfg(ik_start_enabled=True)`. It sets the initial robot pose +before opening the window, then its ordinary manual-physics loop creates and +updates the controller: ```python sim.open_window() # Safely skipped when Viser is configured. @@ -52,8 +54,9 @@ while True: sim.update(step=1) ``` -In the native window, the first **I** press creates the eligible robot IK -controllers; later presses toggle their visibility. Opening a window alone +By default, the first **I** press creates eligible native robot IK controllers. +With `ik_start_enabled=True`, the first update with an open window activates +them from the current robot pose; **I** then toggles their visibility. Opening a window alone does not construct an IK solver or change existing drive targets. Viser displays the TCP handles and constructs the solver on the first drag. Both paths update through `SimulationManager.update()`. diff --git a/docs/source/features/interaction/window.md b/docs/source/features/interaction/window.md index d339d118e..53cd59576 100644 --- a/docs/source/features/interaction/window.md +++ b/docs/source/features/interaction/window.md @@ -99,10 +99,15 @@ sim.disable_entity_gizmo() ``` Robot TCP IK controls are registered automatically for parts with configured -IK chain/TCP metadata. The first **I** press activates them; later presses show -or hide their targets. **G** continues to control entity roots. Normal +IK chain/TCP metadata. By default, the first **I** press activates them; later +presses show or hide their targets. **G** continues to control entity roots. Normal `sim.update()` calls handle IK updates; no controller-specific call is needed. +For immediate activation, set +`SimulationManagerCfg(robot_ik_gizmo=GizmoCfg(ik_start_enabled=True))`. The +controller starts on the first update with an open window. The robot Gizmo +tutorial enables this option after preparing its initial joint pose. + The entity gizmo is native-window only. The Viser backend offers an analogous **click-to-pick** flow (an *Enable click-to-pick Gizmo* checkbox instead of the **G** hotkey, since browsers do not expose keyboard events); see diff --git a/docs/source/tutorial/gizmo.rst b/docs/source/tutorial/gizmo.rst index 844c65356..9e1592005 100644 --- a/docs/source/tutorial/gizmo.rst +++ b/docs/source/tutorial/gizmo.rst @@ -116,20 +116,32 @@ DexSim Newton IK. Applications may instead set this metadata directly in Automatic Robot Controls ------------------------ -With the robot configuration above, no Gizmo-specific configuration or API -call is needed. SimulationManager discovers each control part with existing -root-link and end-link metadata and uses its configured TCP transform. +SimulationManager discovers each control part with existing root-link and +end-link metadata and uses its configured TCP transform. This tutorial +explicitly activates native IK at startup with one setting: -- In a native window, press **I** to create and show the IK targets. Further - presses toggle their visibility using DexSim's native controller. +.. code-block:: python + + sim_cfg = SimulationManagerCfg( + robot_ik_gizmo=GizmoCfg(ik_start_enabled=True), + ) + +The first update after opening the window creates and displays the controller. +The tutorial sets current joint positions and drive targets before opening the +window; activation initializes the controller from that pose. + +- In this tutorial's native window, IK targets start visible. Press **I** to + hide or show them. Ordinary simulations retain ``ik_start_enabled=False`` + and wait for the first **I** press to activate. - In Viser, the TCP controls are available automatically when commands are allowed. The solver is constructed on the first drag. - Pure headless, read-only Viser, and multi-environment simulations skip automatic registration. -Opening a window or registering a control does not initialize another IK -solver or overwrite existing joint drive targets. DexSim Newton IK is the -default. To use the robot's configured solver instead: +Registration alone does not initialize another IK solver or overwrite drive +targets. Explicit startup activation initializes the controller once; closing +and reopening the window preserves its later visibility state. DexSim Newton IK +is the default. To use the robot's configured solver instead: .. code-block:: python @@ -210,7 +222,7 @@ Command-line options: Once running: -1. **Activate**: Press **I** in the native window, or open the Viser page +1. **Open**: The native IK target starts visible, or open the Viser page 2. **Mouse Interaction**: Click and drag the gizmo to move the robot 3. **Real-time IK**: Watch the robot joints automatically adjust to follow the gizmo 4. **Workspace Limits**: Observe how the robot behaves at workspace boundaries diff --git a/embodichain/lab/sim/objects/gizmo.py b/embodichain/lab/sim/objects/gizmo.py index 174282df5..92b07d9e9 100644 --- a/embodichain/lab/sim/objects/gizmo.py +++ b/embodichain/lab/sim/objects/gizmo.py @@ -103,6 +103,15 @@ class GizmoCfg: ik_toggle_key: InputKey = InputKey.SCANCODE_I """Native-window key used to toggle a DexSim robot IK target.""" + ik_start_enabled: bool = False + """Activate managed native IK on the first update with an open window. + + Defaults to waiting for the toggle key. Explicit startup activation creates + the controller and initializes its drive targets from the current pose. + This is a one-time request: later toggles and window reopening retain their + normal behavior. Viser still builds its solver on the first drag. + """ + class _RobotGizmoAdapter: """Expose one EmbodiChain robot control part to DexSim's IK controller.""" @@ -490,6 +499,7 @@ def __init__( self._native_controller: IKGizmoController | None = None self._native_input_controller: GizmoController | None = None self._native_window = None + self._native_start_pending = self.cfg.ik_start_enabled from dexsim.kit.ik.interactive import KeyPressTracker @@ -559,9 +569,10 @@ def _read_robot_pose(self) -> torch.Tensor: ) def update_native(self, world: dexsim.World) -> None: - """Activate robot IK on the first toggle, then update DexSim's controller. + """Activate robot IK on startup or a toggle, then update its controller. - Registration and window reopening never write robot drive targets. + Startup activation is opt-in; otherwise only the toggle key activates + IK. Registration and window reopening never write robot drive targets. Explicit factory-created controllers remain owned by their caller. Args: @@ -574,7 +585,11 @@ def update_native(self, world: dexsim.World) -> None: return if self._native_controller is None: existing = _NATIVE_IK_CONTROLLERS.get((id(self.target), self._control_part)) - if existing is not None or not self._native_toggle.pressed(world): + if existing is not None: + return + start_enabled = self._native_start_pending + self._native_start_pending = False + if not self._native_toggle.pressed(world) and not start_enabled: return self._native_controller, self._native_input_controller = ( create_robot_ik_gizmo_controller( @@ -582,8 +597,8 @@ def update_native(self, world: dexsim.World) -> None: ) ) self._native_window = window - # The activation press created a visible target. Consume it so the - # controller's own toggle does not immediately hide that target. + # Creation starts with a visible target. Consume any held toggle + # so the controller does not immediately hide the new target. self._native_controller.toggle.pressed(world) return if self._native_window is not window: diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 10a784d54..581428c84 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -151,6 +151,8 @@ class SimulationManagerCfg: headless and multi-environment simulations do not register controls. Defaults to DexSim IK. Use ``GizmoCfg(ik_solver="embodichain")`` to reuse configured solvers such as Pink, or ``None`` to disable automatic setup. + ``GizmoCfg(ik_start_enabled=True)`` explicitly activates native controls on + their first update with an open window, without waiting for I. """ render_cfg: RenderCfg = field(default_factory=RenderCfg) diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index c041e335a..872ddc72d 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -23,6 +23,7 @@ import argparse from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.objects import GizmoCfg from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.cfg import ( @@ -56,6 +57,7 @@ def main(): sim_device=args.device, render_cfg=RenderCfg(renderer=args.renderer), visualization=visualization_cfg_from_args(args), + robot_ik_gizmo=GizmoCfg(ik_start_enabled=True), ) sim = SimulationManager(sim_cfg) @@ -116,7 +118,9 @@ def main(): if native_window_opened or args.viser: logger.log_info("Use the gizmo to drag the robot end-effector (EE)") if native_window_opened: - logger.log_info("Press I to activate or toggle the native robot IK Gizmo") + logger.log_info( + "Native robot IK Gizmo starts enabled; press I to show or hide it" + ) logger.log_info("Press Ctrl+C to stop the simulation") run_simulation(sim) diff --git a/tests/gym/utils/test_gym_utils.py b/tests/gym/utils/test_gym_utils.py index 7ea9a26cb..a5be09d5b 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -1408,7 +1408,9 @@ def test_gym_config_preserves_entity_gizmo_startup_preference( assert cfg.sim_cfg.enable_entity_gizmo is (enabled is not False) @pytest.mark.parametrize("suffix", ["yaml", "json"]) - @pytest.mark.parametrize("settings", [None, {}, {"ik_solver": "embodichain"}]) + @pytest.mark.parametrize( + "settings", [None, {}, {"ik_solver": "embodichain"}, {"ik_start_enabled": True}] + ) def test_gym_config_parses_automatic_robot_gizmo_settings( self, tmp_path: Path, suffix: str, settings: dict | None ) -> None: @@ -1430,6 +1432,9 @@ def test_gym_config_parses_automatic_robot_gizmo_settings( assert cfg.sim_cfg.robot_ik_gizmo.ik_solver == settings.get( "ik_solver", "dexsim" ) + assert cfg.sim_cfg.robot_ik_gizmo.ik_start_enabled is settings.get( + "ik_start_enabled", False + ) def test_yaml_gym_config_parses_to_cfg(self, tmp_path): config = { diff --git a/tests/sim/test_robot_gizmo_lifecycle.py b/tests/sim/test_robot_gizmo_lifecycle.py index 2c45944dd..706e8365e 100644 --- a/tests/sim/test_robot_gizmo_lifecycle.py +++ b/tests/sim/test_robot_gizmo_lifecycle.py @@ -231,6 +231,47 @@ def test_explicit_disable_and_solver_override_take_precedence(managed_robot) -> assert set(sim.list_gizmos()) == {"robot:left", "robot:right"} +def test_start_enabled_waits_for_window_and_preserves_later_visibility( + managed_robot, +) -> None: + sim, _, factory = managed_robot + sim.sim_config.robot_ik_gizmo.ik_start_enabled = True + sim.close_window() + sim.update_gizmos() + factory.assert_not_called() + sim.open_window() + sim.update_gizmos() # No I press is needed. + controls = [g._native_controller for _, g in sim.get_gizmo_items()] + assert factory.call_count == 2 + assert all(control.enabled for control in controls) + sim._world.key_down = True + sim.update_gizmos() + sim.update_gizmos() + assert not any(control.enabled for control in controls) + sim.close_window() + sim.open_window() + sim.update_gizmos() + assert [g._native_controller for _, g in sim.get_gizmo_items()] == controls + assert not any(control.enabled for control in controls) + assert factory.call_count == 2 + + +def test_start_enabled_failure_waits_for_a_key_before_retry(managed_robot) -> None: + sim, _, factory = managed_robot + sim.sim_config.robot_ik_gizmo.ik_start_enabled = True + create = factory.side_effect + factory.side_effect = ValueError("invalid chain") + sim.update_gizmos() + sim.update_gizmos() + assert factory.call_count == 2 + assert not sim._window.controls + factory.side_effect = create + sim._world.key_down = True + sim.update_gizmos() + assert factory.call_count == 4 + assert all(g._native_controller.enabled for _, g in sim.get_gizmo_items()) + + def test_whole_robot_disable_cleans_activated_controls(managed_robot) -> None: sim, _, factory = managed_robot sim._world.key_down = True @@ -289,10 +330,12 @@ def test_robot_removal_cleans_native_controls_and_allows_same_uid_again( assert set(sim.list_gizmos()) == {"robot:left", "robot:right"} +@pytest.mark.parametrize("start_enabled", [False, True]) def test_explicit_native_factory_controller_is_not_duplicated( - managed_robot, monkeypatch: pytest.MonkeyPatch + managed_robot, monkeypatch: pytest.MonkeyPatch, start_enabled: bool ) -> None: sim, robot, factory = managed_robot + sim.sim_config.robot_ik_gizmo.ik_start_enabled = start_enabled explicit = Mock() monkeypatch.setitem( gizmo_module._NATIVE_IK_CONTROLLERS, (id(robot), "left"), explicit From 5f53b499b0aa597e25ff78392cd2f978ee85bca9 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 5 Sep 2026 19:12:13 +0800 Subject: [PATCH 11/14] fix(sim): resolve camera attachment through articulation render nodes --- agent_context/MAP.yaml | 5 + .../topics/sensor-system/sensor-system.md | 29 +++- .../simulation-system/simulation-system.md | 7 + .../embodichain.lab.sim.sensors.rst | 19 +++ embodichain/lab/sim/objects/articulation.py | 37 +++++ embodichain/lab/sim/sensors/attachment.py | 83 +++++++++++ embodichain/lab/sim/sensors/camera.py | 47 ++++--- embodichain/lab/sim/sensors/stereo.py | 2 - embodichain/lab/sim/sim_manager.py | 17 ++- tests/sim/objects/test_articulation.py | 62 ++++++++- tests/sim/sensors/test_attachment.py | 107 +++++++++++++++ tests/sim/sensors/test_camera.py | 129 +++++++++++++++++- tests/sim/test_sim_manager.py | 68 +++++++++ 13 files changed, 587 insertions(+), 25 deletions(-) create mode 100644 embodichain/lab/sim/sensors/attachment.py create mode 100644 tests/sim/sensors/test_attachment.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 43d19b42b..06bb3abd5 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -37,6 +37,7 @@ topics: - simulation lifecycle - ArticulationJointKinematics - get_parent_joint_chain + - get_link_render_nodes - ArticulationCfg - enable_gravity paths: @@ -298,13 +299,17 @@ topics: - rgb - depth - pointcloud + - resolve_parent_nodes paths: - topics/sensor-system/sensor-system.md source_of_truth: - embodichain/lab/sim/sensors/base_sensor.py - embodichain/lab/sim/sensors/camera.py - embodichain/lab/sim/sensors/stereo.py + - embodichain/lab/sim/sensors/attachment.py - embodichain/lab/sim/sensors/contact_sensor.py + - embodichain/lab/sim/sim_manager.py + - embodichain/lab/sim/objects/articulation.py - embodichain/lab/gym/utils/gym_utils.py - embodichain/lab/gym/utils/_component_composition.py - embodichain/lab/task_program/integrations/_configured_composition.py diff --git a/agent_context/topics/sensor-system/sensor-system.md b/agent_context/topics/sensor-system/sensor-system.md index 19bbe5f09..b08a5fa33 100644 --- a/agent_context/topics/sensor-system/sensor-system.md +++ b/agent_context/topics/sensor-system/sensor-system.md @@ -9,6 +9,9 @@ | Camera | `embodichain/lab/sim/sensors/camera.py` → `Camera`, `CameraCfg` | | Stereo camera | `embodichain/lab/sim/sensors/stereo.py` → `StereoCamera`, `StereoCameraCfg` | | Contact sensor | `embodichain/lab/sim/sensors/contact_sensor.py` → `ContactSensor`, `ContactSensorCfg` | +| Sensor creation and attachment coordination | `embodichain/lab/sim/sim_manager.py` → `SimulationManager.add_sensor()` | +| Camera parent resolution | `embodichain/lab/sim/sensors/attachment.py` → `resolve_parent_nodes()` | +| Native link render nodes | `embodichain/lab/sim/objects/articulation.py` → `Articulation.get_link_render_nodes()` | ## Overview @@ -103,6 +106,30 @@ Extends `SensorCfg.OffsetCfg` with look-at support: When `eye` is provided, the transformation is computed via `look_at_to_pose()`. Otherwise falls back to `pos`/`quat`. +### Camera attachment + +- `SimulationManager.add_sensor()` passes the registered Robots/Articulations and + expected environment count to `sensors.attachment.resolve_parent_nodes()` before + allocating camera views. The manager only coordinates creation and attachment. +- The resolver owns `extrinsics.parent` parsing, link-name disambiguation, and + instance-count validation. It uses public asset queries, not a manager singleton + or native handles. A plain canonical link name remains valid; use + `"/"` to disambiguate shared names. +- `Articulation.get_link_render_nodes()` encapsulates per-arena topology checks and + `get_render_body(link_name).render_node()`; Robot inherits this query. Do not + access `asset._entities` from the resolver, use global `Env.find_node()`, or + infer backend clone suffixes such as `.0` and `.1`. +- `Camera.attach_to_parent_nodes()` attaches one resolved node per camera instance, + reapplies parent-relative extrinsics, and then sets `is_attached` to `True`. + Stereo cameras use the same method, forwarding attachment to both views. +- Directly constructed cameras require an explicit `attach_to_parent_nodes()` + call. With `parent=None`, cameras added through the manager remain in arena space. +- Focused validation: `tests/sim/sensors/test_attachment.py` for resolution, + `tests/sim/objects/test_articulation.py` for native queries, + `tests/sim/sensors/test_camera.py` for attachment, and + `tests/sim/test_sim_manager.py` for coordination. Pure logic tests do not + initialize a renderer. + ### StereoCameraCfg Extends `CameraCfg` with stereo-specific fields: @@ -131,7 +158,7 @@ Properties `left_to_right` and `right_to_left` return `4×4` transform tensors. - **`sensor_type` string mismatch** — `SensorCfg.from_dict()` looks up `sensor_type + "Cfg"` in the sensors module. A typo (e.g. `"camera"` instead of `"Camera"`) causes `AttributeError`. - **Depth not enabled** — `enable_depth` defaults to `False`. Accessing depth data without enabling it returns empty tensors. -- **Parent frame not found** — `OffsetCfg.parent` must exactly match a link name in the scene. A wrong name silently places the sensor at the arena origin. +- **Invalid camera parent** — Missing or ambiguous registered links raise `ValueError`; missing per-arena links or render nodes raise `RuntimeError`. Parent nodes must resolve in every arena before attachment. - **Stereo baseline sign** — `left_to_right_pos` defines translation from left to right camera. Flipping the sign inverts the disparity. - **Contact sensor buffer overflow** — `max_contacts_per_env` caps the contact count. Exceeding it silently drops contacts; increase if the scene has dense collisions. - **View attribute flags** — `Camera.get_view_attrib()` computes `dr.ViewFlags` from enabled booleans. Adding a new data type requires both the `enable_*` flag and the corresponding `ViewFlags` bit. diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index d2dd7b620..62ac78b0e 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -111,6 +111,13 @@ for integrations that need link ancestry. It returns immediate-parent-first origin, axis, and optional limits. Consumers must not reach into `BatchEntity._entities` or retain backend-native joint-info objects. +`Articulation.get_link_render_nodes(link_name)` is the explicit render-attachment +query, inherited by Robot. It validates every instance and returns live render +nodes in environment order; those nodes must not be used after asset destruction. +`sensors.attachment.resolve_parent_nodes()` owns camera parent-name resolution +using public asset queries. `SimulationManager.add_sensor()` only supplies the +asset registry and environment count, then coordinates creation and attachment. + `Articulation` also exposes deterministic link meshes through `get_link_vert_face()` and named-state FK through `compute_fk()` with `qpos_joint_names`. Stochastic surface sampling and Atomic Action geometry keys diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.sensors.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.sensors.rst index 06829c641..e93f1dd6e 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.sensors.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.sensors.rst @@ -54,6 +54,25 @@ Camera :show-inheritance: :exclude-members: __init__, copy, replace, to_dict, validate +Attachment resolution +--------------------- + +``SimulationManager.add_sensor()`` delegates parent-name resolution to the +function below before creating camera views. It accepts a canonical link name +or ``"/"`` to distinguish links shared by several assets. +The resolver receives the scene asset mapping explicitly and queries +``Articulation.get_link_render_nodes()``; it does not look up a global manager. + +``Camera.attach_to_parent_nodes()`` then attaches the resolved nodes, reapplies +parent-relative extrinsics, and updates ``is_attached``. Stereo cameras share +this path. Directly constructed cameras require an explicit attachment call. + +.. autosummary:: + + ~attachment.resolve_parent_nodes + +.. autofunction:: embodichain.lab.sim.sensors.attachment.resolve_parent_nodes + Stereo Camera ------------- .. autoclass:: StereoCamera diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index d92825163..88f70be7c 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -1201,6 +1201,43 @@ def get_link_vert_face(self, link_name: str) -> Tuple[torch.Tensor, torch.Tensor verts, faces = self.body_data.link_vert_face[link_name] return verts, faces + def get_link_render_nodes(self, link_name: str) -> list[dexsim.engine.Node]: + """Get a link's native render node for every articulation instance. + + This is the render-attachment boundary for cameras and other scene + integrations. Returned nodes belong to this articulation and must not + be used after the owning asset is destroyed. + + Args: + link_name: Canonical link name, without backend clone suffixes. + + Returns: + Render nodes ordered by environment index. + + Raises: + ValueError: If the link is not part of this articulation. + RuntimeError: If any instance is missing the link or its render node. + """ + if link_name not in self.link_names: + raise ValueError(f"Articulation {self.uid!r} has no link {link_name!r}.") + + nodes: list[dexsim.engine.Node] = [] + for env_idx, entity in enumerate(self._entities): + if link_name not in entity.get_link_names(): + raise RuntimeError( + f"Articulation {self.uid!r} is missing link " + f"{link_name!r} in arena {env_idx}." + ) + render_body = entity.get_render_body(link_name) + node = None if render_body is None else render_body.render_node() + if node is None: + raise RuntimeError( + f"Articulation {self.uid!r} link {link_name!r} has " + f"no render node in arena {env_idx}." + ) + nodes.append(node) + return nodes + def get_link_pose( self, link_name: str, env_ids: Sequence[int] | None = None, to_matrix=False ) -> torch.Tensor: diff --git a/embodichain/lab/sim/sensors/attachment.py b/embodichain/lab/sim/sensors/attachment.py new file mode 100644 index 000000000..6457d2c90 --- /dev/null +++ b/embodichain/lab/sim/sensors/attachment.py @@ -0,0 +1,83 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Resolve camera attachment targets from an explicit scene asset mapping.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from dexsim.engine import Node + + from embodichain.lab.sim.objects import Articulation + +__all__ = ["resolve_parent_nodes"] + + +def resolve_parent_nodes( + parent: str, assets: Mapping[str, Articulation], num_envs: int +) -> list[Node]: + """Resolve a canonical Robot or Articulation link in each environment. + + Args: + parent: A canonical link name, or ``"/"`` to + disambiguate links shared by multiple assets. + assets: Scene asset UIDs mapped to Articulation or Robot instances. + num_envs: Expected number of camera and asset instances. + + Returns: + Parent render nodes ordered by environment index. + + Raises: + ValueError: If the parent link is missing or ambiguous. + RuntimeError: If the asset count differs from ``num_envs``, or an + environment is missing the link or its render node. + """ + asset_uid: str | None = None + link_name = parent + if "/" in parent: + candidate_uid, candidate_link = parent.split("/", maxsplit=1) + if candidate_uid in assets: + asset_uid, link_name = candidate_uid, candidate_link + + matches: list[tuple[str, list[Node]]] = [] + for uid, asset in assets.items(): + if asset_uid is not None and uid != asset_uid: + continue + if link_name not in asset.link_names: + continue + if asset.num_instances != num_envs: + raise RuntimeError( + f"Camera parent asset {uid!r} has {asset.num_instances} instances " + f"for {num_envs} arenas." + ) + matches.append((uid, asset.get_link_render_nodes(link_name))) + + if len(matches) == 1: + return matches[0][1] + if len(matches) > 1: + owners = ", ".join(uid for uid, _ in matches) + raise ValueError( + f"Camera parent link {link_name!r} is ambiguous across assets " + f"[{owners}]; use '/{link_name}'." + ) + scope = f" on asset {asset_uid!r}" if asset_uid is not None else "" + raise ValueError( + f"Camera parent link {link_name!r} was not found{scope} in any " + "registered Robot or Articulation." + ) diff --git a/embodichain/lab/sim/sensors/camera.py b/embodichain/lab/sim/sensors/camera.py index cc9a7aa44..c65b30c28 100644 --- a/embodichain/lab/sim/sensors/camera.py +++ b/embodichain/lab/sim/sensors/camera.py @@ -41,6 +41,9 @@ class ExtrinsicsCfg(SensorCfg.OffsetCfg): The extrinsics define the position and orientation of the camera in the 3D world. If eye, target, and up are provided, they will be used to compute the extrinsics. Otherwise, the position and orientation will be set to the defaults. + + SimulationManager resolves ``parent`` as a Robot or Articulation link + name. Use ``"/"`` when the link name is ambiguous. """ eye: Tuple[float, float, float] | None = None @@ -136,6 +139,7 @@ class Camera(BaseSensor): def __init__( self, config: CameraCfg, device: torch.device = torch.device("cpu") ) -> None: + self._is_attached = False super().__init__(config, device) def _build_sensor_from_config( @@ -202,8 +206,6 @@ def _build_sensor_from_config( ) self.cfg: CameraCfg = config - if self.cfg.extrinsics.parent is not None: - self._attach_to_entity() @cached_property def group_id(self) -> int: @@ -221,7 +223,7 @@ def is_attached(self) -> bool: Returns: bool: True if the camera is attached to a parent entity, False otherwise. """ - return self.cfg.extrinsics.parent is not None + return self._is_attached def update(self, **kwargs) -> None: """Update the sensor data. @@ -268,22 +270,35 @@ def update(self, **kwargs) -> None: self._frame_buffer.get_position_gpu_buffer().to(self.device)[..., :3] ) - def _attach_to_entity(self) -> None: - """Attach the sensor to the parent entity in each environment.""" - env = self._world.get_env() - for i, entity in enumerate(self._entities): + def attach_to_parent_nodes( + self, parent_nodes: Sequence[dexsim.engine.Node] + ) -> None: + """Attach camera views to one resolved parent node per environment. - parent = None - if i == 0: - parent = env.find_node(f"{self.cfg.extrinsics.parent}") - else: - parent = env.find_node(f"{self.cfg.extrinsics.parent}.{i-1}") - if parent is None: - logger.log_error( - f"Failed to find parent entity {self.cfg.extrinsics.parent} for sensor {self.cfg.uid}." - ) + SimulationManager calls this after resolving ``extrinsics.parent``. + Cameras constructed directly can be attached by supplying their parent + render nodes explicitly. + + Args: + parent_nodes: Parent render nodes ordered by environment index. + Raises: + RuntimeError: If the parent count differs from the camera count. + ValueError: If any parent node is missing. + """ + nodes = list(parent_nodes) + if len(nodes) != self.num_instances: + raise RuntimeError( + f"Camera attachment received {len(nodes)} parent nodes for " + f"{self.num_instances} camera instances." + ) + if any(node is None for node in nodes): + raise ValueError("Camera attachment requires a parent node in every arena.") + for entity, parent in zip(self._entities, nodes, strict=True): entity.attach_node(parent) + # Reapply parent-relative extrinsics after reparenting the camera views. + self.reset() + self._is_attached = True def set_local_pose( self, pose: torch.Tensor, env_ids: Sequence[int] | None = None diff --git a/embodichain/lab/sim/sensors/stereo.py b/embodichain/lab/sim/sensors/stereo.py index 999bedca9..2475a3f9e 100644 --- a/embodichain/lab/sim/sensors/stereo.py +++ b/embodichain/lab/sim/sensors/stereo.py @@ -277,8 +277,6 @@ def _build_sensor_from_config( ][:, :, config.width :, :] self.cfg: CameraCfg = config - if self.cfg.extrinsics.parent is not None: - self._attach_to_entity() def update(self, **kwargs) -> None: """Update the sensor data. diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 581428c84..76461bbdd 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -71,9 +71,11 @@ SensorCfg, BaseSensor, Camera, + CameraCfg, StereoCamera, ContactSensor, ) +from embodichain.lab.sim.sensors.attachment import resolve_parent_nodes from embodichain.lab.sim.cfg import ( RenderCfg, PhysicsCfg, @@ -2462,14 +2464,25 @@ def add_sensor(self, sensor_cfg: SensorCfg) -> BaseSensor: logger.log_warning(f"Sensor {sensor_uid} already exists.") return None + parent_nodes = None + if ( + isinstance(sensor_cfg, CameraCfg) + and sensor_cfg.extrinsics.parent is not None + ): + parent_nodes = resolve_parent_nodes( + parent=sensor_cfg.extrinsics.parent, + assets={**self._articulations, **self._robots}, + num_envs=self.num_envs, + ) + sensor = self.SUPPORTED_SENSOR_TYPES[sensor_type](sensor_cfg, self.device) + if isinstance(sensor, Camera) and parent_nodes is not None: + sensor.attach_to_parent_nodes(parent_nodes) self._sensors[sensor_uid] = sensor if isinstance(sensor, Camera): self.notify_visualization_topology_changed() - # Check if the sensor needs to change the parent frame. - return sensor def get_sensor(self, uid: str) -> BaseSensor | None: diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index 6b3288191..0471f8119 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -18,6 +18,7 @@ import os from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -27,7 +28,7 @@ SimulationManagerCfg, VisualMaterialCfg, ) -from embodichain.lab.sim.objects import Articulation, ArticulationJointKinematics +from embodichain.lab.sim.objects import Articulation, ArticulationJointKinematics, Robot from embodichain.lab.sim.cfg import ( ArticulationCfg, JointDrivePropertiesCfg, @@ -183,6 +184,65 @@ def test_get_parent_joint_chain_returns_backend_neutral_child_to_root_values(): assert chain[0].origin_pose[0, 3].item() == 0.0 +def _make_render_node_articulation( + asset_type: type[Articulation] = Articulation, num_envs: int = 2 +) -> tuple[Articulation, list[object]]: + asset = object.__new__(asset_type) + asset.uid = "arm" + asset._data = SimpleNamespace(link_names=["wrist"]) + nodes = [object() for _ in range(num_envs)] + asset._entities = [] + for node in nodes: + entity = MagicMock(spec=["get_link_names", "get_render_body"]) + entity.get_link_names.return_value = ["wrist"] + entity.get_render_body.return_value.render_node.return_value = node + asset._entities.append(entity) + return asset, nodes + + +@pytest.mark.no_sim +@pytest.mark.parametrize("asset_type", [Articulation, Robot]) +@pytest.mark.parametrize("num_envs", [1, 3]) +def test_get_link_render_nodes_uses_canonical_link_in_each_arena( + asset_type: type[Articulation], num_envs: int +) -> None: + asset, nodes = _make_render_node_articulation(asset_type, num_envs) + + assert asset.get_link_render_nodes("wrist") == nodes + for entity in asset._entities: + entity.get_render_body.assert_called_once_with("wrist") + + +@pytest.mark.no_sim +def test_get_link_render_nodes_rejects_unknown_link_before_native_access() -> None: + asset, _ = _make_render_node_articulation() + + with pytest.raises(ValueError, match="has no link 'missing'"): + asset.get_link_render_nodes("missing") + + for entity in asset._entities: + entity.get_render_body.assert_not_called() + + +@pytest.mark.no_sim +@pytest.mark.parametrize("failure", ["missing_link", "missing_body", "missing_node"]) +def test_get_link_render_nodes_rejects_incomplete_arena_topology(failure: str) -> None: + asset, _ = _make_render_node_articulation() + entity = asset._entities[1] + if failure == "missing_link": + entity.get_link_names.return_value = [] + error = "missing link 'wrist' in arena 1" + else: + if failure == "missing_body": + entity.get_render_body.return_value = None + else: + entity.get_render_body.return_value.render_node.return_value = None + error = "no render node in arena 1" + + with pytest.raises(RuntimeError, match=error): + asset.get_link_render_nodes("wrist") + + def _link_static_friction(art: Articulation, link_name: str, env_idx: int = 0) -> float: return art._entities[env_idx].get_physical_attr(link_name).static_friction diff --git a/tests/sim/sensors/test_attachment.py b/tests/sim/sensors/test_attachment.py new file mode 100644 index 000000000..823ecb200 --- /dev/null +++ b/tests/sim/sensors/test_attachment.py @@ -0,0 +1,107 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from embodichain.lab.sim.sensors.attachment import resolve_parent_nodes + +pytestmark = pytest.mark.no_sim + + +def _make_asset( + num_envs: int = 2, link_name: str = "wrist" +) -> tuple[SimpleNamespace, list[object]]: + """Only expose public asset queries, without manager or native handles.""" + nodes = [object() for _ in range(num_envs)] + return ( + SimpleNamespace( + link_names=[link_name], + num_instances=num_envs, + get_link_render_nodes=MagicMock(return_value=nodes), + ), + nodes, + ) + + +@pytest.mark.parametrize("num_envs", [1, 3]) +def test_resolve_parent_nodes_preserves_environment_order(num_envs: int) -> None: + asset, nodes = _make_asset(num_envs) + + assert resolve_parent_nodes("wrist", {"arm": asset}, num_envs) == nodes + asset.get_link_render_nodes.assert_called_once_with("wrist") + + +def test_resolve_parent_nodes_requires_asset_uid_for_ambiguous_links() -> None: + arm, arm_nodes = _make_asset() + tool, tool_nodes = _make_asset() + assets = {"arm": arm, "tool": tool} + + with pytest.raises(ValueError, match="ambiguous.*asset_uid"): + resolve_parent_nodes("wrist", assets, 2) + + arm.get_link_render_nodes.reset_mock() + tool.get_link_render_nodes.reset_mock() + assert resolve_parent_nodes("arm/wrist", assets, 2) == arm_nodes + arm.get_link_render_nodes.assert_called_once_with("wrist") + tool.get_link_render_nodes.assert_not_called() + assert resolve_parent_nodes("tool/wrist", assets, 2) == tool_nodes + + +def test_resolve_parent_nodes_preserves_slashes_in_canonical_link_names() -> None: + asset, nodes = _make_asset(link_name="tool/wrist") + + assert resolve_parent_nodes("tool/wrist", {"arm": asset}, 2) == nodes + assert resolve_parent_nodes("arm/tool/wrist", {"arm": asset}, 2) == nodes + + +@pytest.mark.parametrize("parent", ["missing", "arm/missing", "unknown/wrist"]) +def test_resolve_parent_nodes_rejects_unknown_links(parent: str) -> None: + asset, _ = _make_asset() + + with pytest.raises(ValueError, match="was not found"): + resolve_parent_nodes(parent, {"arm": asset}, 2) + + asset.get_link_render_nodes.assert_not_called() + + +def test_resolve_parent_nodes_rejects_empty_asset_mapping() -> None: + with pytest.raises(ValueError, match="was not found"): + resolve_parent_nodes("wrist", {}, 2) + + +def test_resolve_parent_nodes_rejects_instance_count_mismatch() -> None: + asset, _ = _make_asset(num_envs=1) + + with pytest.raises(RuntimeError, match="1 instances for 2 arenas"): + resolve_parent_nodes("wrist", {"arm": asset}, 2) + + asset.get_link_render_nodes.assert_not_called() + + +def test_resolve_parent_nodes_propagates_asset_query_failure() -> None: + asset, _ = _make_asset() + error = RuntimeError("missing render node in arena 1") + asset.get_link_render_nodes.side_effect = error + + with pytest.raises(RuntimeError) as raised: + resolve_parent_nodes("wrist", {"arm": asset}, 2) + + assert raised.value is error diff --git a/tests/sim/sensors/test_camera.py b/tests/sim/sensors/test_camera.py index f9d522e4b..9ac69af2a 100644 --- a/tests/sim/sensors/test_camera.py +++ b/tests/sim/sensors/test_camera.py @@ -20,10 +20,21 @@ import torch import os +from unittest.mock import MagicMock + +import numpy as np from tensordict import TensorDict from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.sensors import Camera, SensorCfg, CameraCfg +from embodichain.lab.sim.sensors import ( + BaseSensor, + Camera, + SensorCfg, + CameraCfg, + StereoCamera, + StereoCameraCfg, +) +from embodichain.lab.sim.sensors.stereo import PairCameraView from embodichain.lab.sim.objects import Articulation from embodichain.lab.sim.cfg import ArticulationCfg, RenderCfg from embodichain.data import get_data_path @@ -140,12 +151,29 @@ def test_attach_to_parent(self): self.art: Articulation = self.sim.add_articulation( cfg=ArticulationCfg.from_dict(cfg_dict) ) - # from IPython import embed; embed() self.camera: Camera = self.sim.add_sensor( sensor_cfg=CameraCfg( - uid="test", extrinsics=CameraCfg.ExtrinsicsCfg(parent="handle_xpos") + uid="test", + extrinsics=CameraCfg.ExtrinsicsCfg( + parent="handle_xpos", pos=(0.1, 0.2, 0.3) + ), ) ) + assert self.camera.is_attached + for view, articulation in zip( + self.camera._entities, self.art._entities, strict=True + ): + parent = articulation.get_render_body("handle_xpos").render_node() + assert ( + view.get_node().path_name().rsplit("/", maxsplit=1)[0] + == parent.path_name() + ) + expected_pose = self.camera.cfg.extrinsics.transformation.unsqueeze(0).repeat( + self.camera.num_instances, 1, 1 + ) + torch.testing.assert_close( + self.camera.get_local_pose(to_matrix=True).cpu(), expected_pose + ) def test_set_intrinsics(self): # Define new intrinsic parameters @@ -208,6 +236,101 @@ def test_camera_backend_smoke(sim_device, renderer): test.teardown_method() +@pytest.mark.no_sim +@pytest.mark.parametrize("stereo", [False, True]) +def test_camera_attachment_reapplies_parent_relative_extrinsics(stereo: bool) -> None: + """Every view attaches to its own arena node before resetting its local pose.""" + cfg_type = StereoCameraCfg if stereo else CameraCfg + camera_type = StereoCamera if stereo else Camera + camera = object.__new__(camera_type) + camera.cfg = cfg_type( + uid="wrist_camera", + extrinsics=CameraCfg.ExtrinsicsCfg(parent="wrist", pos=(0.1, 0.2, 0.3)), + ) + camera.num_instances = 2 + camera._is_attached = False + views = [ + [ + MagicMock(spec=["attach_node", "set_local_pose"]) + for _ in range(2 if stereo else 1) + ] + for _ in range(2) + ] + camera._entities = [ + PairCameraView(*pair, camera.cfg.left_to_right.numpy()) if stereo else pair[0] + for pair in views + ] + nodes = [object(), object()] + + assert not camera.is_attached + camera.attach_to_parent_nodes(nodes) + + assert camera.is_attached + for node, pair in zip(nodes, views, strict=True): + for index, view in enumerate(pair): + view.attach_node.assert_called_once_with(node) + assert [call[0] for call in view.method_calls] == [ + "attach_node", + "set_local_pose", + ] + expected_pose = camera.cfg.extrinsics.transformation.numpy().copy() + if stereo: + expected_pose[0, 3] += (-0.5 if index == 0 else 0.5) * 0.05 + np.testing.assert_allclose( + view.set_local_pose.call_args.args[0], expected_pose + ) + + +@pytest.mark.no_sim +@pytest.mark.parametrize("stereo", [False, True]) +def test_camera_parent_config_does_not_imply_attachment( + monkeypatch: pytest.MonkeyPatch, stereo: bool +) -> None: + """Attachment state reflects actual reparenting, not merely configuration.""" + monkeypatch.setattr( + BaseSensor, + "__init__", + lambda self, config, device: setattr(self, "cfg", config), + ) + cfg_type = StereoCameraCfg if stereo else CameraCfg + camera_type = StereoCamera if stereo else Camera + camera = camera_type(cfg_type(extrinsics=CameraCfg.ExtrinsicsCfg(parent="wrist"))) + assert not camera.is_attached + + +@pytest.mark.no_sim +@pytest.mark.parametrize("parent_count", [0, 1, 3]) +def test_camera_attachment_rejects_mismatched_parent_count(parent_count: int) -> None: + camera = object.__new__(Camera) + camera.num_instances = 2 + camera._entities = [MagicMock(spec=["attach_node"]) for _ in range(2)] + camera._is_attached = False + camera.reset = MagicMock() + + with pytest.raises(RuntimeError, match="parent nodes for 2 camera instances"): + camera.attach_to_parent_nodes([object() for _ in range(parent_count)]) + + assert not camera.is_attached + camera.reset.assert_not_called() + for view in camera._entities: + view.attach_node.assert_not_called() + + +@pytest.mark.no_sim +def test_camera_attachment_rejects_missing_node_before_reparenting() -> None: + camera = object.__new__(Camera) + camera.num_instances = 2 + camera._entities = [MagicMock(spec=["attach_node"]) for _ in range(2)] + camera._is_attached = False + + with pytest.raises(ValueError, match="parent node in every arena"): + camera.attach_to_parent_nodes([object(), None]) + + assert not camera.is_attached + for view in camera._entities: + view.attach_node.assert_not_called() + + if __name__ == "__main__": test = TestCameraHybridCUDA() test.setup_method() diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index 8ac236b1b..f390252dd 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -34,6 +34,7 @@ SimulationManagerCfg, _WindowRecordState, ) +from embodichain.lab.sim.sensors import Camera, CameraCfg, StereoCamera, StereoCameraCfg from embodichain.lab.visualization import ( GizmoCommand, PickCommand, @@ -1008,6 +1009,73 @@ def test_add_stereo_camera_marks_visualization_topology_dirty() -> None: assert sim._visualization_topology_revision == 3 +def _make_camera_parent_asset( + num_envs: int = 2, link_name: str = "wrist" +) -> tuple[SimpleNamespace, list[object]]: + """Expose only the public articulation API used by attachment resolution.""" + nodes = [object() for _ in range(num_envs)] + return ( + SimpleNamespace( + link_names=[link_name], + num_instances=num_envs, + get_link_render_nodes=MagicMock(return_value=nodes), + ), + nodes, + ) + + +def _make_camera_attachment_manager(num_envs: int = 2) -> SimulationManager: + sim = object.__new__(SimulationManager) + sim.num_envs = num_envs + sim.device = torch.device("cpu") + sim._robots = {} + sim._articulations = {} + sim._sensors = {} + sim._visualization_topology_revision = 0 + return sim + + +@pytest.mark.parametrize("registry", ["_robots", "_articulations"]) +@pytest.mark.parametrize("stereo", [False, True]) +@pytest.mark.parametrize("parent", [None, "wrist", "arm/wrist"]) +def test_add_camera_attaches_resolved_nodes_only_when_parent_is_configured( + registry: str, stereo: bool, parent: str | None +) -> None: + sim = _make_camera_attachment_manager() + asset, nodes = _make_camera_parent_asset() + getattr(sim, registry)["arm"] = asset + cfg_type = StereoCameraCfg if stereo else CameraCfg + camera_type = StereoCamera if stereo else Camera + cfg = cfg_type(uid="camera", extrinsics=CameraCfg.ExtrinsicsCfg(parent=parent)) + camera = object.__new__(camera_type) + camera.attach_to_parent_nodes = MagicMock() + sim.SUPPORTED_SENSOR_TYPES = {cfg.sensor_type: lambda cfg, device: camera} + + assert sim.add_sensor(cfg) is camera + assert sim._sensors["camera"] is camera + assert sim._visualization_topology_revision == 1 + if parent is None: + camera.attach_to_parent_nodes.assert_not_called() + asset.get_link_render_nodes.assert_not_called() + else: + camera.attach_to_parent_nodes.assert_called_once_with(nodes) + asset.get_link_render_nodes.assert_called_once_with("wrist") + + +def test_add_camera_validates_parent_before_allocating_views() -> None: + sim = _make_camera_attachment_manager() + factory = MagicMock() + sim.SUPPORTED_SENSOR_TYPES = {"Camera": factory} + cfg = CameraCfg(uid="camera", extrinsics=CameraCfg.ExtrinsicsCfg(parent="missing")) + + with pytest.raises(ValueError, match="was not found"): + sim.add_sensor(cfg) + + factory.assert_not_called() + assert sim._sensors == {} + assert sim._visualization_topology_revision == 0 + + def test_window_camera_pose_to_look_at_uses_dexsim_world_up() -> None: """Captured look-at snippets preserve DexSim's default Z-up controls.""" pose = np.eye(4, dtype=np.float32) From 95689dafb23e9f93c7e3c1426635eb63b82c9e0d Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 5 Sep 2026 19:45:55 +0800 Subject: [PATCH 12/14] wip --- examples/sim/gizmo/gizmo_camera.py | 20 +- examples/sim/gizmo/gizmo_object.py | 9 +- examples/sim/gizmo/gizmo_robot.py | 22 ++- examples/sim/gizmo/gizmo_scene.py | 294 ----------------------------- examples/sim/gizmo/gizmo_w1.py | 233 ----------------------- 5 files changed, 32 insertions(+), 546 deletions(-) delete mode 100644 examples/sim/gizmo/gizmo_scene.py delete mode 100644 examples/sim/gizmo/gizmo_w1.py diff --git a/examples/sim/gizmo/gizmo_camera.py b/examples/sim/gizmo/gizmo_camera.py index c96352365..795ff3379 100644 --- a/examples/sim/gizmo/gizmo_camera.py +++ b/examples/sim/gizmo/gizmo_camera.py @@ -60,7 +60,7 @@ def main(): # Create simulation context sim = SimulationManager(sim_cfg) - sim.set_manual_update(False) + sim.set_manual_update(True) # Add some objects to the scene for camera to observe for i in range(5): @@ -98,8 +98,9 @@ def main(): # Add camera to simulation camera = sim.add_sensor(sensor_cfg=camera_cfg) - # Wait for initialization - time.sleep(0.2) + if sim.is_use_gpu_physics: + sim.init_gpu_physics() + sim.update(step=1) native_window_opened = False if not args.headless: @@ -141,7 +142,7 @@ def run_simulation( ) -> None: """Run the simulation loop with gizmo updates.""" step_count = 0 - last_time = time.time() + last_time = time.perf_counter() last_step = 0 if show_camera_window: @@ -153,9 +154,9 @@ def run_simulation( try: while True: - # Update all gizmos managed by sim (including camera gizmo) - sim.update_gizmos() - sim.capture_visualization_safely() + frame_start = time.perf_counter() + # update() applies Gizmo commands before advancing one physics step. + sim.update(step=1) # Update camera to get latest sensor data camera.update() @@ -199,7 +200,7 @@ def run_simulation( # Print simulation statistics and camera info if step_count % 1000 == 0: - current_time = time.time() + current_time = time.perf_counter() elapsed = current_time - last_time fps = ( sim.num_envs * (step_count - last_step) / elapsed @@ -220,6 +221,9 @@ def run_simulation( last_time = current_time last_step = step_count + elapsed = time.perf_counter() - frame_start + time.sleep(max(0.0, sim.sim_config.physics_dt - elapsed)) + except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: diff --git a/examples/sim/gizmo/gizmo_object.py b/examples/sim/gizmo/gizmo_object.py index cff7a6492..4e315d29c 100644 --- a/examples/sim/gizmo/gizmo_object.py +++ b/examples/sim/gizmo/gizmo_object.py @@ -57,6 +57,7 @@ def main(): # Create the simulation instance sim = SimulationManager(sim_cfg) + sim.set_manual_update(True) # Add two cubes to the scene cube1: RigidObject = sim.add_rigid_object( @@ -134,9 +135,10 @@ def run_simulation(sim: SimulationManager): step_count = 0 gizmo_enabled = True try: - last_time = time.time() + last_time = time.perf_counter() last_step = 0 while True: + frame_start = time.perf_counter() sim.update(step=1) step_count += 1 @@ -153,7 +155,7 @@ def run_simulation(sim: SimulationManager): # Print FPS every second if step_count % 1000 == 0: - current_time = time.time() + current_time = time.perf_counter() elapsed = current_time - last_time fps = ( sim.num_envs * (step_count - last_step) / elapsed @@ -163,6 +165,9 @@ def run_simulation(sim: SimulationManager): logger.log_info(f"Simulation step: {step_count}, FPS: {fps:.2f}") last_time = current_time last_step = step_count + + elapsed = time.perf_counter() - frame_start + time.sleep(max(0.0, sim.sim_config.physics_dt - elapsed)) except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: diff --git a/examples/sim/gizmo/gizmo_robot.py b/examples/sim/gizmo/gizmo_robot.py index 75d47ed75..dadff2c2b 100644 --- a/examples/sim/gizmo/gizmo_robot.py +++ b/examples/sim/gizmo/gizmo_robot.py @@ -68,7 +68,7 @@ def main(): ) sim = SimulationManager(sim_cfg) - sim.set_manual_update(False) + sim.set_manual_update(True) # Get UR10 URDF path ur10_urdf_path = get_data_path("UniversalRobots/UR10/UR10.urdf") @@ -121,17 +121,19 @@ def main(): init_qpos=[0.0, -np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi / 2, 0.0, 0.0, 0.0], ) robot = sim.add_robot(cfg=robot_cfg) + if sim.is_use_gpu_physics: + sim.init_gpu_physics() # Set initial joint positions initial_qpos = torch.tensor( [[0.0, -np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi / 2, 0.0]], dtype=torch.float32, - device="cpu", + device=sim.device, ) joint_ids = robot.get_joint_ids("arm") + robot.set_qpos(qpos=initial_qpos, joint_ids=joint_ids, target=False) robot.set_qpos(qpos=initial_qpos, joint_ids=joint_ids) - - time.sleep(0.2) # Wait for a moment to ensure everything is set up + sim.update(step=1) native_window_opened = False if not args.headless: @@ -172,18 +174,17 @@ def main(): def run_simulation(sim: SimulationManager, native_control=None): step_count = 0 try: - last_time = time.time() + last_time = time.perf_counter() last_step = 0 while True: - time.sleep(0.033) # 30Hz + frame_start = time.perf_counter() if native_control is not None: native_control[0].update() - sim.update_gizmos() - sim.capture_visualization_safely() + sim.update(step=1) step_count += 1 if step_count % 100 == 0: - current_time = time.time() + current_time = time.perf_counter() elapsed = current_time - last_time fps = ( sim.num_envs * (step_count - last_step) / elapsed @@ -193,6 +194,9 @@ def run_simulation(sim: SimulationManager, native_control=None): logger.log_info(f"Simulation step: {step_count}, FPS: {fps:.2f}") last_time = current_time last_step = step_count + + elapsed = time.perf_counter() - frame_start + time.sleep(max(0.0, sim.sim_config.physics_dt - elapsed)) except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: diff --git a/examples/sim/gizmo/gizmo_scene.py b/examples/sim/gizmo/gizmo_scene.py deleted file mode 100644 index 8c3844f41..000000000 --- a/examples/sim/gizmo/gizmo_scene.py +++ /dev/null @@ -1,294 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- -""" -Gizmo Scene Example: Interactive scene with both robot and rigid object gizmos - -This example demonstrates how to create an interactive simulation scene with: -- A UR10 robot with gizmo control for end-effector manipulation -- A rigid object (cube) with gizmo control for direct manipulation -Both objects can be interactively controlled through their respective gizmos. -""" - -from __future__ import annotations - -import time -import torch -import numpy as np -import argparse -import cv2 - -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.cfg import ( - RenderCfg, - RobotCfg, - URDFCfg, - JointDrivePropertiesCfg, - RigidObjectCfg, - RigidBodyAttributesCfg, -) -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.shapes import CubeCfg -from embodichain.lab.sim.sensors import CameraCfg -from embodichain.lab.sim.solvers import PinkSolverCfg -from embodichain.lab.sim.objects import create_robot_ik_gizmo_controller -from embodichain.data import get_data_path -from embodichain.utils import logger - - -def main(): - """Main function to create and run the simulation scene.""" - - parser = argparse.ArgumentParser( - description="Create a simulation scene with SimulationManager" - ) - add_env_launcher_args_to_parser(parser) - args = parser.parse_args() - - # Configure the simulation - sim_cfg = SimulationManagerCfg( - width=1920, - height=1080, - headless=True, - physics_dt=1.0 / 100.0, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), - visualization=visualization_cfg_from_args(args), - ) - - sim = SimulationManager(sim_cfg) - sim.set_manual_update(False) - - # Get DexForce W1 URDF path - urdf_path = get_data_path("DexforceW1V021/DexforceW1_v02_1.urdf") - - # Create DexForce W1 robot - robot_cfg = RobotCfg( - uid="w1_gizmo_test", - urdf_cfg=URDFCfg( - components=[{"component_type": "humanoid", "urdf_path": urdf_path}] - ), - control_parts={"left_arm": ["LEFT_J[1-7]"], "right_arm": ["RIGHT_J[1-7]"]}, - solver_cfg={ - "left_arm": PinkSolverCfg( - urdf_path=urdf_path, - end_link_name="left_ee", - root_link_name="left_arm_base", - pos_eps=1e-2, - rot_eps=5e-2, - max_iterations=300, - dt=0.1, - ), - "right_arm": PinkSolverCfg( - urdf_path=urdf_path, - end_link_name="right_ee", - root_link_name="right_arm_base", - pos_eps=1e-2, - rot_eps=5e-2, - max_iterations=300, - dt=0.1, - ), - }, - drive_pros=JointDrivePropertiesCfg( - stiffness={"LEFT_J[1-7]": 1e4, "RIGHT_J[1-7]": 1e4}, - damping={"LEFT_J[1-7]": 1e3, "RIGHT_J[1-7]": 1e3}, - ), - ) - robot = sim.add_robot(cfg=robot_cfg) - - # Set initial joint positions for both arms - left_arm_qpos = torch.tensor( - [[0, -np.pi / 4, np.pi / 4, -np.pi / 2, 0.0, np.pi / 4, 0.0]], - dtype=torch.float32, - device="cpu", - ) - right_arm_qpos = torch.tensor( - [[0, np.pi / 4, -np.pi / 4, np.pi / 2, 0.0, -np.pi / 4, 0.0]], - dtype=torch.float32, - device="cpu", - ) - - left_joint_ids = robot.get_joint_ids("left_arm") - right_joint_ids = robot.get_joint_ids("right_arm") - - robot.set_qpos(qpos=left_arm_qpos, joint_ids=left_joint_ids) - robot.set_qpos(qpos=right_arm_qpos, joint_ids=right_joint_ids) - - # Create a rigid object (cube) positioned to the side of the robot - cube_cfg = RigidObjectCfg( - uid="interactive_cube", - shape=CubeCfg(size=[0.1, 0.1, 0.1]), - body_type="kinematic", - attrs=RigidBodyAttributesCfg( - mass=1.0, - dynamic_friction=0.5, - static_friction=0.5, - restitution=0.1, - ), - init_pos=[1.0, 0.0, 0.5], # Position to the side of the robot - ) - cube = sim.add_rigid_object(cube_cfg) - - camera_cfg = CameraCfg( - uid="scene_camera", - width=640, - height=480, - intrinsics=(320, 320, 320, 240), # fx, fy, cx, cy - near=0.1, - far=10.0, - enable_color=True, - enable_depth=True, - extrinsics=CameraCfg.ExtrinsicsCfg( - eye=(2.0, 2.0, 2.0), - target=(0.0, 0.0, 0.0), - up=(0.0, 0.0, 1.0), - ), - ) - camera = sim.add_sensor(sensor_cfg=camera_cfg) - - native_window_opened = False - if not args.headless: - native_window_opened = sim.open_window() - - native_controls = [] - if native_window_opened: - for control_part in ("left_arm", "right_arm"): - native_controls.append( - create_robot_ik_gizmo_controller( - robot, - control_part=control_part, - world=sim.get_world(), - ) - ) - elif args.viser: - sim.enable_gizmo( - uid="w1_gizmo_test", - control_part="left_arm", - ) - if not sim.has_gizmo("w1_gizmo_test", control_part="left_arm"): - logger.log_error("Failed to enable left arm gizmo!") - return - - sim.enable_gizmo( - uid="w1_gizmo_test", - control_part="right_arm", - ) - if not sim.has_gizmo("w1_gizmo_test", control_part="right_arm"): - logger.log_error("Failed to enable right arm gizmo!") - return - - sim.enable_gizmo( - uid="interactive_cube", - ) - if not sim.has_gizmo("interactive_cube"): - logger.log_error("Failed to enable gizmo for cube!") - return - - sim.enable_gizmo( - uid="scene_camera", - ) - if not sim.has_gizmo("scene_camera"): - logger.log_error("Failed to enable gizmo for camera!") - return - else: - logger.log_warning( - "Gizmo interaction is disabled in headless mode without Viser." - ) - - logger.log_info("Gizmo Scene example started!") - if args.viser: - logger.log_info("Four gizmos are active in the scene:") - logger.log_info( - "1. Left arm gizmo - Use to drag the left arm end-effector (EE)" - ) - logger.log_info( - "2. Right arm gizmo - Use to drag the right arm end-effector (EE)" - ) - logger.log_info("3. Cube gizmo - Use to drag and position the cube") - logger.log_info("4. Camera gizmo - Use to drag and orient the camera") - elif native_window_opened: - logger.log_info("Press I to show or hide each robot TCP IK Gizmo.") - logger.log_info("Select a scene entity and press G to manipulate its root.") - logger.log_info("Press Ctrl+C to stop the simulation") - - run_simulation( - sim, - native_controls=native_controls, - show_camera_window=native_window_opened, - ) - - -def run_simulation( - sim: SimulationManager, - *, - native_controls=(), - show_camera_window: bool, -) -> None: - step_count = 0 - # Get the camera instance by uid - camera = sim.get_sensor("scene_camera") - try: - last_time = time.time() - last_step = 0 - while True: - time.sleep(0.033) # 30Hz - for controller, _ in native_controls: - controller.update() - sim.update_gizmos() - sim.capture_visualization_safely() - step_count += 1 - - # Display camera view in a window every 5 steps - if show_camera_window and camera is not None and step_count % 5 == 0: - camera.update() - data = camera.get_data() - if "color" in data: - rgb_image = data["color"].cpu().numpy()[0, :, :, :3] - bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR) - cv2.putText( - bgr_image, - "Camera sensor preview", - (10, 30), - cv2.FONT_HERSHEY_SIMPLEX, - 0.6, - (0, 255, 0), - 2, - ) - cv2.imshow("Camera Sensor View", bgr_image) - cv2.waitKey(1) - - if step_count % 100 == 0: - current_time = time.time() - elapsed = current_time - last_time - fps = ( - sim.num_envs * (step_count - last_step) / elapsed - if elapsed > 0 - else 0 - ) - logger.log_info(f"Simulation step: {step_count}, FPS: {fps:.2f}") - last_time = current_time - last_step = step_count - except KeyboardInterrupt: - logger.log_info("\nStopping simulation...") - finally: - if show_camera_window: - cv2.destroyAllWindows() - sim.destroy() - logger.log_info("Simulation terminated successfully") - - -if __name__ == "__main__": - main() diff --git a/examples/sim/gizmo/gizmo_w1.py b/examples/sim/gizmo/gizmo_w1.py deleted file mode 100644 index 394585554..000000000 --- a/examples/sim/gizmo/gizmo_w1.py +++ /dev/null @@ -1,233 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- -""" -Gizmo-Robot Example: Test Gizmo class on a robot (UR10) -""" - -from __future__ import annotations - -import time -import torch -import numpy as np -import argparse - -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.cfg import ( - RenderCfg, - RobotCfg, - URDFCfg, - JointDrivePropertiesCfg, -) -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.solvers import PinkSolverCfg -from embodichain.data import get_data_path -from embodichain.utils import logger -from embodichain.lab.sim.robots.dexforce_w1.cfg import DexforceW1Cfg -from embodichain.lab.sim.objects import create_robot_ik_gizmo_controller - - -def main(): - """Main function to create and run the simulation scene.""" - - # Parse command line arguments - parser = argparse.ArgumentParser( - description="Create a simulation scene with SimulationManager" - ) - add_env_launcher_args_to_parser(parser) - args = parser.parse_args() - - # Configure the simulation - sim_cfg = SimulationManagerCfg( - width=1920, - height=1080, - headless=True, - physics_dt=1.0 / 100.0, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), - visualization=visualization_cfg_from_args(args), - ) - - sim = SimulationManager(sim_cfg) - sim.set_manual_update(False) - - cfg = DexforceW1Cfg.from_dict( - { - "uid": "w1_gizmo_test", - } - ) - cfg.solver_cfg["left_arm"].tcp = np.array( - [ - [1.0, 0.0, 0.0, 0.012], - [0.0, 1.0, 0.0, 0.04], - [0.0, 0.0, 1.0, 0.11], - [0.0, 0.0, 0.0, 1.0], - ] - ) - cfg.solver_cfg["right_arm"].tcp = np.array( - [ - [1.0, 0.0, 0.0, 0.012], - [0.0, 1.0, 0.0, -0.04], - [0.0, 0.0, 1.0, 0.11], - [0.0, 0.0, 0.0, 1.0], - ] - ) - - cfg.init_qpos = [ - 1.0000e00, - -2.0000e00, - 1.0000e00, - 0.0000e00, - -2.6921e-05, - -2.6514e-03, - -1.5708e00, - 1.4575e00, - -7.8540e-01, - 1.2834e-01, - 1.5708e00, - -2.2310e00, - -7.8540e-01, - 1.4461e00, - -1.5708e00, - 1.6716e00, - 7.8540e-01, - 7.6745e-01, - 0.0000e00, - 3.8108e-01, - 0.0000e00, - 0.0000e00, - 0.0000e00, - 0.0000e00, - 1.5000e00, - 0.0000e00, - 0.0000e00, - 0.0000e00, - 0.0000e00, - 1.5000e00, - 6.9974e-02, - 7.3950e-02, - 6.6574e-02, - 6.0923e-02, - 0.0000e00, - 6.7342e-02, - 7.0862e-02, - 6.3684e-02, - 5.7822e-02, - 0.0000e00, - ] - robot = sim.add_robot(cfg=cfg) - - # Set initial joint positions for both arms - # Left arm: 8 joints (WAIST + 7 LEFT_J), Right arm: 8 joints (WAIST + 7 RIGHT_J) - left_arm_qpos = torch.tensor( - [ - [0, 0, -np.pi / 4, np.pi / 4, -np.pi / 2, 0.0, np.pi / 4, 0.0] - ], # WAIST + LEFT_J[1-7] - dtype=torch.float32, - device="cpu", - ) - right_arm_qpos = torch.tensor( - [ - [0, 0, np.pi / 4, -np.pi / 4, np.pi / 2, 0.0, -np.pi / 4, 0.0] - ], # WAIST + RIGHT_J[1-7] - dtype=torch.float32, - device="cpu", - ) - - left_joint_ids = robot.get_joint_ids("left_arm") - right_joint_ids = robot.get_joint_ids("right_arm") - - robot.set_qpos(qpos=left_arm_qpos, joint_ids=left_joint_ids) - robot.set_qpos(qpos=right_arm_qpos, joint_ids=right_joint_ids) - - time.sleep(0.2) # Wait for a moment to ensure everything is set up - - native_window_opened = False - if not args.headless: - native_window_opened = sim.open_window() - - native_controls = [] - if native_window_opened: - for control_part in ("left_arm", "right_arm"): - native_controls.append( - create_robot_ik_gizmo_controller( - robot, - control_part=control_part, - world=sim.get_world(), - ) - ) - elif args.viser: - sim.enable_gizmo( - uid="w1_gizmo_test", - control_part="left_arm", - ) - if not sim.has_gizmo("w1_gizmo_test", control_part="left_arm"): - logger.log_error("Failed to enable left arm gizmo!") - return - - sim.enable_gizmo( - uid="w1_gizmo_test", - control_part="right_arm", - ) - if not sim.has_gizmo("w1_gizmo_test", control_part="right_arm"): - logger.log_error("Failed to enable right arm gizmo!") - return - else: - logger.log_warning( - "Gizmo interaction is disabled in headless mode without Viser." - ) - - logger.log_info("Gizmo-DexForce W1 example started!") - if native_window_opened or args.viser: - logger.log_info("Use the gizmos to drag both robot arms' end-effectors") - logger.log_info("Press Ctrl+C to stop the simulation") - - run_simulation(sim, native_controls) - - -def run_simulation(sim: SimulationManager, native_controls=()): - step_count = 0 - try: - last_time = time.time() - last_step = 0 - while True: - time.sleep(0.033) # 30Hz - for controller, _ in native_controls: - controller.update() - sim.update_gizmos() - sim.capture_visualization_safely() - step_count += 1 - - if step_count % 100 == 0: - current_time = time.time() - elapsed = current_time - last_time - fps = ( - sim.num_envs * (step_count - last_step) / elapsed - if elapsed > 0 - else 0 - ) - logger.log_info(f"Simulation step: {step_count}, FPS: {fps:.2f}") - last_time = current_time - last_step = step_count - except KeyboardInterrupt: - logger.log_info("\nStopping simulation...") - finally: - sim.destroy() - logger.log_info("Simulation terminated successfully") - - -if __name__ == "__main__": - main() From eeae8f809ceb4d4d8b47b159ae259ba4d5b07118 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 5 Sep 2026 20:00:09 +0800 Subject: [PATCH 13/14] revert(sim): remove misplaced sensor attachment change from gizmo PR Reverts commit 5f53b499b0aa597e25ff78392cd2f978ee85bca9. The sensor attachment update belongs to PR #586; preserve subsequent Gizmo changes on this branch. --- agent_context/MAP.yaml | 5 - .../topics/sensor-system/sensor-system.md | 29 +--- .../simulation-system/simulation-system.md | 7 - .../embodichain.lab.sim.sensors.rst | 19 --- embodichain/lab/sim/objects/articulation.py | 37 ----- embodichain/lab/sim/sensors/attachment.py | 83 ----------- embodichain/lab/sim/sensors/camera.py | 47 +++---- embodichain/lab/sim/sensors/stereo.py | 2 + embodichain/lab/sim/sim_manager.py | 17 +-- tests/sim/objects/test_articulation.py | 62 +-------- tests/sim/sensors/test_attachment.py | 107 --------------- tests/sim/sensors/test_camera.py | 129 +----------------- tests/sim/test_sim_manager.py | 68 --------- 13 files changed, 25 insertions(+), 587 deletions(-) delete mode 100644 embodichain/lab/sim/sensors/attachment.py delete mode 100644 tests/sim/sensors/test_attachment.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 06bb3abd5..43d19b42b 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -37,7 +37,6 @@ topics: - simulation lifecycle - ArticulationJointKinematics - get_parent_joint_chain - - get_link_render_nodes - ArticulationCfg - enable_gravity paths: @@ -299,17 +298,13 @@ topics: - rgb - depth - pointcloud - - resolve_parent_nodes paths: - topics/sensor-system/sensor-system.md source_of_truth: - embodichain/lab/sim/sensors/base_sensor.py - embodichain/lab/sim/sensors/camera.py - embodichain/lab/sim/sensors/stereo.py - - embodichain/lab/sim/sensors/attachment.py - embodichain/lab/sim/sensors/contact_sensor.py - - embodichain/lab/sim/sim_manager.py - - embodichain/lab/sim/objects/articulation.py - embodichain/lab/gym/utils/gym_utils.py - embodichain/lab/gym/utils/_component_composition.py - embodichain/lab/task_program/integrations/_configured_composition.py diff --git a/agent_context/topics/sensor-system/sensor-system.md b/agent_context/topics/sensor-system/sensor-system.md index b08a5fa33..19bbe5f09 100644 --- a/agent_context/topics/sensor-system/sensor-system.md +++ b/agent_context/topics/sensor-system/sensor-system.md @@ -9,9 +9,6 @@ | Camera | `embodichain/lab/sim/sensors/camera.py` → `Camera`, `CameraCfg` | | Stereo camera | `embodichain/lab/sim/sensors/stereo.py` → `StereoCamera`, `StereoCameraCfg` | | Contact sensor | `embodichain/lab/sim/sensors/contact_sensor.py` → `ContactSensor`, `ContactSensorCfg` | -| Sensor creation and attachment coordination | `embodichain/lab/sim/sim_manager.py` → `SimulationManager.add_sensor()` | -| Camera parent resolution | `embodichain/lab/sim/sensors/attachment.py` → `resolve_parent_nodes()` | -| Native link render nodes | `embodichain/lab/sim/objects/articulation.py` → `Articulation.get_link_render_nodes()` | ## Overview @@ -106,30 +103,6 @@ Extends `SensorCfg.OffsetCfg` with look-at support: When `eye` is provided, the transformation is computed via `look_at_to_pose()`. Otherwise falls back to `pos`/`quat`. -### Camera attachment - -- `SimulationManager.add_sensor()` passes the registered Robots/Articulations and - expected environment count to `sensors.attachment.resolve_parent_nodes()` before - allocating camera views. The manager only coordinates creation and attachment. -- The resolver owns `extrinsics.parent` parsing, link-name disambiguation, and - instance-count validation. It uses public asset queries, not a manager singleton - or native handles. A plain canonical link name remains valid; use - `"/"` to disambiguate shared names. -- `Articulation.get_link_render_nodes()` encapsulates per-arena topology checks and - `get_render_body(link_name).render_node()`; Robot inherits this query. Do not - access `asset._entities` from the resolver, use global `Env.find_node()`, or - infer backend clone suffixes such as `.0` and `.1`. -- `Camera.attach_to_parent_nodes()` attaches one resolved node per camera instance, - reapplies parent-relative extrinsics, and then sets `is_attached` to `True`. - Stereo cameras use the same method, forwarding attachment to both views. -- Directly constructed cameras require an explicit `attach_to_parent_nodes()` - call. With `parent=None`, cameras added through the manager remain in arena space. -- Focused validation: `tests/sim/sensors/test_attachment.py` for resolution, - `tests/sim/objects/test_articulation.py` for native queries, - `tests/sim/sensors/test_camera.py` for attachment, and - `tests/sim/test_sim_manager.py` for coordination. Pure logic tests do not - initialize a renderer. - ### StereoCameraCfg Extends `CameraCfg` with stereo-specific fields: @@ -158,7 +131,7 @@ Properties `left_to_right` and `right_to_left` return `4×4` transform tensors. - **`sensor_type` string mismatch** — `SensorCfg.from_dict()` looks up `sensor_type + "Cfg"` in the sensors module. A typo (e.g. `"camera"` instead of `"Camera"`) causes `AttributeError`. - **Depth not enabled** — `enable_depth` defaults to `False`. Accessing depth data without enabling it returns empty tensors. -- **Invalid camera parent** — Missing or ambiguous registered links raise `ValueError`; missing per-arena links or render nodes raise `RuntimeError`. Parent nodes must resolve in every arena before attachment. +- **Parent frame not found** — `OffsetCfg.parent` must exactly match a link name in the scene. A wrong name silently places the sensor at the arena origin. - **Stereo baseline sign** — `left_to_right_pos` defines translation from left to right camera. Flipping the sign inverts the disparity. - **Contact sensor buffer overflow** — `max_contacts_per_env` caps the contact count. Exceeding it silently drops contacts; increase if the scene has dense collisions. - **View attribute flags** — `Camera.get_view_attrib()` computes `dr.ViewFlags` from enabled booleans. Adding a new data type requires both the `enable_*` flag and the corresponding `ViewFlags` bit. diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index 62ac78b0e..d2dd7b620 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -111,13 +111,6 @@ for integrations that need link ancestry. It returns immediate-parent-first origin, axis, and optional limits. Consumers must not reach into `BatchEntity._entities` or retain backend-native joint-info objects. -`Articulation.get_link_render_nodes(link_name)` is the explicit render-attachment -query, inherited by Robot. It validates every instance and returns live render -nodes in environment order; those nodes must not be used after asset destruction. -`sensors.attachment.resolve_parent_nodes()` owns camera parent-name resolution -using public asset queries. `SimulationManager.add_sensor()` only supplies the -asset registry and environment count, then coordinates creation and attachment. - `Articulation` also exposes deterministic link meshes through `get_link_vert_face()` and named-state FK through `compute_fk()` with `qpos_joint_names`. Stochastic surface sampling and Atomic Action geometry keys diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.sensors.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.sensors.rst index e93f1dd6e..06829c641 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.sensors.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.sensors.rst @@ -54,25 +54,6 @@ Camera :show-inheritance: :exclude-members: __init__, copy, replace, to_dict, validate -Attachment resolution ---------------------- - -``SimulationManager.add_sensor()`` delegates parent-name resolution to the -function below before creating camera views. It accepts a canonical link name -or ``"/"`` to distinguish links shared by several assets. -The resolver receives the scene asset mapping explicitly and queries -``Articulation.get_link_render_nodes()``; it does not look up a global manager. - -``Camera.attach_to_parent_nodes()`` then attaches the resolved nodes, reapplies -parent-relative extrinsics, and updates ``is_attached``. Stereo cameras share -this path. Directly constructed cameras require an explicit attachment call. - -.. autosummary:: - - ~attachment.resolve_parent_nodes - -.. autofunction:: embodichain.lab.sim.sensors.attachment.resolve_parent_nodes - Stereo Camera ------------- .. autoclass:: StereoCamera diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index 88f70be7c..d92825163 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -1201,43 +1201,6 @@ def get_link_vert_face(self, link_name: str) -> Tuple[torch.Tensor, torch.Tensor verts, faces = self.body_data.link_vert_face[link_name] return verts, faces - def get_link_render_nodes(self, link_name: str) -> list[dexsim.engine.Node]: - """Get a link's native render node for every articulation instance. - - This is the render-attachment boundary for cameras and other scene - integrations. Returned nodes belong to this articulation and must not - be used after the owning asset is destroyed. - - Args: - link_name: Canonical link name, without backend clone suffixes. - - Returns: - Render nodes ordered by environment index. - - Raises: - ValueError: If the link is not part of this articulation. - RuntimeError: If any instance is missing the link or its render node. - """ - if link_name not in self.link_names: - raise ValueError(f"Articulation {self.uid!r} has no link {link_name!r}.") - - nodes: list[dexsim.engine.Node] = [] - for env_idx, entity in enumerate(self._entities): - if link_name not in entity.get_link_names(): - raise RuntimeError( - f"Articulation {self.uid!r} is missing link " - f"{link_name!r} in arena {env_idx}." - ) - render_body = entity.get_render_body(link_name) - node = None if render_body is None else render_body.render_node() - if node is None: - raise RuntimeError( - f"Articulation {self.uid!r} link {link_name!r} has " - f"no render node in arena {env_idx}." - ) - nodes.append(node) - return nodes - def get_link_pose( self, link_name: str, env_ids: Sequence[int] | None = None, to_matrix=False ) -> torch.Tensor: diff --git a/embodichain/lab/sim/sensors/attachment.py b/embodichain/lab/sim/sensors/attachment.py deleted file mode 100644 index 6457d2c90..000000000 --- a/embodichain/lab/sim/sensors/attachment.py +++ /dev/null @@ -1,83 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Resolve camera attachment targets from an explicit scene asset mapping.""" - -from __future__ import annotations - -from collections.abc import Mapping -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from dexsim.engine import Node - - from embodichain.lab.sim.objects import Articulation - -__all__ = ["resolve_parent_nodes"] - - -def resolve_parent_nodes( - parent: str, assets: Mapping[str, Articulation], num_envs: int -) -> list[Node]: - """Resolve a canonical Robot or Articulation link in each environment. - - Args: - parent: A canonical link name, or ``"/"`` to - disambiguate links shared by multiple assets. - assets: Scene asset UIDs mapped to Articulation or Robot instances. - num_envs: Expected number of camera and asset instances. - - Returns: - Parent render nodes ordered by environment index. - - Raises: - ValueError: If the parent link is missing or ambiguous. - RuntimeError: If the asset count differs from ``num_envs``, or an - environment is missing the link or its render node. - """ - asset_uid: str | None = None - link_name = parent - if "/" in parent: - candidate_uid, candidate_link = parent.split("/", maxsplit=1) - if candidate_uid in assets: - asset_uid, link_name = candidate_uid, candidate_link - - matches: list[tuple[str, list[Node]]] = [] - for uid, asset in assets.items(): - if asset_uid is not None and uid != asset_uid: - continue - if link_name not in asset.link_names: - continue - if asset.num_instances != num_envs: - raise RuntimeError( - f"Camera parent asset {uid!r} has {asset.num_instances} instances " - f"for {num_envs} arenas." - ) - matches.append((uid, asset.get_link_render_nodes(link_name))) - - if len(matches) == 1: - return matches[0][1] - if len(matches) > 1: - owners = ", ".join(uid for uid, _ in matches) - raise ValueError( - f"Camera parent link {link_name!r} is ambiguous across assets " - f"[{owners}]; use '/{link_name}'." - ) - scope = f" on asset {asset_uid!r}" if asset_uid is not None else "" - raise ValueError( - f"Camera parent link {link_name!r} was not found{scope} in any " - "registered Robot or Articulation." - ) diff --git a/embodichain/lab/sim/sensors/camera.py b/embodichain/lab/sim/sensors/camera.py index c65b30c28..cc9a7aa44 100644 --- a/embodichain/lab/sim/sensors/camera.py +++ b/embodichain/lab/sim/sensors/camera.py @@ -41,9 +41,6 @@ class ExtrinsicsCfg(SensorCfg.OffsetCfg): The extrinsics define the position and orientation of the camera in the 3D world. If eye, target, and up are provided, they will be used to compute the extrinsics. Otherwise, the position and orientation will be set to the defaults. - - SimulationManager resolves ``parent`` as a Robot or Articulation link - name. Use ``"/"`` when the link name is ambiguous. """ eye: Tuple[float, float, float] | None = None @@ -139,7 +136,6 @@ class Camera(BaseSensor): def __init__( self, config: CameraCfg, device: torch.device = torch.device("cpu") ) -> None: - self._is_attached = False super().__init__(config, device) def _build_sensor_from_config( @@ -206,6 +202,8 @@ def _build_sensor_from_config( ) self.cfg: CameraCfg = config + if self.cfg.extrinsics.parent is not None: + self._attach_to_entity() @cached_property def group_id(self) -> int: @@ -223,7 +221,7 @@ def is_attached(self) -> bool: Returns: bool: True if the camera is attached to a parent entity, False otherwise. """ - return self._is_attached + return self.cfg.extrinsics.parent is not None def update(self, **kwargs) -> None: """Update the sensor data. @@ -270,35 +268,22 @@ def update(self, **kwargs) -> None: self._frame_buffer.get_position_gpu_buffer().to(self.device)[..., :3] ) - def attach_to_parent_nodes( - self, parent_nodes: Sequence[dexsim.engine.Node] - ) -> None: - """Attach camera views to one resolved parent node per environment. - - SimulationManager calls this after resolving ``extrinsics.parent``. - Cameras constructed directly can be attached by supplying their parent - render nodes explicitly. + def _attach_to_entity(self) -> None: + """Attach the sensor to the parent entity in each environment.""" + env = self._world.get_env() + for i, entity in enumerate(self._entities): - Args: - parent_nodes: Parent render nodes ordered by environment index. + parent = None + if i == 0: + parent = env.find_node(f"{self.cfg.extrinsics.parent}") + else: + parent = env.find_node(f"{self.cfg.extrinsics.parent}.{i-1}") + if parent is None: + logger.log_error( + f"Failed to find parent entity {self.cfg.extrinsics.parent} for sensor {self.cfg.uid}." + ) - Raises: - RuntimeError: If the parent count differs from the camera count. - ValueError: If any parent node is missing. - """ - nodes = list(parent_nodes) - if len(nodes) != self.num_instances: - raise RuntimeError( - f"Camera attachment received {len(nodes)} parent nodes for " - f"{self.num_instances} camera instances." - ) - if any(node is None for node in nodes): - raise ValueError("Camera attachment requires a parent node in every arena.") - for entity, parent in zip(self._entities, nodes, strict=True): entity.attach_node(parent) - # Reapply parent-relative extrinsics after reparenting the camera views. - self.reset() - self._is_attached = True def set_local_pose( self, pose: torch.Tensor, env_ids: Sequence[int] | None = None diff --git a/embodichain/lab/sim/sensors/stereo.py b/embodichain/lab/sim/sensors/stereo.py index 2475a3f9e..999bedca9 100644 --- a/embodichain/lab/sim/sensors/stereo.py +++ b/embodichain/lab/sim/sensors/stereo.py @@ -277,6 +277,8 @@ def _build_sensor_from_config( ][:, :, config.width :, :] self.cfg: CameraCfg = config + if self.cfg.extrinsics.parent is not None: + self._attach_to_entity() def update(self, **kwargs) -> None: """Update the sensor data. diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 76461bbdd..581428c84 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -71,11 +71,9 @@ SensorCfg, BaseSensor, Camera, - CameraCfg, StereoCamera, ContactSensor, ) -from embodichain.lab.sim.sensors.attachment import resolve_parent_nodes from embodichain.lab.sim.cfg import ( RenderCfg, PhysicsCfg, @@ -2464,25 +2462,14 @@ def add_sensor(self, sensor_cfg: SensorCfg) -> BaseSensor: logger.log_warning(f"Sensor {sensor_uid} already exists.") return None - parent_nodes = None - if ( - isinstance(sensor_cfg, CameraCfg) - and sensor_cfg.extrinsics.parent is not None - ): - parent_nodes = resolve_parent_nodes( - parent=sensor_cfg.extrinsics.parent, - assets={**self._articulations, **self._robots}, - num_envs=self.num_envs, - ) - sensor = self.SUPPORTED_SENSOR_TYPES[sensor_type](sensor_cfg, self.device) - if isinstance(sensor, Camera) and parent_nodes is not None: - sensor.attach_to_parent_nodes(parent_nodes) self._sensors[sensor_uid] = sensor if isinstance(sensor, Camera): self.notify_visualization_topology_changed() + # Check if the sensor needs to change the parent frame. + return sensor def get_sensor(self, uid: str) -> BaseSensor | None: diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index 0471f8119..6b3288191 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -18,7 +18,6 @@ import os from types import SimpleNamespace -from unittest.mock import MagicMock import pytest import torch @@ -28,7 +27,7 @@ SimulationManagerCfg, VisualMaterialCfg, ) -from embodichain.lab.sim.objects import Articulation, ArticulationJointKinematics, Robot +from embodichain.lab.sim.objects import Articulation, ArticulationJointKinematics from embodichain.lab.sim.cfg import ( ArticulationCfg, JointDrivePropertiesCfg, @@ -184,65 +183,6 @@ def test_get_parent_joint_chain_returns_backend_neutral_child_to_root_values(): assert chain[0].origin_pose[0, 3].item() == 0.0 -def _make_render_node_articulation( - asset_type: type[Articulation] = Articulation, num_envs: int = 2 -) -> tuple[Articulation, list[object]]: - asset = object.__new__(asset_type) - asset.uid = "arm" - asset._data = SimpleNamespace(link_names=["wrist"]) - nodes = [object() for _ in range(num_envs)] - asset._entities = [] - for node in nodes: - entity = MagicMock(spec=["get_link_names", "get_render_body"]) - entity.get_link_names.return_value = ["wrist"] - entity.get_render_body.return_value.render_node.return_value = node - asset._entities.append(entity) - return asset, nodes - - -@pytest.mark.no_sim -@pytest.mark.parametrize("asset_type", [Articulation, Robot]) -@pytest.mark.parametrize("num_envs", [1, 3]) -def test_get_link_render_nodes_uses_canonical_link_in_each_arena( - asset_type: type[Articulation], num_envs: int -) -> None: - asset, nodes = _make_render_node_articulation(asset_type, num_envs) - - assert asset.get_link_render_nodes("wrist") == nodes - for entity in asset._entities: - entity.get_render_body.assert_called_once_with("wrist") - - -@pytest.mark.no_sim -def test_get_link_render_nodes_rejects_unknown_link_before_native_access() -> None: - asset, _ = _make_render_node_articulation() - - with pytest.raises(ValueError, match="has no link 'missing'"): - asset.get_link_render_nodes("missing") - - for entity in asset._entities: - entity.get_render_body.assert_not_called() - - -@pytest.mark.no_sim -@pytest.mark.parametrize("failure", ["missing_link", "missing_body", "missing_node"]) -def test_get_link_render_nodes_rejects_incomplete_arena_topology(failure: str) -> None: - asset, _ = _make_render_node_articulation() - entity = asset._entities[1] - if failure == "missing_link": - entity.get_link_names.return_value = [] - error = "missing link 'wrist' in arena 1" - else: - if failure == "missing_body": - entity.get_render_body.return_value = None - else: - entity.get_render_body.return_value.render_node.return_value = None - error = "no render node in arena 1" - - with pytest.raises(RuntimeError, match=error): - asset.get_link_render_nodes("wrist") - - def _link_static_friction(art: Articulation, link_name: str, env_idx: int = 0) -> float: return art._entities[env_idx].get_physical_attr(link_name).static_friction diff --git a/tests/sim/sensors/test_attachment.py b/tests/sim/sensors/test_attachment.py deleted file mode 100644 index 823ecb200..000000000 --- a/tests/sim/sensors/test_attachment.py +++ /dev/null @@ -1,107 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest - -from embodichain.lab.sim.sensors.attachment import resolve_parent_nodes - -pytestmark = pytest.mark.no_sim - - -def _make_asset( - num_envs: int = 2, link_name: str = "wrist" -) -> tuple[SimpleNamespace, list[object]]: - """Only expose public asset queries, without manager or native handles.""" - nodes = [object() for _ in range(num_envs)] - return ( - SimpleNamespace( - link_names=[link_name], - num_instances=num_envs, - get_link_render_nodes=MagicMock(return_value=nodes), - ), - nodes, - ) - - -@pytest.mark.parametrize("num_envs", [1, 3]) -def test_resolve_parent_nodes_preserves_environment_order(num_envs: int) -> None: - asset, nodes = _make_asset(num_envs) - - assert resolve_parent_nodes("wrist", {"arm": asset}, num_envs) == nodes - asset.get_link_render_nodes.assert_called_once_with("wrist") - - -def test_resolve_parent_nodes_requires_asset_uid_for_ambiguous_links() -> None: - arm, arm_nodes = _make_asset() - tool, tool_nodes = _make_asset() - assets = {"arm": arm, "tool": tool} - - with pytest.raises(ValueError, match="ambiguous.*asset_uid"): - resolve_parent_nodes("wrist", assets, 2) - - arm.get_link_render_nodes.reset_mock() - tool.get_link_render_nodes.reset_mock() - assert resolve_parent_nodes("arm/wrist", assets, 2) == arm_nodes - arm.get_link_render_nodes.assert_called_once_with("wrist") - tool.get_link_render_nodes.assert_not_called() - assert resolve_parent_nodes("tool/wrist", assets, 2) == tool_nodes - - -def test_resolve_parent_nodes_preserves_slashes_in_canonical_link_names() -> None: - asset, nodes = _make_asset(link_name="tool/wrist") - - assert resolve_parent_nodes("tool/wrist", {"arm": asset}, 2) == nodes - assert resolve_parent_nodes("arm/tool/wrist", {"arm": asset}, 2) == nodes - - -@pytest.mark.parametrize("parent", ["missing", "arm/missing", "unknown/wrist"]) -def test_resolve_parent_nodes_rejects_unknown_links(parent: str) -> None: - asset, _ = _make_asset() - - with pytest.raises(ValueError, match="was not found"): - resolve_parent_nodes(parent, {"arm": asset}, 2) - - asset.get_link_render_nodes.assert_not_called() - - -def test_resolve_parent_nodes_rejects_empty_asset_mapping() -> None: - with pytest.raises(ValueError, match="was not found"): - resolve_parent_nodes("wrist", {}, 2) - - -def test_resolve_parent_nodes_rejects_instance_count_mismatch() -> None: - asset, _ = _make_asset(num_envs=1) - - with pytest.raises(RuntimeError, match="1 instances for 2 arenas"): - resolve_parent_nodes("wrist", {"arm": asset}, 2) - - asset.get_link_render_nodes.assert_not_called() - - -def test_resolve_parent_nodes_propagates_asset_query_failure() -> None: - asset, _ = _make_asset() - error = RuntimeError("missing render node in arena 1") - asset.get_link_render_nodes.side_effect = error - - with pytest.raises(RuntimeError) as raised: - resolve_parent_nodes("wrist", {"arm": asset}, 2) - - assert raised.value is error diff --git a/tests/sim/sensors/test_camera.py b/tests/sim/sensors/test_camera.py index 9ac69af2a..f9d522e4b 100644 --- a/tests/sim/sensors/test_camera.py +++ b/tests/sim/sensors/test_camera.py @@ -20,21 +20,10 @@ import torch import os -from unittest.mock import MagicMock - -import numpy as np from tensordict import TensorDict from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.sensors import ( - BaseSensor, - Camera, - SensorCfg, - CameraCfg, - StereoCamera, - StereoCameraCfg, -) -from embodichain.lab.sim.sensors.stereo import PairCameraView +from embodichain.lab.sim.sensors import Camera, SensorCfg, CameraCfg from embodichain.lab.sim.objects import Articulation from embodichain.lab.sim.cfg import ArticulationCfg, RenderCfg from embodichain.data import get_data_path @@ -151,29 +140,12 @@ def test_attach_to_parent(self): self.art: Articulation = self.sim.add_articulation( cfg=ArticulationCfg.from_dict(cfg_dict) ) + # from IPython import embed; embed() self.camera: Camera = self.sim.add_sensor( sensor_cfg=CameraCfg( - uid="test", - extrinsics=CameraCfg.ExtrinsicsCfg( - parent="handle_xpos", pos=(0.1, 0.2, 0.3) - ), + uid="test", extrinsics=CameraCfg.ExtrinsicsCfg(parent="handle_xpos") ) ) - assert self.camera.is_attached - for view, articulation in zip( - self.camera._entities, self.art._entities, strict=True - ): - parent = articulation.get_render_body("handle_xpos").render_node() - assert ( - view.get_node().path_name().rsplit("/", maxsplit=1)[0] - == parent.path_name() - ) - expected_pose = self.camera.cfg.extrinsics.transformation.unsqueeze(0).repeat( - self.camera.num_instances, 1, 1 - ) - torch.testing.assert_close( - self.camera.get_local_pose(to_matrix=True).cpu(), expected_pose - ) def test_set_intrinsics(self): # Define new intrinsic parameters @@ -236,101 +208,6 @@ def test_camera_backend_smoke(sim_device, renderer): test.teardown_method() -@pytest.mark.no_sim -@pytest.mark.parametrize("stereo", [False, True]) -def test_camera_attachment_reapplies_parent_relative_extrinsics(stereo: bool) -> None: - """Every view attaches to its own arena node before resetting its local pose.""" - cfg_type = StereoCameraCfg if stereo else CameraCfg - camera_type = StereoCamera if stereo else Camera - camera = object.__new__(camera_type) - camera.cfg = cfg_type( - uid="wrist_camera", - extrinsics=CameraCfg.ExtrinsicsCfg(parent="wrist", pos=(0.1, 0.2, 0.3)), - ) - camera.num_instances = 2 - camera._is_attached = False - views = [ - [ - MagicMock(spec=["attach_node", "set_local_pose"]) - for _ in range(2 if stereo else 1) - ] - for _ in range(2) - ] - camera._entities = [ - PairCameraView(*pair, camera.cfg.left_to_right.numpy()) if stereo else pair[0] - for pair in views - ] - nodes = [object(), object()] - - assert not camera.is_attached - camera.attach_to_parent_nodes(nodes) - - assert camera.is_attached - for node, pair in zip(nodes, views, strict=True): - for index, view in enumerate(pair): - view.attach_node.assert_called_once_with(node) - assert [call[0] for call in view.method_calls] == [ - "attach_node", - "set_local_pose", - ] - expected_pose = camera.cfg.extrinsics.transformation.numpy().copy() - if stereo: - expected_pose[0, 3] += (-0.5 if index == 0 else 0.5) * 0.05 - np.testing.assert_allclose( - view.set_local_pose.call_args.args[0], expected_pose - ) - - -@pytest.mark.no_sim -@pytest.mark.parametrize("stereo", [False, True]) -def test_camera_parent_config_does_not_imply_attachment( - monkeypatch: pytest.MonkeyPatch, stereo: bool -) -> None: - """Attachment state reflects actual reparenting, not merely configuration.""" - monkeypatch.setattr( - BaseSensor, - "__init__", - lambda self, config, device: setattr(self, "cfg", config), - ) - cfg_type = StereoCameraCfg if stereo else CameraCfg - camera_type = StereoCamera if stereo else Camera - camera = camera_type(cfg_type(extrinsics=CameraCfg.ExtrinsicsCfg(parent="wrist"))) - assert not camera.is_attached - - -@pytest.mark.no_sim -@pytest.mark.parametrize("parent_count", [0, 1, 3]) -def test_camera_attachment_rejects_mismatched_parent_count(parent_count: int) -> None: - camera = object.__new__(Camera) - camera.num_instances = 2 - camera._entities = [MagicMock(spec=["attach_node"]) for _ in range(2)] - camera._is_attached = False - camera.reset = MagicMock() - - with pytest.raises(RuntimeError, match="parent nodes for 2 camera instances"): - camera.attach_to_parent_nodes([object() for _ in range(parent_count)]) - - assert not camera.is_attached - camera.reset.assert_not_called() - for view in camera._entities: - view.attach_node.assert_not_called() - - -@pytest.mark.no_sim -def test_camera_attachment_rejects_missing_node_before_reparenting() -> None: - camera = object.__new__(Camera) - camera.num_instances = 2 - camera._entities = [MagicMock(spec=["attach_node"]) for _ in range(2)] - camera._is_attached = False - - with pytest.raises(ValueError, match="parent node in every arena"): - camera.attach_to_parent_nodes([object(), None]) - - assert not camera.is_attached - for view in camera._entities: - view.attach_node.assert_not_called() - - if __name__ == "__main__": test = TestCameraHybridCUDA() test.setup_method() diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index f390252dd..8ac236b1b 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -34,7 +34,6 @@ SimulationManagerCfg, _WindowRecordState, ) -from embodichain.lab.sim.sensors import Camera, CameraCfg, StereoCamera, StereoCameraCfg from embodichain.lab.visualization import ( GizmoCommand, PickCommand, @@ -1009,73 +1008,6 @@ def test_add_stereo_camera_marks_visualization_topology_dirty() -> None: assert sim._visualization_topology_revision == 3 -def _make_camera_parent_asset( - num_envs: int = 2, link_name: str = "wrist" -) -> tuple[SimpleNamespace, list[object]]: - """Expose only the public articulation API used by attachment resolution.""" - nodes = [object() for _ in range(num_envs)] - return ( - SimpleNamespace( - link_names=[link_name], - num_instances=num_envs, - get_link_render_nodes=MagicMock(return_value=nodes), - ), - nodes, - ) - - -def _make_camera_attachment_manager(num_envs: int = 2) -> SimulationManager: - sim = object.__new__(SimulationManager) - sim.num_envs = num_envs - sim.device = torch.device("cpu") - sim._robots = {} - sim._articulations = {} - sim._sensors = {} - sim._visualization_topology_revision = 0 - return sim - - -@pytest.mark.parametrize("registry", ["_robots", "_articulations"]) -@pytest.mark.parametrize("stereo", [False, True]) -@pytest.mark.parametrize("parent", [None, "wrist", "arm/wrist"]) -def test_add_camera_attaches_resolved_nodes_only_when_parent_is_configured( - registry: str, stereo: bool, parent: str | None -) -> None: - sim = _make_camera_attachment_manager() - asset, nodes = _make_camera_parent_asset() - getattr(sim, registry)["arm"] = asset - cfg_type = StereoCameraCfg if stereo else CameraCfg - camera_type = StereoCamera if stereo else Camera - cfg = cfg_type(uid="camera", extrinsics=CameraCfg.ExtrinsicsCfg(parent=parent)) - camera = object.__new__(camera_type) - camera.attach_to_parent_nodes = MagicMock() - sim.SUPPORTED_SENSOR_TYPES = {cfg.sensor_type: lambda cfg, device: camera} - - assert sim.add_sensor(cfg) is camera - assert sim._sensors["camera"] is camera - assert sim._visualization_topology_revision == 1 - if parent is None: - camera.attach_to_parent_nodes.assert_not_called() - asset.get_link_render_nodes.assert_not_called() - else: - camera.attach_to_parent_nodes.assert_called_once_with(nodes) - asset.get_link_render_nodes.assert_called_once_with("wrist") - - -def test_add_camera_validates_parent_before_allocating_views() -> None: - sim = _make_camera_attachment_manager() - factory = MagicMock() - sim.SUPPORTED_SENSOR_TYPES = {"Camera": factory} - cfg = CameraCfg(uid="camera", extrinsics=CameraCfg.ExtrinsicsCfg(parent="missing")) - - with pytest.raises(ValueError, match="was not found"): - sim.add_sensor(cfg) - - factory.assert_not_called() - assert sim._sensors == {} - assert sim._visualization_topology_revision == 0 - - def test_window_camera_pose_to_look_at_uses_dexsim_world_up() -> None: """Captured look-at snippets preserve DexSim's default Z-up controls.""" pose = np.eye(4, dtype=np.float32) From 6674732c6e1acbdbeabf3fa6bfb1653dbd4eaf95 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 6 Sep 2026 00:21:11 +0800 Subject: [PATCH 14/14] test(sim): align gizmo stepping cases with current examples --- tests/sim/test_interactive_stepping.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/tests/sim/test_interactive_stepping.py b/tests/sim/test_interactive_stepping.py index f8da01e40..4ac383333 100644 --- a/tests/sim/test_interactive_stepping.py +++ b/tests/sim/test_interactive_stepping.py @@ -29,8 +29,7 @@ "module_name", [ "examples.sim.gizmo.gizmo_robot", - "examples.sim.gizmo.gizmo_w1", - "examples.sim.gizmo.gizmo_scene", + "examples.sim.gizmo.gizmo_object", "examples.sim.gizmo.gizmo_camera", "scripts.tutorials.sim.gizmo_robot", ], @@ -63,8 +62,8 @@ def sleep(duration: float) -> None: sim = SimpleNamespace( sim_config=SimpleNamespace(physics_dt=physics_dt), num_envs=1, + is_use_gpu_physics=False, update=update, - get_sensor=lambda _uid: None, has_gizmo=lambda _uid: False, destroy=Mock(), ) @@ -79,8 +78,6 @@ def sleep(duration: float) -> None: if module_name.endswith("gizmo_camera"): module.run_simulation(sim, Mock(), show_camera_window=False) - elif module_name.endswith("gizmo_scene"): - module.run_simulation(sim, show_camera_window=False) else: module.run_simulation(sim)