diff --git a/docs/source/features/interaction/window.md b/docs/source/features/interaction/window.md index 10faf2110..da7279ce0 100644 --- a/docs/source/features/interaction/window.md +++ b/docs/source/features/interaction/window.md @@ -44,6 +44,52 @@ Recording hotkey registration is controlled by `SimConfig.window_record.enable_h The camera-pose hotkey is controlled by `SimulationManagerCfg.window_camera_pose.enable_hotkey` and prints look-at form by default. Set `SimulationManagerCfg.window_camera_pose.convert_to_look_at=False` to print the raw 4x4 pose matrix instead. The same output can be requested programmatically with `SimulationManager.print_window_camera_pose()`. +### Entity Gizmo Control + +Opening a non-headless `SimulationManager` window enables dexsim's world-owned +`EntityGizmoManipulator` by default: + +```python +import dexsim + +gizmo_config = dexsim.interaction.EntityGizmoConfig() +gizmo_config.max_gizmos = 0 # Unlimited simultaneous bindings. +sim.open_window(entity_gizmo_config=gizmo_config) +``` + +While enabled, left-click a render mesh, dynamic/kinematic rigid body, or +articulation link and press **G** to attach or detach its root gizmo. The +controller supports multiple simultaneous bindings and owns selection, +temporary physics-state changes, and cleanup. No `sim.update_gizmos()` call is +needed for this world-level controller. + +EmbodiChain's built-in `default_plane` is registered as an immovable target and +cannot receive an entity gizmo. Other supported scene entities remain +selectable normally. + +For a view-only window, opt out explicitly: + +```python +sim.open_window(enable_entity_gizmo=False) +``` + +Set `SimulationManagerCfg.enable_entity_gizmo_on_window_open=False` to change +the default for constructor-opened and subsequently opened windows. Headless +simulations do not create or enable the controller. + +`sim.enable_entity_gizmo(config)` can reconfigure or reactivate the controller +at any time, and `sim.disable_entity_gizmo()` cancels it without closing the +window. The last explicit configuration is restored if the window is closed +and reopened. + +Use `sim.get_entity_gizmo()` to access the native controller and +`sim.has_entity_gizmo()` to query its lifecycle state. Closing the window or +destroying the `SimulationManager` disables it automatically. + +This controller is distinct from the target-specific Robot TCP IK gizmo. When +both are active, **G** controls entity roots and **I** shows or hides the Robot +TCP IK gizmo. + ## Customizing Window Events Users can create their own custom window interaction controls by subclassing the `ObjectManipulator` class (provided by `dexsim`). This allows for the implementation of specific behaviors and responses to user inputs. diff --git a/docs/source/tutorial/gizmo.rst b/docs/source/tutorial/gizmo.rst index 6f2a5b7a0..49404a358 100644 --- a/docs/source/tutorial/gizmo.rst +++ b/docs/source/tutorial/gizmo.rst @@ -5,7 +5,7 @@ Interactive Robot Control with Gizmo .. currentmodule:: embodichain.lab.sim -This tutorial demonstrates how to use the Gizmo class for interactive robot manipulation in SimulationManager. You'll learn how to create a gizmo attached to a robot's end-effector and use it for real-time inverse kinematics (IK) control, allowing intuitive manipulation of robot poses through visual interaction. +This tutorial demonstrates how to use the Gizmo class for interactive robot manipulation in SimulationManager. Robot gizmos delegate interactive inverse kinematics (IK) to dexsim's Newton IK controller while all joint-state reads and drive-target writes continue to pass through the EmbodiChain ``Robot`` abstraction. The Code ~~~~~~~~ @@ -28,7 +28,8 @@ Similar to the previous tutorial on robot simulation, we use the :class:`Simulat -**Important:** Gizmo only supports single environment mode (`num_envs=1`). Using multiple environments will raise an exception. +**Important:** The target-specific Robot TCP, rigid-object, and Camera +``Gizmo`` wrapper supports only single-environment mode (``num_envs=1``). All gizmo creation, visibility, and destruction operations must be managed via the SimulationManager API: @@ -57,21 +58,34 @@ The :class:`objects.Gizmo` class provides a unified interface for interactive co Setting up Robot Configuration ------------------------------ -First, we configure a UR10 robot with an IK solver for end-effector control: +First, configure a UR10 robot and its controllable arm joints: .. literalinclude:: ../../../scripts/tutorials/sim/gizmo_robot.py :language: python - :start-at: # Create UR10 robot configuration + :start-at: # Create UR10 robot :end-at: robot = sim.add_robot(cfg=robot_cfg) Key components of the robot configuration: - **URDF Configuration**: Loads the robot's kinematic and visual model - **Control Parts**: Defines which joints can be controlled (``"Joint[1-6]"`` for UR10) -- **IK Solver**: :class:`solvers.PinkSolverCfg` provides inverse kinematics capabilities - **Drive Properties**: Sets stiffness and damping for joint control -The IK solver is crucial for gizmo functionality, as it enables the robot to automatically calculate joint angles needed to reach gizmo target positions. +An EmbodiChain kinematics solver is not required by the gizmo. The IK chain is declared when enabling it: + +.. code-block:: python + + gizmo_cfg = GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ) + sim.enable_gizmo( + uid="ur10_gizmo_test", + control_part="arm", + gizmo_cfg=gizmo_cfg, + ) + +For existing robot configurations, these link names and the TCP can instead be inherited from the selected control part's configured EmbodiChain solver. The solver supplies metadata only; interactive IK is still performed by dexsim's ``NewtonChainIK``. Creating and Attaching a Gizmo ------------------------------- @@ -83,7 +97,14 @@ After configuring the robot, enable the gizmo for interactive control using the .. code-block:: python # Enable gizmo for the robot's arm - sim.enable_gizmo(uid="ur10_gizmo_test", control_part="arm") + sim.enable_gizmo( + uid="ur10_gizmo_test", + control_part="arm", + gizmo_cfg=GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ), + ) if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): logger.log_error("Failed to enable gizmo!") return @@ -102,22 +123,23 @@ The Gizmo system will automatically: 1. **Detect Target Type**: Identify that the target is a robot (vs. rigid object or camera) 2. **Find End-Effector**: Locate the robot's end-effector link (``ee_link`` for UR10) -3. **Create Proxy Object**: Generate a small invisible cube at the end-effector position -4. **Set Up IK Callback**: Configure the gizmo to trigger IK solving when moved +3. **Build Newton IK Chain**: Build a reduced start-link-to-end-link model from the robot URDF +4. **Bind dexsim Controller**: Attach ``IKGizmoController`` directly to the articulation adapter How Gizmo-Robot Interaction Works ---------------------------------- -The gizmo-robot interaction follows this efficient workflow: +The gizmo-robot interaction follows this workflow: -1. **Gizmo Callback**: When the user drags the gizmo, a callback function updates the proxy object's transform -2. **Deferred IK Solving**: Instead of solving IK immediately in the callback (which causes UI lag), the target transform is stored -3. **Update Loop**: During each simulation step, ``gizmo.update()`` solves IK and applies joint commands -4. **Robot Motion**: The robot smoothly moves to follow the gizmo position +1. **Target Update**: Dragging the dexsim target gizmo updates the Newton IK target state +2. **Deferred Solve**: ``sim.update_gizmos()`` asks ``IKGizmoController`` to solve only when the target changed +3. **State Bridge**: The adapter reads the selected EmbodiChain control-part joints as the solve seed +4. **Drive Target**: The solved positions are written through ``Robot.set_qpos(..., target=True)`` so CPU and CUDA state paths stay synchronized +5. **Robot Motion**: Joint drives move the robot toward the target without teleporting its current state -This design separates UI responsiveness from computational IK solving, ensuring smooth interaction even with complex robots. +Robot gizmos no longer create or maintain an EmbodiChain proxy cube. Camera gizmos continue to use their existing proxy path, and rigid-object gizmos continue to follow the selected object directly. The Simulation Loop ------------------- @@ -163,19 +185,57 @@ Gizmo Lifecycle Management Gizmo lifecycle is managed by SimulationManager: -- Enable: `sim.enable_gizmo(...)` +- Enable a target-specific gizmo: `sim.enable_gizmo(...)` - Update: Main loop automatically calls `sim.update_gizmos()` - Destroy/disable: `sim.disable_gizmo(...)` or `sim.destroy()` (recommended) There is no need to manually create or destroy Gizmo instances. All resources are managed by SimulationManager. +World-Level Entity Gizmo +------------------------ + +For selection-based root manipulation, opening a non-headless window enables +dexsim's world-level entity gizmo by default: + +.. code-block:: python + + import dexsim + + config = dexsim.interaction.EntityGizmoConfig() + config.max_gizmos = 0 + sim.open_window(entity_gizmo_config=config) + + # Left-click an entity and press G to attach or detach a gizmo. + # Multiple entities may remain attached. + + sim.disable_entity_gizmo() + +This path supports render meshes, eligible rigid bodies, and articulation +roots. dexsim owns raycast selection, temporary body-state changes, multiple +bindings, and cleanup. It requires no ``sim.update_gizmos()`` call. + +EmbodiChain's built-in ``default_plane`` is excluded from manipulation. +Selecting it and pressing **G** does not create a gizmo. + +Use ``sim.open_window(enable_entity_gizmo=False)`` for a view-only window, or +set ``SimulationManagerCfg.enable_entity_gizmo_on_window_open=False`` to +change the default. Headless simulations do not create the controller. + +``sim.get_entity_gizmo()`` returns the native +``EntityGizmoManipulator`` and ``sim.has_entity_gizmo()`` reports whether it is +enabled. Closing the window or destroying the simulation also disables it. + +The Robot end-effector controller remains target-specific because it solves a +TCP pose rather than editing the articulation root. By default, **G** controls +the entity gizmo and **I** toggles Robot TCP IK gizmo visibility. + Available Gizmo Methods ----------------------- -If you need to access the underlying Gizmo instance (via `sim.get_gizmo`), you can use the following methods: +If you need to access the underlying Gizmo instance (via `sim.get_gizmo`), you can use the following methods. For robot targets these methods operate on dexsim's IK target gizmo: **Transform Control:** @@ -245,8 +305,8 @@ Tips and Best Practices **Robot compatibility:** -- Ensure your robot is configured with a correct IK solver -- Check the end-effector (EE) link name +- Set valid ``ik_root_link_name`` and ``ik_end_link_name`` values, or configure an EmbodiChain solver whose chain metadata can be inherited +- Set ``ik_tcp_pose`` when the desired tool center point differs from the end-link frame - Test joint limits and workspace boundaries @@ -254,7 +314,7 @@ Tips and Best Practices **Visualization customization:** - Adjust gizmo appearance via Gizmo config (e.g., ``set_line_width()``; requires access to the instance via `sim.get_gizmo`) -- Adjust gizmo scale according to robot size +- Adjust robot target size with ``GizmoCfg.ik_gizmo_scale`` - Enable collision for debugging if needed Next Steps @@ -265,6 +325,6 @@ After mastering basic gizmo usage, you can explore: - **Multi-robot Gizmos**: Attach gizmos to multiple robots simultaneously - **Custom Gizmo Callbacks**: Implement application-specific interaction logic - **Gizmo with Rigid Objects**: Use gizmos for interactive object manipulation -- **Advanced IK Configuration**: Fine-tune solver parameters for specific robots +- **Advanced IK Configuration**: Tune ``GizmoCfg.ik_iterations``, ``ik_device``, and the TCP pose -For more advanced robot control and simulation features, refer to the complete :doc:`robot` tutorial and the API documentation for :class:`objects.Gizmo` and :class:`solvers.PinkSolverCfg`. +For more advanced robot control and simulation features, refer to the complete :doc:`robot` tutorial and the API documentation for :class:`objects.Gizmo`. diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index eb0dd7f58..b30257294 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -1735,8 +1735,8 @@ class RobotCfg(ArticulationCfg): If no control part is specified, the robot will use all joints as a single control part. Note: - - if `control_parts` is specified, `solver_cfg` must be a dict with part names as - keys corresponding to the control parts name. + - `control_parts` can be used without `solver_cfg`. If `solver_cfg` is a + dictionary, its keys must correspond to control-part names. - The joint names in the control parts support regular expressions, e.g., 'joint[1-6]'. After initialization of robot, the names will be expanded to a list of full joint names. - `Robot` is a derived class of `Articulation`, with control parts support. So the `drive_pros` diff --git a/embodichain/lab/sim/objects/__init__.py b/embodichain/lab/sim/objects/__init__.py index 2254f1008..06014aef1 100644 --- a/embodichain/lab/sim/objects/__init__.py +++ b/embodichain/lab/sim/objects/__init__.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + from ..common import BatchEntity from .rigid_object import RigidObject, RigidBodyData, RigidObjectCfg from .rigid_object_group import ( @@ -26,7 +28,7 @@ from .articulation import Articulation, ArticulationData, ArticulationCfg from .robot import Robot, RobotCfg from .light import Light, LightCfg -from .gizmo import Gizmo +from .gizmo import Gizmo, GizmoCfg from .constraint import RigidConstraint diff --git a/embodichain/lab/sim/objects/gizmo.py b/embodichain/lab/sim/objects/gizmo.py index 9fb370c83..42ef89da6 100644 --- a/embodichain/lab/sim/objects/gizmo.py +++ b/embodichain/lab/sim/objects/gizmo.py @@ -13,71 +13,110 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -""" -Gizmo: A reusable controller for interactive manipulation of simulation elements (object, robot, camera, etc.) -""" +"""Interactive gizmos for simulation objects, robots, and cameras.""" -import numpy as np -import torch -import dexsim -from typing import Callable -from scipy.spatial.transform import Rotation as R +from __future__ import annotations -from embodichain.lab.sim.common import BatchEntity -from embodichain.lab.sim.objects import RigidObject, Robot -from embodichain.lab.sim.sensors import Camera -from embodichain.utils import configclass, logger +from collections.abc import Callable +from typing import TYPE_CHECKING, Any +import dexsim +import numpy as np +import torch +import warp as wp from dexsim.types import ( - AxisOption, - RotationRingsOption, AxisArrowType, AxisCornerType, + AxisOption, AxisTagType, - TransformMask, - ActorType, - RigidBodyShape, - PhysicalAttr, + InputKey, + RotationRingsOption, ) +from scipy.spatial.transform import Rotation as R +from embodichain.lab.sim.common import BatchEntity +from embodichain.lab.sim.objects.rigid_object import RigidObject +from embodichain.lab.sim.objects.robot import Robot +from embodichain.lab.sim.sensors import Camera from embodichain.lab.sim.utility.gizmo_utils import create_gizmo_callback +from embodichain.utils import configclass, logger + +if TYPE_CHECKING: + from dexsim.kit.ik import IKGizmoController, NewtonChainIK + +__all__ = ["Gizmo", "GizmoCfg"] @configclass class GizmoCfg: - """Configuration class for Gizmo parameters. + """Configure gizmo appearance and robot Newton IK behavior.""" - This class defines the visual and interaction parameters for gizmo controllers, - including axis appearance and rotation rings settings. - """ - - # Axis configuration axis_length_x: float = 0.2 - """Length of X-axis arrow.""" + """Length of the X-axis arrow.""" + axis_length_y: float = 0.2 - """Length of Y-axis arrow.""" + """Length of the Y-axis arrow.""" + axis_length_z: float = 0.2 - """Length of Z-axis arrow.""" + """Length of the Z-axis arrow.""" + axis_size: float = 0.01 - """Thickness of axis lines.""" + """Thickness of the axis lines.""" + arrow_type: AxisArrowType = AxisArrowType.CONE """Type of arrow head.""" + corner_type: AxisCornerType = AxisCornerType.SPHERE """Type of axis corner.""" + tag_type: AxisTagType = AxisTagType.PLANE """Type of axis label.""" - # Rotation rings configuration rings_radius: float = 0.15 - """Radius of rotation rings.""" + """Radius of the rotation rings.""" + rings_size: float = 0.01 - """Thickness of rotation rings.""" + """Thickness of the rotation rings.""" + + ik_root_link_name: str | None = None + """Robot IK chain root link. + + When omitted, the value is read from the selected control part's configured + EmbodiChain solver. + """ + + ik_end_link_name: str | None = None + """Robot IK chain end link. + + When omitted, the value is read from the selected control part's configured + EmbodiChain solver. + """ - def to_options_dict(self) -> dict: - """Convert configuration to options dictionary format expected by gizmo creation. + ik_tcp_pose: torch.Tensor | np.ndarray | list[list[float]] | None = None + """End-link-to-TCP transform for robot IK. + + When omitted, the configured EmbodiChain solver TCP is used if available; + otherwise the identity transform is used. + """ + + ik_iterations: int = 24 + """Number of Newton IK iterations per changed target.""" + + ik_device: str | None = None + """Warp device for the Newton IK model, or the robot device when omitted.""" + + ik_gizmo_scale: float = 1.5 + """Isotropic scale of dexsim's robot IK target gizmo.""" + + ik_toggle_key: InputKey = InputKey.SCANCODE_I + """Window key used by dexsim to toggle the robot IK gizmo.""" + + def to_options_dict(self) -> dict[str, object]: + """Convert the visual configuration to dexsim gizmo options. Returns: - Dictionary containing AxisOption and RotationRingsOption objects. + The axis and rotation-ring options used by rigid-object and camera + gizmos. """ return { "axis": AxisOption( @@ -90,19 +129,125 @@ def to_options_dict(self) -> dict: tag_type=self.tag_type, ), "rings": RotationRingsOption( - radius=self.rings_radius, size=self.rings_size + radius=self.rings_radius, + size=self.rings_size, ), } +class _RobotGizmoAdapter: + """Expose one EmbodiChain robot control part to dexsim's IK controller.""" + + def __init__(self, robot: Robot, control_part: str, env_id: int = 0) -> None: + """Create the adapter. + + Args: + robot: EmbodiChain robot whose state is synchronized. + control_part: Robot control part driven by the IK solution. + env_id: Environment instance exposed to the interactive controller. + + Raises: + ValueError: If the control part, environment, or joint selection is + invalid. + """ + if not robot.control_parts or control_part not in robot.control_parts: + raise ValueError( + f"Control part {control_part!r} is not defined. Available parts: " + f"{list(robot.control_parts or {})}." + ) + if env_id < 0 or env_id >= robot.num_instances: + raise ValueError( + f"Robot gizmo env_id={env_id} is outside [0, {robot.num_instances})." + ) + + joint_ids = robot.get_joint_ids(control_part, remove_mimic=True) + if not joint_ids: + raise ValueError( + f"Control part {control_part!r} has no non-mimic active joints." + ) + + self.robot = robot + self.control_part = control_part + self.env_id = env_id + self.joint_ids = list(joint_ids) + self.joint_names = [robot.joint_names[index] for index in self.joint_ids] + + def get_current_qpos(self) -> np.ndarray: + """Return current selected joint positions in dexsim joint-name order.""" + return self._selected_qpos(target=False) + + def get_target_qpos(self) -> np.ndarray: + """Return target selected joint positions in dexsim joint-name order.""" + return self._selected_qpos(target=True) + + def set_current_qpos(self, qpos: np.ndarray) -> None: + """Write selected current positions through the EmbodiChain abstraction.""" + self._set_qpos(qpos, target=False) + + def set_target_qpos(self, qpos: np.ndarray) -> None: + """Write selected drive targets through the EmbodiChain abstraction.""" + self._set_qpos(qpos, target=True) + + def get_actived_joint_names(self) -> list[str]: + """Return selected active joint names using dexsim's API spelling.""" + return self.joint_names.copy() + + def get_world_pose(self) -> np.ndarray: + """Return the selected robot instance's root pose as a matrix.""" + pose = self.robot.get_local_pose(to_matrix=True)[self.env_id] + return pose.detach().cpu().numpy().astype(np.float32, copy=True) + + def get_link_names(self, include_fixed: bool = True) -> list[str]: + """Return all runtime link names. + + Args: + include_fixed: Kept for compatibility with the dexsim articulation + API. EmbodiChain's link list already includes fixed links. + """ + del include_fixed + return list(self.robot.link_names) + + def get_link_pose(self, link_name: str) -> np.ndarray: + """Return one runtime link pose as a world-space matrix.""" + pose = self.robot.get_link_pose( + link_name, + env_ids=[self.env_id], + to_matrix=True, + )[0] + return pose.detach().cpu().numpy().astype(np.float32, copy=True) + + def _selected_qpos(self, target: bool) -> np.ndarray: + qpos = self.robot.get_qpos(target=target)[self.env_id, self.joint_ids] + return qpos.detach().cpu().numpy().astype(np.float32, copy=True) + + def _set_qpos(self, qpos: np.ndarray, target: bool) -> None: + values = np.asarray(qpos, dtype=np.float32) + if values.shape != (len(self.joint_ids),): + raise ValueError( + f"Expected qpos shape ({len(self.joint_ids)},), got {values.shape}." + ) + self.robot.set_qpos( + qpos=torch.as_tensor( + values, + dtype=torch.float32, + device=self.robot.device, + ).unsqueeze(0), + joint_ids=self.joint_ids, + env_ids=[self.env_id], + target=target, + ) + + class Gizmo: - """ - Generic Gizmo controller for simulation elements. - Supports RigidObject, Robot, and Camera with type-specific handling. + """Control one rigid object, robot end effector, or camera interactively. + + Robot targets use dexsim's :class:`IKGizmoController` and + :class:`NewtonChainIK`. Rigid-object and camera behavior remains on the + existing direct/proxy paths. - Note: - Gizmo can only be used in single environment mode (num_envs=1). - Will raise RuntimeError if used with multiple environments. + .. attention:: + Gizmos currently expose only one environment instance. Create them only + when ``num_envs=1``. """ def __init__( @@ -110,441 +255,426 @@ def __init__( target: BatchEntity, cfg: GizmoCfg | None = None, control_part: str | None = "arm", - ): - """ + ) -> None: + """Create and attach a gizmo. + Args: - target: The simulation element to control (RigidObject, Robot, or Camera) - cfg: Gizmo configuration parameters (optional, uses default if None) - control_part: For robots, specifies which control part to use (optional, default: "arm") + target: Simulation element to control. + cfg: Gizmo appearance and robot IK configuration. + control_part: Robot control part. When omitted, the first configured + part is selected. """ - self.target = target - self._target_type = self._detect_target_type(target) - self._control_part = control_part - self._env = dexsim.default_world().get_env() - self._windows = dexsim.default_world().get_windows() + world = dexsim.default_world() + if world is None: + raise RuntimeError("A dexsim world must exist before creating a gizmo.") - # Check if running in single environment (num_env must be 1) - num_envs = dexsim.get_world_num() - if num_envs > 1: + self.cfg = cfg if cfg is not None else GizmoCfg() + self._world = world + self._env = world.get_env() + self._control_part = control_part + self._callback: Callable[..., Any] | None = None + self._state = "active" + self._is_visible = True + self._gizmo: object | None = None + self._proxy_cube: object | None = None + self._pending_target_transform: torch.Tensor | None = None + self._ik_model: object | None = None + self._ik_solver: NewtonChainIK | None = None + self._ik_controller: IKGizmoController | None = None + self._robot_adapter: _RobotGizmoAdapter | None = None + self.target: BatchEntity | None = None + self._target_type = "" + self._attach_target(target) + + def _attach_target(self, target: BatchEntity) -> None: + num_instances = int(getattr(target, "num_instances", dexsim.get_world_num())) + if num_instances > 1: raise RuntimeError( - f"Gizmo can only be used in single environment mode (num_env=1), " - f"but current num_envs={num_envs}. Please create simulation with num_envs=1." + "Gizmo can only be used in single environment mode " + f"(num_envs=1), but target has {num_instances} instances." ) - # Use provided config or get default - if cfg is None: - cfg = self._get_default_cfg() - self.cfg = cfg + self.target = target + self._target_type = self._detect_target_type(target) + if self._target_type == "robot": + self._setup_robot_gizmo() + return + self._gizmo = self._create_gizmo(self.cfg) - self._callback = None - self._state = "active" - self._setup_gizmo_follow() + if self._target_type == "rigidobject": + self._setup_rigid_object_gizmo() + else: + self._setup_camera_gizmo() - def _detect_target_type(self, target: BatchEntity) -> str: - """Detect target type: 'rigidobject', 'robot', or 'camera' using isinstance only.""" - if Robot is not None and isinstance(target, Robot): + @staticmethod + def _detect_target_type(target: BatchEntity) -> str: + if isinstance(target, Robot): return "robot" - if Camera is not None and isinstance(target, Camera): + if isinstance(target, Camera): return "camera" - if RigidObject is not None and isinstance(target, RigidObject): + if isinstance(target, RigidObject): return "rigidobject" - raise ValueError( - f"Unsupported target type: {type(target)}. Only RigidObject, Robot, and Camera are supported." + f"Unsupported target type: {type(target)}. Only RigidObject, Robot, " + "and Camera are supported." ) - def _get_default_cfg(self) -> GizmoCfg: - """Get default gizmo configuration (same for all target types)""" - return GizmoCfg() - - def _create_gizmo(self, cfg: GizmoCfg): - """Create gizmo using configuration object""" + def _create_gizmo(self, cfg: GizmoCfg) -> object: options = cfg.to_options_dict() - axis = options["axis"] - rings = options["rings"] - return self._env.create_gizmo(axis, rings) - - def _compute_ee_pose_fk(self): - """Compute end-effector pose using forward kinematics""" - # Get current joint positions for this arm - proprioception = self.target.get_proprioception() - current_qpos_full = proprioception["qpos"] - current_joint_ids = self.target.get_joint_ids(self._robot_arm_name) - - joint_positions = current_qpos_full[:, current_joint_ids] - if joint_positions.dim() > 1: - joint_positions = joint_positions[0] - - # Compute forward kinematics - ee_pose = self.target.compute_fk( - joint_positions, name=self._control_part, to_matrix=True - ) + return self._env.create_gizmo(options["axis"], options["rings"]) - return ee_pose + def _setup_rigid_object_gizmo(self) -> None: + target = self._require_target() + target_node = target._entities[0].node + self._require_gizmo().follow(target_node) + self._require_gizmo().set_flush_localpose_callback(create_gizmo_callback()) - def _create_proxy_cube( - self, position: np.ndarray, rotation_matrix: np.ndarray, name: str - ): - """Create a proxy cube for gizmo tracking""" - # Convert rotation matrix to euler angles - euler = R.from_matrix(rotation_matrix).as_euler("xyz", degrees=False) + def _setup_robot_gizmo(self) -> None: + try: + from dexsim.kit.ik import ( + IKApplyMode, + IKGizmoController, + NewtonChainIK, + build_newton_model_from_urdf, + ) + except ImportError as error: + raise RuntimeError( + "Robot gizmo requires a dexsim build that exports " + "IKGizmoController, NewtonChainIK, and " + "build_newton_model_from_urdf." + ) from error + + target = self._require_robot() + control_parts = list(target.control_parts or {}) + if not control_parts: + raise ValueError("Robot has no control parts defined.") + if self._control_part is None: + self._control_part = control_parts[0] + if self._control_part not in control_parts: + raise ValueError( + f"Control part {self._control_part!r} is not defined. Available " + f"parts: {control_parts}." + ) - # Create small proxy cube at specified position - proxy_cube = self._env.create_cube(0.02, 0.02, 0.02) # 2cm cube - proxy_cube.set_location(position[0], position[1], position[2]) - proxy_cube.set_rotation_euler(euler[0], euler[1], euler[2]) + root_link, end_link, tcp_pose = self._resolve_robot_ik_chain(target) + if self.cfg.ik_iterations <= 0: + raise ValueError("ik_iterations must be greater than zero.") + if not np.isfinite(self.cfg.ik_gizmo_scale) or self.cfg.ik_gizmo_scale <= 0: + raise ValueError("ik_gizmo_scale must be positive and finite.") + + adapter = _RobotGizmoAdapter(target, self._control_part) + ik_device = self.cfg.ik_device or str(target.device) + with wp.ScopedDevice(ik_device): + ik_model = build_newton_model_from_urdf( + target.cfg.fpath, + hide_visuals=True, + ) + ik_solver = NewtonChainIK( + ik_model, + start_link=root_link, + end_link=end_link, + iterations=self.cfg.ik_iterations, + tcp_pose=tcp_pose, + ) - # Connect gizmo to proxy cube. - self._gizmo.follow(proxy_cube.node) + current_qpos = adapter.get_current_qpos() + ik_solver.set_qpos_from_joint_names( + adapter.get_actived_joint_names(), + current_qpos, + ) + base_pose = adapter.get_world_pose() + ik_solver.sync_target_state_from_link(adapter, base_pose) + + target_name = getattr(target.cfg, "uid", "robot") + ik_controller = IKGizmoController( + self._world, + adapter, + ik_solver, + base_state={"pose": base_pose}, + toggle_key=self.cfg.ik_toggle_key, + follow_robot_base=True, + apply_mode=IKApplyMode.DRIVE_TARGET, + gizmo_scale=self.cfg.ik_gizmo_scale, + name=f"{target_name}_{self._control_part}_ik", + ) - logger.log_info(f"{name} gizmo proxy created at position: {position}") - return proxy_cube + self._robot_adapter = adapter + self._ik_model = ik_model + self._ik_solver = ik_solver + self._ik_controller = ik_controller + self._gizmo = ik_controller.target_gizmo.gizmo + logger.log_info( + f"Robot gizmo uses dexsim Newton IK for control part " + f"{self._control_part!r} ({root_link} -> {end_link})." + ) - def _setup_camera_gizmo(self): - """Setup gizmo for Camera by creating a proxy RigidObject at camera position""" - # Get current camera pose - camera_pose = self.target.get_local_pose(to_matrix=True)[0] # Get first camera - camera_pos = camera_pose[:3, 3].cpu().numpy() - camera_rot_matrix = camera_pose[:3, :3].cpu().numpy() + def _resolve_robot_ik_chain( + self, + target: Robot, + ) -> tuple[str, str, np.ndarray]: + solver = ( + target.get_solver(self._control_part) + if target.cfg.solver_cfg is not None + else None + ) + root_link = self.cfg.ik_root_link_name or getattr( + solver, + "root_link_name", + None, + ) + end_link = self.cfg.ik_end_link_name or getattr( + solver, + "end_link_name", + None, + ) + if not root_link or not end_link: + raise ValueError( + "Robot gizmo needs an IK chain. Set GizmoCfg.ik_root_link_name " + "and ik_end_link_name, or configure a solver for the selected " + "robot control part." + ) - # Create proxy cube and set callback + tcp_pose = self.cfg.ik_tcp_pose + if tcp_pose is None and solver is not None: + tcp_pose = solver.get_tcp() + if tcp_pose is None: + tcp_pose = np.eye(4, dtype=np.float32) + if isinstance(tcp_pose, torch.Tensor): + tcp_pose = tcp_pose.detach().cpu().numpy() + return root_link, end_link, np.asarray(tcp_pose, dtype=np.float32) + + def _setup_camera_gizmo(self) -> None: + target = self._require_target() + camera_pose = target.get_local_pose(to_matrix=True)[0] + camera_pos = camera_pose[:3, 3].detach().cpu().numpy() + camera_rotation = camera_pose[:3, :3].detach().cpu().numpy() self._proxy_cube = self._create_proxy_cube( - camera_pos, camera_rot_matrix, "Camera" + camera_pos, + camera_rotation, + "Camera", ) - # New API uses set_flush_localpose_callback - self._gizmo.set_flush_localpose_callback(self._proxy_gizmo_callback) + self._require_gizmo().set_flush_localpose_callback(self._proxy_gizmo_callback) - def _proxy_gizmo_callback(self, *args): - """Generic callback for proxy-based gizmo. + def _create_proxy_cube( + self, + position: np.ndarray, + rotation_matrix: np.ndarray, + name: str, + ) -> object: + euler = R.from_matrix(rotation_matrix).as_euler("xyz", degrees=False) + proxy_cube = self._env.create_cube(0.02, 0.02, 0.02) + proxy_cube.set_location(*position) + proxy_cube.set_rotation_euler(*euler) + self._require_gizmo().follow(proxy_cube.node) + logger.log_info(f"{name} gizmo proxy created at position: {position}.") + return proxy_cube - Supports both old signature: (node, translation, rotation, flag) - and new signature: (node, local_pose, flag) where local_pose is a 4x4 matrix. - Updates the proxy cube transform and sets `_pending_target_transform`. - """ - # New API callback signature: (node, local_pose, flag) - if len(args) != 3: + def _proxy_gizmo_callback(self, *args: object) -> None: + if len(args) != 3 or self._proxy_cube is None: return node, local_pose, flag = args if node is None: return - # Check if proxy cube still exists - if not hasattr(self, "_proxy_cube") or self._proxy_cube is None: - return - - # convert to numpy 4x4 matrix if isinstance(local_pose, torch.Tensor): - lp = local_pose.cpu().numpy() + pose = local_pose.detach().cpu().numpy() else: - lp = np.asarray(local_pose) - - if lp.shape != (4, 4): + pose = np.asarray(local_pose) + if pose.shape != (4, 4): return - trans = lp[:3, 3] - rot_mat = lp[:3, :3] - euler = R.from_matrix(rot_mat).as_euler("xyz", degrees=False) - - self._proxy_cube.set_location(float(trans[0]), float(trans[1]), float(trans[2])) - self._proxy_cube.set_rotation_euler( - float(euler[0]), float(euler[1]), float(euler[2]) - ) - - # Build pending target transform (1,4,4) - target_transform = torch.eye(4, dtype=torch.float32) - target_transform[:3, 3] = torch.tensor( - [trans[0], trans[1], trans[2]], dtype=torch.float32 - ) - target_transform[:3, :3] = torch.tensor(rot_mat, dtype=torch.float32) - self._pending_target_transform = target_transform.unsqueeze(0) - - def _update_camera_pose(self, target_transform: torch.Tensor): - """Update camera pose to match target transform""" + node.set_transform(pose, flag) + position = pose[:3, 3] + euler = R.from_matrix(pose[:3, :3]).as_euler("xyz", degrees=False) + self._proxy_cube.set_location(*position) + self._proxy_cube.set_rotation_euler(*euler) + self._pending_target_transform = torch.as_tensor( + pose, + dtype=torch.float32, + ).unsqueeze(0) + + def _update_camera_pose(self, target_transform: torch.Tensor) -> bool: try: - # Set camera pose using set_local_pose method - self.target.set_local_pose(target_transform) + self._require_target().set_local_pose(target_transform) return True - except Exception as e: - logger.log_error(f"Error updating camera pose: {e}") + except Exception as error: + logger.log_error(f"Error updating camera pose: {error}") return False - def _setup_robot_gizmo(self): - """Setup gizmo for Robot by creating a proxy RigidObject at end-effector""" - # Get end-effector pose using specified control part - if self.target.cfg.solver_cfg is None: - raise ValueError( - "Robot has no solver configured for IK/FK computations for gizmo" - ) - - arm_names = list(self.target.control_parts.keys()) - if not arm_names: - raise ValueError("Robot has no control parts defined") + def attach(self, target: BatchEntity) -> None: + """Attach this controller to another supported single-instance target.""" + self._release_resources() + self._attach_target(target) - # Use specified control part or fall back to first available - if self._control_part and self._control_part in arm_names: - self._robot_arm_name = self._control_part - else: - logger.log_error(f"Control part '{self._control_part}' not found.") - - logger.log_info(f"Using control part: {self._robot_arm_name}") - - # Get end-effector pose using forward kinematics - ee_pose = self._compute_ee_pose_fk()[0] # remove batch dimension - - ee_pos = ee_pose[:3, 3].cpu().numpy() - ee_rot_matrix = ee_pose[:3, :3].cpu().numpy() - - # Create proxy cube and set callback (use new callback API) - self._proxy_cube = self._create_proxy_cube(ee_pos, ee_rot_matrix, "Robot") - self._gizmo.set_flush_localpose_callback(self._proxy_gizmo_callback) - - def _update_robot_ik(self, target_transform: torch.Tensor): - """Update robot joints using IK to reach target transform""" - try: - # Get current joint positions as seed using proprioception - proprioception = self.target.get_proprioception() - current_qpos_full = proprioception["qpos"] # Full joint positions - - # Get joint IDs for this arm - current_joint_ids = self.target.get_joint_ids(self._robot_arm_name) - - # Extract joint positions for this specific arm - if len(current_joint_ids) > 0: - joint_seed = current_qpos_full[ - :, current_joint_ids - ] # Select arm joints - if joint_seed.dim() > 1: - joint_seed = joint_seed[0] # Take first batch element - else: - logger.log_warning( - f"No joint IDs found for arm: {self._robot_arm_name}" - ) - return False - - # Solve IK - ik_success, new_qpos = self.target.compute_ik( - pose=target_transform, name=self._robot_arm_name, joint_seed=joint_seed - ) - - if ik_success: - # Ensure correct dimensions for setting qpos - # new_qpos from IK solver may be (1, N, dof) or (N, dof), flatten to (dof,) for single env - if new_qpos.dim() > 1: - new_qpos = new_qpos.squeeze() # Remove all singleton dimensions - if new_qpos.dim() == 1: - new_qpos = new_qpos.unsqueeze(0) # Make it (1, dof) for set_qpos - - # Update robot joint positions - self.target.set_qpos(qpos=new_qpos[0], joint_ids=current_joint_ids) - return True - else: - logger.log_warning("IK solution not found") - return False - - except Exception as e: - logger.log_error(f"Error in robot IK: {e}") - return False - - def _setup_gizmo_follow(self): - """Setup gizmo based on target type""" - if self._target_type == "rigidobject": - # RigidObject: direct node access through MeshObject — use follow/attach - tgt_node = self.target._entities[0].node - self._gizmo.follow(tgt_node) - # set callback (localpose-style) - self._gizmo.set_flush_localpose_callback(create_gizmo_callback()) - - elif self._target_type == "robot": - # Robot: create proxy object at end-effector position - self._setup_robot_gizmo() - elif self._target_type == "camera": - # Camera: create proxy object at camera position - self._setup_camera_gizmo() - - def attach(self, target: BatchEntity): - """Attach gizmo to a new simulation element.""" - self.target = target - self._target_type = self._detect_target_type(target) - self._setup_gizmo_follow() - - def detach(self): - """Detach gizmo from current element.""" + def detach(self) -> None: + """Detach the gizmo and release target-specific controller resources.""" + self._release_resources() self.target = None - # Detach gizmo using new API - self._gizmo.detach_parent() + self._target_type = "" - def set_transform_callback(self, callback: Callable): - """Set callback for gizmo transform events (translation/rotation).""" + def set_transform_callback(self, callback: Callable[..., Any]) -> None: + """Set an additional raw gizmo transform callback.""" self._callback = callback - self._gizmo.set_transform_flush_callback(callback) + self._require_gizmo().set_transform_flush_callback(callback) - def set_world_pose(self, pose): - """Set gizmo's world pose.""" - self._gizmo.set_world_pose(pose) + def set_world_pose(self, pose: np.ndarray) -> None: + """Set the underlying gizmo's world pose.""" + self._require_gizmo().set_world_pose(pose) - def set_local_pose(self, pose): - """Set gizmo's local pose.""" - self._gizmo.set_local_pose(pose) + def set_local_pose(self, pose: np.ndarray) -> None: + """Set the underlying gizmo's local pose.""" + self._require_gizmo().set_local_pose(pose) - def set_line_width(self, width: float): - """Set gizmo line width.""" - self._gizmo.set_line_width(width) + def set_line_width(self, width: float) -> None: + """Set the underlying gizmo line width.""" + self._require_gizmo().set_line_width(width) - def enable_collision(self, enabled: bool): + def enable_collision(self, enabled: bool) -> None: """Enable or disable gizmo collision.""" - self._gizmo.enable_collision(enabled) + self._require_gizmo().enable_collision(enabled) - def get_world_pose(self): - """Get gizmo's world pose.""" - return self._gizmo.get_world_pose() + def get_world_pose(self) -> np.ndarray: + """Return the underlying gizmo's world pose.""" + return self._require_gizmo().get_world_pose() - def get_local_pose(self): - """Get gizmo's local pose.""" - return self._gizmo.get_local_pose() + def get_local_pose(self) -> np.ndarray: + """Return the underlying gizmo's local pose.""" + return self._require_gizmo().get_local_pose() - def get_name(self): - """Get gizmo node name.""" - return self._gizmo.get_name() + def get_name(self) -> str: + """Return the underlying gizmo name.""" + return self._require_gizmo().get_name() - def get_parent(self): - """Get gizmo's parent node.""" - return self._gizmo.get_parent() + def get_parent(self) -> object: + """Return the underlying gizmo parent.""" + return self._require_gizmo().get_parent() def toggle_visibility(self) -> bool: - """ - Toggle the visibility of the gizmo. - - Returns: - bool: The new visibility state (True = visible, False = hidden) - """ - if not hasattr(self, "_is_visible"): - self._is_visible = True # Default to visible - - # Toggle the state - self._is_visible = not self._is_visible - - # Apply the visibility setting to the gizmo node - if self._gizmo: - self._gizmo.set_visible(self._is_visible) - - return self._is_visible - - def set_visible(self, visible: bool): - """ - Set the visibility of the gizmo. - - Args: - visible (bool): True to show, False to hide the gizmo - """ - self._is_visible = visible - - # Apply the visibility setting to the gizmo node - if self._gizmo: - self._gizmo.set_visible(self._is_visible) + """Toggle gizmo visibility and return the new state.""" + visible = not self.is_visible() + self.set_visible(visible) + return visible + + def set_visible(self, visible: bool) -> None: + """Set gizmo visibility.""" + self._is_visible = bool(visible) + if self._ik_controller is not None: + self._ik_controller.enabled = self._is_visible + gizmo = self._gizmo + if gizmo is not None: + gizmo.set_visible(self._is_visible) def is_visible(self) -> bool: - """ - Check if the gizmo is currently visible. - - Returns: - bool: True if visible, False if hidden - """ - return getattr(self, "_is_visible", True) + """Return whether the gizmo is visible.""" + if self._ik_controller is not None: + return bool(self._ik_controller.enabled) + return self._is_visible - def update(self): - """Synchronize gizmo with target's current transform, and handle IK solving here.""" + def update(self) -> None: + """Synchronize the gizmo and apply pending target changes.""" + if self.target is None: + return if self._target_type == "rigidobject": - tgt_node = self.target._entities[0].node - self._gizmo.follow(tgt_node) - + target_node = self.target._entities[0].node + self._require_gizmo().follow(target_node) elif self._target_type == "robot": - # If there is a pending target, solve IK and clear it - if ( - hasattr(self, "_pending_target_transform") - and self._pending_target_transform is not None - ): - self._update_robot_ik(self._pending_target_transform) - self._pending_target_transform = None + if self._ik_controller is not None: + self._ik_controller.update(iterations=self.cfg.ik_iterations) elif self._target_type == "camera": - # Update proxy cube position to match current camera pose - if hasattr(self, "_proxy_cube") and self._proxy_cube: + if self._proxy_cube is not None: camera_pose = self.target.get_local_pose(to_matrix=True)[0] - camera_pos = camera_pose[:3, 3].cpu().numpy() - self._proxy_cube.set_location( - camera_pos[0], camera_pos[1], camera_pos[2] - ) - - # If there is a pending camera target, update camera pose and clear it - if ( - hasattr(self, "_pending_target_transform") - and self._pending_target_transform is not None - ): + position = camera_pose[:3, 3].detach().cpu().numpy() + self._proxy_cube.set_location(*position) + if self._pending_target_transform is not None: self._update_camera_pose(self._pending_target_transform) self._pending_target_transform = None - def apply_transform(self, translation, rotation): - """Apply transform based on target type""" + def apply_transform( + self, + translation: np.ndarray, + rotation: np.ndarray, + ) -> None: + """Apply a direct transform where the target path supports it.""" + if self.target is None: + return if self._target_type == "rigidobject": self.target.set_location(*translation) self.target.set_rotation_euler(*rotation) - elif self._target_type == "robot": - # Robot transforms are handled by IK in the gizmo callback - if hasattr(self, "_proxy_cube") and self._proxy_cube: - self._proxy_cube.set_location(*translation) - self._proxy_cube.set_rotation_euler(*rotation) - elif self._target_type == "camera": - # Camera transforms are handled by pose update in the gizmo callback - if hasattr(self, "_proxy_cube") and self._proxy_cube: - self._proxy_cube.set_location(*translation) - self._proxy_cube.set_rotation_euler(*rotation) - else: - # Other target types - pass + elif self._target_type == "camera" and self._proxy_cube is not None: + self._proxy_cube.set_location(*translation) + self._proxy_cube.set_rotation_euler(*rotation) - def destroy(self): - """Clean up gizmo resources and release references.""" - # Clear transform callback first to avoid bad_function_call - if hasattr(self, "_gizmo") and self._gizmo and hasattr(self._gizmo, "node"): + def destroy(self) -> None: + """Release gizmo resources and target references.""" + self._release_resources() + self.target = None + self._target_type = "" + + def _release_resources(self) -> None: + gizmo = self._gizmo + if gizmo is not None: + for method_name in ( + "set_flush_localpose_callback", + "set_transform_flush_callback", + ): + method = getattr(gizmo, method_name, None) + if callable(method): + try: + method(None) + except (TypeError, RuntimeError): + pass try: - # Clear transform callback before any other cleanup - self._gizmo.node.set_flush_transform_callback(None) - logger.log_info("Cleared gizmo transform callback") - except Exception as e: - logger.log_warning(f"Failed to clear gizmo callback: {e}") - - # Remove proxy cube if exists (before detaching gizmo) - if hasattr(self, "_proxy_cube") and self._proxy_cube: + gizmo.set_visible(False) + except (AttributeError, TypeError, RuntimeError): + pass try: - # Detach gizmo from proxy cube first - if ( - hasattr(self, "_gizmo") - and self._gizmo - and hasattr(self._gizmo, "node") - ): - self._gizmo.detach_parent() - # Then remove the proxy cube - self._env.remove_actor(self._proxy_cube) - logger.log_info("Successfully removed proxy cube from environment") - except Exception as e: - logger.log_warning(f"Failed to remove proxy cube: {e}") - self._proxy_cube = None + gizmo.detach_parent() + except (AttributeError, TypeError, RuntimeError): + pass + + if self._ik_controller is not None: + target_node = self._ik_controller.target_gizmo.target_node + try: + target_node.detach_parent() + except (AttributeError, TypeError, RuntimeError): + pass - # Final gizmo cleanup - if hasattr(self, "_gizmo") and self._gizmo and hasattr(self._gizmo, "node"): + if self._proxy_cube is not None: try: - # Ensure detach_parent is called if not done above - if self._target_type in ["robot", "camera"]: - pass # Already detached above - else: - self._gizmo.node.detach_parent() - logger.log_info("Successfully cleaned up gizmo node") - except Exception as e: - logger.log_warning(f"Failed to cleanup gizmo node: {e}") - - # Clear pending transform - if hasattr(self, "_pending_target_transform"): - self._pending_target_transform = None - - # Directly release references + self._env.remove_actor(self._proxy_cube) + except (AttributeError, TypeError, RuntimeError) as error: + logger.log_warning(f"Failed to remove gizmo proxy cube: {error}") + + if gizmo is not None: + remove_gizmo = getattr(self._env, "remove_gizmo", None) + if callable(remove_gizmo): + try: + remove_gizmo(gizmo) + except (AttributeError, TypeError, RuntimeError) as error: + logger.log_warning( + f"Failed to remove gizmo from dexsim environment: {error}" + ) + + self._pending_target_transform = None + self._proxy_cube = None self._gizmo = None - self.target = None + self._ik_controller = None + self._ik_solver = None + self._ik_model = None + self._robot_adapter = None + + def _require_gizmo(self) -> object: + if self._gizmo is None: + raise RuntimeError("Gizmo is not attached.") + return self._gizmo + + def _require_target(self) -> BatchEntity: + if self.target is None: + raise RuntimeError("Gizmo has no target.") + return self.target + + def _require_robot(self) -> Robot: + target = self._require_target() + if not isinstance(target, Robot): + raise TypeError(f"Expected Robot target, got {type(target)}.") + return target diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 86be05455..65a65adaf 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -32,7 +32,7 @@ from copy import deepcopy from datetime import datetime from functools import cached_property -from typing import Callable, Dict, List, Sequence, Union +from typing import TYPE_CHECKING, Callable, Dict, List, Sequence, Union from dataclasses import dataclass, asdict, field, MISSING # Global cache directories @@ -66,7 +66,7 @@ Light, RigidConstraint, ) -from embodichain.lab.sim.objects.gizmo import Gizmo +from embodichain.lab.sim.objects.gizmo import Gizmo, GizmoCfg from embodichain.lab.sim.sensors import ( SensorCfg, BaseSensor, @@ -94,6 +94,9 @@ from embodichain.utils import configclass, logger from embodichain.utils.math import look_at_to_pose, pose_inv +if TYPE_CHECKING: + from dexsim.interaction import EntityGizmoConfig, EntityGizmoManipulator + __all__ = [ "SimulationManager", "SimulationManagerCfg", @@ -159,6 +162,9 @@ class SimulationManagerCfg: window_camera_pose: WindowCameraPoseCfg = field(default_factory=WindowCameraPoseCfg) """Interactive viewer camera-pose printing settings.""" + enable_entity_gizmo_on_window_open: bool = True + """Whether opening a viewer window automatically enables entity gizmo control.""" + @dataclass class _WindowRecordState: @@ -200,6 +206,7 @@ class SimulationManager: _instances = {} _cleanup_queue: queue.Queue = queue.Queue() + _DEFAULT_PLANE_GIZMO_TARGET_ID = (1 << 64) - 1 SUPPORTED_SENSOR_TYPES = { "Camera": Camera, @@ -249,6 +256,7 @@ def __init__( self._world: dexsim.World = dexsim.World(world_config) self._window: Windows | None = None + self._entity_gizmo_config: EntityGizmoConfig | None = None self._window_record_state: _WindowRecordState | None = None self._window_record_camera: object | None = None wr = sim_config.window_record @@ -326,6 +334,7 @@ def __init__( if sim_config.headless is False: self._window = self._world.get_windows() + self._on_window_opened() @classmethod def get_instance(cls, instance_id: int = 0) -> SimulationManager: @@ -616,11 +625,43 @@ def get_env(self, arena_index: int = -1) -> dexsim.environment.Arena: def get_world(self) -> dexsim.World: return self._world - def open_window(self) -> None: - """Open the simulation window.""" - self._world.open_window() + def open_window( + self, + *, + enable_entity_gizmo: bool | None = None, + entity_gizmo_config: EntityGizmoConfig | None = None, + ) -> None: + """Open the simulation window and initialize its interaction controls. + + Entity gizmo control is enabled by default. Set + ``enable_entity_gizmo=False`` for a view-only window. When the argument + is omitted, :attr:`SimulationManagerCfg.enable_entity_gizmo_on_window_open` + determines the behavior. + + Args: + enable_entity_gizmo: Whether to enable world-level entity gizmo + control for this window. ``None`` uses the simulation + configuration default. + entity_gizmo_config: Optional native dexsim configuration. Passing + a configuration implies entity gizmo control unless explicitly + disabled. + """ + if not self.is_window_opened or self._window is None: + self._world.open_window() self._window = self._world.get_windows() + self.is_window_opened = True + self._on_window_opened( + enable_entity_gizmo=enable_entity_gizmo, + entity_gizmo_config=entity_gizmo_config, + ) + def _on_window_opened( + self, + *, + enable_entity_gizmo: bool | None = None, + entity_gizmo_config: EntityGizmoConfig | None = None, + ) -> None: + """Initialize controls shared by constructor-opened and reopened windows.""" if ( self._window_record_hotkey_cfg is not None and self._window_record_input_control is None @@ -631,10 +672,31 @@ def open_window(self) -> None: and self._window_camera_pose_input_control is None ): self.enable_window_camera_pose_hotkey(**self._window_camera_pose_hotkey_cfg) - self.is_window_opened = True + + if enable_entity_gizmo is None: + enable_entity_gizmo = entity_gizmo_config is not None or getattr( + self.sim_config, + "enable_entity_gizmo_on_window_open", + True, + ) + + try: + if enable_entity_gizmo: + if entity_gizmo_config is not None: + self.enable_entity_gizmo(entity_gizmo_config) + elif not self.has_entity_gizmo(): + self.enable_entity_gizmo(self._entity_gizmo_config) + elif self.has_entity_gizmo(): + self.disable_entity_gizmo() + except RuntimeError as error: + logger.log_warning( + f"Entity gizmo control could not be initialized for the window: {error}" + ) def close_window(self) -> None: """Close the simulation window.""" + if self.has_entity_gizmo(): + self.disable_entity_gizmo() if self.is_window_recording(): self.stop_window_record() self._world.close_window() @@ -1620,15 +1682,140 @@ def get_robot_uid_list(self) -> List[str]: """ return list(self._robots.keys()) + def enable_entity_gizmo( + self, + config: EntityGizmoConfig | None = None, + ) -> EntityGizmoManipulator: + """Enable dexsim's world-level entity gizmo controller. + + This is a thin lifecycle wrapper around + :meth:`dexsim.World.enable_entity_gizmo`. The returned controller owns + window selection, hotkey handling, multiple gizmo bindings, temporary + physics-state changes, and rigid-body/articulation-root manipulation. + + Args: + config: Native dexsim entity-gizmo configuration. When omitted, + dexsim's defaults are used. + + Returns: + The world-owned dexsim entity gizmo manipulator. + + Raises: + RuntimeError: If the installed dexsim build does not provide entity + gizmo support or fails to create the controller. + """ + world = getattr(self, "_world", None) + enable = getattr(world, "enable_entity_gizmo", None) + if not callable(enable): + raise RuntimeError( + "The installed dexsim build does not provide " + "World.enable_entity_gizmo()." + ) + + controller = enable() if config is None else enable(config) + if controller is None: + raise RuntimeError("dexsim failed to enable the entity gizmo controller.") + self._exclude_default_plane_from_entity_gizmo(controller) + self._entity_gizmo_config = config + logger.log_info("Dexsim entity gizmo control enabled.") + return controller + + def _exclude_default_plane_from_entity_gizmo( + self, + controller: EntityGizmoManipulator, + ) -> None: + """Register the EmbodiChain ground as an immovable gizmo target. + + dexsim resolves registered external targets before its generic + render-mesh path. Registering the default plane as a static rigid body + therefore makes both raycast toggles and programmatic attachment return + ``STATIC_RIGID_BODY`` without adding physics to the visual plane. + """ + default_plane = getattr(self, "_default_plane", None) + register = getattr(controller, "register_external_target", None) + if default_plane is None: + return + if not callable(register): + logger.log_warning( + "The installed dexsim build cannot exclude the default plane " + "from entity gizmo control." + ) + return + + try: + result = register( + self._DEFAULT_PLANE_GIZMO_TARGET_ID, + dexsim.interaction.EntityGizmoTargetType.RIGID_BODY, + default_plane, + ActorType.STATIC, + ) + except (AttributeError, TypeError, RuntimeError) as error: + logger.log_warning( + "Failed to exclude the default plane from entity gizmo " + f"control: {error}." + ) + return + if result != dexsim.interaction.EntityGizmoResult.SUCCESS: + logger.log_warning( + "Failed to exclude the default plane from entity gizmo " + f"control: {result}." + ) + + def disable_entity_gizmo(self) -> bool: + """Disable dexsim's world-level entity gizmo controller. + + Returns: + ``True`` when an active controller was disabled, or ``False`` when + entity gizmo control was already disabled. + + Raises: + RuntimeError: If the installed dexsim build does not provide entity + gizmo support. + """ + world = getattr(self, "_world", None) + get_controller = getattr(world, "get_entity_gizmo", None) + disable = getattr(world, "disable_entity_gizmo", None) + if not callable(get_controller) or not callable(disable): + raise RuntimeError( + "The installed dexsim build does not provide entity gizmo " + "lifecycle APIs." + ) + if get_controller() is None: + return False + + disable() + logger.log_info("Dexsim entity gizmo control disabled.") + return True + + def get_entity_gizmo(self) -> EntityGizmoManipulator | None: + """Return dexsim's active world-level entity gizmo controller.""" + world = getattr(self, "_world", None) + get_controller = getattr(world, "get_entity_gizmo", None) + if not callable(get_controller): + raise RuntimeError( + "The installed dexsim build does not provide " + "World.get_entity_gizmo()." + ) + return get_controller() + + def has_entity_gizmo(self) -> bool: + """Return whether world-level entity gizmo control is enabled.""" + world = getattr(self, "_world", None) + get_controller = getattr(world, "get_entity_gizmo", None) + return callable(get_controller) and get_controller() is not None + def enable_gizmo( - self, uid: str, control_part: str | None = None, gizmo_cfg: object = None - ) -> Gizmo: + self, + uid: str, + control_part: str | None = None, + gizmo_cfg: GizmoCfg | None = None, + ) -> Gizmo | None: """Enable gizmo control for any simulation object (Robot, RigidObject, Camera, etc.). Args: uid (str): UID of the object to attach gizmo to (searches in robots, rigid_objects, sensors, etc.) control_part (str | None, optional): Control part name for robots. Defaults to "arm". - gizmo_cfg (object, optional): Gizmo configuration object. Defaults to None. + gizmo_cfg: Gizmo configuration. Defaults to None. """ # Create gizmo key combining uid and control_part gizmo_key = f"{uid}:{control_part}" if control_part else uid @@ -1638,7 +1825,7 @@ def enable_gizmo( logger.log_warning( f"Gizmo for '{uid}' with control_part '{control_part}' already exists." ) - return + return self._gizmos[gizmo_key] # Search for target object in different collections target = None @@ -1656,17 +1843,13 @@ def enable_gizmo( else: logger.log_error( - f"Object with uid '{uid}' not found in any collection (robots, rigid_objects, sensors, articulations)." + f"Object with uid '{uid}' not found in any supported collection " + "(robots, rigid_objects, sensors)." ) - return + return None + gizmo = None try: - gizmo = Gizmo(target, gizmo_cfg, control_part) - self._gizmos[gizmo_key] = gizmo - logger.log_info( - f"Gizmo enabled for {object_type} '{uid}' with control_part '{control_part}'" - ) - # Initialize GizmoController if not already done. if not hasattr(self, "_gizmo_controller") or self._gizmo_controller is None: window = ( @@ -1674,9 +1857,17 @@ def enable_gizmo( if hasattr(self._world, "get_windows") else None ) + if window is None: + raise RuntimeError("Gizmo requires a simulation window.") self._gizmo_controller = GizmoController() window.add_input_control(self._gizmo_controller) + gizmo = Gizmo(target, gizmo_cfg, control_part) + self._gizmos[gizmo_key] = gizmo + logger.log_info( + f"Gizmo enabled for {object_type} '{uid}' with control_part '{control_part}'" + ) + except Exception as e: logger.log_error( f"Failed to create gizmo for {object_type} '{uid}' with control_part '{control_part}': {e}" @@ -2699,6 +2890,9 @@ def destroy(self, exit_process: bool | None = None) -> None: def _deferred_destroy(self) -> None: """Destroy all simulated assets and release resources.""" + if self.has_entity_gizmo(): + self.disable_entity_gizmo() + # Clean up all gizmos before destroying the simulation for uid in list(self._gizmos.keys()): self.disable_gizmo(uid) diff --git a/embodichain/lab/sim/utility/gizmo_utils.py b/embodichain/lab/sim/utility/gizmo_utils.py index 3ff1c7de7..177612023 100644 --- a/embodichain/lab/sim/utility/gizmo_utils.py +++ b/embodichain/lab/sim/utility/gizmo_utils.py @@ -14,31 +14,33 @@ # limitations under the License. # ---------------------------------------------------------------------------- -""" -Gizmo utility functions for EmbodiSim. +"""Gizmo utility functions for EmbodiChain. This module provides utility functions for creating gizmo transform callbacks. """ -from typing import Callable -from typing import TYPE_CHECKING -from dexsim.types import TransformMask +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from embodichain.lab.sim.objects import Robot +__all__ = ["create_gizmo_callback", "run_gizmo_robot_control_loop"] + -def create_gizmo_callback() -> Callable: +def create_gizmo_callback() -> Callable[[Any, Any, Any], None]: """Create a standard gizmo transform callback function. This callback handles local pose for gizmo controls. It applies transformations directly to the node when gizmo controls are manipulated. Returns: - Callable: A callback function that can be used with gizmo.node.set_flush_transform_callback() + A callback compatible with dexsim's gizmo local-pose flush hook. """ - def gizmo_transform_callback(node, local_pose, flag): + def gizmo_transform_callback(node: Any, local_pose: Any, flag: Any) -> None: if node is not None: node.set_transform(local_pose, flag) @@ -46,8 +48,10 @@ def gizmo_transform_callback(node, local_pose, flag): def run_gizmo_robot_control_loop( - robot: object | str, control_part: str = "arm", end_link_name: str | None = None -): + robot: Robot | str, + control_part: str = "arm", + end_link_name: str | None = None, +) -> None: """Run a control loop for testing gizmo controls on a robot. This function implements a control loop that allows users to manipulate a robot @@ -75,38 +79,62 @@ def run_gizmo_robot_control_loop( np.set_printoptions(precision=5, suppress=True) from embodichain.lab.sim import SimulationManager - from embodichain.lab.sim.objects import Robot - from embodichain.lab.sim.solvers import PinkSolverCfg + from embodichain.lab.sim.objects import GizmoCfg - from embodichain.utils.logger import log_info, log_warning, log_error + from embodichain.utils.logger import log_error, log_info sim = SimulationManager.get_instance() if isinstance(robot, str): - robot = sim.get_robot(uid=robot) + robot_uid = robot + robot = sim.get_robot(uid=robot_uid) + if robot is None: + log_error(f"Robot {robot_uid!r} was not found.") + return # Enter auto-update mode. sim.set_manual_update(False) - # Replace robot's default solver with PinkSolver for gizmo control. - robot_solver = robot.get_solver(name=control_part) + # Resolve only the chain metadata. dexsim owns the Newton IK solver and + # writes its drive targets back through the EmbodiChain Robot API. + robot_solver = ( + robot.get_solver(name=control_part) + if robot.cfg.solver_cfg is not None + else None + ) control_part_link_names = robot.get_control_part_link_names(name=control_part) + if not control_part_link_names: + raise ValueError(f"Control part {control_part!r} has no links.") + root_link_name = ( + robot_solver.root_link_name + if robot_solver is not None + else control_part_link_names[0] + ) end_link_name = ( - control_part_link_names[-1] if end_link_name is None else end_link_name + ( + robot_solver.end_link_name + if robot_solver is not None + else control_part_link_names[-1] + ) + if end_link_name is None + else end_link_name ) - pink_solver_cfg = PinkSolverCfg( - urdf_path=robot.cfg.fpath, - end_link_name=end_link_name, - root_link_name=robot_solver.root_link_name, - pos_eps=1e-2, - rot_eps=5e-2, - max_iterations=300, - dt=0.1, + tcp_pose = robot_solver.get_tcp() if robot_solver is not None else None + gizmo_cfg = GizmoCfg( + ik_root_link_name=root_link_name, + ik_end_link_name=end_link_name, + ik_tcp_pose=tcp_pose, ) - robot.init_solver(cfg={control_part: pink_solver_cfg}) # Enable gizmo for the robot - gizmo = sim.enable_gizmo(uid=robot.uid, control_part=control_part) + gizmo = sim.enable_gizmo( + uid=robot.uid, + control_part=control_part, + gizmo_cfg=gizmo_cfg, + ) + if gizmo is None: + log_error(f"Failed to enable gizmo for control part {control_part!r}.") + return # Store initial robot configuration initial_qpos = robot.get_qpos(name=control_part) @@ -127,7 +155,7 @@ def run_gizmo_robot_control_loop( old_settings = termios.tcgetattr(sys.stdin) tty.setcbreak(sys.stdin.fileno()) - def get_key(): + def get_key() -> str | None: """Non-blocking keyboard input.""" if select.select([sys.stdin], [], [], 0)[0]: return sys.stdin.read(1) @@ -146,29 +174,20 @@ def get_key(): if key in ["q", "Q", "\x1b"]: # Q or ESC log_info("Exiting gizmo control loop...") sim.disable_gizmo(uid=robot.uid, control_part=control_part) - if robot_solver: - robot.init_solver( - cfg={control_part: robot_solver.cfg} - ) # Restore original solver break # Print robot state elif key in ["p", "P"]: current_qpos = robot.get_qpos(name=control_part) - eef_pose = robot.compute_fk(name=control_part, qpos=current_qpos) + eef_pose = robot.get_link_pose(end_link_name, to_matrix=True) + if tcp_pose is not None: + tcp_tensor = np.asarray(tcp_pose, dtype=np.float32) + eef_pose = eef_pose @ eef_pose.new_tensor(tcp_tensor) log_info(f"\n=== Robot State ===") log_info(f"Control part: {control_part}") log_info(f"Joint positions: {current_qpos.squeeze().tolist()}") - log_info(f"End-effector pose:\n{eef_pose.squeeze().numpy()}") - - if eef_pose is None: - log_info( - "End-effector pose unavailable: compute_fk returned None " - f"for control part '{control_part}'." - ) - else: - eef_pose_np = eef_pose.detach().cpu().numpy().squeeze() - log_info(f"End-effector pose:\n{eef_pose_np}") + eef_pose_np = eef_pose.detach().cpu().numpy().squeeze() + log_info(f"End-effector pose:\n{eef_pose_np}") elif key in ["g", "G"]: if gizmo_visible: sim.set_gizmo_visibility( @@ -189,7 +208,11 @@ def get_key(): sim.disable_gizmo(uid=robot.uid, control_part=control_part) robot.clear_dynamics() robot.set_qpos(qpos=initial_qpos, name=control_part, target=False) - sim.enable_gizmo(uid=robot.uid, control_part=control_part) + sim.enable_gizmo( + uid=robot.uid, + control_part=control_part, + gizmo_cfg=gizmo_cfg, + ) log_info("Robot reset to initial pose") # Print info @@ -206,10 +229,6 @@ def get_key(): except KeyboardInterrupt: sim.disable_gizmo(uid=robot.uid, control_part=control_part) - if robot_solver: - robot.init_solver( - cfg={control_part: robot_solver.cfg} - ) # Restore original solver log_info("\nControl loop interrupted by user (Ctrl+C)") finally: diff --git a/examples/sim/gizmo/gizmo_object.py b/examples/sim/gizmo/gizmo_object.py index b0931f241..62a9f123b 100644 --- a/examples/sim/gizmo/gizmo_object.py +++ b/examples/sim/gizmo/gizmo_object.py @@ -14,14 +14,15 @@ # limitations under the License. # ---------------------------------------------------------------------------- -""" -This script demonstrates how to create a simulation scene using SimulationManager. -It shows the basic setup of simulation context, adding objects, and sensors. -""" +"""Manipulate raycast-selected entities with dexsim's world-level gizmo.""" + +from __future__ import annotations import argparse import time +import dexsim + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg from embodichain.lab.sim.shapes import CubeCfg @@ -60,7 +61,7 @@ def main(): cfg=RigidObjectCfg( uid="cube1", shape=CubeCfg(size=[0.1, 0.1, 0.1]), - body_type="kinematic", + body_type="dynamic", attrs=RigidBodyAttributesCfg( mass=1.0, dynamic_friction=0.5, @@ -85,24 +86,20 @@ def main(): ) ) - # Enable Gizmo for both cubes using the new API (only in window mode) + # Opening a window enables the world-level controller by default. Passing + # a config here reconfigures it for unlimited simultaneous bindings. if not args.headless: - sim.enable_gizmo(uid="cube1") - sim.enable_gizmo(uid="cube2") + gizmo_config = dexsim.interaction.EntityGizmoConfig() + gizmo_config.max_gizmos = 0 + sim.open_window(entity_gizmo_config=gizmo_config) logger.log_info("Scene setup complete!") logger.log_info(f"Running simulation with 1 environment(s)") if not args.headless: - if sim.has_gizmo("cube1"): - logger.log_info("Gizmo enabled for cube1 - you can drag it around!") - if sim.has_gizmo("cube2"): - logger.log_info("Gizmo enabled for cube2 - you can drag it around!") + logger.log_info("Left-click an entity and press G to attach/detach its gizmo.") + logger.log_info("Multiple selected entities can keep gizmos simultaneously.") logger.log_info("Press Ctrl+C to stop the simulation") - # Open window when the scene has been set up - if not args.headless: - sim.open_window() - # Run the simulation run_simulation(sim) @@ -113,22 +110,19 @@ def run_simulation(sim: SimulationManager): sim.init_gpu_physics() step_count = 0 - gizmo_enabled = True + gizmo_enabled = sim.has_entity_gizmo() try: last_time = time.time() last_step = 0 while True: sim.update(step=1) - # Update all gizmos if any are enabled - sim.update_gizmos() - step_count += 1 - # Disable gizmo after 200000 steps (example) + # Demonstrate programmatic cancellation after 200000 steps. if step_count == 200000 and gizmo_enabled: - logger.log_info("Disabling gizmo at step 200000") - sim.disable_gizmo("cube") + logger.log_info("Disabling entity gizmo control at step 200000") + sim.disable_entity_gizmo() gizmo_enabled = False # Print FPS every second @@ -146,6 +140,8 @@ def run_simulation(sim: SimulationManager): except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: + if sim.has_entity_gizmo(): + sim.disable_entity_gizmo() sim.destroy() logger.log_info("Simulation terminated successfully") diff --git a/examples/sim/gizmo/gizmo_robot.py b/examples/sim/gizmo/gizmo_robot.py index 6b8b9effb..06efff88a 100644 --- a/examples/sim/gizmo/gizmo_robot.py +++ b/examples/sim/gizmo/gizmo_robot.py @@ -13,9 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -""" -Gizmo-Robot Example: Test Gizmo class on a robot (UR10) -""" +"""Control a UR10 end effector with dexsim's Newton IK gizmo.""" + +from __future__ import annotations import time import torch @@ -23,15 +23,14 @@ import argparse from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.solvers import PytorchSolverCfg from embodichain.lab.sim.cfg import ( RenderCfg, RobotCfg, URDFCfg, JointDrivePropertiesCfg, ) +from embodichain.lab.sim.objects import GizmoCfg from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.solvers import PinkSolverCfg from embodichain.data import get_data_path from embodichain.utils import logger @@ -75,19 +74,6 @@ def main(): "arm": ["JOINT[0-9]"], "hand": ["FINGER[1-2]"], }, - solver_cfg={ - "arm": PytorchSolverCfg( - end_link_name="ee_link", - root_link_name="base_link", - tcp=[ - [0.0, 1.0, 0.0, 0.0], - [-1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.12], - [0.0, 0.0, 0.0, 1.0], - ], - num_samples=30, - ) - }, drive_pros=JointDrivePropertiesCfg( stiffness={"JOINT[0-9]": 1e4, "FINGER[1-2]": 1e2}, damping={"JOINT[0-9]": 1e3, "FINGER[1-2]": 1e1}, @@ -109,8 +95,23 @@ def main(): time.sleep(0.2) # Wait for a moment to ensure everything is set up - # Enable gizmo using the new API - sim.enable_gizmo(uid="ur10_gizmo_test", control_part="arm") + # The robot needs no EmbodiChain IK solver for interactive gizmo control. + # dexsim builds and owns the Newton IK chain from this metadata. + gizmo_cfg = GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ik_tcp_pose=[ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.12], + [0.0, 0.0, 0.0, 1.0], + ], + ) + sim.enable_gizmo( + uid="ur10_gizmo_test", + control_part="arm", + gizmo_cfg=gizmo_cfg, + ) if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): logger.log_error("Failed to enable gizmo!") return @@ -119,6 +120,7 @@ def main(): logger.log_info("Gizmo-Robot example started!") logger.log_info("Use the gizmo to drag the robot end-effector (EE)") + logger.log_info("Press I to show or hide the Robot TCP IK gizmo") logger.log_info("Press Ctrl+C to stop the simulation") run_simulation(sim) diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index 6d6613f9a..7d8432079 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -13,9 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -""" -Gizmo-Robot Example: Test Gizmo class on a robot (UR10) -""" +"""Control a UR10 end effector with dexsim's Newton IK gizmo.""" + +from __future__ import annotations import time import torch @@ -30,8 +30,7 @@ URDFCfg, JointDrivePropertiesCfg, ) - -from embodichain.lab.sim.solvers import PinkSolverCfg +from embodichain.lab.sim.objects import GizmoCfg from embodichain.data import get_data_path from embodichain.utils import logger @@ -68,17 +67,6 @@ def main(): components=[{"component_type": "arm", "urdf_path": urdf_path}] ), control_parts={"arm": ["Joint[1-6]"]}, - solver_cfg={ - "arm": PinkSolverCfg( - urdf_path=urdf_path, - end_link_name="ee_link", - root_link_name="base_link", - pos_eps=1e-2, - rot_eps=5e-2, - max_iterations=300, - dt=0.1, - ) - }, drive_pros=JointDrivePropertiesCfg( stiffness={"Joint[1-6]": 1e4}, damping={"Joint[1-6]": 1e3}, @@ -97,8 +85,15 @@ def main(): time.sleep(0.2) # Wait for a moment to ensure everything is set up - # Enable gizmo using the new API - sim.enable_gizmo(uid="ur10_gizmo_test", control_part="arm") + # dexsim owns the Newton IK solver used by the interactive controller. + sim.enable_gizmo( + uid="ur10_gizmo_test", + control_part="arm", + gizmo_cfg=GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ), + ) if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): logger.log_error("Failed to enable gizmo!") return @@ -107,6 +102,7 @@ def main(): logger.log_info("Gizmo-Robot example started!") logger.log_info("Use the gizmo to drag the robot end-effector (EE)") + logger.log_info("Press I to show or hide the Robot TCP IK gizmo") logger.log_info("Press Ctrl+C to stop the simulation") run_simulation(sim) diff --git a/tests/sim/objects/test_gizmo.py b/tests/sim/objects/test_gizmo.py new file mode 100644 index 000000000..faddcc717 --- /dev/null +++ b/tests/sim/objects/test_gizmo.py @@ -0,0 +1,189 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from embodichain.lab.sim.objects.gizmo import ( + Gizmo, + GizmoCfg, + _RobotGizmoAdapter, +) + + +class _FakeRobot: + """Small Robot-compatible state holder for adapter tests.""" + + def __init__(self) -> None: + self.control_parts = {"arm": ["joint_a", "joint_mimic", "joint_b"]} + self.num_instances = 1 + self.joint_names = ["joint_a", "joint_mimic", "joint_b"] + self.link_names = ["base_link", "tool_link"] + self.device = torch.device("cpu") + self.cfg = SimpleNamespace(solver_cfg=None) + self.current_qpos = torch.tensor([[0.1, 0.2, 0.3]], dtype=torch.float32) + self.target_qpos = torch.tensor([[0.4, 0.5, 0.6]], dtype=torch.float32) + self.write_calls: list[dict[str, object]] = [] + + def get_joint_ids( + self, + name: str, + remove_mimic: bool = False, + ) -> list[int]: + assert name == "arm" + return [0, 2] if remove_mimic else [0, 1, 2] + + def get_qpos(self, target: bool = False) -> torch.Tensor: + return self.target_qpos if target else self.current_qpos + + def set_qpos(self, **kwargs: object) -> None: + self.write_calls.append(kwargs) + + def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: + assert to_matrix + return torch.eye(4, dtype=torch.float32).unsqueeze(0) + + def get_link_pose( + self, + link_name: str, + env_ids: list[int], + to_matrix: bool = False, + ) -> torch.Tensor: + assert link_name in self.link_names + assert env_ids == [0] + assert to_matrix + pose = torch.eye(4, dtype=torch.float32) + pose[2, 3] = 0.8 + return pose.unsqueeze(0) + + +def test_robot_adapter_synchronizes_selected_joint_state() -> None: + robot = _FakeRobot() + adapter = _RobotGizmoAdapter(robot, "arm") + + assert adapter.get_actived_joint_names() == ["joint_a", "joint_b"] + np.testing.assert_allclose(adapter.get_current_qpos(), [0.1, 0.3]) + np.testing.assert_allclose(adapter.get_target_qpos(), [0.4, 0.6]) + + adapter.set_target_qpos(np.array([0.7, 0.9], dtype=np.float32)) + + write = robot.write_calls[-1] + assert write["joint_ids"] == [0, 2] + assert write["env_ids"] == [0] + assert write["target"] is True + torch.testing.assert_close( + write["qpos"], + torch.tensor([[0.7, 0.9]], dtype=torch.float32), + ) + + +def test_robot_adapter_reads_root_and_link_pose_through_robot() -> None: + adapter = _RobotGizmoAdapter(_FakeRobot(), "arm") + + np.testing.assert_allclose(adapter.get_world_pose(), np.eye(4)) + link_pose = adapter.get_link_pose("tool_link") + assert link_pose[2, 3] == pytest.approx(0.8) + assert adapter.get_link_names(True) == ["base_link", "tool_link"] + + +def test_robot_adapter_rejects_wrong_qpos_shape() -> None: + adapter = _RobotGizmoAdapter(_FakeRobot(), "arm") + + with pytest.raises(ValueError, match="Expected qpos shape"): + adapter.set_target_qpos(np.zeros(3, dtype=np.float32)) + + +def test_robot_ik_chain_can_be_configured_without_embodichain_solver() -> None: + gizmo = object.__new__(Gizmo) + gizmo.cfg = GizmoCfg( + ik_root_link_name="base_link", + ik_end_link_name="tool_link", + ) + gizmo._control_part = "arm" + robot = _FakeRobot() + + root_link, end_link, tcp_pose = gizmo._resolve_robot_ik_chain(robot) + + assert (root_link, end_link) == ("base_link", "tool_link") + np.testing.assert_allclose(tcp_pose, np.eye(4)) + + +def test_robot_update_delegates_to_dexsim_ik_controller() -> None: + calls: list[int] = [] + + class _Controller: + def update(self, *, iterations: int) -> None: + calls.append(iterations) + + gizmo = object.__new__(Gizmo) + gizmo.target = object() + gizmo._target_type = "robot" + gizmo._ik_controller = _Controller() + gizmo.cfg = GizmoCfg(ik_iterations=12) + + gizmo.update() + + assert calls == [12] + + +def test_destroy_removes_gizmo_from_dexsim_environment() -> None: + class _DexsimGizmo: + def __init__(self) -> None: + self.detached = False + + def set_flush_localpose_callback(self, callback: object | None) -> None: + pass + + def set_transform_flush_callback(self, callback: object | None) -> None: + pass + + def set_visible(self, visible: bool) -> None: + pass + + def detach_parent(self) -> None: + self.detached = True + + class _Environment: + def __init__(self) -> None: + self.removed: object | None = None + + def remove_gizmo(self, gizmo: object) -> None: + self.removed = gizmo + + dexsim_gizmo = _DexsimGizmo() + environment = _Environment() + gizmo = object.__new__(Gizmo) + gizmo._env = environment + gizmo._gizmo = dexsim_gizmo + gizmo._proxy_cube = None + gizmo._ik_controller = None + gizmo._ik_solver = None + gizmo._ik_model = None + gizmo._robot_adapter = None + gizmo._pending_target_transform = None + gizmo.target = object() + gizmo._target_type = "rigidobject" + + gizmo.destroy() + + assert environment.removed is dexsim_gizmo + assert dexsim_gizmo.detached is True + assert gizmo._gizmo is None diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index b78a63b24..2f1e994e0 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -19,6 +19,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock +import dexsim import numpy as np import pytest @@ -69,15 +70,61 @@ def add_loop(self, callback, time_step: float) -> str: return "loop_handle" +class FakeEntityGizmo: + """Entity-gizmo stub with external-target registration.""" + + def __init__(self) -> None: + self.active = True + self.external_targets: list[tuple[int, object, object, object]] = [] + + def register_external_target( + self, + target_id: int, + target_type: object, + target: object, + actor_type: object, + ) -> object: + self.external_targets.append((target_id, target_type, target, actor_type)) + return dexsim.interaction.EntityGizmoResult.SUCCESS + + class FakeWorld: """World stub exposing the render-thread loop API.""" def __init__(self) -> None: self.thread_runtime = FakeThreadRuntime() + self.entity_gizmo: object | None = None + self.entity_gizmo_configs: list[object | None] = [] + self.window = SimpleNamespace(add_input_control=lambda control: None) + self.window_open_count = 0 + self.window_closed = False def thread_rt(self) -> FakeThreadRuntime: return self.thread_runtime + def enable_entity_gizmo(self, config: object | None = None) -> object: + self.entity_gizmo_configs.append(config) + self.entity_gizmo = FakeEntityGizmo() + return self.entity_gizmo + + def disable_entity_gizmo(self) -> None: + if self.entity_gizmo is not None: + self.entity_gizmo.active = False + self.entity_gizmo = None + + def get_entity_gizmo(self) -> object | None: + return self.entity_gizmo + + def open_window(self) -> None: + self.window_open_count += 1 + self.window_closed = False + + def get_windows(self) -> object: + return self.window + + def close_window(self) -> None: + self.window_closed = True + class FakeEnv: """Environment stub that creates fake cameras.""" @@ -95,13 +142,24 @@ def _make_sim_manager(window: object | None = None) -> SimulationManager: """Create a minimally initialized simulation manager for recorder tests.""" sim = object.__new__(SimulationManager) sim.instance_id = 0 - sim.sim_config = SimpleNamespace(width=64, height=48) + sim.sim_config = SimpleNamespace( + width=64, + height=48, + enable_entity_gizmo_on_window_open=True, + ) sim._window = window + sim._entity_gizmo_config = None sim._window_record_state = None sim._window_record_camera = None sim._window_record_save_threads = [] + sim._window_record_hotkey_cfg = None + sim._window_camera_pose_hotkey_cfg = None + sim._window_record_input_control = None + sim._window_camera_pose_input_control = None sim._env = FakeEnv() sim._world = FakeWorld() + sim._default_plane = object() + sim.is_window_opened = window is not None return sim @@ -198,6 +256,129 @@ def fake_save_window_record_worker( assert sim._window_record_save_threads == [] +def test_entity_gizmo_lifecycle_delegates_to_dexsim_world() -> None: + sim = _make_sim_manager() + config = object() + + controller = sim.enable_entity_gizmo(config) + + assert controller is sim._world.get_entity_gizmo() + assert sim._world.entity_gizmo_configs == [config] + assert sim.get_entity_gizmo() is controller + assert sim.has_entity_gizmo() is True + assert sim.disable_entity_gizmo() is True + assert controller.active is False + assert sim.has_entity_gizmo() is False + assert sim.disable_entity_gizmo() is False + + +def test_entity_gizmo_registers_default_plane_as_static_exclusion() -> None: + sim = _make_sim_manager() + + controller = sim.enable_entity_gizmo() + + assert controller.external_targets == [ + ( + SimulationManager._DEFAULT_PLANE_GIZMO_TARGET_ID, + dexsim.interaction.EntityGizmoTargetType.RIGID_BODY, + sim._default_plane, + dexsim.types.ActorType.STATIC, + ) + ] + + +def test_open_window_enables_entity_gizmo_by_default() -> None: + sim = _make_sim_manager() + + sim.open_window() + + assert sim.is_window_opened is True + assert sim._world.window_open_count == 1 + assert sim.has_entity_gizmo() is True + assert sim._world.entity_gizmo_configs == [None] + + +def test_open_window_supports_view_only_opt_out() -> None: + sim = _make_sim_manager() + + sim.open_window(enable_entity_gizmo=False) + + assert sim.is_window_opened is True + assert sim.has_entity_gizmo() is False + assert sim._world.entity_gizmo_configs == [] + + +def test_open_window_view_only_opt_out_disables_active_controller() -> None: + sim = _make_sim_manager() + controller = sim.enable_entity_gizmo() + + sim.open_window(enable_entity_gizmo=False) + + assert controller.active is False + assert sim.has_entity_gizmo() is False + + +def test_open_window_respects_configured_entity_gizmo_default() -> None: + sim = _make_sim_manager() + sim.sim_config.enable_entity_gizmo_on_window_open = False + + sim.open_window() + + assert sim.is_window_opened is True + assert sim.has_entity_gizmo() is False + + +def test_open_window_tolerates_dexsim_without_entity_gizmo_api() -> None: + sim = _make_sim_manager() + window = object() + sim._world = SimpleNamespace( + open_window=lambda: None, + get_windows=lambda: window, + ) + + sim.open_window() + + assert sim.is_window_opened is True + assert sim._window is window + assert sim.has_entity_gizmo() is False + + +def test_open_window_preserves_active_entity_gizmo_configuration() -> None: + sim = _make_sim_manager(window=object()) + config = object() + controller = sim.enable_entity_gizmo(config) + + sim.open_window() + + assert sim.get_entity_gizmo() is controller + assert sim._world.entity_gizmo_configs == [config] + assert sim._world.window_open_count == 0 + + +def test_reopened_window_restores_last_entity_gizmo_configuration() -> None: + sim = _make_sim_manager(window=object()) + config = object() + sim.enable_entity_gizmo(config) + sim.close_window() + + sim.open_window() + + assert sim.has_entity_gizmo() is True + assert sim._world.entity_gizmo_configs == [config, config] + + +def test_close_window_disables_entity_gizmo() -> None: + sim = _make_sim_manager(window=object()) + controller = sim.enable_entity_gizmo() + + sim.close_window() + + assert controller.active is False + assert sim.has_entity_gizmo() is False + assert sim._world.window_closed is True + assert sim.is_window_opened is False + + def test_reset_objects_state_includes_soft_and_cloth_assets() -> None: sim = object.__new__(SimulationManager) sim._robots = {}