From 768ce4340ab24bee354bfb1fd4e5d2d3ca4d826d Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 27 Jul 2026 19:44:47 +0800 Subject: [PATCH 1/4] 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 2/4] 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 3/4] 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 4/4] 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