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 f9afed48c..70404f376 100644 --- a/docs/source/features/interaction/window.md +++ b/docs/source/features/interaction/window.md @@ -51,6 +51,49 @@ 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 + +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() +sim.enable_entity_gizmo(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. + +`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 +world = sim.get_world() +controller = world.get_entity_gizmo() +world.disable_entity_gizmo() +``` + +This controller is distinct from DexSim's target-specific Robot TCP IK +controller. When both are active, **G** controls entity roots and **I** shows or +hides the Robot TCP target. + +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/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/gizmo.rst b/docs/source/tutorial/gizmo.rst index f81ca8919..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,10 +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. -The standard Viser mode includes interactive Gizmo control: +Native controls use DexSim directly. The standard Viser mode includes +interactive Gizmo control: .. code-block:: bash @@ -52,6 +53,27 @@ The standard Viser mode includes 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? ----------------- @@ -62,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 ------------------------------ @@ -81,77 +104,94 @@ 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 IK solver is crucial for gizmo functionality, as it enables the robot to automatically calculate joint angles needed to reach gizmo target positions. +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 - # Enable gizmo for the robot's arm - sim.enable_gizmo( - uid="ur10_gizmo_test", - control_part="arm", - enable_native=False, # Pure Viser; use True for a DexSim window too. + from embodichain.lab.sim.objects import ( + GizmoCfg, + create_robot_ik_gizmo_controller, ) - if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): - logger.log_error("Failed to enable gizmo!") - return - + ik_controller, input_controller = create_robot_ik_gizmo_controller( + robot, + control_part="arm", + cfg=GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ), + world=sim.get_world(), + ) -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. **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 +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 ---------------------------------- -The gizmo-robot interaction follows this efficient workflow: +The gizmo-robot interaction follows this workflow: -1. **Gizmo Callback**: DexSim or Viser records the requested transform -2. **Deferred IK Solving**: Instead of solving IK in the UI callback, the target transform is queued -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 - -This design separates UI responsiveness from computational IK solving, ensuring smooth interaction even with complex robots. +1. **Target Update**: DexSim or Viser records the requested TCP transform +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 paths use ``Robot.set_qpos(..., target=True)`` +5. **Robot Motion**: Joint drives move the robot toward the target without teleporting its current state 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... @@ -165,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 @@ -177,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 -------------------- @@ -233,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 @@ -245,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 @@ -254,14 +271,14 @@ 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 **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 @@ -269,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 @@ -279,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/sim/cfg.py b/embodichain/lab/sim/cfg.py index c36fabbcd..fc6d4e52b 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -1745,8 +1745,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 52c24fefe..838023f8f 100644 --- a/embodichain/lab/sim/objects/__init__.py +++ b/embodichain/lab/sim/objects/__init__.py @@ -19,6 +19,8 @@ Covers lights, rigid bodies (and groups), articulations, robots, deformables (soft/cloth), gizmos, and rigid constraints; every object derives from ``BatchEntity``. """ +from __future__ import annotations + from ..common import BatchEntity from .rigid_object import RigidObject, RigidBodyData, RigidObjectCfg from .rigid_object_group import ( @@ -31,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 dd3176662..cb0352c7e 100644 --- a/embodichain/lab/sim/objects/gizmo.py +++ b/embodichain/lab/sim/objects/gizmo.py @@ -14,37 +14,35 @@ # 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 Any +from typing import TYPE_CHECKING import dexsim import numpy as np import torch -from dexsim.types import ( - AxisArrowType, - AxisCornerType, - AxisOption, - AxisTagType, - RotationRingsOption, -) -from scipy.spatial.transform import Rotation +import warp as wp +from dexsim.types import InputKey from embodichain.lab.sim.common import BatchEntity -from embodichain.lab.sim.objects import RigidObject, Robot +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.utils import configclass, logger -__all__ = ["Gizmo", "GizmoCfg"] +if TYPE_CHECKING: + from dexsim.engine import GizmoController + from dexsim.kit.ik import IKGizmoController, NewtonChainIK + +__all__ = ["Gizmo", "GizmoCfg", "create_robot_ik_gizmo_controller"] @configclass class GizmoCfg: - """Configure native and Viser Gizmo appearance.""" + """Configure Viser Gizmo appearance and robot IK behavior.""" axis_length_x: float = 0.2 """Length of the X-axis arrow.""" @@ -56,97 +54,310 @@ 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.""" - - 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, - ), - } + """Thickness of the rotation rings.""" + + ik_root_link_name: str | None = None + """Robot IK chain root link, or the configured solver root when omitted.""" + + ik_end_link_name: str | None = None + """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.""" + + ik_iterations: int = 24 + """Number of Newton IK iterations per changed target.""" + + ik_device: str | None = None + """Warp device for Newton IK, or the robot device when omitted.""" + + ik_gizmo_scale: float = 1.5 + """Isotropic scale of a native DexSim robot IK target.""" + + ik_toggle_key: InputKey = InputKey.SCANCODE_I + """Native-window key used to toggle a DexSim robot IK target.""" + + +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: + 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 robot abstraction.""" + self._set_qpos(qpos, target=False) + + def set_target_qpos(self, qpos: np.ndarray) -> None: + """Write selected drive targets through the robot abstraction.""" + self._set_qpos(qpos, target=True) + + def get_actived_joint_names(self) -> list[str]: + """Return 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.""" + 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, + ) + + +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 DexSim Gizmo and proxy actor. """ def __init__( self, target: BatchEntity, cfg: GizmoCfg | None = None, - control_part: str | None = "arm", - *, - enable_native: bool = True, + control_part: str | None = None, ) -> None: - num_envs = 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 current num_envs={num_envs}." + "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._control_part = control_part self._target_type = self._detect_target_type(target) - self._env = dexsim.default_world().get_env() - 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_solver: NewtonChainIK | None = None + self._robot_adapter: _RobotGizmoAdapter | 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._desired_target_transform = self._read_target_pose() - - if enable_native: - self._gizmo = self._create_native_gizmo(self.cfg) - self._setup_native_gizmo() + self._control_part = _resolve_control_part(target, control_part) + self._setup_robot_ik_solver() + self._desired_target_transform = self._read_robot_pose() + else: + self._desired_target_transform = self._read_target_pose() @property def target_type(self) -> str: @@ -158,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" @@ -175,24 +381,33 @@ def _detect_target_type(self, target: BatchEntity) -> str: "RigidObject, Robot, or Camera." ) - def _configure_robot(self) -> None: + def _setup_robot_ik_solver(self) -> None: if self.target is None or not isinstance(self.target, Robot): raise RuntimeError("Robot Gizmo has no attached Robot.") - if 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.") 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}." - ) + raise RuntimeError("Robot Gizmo control part is not configured.") + adapter, solver, end_link, tcp_pose = _build_robot_ik( + self.target, + self._control_part, + self.cfg, + ) + self._robot_adapter = adapter + 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 _target_device(self) -> torch.device: if self.target is None: @@ -213,29 +428,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_robot_pose() pose = self.target.get_local_pose(to_matrix=True) return self._as_pose_matrix(pose[0], self._target_device()) @@ -250,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: @@ -260,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}: @@ -274,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 @@ -292,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 @@ -359,200 +502,79 @@ 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: 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], + + base_pose = self._robot_adapter.get_world_pose() + target_pose = target_transform[0].detach().cpu().numpy().astype(np.float32) + base_local = local_pose_from_world(base_pose, target_pose) + joint_names = self._robot_adapter.get_actived_joint_names() + current_qpos = self._robot_adapter.get_current_qpos() + self._ik_solver.set_target_pose( + np.asarray(base_local[:3, 3], dtype=np.float32), + rotation_matrix_to_quat_xyzw(base_local[:3, :3]), + ) + self._ik_solver.solve( + joint_names, + current_qpos, + iterations=self.cfg.ik_iterations, ) + solved_qpos = self._ik_solver.qpos_for_joint_names( + joint_names, + current_qpos, + ) + self._robot_adapter.set_target_qpos(solved_qpos) return True except Exception as error: - logger.log_error(f"Error in 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.""" + """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.""" - self.target = target - self._target_type = self._detect_target_type(target) - self._robot_arm_name = None - if self._target_type == "robot": - self._configure_robot() - with self._state_lock: - self._interaction_owner = None - self._pending_target_transform = None - self._desired_target_transform = self._read_target_pose() - if self._gizmo is not None: - self._remove_proxy_cube() - self._setup_native_gizmo() - - def detach(self) -> None: - """Detach this Gizmo from its current target.""" - if self._gizmo is not None: - self._gizmo.detach_parent() - self._remove_proxy_cube() - with self._state_lock: - self._interaction_owner = None - self._pending_target_transform = None - self.target = None - - 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._gizmo is not None: - self._gizmo.set_visible(self._is_visible) def is_visible(self) -> bool: - """Return whether this Gizmo should be visible.""" + """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 destroy(self) -> None: - """Release native resources and target references.""" - if self._gizmo is not None and hasattr(self._gizmo, "node"): - try: - self._gizmo.node.set_flush_transform_callback(None) - except Exception as error: - logger.log_warning(f"Failed to clear Gizmo callback: {error}") - self._remove_proxy_cube() - if self._gizmo is not None: - try: - self._gizmo.detach_parent() - except Exception as error: - logger.log_warning(f"Failed to detach Gizmo: {error}") + """Release target and IK references.""" with self._state_lock: self._interaction_owner = None self._pending_target_transform = None self._desired_target_transform = None - self._gizmo = 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 6b71f2adf..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, @@ -97,6 +97,7 @@ from embodichain.utils.math import look_at_to_pose, matrix_from_quat, pose_inv if TYPE_CHECKING: + from dexsim.interaction import EntityGizmoConfig, EntityGizmoManipulator from embodichain.lab.visualization import ( RuntimeHealth, RuntimeStats, @@ -232,6 +233,7 @@ class SimulationManager: _instances = {} _cleanup_queue: queue.Queue = queue.Queue() + _DEFAULT_PLANE_GIZMO_TARGET_ID = (1 << 64) - 1 SUPPORTED_SENSOR_TYPES = { "Camera": Camera, @@ -324,6 +326,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() @@ -676,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: @@ -1905,22 +1913,51 @@ 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 entity control and exclude the EmbodiChain ground. + + Args: + config: Native DexSim entity-Gizmo configuration. + + Returns: + The active world-owned entity Gizmo manipulator. + """ + controller = ( + self._world.enable_entity_gizmo() + if config is None + else self._world.enable_entity_gizmo(config) + ) + default_plane = getattr(self, "_default_plane", None) + if default_plane is None: + 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}." + ) + 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. @@ -1955,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: @@ -2013,6 +2028,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() @@ -2130,6 +2147,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() @@ -2140,6 +2158,65 @@ 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() + 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) + 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/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 3ff1c7de7..000000000 --- a/embodichain/lab/sim/utility/gizmo_utils.py +++ /dev/null @@ -1,221 +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 EmbodiSim. - -This module provides utility functions for creating gizmo transform callbacks. -""" - -from typing import Callable -from typing import TYPE_CHECKING -from dexsim.types import TransformMask - -if TYPE_CHECKING: - from embodichain.lab.sim.objects import Robot - - -def create_gizmo_callback() -> Callable: - """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() - """ - - def gizmo_transform_callback(node, local_pose, flag): - if node is not None: - node.set_transform(local_pose, flag) - - return gizmo_transform_callback - - -def run_gizmo_robot_control_loop( - robot: object | str, control_part: str = "arm", end_link_name: str | 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 Robot - from embodichain.lab.sim.solvers import PinkSolverCfg - - from embodichain.utils.logger import log_info, log_warning, log_error - - sim = SimulationManager.get_instance() - - if isinstance(robot, str): - robot = sim.get_robot(uid=robot) - - # 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) - control_part_link_names = robot.get_control_part_link_names(name=control_part) - end_link_name = ( - 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, - ) - 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) - - # 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(): - """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) - 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) - 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}") - 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) - 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) - 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: - 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/__init__.py b/embodichain/lab/visualization/__init__.py index 73e45aff0..e66b1729b 100644 --- a/embodichain/lab/visualization/__init__.py +++ b/embodichain/lab/visualization/__init__.py @@ -37,6 +37,7 @@ JointControlSpec, JointControlState, MeshGeometry, + PickCommand, PointCloudOverlay, SceneFrame, SceneManifest, @@ -75,6 +76,7 @@ "JointControlState", "LatestFrameQueue", "MeshGeometry", + "PickCommand", "PointCloudOverlay", "RuntimeHealth", "RuntimeStats", diff --git a/embodichain/lab/visualization/backends/base.py b/embodichain/lab/visualization/backends/base.py index be73aa482..655ea728a 100644 --- a/embodichain/lab/visualization/backends/base.py +++ b/embodichain/lab/visualization/backends/base.py @@ -23,6 +23,7 @@ CameraImageFrame, GizmoCommand, JointControlCommand, + PickCommand, SceneFrame, SceneManifest, ) @@ -40,6 +41,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 + def set_joint_control_command_sink( self, sink: Callable[[JointControlCommand], None] | None, diff --git a/embodichain/lab/visualization/backends/viser.py b/embodichain/lab/visualization/backends/viser.py index dc9862ea3..931293271 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, @@ -36,6 +37,7 @@ JointControlSpec, JointControlState, MeshGeometry, + PickCommand, PointCloudOverlay, SceneFrame, SceneManifest, @@ -140,6 +142,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._joint_control_handles: dict[str, _JointControlHandle] = {} self._joint_control_specs: dict[str, JointControlSpec] = {} self._joint_control_states: dict[str, JointControlState] = {} @@ -231,6 +240,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: @@ -342,6 +359,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) @@ -352,6 +385,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, @@ -1024,6 +1139,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() @@ -1181,6 +1298,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) @@ -1405,6 +1526,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]) @@ -1556,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/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 3fd30b5e8..5da1eed1b 100644 --- a/embodichain/lab/visualization/protocol.py +++ b/embodichain/lab/visualization/protocol.py @@ -39,6 +39,7 @@ "JointControlSpec", "JointControlState", "MeshGeometry", + "PickCommand", "PointCloudOverlay", "SceneFrame", "SceneManifest", @@ -355,6 +356,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 JointControlSpec: """Static description of one scalar articulation joint control. diff --git a/embodichain/lab/visualization/runtime.py b/embodichain/lab/visualization/runtime.py index 482e52fc0..9fa020e98 100644 --- a/embodichain/lab/visualization/runtime.py +++ b/embodichain/lab/visualization/runtime.py @@ -30,6 +30,7 @@ GizmoCommand, JointControlCommand, JointControlProvider, + PickCommand, SceneFrame, SceneManifest, SceneOverlays, @@ -141,6 +142,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: + del self._commands[index] + break + 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() + + class JointControlCommandQueue: """Bounded queue that keeps only the newest value for each joint control.""" @@ -243,6 +282,8 @@ def __init__( self._gizmo_commands = GizmoCommandQueue() self._joint_control_commands = JointControlCommandQueue() 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._backend.set_joint_control_command_sink( self._enqueue_joint_control_command ) @@ -277,6 +318,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() + def _enqueue_joint_control_command(self, command: JointControlCommand) -> None: if self.cfg.allow_commands: self._joint_control_commands.put(command) @@ -611,6 +662,7 @@ def stop(self, timeout: float = 10.0) -> None: self._replay_control_states.clear() self._replay_control_commands.clear() self._gizmo_commands.clear() + self._pick_commands.clear() self._joint_control_commands.clear() self._raise_worker_error() diff --git a/embodichain/lab/visualization/scene_exporter.py b/embodichain/lab/visualization/scene_exporter.py index cf77446b8..415ba4a72 100644 --- a/embodichain/lab/visualization/scene_exporter.py +++ b/embodichain/lab/visualization/scene_exporter.py @@ -483,6 +483,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/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 690c7a6de..b01175738 100644 --- a/examples/sim/gizmo/gizmo_object.py +++ b/examples/sim/gizmo/gizmo_object.py @@ -14,16 +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 objects with native DexSim or browser-based Viser Gizmos.""" from __future__ import annotations import argparse import time +import dexsim + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg @@ -91,18 +90,24 @@ def main(): native_window_opened = False if not args.headless: + entity_gizmo_config = dexsim.interaction.EntityGizmoConfig() + entity_gizmo_config.max_gizmos = 0 native_window_opened = sim.open_window() + if native_window_opened: + sim.enable_entity_gizmo(entity_gizmo_config) - # Enable native-window or Viser Gizmo control. - if native_window_opened or args.viser: + # 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=native_window_opened, ) sim.enable_gizmo( uid="cube2", - enable_native=native_window_opened, ) + elif native_window_opened: + 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.") else: logger.log_warning( "Gizmo interaction is disabled in headless mode without Viser." @@ -111,9 +116,9 @@ def main(): logger.log_info("Scene setup complete!") logger.log_info(f"Running simulation with 1 environment(s)") if native_window_opened or args.viser: - if sim.has_gizmo("cube1"): + if args.viser and sim.has_gizmo("cube1"): logger.log_info("Gizmo enabled for cube1 - you can drag it around!") - if sim.has_gizmo("cube2"): + if args.viser and sim.has_gizmo("cube2"): logger.log_info("Gizmo enabled for cube2 - you can drag it around!") logger.log_info("Press Ctrl+C to stop the simulation") @@ -136,11 +141,14 @@ def run_simulation(sim: SimulationManager): step_count += 1 - # Disable gizmo after 200000 steps (example) + # Disable Gizmo control after 200000 steps (example). if step_count == 200000 and gizmo_enabled: - logger.log_info("Disabling gizmo at step 200000") - sim.disable_gizmo("cube1") - sim.disable_gizmo("cube2") + logger.log_info("Disabling Gizmo control at step 200000") + if sim.get_world().get_entity_gizmo() is not None: + sim.get_world().disable_entity_gizmo() + else: + sim.disable_gizmo("cube1") + sim.disable_gizmo("cube2") gizmo_enabled = False # Print FPS every second diff --git a/examples/sim/gizmo/gizmo_robot.py b/examples/sim/gizmo/gizmo_robot.py index cfccfd566..477abb924 100644 --- a/examples/sim/gizmo/gizmo_robot.py +++ b/examples/sim/gizmo/gizmo_robot.py @@ -13,9 +13,7 @@ # 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 a native DexSim or Viser Gizmo.""" from __future__ import annotations @@ -27,6 +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, + create_robot_ik_gizmo_controller, +) from embodichain.lab.sim.cfg import ( RenderCfg, RobotCfg, @@ -34,7 +36,6 @@ JointDrivePropertiesCfg, ) from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.solvers import PinkSolverCfg from embodichain.data import get_data_path from embodichain.utils import logger @@ -118,12 +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], + ], + ) + 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", - 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!") @@ -136,19 +154,22 @@ def main(): logger.log_info("Gizmo-Robot example started!") if native_window_opened or args.viser: logger.log_info("Use the gizmo to drag the robot end-effector (EE)") + if native_window_opened: + logger.log_info("Press I to show or hide the native robot IK Gizmo") logger.log_info("Press 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 3c04f4eb4..c14d9dd3b 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -13,9 +13,7 @@ # 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 a native DexSim or Viser Gizmo.""" from __future__ import annotations @@ -26,6 +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, + 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, @@ -106,12 +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", - 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!") @@ -124,19 +137,22 @@ def main(): logger.log_info("Gizmo-Robot example started!") if native_window_opened or args.viser: logger.log_info("Use the gizmo to drag the robot end-effector (EE)") + if native_window_opened: + logger.log_info("Press I to show or hide the native robot IK Gizmo") logger.log_info("Press 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/sim/objects/test_gizmo.py b/tests/sim/objects/test_gizmo.py index 0b1f3c6a5..2cf91a13a 100644 --- a/tests/sim/objects/test_gizmo.py +++ b/tests/sim/objects/test_gizmo.py @@ -18,85 +18,192 @@ from types import SimpleNamespace +import numpy as np +import pytest import torch import embodichain.lab.sim.objects.gizmo as gizmo_module -from embodichain.lab.sim.objects.gizmo import Gizmo +from embodichain.lab.sim.objects.gizmo import ( + Gizmo, + GizmoCfg, + _RobotGizmoAdapter, + create_robot_ik_gizmo_controller, +) -class _RigidObject: +class _FakeAdapterRobot: + """Small Robot-compatible state holder for native 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(uid="cube") - self.pose = torch.eye(4, dtype=torch.float32).unsqueeze(0) - self.set_calls: list[tuple[torch.Tensor, list[int]]] = [] + 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]] = [] + + 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 self.pose.clone() + return torch.eye(4, dtype=torch.float32).unsqueeze(0) - def set_local_pose( + def get_link_pose( self, - pose: torch.Tensor, + link_name: str, env_ids: list[int], - ) -> None: - self.pose = pose.clone() - self.set_calls.append((pose.clone(), env_ids)) + 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) -class _Camera(_RigidObject): - pass +def test_robot_adapter_synchronizes_selected_joint_state() -> None: + robot = _FakeAdapterRobot() + 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]) -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]]] = [] + adapter.set_target_qpos(np.array([0.7, 0.9], dtype=np.float32)) - def get_proprioception(self) -> dict[str, torch.Tensor]: - return {"qpos": torch.zeros((1, 2), dtype=torch.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 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 test_robot_adapter_reads_root_and_link_pose_through_robot() -> None: + adapter = _RobotGizmoAdapter(_FakeAdapterRobot(), "arm") - def compute_ik( - self, - *args: object, - **kwargs: object, - ) -> tuple[torch.Tensor, torch.Tensor]: - return torch.tensor([True]), torch.tensor([[0.4, -0.2]]) + 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(_FakeAdapterRobot(), "arm") + + with pytest.raises(ValueError, match="Expected qpos shape"): + adapter.set_target_qpos(np.zeros(3, dtype=np.float32)) - def set_qpos( + +def test_robot_native_ik_chain_can_be_configured_without_solver() -> None: + cfg = GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="tool_link", + ) + + 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_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), + ), + ) + + class _InputController: + pass + + class _IKController: + def __init__(self, *args, **kwargs) -> None: + self.args = args + self.kwargs = kwargs + + import dexsim.engine + import dexsim.kit.ik + + 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) + + controller, input_controller = create_robot_ik_gizmo_controller( + robot, + world=world, + ) + + 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) + self.set_calls: list[tuple[torch.Tensor, list[int]]] = [] + + def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: + assert to_matrix + return self.pose.clone() + + def set_local_pose( self, - qpos: torch.Tensor, - joint_ids: list[int], + pose: torch.Tensor, env_ids: list[int], ) -> None: - self.set_calls.append((qpos.clone(), joint_ids, env_ids)) + self.pose = pose.clone() + self.set_calls.append((pose.clone(), env_ids)) -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()), - ) +class _Camera(_RigidObject): + pass 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]) @@ -106,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 @@ -125,19 +230,61 @@ 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) - _patch_headless_dexsim(monkeypatch) - target = _Robot() - gizmo = Gizmo(target, control_part="arm", enable_native=False) +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) + 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._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") pose = torch.eye(4, dtype=torch.float32).unsqueeze(0) pose[0, 0, 3] = 0.5 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 699acb0e2..6074e1cb2 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -22,6 +22,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock +import dexsim import numpy as np import pytest import torch @@ -35,6 +36,7 @@ ) from embodichain.lab.visualization import ( GizmoCommand, + PickCommand, PointCloudOverlay, SceneOverlays, VisualizationCfg, @@ -87,12 +89,35 @@ 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.physics_updates: list[float] = [] + 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 @@ -103,6 +128,21 @@ def is_physics_manually_update(self) -> bool: def update(self, physics_dt: float) -> None: self.physics_updates.append(physics_dt) + 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 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.""" @@ -168,13 +208,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, + visualization=SimpleNamespace(backend="none"), + ) sim._window = window sim._window_record_state = None sim._window_record_camera = None sim._window_record_save_threads = [] + sim._window_record_hotkey_cfg = None + sim._window_record_input_control = None + sim._window_camera_pose_hotkey_cfg = None + sim._window_camera_pose_input_control = None sim._env = FakeEnv() sim._world = FakeWorld() + sim._default_plane = object() + sim._visualization_runtime = None + sim.is_window_opened = window is not None return sim @@ -317,6 +368,219 @@ 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): + enabled.append((uid, 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: ( + f"{uid}:{control_part}" if control_part else uid + ) + in sim._gizmos + ) + 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) + + +@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( + 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() @@ -410,6 +674,45 @@ def test_open_window_is_idempotent() -> None: sim._world.open_window.assert_not_called() +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.entity_gizmo + assert sim._world.entity_gizmo_configs == [config] + 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_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._world.entity_gizmo_configs == [] + + +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 True + assert sim._world.window_closed is True + assert sim.is_window_opened is False + + def test_start_visualization_rejects_open_native_window() -> None: sim = object.__new__(SimulationManager) sim.sim_config = SimpleNamespace( @@ -508,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_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_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 diff --git a/tests/visualization/test_viser_backend.py b/tests/visualization/test_viser_backend.py index b8af652f3..5b7dcf8a6 100644 --- a/tests/visualization/test_viser_backend.py +++ b/tests/visualization/test_viser_backend.py @@ -275,6 +275,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: @@ -838,6 +849,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(