diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 657b67c32..8a6464c77 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -21,6 +21,10 @@ topics: - simulation - SimulationManager - SimulationManagerCfg + - EntityGizmoManipulator + - enable_entity_gizmo + - disable_entity_gizmo + - robot_ik_gizmo - DexSim - world - arena @@ -349,11 +353,18 @@ topics: - deformable - soft body - cloth + - gizmo + - IKGizmoController + - create_robot_ik_gizmo_controller + - GizmoCfg + - ScenePicker + - PickCommand paths: - topics/sim-visualization/sim-visualization.md source_of_truth: - embodichain/lab/sim/sim_manager.py - embodichain/lab/visualization/ + - embodichain/lab/sim/objects/gizmo.py - embodichain/lab/gym/utils/gym_utils.py - embodichain/lab/gym/envs/base_env.py - embodichain/lab/scripts/preview_asset.py diff --git a/agent_context/topics/ik-solvers/ik-solvers.md b/agent_context/topics/ik-solvers/ik-solvers.md index ac0ddcbf0..2de569ecc 100644 --- a/agent_context/topics/ik-solvers/ik-solvers.md +++ b/agent_context/topics/ik-solvers/ik-solvers.md @@ -34,6 +34,16 @@ path for performance. ## Solver Hierarchy +Gizmo solver selection belongs to `embodichain/lab/sim/objects/gizmo.py`, not +to a new solver subclass. Native DexSim `IKGizmoController` remains the +controller in both modes: `GizmoCfg.ik_solver="dexsim"` selects Newton IK; +`"embodichain"` adapts `robot.get_solver(control_part)` for the native window +and Viser. It preserves configured solver convergence/limits, maps joint names +between solver and robot order, and holds current qpos on failure/non-finite +results. Gizmo root/end links must match the configured solver; TCP overrides +are converted without mutating that solver. See the robot gizmo example's +`--ik-solver pink` option for an executable Pink configuration. + ``` SolverCfg (@configclass, abstract) ├── SRSSolverCfg diff --git a/agent_context/topics/sim-visualization/sim-visualization.md b/agent_context/topics/sim-visualization/sim-visualization.md index 64418e222..8fb1e335f 100644 --- a/agent_context/topics/sim-visualization/sim-visualization.md +++ b/agent_context/topics/sim-visualization/sim-visualization.md @@ -252,6 +252,32 @@ Soft bodies and cloth require GPU physics. Their live vertices are sampled at ## Browser Controls and Overlays +Gizmo implementation lives in `embodichain/lab/sim/objects/gizmo.py`. Native +windows delegate object picking/manipulation to DexSim's entity gizmo and robot +targets to its `IKGizmoController`. Entity interaction defaults on at the first +native window open; `SimulationManagerCfg.enable_entity_gizmo=False` or +`sim.disable_entity_gizmo()` opts out and reopening preserves the choice. +Pure headless and Viser runs do not automatically create native controls. +`SimulationManagerCfg.robot_ik_gizmo` automatically registers robot parts with +solver chain/TCP metadata: native IK activates on I by default, or on the first +update with an open window when `GizmoCfg(ik_start_enabled=True)`. The robot +tutorial uses this startup option. Viser builds IK on the first drag. The manager updates both and preserves native controls across window +reopen. Explicit configuration and disable calls take precedence; caller-owned +native factory controllers are not duplicated. Automatic Viser registration +requires `allow_commands`, independently of native interaction preferences. +Viser transports poses to the +simulation thread. Both robot paths use Newton IK by default or, with +`GizmoCfg(ik_solver="embodichain")`, reuse the configured control-part solver +(including Pink). Chain roots follow the live robot link, including upstream +joint motion; TCP overrides adapt targets without changing the shared solver. + +Viser click picking runs on the visualization worker via the existing GUI +event queue. A manifest invalidates cached pick poses until its matching frame +arrives; stale clicks are dropped. The manager validates `PickCommand` run and +revision before attaching a gizmo and tracks picker ownership independently +from explicitly created gizmos. Clearing selection releases only the picker +gizmo. Native entity selection remains entirely DexSim-owned. + The browser GUI has: - **Environments** — visibility per exported environment; diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index ce7246c06..6010e0dd1 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -71,6 +71,15 @@ then advances the world for the requested number of physics steps. Each environment control step normally calls it with `sim_steps_per_control`. +`scripts/tutorials/sim/gizmo_robot.py` supports only manual physics. It initializes +GPU physics after robot creation when needed, sets both current and target +joint positions, and advances once before opening the window. It explicitly +sets `GizmoCfg(ik_start_enabled=True)` so the native controller activates on the +first update after opening the window. Its loop only +calls `sim.update(step=1)`; the manager owns native IK updates and Viser +commands/capture. The loop is paced by `physics_dt` and has no automatic +physics polling path. + `ArticulationCfg.enable_gravity` defaults to `True`. During articulation construction, `Articulation` applies this explicit runtime flag to every native entity before the first physics update, including when @@ -118,6 +127,45 @@ do not belong to the simulation object; use ## Configuration Flow +### Gizmo ownership + +Native entity manipulation belongs to DexSim 0.5.0. The first successful native +window open enables its world-owned `EntityGizmoManipulator` by default, after +the scene and default plane are ready. The manager registers the default plane +as a static external target. `SimulationManagerCfg.enable_entity_gizmo=False` +opts out; Gym deployments accept the same top-level JSON/YAML field through +`gym.utils.gym_utils.config_to_cfg()`. + +`sim.enable_entity_gizmo(config)` explicitly enables/configures the controller; +`sim.disable_entity_gizmo()` also cancels pending automatic enablement before +the first window. These explicit calls take precedence over the startup default. +Query through `sim.get_world().get_entity_gizmo()`. DexSim owns window +detach/reopen and controller state; reopening never reapplies the default or +overwrites an explicit native disable. Pure headless and Viser runs do not +automatically create a native entity controller. + +`SimulationManagerCfg.robot_ik_gizmo` defaults to `GizmoCfg()`. During normal +updates the manager registers robot control parts with complete solver chain/TCP +metadata in single-environment interactive runs. Pure headless, read-only Viser, and +multi-environment runs do not register automatic controls. The first native I +press creates DexSim's `IKGizmoController` by default; +`GizmoCfg(ik_start_enabled=True)` opts into activation on the first update with +an open window. The startup attempt is consumed once, including on failure; +later key presses can retry. Viser constructs IK on its first drag. +Registration never writes drive targets. `Gizmo` owns managed native input and +target-node cleanup, detaches input on window close, and reattaches the same +controller on reopen. Robot removal releases all its managed controls. + +Set `robot_ik_gizmo=None` to opt out or supply `GizmoCfg` overrides; Gym +JSON/YAML accepts the same mapping/null. `enable_gizmo()` can override one part, +and `disable_gizmo()` prevents automatic recreation (all parts when omitted). +The explicit `create_robot_ik_gizmo_controller()` factory still returns +caller-owned controllers; a weak registry prevents automatic duplicates. +Both robot paths +default to native Newton IK; `GizmoCfg(ik_solver="embodichain")` adapts the +control part's existing solver, such as PinkSolver. Both support one environment +and write only selected non-mimic joint drive targets through `Robot`. + `SimulationManagerCfg` owns window size, headless mode, rendering, GPU/CPU selection, arena count and spacing, physics timestep, physics and GPU-memory settings, recording, profiling, and browser visualization. diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst index 7f9577485..fc0a040fd 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst @@ -42,6 +42,7 @@ bodies. RobotWorkspaceCfg Gizmo GizmoCfg + create_robot_ik_gizmo_controller RigidConstraint .. currentmodule:: embodichain.lab.sim.objects @@ -191,6 +192,27 @@ Gizmo :show-inheritance: :exclude-members: __init__, copy, replace, to_dict, validate +.. autofunction:: create_robot_ik_gizmo_controller + +SimulationManager automatically discovers robot control parts with configured +IK chain metadata. Native controllers activate on the first I press by default, while +Viser constructs its solver on the first drag. ``sim.update()`` owns updates +and cleanup; ordinary applications do not need the explicit factory. +Use ``SimulationManagerCfg(robot_ik_gizmo=None)`` to disable automatic setup. +Set ``GizmoCfg(ik_start_enabled=True)`` to activate native IK on the first update +with an open window; subsequent visibility toggles and reopening are preserved. + +The native controller defaults to DexSim Newton IK. With a ``PinkSolverCfg`` +(or another EmbodiChain solver) configured for the robot's control part, pass +``GizmoCfg(ik_solver="embodichain")`` to select that solver for either a native +controller or a Viser gizmo. Its iteration limits and convergence settings +remain owned by the configured solver; ``ik_iterations`` applies to Newton IK. +Only the selected control part's drive targets are written, and failed +EmbodiChain IK solutions preserve the current joint positions. + +The runnable example ``examples/sim/gizmo/gizmo_robot.py`` exposes +``--ik-solver dexsim|pytorch|pink`` for both the native window and ``--viser``. + Rigid Constraint ---------------- diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst index 088dcb39e..1131a77f8 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst @@ -13,6 +13,22 @@ configuration. Downstream components (environments, planners, IK solvers, the visualization runtime) look up the active manager through its class-level instance registry instead of passing it around explicitly. +Native entity interaction defaults on when the first native window opens. +Set ``SimulationManagerCfg(enable_entity_gizmo=False)`` to opt out, or call +``sim.disable_entity_gizmo()`` at runtime. Explicit enable/disable calls and +custom DexSim controller settings survive window close/reopen. Pure headless +and Viser runs do not automatically create native gizmos. + +``SimulationManagerCfg.robot_ik_gizmo`` defaults to ``GizmoCfg()`` and registers +robot control parts with configured IK-chain/TCP metadata during normal updates. +Native IK activates on the first **I** press by default; Viser constructs its solver on +the first drag and requires ``visualization.allow_commands``. Registration does +not write drive targets. Set this field to ``None`` to opt out or select +``GizmoCfg(ik_solver="embodichain")`` to reuse configured solvers. Explicit +``enable_gizmo()`` settings override automatic defaults, and ``disable_gizmo()`` +prevents automatic recreation. ``GizmoCfg(ik_start_enabled=True)`` activates +native IK on the first update with an open window, as used by the robot tutorial. + .. rubric:: Classes .. autosummary:: 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/api_reference/embodichain/embodichain.lab.visualization.rst b/docs/source/api_reference/embodichain/embodichain.lab.visualization.rst index 2a52d2432..fc3b01c71 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.visualization.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.visualization.rst @@ -103,6 +103,14 @@ Scene Export :members: :undoc-members: +``PickCommand`` identifies a node in a specific simulation run and scene +revision. An empty selection releases only the gizmo created by click picking; +explicitly configured gizmos retain their ownership. + +.. autoclass:: PickCommand + :members: + :undoc-members: + .. autoclass:: JointControlSpec :members: :undoc-members: diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index ea1aa59e5..1c05b3565 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -865,10 +865,15 @@ embodichain.lab.sim.objects.gizmo .. currentmodule:: embodichain.lab.sim.objects.gizmo +Native robot targets use DexSim's controller with Newton IK by default. +Set ``GizmoCfg.ik_solver="embodichain"`` to reuse the robot control part's +configured solver, including PinkSolver; Viser uses the same solver adapter. + .. autosummary:: Gizmo GizmoCfg + create_robot_ik_gizmo_controller embodichain.lab.sim.objects.rigid_object ---------------------------------------- @@ -1499,6 +1504,22 @@ embodichain.lab.visualization.cli add_viser_args_to_parser visualization_cfg_from_args +embodichain.lab.visualization.picker +------------------------------------ + +.. currentmodule:: embodichain.lab.visualization.picker + +Browser picking caches triangle geometry and returns the closest node hit by +a world-space ray. The Viser worker pairs this geometry with poses from the +same scene revision before producing a pick command. + +.. autosummary:: + + ScenePicker + +.. autoclass:: ScenePicker + :members: + embodichain.lab.visualization.protocol -------------------------------------- @@ -1520,6 +1541,7 @@ embodichain.lab.visualization.protocol JointControlSpec JointControlState MeshGeometry + PickCommand PointCloudOverlay SceneFrame SceneManifest diff --git a/docs/source/features/interaction/gizmo.md b/docs/source/features/interaction/gizmo.md index 53ae07881..f8992e972 100644 --- a/docs/source/features/interaction/gizmo.md +++ b/docs/source/features/interaction/gizmo.md @@ -3,21 +3,31 @@ ```{currentmodule} embodichain.lab.sim ``` -A Gizmo is a registered transform control for manipulating a simulation target -from either the native DexSim window or a trusted Viser browser. The simulation -owns the authoritative state: UI callbacks enqueue requested poses, and -`SimulationManager` applies them on the simulation thread. +A Gizmo is a transform control for manipulating a simulation target. Native +windows use DexSim's entity and robot IK controllers. Viser uses registered +controls whose callbacks enqueue requested poses for `SimulationManager` to +apply on the simulation thread. + +Native entity interaction is enabled by default when the first native window +opens. Select an object and press **G** to attach its gizmo. Set +`SimulationManagerCfg(enable_entity_gizmo=False)` or call +`sim.disable_entity_gizmo()` to opt out; reopening preserves that choice. +Robot IK controls are discovered automatically from configured control parts +and their solver's root link, end link, and TCP. +See {doc}`native window controls ` for native entity configuration. ## Supported Targets | Target | Gizmo behavior | |---|---| -| Robot | Moves the selected control part through its configured FK/IK solver. | +| Robot | Moves the selected control part through native DexSim IK by default, or its configured EmbodiChain solver. | | Rigid object | Sets the object's local arena pose. | | Camera | Sets the camera's local pose. | -Gizmo interaction currently requires `num_envs=1`. Robot targets also require -a valid `control_part` and solver configuration. +Registered Viser controls and robot IK controls currently require `num_envs=1`. +Automatic discovery requires robot solver metadata for each supported control +part. IK computation still defaults to DexSim. Robots without this metadata +can use explicit `GizmoCfg` chain settings instead; no end link or TCP is guessed. ## Quick Start @@ -33,39 +43,46 @@ Use the same target through Viser: python scripts/tutorials/sim/gizmo_robot.py --viser ``` -Register controls through the manager rather than constructing or destroying -Gizmo instances directly: +The tutorial opts into immediate native activation with +`robot_ik_gizmo=GizmoCfg(ik_start_enabled=True)`. It sets the initial robot pose +before opening the window, then its ordinary manual-physics loop creates and +updates the controller: ```python -sim.enable_gizmo( - uid="robot", - control_part="arm", -) - -if not sim.has_gizmo("robot", control_part="arm"): - raise RuntimeError("Gizmo setup failed") +sim.open_window() # Safely skipped when Viser is configured. +while True: + sim.update(step=1) ``` -`SimulationManager.update()` drains pending Gizmo commands, steps physics, -and publishes the resulting pose to Viser. Interactive applications call it -explicitly in their main loop: +By default, the first **I** press creates eligible native robot IK controllers. +With `ik_start_enabled=True`, the first update with an open window activates +them from the current robot pose; **I** then toggles their visibility. Opening a window alone +does not construct an IK solver or change existing drive targets. Viser displays +the TCP handles and constructs the solver on the first drag. Both paths update +through `SimulationManager.update()`. + +The default `SimulationManagerCfg.robot_ik_gizmo` is `GizmoCfg()`. To use each +control part's configured solver (including Pink), or disable automatic setup: ```python -while True: - sim.update(step=1) +SimulationManagerCfg(robot_ik_gizmo=GizmoCfg(ik_solver="embodichain")) +SimulationManagerCfg(robot_ik_gizmo=None) ``` -The caller can pace this loop against wall time using the configured physics -timestep. For editing while physics is paused, call `update_gizmos()` followed -by `capture_visualization_safely()` without stepping the world. - -Use `disable_gizmo()`, `set_gizmo_visibility()`, and -`toggle_gizmo_visibility()` for lifecycle and visibility changes. +Gym JSON/YAML deployments accept the same `robot_ik_gizmo` mapping or `null`. +`sim.enable_gizmo(uid, control_part, gizmo_cfg)` provides an explicit per-part +override. `sim.disable_gizmo(uid, control_part)` prevents automatic recreation, +including across window reopen; omit the part to disable all parts of a robot. +Use `set_gizmo_visibility()` and `toggle_gizmo_visibility()` for visibility. +Activated native controls retain their IK and visibility state across window +close/reopen. Removing a robot or destroying the manager releases their input +handlers and native target nodes. ## Frontend Behavior -- The native window uses the DexSim Gizmo controller and requires an open - window. +- Native entity interaction starts with the first native window; pure headless + and Viser runs do not automatically create native controllers. Native robot + targets activate on **I** without a script-level controller call. - Viser exports browser-native transform controls when `VisualizationCfg.allow_commands=True`; otherwise it shows read-only frames. - The standard `--viser` launcher enables registered commands for trusted @@ -73,6 +90,10 @@ Use `disable_gizmo()`, `set_gizmo_visibility()`, and - A Viser Gizmo is owned by one browser client from drag start until drag end or disconnect. Other clients continue receiving authoritative poses. +Advanced callers may still use `create_robot_ik_gizmo_controller()` directly +and retain/update its returned objects. Automatic setup detects an existing +factory-created controller and does not create or update a duplicate. + For the complete robot setup and IK walkthrough, continue with the {doc}`Gizmo tutorial `. See {doc}`Viser browser visualization ` for diff --git a/docs/source/features/interaction/window.md b/docs/source/features/interaction/window.md index eda9a35a3..53cd59576 100644 --- a/docs/source/features/interaction/window.md +++ b/docs/source/features/interaction/window.md @@ -53,6 +53,66 @@ 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. EmbodiChain enables it +automatically when the first native window opens, including a window created +with `SimulationManagerCfg(headless=False)`. Pure headless and Viser runs do not +automatically create native entity gizmos. + +To start with native interaction disabled, use +`SimulationManagerCfg(enable_entity_gizmo=False)`. Gym JSON/YAML deployments +accept the top-level field `enable_entity_gizmo: false` as well. At runtime, +`sim.disable_entity_gizmo()` disables interaction and cancels automatic +enablement even before the first window opens. Closing and reopening a window +preserves the controller's current enabled state and custom configuration. + +For custom DexSim settings, enable or reconfigure the controller explicitly: + +```python +import dexsim + +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. Query the controller through DexSim's +world object; use the manager's disable helper to preserve your preference +across future window opens: + +```python +controller = sim.get_world().get_entity_gizmo() +sim.disable_entity_gizmo() +``` + +Robot TCP IK controls are registered automatically for parts with configured +IK chain/TCP metadata. By default, the first **I** press activates them; later +presses show or hide their targets. **G** continues to control entity roots. Normal +`sim.update()` calls handle IK updates; no controller-specific call is needed. + +For immediate activation, set +`SimulationManagerCfg(robot_ik_gizmo=GizmoCfg(ik_start_enabled=True))`. The +controller starts on the first update with an open window. The robot Gizmo +tutorial enables this option after preparing its initial joint pose. + +The entity gizmo is native-window only. The Viser backend offers an analogous +**click-to-pick** flow (an *Enable click-to-pick Gizmo* checkbox instead of the +**G** hotkey, since browsers do not expose keyboard events); see +: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 619715642..407054e6c 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 9e7003ad9..9e1592005 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, Viser commands, and controller lifecycle management. For the cross-frontend capability summary, supported targets, lifecycle rules, and security boundary, see :doc:`/features/interaction/gizmo`. @@ -31,9 +33,11 @@ Similar to the previous tutorial on robot simulation, we use the :class:`Simulat -**Important:** Gizmo only supports single environment mode (`num_envs=1`). Using multiple environments will raise an exception. +**Important:** Gizmo supports a single environment (``num_envs=1``). Automatic +registration is skipped for multi-environment simulations. -All gizmo creation, visibility, and destruction operations must be managed via the SimulationManager API: +Robot Gizmo registration, updates, visibility, and destruction are managed by +SimulationManager: .. code-block:: python @@ -43,10 +47,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 interaction uses DexSim controllers. The standard Viser mode includes +interactive Gizmo control: .. code-block:: bash @@ -55,6 +57,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? ----------------- @@ -65,7 +88,8 @@ A Gizmo is an interactive visual tool that allows users to manipulate simulation - **Real-time Manipulation**: Provide immediate visual feedback during robot motion planning - **Debugging and Visualization**: Test robot reachability and workspace limits -The :class:`objects.Gizmo` class provides a unified interface for interactive control of different simulation elements including robots, rigid objects, and cameras. +The :class:`objects.Gizmo` class manages native robot controllers and Viser +targets. Native entity manipulation remains owned by DexSim. Setting up Robot Configuration ------------------------------ @@ -81,142 +105,102 @@ Key components of the robot configuration: - **URDF Configuration**: Loads the robot's kinematic and visual model - **Control Parts**: Defines which joints can be controlled (``"Joint[1-6]"`` for UR10) -- **IK Solver**: :class:`solvers.PinkSolverCfg` provides inverse kinematics capabilities +- **IK Solver**: :class:`solvers.PinkSolverCfg` supplies chain metadata and an optional solver override - **Drive Properties**: Sets stiffness and damping for joint control -The IK solver is crucial for gizmo functionality, as it enables the robot to automatically calculate joint angles needed to reach gizmo target positions. - -Creating and Attaching a Gizmo -------------------------------- - +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`. +Automatic Robot Controls +------------------------ -After configuring the robot, enable the gizmo for interactive control using the SimulationManager API (supports robot, rigid object, camera; key is `uid:control_part`): +SimulationManager discovers each control part with existing root-link and +end-link metadata and uses its configured TCP transform. This tutorial +explicitly activates native IK at startup with one setting: .. 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. + sim_cfg = SimulationManagerCfg( + robot_ik_gizmo=GizmoCfg(ik_start_enabled=True), ) - if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): - logger.log_error("Failed to enable gizmo!") - return +The first update after opening the window creates and displays the controller. +The tutorial sets current joint positions and drive targets before opening the +window; activation initializes the controller from that pose. +- In this tutorial's native window, IK targets start visible. Press **I** to + hide or show them. Ordinary simulations retain ``ik_start_enabled=False`` + and wait for the first **I** press to activate. +- In Viser, the TCP controls are available automatically when commands are + allowed. The solver is constructed on the first drag. +- Pure headless, read-only Viser, and multi-environment simulations skip + automatic registration. -The Gizmo instance is managed internally by SimulationManager. If you need to access it: +Registration alone does not initialize another IK solver or overwrite drive +targets. Explicit startup activation initializes the controller once; closing +and reopening the window preserves its later visibility state. DexSim Newton IK +is the default. To use the robot's configured solver instead: .. code-block:: python - gizmo = sim.get_gizmo("ur10_gizmo_test", control_part="arm") - + from embodichain.lab.sim.objects import GizmoCfg + sim_cfg = SimulationManagerCfg( + robot_ik_gizmo=GizmoCfg(ik_solver="embodichain"), + ) -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 +Set ``robot_ik_gizmo=None`` to disable automatic setup. Robots without chain +metadata can still use ``sim.enable_gizmo(...)`` with explicit +:class:`objects.GizmoCfg` link settings. Advanced callers can use +:func:`objects.create_robot_ik_gizmo_controller` and manage its updates directly; +SimulationManager will not create a duplicate native controller for that part. How Gizmo-Robot Interaction Works ---------------------------------- -The gizmo-robot interaction follows this efficient 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 +The gizmo-robot interaction follows this workflow: -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 ------------------- +The tutorial uses manual physics only. After setting initial joint positions +and drive targets, each iteration advances one physics step: +.. literalinclude:: ../../../scripts/tutorials/sim/gizmo_robot.py + :language: python + :start-at: def run_simulation( + :end-at: sim.update(step=1) -In the main loop, call ``sim.update(step=1)`` to apply pending Gizmo controls, -advance physics, and publish the resulting state to Viser. - - - -.. code-block:: python - - def run_simulation(sim: SimulationManager): - step_count = 0 - try: - last_time = time.time() - last_step = 0 - while True: - step_start = time.perf_counter() - sim.update(step=1) - step_count += 1 - # ...performance statistics, etc... - time.sleep(max(0.0, sim.sim_config.physics_dt - (time.perf_counter() - step_start))) - except KeyboardInterrupt: - logger.log_info("\nStopping simulation...") - finally: - sim.destroy() # Release all resources - logger.log_info("Simulation terminated successfully") - - - -Main loop highlights: - -- **Gizmo update**: ``sim.update(step=1)`` processes controls before stepping physics -- **Viser update**: Each explicit physics step publishes the resulting state when Viser is enabled -- **Performance monitoring**: Optional FPS statistics -- **Resource cleanup**: Only `sim.destroy()` is needed, no manual Gizmo destruction -- **Graceful shutdown**: Supports Ctrl+C interruption +``sim.update()`` processes native and Viser interaction, advances physics, and +publishes visualization state. No separate controller update is required. +The tutorial paces the loop using ``physics_dt`` and releases resources with +``sim.destroy()`` on Ctrl+C. Gizmo Lifecycle Management --------------------------- - - - - -Gizmo lifecycle is managed by SimulationManager: - -- Enable: `sim.enable_gizmo(...)` -- Update: ``sim.update(step=1)`` processes Gizmos in 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):** +------------------------- -- ``sim.toggle_gizmo_visibility(uid, control_part=None)``: Toggle gizmo visibility -- ``sim.set_gizmo_visibility(uid, visible, control_part=None)``: Set gizmo visibility +SimulationManager handles automatic robot control registration and cleanup. +Closing a native window detaches input handlers; reopening it reuses existing +controllers and preserves their visibility without writing new drive targets. +Removing a robot also removes its managed controls. -**Hierarchy Management:** +For explicit overrides: -- ``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 +- ``sim.enable_gizmo(uid, control_part, gizmo_cfg)`` replaces that part's settings. +- ``sim.disable_gizmo(uid, control_part)`` disables one part and prevents automatic + recreation; omitting the part disables every part of the robot. +- ``sim.toggle_gizmo_visibility(uid, control_part)`` and + ``sim.set_gizmo_visibility(uid, visible, control_part)`` control visibility. Running the Tutorial -------------------- @@ -233,14 +217,16 @@ Command-line options: - ``--device cpu|cuda``: Choose simulation device - ``--num_envs N``: Number of parallel environments - ``--headless``: Run without GUI for automated testing -- ``--renderer``: Enable ray tracing for better visuals +- ``--renderer auto|hybrid|fast-rt|rt``: Select the renderer +- ``--viser``: Use browser-based interaction Once running: -1. **Mouse Interaction**: Click and drag the gizmo (colorful axes) to move the robot -2. **Real-time IK**: Watch the robot joints automatically adjust to follow the gizmo -3. **Workspace Limits**: Observe how the robot behaves at workspace boundaries -4. **Performance**: Monitor FPS in the console output +1. **Open**: The native IK target starts visible, or open the Viser page +2. **Mouse Interaction**: Click and drag the gizmo to move the robot +3. **Real-time IK**: Watch the robot joints automatically adjust to follow the gizmo +4. **Workspace Limits**: Observe how the robot behaves at workspace boundaries +5. **Performance**: Monitor FPS in the console output Tips and Best Practices ------------------------ @@ -249,23 +235,23 @@ Tips and Best Practices **Performance optimization:** -- Call ``sim.update(step=1)`` in the main loop; do not update Gizmos a second time +- Use ``sim.update(step=1)`` to service interaction and advance manual physics - Reduce IK solver iterations for better real-time performance if needed -- Pace the loop with the configured physics timestep for interactive playback +- Pace manual steps using ``physics_dt`` **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 @@ -273,7 +259,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 @@ -283,7 +269,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/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 34ecc7574..15c714737 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -620,6 +620,8 @@ class ComponentCfg: env_cfg.sim_cfg = SimulationManagerCfg( headless=config.get("headless", False), + enable_entity_gizmo=config.get("enable_entity_gizmo", True), + robot_ik_gizmo=config.get("robot_ik_gizmo", {}), sim_device=config.get("device", "cpu"), render_cfg=RenderCfg(**render_config), gpu_id=config.get("gpu_id", 0), diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index 0c51b56ef..87549116f 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -1751,8 +1751,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 f74767b3c..a6220b1d3 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 ( @@ -36,7 +38,7 @@ ) 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..92b07d9e9 100644 --- a/embodichain/lab/sim/objects/gizmo.py +++ b/embodichain/lab/sim/objects/gizmo.py @@ -14,37 +14,47 @@ # 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, Literal +from weakref import WeakValueDictionary import dexsim import numpy as np import torch -from dexsim.types import ( - AxisArrowType, - AxisCornerType, - AxisOption, - AxisTagType, - RotationRingsOption, +from dexsim.kit.ik.pose import ( + local_pose_from_world, + pose_from_position_rotation, + rotation_matrix_to_quat_xyzw, ) -from scipy.spatial.transform import Rotation +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 + from embodichain.lab.sim.solvers import BaseSolver + +__all__ = ["Gizmo", "GizmoCfg", "create_robot_ik_gizmo_controller"] + +# Explicit factory users retain ownership; automatic controls must not create +# a second controller for the same robot/control part. +_NATIVE_IK_CONTROLLERS: WeakValueDictionary[tuple[int, str], IKGizmoController] = ( + WeakValueDictionary() +) @configclass class GizmoCfg: - """Configure native and Viser Gizmo appearance.""" + """Configure Gizmo appearance and native or Viser robot IK behavior.""" axis_length_x: float = 0.2 """Length of the X-axis arrow.""" @@ -56,97 +66,452 @@ 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_solver: Literal["dexsim", "embodichain"] = "dexsim" + """Use native Newton IK, or the control part's configured EmbodiChain solver. + + The native window always uses DexSim's ``IKGizmoController``. Selecting + ``embodichain`` reuses ``robot.get_solver(control_part)`` (for example, + PinkSolver), including that solver's convergence and joint-limit settings. + """ + + ik_root_link_name: str | None = None + """Robot IK chain root link, or the configured solver root when omitted.""" + + 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.""" + + ik_start_enabled: bool = False + """Activate managed native IK on the first update with an open window. + + Defaults to waiting for the toggle key. Explicit startup activation creates + the controller and initializes its drive targets from the current pose. + This is a one-time request: later toggles and window reopening retain their + normal behavior. Viser still builds its solver on the first drag. + """ + + +class _RobotGizmoAdapter: + """Expose one EmbodiChain robot control part to DexSim's IK controller.""" + + 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] + self.root_link_name: str | None = None + self.model_root_inverse = np.eye(4, dtype=np.float32) + + def get_current_qpos(self) -> np.ndarray: + """Return current selected joint positions in DexSim joint-name order.""" + 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: + """Map the IK model frame into the selected instance's arena frame.""" + if self.root_link_name is not None: + return self.get_link_pose(self.root_link_name) @ self.model_root_inverse + pose = self.robot.get_local_pose(to_matrix=True)[self.env_id] + return pose.detach().cpu().numpy().astype(np.float32, copy=True) + + 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 + + +class _EmbodiChainIK: + """Adapt BaseSolver to the solver contract consumed by DexSim and Viser.""" + + def __init__( + self, adapter: _RobotGizmoAdapter, solver: BaseSolver, tcp_pose: np.ndarray + ) -> None: + self.solver = solver + self.joint_names = list(solver.joint_names or adapter.joint_names) + if set(self.joint_names) != set(adapter.joint_names): + raise ValueError( + "IK solver joints must match the control part's active joints." + ) + self._tcp_inverse = np.linalg.inv(tcp_pose) + self._solution: dict[str, float] = {} + pose = ( + local_pose_from_world( + adapter.get_world_pose(), adapter.get_link_pose(solver.end_link_name) + ) + @ tcp_pose + ) + self._target_state: dict[str, np.ndarray] = {} + self.set_target_pose(pose[:3, 3], rotation_matrix_to_quat_xyzw(pose[:3, :3])) + self.reset_target_state_changed() + + def target_state(self) -> dict[str, np.ndarray]: + return self._target_state + + def set_target_pose(self, position: np.ndarray, rotation: np.ndarray) -> None: + self._target_state.update( + position=np.array(position), rotation=np.array(rotation) + ) + + def target_state_changed(self) -> bool: + return any( + not np.allclose(value, self._snapshot[key], atol=1e-5) + for key, value in self._target_state.items() + ) + + def reset_target_state_changed(self) -> None: + self._snapshot = { + key: value.copy() for key, value in self._target_state.items() } + def solve( + self, + joint_names: list[str], + current_qpos: np.ndarray, + *, + iterations: int | None = None, + ) -> None: + # EmbodiChain solvers retain their own iteration/convergence configuration. + del iterations + self._solution.clear() + seed = dict(zip(joint_names, current_qpos)) + target = pose_from_position_rotation(**self._target_state) + # The gizmo TCP may differ from the shared solver's TCP. Adapt the + # target without mutating the solver used by the rest of the robot. + target = target @ self._tcp_inverse @ self.solver.get_tcp() + success, qpos = self.solver.get_ik( + torch.as_tensor( + target, dtype=torch.float32, device=self.solver.device + ).unsqueeze(0), + qpos_seed=torch.tensor( + [[seed[name] for name in self.joint_names]], + dtype=torch.float32, + device=self.solver.device, + ), + return_all_solutions=False, + ) + if bool(success.all()) and bool(torch.isfinite(qpos).all()): + values = qpos.detach().cpu().numpy().reshape(len(self.joint_names)) + self._solution = dict(zip(self.joint_names, values)) + + def qpos_for_joint_names( + self, joint_names: list[str], fallback_qpos: np.ndarray + ) -> np.ndarray: + return np.array( + [ + self._solution.get(name, value) + for name, value in zip(joint_names, fallback_qpos) + ], + dtype=np.float32, + ) + + +def _build_robot_ik( + robot: Robot, + control_part: str, + cfg: GizmoCfg, +) -> tuple[_RobotGizmoAdapter, NewtonChainIK | _EmbodiChainIK, str, np.ndarray]: + if robot.num_instances != 1: + raise RuntimeError( + "Robot Gizmo supports exactly one environment, " + f"but the robot has {robot.num_instances} instances." + ) + if cfg.ik_solver not in {"dexsim", "embodichain"}: + raise ValueError("ik_solver must be 'dexsim' or 'embodichain'.") + if cfg.ik_solver == "dexsim" and cfg.ik_iterations <= 0: + raise ValueError("ik_iterations must be greater than zero.") + + root_link, end_link, tcp_pose = _resolve_robot_ik_chain( + robot, + control_part, + cfg, + ) + adapter = _RobotGizmoAdapter(robot, control_part) + adapter.root_link_name = root_link + if cfg.ik_solver == "embodichain": + solver = ( + robot.get_solver(control_part) if robot.cfg.solver_cfg is not None else None + ) + if solver is None: + raise ValueError( + f"Control part {control_part!r} needs a configured EmbodiChain solver." + ) + if (root_link, end_link) != (solver.root_link_name, solver.end_link_name): + raise ValueError( + "Gizmo IK links must match the configured EmbodiChain solver." + ) + return adapter, _EmbodiChainIK(adapter, solver, tcp_pose), end_link, tcp_pose + + import warp as wp + from dexsim.kit.ik import NewtonChainIK, build_newton_model_from_urdf + + with wp.ScopedDevice(cfg.ik_device or str(robot.device)): + model = build_newton_model_from_urdf(robot.cfg.fpath, hide_visuals=True) + solver = NewtonChainIK( + model, + start_link=root_link, + end_link=end_link, + iterations=cfg.ik_iterations, + tcp_pose=tcp_pose, + ) + + # Newton preserves the start link's URDF rest transform in its reduced + # model. Follow the live chain root (including upstream joints) while + # retaining that model frame, rather than assuming the robot root is it. + root_index = list(solver.model.body_label).index(solver.info.start_link) + root_pose = solver.state.body_q.numpy()[root_index] + adapter.model_root_inverse = np.linalg.inv( + pose_from_position_rotation(root_pose[:3], root_pose[3:]) + ) + solver.set_qpos_from_joint_names( + adapter.get_actived_joint_names(), + adapter.get_current_qpos(), + ) + 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, solver selection, 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() + 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", + ) + window.add_input_control(input_controller) + _NATIVE_IK_CONTROLLERS[id(robot), control_part] = controller + return controller, input_controller + class Gizmo: - """Control a rigid object, robot end-effector, or camera. + """Manage native robot IK or Viser control for 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. The manager + discovers robot TCP controls and owns their updates and cleanup. .. attention:: - Gizmo control currently supports exactly one simulation environment. + Gizmo control supports exactly one simulation environment. Args: - target: Simulation element controlled by this Gizmo. - cfg: Appearance configuration. + target: Robot, or a rigid object or camera controlled from Viser. + cfg: Gizmo appearance and robot IK configuration. control_part: Robot control part used for FK and IK. - 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 | _EmbodiChainIK | None = None + self._robot_adapter: _RobotGizmoAdapter | None = None + self._robot_end_link: str | None = None + self._robot_tcp_pose: np.ndarray | None = None + self._native_controller: IKGizmoController | None = None + self._native_input_controller: GizmoController | None = None + self._native_window = None + self._native_start_pending = self.cfg.ik_start_enabled - if self._target_type == "robot": - self._configure_robot() - self._desired_target_transform = self._read_target_pose() + from dexsim.kit.ik.interactive import KeyPressTracker - if enable_native: - self._gizmo = self._create_native_gizmo(self.cfg) - self._setup_native_gizmo() + self._native_toggle = KeyPressTracker(self.cfg.ik_toggle_key) + + if self._target_type == "robot": + self._control_part = _resolve_control_part(target, control_part) + _, self._robot_end_link, self._robot_tcp_pose = _resolve_robot_ik_chain( + target, self._control_part, self.cfg + ) + else: + self._desired_target_transform = self._read_target_pose() @property def target_type(self) -> str: @@ -158,11 +523,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 +535,83 @@ 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.target is None + or self._robot_end_link is None + or self._robot_tcp_pose is None + ): + raise RuntimeError("Robot Gizmo IK is not configured.") + link_pose = self.target.get_link_pose( + self._robot_end_link, env_ids=[0], to_matrix=True + )[0] + return self._as_pose_matrix( + link_pose + @ torch.as_tensor( + self._robot_tcp_pose, dtype=link_pose.dtype, device=link_pose.device + ), + self._target_device(), + ) + + def update_native(self, world: dexsim.World) -> None: + """Activate robot IK on startup or a toggle, then update its controller. + + Startup activation is opt-in; otherwise only the toggle key activates + IK. Registration and window reopening never write robot drive targets. + Explicit factory-created controllers remain owned by their caller. + + Args: + world: Simulation world with an open native window. + """ + if self._target_type != "robot" or self.target is None: + return + window = world.get_windows() + if window is None: + return + if self._native_controller is None: + existing = _NATIVE_IK_CONTROLLERS.get((id(self.target), self._control_part)) + if existing is not None: + return + start_enabled = self._native_start_pending + self._native_start_pending = False + if not self._native_toggle.pressed(world) and not start_enabled: + return + self._native_controller, self._native_input_controller = ( + create_robot_ik_gizmo_controller( + self.target, self._control_part, self.cfg, world=world + ) ) + self._native_window = window + # Creation starts with a visible target. Consume any held toggle + # so the controller does not immediately hide the new target. + self._native_controller.toggle.pressed(world) + return + if self._native_window is not window: + self.detach_native_window() + window.add_input_control(self._native_input_controller) + self._native_window = window + self._native_controller.update() + + def detach_native_window(self) -> None: + """Detach input before window close while retaining IK state and visibility.""" + if self._native_window is not None: + self._native_window.remove_input_control(self._native_input_controller) + self._native_window = None def _target_device(self) -> torch.device: if self.target is None: @@ -213,29 +632,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 +651,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 +661,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 +671,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,267 +689,67 @@ 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 - try: - self.target.set_local_pose(target_transform, env_ids=[0]) - return True - except Exception as error: - logger.log_error(f"Error updating camera pose: {error}") - return False - - def _update_rigid_object_pose(self, target_transform: torch.Tensor) -> bool: - if self.target is None or not isinstance(self.target, RigidObject): - return False - try: - self.target.set_local_pose(target_transform, env_ids=[0]) - return True - except Exception as error: - logger.log_error(f"Error updating rigid object pose: {error}") - return False - - def _update_robot_ik(self, target_transform: torch.Tensor) -> bool: - if self.target is None or not isinstance(self.target, Robot): - return False - if self._robot_arm_name 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], - ) - 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], - ) - return True - except Exception as error: - logger.log_error(f"Error in 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 or self.target 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() + if self._target_type != "robot": + self.target.set_local_pose(pending, env_ids=[0]) + return + if self._ik_solver is None: + self._setup_robot_ik_solver() + adapter, solver = self._robot_adapter, self._ik_solver + base_local = local_pose_from_world( + adapter.get_world_pose(), pending[0].detach().cpu().numpy() + ) + joint_names = adapter.get_actived_joint_names() + current_qpos = adapter.get_current_qpos() + solver.set_target_pose( + base_local[:3, 3], rotation_matrix_to_quat_xyzw(base_local[:3, :3]) + ) + solver.solve(joint_names, current_qpos, iterations=self.cfg.ik_iterations) + adapter.set_target_qpos(solver.qpos_for_joint_names(joint_names, current_qpos)) def toggle_visibility(self) -> bool: - """Toggle visibility and return the new state.""" - self.set_visible(not self._is_visible) - return self._is_visible + """Toggle Gizmo visibility and return the new state.""" + visible = not self.is_visible() + self.set_visible(visible) + return visible def set_visible(self, visible: bool) -> None: - """Set native and Viser Gizmo visibility.""" + """Set Gizmo visibility, including an activated native controller.""" self._is_visible = bool(visible) - if self._gizmo is not None: - self._gizmo.set_visible(self._is_visible) + if self._native_controller is not None: + self._native_controller._set_visible(self._is_visible) def is_visible(self) -> bool: - """Return whether this Gizmo should be visible.""" + """Return the active native visibility or configured Viser visibility.""" + if self._native_controller is not None: + return self._native_controller.enabled return self._is_visible - def 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.""" + self.detach_native_window() + if self._native_controller is not None: + controller = self._native_controller + env = controller.world.get_env() + env.remove_gizmo(controller.target_gizmo.gizmo) + env.remove_dummy_node(controller.target_gizmo.target_node) + key = (id(self.target), self._control_part) + if _NATIVE_IK_CONTROLLERS.get(key) is controller: + _NATIVE_IK_CONTROLLERS.pop(key) + self._native_controller = None + self._native_input_controller = None with self._state_lock: self._interaction_owner = None self._pending_target_transform = None 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 3d2dc2c54..1bc8fdd9d 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, @@ -99,6 +99,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, @@ -135,6 +136,27 @@ class SimulationManagerCfg: require a native window. """ + enable_entity_gizmo: bool = True + """Enable native object interaction when the first native window opens. + + Headless and Viser simulations do not automatically create native gizmos. + Explicit runtime enable/disable calls take precedence over this default; + reopening a window preserves the current DexSim controller state. + Robot IK interaction is controlled separately by ``robot_ik_gizmo``. + """ + + robot_ik_gizmo: GizmoCfg | None = field(default_factory=GizmoCfg) + """Automatically register robot parts with configured IK chain/TCP metadata. + + Native controls activate on the first I press; Viser creates its IK solver + on the first drag and requires ``visualization.allow_commands``. Pure + headless and multi-environment simulations do not register controls. + Defaults to DexSim IK. Use ``GizmoCfg(ik_solver="embodichain")`` to reuse + configured solvers such as Pink, or ``None`` to disable automatic setup. + ``GizmoCfg(ik_start_enabled=True)`` explicitly activates native controls on + their first update with an open window, without waiting for I. + """ + render_cfg: RenderCfg = field(default_factory=RenderCfg) """The rendering configuration parameters.""" @@ -190,6 +212,12 @@ class SimulationManagerCfg: def __post_init__(self) -> None: """Apply visualization-dependent simulation defaults.""" + if isinstance(self.robot_ik_gizmo, dict): + self.robot_ik_gizmo = GizmoCfg(**self.robot_ik_gizmo) + if self.robot_ik_gizmo is not None and not isinstance( + self.robot_ik_gizmo, GizmoCfg + ): + raise TypeError("robot_ik_gizmo must be a GizmoCfg, mapping, or None.") if self.visualization.backend == "viser": self.headless = True @@ -244,6 +272,7 @@ class SimulationManager: _instances = {} _cleanup_queue: queue.Queue = queue.Queue() + _DEFAULT_PLANE_GIZMO_TARGET_ID = (1 << 64) - 1 SUPPORTED_SENSOR_TYPES = { "Camera": Camera, @@ -297,6 +326,7 @@ def __init__( self._world.set_manual_update(True) self._window: Windows | None = None + self._auto_entity_gizmo_pending = sim_config.enable_entity_gizmo self._window_record_state: _WindowRecordState | None = None self._window_record_camera: object | None = None wr = sim_config.window_record @@ -339,6 +369,11 @@ def __init__( # gizmo management self._gizmos: Dict[str, object] = dict() # Store active gizmos + self._disabled_robot_gizmos: set[str] = set() + # ``(uid, control_part)`` of the Gizmo currently owned by the Viser + # click-to-pick feature, or ``None``. Only one picker Gizmo is kept at a + # time and user-created Gizmos are never touched. + self._picker_gizmo: tuple[str, str | None] | None = None # marker management self._markers: dict[str, _AxisMarkerGroup] = {} @@ -380,6 +415,8 @@ def __init__( if sim_config.headless is False: self._window = self._world.get_windows() + self.is_window_opened = self._window is not None + self._enable_default_entity_gizmo() self._is_constructed = True @@ -687,6 +724,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: @@ -971,6 +1010,9 @@ def open_window(self) -> bool: return True self._world.open_window() self._window = self._world.get_windows() + if self._window is None: + return False + self._enable_default_entity_gizmo() if ( self._window_record_hotkey_cfg is not None @@ -989,6 +1031,8 @@ def close_window(self) -> None: """Close the simulation window.""" if self.is_window_recording(): self.stop_window_record() + for _, gizmo in self.get_gizmo_items(): + gizmo.detach_native_window() self._world.close_window() self._window = None self._window_record_input_control = None @@ -1978,22 +2022,76 @@ 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. + + Called automatically on the first native window unless disabled in + :class:`SimulationManagerCfg`. Explicit calls also work headlessly. + DexSim retains this controller and its configuration across window + close/reopen; reopening does not reapply the default configuration. + + Args: + config: Native DexSim entity-Gizmo configuration. + + 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) + ) + self._auto_entity_gizmo_pending = False + 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 disable_entity_gizmo(self) -> None: + """Disable native object interaction, including future window opens. + + Call :meth:`enable_entity_gizmo` to explicitly re-enable interaction. + Robot IK controllers and Viser command permissions are independent. + """ + self._world.disable_entity_gizmo() + self._auto_entity_gizmo_pending = False + + def _enable_default_entity_gizmo(self) -> None: + # Apply the startup default once; DexSim owns subsequent controller + # state, including explicit native disable calls and window reopens. + if self._window is not None and self._auto_entity_gizmo_pending: + self.enable_entity_gizmo() + def enable_gizmo( self, uid: str, 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.). + """Register Gizmo control for a simulation target. + + Robot controls support native windows and Viser; native IK activates + on the configured toggle key. Explicit configuration takes precedence + over automatically registered controls. Args: uid: UID of the robot, rigid object, or camera sensor. control_part: Robot control part used for IK/FK. - gizmo_cfg: 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: Gizmo appearance and robot IK configuration. Returns: The created Gizmo, or ``None`` if setup failed. @@ -2001,12 +2099,17 @@ def enable_gizmo( # Create gizmo key combining uid and control_part gizmo_key = f"{uid}:{control_part}" if control_part else uid + if hasattr(self, "_disabled_robot_gizmos"): + self._disabled_robot_gizmos.discard(gizmo_key) + # Check if gizmo already exists if gizmo_key in self._gizmos: - logger.log_warning( - f"Gizmo for '{uid}' with control_part '{control_part}' already exists." - ) - return self._gizmos[gizmo_key] + if gizmo_cfg is not None: + self.disable_gizmo(uid, control_part) + if hasattr(self, "_disabled_robot_gizmos"): + self._disabled_robot_gizmos.discard(gizmo_key) + else: + return self._gizmos[gizmo_key] # Search for target object in different collections target = None @@ -2028,36 +2131,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"Gizmo registered for {object_type} '{uid}' with " + f"control_part '{control_part}'." ) except Exception as e: @@ -2071,14 +2152,23 @@ def enable_gizmo( return gizmo def disable_gizmo(self, uid: str, control_part: str | None = None) -> None: - """Disable and remove a Gizmo. + """Disable a Gizmo and prevent its automatic recreation. Args: uid: Target asset UID. - control_part: Robot control part, if applicable. + control_part: Robot control part. Omit to disable every part of a robot. """ gizmo_key = f"{uid}:{control_part}" if control_part else uid + if hasattr(self, "_disabled_robot_gizmos"): + self._disabled_robot_gizmos.add(gizmo_key) + robot = getattr(self, "_robots", {}).get(uid) + if robot is not None and control_part is None: + for key, gizmo in self.get_gizmo_items(): + if key != gizmo_key and gizmo.target is robot: + self.disable_gizmo(key) if gizmo_key not in self._gizmos: + if robot is not None: + return logger.log_warning( f"No gizmo found for '{uid}' with control_part '{control_part}'." ) @@ -2086,6 +2176,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() @@ -2201,17 +2293,120 @@ def process_visualization_commands(self) -> int: gizmo.end_interaction(source_id) return accepted + def _register_robot_gizmos(self) -> None: + """Discover IK-capable parts without constructing another solver.""" + cfg = getattr(self.sim_config, "robot_ik_gizmo", None) + if cfg is None or self.num_envs != 1: + return + native = self._window is not None + viser = ( + self.sim_config.visualization.backend == "viser" + and self.sim_config.visualization.allow_commands + ) + if not native and not viser: + return + for uid, robot in self._robots.items(): + solver_cfgs = robot.cfg.solver_cfg + if not isinstance(solver_cfgs, dict): + continue + for part, solver_cfg in solver_cfgs.items(): + key = f"{uid}:{part}" + if ( + uid in self._disabled_robot_gizmos + or key in self._disabled_robot_gizmos + or key in self._gizmos + or part not in (robot.control_parts or {}) + or not getattr(solver_cfg, "root_link_name", None) + or not getattr(solver_cfg, "end_link_name", None) + ): + continue + if any( + gizmo.target is robot and gizmo.control_part == part + for _, gizmo in self.get_gizmo_items() + ): + continue + try: + self.enable_gizmo(uid, part, deepcopy(cfg)) + except Exception as error: + self._disabled_robot_gizmos.add(key) + logger.log_warning( + f"Could not register robot Gizmo '{key}': {error}" + ) + def update_gizmos(self) -> None: - """Apply Viser commands and update all active Gizmos.""" + """Register robot controls, process interaction, and update active Gizmos.""" + self._register_robot_gizmos() + self.process_pick_commands() self.process_visualization_commands() for gizmo_key, gizmo in list( getattr(self, "_gizmos", {}).items() ): # Use list() to avoid modification during iteration if gizmo is not None: try: + if getattr(self, "_window", None) is not None: + gizmo.update_native(self._world) gizmo.update() except Exception as error: - logger.log_error(f"Error updating gizmo '{gizmo_key}': {error}") + logger.log_warning(f"Error updating gizmo '{gizmo_key}': {error}") + + def process_pick_commands(self) -> int: + """Apply queued Viser click-pick commands on the simulation thread. + + 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 @@ -2341,6 +2536,14 @@ def remove_asset(self, uid: str) -> bool: if uid in self._robots: robot = self._robots.pop(uid) + for key, gizmo in self.get_gizmo_items(): + if gizmo.target is robot: + self.disable_gizmo(key) + self._disabled_robot_gizmos = { + key + for key in self._disabled_robot_gizmos + if key != uid and not key.startswith(f"{uid}:") + } robot.destroy() self.notify_visualization_topology_changed() return True diff --git a/embodichain/lab/sim/utility/__init__.py b/embodichain/lab/sim/utility/__init__.py index 02f142a69..ae4c254cf 100644 --- a/embodichain/lab/sim/utility/__init__.py +++ b/embodichain/lab/sim/utility/__init__.py @@ -16,8 +16,9 @@ """Helper utilities for simulation state conversion, mesh/geometry handling, configuration transforms, keyboard interaction, and action/solver adaptation.""" +from __future__ import annotations + from .sim_utils import * from .mesh_utils import * -from .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 9f897a56b..000000000 --- a/embodichain/lab/sim/utility/gizmo_utils.py +++ /dev/null @@ -1,227 +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 __future__ import annotations - -from typing import Callable -from typing import TYPE_CHECKING -from dexsim.types import TransformMask - -if TYPE_CHECKING: - from embodichain.lab.sim.objects import Robot - -__all__ = ["create_gizmo_callback", "run_gizmo_robot_control_loop"] - - -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 -) -> 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. The loop - explicitly steps physics and paces execution with the configured timestep. - - 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) - - # 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: - step_start = time.perf_counter() - sim.update(step=1) - - # 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") - - time.sleep( - max(0.0, sim.sim_config.physics_dt - (time.perf_counter() - step_start)) - ) - - 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..2fde15f7f 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,16 @@ 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._gui_events.put( + _GuiEvent("pick", (self._run_id, self._scene_revision, 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 +361,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 +387,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 or self._frame_positions is None: + 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, @@ -1008,6 +1125,11 @@ def publish_manifest(self, manifest: SceneManifest) -> None: geometry_by_id = { geometry.geometry_id: geometry for geometry in manifest.geometries } + # Pick topology and poses must always come from the same revision. + # Defer clicks until a matching frame arrives after this manifest. + self._frame_positions = None + self._frame_wxyz = None + self._frame_visible = None self._frame_node_ids = tuple(node.node_id for node in manifest.nodes) node_indices = { node_id: index for index, node_id in enumerate(self._frame_node_ids) @@ -1024,6 +1146,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() @@ -1162,6 +1286,11 @@ def _apply_gui_events(self) -> None: event = self._gui_events.get_nowait() except queue.Empty: break + if event.category == "pick": + run_id, revision, click = event.value + if (run_id, revision) == (self._run_id, self._scene_revision): + self._handle_pick_click(click) + continue if event.category == "environment": env_id, visible = event.value self._env_visibility[int(env_id)] = bool(visible) @@ -1181,6 +1310,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 +1538,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 +1693,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 31d12c027..f5654b9c8 100644 --- a/embodichain/lab/visualization/scene_exporter.py +++ b/embodichain/lab/visualization/scene_exporter.py @@ -484,6 +484,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 944b258e8..15b7e7fd0 100644 --- a/examples/sim/gizmo/gizmo_camera.py +++ b/examples/sim/gizmo/gizmo_camera.py @@ -97,29 +97,30 @@ def main(): # Add camera to simulation camera = sim.add_sensor(sensor_cfg=camera_cfg) - # Wait for initialization - sim.update(step=round(0.2 / sim.sim_config.physics_dt)) + if sim.is_use_gpu_physics: + sim.init_gpu_physics() + sim.update(step=1) native_window_opened = False 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" ) @@ -140,11 +141,11 @@ def run_simulation( ) -> None: """Run the simulation loop with gizmo updates.""" step_count = 0 - last_time = time.time() + last_time = time.perf_counter() last_step = 0 if show_camera_window: - 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" @@ -152,7 +153,8 @@ def run_simulation( try: while True: - step_start = time.perf_counter() + frame_start = time.perf_counter() + # update() applies Gizmo commands before advancing one physics step. sim.update(step=1) # Update camera to get latest sensor data @@ -175,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"): @@ -202,7 +199,7 @@ def run_simulation( # Print simulation statistics and camera info if step_count % 1000 == 0: - current_time = time.time() + current_time = time.perf_counter() elapsed = current_time - last_time fps = ( sim.num_envs * (step_count - last_step) / elapsed @@ -223,9 +220,8 @@ def run_simulation( last_time = current_time last_step = step_count - time.sleep( - max(0.0, sim.sim_config.physics_dt - (time.perf_counter() - step_start)) - ) + elapsed = time.perf_counter() - frame_start + time.sleep(max(0.0, sim.sim_config.physics_dt - elapsed)) except KeyboardInterrupt: logger.log_info("\nStopping simulation...") diff --git a/examples/sim/gizmo/gizmo_object.py b/examples/sim/gizmo/gizmo_object.py index 690c7a6de..4e315d29c 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 @@ -58,6 +57,7 @@ def main(): # Create the simulation instance sim = SimulationManager(sim_cfg) + sim.set_manual_update(True) # Add two cubes to the scene cube1: RigidObject = sim.add_rigid_object( @@ -91,18 +91,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 +117,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") @@ -129,23 +135,27 @@ def run_simulation(sim: SimulationManager): step_count = 0 gizmo_enabled = True try: - last_time = time.time() + last_time = time.perf_counter() last_step = 0 while True: + frame_start = time.perf_counter() sim.update(step=1) step_count += 1 - # 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.disable_entity_gizmo() + else: + sim.disable_gizmo("cube1") + sim.disable_gizmo("cube2") gizmo_enabled = False # Print FPS every second if step_count % 1000 == 0: - current_time = time.time() + current_time = time.perf_counter() elapsed = current_time - last_time fps = ( sim.num_envs * (step_count - last_step) / elapsed @@ -155,6 +165,9 @@ def run_simulation(sim: SimulationManager): logger.log_info(f"Simulation step: {step_count}, FPS: {fps:.2f}") last_time = current_time last_step = step_count + + elapsed = time.perf_counter() - frame_start + time.sleep(max(0.0, sim.sim_config.physics_dt - elapsed)) except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: diff --git a/examples/sim/gizmo/gizmo_robot.py b/examples/sim/gizmo/gizmo_robot.py index 66c1da508..e8271088e 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 @@ -26,7 +24,11 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.solvers import PytorchSolverCfg +from embodichain.lab.sim.solvers import PinkSolverCfg, PytorchSolverCfg +from embodichain.lab.sim.objects import ( + GizmoCfg, + create_robot_ik_gizmo_controller, +) 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 @@ -47,6 +48,12 @@ def main(): description="Create a simulation scene with SimulationManager" ) add_env_launcher_args_to_parser(parser) + parser.add_argument( + "--ik-solver", + choices=("dexsim", "pytorch", "pink"), + default="dexsim", + help="IK solver; the native window always uses DexSim's IKGizmoController.", + ) args = parser.parse_args() # Configure the simulation @@ -66,6 +73,30 @@ def main(): ur10_urdf_path = get_data_path("UniversalRobots/UR10/UR10.urdf") gripper_urdf_path = get_data_path("DH_PGC_140_50_M/DH_PGC_140_50_M.urdf") + # Native IK needs only chain metadata. Build an EmbodiChain solver only + # when explicitly selected, and share the same TCP configuration. + gizmo_cfg = GizmoCfg( + ik_solver="dexsim" if args.ik_solver == "dexsim" else "embodichain", + ik_root_link_name="base_link", + ik_end_link_name="ee_link", + ik_tcp_pose=[ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.12], + [0.0, 0.0, 0.0, 1.0], + ], + ) + solver_cfg = None + if args.ik_solver != "dexsim": + solver_type = PinkSolverCfg if args.ik_solver == "pink" else PytorchSolverCfg + solver_cfg = { + "arm": solver_type( + root_link_name=gizmo_cfg.ik_root_link_name, + end_link_name=gizmo_cfg.ik_end_link_name, + tcp=gizmo_cfg.ik_tcp_pose, + ) + } + # Create UR10 robot robot_cfg = RobotCfg( uid="ur10_gizmo_test", @@ -76,53 +107,50 @@ def main(): ] ), control_parts={ - "arm": ["JOINT[0-9]"], + "arm": ["Joint[0-9]"], "hand": ["FINGER[1-2]"], }, - solver_cfg={ - "arm": PytorchSolverCfg( - end_link_name="ee_link", - root_link_name="base_link", - tcp=[ - [0.0, 1.0, 0.0, 0.0], - [-1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.12], - [0.0, 0.0, 0.0, 1.0], - ], - num_samples=30, - ) - }, + solver_cfg=solver_cfg, drive_pros=JointDrivePropertiesCfg( - stiffness={"JOINT[0-9]": 1e4, "FINGER[1-2]": 1e2}, - damping={"JOINT[0-9]": 1e3, "FINGER[1-2]": 1e1}, - max_effort={"JOINT[0-9]": 1e5, "FINGER[1-2]": 1e3}, + stiffness={"Joint[0-9]": 1e4, "FINGER[1-2]": 1e2}, + damping={"Joint[0-9]": 1e3, "FINGER[1-2]": 1e1}, + max_effort={"Joint[0-9]": 1e5, "FINGER[1-2]": 1e3}, drive_type="force", ), init_qpos=[0.0, -np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi / 2, 0.0, 0.0, 0.0], ) robot = sim.add_robot(cfg=robot_cfg) + if sim.is_use_gpu_physics: + sim.init_gpu_physics() # Set initial joint positions initial_qpos = torch.tensor( [[0.0, -np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi / 2, 0.0]], dtype=torch.float32, - device="cpu", + device=sim.device, ) joint_ids = robot.get_joint_ids("arm") + robot.set_qpos(qpos=initial_qpos, joint_ids=joint_ids, target=False) robot.set_qpos(qpos=initial_qpos, joint_ids=joint_ids) - - sim.update(step=round(0.2 / sim.sim_config.physics_dt)) + sim.update(step=1) native_window_opened = False if not args.headless: native_window_opened = sim.open_window() - # Enable gizmo using the new API - if native_window_opened or args.viser: + 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!") @@ -135,23 +163,27 @@ 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_time = time.perf_counter() last_step = 0 while True: - step_start = time.perf_counter() + frame_start = time.perf_counter() + if native_control is not None: + native_control[0].update() sim.update(step=1) step_count += 1 if step_count % 100 == 0: - current_time = time.time() + current_time = time.perf_counter() elapsed = current_time - last_time fps = ( sim.num_envs * (step_count - last_step) / elapsed @@ -162,9 +194,8 @@ def run_simulation(sim: SimulationManager): last_time = current_time last_step = step_count - time.sleep( - max(0.0, sim.sim_config.physics_dt - (time.perf_counter() - step_start)) - ) + elapsed = time.perf_counter() - frame_start + time.sleep(max(0.0, sim.sim_config.physics_dt - elapsed)) except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: diff --git a/examples/sim/gizmo/gizmo_scene.py b/examples/sim/gizmo/gizmo_scene.py deleted file mode 100644 index dd1ab7388..000000000 --- a/examples/sim/gizmo/gizmo_scene.py +++ /dev/null @@ -1,283 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- -""" -Gizmo Scene Example: Interactive scene with both robot and rigid object gizmos - -This example demonstrates how to create an interactive simulation scene with: -- A UR10 robot with gizmo control for end-effector manipulation -- A rigid object (cube) with gizmo control for direct manipulation -Both objects can be interactively controlled through their respective gizmos. -""" - -from __future__ import annotations - -import time -import torch -import numpy as np -import argparse -import cv2 - -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.cfg import ( - RenderCfg, - RobotCfg, - URDFCfg, - JointDrivePropertiesCfg, - RigidObjectCfg, - RigidBodyAttributesCfg, -) -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.shapes import CubeCfg -from embodichain.lab.sim.sensors import CameraCfg -from embodichain.lab.sim.solvers import PinkSolverCfg -from embodichain.data import get_data_path -from embodichain.utils import logger - - -def main(): - """Main function to create and run the simulation scene.""" - - parser = argparse.ArgumentParser( - description="Create a simulation scene with SimulationManager" - ) - add_env_launcher_args_to_parser(parser) - args = parser.parse_args() - - # Configure the simulation - sim_cfg = SimulationManagerCfg( - width=1920, - height=1080, - headless=True, - physics_dt=1.0 / 100.0, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), - visualization=visualization_cfg_from_args(args), - ) - - sim = SimulationManager(sim_cfg) - - # Get DexForce W1 URDF path - urdf_path = get_data_path("DexforceW1V021/DexforceW1_v02_1.urdf") - - # Create DexForce W1 robot - robot_cfg = RobotCfg( - uid="w1_gizmo_test", - urdf_cfg=URDFCfg( - components=[{"component_type": "humanoid", "urdf_path": urdf_path}] - ), - control_parts={"left_arm": ["LEFT_J[1-7]"], "right_arm": ["RIGHT_J[1-7]"]}, - solver_cfg={ - "left_arm": PinkSolverCfg( - urdf_path=urdf_path, - end_link_name="left_ee", - root_link_name="left_arm_base", - pos_eps=1e-2, - rot_eps=5e-2, - max_iterations=300, - dt=0.1, - ), - "right_arm": PinkSolverCfg( - urdf_path=urdf_path, - end_link_name="right_ee", - root_link_name="right_arm_base", - pos_eps=1e-2, - rot_eps=5e-2, - max_iterations=300, - dt=0.1, - ), - }, - drive_pros=JointDrivePropertiesCfg( - stiffness={"LEFT_J[1-7]": 1e4, "RIGHT_J[1-7]": 1e4}, - damping={"LEFT_J[1-7]": 1e3, "RIGHT_J[1-7]": 1e3}, - ), - ) - robot = sim.add_robot(cfg=robot_cfg) - - # Set initial joint positions for both arms - left_arm_qpos = torch.tensor( - [[0, -np.pi / 4, np.pi / 4, -np.pi / 2, 0.0, np.pi / 4, 0.0]], - dtype=torch.float32, - device="cpu", - ) - right_arm_qpos = torch.tensor( - [[0, np.pi / 4, -np.pi / 4, np.pi / 2, 0.0, -np.pi / 4, 0.0]], - dtype=torch.float32, - device="cpu", - ) - - left_joint_ids = robot.get_joint_ids("left_arm") - right_joint_ids = robot.get_joint_ids("right_arm") - - robot.set_qpos(qpos=left_arm_qpos, joint_ids=left_joint_ids) - robot.set_qpos(qpos=right_arm_qpos, joint_ids=right_joint_ids) - - # Create a rigid object (cube) positioned to the side of the robot - cube_cfg = RigidObjectCfg( - uid="interactive_cube", - shape=CubeCfg(size=[0.1, 0.1, 0.1]), - body_type="kinematic", - attrs=RigidBodyAttributesCfg( - mass=1.0, - dynamic_friction=0.5, - static_friction=0.5, - restitution=0.1, - ), - init_pos=[1.0, 0.0, 0.5], # Position to the side of the robot - ) - cube = sim.add_rigid_object(cube_cfg) - - camera_cfg = CameraCfg( - uid="scene_camera", - width=640, - height=480, - intrinsics=(320, 320, 320, 240), # fx, fy, cx, cy - near=0.1, - far=10.0, - enable_color=True, - enable_depth=True, - extrinsics=CameraCfg.ExtrinsicsCfg( - eye=(2.0, 2.0, 2.0), - target=(0.0, 0.0, 0.0), - up=(0.0, 0.0, 1.0), - ), - ) - camera = sim.add_sensor(sensor_cfg=camera_cfg) - - native_window_opened = False - if not args.headless: - native_window_opened = sim.open_window() - - # Enable gizmo for all assets after all are created and initialized - if native_window_opened or 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!") - return - - 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!") - return - - 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!") - return - - 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!") - return - else: - logger.log_warning( - "Gizmo interaction is disabled in headless mode without Viser." - ) - - logger.log_info("Gizmo Scene example started!") - if native_window_opened or args.viser: - logger.log_info("Four gizmos are active in the scene:") - logger.log_info( - "1. Left arm gizmo - Use to drag the left arm end-effector (EE)" - ) - logger.log_info( - "2. Right arm gizmo - Use to drag the right arm end-effector (EE)" - ) - logger.log_info("3. Cube gizmo - Use to drag and position the cube") - logger.log_info("4. Camera gizmo - Use to drag and orient the camera") - logger.log_info("Press Ctrl+C to stop the simulation") - - run_simulation(sim, show_camera_window=native_window_opened) - - -def run_simulation( - sim: SimulationManager, - *, - show_camera_window: bool, -) -> None: - step_count = 0 - # Get the camera instance by uid - camera = sim.get_sensor("scene_camera") - try: - last_time = time.time() - last_step = 0 - while True: - step_start = time.perf_counter() - sim.update(step=1) - step_count += 1 - - # Display camera view in a window every 5 steps - if show_camera_window and camera is not None and step_count % 5 == 0: - camera.update() - data = camera.get_data() - if "color" in data: - rgb_image = data["color"].cpu().numpy()[0, :, :, :3] - bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR) - cv2.putText( - bgr_image, - "Press 'h' to toggle camera gizmo visibility", - (10, 30), - cv2.FONT_HERSHEY_SIMPLEX, - 0.6, - (0, 255, 0), - 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") - - if step_count % 100 == 0: - current_time = time.time() - elapsed = current_time - last_time - fps = ( - sim.num_envs * (step_count - last_step) / elapsed - if elapsed > 0 - else 0 - ) - logger.log_info(f"Simulation step: {step_count}, FPS: {fps:.2f}") - last_time = current_time - last_step = step_count - - time.sleep( - max(0.0, sim.sim_config.physics_dt - (time.perf_counter() - step_start)) - ) - except KeyboardInterrupt: - logger.log_info("\nStopping simulation...") - finally: - if show_camera_window: - cv2.destroyAllWindows() - sim.destroy() - logger.log_info("Simulation terminated successfully") - - -if __name__ == "__main__": - main() diff --git a/examples/sim/gizmo/gizmo_w1.py b/examples/sim/gizmo/gizmo_w1.py deleted file mode 100644 index 49caa8c31..000000000 --- a/examples/sim/gizmo/gizmo_w1.py +++ /dev/null @@ -1,225 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- -""" -Gizmo-Robot Example: Test Gizmo class on a robot (UR10) -""" - -from __future__ import annotations - -import time -import torch -import numpy as np -import argparse - -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.cfg import ( - RenderCfg, - RobotCfg, - URDFCfg, - JointDrivePropertiesCfg, -) -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.solvers import PinkSolverCfg -from embodichain.data import get_data_path -from embodichain.utils import logger -from embodichain.lab.sim.robots.dexforce_w1.cfg import DexforceW1Cfg - - -def main(): - """Main function to create and run the simulation scene.""" - - # Parse command line arguments - parser = argparse.ArgumentParser( - description="Create a simulation scene with SimulationManager" - ) - add_env_launcher_args_to_parser(parser) - args = parser.parse_args() - - # Configure the simulation - sim_cfg = SimulationManagerCfg( - width=1920, - height=1080, - headless=True, - physics_dt=1.0 / 100.0, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), - visualization=visualization_cfg_from_args(args), - ) - - sim = SimulationManager(sim_cfg) - - cfg = DexforceW1Cfg.from_dict( - { - "uid": "w1_gizmo_test", - } - ) - cfg.solver_cfg["left_arm"].tcp = np.array( - [ - [1.0, 0.0, 0.0, 0.012], - [0.0, 1.0, 0.0, 0.04], - [0.0, 0.0, 1.0, 0.11], - [0.0, 0.0, 0.0, 1.0], - ] - ) - cfg.solver_cfg["right_arm"].tcp = np.array( - [ - [1.0, 0.0, 0.0, 0.012], - [0.0, 1.0, 0.0, -0.04], - [0.0, 0.0, 1.0, 0.11], - [0.0, 0.0, 0.0, 1.0], - ] - ) - - cfg.init_qpos = [ - 1.0000e00, - -2.0000e00, - 1.0000e00, - 0.0000e00, - -2.6921e-05, - -2.6514e-03, - -1.5708e00, - 1.4575e00, - -7.8540e-01, - 1.2834e-01, - 1.5708e00, - -2.2310e00, - -7.8540e-01, - 1.4461e00, - -1.5708e00, - 1.6716e00, - 7.8540e-01, - 7.6745e-01, - 0.0000e00, - 3.8108e-01, - 0.0000e00, - 0.0000e00, - 0.0000e00, - 0.0000e00, - 1.5000e00, - 0.0000e00, - 0.0000e00, - 0.0000e00, - 0.0000e00, - 1.5000e00, - 6.9974e-02, - 7.3950e-02, - 6.6574e-02, - 6.0923e-02, - 0.0000e00, - 6.7342e-02, - 7.0862e-02, - 6.3684e-02, - 5.7822e-02, - 0.0000e00, - ] - robot = sim.add_robot(cfg=cfg) - - # Set initial joint positions for both arms - # Left arm: 8 joints (WAIST + 7 LEFT_J), Right arm: 8 joints (WAIST + 7 RIGHT_J) - left_arm_qpos = torch.tensor( - [ - [0, 0, -np.pi / 4, np.pi / 4, -np.pi / 2, 0.0, np.pi / 4, 0.0] - ], # WAIST + LEFT_J[1-7] - dtype=torch.float32, - device="cpu", - ) - right_arm_qpos = torch.tensor( - [ - [0, 0, np.pi / 4, -np.pi / 4, np.pi / 2, 0.0, -np.pi / 4, 0.0] - ], # WAIST + RIGHT_J[1-7] - dtype=torch.float32, - device="cpu", - ) - - left_joint_ids = robot.get_joint_ids("left_arm") - right_joint_ids = robot.get_joint_ids("right_arm") - - robot.set_qpos(qpos=left_arm_qpos, joint_ids=left_joint_ids) - robot.set_qpos(qpos=right_arm_qpos, joint_ids=right_joint_ids) - - sim.update(step=round(0.2 / sim.sim_config.physics_dt)) - - native_window_opened = False - 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: - 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!") - return - - 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!") - return - else: - logger.log_warning( - "Gizmo interaction is disabled in headless mode without Viser." - ) - - logger.log_info("Gizmo-DexForce W1 example started!") - if native_window_opened or args.viser: - logger.log_info("Use the gizmos to drag both robot arms' end-effectors") - logger.log_info("Press Ctrl+C to stop the simulation") - - run_simulation(sim) - - -def run_simulation(sim: SimulationManager): - step_count = 0 - try: - last_time = time.time() - last_step = 0 - while True: - step_start = time.perf_counter() - sim.update(step=1) - step_count += 1 - - if step_count % 100 == 0: - current_time = time.time() - elapsed = current_time - last_time - fps = ( - sim.num_envs * (step_count - last_step) / elapsed - if elapsed > 0 - else 0 - ) - logger.log_info(f"Simulation step: {step_count}, FPS: {fps:.2f}") - last_time = current_time - last_step = step_count - - time.sleep( - max(0.0, sim.sim_config.physics_dt - (time.perf_counter() - step_start)) - ) - except KeyboardInterrupt: - logger.log_info("\nStopping simulation...") - finally: - sim.destroy() - logger.log_info("Simulation terminated successfully") - - -if __name__ == "__main__": - main() diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index 938e1c0bc..98fd43e53 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 Gizmo and manual physics stepping.""" from __future__ import annotations @@ -25,6 +23,7 @@ import argparse from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.objects import GizmoCfg from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.cfg import ( @@ -58,6 +57,7 @@ def main(): sim_device=args.device, render_cfg=RenderCfg(renderer=args.renderer), visualization=visualization_cfg_from_args(args), + robot_ik_gizmo=GizmoCfg(ik_start_enabled=True), ) sim = SimulationManager(sim_cfg) @@ -89,33 +89,26 @@ def main(): ), ) robot = sim.add_robot(cfg=robot_cfg) + if sim.is_use_gpu_physics: + sim.init_gpu_physics() # Set initial joint positions initial_qpos = torch.tensor( [[0, -np.pi / 2, np.pi / 2, 0.0, np.pi / 2, 0.0]], dtype=torch.float32, - device="cpu", + device=sim.device, ) joint_ids = robot.get_joint_ids("arm") + robot.set_qpos(qpos=initial_qpos, joint_ids=joint_ids, target=False) robot.set_qpos(qpos=initial_qpos, joint_ids=joint_ids) - sim.update(step=round(0.2 / sim.sim_config.physics_dt)) + sim.update(step=1) # Refresh link poses before creating the IK target. native_window_opened = False if not args.headless: native_window_opened = sim.open_window() - # Enable gizmo using the new API - if native_window_opened or args.viser: - sim.enable_gizmo( - uid="ur10_gizmo_test", - control_part="arm", - enable_native=native_window_opened, - ) - if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): - logger.log_error("Failed to enable gizmo!") - return - else: + if not native_window_opened and not args.viser: logger.log_warning( "Gizmo interaction is disabled in headless mode without Viser." ) @@ -123,23 +116,29 @@ 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( + "Native robot IK Gizmo starts enabled; press I to show or hide it" + ) logger.log_info("Press Ctrl+C to stop the simulation") run_simulation(sim) -def run_simulation(sim: SimulationManager): +def run_simulation(sim: SimulationManager) -> None: + """Advance physics; the manager owns native and Viser IK interaction.""" step_count = 0 try: - last_time = time.time() + last_time = time.perf_counter() last_step = 0 while True: - step_start = time.perf_counter() + frame_start = time.perf_counter() + # update() owns IK interaction, physics stepping, and Viser capture. sim.update(step=1) step_count += 1 if step_count % 100 == 0: - current_time = time.time() + current_time = time.perf_counter() elapsed = current_time - last_time fps = ( sim.num_envs * (step_count - last_step) / elapsed @@ -150,9 +149,8 @@ def run_simulation(sim: SimulationManager): last_time = current_time last_step = step_count - time.sleep( - max(0.0, sim.sim_config.physics_dt - (time.perf_counter() - step_start)) - ) + elapsed = time.perf_counter() - frame_start + time.sleep(max(0.0, sim.sim_config.physics_dt - elapsed)) except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: diff --git a/tests/gym/utils/test_gym_utils.py b/tests/gym/utils/test_gym_utils.py index a59bf74bd..a5be09d5b 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -1385,6 +1385,57 @@ def test_task_integration_rejects_removed_version_field( source_path=tmp_path / "env.yaml", ) + @pytest.mark.parametrize("suffix", ["yaml", "json"]) + @pytest.mark.parametrize("enabled", [None, False, True]) + def test_gym_config_preserves_entity_gizmo_startup_preference( + self, tmp_path: Path, suffix: str, enabled: bool | None + ) -> None: + """Task deployments default to native interaction and can opt out.""" + config = { + "id": "EmbodiedEnv-v1", + "env": {}, + "robot": {"uid": "TestRobot"}, + } + if enabled is not None: + config["enable_entity_gizmo"] = enabled + config_path = tmp_path / f"gym_config.{suffix}" + save_config(config_path, config) + + cfg = config_to_cfg( + load_config(config_path), manager_modules=DEFAULT_MANAGER_MODULES + ) + + assert cfg.sim_cfg.enable_entity_gizmo is (enabled is not False) + + @pytest.mark.parametrize("suffix", ["yaml", "json"]) + @pytest.mark.parametrize( + "settings", [None, {}, {"ik_solver": "embodichain"}, {"ik_start_enabled": True}] + ) + def test_gym_config_parses_automatic_robot_gizmo_settings( + self, tmp_path: Path, suffix: str, settings: dict | None + ) -> None: + """Deployments may disable automatic IK controls or select their solver.""" + path = tmp_path / f"gym_config.{suffix}" + save_config( + path, + { + "id": "EmbodiedEnv-v1", + "env": {}, + "robot": {"uid": "robot"}, + "robot_ik_gizmo": settings, + }, + ) + cfg = config_to_cfg(load_config(path), manager_modules=DEFAULT_MANAGER_MODULES) + if settings is None: + assert cfg.sim_cfg.robot_ik_gizmo is None + else: + assert cfg.sim_cfg.robot_ik_gizmo.ik_solver == settings.get( + "ik_solver", "dexsim" + ) + assert cfg.sim_cfg.robot_ik_gizmo.ik_start_enabled is settings.get( + "ik_start_enabled", False + ) + def test_yaml_gym_config_parses_to_cfg(self, tmp_path): config = { "id": "EmbodiedEnv-v1", diff --git a/tests/sim/objects/test_gizmo.py b/tests/sim/objects/test_gizmo.py index 0b1f3c6a5..2a525f025 100644 --- a/tests/sim/objects/test_gizmo.py +++ b/tests/sim/objects/test_gizmo.py @@ -17,86 +17,194 @@ from __future__ import annotations from types import SimpleNamespace +from pathlib import Path +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 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 set_qpos( + 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 +214,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 +231,295 @@ 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, + GizmoCfg(ik_root_link_name="base_link", ik_end_link_name="tool_link"), + control_part="arm", + ) + assert gizmo._ik_solver is None pose = torch.eye(4, dtype=torch.float32).unsqueeze(0) pose[0, 0, 3] = 0.5 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), ) + + +class _ConfiguredRobot(_FakeAdapterRobot): + def __init__(self) -> None: + super().__init__() + self.solver = SimpleNamespace( + joint_names=["joint_b", "joint_a"], + device=torch.device("cpu"), + root_link_name="base_link", + end_link_name="tool_link", + get_tcp=lambda: np.eye(4, dtype=np.float32), + ) + self.cfg.solver_cfg = {"arm": object()} + + def get_solver(self, name: str): + assert name == "arm" + return self.solver + + +@pytest.mark.parametrize("success,finite", [(True, True), (False, True), (True, False)]) +def test_configured_solver_preserves_joint_order_tcp_and_failed_seed( + success, finite +) -> None: + robot = _ConfiguredRobot() + calls = [] + + def solve(target, qpos_seed, return_all_solutions): + calls.append((target, qpos_seed)) + assert return_all_solutions is False + return torch.tensor([success]), torch.tensor( + [[[0.9, 0.7 if finite else float("nan")]]] + ) + + robot.solver.get_ik = solve + tcp = np.eye(4, dtype=np.float32) + tcp[0, 3] = 0.2 + adapter, solver, _, _ = gizmo_module._build_robot_ik( + robot, "arm", GizmoCfg(ik_solver="embodichain", ik_tcp_pose=tcp) + ) + solver.set_target_pose(np.array([0.5, 0.2, 0.1]), np.array([0.0, 0.0, 0.0, 1.0])) + assert solver.target_state_changed() + names, seed = adapter.get_actived_joint_names(), adapter.get_current_qpos() + solver.solve(names, seed) + torch.testing.assert_close(calls[0][1], torch.tensor([[0.3, 0.1]])) + torch.testing.assert_close(calls[0][0][0, :3, 3], torch.tensor([0.3, 0.2, 0.1])) + np.testing.assert_allclose( + solver.qpos_for_joint_names(names, seed), + [0.7, 0.9] if success and finite else seed, + ) + solver.reset_target_state_changed() + assert not solver.target_state_changed() + solver.target_state()["position"][0] += 0.1 + assert solver.target_state_changed() + np.testing.assert_array_equal(robot.solver.get_tcp(), np.eye(4)) + + +@pytest.mark.parametrize( + "configuration,match", + [ + ({"ik_solver": "unknown"}, "ik_solver must"), + ({"ik_solver": "embodichain"}, "configured EmbodiChain solver"), + ({"ik_iterations": 0}, "ik_iterations"), + ], +) +def test_robot_ik_rejects_invalid_configuration(configuration, match) -> None: + cfg = GizmoCfg( + ik_root_link_name="base_link", ik_end_link_name="tool_link", **configuration + ) + with pytest.raises(ValueError, match=match): + gizmo_module._build_robot_ik(_FakeAdapterRobot(), "arm", cfg) + + +def test_configured_solver_rejects_chain_and_joint_mismatch() -> None: + robot = _ConfiguredRobot() + with pytest.raises(ValueError, match="links must match"): + gizmo_module._build_robot_ik( + robot, "arm", GizmoCfg(ik_solver="embodichain", ik_root_link_name="other") + ) + robot.solver.joint_names = ["other_joint"] + with pytest.raises(ValueError, match="joints must match"): + gizmo_module._build_robot_ik(robot, "arm", GizmoCfg(ik_solver="embodichain")) + + +_PLANAR_URDF = """ + + + + + + + + + + + + + + + + + +""" + + +def _planar_pose(angle: float, x: float = 0.0) -> np.ndarray: + pose = np.eye(4, dtype=np.float32) + c, s = np.cos(angle), np.sin(angle) + pose[:2, :2] = [[c, -s], [s, c]] + pose[0, 3] = x + return pose + + +class _PlanarRobot(_ConfiguredRobot): + def __init__(self, urdf_path: Path, backend: str) -> None: + super().__init__() + urdf_path.write_text(_PLANAR_URDF) + self.cfg.fpath = str(urdf_path) + self.cfg.solver_cfg = None + # Move the runtime chain root away from both the robot root and the + # URDF rest transform, as happens when a torso or upstream joint moves. + self.root_pose = _planar_pose(0.8, 0.7) + self.root_pose[1:3, 3] = [-0.2, 0.6] + if backend == "embodichain": + pytest.importorskip("pink") + from embodichain.lab.sim.solvers import PinkSolverCfg + + self.solver = PinkSolverCfg( + urdf_path=str(urdf_path), + joint_names=["joint_a", "joint_b"], + root_link_name="base_link", + end_link_name="tool_link", + max_iterations=150, + pos_eps=1e-5, + rot_eps=1e-5, + show_ik_warnings=False, + ).init_solver(device=torch.device("cpu")) + self.cfg.solver_cfg = {"arm": self.solver.cfg} + + def tool_pose(self, qpos: np.ndarray) -> np.ndarray: + return ( + self.root_pose + @ _planar_pose(qpos[0]) + @ _planar_pose(qpos[1], 0.5) + @ _planar_pose(0.0, 0.5) + ) + + def get_link_pose(self, link_name, env_ids, to_matrix=False): + assert env_ids == [0] and to_matrix + pose = ( + self.root_pose + if link_name == "base_link" + else self.tool_pose(self.current_qpos[0, [0, 2]].numpy()) + ) + return torch.tensor(pose).unsqueeze(0) + + +@pytest.mark.parametrize( + "backend,ik_device", + [ + ("dexsim", "cpu"), + ("embodichain", "cpu"), + pytest.param("dexsim", "cuda:0", marks=pytest.mark.gpu), + ], +) +@pytest.mark.parametrize("frontend", ["native", "viser"]) +def test_real_ik_solvers_follow_live_chain_root_and_tcp( + tmp_path, monkeypatch, backend, ik_device, frontend +) -> None: + """Run real Newton/Pink IK; replace only native window rendering and input.""" + robot = _PlanarRobot(tmp_path / "robot.urdf", backend) + tcp = np.eye(4, dtype=np.float32) + tcp[0, 3] = 0.1 + cfg = GizmoCfg( + ik_solver=backend, + ik_root_link_name="base_link", + ik_end_link_name="tool_link", + ik_tcp_pose=tcp, + ik_iterations=100, + ik_device=ik_device, + ) + if frontend == "native": + import dexsim.kit.ik.interactive as interactive + import dexsim.engine + + monkeypatch.setattr( + interactive, + "setup_target_gizmo", + lambda *args, **kwargs: SimpleNamespace( + gizmo=SimpleNamespace( + set_visible=lambda _: None, follow=lambda _: None + ), + target_node=SimpleNamespace(set_world_pose=lambda _: None), + ), + ) + monkeypatch.setattr(dexsim.engine, "GizmoController", lambda: object()) + world = SimpleNamespace( + get_windows=lambda: SimpleNamespace(add_input_control=lambda _: None), + key_state=lambda _: False, + ) + controller, _ = create_robot_ik_gizmo_controller(robot, cfg=cfg, world=world) + assert isinstance(controller, interactive.IKGizmoController) + assert not controller.update() + adapter = controller.robot + else: + monkeypatch.setattr(gizmo_module, "Robot", _PlanarRobot) + controller = Gizmo(robot, cfg, "arm") + adapter = controller._robot_adapter + + # Move again after construction to verify live root following on update. + robot.root_pose = _planar_pose(-0.2, 0.3) @ robot.root_pose + target = robot.tool_pose(np.array([0.4, -0.2])) @ tcp + if frontend == "native": + local_target = gizmo_module.local_pose_from_world( + adapter.get_world_pose(), target + ) + controller.ik.set_target_pose( + local_target[:3, 3], + gizmo_module.rotation_matrix_to_quat_xyzw(local_target[:3, :3]), + ) + assert controller.update() + else: + assert controller.request_local_pose(target, source_id="viser:test") + controller.update() + write = robot.write_calls[-1] + assert write["target"] is True + assert write["joint_ids"] == [0, 2] + assert write["env_ids"] == [0] + actual = robot.tool_pose(write["qpos"][0].numpy()) @ tcp + np.testing.assert_allclose(actual, target, atol=1e-3) diff --git a/tests/sim/test_interactive_stepping.py b/tests/sim/test_interactive_stepping.py index f8da01e40..4ac383333 100644 --- a/tests/sim/test_interactive_stepping.py +++ b/tests/sim/test_interactive_stepping.py @@ -29,8 +29,7 @@ "module_name", [ "examples.sim.gizmo.gizmo_robot", - "examples.sim.gizmo.gizmo_w1", - "examples.sim.gizmo.gizmo_scene", + "examples.sim.gizmo.gizmo_object", "examples.sim.gizmo.gizmo_camera", "scripts.tutorials.sim.gizmo_robot", ], @@ -63,8 +62,8 @@ def sleep(duration: float) -> None: sim = SimpleNamespace( sim_config=SimpleNamespace(physics_dt=physics_dt), num_envs=1, + is_use_gpu_physics=False, update=update, - get_sensor=lambda _uid: None, has_gizmo=lambda _uid: False, destroy=Mock(), ) @@ -79,8 +78,6 @@ def sleep(duration: float) -> None: if module_name.endswith("gizmo_camera"): module.run_simulation(sim, Mock(), show_camera_window=False) - elif module_name.endswith("gizmo_scene"): - module.run_simulation(sim, show_camera_window=False) else: module.run_simulation(sim) diff --git a/tests/sim/test_robot_gizmo_lifecycle.py b/tests/sim/test_robot_gizmo_lifecycle.py new file mode 100644 index 000000000..706e8365e --- /dev/null +++ b/tests/sim/test_robot_gizmo_lifecycle.py @@ -0,0 +1,347 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from dexsim.kit.ik.interactive import KeyPressTracker + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.objects import GizmoCfg +from embodichain.lab.sim.objects import gizmo as gizmo_module +from embodichain.lab.visualization import VisualizationCfg + +pytestmark = pytest.mark.no_sim + + +class _Robot: + __hash__ = None # Robot inherits dataclass equality and is not hashable. + num_instances = 1 + device = torch.device("cpu") + + def __init__(self) -> None: + self.control_parts = { + "left": ["left_joint"], + "right": ["right_joint"], + "hand": ["finger"], + } + self.cfg = SimpleNamespace( + uid="robot", + solver_cfg={ + part: SimpleNamespace( + root_link_name="base", end_link_name=f"{part}_tool" + ) + for part in ("left", "right") + }, + ) + self.pose = torch.eye(4).unsqueeze(0) + self.destroy = Mock() + + def get_solver(self, part: str) -> SimpleNamespace: + cfg = self.cfg.solver_cfg[part] + return SimpleNamespace(**vars(cfg), get_tcp=lambda: np.eye(4)) + + def get_link_pose( + self, link: str, *, env_ids: list[int], to_matrix: bool + ) -> torch.Tensor: + return self.pose.clone() + + +class _Window: + def __init__(self) -> None: + self.controls: list[object] = [] + + def add_input_control(self, control: object) -> None: + self.controls.append(control) + + def remove_input_control(self, control: object) -> None: + self.controls.remove(control) + + +class _World: + def __init__(self) -> None: + self.window = _Window() + self.key_down = False + self.env = SimpleNamespace(remove_gizmo=Mock(), remove_dummy_node=Mock()) + + def get_windows(self) -> _Window | None: + return self.window + + def get_env(self) -> SimpleNamespace: + return self.env + + def key_state(self, key: object) -> bool: + return self.key_down + + def close_window(self) -> None: + self.window = None + + def open_window(self) -> None: + self.window = _Window() + + +class _Controller: + def __init__(self, world: _World, cfg: GizmoCfg) -> None: + self.world = world + self.toggle = KeyPressTracker(cfg.ik_toggle_key) + self.enabled = True + self.target_gizmo = SimpleNamespace(gizmo=object(), target_node=object()) + + def update(self) -> None: + if self.toggle.pressed(self.world): + self.enabled = not self.enabled + + def _set_visible(self, visible: bool) -> None: + self.enabled = visible + + +@pytest.fixture +def managed_robot(monkeypatch: pytest.MonkeyPatch): + """Use the real manager/Gizmo lifecycle with rendering and IK construction stubbed.""" + monkeypatch.setattr(gizmo_module, "Robot", _Robot) + robot = _Robot() + sim = object.__new__(SimulationManager) + sim.sim_config = SimulationManagerCfg(headless=True) + sim.num_envs = 1 + sim._world = _World() + sim._window = sim._world.window + sim.is_window_opened = True + sim._window_record_state = None + sim._window_record_hotkey_cfg = None + sim._window_camera_pose_hotkey_cfg = None + sim._auto_entity_gizmo_pending = False + sim._visualization_runtime = None + sim._visualization_topology_revision = 0 + sim._gizmos = {} + sim._disabled_robot_gizmos = set() + sim._picker_gizmo = None + sim._robots = {robot.cfg.uid: robot} + for registry in ( + "_rigid_objects", + "_rigid_object_groups", + "_articulations", + "_soft_objects", + "_cloth_objects", + "_sensors", + ): + setattr(sim, registry, {}) + sim.process_pick_commands = Mock(return_value=0) + sim.process_visualization_commands = Mock(return_value=0) + + def create(robot, part, cfg, *, world): + controller = _Controller(world, cfg) + control = object() + world.get_windows().add_input_control(control) + return controller, control + + factory = Mock(side_effect=create) + monkeypatch.setattr(gizmo_module, "create_robot_ik_gizmo_controller", factory) + yield sim, robot, factory + for _, gizmo in sim.get_gizmo_items(): + gizmo.destroy() + + +@pytest.mark.parametrize( + "frontend", ["native", "viser", "read_only", "headless", "batch", "disabled"] +) +def test_automatic_registration_respects_interaction_boundaries( + managed_robot, frontend: str +) -> None: + sim, robot, factory = managed_robot + if frontend in {"viser", "read_only", "headless"}: + sim._window = None + if frontend in {"viser", "read_only"}: + sim.sim_config.visualization = VisualizationCfg( + backend="viser", allow_commands=frontend == "viser" + ) + if frontend == "batch": + sim.num_envs = 2 + if frontend == "disabled": + sim.sim_config.robot_ik_gizmo = None + sim.update_gizmos() + sim.update_gizmos() + expected = ( + {"robot:left", "robot:right"} if frontend in {"native", "viser"} else set() + ) + assert set(sim.list_gizmos()) == expected + factory.assert_not_called() + for _, gizmo in sim.get_gizmo_items(): + assert gizmo._ik_solver is None + robot.pose[0, 0, 3] = 0.5 + assert gizmo.get_control_pose()[0, 0, 3] == pytest.approx(0.5) + + +def test_native_activation_and_reopen_preserve_controller_and_visibility( + managed_robot, +) -> None: + sim, _, factory = managed_robot + sim.update_gizmos() + sim._world.key_down = True + sim.update_gizmos() + controls = [gizmo._native_controller for _, gizmo in sim.get_gizmo_items()] + assert factory.call_count == 2 + sim.update_gizmos() # Holding I does not toggle a second time. + assert all(control.enabled for control in controls) + sim._world.key_down = False + sim.update_gizmos() + sim._world.key_down = True + sim.update_gizmos() + assert not any(control.enabled for control in controls) + + old_window = sim._window + sim.close_window() + assert old_window.controls == [] + assert sim.open_window() + sim.update_gizmos() + assert factory.call_count == 2 + assert len(sim._window.controls) == 2 + assert [g._native_controller for _, g in sim.get_gizmo_items()] == controls + assert not any(control.enabled for control in controls) + + +def test_explicit_disable_and_solver_override_take_precedence(managed_robot) -> None: + sim, _, _ = managed_robot + sim.disable_gizmo("robot", "left") + sim.update_gizmos() + assert set(sim.list_gizmos()) == {"robot:right"} + sim.enable_gizmo("robot", "right", GizmoCfg(ik_solver="embodichain")) + right = sim.get_gizmo("robot", "right") + sim.update_gizmos() + assert sim.get_gizmo("robot", "right") is right + assert right.cfg.ik_solver == "embodichain" + sim.enable_gizmo("robot", "left") + assert set(sim.list_gizmos()) == {"robot:left", "robot:right"} + + +def test_start_enabled_waits_for_window_and_preserves_later_visibility( + managed_robot, +) -> None: + sim, _, factory = managed_robot + sim.sim_config.robot_ik_gizmo.ik_start_enabled = True + sim.close_window() + sim.update_gizmos() + factory.assert_not_called() + sim.open_window() + sim.update_gizmos() # No I press is needed. + controls = [g._native_controller for _, g in sim.get_gizmo_items()] + assert factory.call_count == 2 + assert all(control.enabled for control in controls) + sim._world.key_down = True + sim.update_gizmos() + sim.update_gizmos() + assert not any(control.enabled for control in controls) + sim.close_window() + sim.open_window() + sim.update_gizmos() + assert [g._native_controller for _, g in sim.get_gizmo_items()] == controls + assert not any(control.enabled for control in controls) + assert factory.call_count == 2 + + +def test_start_enabled_failure_waits_for_a_key_before_retry(managed_robot) -> None: + sim, _, factory = managed_robot + sim.sim_config.robot_ik_gizmo.ik_start_enabled = True + create = factory.side_effect + factory.side_effect = ValueError("invalid chain") + sim.update_gizmos() + sim.update_gizmos() + assert factory.call_count == 2 + assert not sim._window.controls + factory.side_effect = create + sim._world.key_down = True + sim.update_gizmos() + assert factory.call_count == 4 + assert all(g._native_controller.enabled for _, g in sim.get_gizmo_items()) + + +def test_whole_robot_disable_cleans_activated_controls(managed_robot) -> None: + sim, _, factory = managed_robot + sim._world.key_down = True + sim.update_gizmos() + assert not sim.toggle_gizmo_visibility("robot", "left") + assert not sim.get_gizmo("robot", "left")._native_controller.enabled + sim.set_gizmo_visibility("robot", True, "left") + assert sim.get_gizmo("robot", "left")._native_controller.enabled + sim.disable_gizmo("robot") + sim.update_gizmos() + assert not sim.list_gizmos() + assert not sim._window.controls + assert factory.call_count == 2 + + +def test_native_activation_failure_does_not_interrupt_other_gizmos( + managed_robot, +) -> None: + sim, _, factory = managed_robot + create = factory.side_effect + + def fail_left(robot, part, cfg, *, world): + if part == "left": + raise ValueError("invalid chain") + return create(robot, part, cfg, world=world) + + factory.side_effect = fail_left + sim._world.key_down = True + sim.update_gizmos() + sim.update_gizmos() + assert factory.call_count == 2 # A held key does not repeatedly retry a failure. + assert sim.get_gizmo("robot", "left")._native_controller is None + assert sim.get_gizmo("robot", "right")._native_controller is not None + + +def test_robot_gizmo_setting_rejects_boolean() -> None: + with pytest.raises(TypeError, match="robot_ik_gizmo"): + SimulationManagerCfg(robot_ik_gizmo=False) + + +def test_robot_removal_cleans_native_controls_and_allows_same_uid_again( + managed_robot, +) -> None: + sim, robot, _ = managed_robot + sim._world.key_down = True + sim.update_gizmos() + assert sim.remove_asset("robot") + assert not sim.list_gizmos() + assert not sim._window.controls + assert sim._world.env.remove_gizmo.call_count == 2 + assert sim._world.env.remove_dummy_node.call_count == 2 + robot.destroy.assert_called_once() + sim._robots["robot"] = _Robot() + sim._world.key_down = False + sim.update_gizmos() + assert set(sim.list_gizmos()) == {"robot:left", "robot:right"} + + +@pytest.mark.parametrize("start_enabled", [False, True]) +def test_explicit_native_factory_controller_is_not_duplicated( + managed_robot, monkeypatch: pytest.MonkeyPatch, start_enabled: bool +) -> None: + sim, robot, factory = managed_robot + sim.sim_config.robot_ik_gizmo.ik_start_enabled = start_enabled + explicit = Mock() + monkeypatch.setitem( + gizmo_module._NATIVE_IK_CONTROLLERS, (id(robot), "left"), explicit + ) + sim._world.key_down = True + sim.update_gizmos() + assert factory.call_count == 1 + assert factory.call_args.args[1] == "right" + explicit.update.assert_not_called() diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index 0616eb451..fd20ebfcd 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 @@ -37,6 +38,7 @@ from embodichain.lab.sim.sensors import Camera, CameraCfg, StereoCamera, StereoCameraCfg from embodichain.lab.visualization import ( GizmoCommand, + PickCommand, PointCloudOverlay, SceneOverlays, VisualizationCfg, @@ -108,12 +110,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 @@ -121,6 +146,32 @@ def thread_rt(self) -> FakeThreadRuntime: 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) + if self.entity_gizmo is None: + self.entity_gizmo = FakeEntityGizmo() + self.entity_gizmo.active = True + return self.entity_gizmo + + def disable_entity_gizmo(self) -> None: + if self.entity_gizmo is not None: + self.entity_gizmo.active = False + + def get_entity_gizmo(self) -> object | None: + if self.entity_gizmo is not None and self.entity_gizmo.active: + return self.entity_gizmo + return None + + def open_window(self) -> None: + self.window_open_count += 1 + self.window_closed = False + + 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 and point clouds.""" @@ -188,17 +239,32 @@ def end_interaction(self, source_id: str) -> bool: return True -def _make_sim_manager(window: object | None = None) -> SimulationManager: +def _make_sim_manager( + window: object | None = None, *, enable_entity_gizmo: bool = True +) -> SimulationManager: """Create a minimally initialized simulation manager for recorder tests.""" sim = object.__new__(SimulationManager) sim.instance_id = 0 - sim.sim_config = SimpleNamespace(width=64, height=48) + sim.sim_config = SimpleNamespace( + width=64, + height=48, + enable_entity_gizmo=enable_entity_gizmo, + visualization=SimpleNamespace(backend="none"), + ) sim._window = window + sim._auto_entity_gizmo_pending = enable_entity_gizmo sim._window_record_state = None sim._window_record_camera = None sim._window_record_save_threads = [] + 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 @@ -414,6 +480,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() @@ -483,6 +762,7 @@ def test_open_window_allows_native_backend() -> None: sim._window_record_input_control = None sim._window_camera_pose_hotkey_cfg = None sim._window_camera_pose_input_control = None + sim._auto_entity_gizmo_pending = True sim.is_window_opened = False opened = sim.open_window() @@ -507,6 +787,111 @@ 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_enables_entity_gizmo_by_default() -> None: + sim = _make_sim_manager() + + assert sim.open_window() + + assert sim.is_window_opened is True + assert sim._world.window_open_count == 1 + assert sim._world.entity_gizmo_configs == [None] + assert sim._world.get_entity_gizmo().external_targets[0][2] is sim._default_plane + + +def test_entity_gizmo_can_be_disabled_in_startup_configuration() -> None: + cfg = SimulationManagerCfg() + assert cfg.enable_entity_gizmo is True + cfg = SimulationManagerCfg(enable_entity_gizmo=False) + sim = _make_sim_manager(enable_entity_gizmo=cfg.enable_entity_gizmo) + + assert sim.open_window() + sim.close_window() + assert sim.open_window() + assert sim._world.get_entity_gizmo() is None + assert sim._world.entity_gizmo_configs == [] + + +@pytest.mark.parametrize("before_first_window", [True, False]) +def test_explicit_entity_gizmo_disable_survives_window_reopen( + before_first_window: bool, +) -> None: + sim = _make_sim_manager() + if not before_first_window: + assert sim.open_window() + sim.disable_entity_gizmo() + sim.close_window() + assert sim.open_window() + assert sim._world.get_entity_gizmo() is None + assert len(sim._world.entity_gizmo_configs) == (0 if before_first_window else 1) + + controller = sim.enable_entity_gizmo() + sim.close_window() + assert sim.open_window() + assert sim._world.get_entity_gizmo() is controller + + +def test_native_entity_gizmo_disable_is_not_overridden_on_reopen() -> None: + sim = _make_sim_manager() + assert sim.open_window() + sim._world.disable_entity_gizmo() + sim.close_window() + assert sim.open_window() + assert sim._world.get_entity_gizmo() is None + assert sim._world.entity_gizmo_configs == [None] + + +def test_window_reopen_preserves_explicit_entity_gizmo_configuration() -> None: + sim = _make_sim_manager() + config = object() + controller = sim.enable_entity_gizmo(config) + assert sim.open_window() + sim.close_window() + assert sim.open_window() + assert sim._world.get_entity_gizmo() is controller + assert sim._world.entity_gizmo_configs == [config] + + +def test_failed_window_open_does_not_consume_gizmo_startup_default() -> None: + sim = _make_sim_manager() + window = sim._world.window + sim._world.window = None + assert not sim.open_window() + assert not sim.is_window_opened + assert sim._world.entity_gizmo_configs == [] + sim._world.window = window + assert sim.open_window() + assert sim._world.entity_gizmo_configs == [None] + + +def test_close_window_leaves_entity_gizmo_lifecycle_to_dexsim() -> None: + sim = _make_sim_manager(window=object()) + controller = sim.enable_entity_gizmo() + + 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( @@ -519,7 +904,22 @@ def test_start_visualization_rejects_open_native_window() -> None: sim.start_visualization() -def test_constructor_starts_visualization_after_default_scene(monkeypatch) -> None: +@pytest.mark.parametrize( + "headless,entity_gizmo,backend,expected_gizmo", + [ + (False, True, "none", True), + (False, False, "none", False), + (True, True, "none", False), + (False, True, "viser", False), + ], +) +def test_constructor_starts_visualization_after_default_scene( + monkeypatch: pytest.MonkeyPatch, + headless: bool, + entity_gizmo: bool, + backend: str, + expected_gizmo: bool, +) -> None: lifecycle: list[str] = [] world = MagicMock() world.set_manual_update.side_effect = lambda _enable: lifecycle.append( @@ -583,9 +983,22 @@ def start_visualization(sim: SimulationManager) -> None: "start_visualization", start_visualization, ) + monkeypatch.setattr( + SimulationManager, + "enable_entity_gizmo", + lambda _self: lifecycle.append("entity_gizmo"), + ) sim = object.__new__(SimulationManager) - SimulationManager.__init__(sim, SimulationManagerCfg(num_envs=3)) + SimulationManager.__init__( + sim, + SimulationManagerCfg( + num_envs=3, + headless=headless, + enable_entity_gizmo=entity_gizmo, + visualization=VisualizationCfg(backend=backend), + ), + ) world.set_manual_update.assert_called_once_with(True) assert lifecycle == [ @@ -597,7 +1010,7 @@ def start_visualization(sim: SimulationManager) -> None: "lighting", "arenas", "visualization:3", - ] + ] + (["entity_gizmo"] if expected_gizmo else []) def test_remove_asset_marks_visualization_topology_dirty() -> None: @@ -613,6 +1026,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_gizmo_robot_tutorial.py b/tests/visualization/test_gizmo_robot_tutorial.py new file mode 100644 index 000000000..0b30f34d4 --- /dev/null +++ b/tests/visualization/test_gizmo_robot_tutorial.py @@ -0,0 +1,48 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock, call + +import pytest + +from scripts.tutorials.sim import gizmo_robot + +pytestmark = pytest.mark.no_sim + + +@pytest.mark.parametrize("work_duration", [0.002, 0.050]) +def test_gizmo_loop_advances_manual_physics_and_paces_frames( + monkeypatch: pytest.MonkeyPatch, + work_duration: float, +) -> None: + """Apply IK before stepping, avoid duplicate updates, and clean up on exit.""" + physics_dt = 0.01 + calls = Mock() + sim = calls.sim + sim.sim_config = SimpleNamespace(physics_dt=physics_dt) + clock = Mock(side_effect=[0.0, 0.0, work_duration]) + sleep = Mock(side_effect=KeyboardInterrupt) + monkeypatch.setattr(gizmo_robot.time, "perf_counter", clock) + monkeypatch.setattr(gizmo_robot.time, "sleep", sleep) + + gizmo_robot.run_simulation(sim) + + expected_calls = [call.sim.update(step=1), call.sim.destroy()] + assert calls.mock_calls == expected_calls + sleep.assert_called_once_with(pytest.approx(max(0.0, physics_dt - work_duration))) 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 aac24485c..ee8140b7a 100644 --- a/tests/visualization/test_viser_backend.py +++ b/tests/visualization/test_viser_backend.py @@ -17,8 +17,10 @@ from __future__ import annotations from types import SimpleNamespace +from dataclasses import replace import numpy as np +import pytest from embodichain.lab.visualization import ( CameraImage, @@ -280,6 +282,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: @@ -884,6 +897,216 @@ 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), + ) + ) + backend.poll() + 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), + ) + ) + + backend.poll() + 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), + ) + ) + + backend.poll() + assert len(pick_commands) == 1 + assert pick_commands[0].node_id is None + backend.stop() + + +@pytest.mark.parametrize("node_count", [1, 2]) +def test_picker_waits_for_matching_frame_after_topology_change(node_count: int) -> None: + backend, server, commands = _make_pick_backend() + manifest, frame = _make_pickable_scene() + backend.start() + try: + backend.publish_manifest(manifest) + assert backend.publish_frame(frame) + backend._pick_enabled = True + event = SimpleNamespace( + client_id="client-a", + ray_origin=np.array([0.0, 0.0, 5.0]), + ray_direction=np.array([0.0, 0.0, -1.0]), + ) + nodes = tuple( + replace(manifest.nodes[0], node_id=f"env:0/rigid:new{index}") + for index in range(node_count) + ) + backend.publish_manifest(replace(manifest, scene_revision=2, nodes=nodes)) + server.scene.pointer_callbacks[0][1](event) + backend.poll() + # Neither select a new node using old transforms nor detach an active + # gizmo as if this temporary absence of pose data were an empty click. + assert commands == [] + assert not backend.publish_frame(frame) + backend.publish_frame( + replace( + frame, + scene_revision=2, + node_ids=tuple(node.node_id for node in nodes), + positions=np.zeros((node_count, 3), dtype=np.float32), + wxyz=np.tile([1.0, 0.0, 0.0, 0.0], (node_count, 1)), + visible=np.ones(node_count, dtype=np.bool_), + ) + ) + server.scene.pointer_callbacks[0][1](event) + backend.poll() + assert commands[-1].node_id == nodes[0].node_id + assert commands[-1].scene_revision == 2 + finally: + backend.stop() + + +def test_viser_backend_disabling_pick_clears_picker_gizmo() -> None: + backend, server, pick_commands = _make_pick_backend() + manifest, frame = _make_pickable_scene() + + 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. + backend.poll() + 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(