diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.demo_base.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.demo_base.rst new file mode 100644 index 000000000..ca38cadcb --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.demo_base.rst @@ -0,0 +1,12 @@ +embodichain.lab.sim.demo\_base +============================== + +.. automodule:: embodichain.lab.sim.demo_base + + + .. rubric:: Classes + + .. autosummary:: + + DemoBase + \ No newline at end of file diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.rst index 412f570d1..6927312c7 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.rst @@ -15,6 +15,7 @@ management, materials, sensors, planning/IK utilities, and action helpers. :toctree: . sim_manager + demo_base cfg common material @@ -44,6 +45,14 @@ Simulation Manager :show-inheritance: :exclude-members: __init__, copy, replace, to_dict, validate +Demo Base +--------- + +.. automodule:: embodichain.lab.sim.demo_base + :members: + :undoc-members: + :show-inheritance: + Configuration ------------- 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..f1a762a9b 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.utility.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.utility.rst @@ -17,6 +17,7 @@ action/solver adaptation. action_utils atom_action_utils cfg_utils + demo_utils gizmo_utils import_utils io_utils @@ -46,6 +47,12 @@ Configuration Utilities .. automodule:: embodichain.lab.sim.utility.cfg_utils :members: +Demo Utilities +~~~~~~~~~~~~~~ + +.. automodule:: embodichain.lab.sim.utility.demo_utils + :members: + Gizmo Utilities ~~~~~~~~~~~~~~~ diff --git a/docs/source/tutorial/index.rst b/docs/source/tutorial/index.rst index 7ccb59a2e..624d9cd22 100644 --- a/docs/source/tutorial/index.rst +++ b/docs/source/tutorial/index.rst @@ -20,17 +20,18 @@ Follow the tutorials in this order for the best learning experience: 8. :doc:`motion_gen` — Generate smooth trajectories with motion planners. 9. :doc:`atomic_actions` — Use built-in action primitives (move, move joints, pick, move held object, place). 10. :doc:`gizmo` — Interactively control robots with on-screen gizmos. +11. :doc:`writing_demo_scripts` — Write reusable demo scripts with ``DemoBase`` and shared utilities. **Phase 2: Environments** -11. :doc:`basic_env` — Create a simple Gymnasium environment with ``BaseEnv``. Prerequisite: Phase 1 basics. -12. :doc:`modular_env` — Build a config-driven environment with ``EmbodiedEnv``, managers, and randomization. Prerequisite: :doc:`basic_env`. -13. :doc:`data_generation` — Generate expert demonstration datasets for imitation learning. Prerequisite: :doc:`modular_env`. -14. :doc:`rl` — Train RL agents with PPO or GRPO. Prerequisite: :doc:`basic_env`. +12. :doc:`basic_env` — Create a simple Gymnasium environment with ``BaseEnv``. Prerequisite: Phase 1 basics. +13. :doc:`modular_env` — Build a config-driven environment with ``EmbodiedEnv``, managers, and randomization. Prerequisite: :doc:`basic_env`. +14. :doc:`data_generation` — Generate expert demonstration datasets for imitation learning. Prerequisite: :doc:`modular_env`. +15. :doc:`rl` — Train RL agents with PPO or GRPO. Prerequisite: :doc:`basic_env`. **Phase 3: Extending the Framework** -15. :doc:`add_robot` — Add a new robot model to EmbodiChain. +16. :doc:`add_robot` — Add a new robot model to EmbodiChain. .. toctree:: :maxdepth: 1 @@ -48,6 +49,7 @@ Follow the tutorials in this order for the best learning experience: motion_gen atomic_actions gizmo + writing_demo_scripts basic_env modular_env data_generation diff --git a/docs/source/tutorial/writing_demo_scripts.rst b/docs/source/tutorial/writing_demo_scripts.rst new file mode 100644 index 000000000..b7351fe56 --- /dev/null +++ b/docs/source/tutorial/writing_demo_scripts.rst @@ -0,0 +1,185 @@ +.. _tutorial_writing_demo_scripts: + +Writing Demo Scripts +==================== + +Simulation demo scripts (under ``scripts/tutorials/`` and ``examples/sim/demo/``) share a +lot of boilerplate: argument parsing, simulation setup, window/recording management, +trajectory replay, and cleanup. EmbodiChain ships a small **demo utility layer** that +factors this boilerplate out so new demos can focus on their core *setup* and *run* logic. + +The layer has two parts: + +- :mod:`embodichain.lab.sim.utility.demo_utils` - standalone helper functions and the + ``DemoRecording`` context manager. Usable from any script, flat or structured. +- :mod:`embodichain.lab.sim.demo_base` - an optional ``DemoBase`` lifecycle base class + for demos with a clear setup / run / cleanup structure. + +.. note:: + These helpers build on top of + :func:`~embodichain.lab.gym.utils.gym_utils.add_env_launcher_args_to_parser`; they + add demo-specific flags (``--auto_play``, ``--record_steps``, ``--record_fps``, + ``--record_save_path``, ``--no_vis_eef_axis``) rather than replacing it. Common + multi-word options accept both underscore and hyphen spellings, for example + ``--record_steps`` and ``--record-steps``. + +Command-line Arguments +---------------------- + +:func:`~embodichain.lab.sim.utility.demo_utils.add_demo_args` extends a parser with the +standard launcher arguments plus the demo flags above: + +.. code-block:: python + + import argparse + + from embodichain.lab.sim.utility.demo_utils import add_demo_args + + parser = argparse.ArgumentParser(description="My demo.") + parser = add_demo_args(parser) + args = parser.parse_args() + +``--auto_play`` skips every interactive ``input()`` prompt (see +:ref:`the helpers `) so the demo can run end-to-end in CI or +headless. For continuous demos it also selects a finite default run length. +``--record_steps`` enables video recording and gives continuous demos an explicit +simulation-update limit (see :ref:`recording `). + +Creating and Tearing Down the Simulation +----------------------------------------- + +:func:`~embodichain.lab.sim.utility.demo_utils.create_default_sim` builds a +:class:`~embodichain.lab.sim.SimulationManager` from the parsed ``args`` +(``headless``, ``device``, ``renderer``, ``gpu_id``) with sensible defaults and +optionally a main light. Pass ``num_envs`` for parallel-environment demos, and +``add_default_light=False`` if you want to set up your own lighting. + +:func:`~embodichain.lab.sim.utility.demo_utils.shutdown_sim` finishes any active +recording and calls ``sim.destroy()``. +**Always** destroy the simulation before the process exits, otherwise the interpreter +segfaults (exit 139) during teardown. + +.. _demo-window-helpers: + +Window, User and GPU Helpers +---------------------------- + +These read the parsed ``args`` and become no-ops when the relevant flag is set, so the +same code path works for interactive and ``--auto_play`` runs: + +- :func:`~embodichain.lab.sim.utility.demo_utils.maybe_open_window` - opens the viewer + unless ``--headless``. +- :func:`~embodichain.lab.sim.utility.demo_utils.maybe_init_gpu_physics` - calls + ``sim.init_gpu_physics()`` only when the sim is configured for GPU physics. +- :func:`~embodichain.lab.sim.utility.demo_utils.maybe_wait_for_user` - blocks on + ``input(prompt)`` unless ``--auto_play``. +- :func:`~embodichain.lab.sim.utility.demo_utils.maybe_pause_for_inspection` - the same, + with an end-of-demo prompt. +- :func:`~embodichain.lab.sim.utility.demo_utils.resolve_demo_steps` - resolves an + explicit recording limit or a finite ``--auto_play`` limit while leaving interactive + runs unbounded. +- :func:`~embodichain.lab.sim.utility.demo_utils.run_simulation_loop` - runs the common + update/FPS loop with optional sleep, step callback, and maximum update count. + +.. _demo-recording: + +Recording +--------- + +:class:`~embodichain.lab.sim.utility.demo_utils.DemoRecording` is a context manager that +starts window recording when ``args.record_steps`` is set, generates a timestamped file +name under ``--record_save_path`` (default ``./recordings``), and stops + flushes the +video on exit. A path ending in ``.mp4`` is used as the exact output path; other paths are +treated as output directories. Headless runs use a default fixed camera when the caller +does not provide ``look_at``. If recording fails to start it warns and continues instead +of aborting the demo. + +.. code-block:: python + + with DemoRecording(sim, args, prefix="my_demo"): + # ... run the demo; frames are captured here ... + replay_trajectory(sim, robot, traj) + +Replaying a Trajectory +---------------------- + +:func:`~embodichain.lab.sim.utility.demo_utils.replay_trajectory` steps a joint-space +trajectory through the simulator, applying each waypoint with ``robot.set_qpos`` and +``sim.update(step=...)``, then holds the final configuration for ``post_steps``. It +accepts 1-D ``(num_joints,)``, 2-D ``(num_steps, num_joints)`` or 3-D +``(batch, num_steps, num_joints)`` tensors. + +The ``DemoBase`` Lifecycle Class +-------------------------------- + +For demos with a clear structure, subclass +:class:`~embodichain.lab.sim.demo_base.DemoBase` and implement :meth:`setup` and +:meth:`run`. The base class stores ``args``, guarantees :meth:`cleanup` (``sim.destroy()``) +runs in a ``finally`` block, and exposes :meth:`main` to drive the lifecycle: + +.. code-block:: python + + import argparse + + import torch + + from embodichain.lab.sim.demo_base import DemoBase + from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_pause_for_inspection, + maybe_wait_for_user, + replay_trajectory, + setup_print_options, + ) + + + class MyDemo(DemoBase): + def setup(self) -> None: + self.sim = create_default_sim(self.args) + maybe_open_window(self.sim, self.args) + self.robot = self.sim.add_robot(cfg=...) + # ... build the rest of the scene ... + + def run(self) -> None: + maybe_wait_for_user(self.args, "Press Enter to plan...") + traj = ... # plan the trajectory + with DemoRecording(self.sim, self.args, prefix="my_demo"): + replay_trajectory(self.sim, self.robot, traj) + maybe_pause_for_inspection(self.args) + + + def main() -> None: + setup_print_options() + parser = add_demo_args(argparse.ArgumentParser(description="My demo.")) + args = parser.parse_args() + MyDemo(args).main() + + + if __name__ == "__main__": + main() + +Because :meth:`~embodichain.lab.sim.demo_base.DemoBase.main` wraps both :meth:`setup` +and :meth:`run` in ``try / finally``, the simulation is always destroyed - even if setup +fails after allocating the simulation, :meth:`run` raises, or the user interrupts with +``Ctrl+C``. + +.. tip:: + Demos that keep the viewer open until ``Ctrl+C`` should use + ``run_simulation_loop(..., max_steps=resolve_demo_steps(args))`` so interactive runs + stay open while ``--auto_play`` and ``--record_steps`` runs terminate automatically. + +Reference Implementations +------------------------- + +The migrated demo scripts are good references: + +- ``scripts/tutorials/atomic_action/`` - the atomic-action tutorials use the shared + argument, recording, and cleanup conventions; the single-arm demos use ``DemoBase``. +- ``examples/sim/demo/`` - ``press_softbody``, ``pick_up_cloth``, ``grasp_cup_to_caffe`` + and ``scoop_ice`` use ``DemoBase`` for lifecycle and the shared helpers for args, + recording and cleanup. +- ``scripts/tutorials/sim/`` and the focused examples under ``examples/sim/`` show the + flat-function style for short scene, sensor, solver, gizmo, and planner demos. diff --git a/docs/superpowers/plans/2026-07-04-demo-base-class.md b/docs/superpowers/plans/2026-07-04-demo-base-class.md new file mode 100644 index 000000000..7d0fd9bed --- /dev/null +++ b/docs/superpowers/plans/2026-07-04-demo-base-class.md @@ -0,0 +1,1284 @@ +# Demo Base Class / Shared Utilities Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Create a reusable demo utility layer (`demo_utils.py`) and an optional lifecycle base class (`DemoBase`) for simulation demo scripts, then migrate the `scripts/tutorials/atomic_action/` tutorials to use them. + +**Architecture:** Add generic helpers to `embodichain/lab/sim/utility/demo_utils.py` and an abstract `DemoBase` to `embodichain/lab/sim/demo_base.py`. Keep robot-specific UR5+gripper configuration in `scripts/tutorials/atomic_action/tutorial_utils.py` but re-export generic helpers from `demo_utils.py`. Refactor the six `atomic_action` tutorial scripts incrementally, preserving behavior. + +**Tech Stack:** Python 3.10+, PyTorch, `embodichain.lab.sim`, pytest, black. + +## Global Constraints + +- Branch: `feature/demo-base-class` +- Formatter: `black==26.3.1` — run before every commit. +- Every source file begins with the Apache 2.0 copyright header. +- Use `from __future__ import annotations` at the top of every file. +- Use full type hints on all public APIs. +- Define `__all__` in every public module. +- Use Google-style docstrings with Sphinx directives. +- Do not break existing imports from `scripts.tutorials.atomic_action.tutorial_utils`. +- Commit after every independently testable deliverable. + +--- + +## File Structure + +| File | Responsibility | +|---|---| +| `embodichain/lab/sim/utility/demo_utils.py` | Generic demo helpers: argument parsing, sim creation, cleanup, recording context manager, trajectory replay, print options, tensor formatting. | +| `embodichain/lab/sim/utility/__init__.py` | Re-export `demo_utils` public API. | +| `embodichain/lab/sim/demo_base.py` | Optional `DemoBase` abstract lifecycle class. | +| `scripts/tutorials/atomic_action/tutorial_utils.py` | Keep UR5-specific robot/solver helpers; re-export generic helpers from `demo_utils.py`. | +| `scripts/tutorials/atomic_action/move_joints.py` | Refactored to use shared utilities / `DemoBase`. | +| `scripts/tutorials/atomic_action/move_end_effector.py` | Refactored to use shared utilities / `DemoBase`. | +| `scripts/tutorials/atomic_action/pickup.py` | Refactored to use shared utilities / `DemoBase`. | +| `scripts/tutorials/atomic_action/place.py` | Refactored to use shared utilities / `DemoBase`. | +| `scripts/tutorials/atomic_action/move_held_object.py` | Refactored to use shared utilities / `DemoBase`. | +| `scripts/tutorials/atomic_action/press.py` | Refactored to use shared utilities / `DemoBase` (keeps custom robot/table config). | +| `tests/sim/utility/test_demo_utils.py` | Unit tests for `demo_utils.py` helpers. | + +--- + +## Task 1: Core demo utilities — argument parsing, sim setup, cleanup + +**Files:** +- Create: `embodichain/lab/sim/utility/demo_utils.py` +- Modify: `embodichain/lab/sim/utility/__init__.py` +- Test: `tests/sim/utility/test_demo_utils.py` + +**Interfaces:** +- Consumes: `embodichain.lab.gym.utils.gym_utils.add_env_launcher_args_to_parser`, `embodichain.lab.sim.SimulationManagerCfg`, `embodichain.lab.sim.cfg.LightCfg`, `embodichain.lab.sim.cfg.RenderCfg` +- Produces: `add_demo_args(parser)`, `create_default_sim(args, ...)`, `shutdown_sim(sim)`, `setup_print_options()`, `format_tensor(tensor)`, `maybe_init_gpu_physics(sim)` + +- [ ] **Step 1: Write the failing test** + +Create `tests/sim/utility/test_demo_utils.py`: + +```python +# ---------------------------------------------------------------------------- +# 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 argparse +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from unittest.mock import Mock + +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + format_tensor, + setup_print_options, + shutdown_sim, +) + + +def test_add_demo_args_adds_expected_flags(): + parser = argparse.ArgumentParser() + parser = add_demo_args(parser) + args = parser.parse_args(["--headless", "--auto_play", "--record_fps", "60"]) + assert args.headless is True + assert args.auto_play is True + assert args.record_fps == 60 + assert args.record_steps is None + assert args.no_vis_eef_axis is False + + +def test_format_tensor_rounds_and_moves_to_cpu(): + tensor = torch.tensor([1.23456789, 2.34567891]) + result = format_tensor(tensor) + assert result == "[1.2346, 2.3457]" + + +def test_setup_print_options_sets_numpy_and_torch(): + setup_print_options() + assert np.get_printoptions()["precision"] == 5 + assert np.get_printoptions()["suppress"] is True + assert torch.get_printoptions()["precision"] == 5 + assert torch.get_printoptions()["sci_mode"] is False + + +def test_shutdown_sim_calls_destroy(): + sim = Mock(spec=["destroy"]) + shutdown_sim(sim) + sim.destroy.assert_called_once() +``` + +- [ ] **Step 2: Run test to verify it fails** + +```bash +pytest tests/sim/utility/test_demo_utils.py -v +``` + +Expected: FAIL with `ModuleNotFoundError: No module named 'embodichain.lab.sim.utility.demo_utils'`. + +- [ ] **Step 3: Write minimal implementation** + +Create `embodichain/lab/sim/utility/demo_utils.py`: + +```python +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Shared utilities for simulation demo scripts.""" + +from __future__ import annotations + +import argparse +from typing import TYPE_CHECKING + +import numpy as np +import torch + +from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import LightCfg, RenderCfg + +if TYPE_CHECKING: + from collections.abc import Sequence + + +__all__ = [ + "add_demo_args", + "create_default_sim", + "shutdown_sim", + "setup_print_options", + "format_tensor", + "maybe_init_gpu_physics", +] + + +def add_demo_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: + """Add common demo arguments to an environment launcher parser. + + Args: + parser: The parser to extend. + + Returns: + The same parser with demo flags added. + """ + add_env_launcher_args_to_parser(parser) + parser.add_argument( + "--auto_play", + action="store_true", + help="Skip interactive prompts and run the demo automatically.", + ) + parser.add_argument( + "--record_steps", + type=int, + default=None, + help="Number of simulation steps to record. If None, no recording is started.", + ) + parser.add_argument( + "--record_fps", + type=int, + default=30, + help="Frames per second for the recorded video.", + ) + parser.add_argument( + "--record_save_path", + type=str, + default=None, + help="Directory to save recorded videos. Defaults to ./recordings.", + ) + parser.add_argument( + "--no_vis_eef_axis", + action="store_true", + help="Disable end-effector axis visualization.", + ) + return parser + + +def create_default_sim( + args: argparse.Namespace, + *, + width: int = 1920, + height: int = 1080, + physics_dt: float = 1.0 / 100.0, + arena_space: float = 2.5, + add_default_light: bool = True, +) -> SimulationManager: + """Create a SimulationManager with common demo defaults. + + Args: + args: Parsed command-line arguments. Expected to contain ``headless``, + ``device``, ``renderer`` and ``arena_space``. + width: Window/render width. + height: Window/render height. + physics_dt: Physics simulation timestep. + arena_space: Arena space size. + add_default_light: Whether to add a default point light. + + Returns: + Configured simulation manager instance. + """ + cfg = SimulationManagerCfg( + width=width, + height=height, + headless=args.headless, + sim_device=args.device, + render_cfg=RenderCfg(renderer=args.renderer), + physics_dt=physics_dt, + arena_space=arena_space, + ) + sim = SimulationManager(cfg) + if add_default_light: + sim.add_light( + cfg=LightCfg( + uid="main_light", + color=(0.6, 0.6, 0.6), + intensity=30.0, + init_pos=(1.0, 0.0, 3.0), + ) + ) + return sim + + +def shutdown_sim(sim: SimulationManager) -> None: + """Safely destroy a simulation manager. + + Args: + sim: The simulation manager to destroy. + """ + sim.destroy() + + +def setup_print_options() -> None: + """Set common numpy and torch print options for demos.""" + np.set_printoptions(precision=5, suppress=True) + torch.set_printoptions(precision=5, sci_mode=False) + + +def format_tensor(tensor: torch.Tensor) -> str: + """Return a compact, rounded string representation of a tensor. + + Args: + tensor: Input tensor. + + Returns: + Rounded string with 4 decimal places. + """ + rounded = (tensor.detach().cpu() * 10000.0).round() / 10000.0 + return str(rounded.tolist()) + + +def maybe_init_gpu_physics(sim: SimulationManager) -> None: + """Initialize GPU physics if the simulation is configured to use it. + + Args: + sim: The simulation manager. + """ + if sim.is_use_gpu_physics: + sim.init_gpu_physics() +``` + +Modify `embodichain/lab/sim/utility/__init__.py` to add: + +```python +from .demo_utils import * +``` + +Keep the existing imports. + +- [ ] **Step 4: Run test to verify it passes** + +```bash +pytest tests/sim/utility/test_demo_utils.py -v +``` + +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +black . +git add embodichain/lab/sim/utility/demo_utils.py embodichain/lab/sim/utility/__init__.py tests/sim/utility/test_demo_utils.py +git commit -m "feat(demo): add core demo utilities + +Co-Authored-By: Claude " +``` + +--- + +## Task 2: Recording context manager + +**Files:** +- Modify: `embodichain/lab/sim/utility/demo_utils.py` +- Test: `tests/sim/utility/test_demo_utils.py` + +**Interfaces:** +- Consumes: `SimulationManager.start_window_record`, `stop_window_record`, `wait_window_record_saves`, `is_window_recording`, `sim_config` +- Produces: `DemoRecording(sim, args, prefix)` context manager + +- [ ] **Step 1: Write the failing test** + +Append to `tests/sim/utility/test_demo_utils.py`: + +```python +import datetime as _datetime +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest + +from embodichain.lab.sim.utility.demo_utils import DemoRecording + + +def _make_sim_manager(): + sim = Mock(spec=["start_window_record", "stop_window_record", "wait_window_record_saves", "is_window_recording", "sim_config"]) + sim.sim_config = SimpleNamespace(width=1920, height=1080) + sim.start_window_record.return_value = True + sim.is_window_recording.return_value = False + return sim + + +def test_demo_recording_does_nothing_when_record_steps_is_none(): + sim = _make_sim_manager() + args = SimpleNamespace(record_steps=None, record_fps=30, record_save_path="/tmp", auto_play=False, headless=True) + with DemoRecording(sim, args, prefix="demo"): + pass + sim.start_window_record.assert_not_called() + + +def test_demo_recording_starts_and_stops_window_record(): + sim = _make_sim_manager() + args = SimpleNamespace(record_steps=10, record_fps=30, record_save_path="/tmp/recordings", auto_play=False, headless=True) + with DemoRecording(sim, args, prefix="demo") as rec: + assert rec.is_active is True + sim.start_window_record.assert_called_once() + call_kwargs = sim.start_window_record.call_args.kwargs + assert call_kwargs["fps"] == 30 + assert call_kwargs["video_prefix"] == "demo" + assert "/tmp/recordings" in call_kwargs["save_path"] + assert call_kwargs["save_path"].endswith(".mp4") + sim.stop_window_record.assert_called_once() + sim.wait_window_record_saves.assert_called_once() + + +def test_demo_recording_warns_and_skips_on_start_failure(): + sim = _make_sim_manager() + sim.start_window_record.return_value = False + args = SimpleNamespace(record_steps=10, record_fps=30, record_save_path="/tmp", auto_play=False, headless=True) + with pytest.warns(UserWarning, match="Failed to start recording"): + with DemoRecording(sim, args, prefix="demo"): + pass + sim.stop_window_record.assert_not_called() +``` + +- [ ] **Step 2: Run test to verify it fails** + +```bash +pytest tests/sim/utility/test_demo_utils.py::test_demo_recording_does_nothing_when_record_steps_is_none -v +``` + +Expected: FAIL with `ImportError: cannot import name 'DemoRecording'`. + +- [ ] **Step 3: Write minimal implementation** + +Append to `embodichain/lab/sim/utility/demo_utils.py`: + +```python +import datetime +import warnings +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Sequence + + +__all__.append("DemoRecording") + + +class DemoRecording: + """Context manager that handles demo video recording. + + Recording is only started when ``args.record_steps`` is not ``None``. + On exit the window record is stopped and the framework is asked to finish + saving the video file. + + Args: + sim: The simulation manager. + args: Parsed command-line arguments. Expected to contain + ``record_steps``, ``record_fps``, ``record_save_path`` and + ``headless``. + prefix: Prefix used for the generated video filename. + """ + + def __init__( + self, + sim: SimulationManager, + args: argparse.Namespace, + prefix: str = "demo", + look_at: tuple[Sequence[float], Sequence[float], Sequence[float]] | None = None, + ): + self.sim = sim + self.args = args + self.prefix = prefix + self.look_at = look_at + self.is_active = False + + def __enter__(self) -> DemoRecording: + """Start recording if requested.""" + if self.args.record_steps is None: + return self + + save_dir = Path(self.args.record_save_path or "./recordings") + save_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + save_path = str(save_dir / f"{self.prefix}_{timestamp}.mp4") + + original_width = self.sim.sim_config.width + original_height = self.sim.sim_config.height + try: + # Use a smaller resolution for recording to keep files small. + self.sim.sim_config.width = 640 + self.sim.sim_config.height = 480 + started = self.sim.start_window_record( + save_path=save_path, + fps=self.args.record_fps, + max_memory=2048, + video_prefix=self.prefix, + look_at=self.look_at, + use_sim_time=True, + ) + finally: + self.sim.sim_config.width = original_width + self.sim.sim_config.height = original_height + + if not started: + warnings.warn( + f"Failed to start recording for prefix '{self.prefix}'. Continuing without recording.", + UserWarning, + stacklevel=2, + ) + return self + + self.is_active = True + return self + + def __exit__(self, exc_type, exc, tb) -> None: + """Stop recording and wait for the file to be written.""" + if not self.is_active: + return + if self.sim.is_window_recording(): + self.sim.stop_window_record() + self.sim.wait_window_record_saves() +``` + +- [ ] **Step 4: Run test to verify it passes** + +```bash +pytest tests/sim/utility/test_demo_utils.py -v +``` + +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +black . +git add embodichain/lab/sim/utility/demo_utils.py tests/sim/utility/test_demo_utils.py +git commit -m "feat(demo): add DemoRecording context manager + +Co-Authored-By: Claude " +``` + +--- + +## Task 3: Trajectory replay and window/user helpers + +**Files:** +- Modify: `embodichain/lab/sim/utility/demo_utils.py` +- Test: `tests/sim/utility/test_demo_utils.py` + +**Interfaces:** +- Consumes: `SimulationManager.update`, `Robot.set_qpos`, `Robot.get_joint_ids` +- Produces: `replay_trajectory(...)`, `maybe_open_window(...)`, `maybe_wait_for_user(...)`, `maybe_pause_for_inspection(...)` + +- [ ] **Step 1: Write the failing test** + +Append to `tests/sim/utility/test_demo_utils.py`: + +```python +import time +from unittest.mock import Mock, call + +from embodichain.lab.sim.utility.demo_utils import ( + maybe_open_window, + maybe_wait_for_user, + replay_trajectory, +) + + +def test_maybe_open_window_opens_when_not_headless(): + sim = Mock(spec=["open_window"]) + args = SimpleNamespace(headless=False) + maybe_open_window(sim, args) + sim.open_window.assert_called_once() + + +def test_maybe_open_window_does_nothing_when_headless(): + sim = Mock(spec=["open_window"]) + args = SimpleNamespace(headless=True) + maybe_open_window(sim, args) + sim.open_window.assert_not_called() + + +def test_maybe_wait_for_user_prompts_when_not_auto_play(): + args = SimpleNamespace(auto_play=False) + with patch("builtins.input", return_value="") as mock_input: + maybe_wait_for_user(args, "Press enter") + mock_input.assert_called_once_with("Press enter") + + +def test_maybe_wait_for_user_skips_when_auto_play(): + args = SimpleNamespace(auto_play=True) + with patch("builtins.input") as mock_input: + maybe_wait_for_user(args, "Press enter") + mock_input.assert_not_called() + + +def test_replay_trajectory_sets_qpos_and_updates_sim(): + robot = Mock(spec=["set_qpos"]) + sim = Mock(spec=["update"]) + # Shape: (batch=1, num_steps=2, num_joints=3) + traj = torch.tensor([ + [[0.0, 0.1, 0.2], [0.3, 0.4, 0.5]], + ]) + replay_trajectory(sim, robot, traj, post_steps=1, step_size=4, sleep=0.0) + assert robot.set_qpos.call_count == 3 # 2 traj + 1 post + assert sim.update.call_count == 3 + sim.update.assert_has_calls([call(step=4), call(step=4), call(step=2)]) +``` + +- [ ] **Step 2: Run test to verify it fails** + +```bash +pytest tests/sim/utility/test_demo_utils.py::test_replay_trajectory_sets_qpos_and_updates_sim -v +``` + +Expected: FAIL with `ImportError`. + +- [ ] **Step 3: Write minimal implementation** + +Append to `embodichain/lab/sim/utility/demo_utils.py`: + +```python +import time + +from embodichain.lab.sim.objects import Robot + + +__all__.extend([ + "maybe_open_window", + "maybe_wait_for_user", + "maybe_pause_for_inspection", + "replay_trajectory", +]) + + +def maybe_open_window(sim: SimulationManager, args: argparse.Namespace) -> None: + """Open the viewer window unless running headless. + + Args: + sim: The simulation manager. + args: Parsed arguments containing ``headless``. + """ + if not args.headless: + sim.open_window() + + +def maybe_wait_for_user(args: argparse.Namespace, prompt: str) -> None: + """Wait for user input unless auto_play is enabled. + + Args: + args: Parsed arguments containing ``auto_play``. + prompt: Message to display when waiting. + """ + if not args.auto_play: + input(prompt) + + +def maybe_pause_for_inspection(args: argparse.Namespace) -> None: + """Pause at the end of a demo for visual inspection. + + Args: + args: Parsed arguments containing ``auto_play``. + """ + maybe_wait_for_user(args, "Demo finished. Press Enter to exit...") + + +def replay_trajectory( + sim: SimulationManager, + robot: Robot, + traj: torch.Tensor, + *, + post_steps: int = 60, + step_size: int = 4, + sleep: float = 1e-2, + arm_name: str | None = None, +) -> None: + """Replay a joint-space trajectory on a robot. + + ``traj`` may be either a 2-D tensor of shape ``(num_joints,)`` or a 3-D + tensor of shape ``(batch, num_steps, num_joints)``. For 2-D input the + single configuration is held for ``post_steps``. For 3-D input each step + is applied sequentially and the final configuration is held. + + Args: + sim: The simulation manager. + robot: The robot instance. + traj: Joint position trajectory tensor. + post_steps: Number of steps to hold the final configuration. + step_size: Number of physics steps per ``sim.update()`` call. + sleep: Sleep duration between steps (seconds). + arm_name: Optional arm name passed to ``robot.set_qpos``. + """ + if traj.dim() == 1: + traj = traj.unsqueeze(0).unsqueeze(0) + elif traj.dim() == 2: + traj = traj.unsqueeze(0) + + joint_ids = robot.get_joint_ids(arm_name) if arm_name is not None else None + + for i in range(traj.shape[1]): + robot.set_qpos(qpos=traj[:, i, :], joint_ids=joint_ids) + sim.update(step=step_size) + time.sleep(sleep) + + final_qpos = traj[:, -1, :] + for _ in range(post_steps): + robot.set_qpos(qpos=final_qpos, joint_ids=joint_ids) + sim.update(step=2) + time.sleep(sleep) +``` + +- [ ] **Step 4: Run test to verify it passes** + +```bash +pytest tests/sim/utility/test_demo_utils.py -v +``` + +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +black . +git add embodichain/lab/sim/utility/demo_utils.py tests/sim/utility/test_demo_utils.py +git commit -m "feat(demo): add trajectory replay and window helpers + +Co-Authored-By: Claude " +``` + +--- + +## Task 4: Optional DemoBase lifecycle class + +**Files:** +- Create: `embodichain/lab/sim/demo_base.py` +- Test: `tests/sim/test_demo_base.py` + +**Interfaces:** +- Consumes: `embodichain.lab.sim.utility.demo_utils.shutdown_sim` +- Produces: `DemoBase` abstract class with `setup()`, `run()`, `cleanup()`, `main()` + +- [ ] **Step 1: Write the failing test** + +Create `tests/sim/test_demo_base.py`: + +```python +# ---------------------------------------------------------------------------- +# 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 argparse +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from embodichain.lab.sim.demo_base import DemoBase + + +def test_demo_base_runs_setup_run_cleanup(): + class SimpleDemo(DemoBase): + def setup(self): + self.sim = Mock(spec=["destroy"]) + + def run(self): + self.ran = True + + demo = SimpleDemo(SimpleNamespace()) + demo.main() + assert demo.ran is True + demo.sim.destroy.assert_called_once() + + +def test_demo_base_cleanup_runs_even_if_run_raises(): + class BrokenDemo(DemoBase): + def setup(self): + self.sim = Mock(spec=["destroy"]) + + def run(self): + raise RuntimeError("boom") + + demo = BrokenDemo(SimpleNamespace()) + with pytest.raises(RuntimeError, match="boom"): + demo.main() + demo.sim.destroy.assert_called_once() +``` + +- [ ] **Step 2: Run test to verify it fails** + +```bash +pytest tests/sim/test_demo_base.py -v +``` + +Expected: FAIL with `ModuleNotFoundError`. + +- [ ] **Step 3: Write minimal implementation** + +Create `embodichain/lab/sim/demo_base.py`: + +```python +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Optional base class for simulation demos.""" + +from __future__ import annotations + +import argparse +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +from embodichain.lab.sim.utility.demo_utils import shutdown_sim + +if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager + + +__all__ = ["DemoBase"] + + +class DemoBase(ABC): + """Lightweight lifecycle base class for simulation demos. + + Subclasses implement :meth:`setup` and :meth:`run`; the base class handles + argument injection and guaranteed cleanup. + + Args: + args: Parsed command-line arguments. + """ + + def __init__(self, args: argparse.Namespace): + self.args = args + self.sim: SimulationManager | None = None + + @abstractmethod + def setup(self) -> None: + """Create simulation, robot, cameras, etc.""" + + @abstractmethod + def run(self) -> None: + """Execute the demo logic.""" + + def cleanup(self) -> None: + """Release simulation resources. Called automatically by :meth:`main`.""" + if self.sim is not None: + shutdown_sim(self.sim) + + def main(self) -> None: + """Run the full demo lifecycle.""" + self.setup() + try: + self.run() + finally: + self.cleanup() +``` + +- [ ] **Step 4: Run test to verify it passes** + +```bash +pytest tests/sim/test_demo_base.py -v +``` + +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +black . +git add embodichain/lab/sim/demo_base.py tests/sim/test_demo_base.py +git commit -m "feat(demo): add DemoBase lifecycle class + +Co-Authored-By: Claude " +``` + +--- + +## Task 5: Update tutorial_utils.py to re-export generic helpers + +**Files:** +- Modify: `scripts/tutorials/atomic_action/tutorial_utils.py` + +**Interfaces:** +- Consumes: `embodichain.lab.sim.utility.demo_utils` helpers +- Produces: Same public API as before for backward compatibility + +- [ ] **Step 1: Update imports and __all__** + +Modify `scripts/tutorials/atomic_action/tutorial_utils.py`: + +1. Keep the existing UR5-specific helpers and `start_auto_play_recording` / `stop_auto_play_recording` implementations unchanged (these are tutorial-specific wrappers with fixed resolution and look-at defaults). + +2. Add imports from `embodichain.lab.sim.utility.demo_utils` for the generic helpers used by refactored scripts: + +```python +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + format_tensor, + maybe_init_gpu_physics, + maybe_open_window, + maybe_pause_for_inspection, + maybe_wait_for_user, + replay_trajectory, + setup_print_options, + shutdown_sim, +) +``` + +3. Update `__all__` to include the re-exported names while keeping existing entries. + +- [ ] **Step 2: Verify backward compatibility** + +Run a quick import smoke test: + +```bash +python -c "from scripts.tutorials.atomic_action.tutorial_utils import create_ur5_gripper_robot_cfg, start_auto_play_recording, draw_axis_marker, add_demo_args; print('ok')" +``` + +Expected: prints `ok`. + +Modify `DemoRecording` in `embodichain/lab/sim/utility/demo_utils.py`: + +```python +def __init__( + self, + sim: SimulationManager, + args: argparse.Namespace, + prefix: str = "demo", + look_at: tuple[Sequence[float], Sequence[float], Sequence[float]] | None = None, +): + ... + self.look_at = look_at +``` + +And in `__enter__`, pass `look_at=self.look_at` to `start_window_record`. + +Then update the compatibility wrappers. + +- [ ] **Step 3: Run tests** + +```bash +pytest tests/sim/utility/test_demo_utils.py tests/sim/test_demo_base.py -v +``` + +Expected: PASS. + +- [ ] **Step 4: Commit** + +```bash +black . +git add scripts/tutorials/atomic_action/tutorial_utils.py embodichain/lab/sim/utility/demo_utils.py +git commit -m "refactor(tutorial): re-export generic helpers from demo_utils + +Co-Authored-By: Claude " +``` + +--- + +## Task 6: Refactor move_joints.py to use DemoBase + +**Files:** +- Modify: `scripts/tutorials/atomic_action/move_joints.py` + +**Interfaces:** +- Consumes: `DemoBase`, `add_demo_args`, `create_default_sim`, `maybe_open_window`, `maybe_wait_for_user`, `DemoRecording`, `replay_trajectory`, `shutdown_sim` + +- [ ] **Step 1: Read current file** + +Read `scripts/tutorials/atomic_action/move_joints.py` in full. + +- [ ] **Step 2: Rewrite using DemoBase** + +Keep behavior identical. Replace the `_REPO_ROOT` sys.path hack with normal package imports. Use `DemoBase` for lifecycle. + +Example target structure: + +```python +from __future__ import annotations + +import argparse + +import torch + +from embodichain.lab.sim import SimulationManager +from embodichain.lab.sim.demo_base import DemoBase +from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator, ToppraPlannerCfg +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + replay_trajectory, + setup_print_options, +) +from scripts.tutorials.atomic_action.tutorial_utils import create_ur5_gripper_robot_cfg + + +class MoveJointsDemo(DemoBase): + def setup(self) -> None: + self.sim = create_default_sim(self.args, width=1600, height=900) + maybe_open_window(self.sim, self.args) + robot = self.sim.add_robot(cfg=create_ur5_gripper_robot_cfg()) + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + self.robot = robot + self.motion_gen = motion_gen + self.engine = ... # AtomicActionEngine registration unchanged + + def run(self) -> None: + maybe_wait_for_user(self.args, "Press Enter to plan...") + traj = ... # existing plan call + with DemoRecording(self.sim, self.args, prefix="move_joints"): + replay_trajectory(self.sim, self.robot, traj) + + +def main() -> None: + setup_print_options() + parser = argparse.ArgumentParser() + parser = add_demo_args(parser) + parser.add_argument("--target_qpos", nargs="+", type=float, default=[...]) + args = parser.parse_args() + MoveJointsDemo(args).main() + + +if __name__ == "__main__": + main() +``` + +The exact body should match the original script's behavior. Keep the original `POST_TRAJECTORY_STEPS`, action registration, and target defaults. + +- [ ] **Step 3: Run integration test** + +```bash +python scripts/tutorials/atomic_action/move_joints.py --headless --auto_play +``` + +Expected: script runs to completion without error. + +- [ ] **Step 4: Commit** + +```bash +black . +git add scripts/tutorials/atomic_action/move_joints.py +git commit -m "refactor(tutorial): move_joints uses DemoBase + +Co-Authored-By: Claude " +``` + +--- + +## Task 7: Refactor move_end_effector.py to use DemoBase + +**Files:** +- Modify: `scripts/tutorials/atomic_action/move_end_effector.py` + +**Interfaces:** Same as Task 6. + +- [ ] **Step 1: Read current file** + +- [ ] **Step 2: Rewrite using DemoBase** + +Follow the same pattern as Task 6. This script has a multi-waypoint EEF trajectory; keep the waypoint logic unchanged. + +- [ ] **Step 3: Run integration test** + +```bash +python scripts/tutorials/atomic_action/move_end_effector.py --headless --auto_play +``` + +Expected: runs to completion. + +- [ ] **Step 4: Commit** + +```bash +black . +git add scripts/tutorials/atomic_action/move_end_effector.py +git commit -m "refactor(tutorial): move_end_effector uses DemoBase + +Co-Authored-By: Claude " +``` + +--- + +## Task 8: Refactor pickup.py to use DemoBase + +**Files:** +- Modify: `scripts/tutorials/atomic_action/pickup.py` + +**Interfaces:** Same as Task 6, plus object preset dictionary and grasp generator. + +- [ ] **Step 1: Read current file** + +- [ ] **Step 2: Rewrite using DemoBase** + +Follow the same pattern. Keep object creation, grasp generator, and `clear_dynamics` check inside the replay loop. + +- [ ] **Step 3: Run integration test** + +```bash +python scripts/tutorials/atomic_action/pickup.py --headless --auto_play +``` + +Expected: runs to completion. + +- [ ] **Step 4: Commit** + +```bash +black . +git add scripts/tutorials/atomic_action/pickup.py +git commit -m "refactor(tutorial): pickup uses DemoBase + +Co-Authored-By: Claude " +``` + +--- + +## Task 9: Refactor place.py to use DemoBase + +**Files:** +- Modify: `scripts/tutorials/atomic_action/place.py` + +**Interfaces:** Same as Task 8. + +- [ ] **Step 1: Read current file** + +- [ ] **Step 2: Rewrite using DemoBase** + +Keep PickUp + Place action registration and multi-waypoint place target. + +- [ ] **Step 3: Run integration test** + +```bash +python scripts/tutorials/atomic_action/place.py --headless --auto_play +``` + +Expected: runs to completion. + +- [ ] **Step 4: Commit** + +```bash +black . +git add scripts/tutorials/atomic_action/place.py +git commit -m "refactor(tutorial): place uses DemoBase + +Co-Authored-By: Claude " +``` + +--- + +## Task 10: Refactor move_held_object.py to use DemoBase + +**Files:** +- Modify: `scripts/tutorials/atomic_action/move_held_object.py` + +**Interfaces:** Same as Task 8. + +- [ ] **Step 1: Read current file** + +- [ ] **Step 2: Rewrite using DemoBase** + +Keep MoveEndEffector + PickUp + MoveHeldObject registration and HeldObjectPoseTarget logic. + +- [ ] **Step 3: Run integration test** + +```bash +python scripts/tutorials/atomic_action/move_held_object.py --headless --auto_play +``` + +Expected: runs to completion. + +- [ ] **Step 4: Commit** + +```bash +black . +git add scripts/tutorials/atomic_action/move_held_object.py +git commit -m "refactor(tutorial): move_held_object uses DemoBase + +Co-Authored-By: Claude " +``` + +--- + +## Task 11: Refactor press.py to use DemoBase + +**Files:** +- Modify: `scripts/tutorials/atomic_action/press.py` + +**Interfaces:** Same as Task 6, plus custom robot/table/block creation. + +- [ ] **Step 1: Read current file** + +- [ ] **Step 2: Rewrite using DemoBase** + +Keep the custom robot configuration, table+block creation, press center verification, and `run_press_demo` logic. Use `DemoBase` for lifecycle and shared helpers for recording/trajectory replay where applicable. + +- [ ] **Step 3: Run integration test** + +```bash +python scripts/tutorials/atomic_action/press.py --headless --auto_play +``` + +Expected: runs to completion. + +- [ ] **Step 4: Commit** + +```bash +black . +git add scripts/tutorials/atomic_action/press.py +git commit -m "refactor(tutorial): press uses DemoBase + +Co-Authored-By: Claude " +``` + +--- + +## Task 12: Run full integration test suite for atomic_action tutorials + +**Files:** +- All modified atomic_action scripts + +- [ ] **Step 1: Run all six scripts headless + auto_play** + +```bash +for script in move_joints move_end_effector pickup place move_held_object press; do + echo "Running $script..." + python "scripts/tutorials/atomic_action/${script}.py" --headless --auto_play || exit 1 +done +``` + +Expected: All six scripts exit 0. + +- [ ] **Step 2: Run unit tests** + +```bash +pytest tests/sim/utility/test_demo_utils.py tests/sim/test_demo_base.py -v +``` + +Expected: PASS. + +- [ ] **Step 3: Run pre-commit checks** + +Use the `/pre-commit-check` skill or run: + +```bash +black . +``` + +- [ ] **Step 4: Final commit if any fixes** + +```bash +git commit -am "chore: final formatting and integration fixes + +Co-Authored-By: Claude " || true +``` + +--- + +## Self-Review + +**1. Spec coverage:** +- `demo_utils.py` covers argument parsing, sim creation, cleanup, recording, trajectory replay, print options, tensor formatting, GPU physics helper. ✓ +- `DemoBase` covers optional lifecycle base class. ✓ +- `tutorial_utils.py` backward compatibility covered by re-exports and compatibility wrappers. ✓ +- All six atomic_action scripts are migrated. ✓ +- Tests for `demo_utils.py` and `DemoBase` included. ✓ +- Integration tests for all scripts included. ✓ + +**2. Placeholder scan:** +- No TBD/TODO. ✓ +- No vague "add error handling" steps. ✓ +- Test code is concrete. ✓ + +**3. Type consistency:** +- `DemoRecording.__init__` signature updated to include `look_at` in Task 5; all references match. ✓ +- `shutdown_sim(sim)` used consistently. ✓ +- `SimulationManager` type hints consistent. ✓ + +**Gap identified:** `create_robot_from_preset` was mentioned in the design doc but is intentionally scoped out of this plan. The six atomic_action scripts all use `create_ur5_gripper_robot_cfg` directly, so a generic preset factory is not needed for the pilot migration. It can be added later when migrating `examples/sim/demo/`. diff --git a/docs/superpowers/specs/2026-07-04-demo-base-design.md b/docs/superpowers/specs/2026-07-04-demo-base-design.md new file mode 100644 index 000000000..8c2947139 --- /dev/null +++ b/docs/superpowers/specs/2026-07-04-demo-base-design.md @@ -0,0 +1,188 @@ +# Demo 脚本标准类 / 共享工具设计文档 + +**Date:** 2026-07-04 +**Author:** Claude Code (brainstorming session) +**Status:** Approved + +--- + +## 1. 背景与问题 + +`scripts/tutorials/` 与 `examples/` 中包含大量 demo 脚本，它们在完成不同任务的同时，重复了相当多的样板代码： + +- `np.set_printoptions` / `torch.set_printoptions` +- `SimulationManager` 配置、灯光、窗口打开 +- 命令行参数解析（`--headless`、`--device`、`--renderer`、`--auto_play`、`--record_*`） +- 视频录制（`start_window_record` / `stop_window_record` / `wait_window_record_saves`） +- 轨迹回放循环 +- `KeyboardInterrupt` + `sim.destroy()` 清理 + +`scripts/tutorials/atomic_action/tutorial_utils.py` 已经封装了一部分功能，但： + +1. 位置偏（仅 atomic_action 使用）； +2. 缺少 headless / 离屏 camera 录制； +3. 没有可复用的生命周期基类。 + +本设计旨在抽象出一个**可复用的 demo 工具层 + 可选生命周期基类**，降低新 demo 的书写成本，统一录制、清理等行为。 + +--- + +## 2. 设计目标 + +1. **减少样板代码**：新 demo 只需要关注核心逻辑（setup + run）。 +2. **统一参数与录制**：所有 demo 共享一致的 CLI 参数和视频录制行为。 +3. **向后兼容**：不破坏现有 `tutorial_utils.py` 的导入。 +4. **渐进式迁移**：先改造 `atomic_action` 作为试点，再推广到 `examples/sim/demo/`。 +5. **不强迫短脚本**：gizmo / sensor / solver 等生命周期短的脚本可继续扁平编写。 + +--- + +## 3. 非目标 + +- 不替代 `embodichain.lab.gym.utils.gym_utils.add_env_launcher_args_to_parser()`，而是在其之上追加 demo 专用参数。 +- 不封装所有可能的机器人配置；只提供常见 preset 工厂，复杂配置仍由脚本自行构造。 +- 不引入重量级插件系统或配置化 DSL。 + +--- + +## 4. 架构 + +新增两个共享模块： + +``` +embodichain/lab/sim/ +├── utility/ +│ └── demo_utils.py # 可独立调用的工具函数 / 小类 +└── demo_base.py # 可选的轻量 DemoBase 生命周期基类 +``` + +### 4.1 `demo_utils.py` — 工具函数层（核心） + +所有 demo 脚本均可按需调用，保持扁平自由。 + +主要 API： + +| 函数 / 类 | 说明 | +|---|---| +| `add_demo_args(parser)` | 在 `add_env_launcher_args_to_parser()` 基础上追加 `--auto_play`、`--record_steps`、`--record_fps`、`--record_save_path`、`--no_vis_eef_axis`。 | +| `create_default_sim(args, *, width, height, physics_dt, arena_space, add_default_light)` | 返回配置好的 `SimulationManager`，默认加一盏主灯。 | +| `shutdown_sim(sim)` | 安全调用 `sim.destroy()`。 | +| `DemoRecording(sim, args, prefix)` | 上下文管理器，统一处理 window record / headless offscreen camera 录制、文件名生成、停止与等待保存。 | +| `replay_trajectory(sim, robot, traj, *, post_steps, step_size, sleep, arm_name)` | 统一轨迹回放循环。 | +| `maybe_open_window(sim, args)` | `if not args.headless: sim.open_window()`。 | +| `maybe_wait_for_user(args, prompt)` | `if not args.auto_play: input(prompt)`。 | +| `maybe_pause_for_inspection(args)` | `if not args.auto_play: input("Press Enter to finish...")`。 | +| `maybe_init_gpu_physics(sim)` | `if sim.is_use_gpu_physics: sim.init_gpu_physics()`。 | +| `setup_print_options()` | 设置 `np` / `torch` 打印格式。 | +| `format_tensor(tensor)` | 标准化 tensor 字符串输出。 | +| `create_robot_from_preset(sim, preset, **kwargs)` | 常见机器人 preset 工厂（如 `"ur5_gripper"`、`"ur10"`、`"dexforce_w1"`），复杂配置仍由脚本自行构造。 | + +### 4.2 `demo_base.py` — 可选生命周期基类 + +适合有完整 setup/run/cleanup 的 demo。 + +```python +class DemoBase(ABC): + def __init__(self, args: argparse.Namespace): + self.args = args + self.sim: SimulationManager | None = None + + @abstractmethod + def setup(self) -> None: + """Create sim, robot, camera, etc.""" + + @abstractmethod + def run(self) -> None: + """Demo main logic.""" + + def cleanup(self) -> None: + if self.sim is not None: + shutdown_sim(self.sim) + + def main(self) -> None: + self.setup() + try: + self.run() + finally: + self.cleanup() +``` + +子类只需实现 `setup()` 与 `run()`，生命周期由基类保证。 + +--- + +## 5. 数据流 + +### 5.1 使用 `DemoBase` 的脚本（如 `atomic_action/pickup.py`） + +``` +main() + └── DemoBase.main() + ├── add_demo_args(parser) + parse_args() + ├── setup() + │ ├── create_default_sim(args) → SimulationManager + light + │ ├── create_robot_from_preset(sim, "ur5_gripper") + │ └── create_default_motion_generator(robot) + ├── run() + │ ├── maybe_wait_for_user(args, "Press Enter to plan...") + │ ├── with DemoRecording(sim, args, prefix="pickup"): + │ │ └── replay_trajectory(sim, robot, traj) + │ └── maybe_pause_for_inspection(args) + └── cleanup() + └── shutdown_sim(sim) +``` + +### 5.2 使用纯工具函数的短脚本（如 `motion_generator.py`） + +```python +args = add_demo_args(parser).parse_args() +sim = create_default_sim(args) +try: + # ... demo-specific logic ... + with DemoRecording(sim, args, prefix="motion_generator"): + # ... +finally: + shutdown_sim(sim) +``` + +--- + +## 6. 错误处理 + +1. **录制失败不中断 demo**：`DemoRecording` 内部捕获 `start_window_record` 失败，打印 warning 并跳过录制，主流程继续。 +2. **清理必执行**：`DemoBase.main()` 使用 `try/finally`；纯函数场景在文档中给出 `try/finally` 示例。 +3. **参数缺失降级**：`--record_save_path` 未指定时，默认保存到 `./recordings/_.mp4`。 +4. **GPU 物理初始化显式化**：`create_default_sim()` 不自动调用 `init_gpu_physics()`，保留 `maybe_init_gpu_physics(sim)` 让用户在 setup 后按需调用，避免隐藏副作用。 + +--- + +## 7. 迁移策略 + +1. **保留 `tutorial_utils.py`**：`scripts/tutorials/atomic_action/tutorial_utils.py` 保留 UR5+gripper 相关的专用配置函数（如 `create_ur5_gripper_robot_cfg`），并把通用工具函数（如 `start_auto_play_recording`、`stop_auto_play_recording`、`draw_axis_marker`）改为从新的 `demo_utils.py` re-export，现有导入继续工作。 +2. **试点**：优先改造 `scripts/tutorials/atomic_action/` 6 个脚本，验证 API 设计。 +3. **推广**：第二步改造 `examples/sim/demo/` 4 个脚本（`grasp_cup_to_caffe`、`pick_up_cloth`、`press_softbody`、`scoop_ice`）。 +4. **不强制迁移**：gizmo / sensor / solver 等短脚本可继续使用原有写法或仅引用部分工具函数。 + +--- + +## 8. 测试策略 + +1. **单元测试**：`tests/lab/sim/test_demo_utils.py` + - `add_demo_args` 添加了预期参数； + - `format_tensor` 输出格式正确； + - `DemoRecording` 在 `auto_play=False` 时不启动录制； + - `shutdown_sim` 正确调用 `sim.destroy()`。 + +2. **集成测试**：改造 `atomic_action/pickup.py` 后，以 `--headless --auto_play` 运行，确认不 crash 且生成视频文件。 + +3. **向后兼容测试**：确认旧脚本继续导入 `tutorial_utils.py` 无异常。 + +--- + +## 9. 后续工作 + +1. 实现 `embodichain/lab/sim/utility/demo_utils.py`。 +2. 实现 `embodichain/lab/sim/demo_base.py`。 +3. 更新 `scripts/tutorials/atomic_action/tutorial_utils.py` 为 re-export 模式。 +4. 改造 `atomic_action` 脚本并跑通集成测试。 +5. 视情况推广到 `examples/sim/demo/`。 diff --git a/embodichain/data_pipeline/engine/data.py b/embodichain/data_pipeline/engine/data.py index 1b330b914..eb7329d87 100644 --- a/embodichain/data_pipeline/engine/data.py +++ b/embodichain/data_pipeline/engine/data.py @@ -95,14 +95,30 @@ def _sim_worker_fn( Remains set permanently thereafter. close_signal: Event set by the main process to request a graceful shutdown. """ + from threading import RLock + import gymnasium as gym from embodichain.lab.gym.utils.gym_utils import ( config_to_cfg, get_manager_modules, ) + from embodichain.lab.gym.utils.registration import ( + discover_task_packages, + execute_init_hooks, + ) from embodichain.lab.sim import SimulationManagerCfg from embodichain.utils.logger import log_info, log_warning, log_error + # The worker starts in a fresh interpreter. Discover task packages and + # extension hooks here so custom environments and manager functors are + # registered before gym.make() is called. + discover_task_packages() + execute_init_hooks() + # This worker is the only tqdm writer. A thread-only lock avoids leaking + # tqdm's multiprocessing semaphore when SimulationManager exits the worker + # with os._exit() during native resource teardown. + tqdm.set_lock(RLock()) + gym_config: dict = cfg.gym_config action_config: dict = cfg.action_config @@ -351,6 +367,11 @@ def start(self) -> None: self._fill_signal.set() while not self.is_init: + if not self._sim_process.is_alive(): + raise RuntimeError( + "OnlineDataEngine simulation subprocess exited before " + "initializing the shared buffer." + ) time.sleep(0.5) # ----------------------------------------------------------------------- @@ -525,21 +546,34 @@ def stop(self) -> None: Safe to call multiple times — subsequent calls are no-ops if the subprocess has already been terminated. """ - if self._sim_process is None or not self._sim_process.is_alive(): + process = self._sim_process + if process is None: return - # Ask the subprocess to stop and unblock it if it is waiting on fill_signal. - self._close_signal.set() - self._fill_signal.set() + was_alive = process.is_alive() + if was_alive: + # Ask the subprocess to stop and unblock it if it is waiting on + # fill_signal. + self._close_signal.set() + self._fill_signal.set() + + # Allow time for a graceful exit (close_signal is checked between + # steps). + process.join(timeout=5.0) - # Allow time for a graceful exit (close_signal is checked between steps). - self._sim_process.join(timeout=5.0) + if process.is_alive(): + process.terminate() + process.join(timeout=3.0) - if self._sim_process.is_alive(): - self._sim_process.terminate() - self._sim_process.join(timeout=3.0) + if not process.is_alive(): + process.close() + self._sim_process = None - log_info("[OnlineDataEngine] Simulation subprocess terminated.", color="green") + if was_alive: + log_info( + "[OnlineDataEngine] Simulation subprocess terminated.", + color="green", + ) def __del__(self) -> None: self.stop() diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 5d9c34024..a5d04bc70 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -816,6 +816,7 @@ def add_env_launcher_args_to_parser( """ parser.add_argument( "--num_envs", + "--num-envs", help="The number of environments to run in parallel. " "If not given, falls back to the gym config's `num_envs` (default 1).", default=1, @@ -844,12 +845,14 @@ def add_env_launcher_args_to_parser( ) parser.add_argument( "--arena_space", + "--arena-space", help="The size of the arena space.", default=5.0, type=float, ) parser.add_argument( "--gpu_id", + "--gpu-id", help="The GPU ID to use for the simulation.", default=0, type=int, diff --git a/embodichain/lab/sim/demo_base.py b/embodichain/lab/sim/demo_base.py new file mode 100644 index 000000000..d7593a4ad --- /dev/null +++ b/embodichain/lab/sim/demo_base.py @@ -0,0 +1,67 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Optional base class for simulation demos.""" + +from __future__ import annotations + +import argparse +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +from embodichain.lab.sim.utility.demo_utils import shutdown_sim + +if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager + + +__all__ = ["DemoBase"] + + +class DemoBase(ABC): + """Lightweight lifecycle base class for simulation demos. + + Subclasses implement :meth:`setup` and :meth:`run`; the base class handles + argument injection and guaranteed cleanup. + + Args: + args: Parsed command-line arguments. + """ + + def __init__(self, args: argparse.Namespace): + self.args = args + self.sim: SimulationManager | None = None + + @abstractmethod + def setup(self) -> None: + """Create simulation, robot, cameras, etc.""" + + @abstractmethod + def run(self) -> None: + """Execute the demo logic.""" + + def cleanup(self) -> None: + """Release simulation resources. Called automatically by :meth:`main`.""" + if self.sim is not None: + shutdown_sim(self.sim) + + def main(self) -> None: + """Run the full demo lifecycle.""" + try: + self.setup() + self.run() + finally: + self.cleanup() diff --git a/embodichain/lab/sim/sensors/contact_sensor.py b/embodichain/lab/sim/sensors/contact_sensor.py index 49ebbe1c8..7624584f0 100644 --- a/embodichain/lab/sim/sensors/contact_sensor.py +++ b/embodichain/lab/sim/sensors/contact_sensor.py @@ -166,15 +166,18 @@ def _precompute_filter_ids(self, config: ContactSensorCfg): self.item_user_ids = torch.cat( (self.item_user_ids, rigid_object.get_user_ids()) ) - env_ids = torch.tensor( + env_ids = torch.as_tensor( rigid_object._all_indices, dtype=torch.int32, device=self.device ) self.item_env_ids = torch.cat((self.item_env_ids, env_ids)) for articulation_cfg in config.articulation_cfg_list: - articulation = self._sim.get_articulation(articulation_cfg.articulation_uid) - if articulation is None: + if articulation_cfg.articulation_uid in self._sim.get_robot_uid_list(): articulation = self._sim.get_robot(articulation_cfg.articulation_uid) + else: + articulation = self._sim.get_articulation( + articulation_cfg.articulation_uid + ) if articulation is None: logger.log_warning( f"Articulation with uid '{articulation_cfg.articulation_uid}' not found in simulation." @@ -189,12 +192,13 @@ def _precompute_filter_ids(self, config: ContactSensorCfg): for link_name in link_names: if link_name not in all_link_names: logger.log_warning( - f"Link {link_name} not found in articulation {articulation_cfg.uid}." + f"Link {link_name} not found in articulation " + f"{articulation_cfg.articulation_uid}." ) continue link_user_ids = articulation.get_user_ids(link_name).reshape(-1) self.item_user_ids = torch.cat((self.item_user_ids, link_user_ids)) - env_ids = torch.tensor( + env_ids = torch.as_tensor( articulation._all_indices, dtype=torch.int32, device=self.device ) self.item_env_ids = torch.cat((self.item_env_ids, env_ids)) diff --git a/embodichain/lab/sim/utility/__init__.py b/embodichain/lab/sim/utility/__init__.py index 0570c4510..bb69876a8 100644 --- a/embodichain/lab/sim/utility/__init__.py +++ b/embodichain/lab/sim/utility/__init__.py @@ -19,3 +19,4 @@ from .gizmo_utils import * from .keyboard_utils import * from .render_utils import * +from .demo_utils import * diff --git a/embodichain/lab/sim/utility/demo_utils.py b/embodichain/lab/sim/utility/demo_utils.py new file mode 100644 index 000000000..de1a62e15 --- /dev/null +++ b/embodichain/lab/sim/utility/demo_utils.py @@ -0,0 +1,484 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Shared utilities for simulation demo scripts.""" + +from __future__ import annotations + +import argparse +import time +from typing import TYPE_CHECKING + +import numpy as np +import torch + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.objects import Robot + + +__all__ = [ + "add_demo_args", + "create_default_sim", + "shutdown_sim", + "setup_print_options", + "format_tensor", + "maybe_init_gpu_physics", + "DemoRecording", + "maybe_open_window", + "maybe_wait_for_user", + "maybe_pause_for_inspection", + "DEFAULT_DEMO_LOOK_AT", + "resolve_demo_steps", + "run_simulation_loop", + "replay_trajectory", +] + +DEFAULT_DEMO_LOOK_AT = ( + (2.6, -2.2, 1.6), + (0.0, 0.0, 0.45), + (0.0, 0.0, 1.0), +) + + +def add_demo_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: + """Add common demo arguments to an environment launcher parser. + + Args: + parser: The parser to extend. + + Returns: + The same parser with demo flags added. + """ + from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser + + add_env_launcher_args_to_parser(parser) + parser.add_argument( + "--auto_play", + "--auto-play", + action="store_true", + help="Skip interactive prompts and run the demo automatically.", + ) + parser.add_argument( + "--record_steps", + "--record-steps", + type=_positive_int, + default=None, + help=( + "Number of simulation updates to record. Continuous demos also use " + "this as their run limit. If omitted, recording is disabled." + ), + ) + parser.add_argument( + "--record_fps", + "--record-fps", + type=_positive_int, + default=30, + help="Frames per second for the recorded video.", + ) + parser.add_argument( + "--record_save_path", + "--record-save-path", + type=str, + default=None, + help=( + "Output .mp4 path or directory for recorded videos. " + "Defaults to ./recordings." + ), + ) + parser.add_argument( + "--no_vis_eef_axis", + action="store_true", + help="Disable end-effector axis visualization.", + ) + return parser + + +def _positive_int(value: str) -> int: + """Parse a strictly positive integer for an argparse option.""" + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError("must be at least 1") + return parsed + + +def create_default_sim( + args: argparse.Namespace, + *, + width: int = 1920, + height: int = 1080, + physics_dt: float = 1.0 / 100.0, + arena_space: float = 2.5, + num_envs: int = 1, + add_default_light: bool = True, +) -> SimulationManager: + """Create a SimulationManager with common demo defaults. + + Args: + args: Parsed command-line arguments. Expected to contain ``headless``, + ``device`` and ``renderer``. + width: Window/render width. + height: Window/render height. + physics_dt: Physics simulation timestep. + arena_space: Arena space size. + num_envs: Number of parallel environments to simulate. + add_default_light: Whether to add a default point light. + + Returns: + Configured simulation manager instance. + """ + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg + from embodichain.lab.sim.cfg import LightCfg, RenderCfg + + cfg = SimulationManagerCfg( + width=width, + height=height, + headless=args.headless, + sim_device=args.device, + render_cfg=RenderCfg(renderer=args.renderer), + physics_dt=physics_dt, + arena_space=arena_space, + num_envs=num_envs, + gpu_id=getattr(args, "gpu_id", 0), + ) + sim = SimulationManager(cfg) + if add_default_light: + sim.add_light( + cfg=LightCfg( + uid="main_light", + color=(0.6, 0.6, 0.6), + intensity=30.0, + init_pos=(1.0, 0.0, 3.0), + ) + ) + return sim + + +def shutdown_sim(sim: SimulationManager) -> None: + """Safely destroy a simulation manager. + + Args: + sim: The simulation manager to destroy. + """ + # Recording owns renderer resources and must finish before teardown. Use + # attribute checks so this helper remains useful with lightweight test + # doubles and older SimulationManager implementations. + is_recording = getattr(sim, "is_window_recording", None) + try: + if callable(is_recording) and is_recording(): + sim.stop_window_record() + sim.wait_window_record_saves() + finally: + sim.destroy() + + +def setup_print_options() -> None: + """Set common numpy and torch print options for demos.""" + np.set_printoptions(precision=5, suppress=True) + torch.set_printoptions(precision=5, sci_mode=False) + + +def format_tensor(tensor: torch.Tensor) -> str: + """Return a compact, rounded string representation of a tensor. + + Args: + tensor: Input tensor. + + Returns: + Rounded string with 4 decimal places. + """ + values = tensor.detach().cpu().tolist() + return "[" + ", ".join(f"{v:.4f}" for v in values) + "]" + + +def maybe_init_gpu_physics(sim: SimulationManager) -> None: + """Initialize GPU physics if the simulation is configured to use it. + + Args: + sim: The simulation manager. + """ + if sim.is_use_gpu_physics: + sim.init_gpu_physics() + + +class DemoRecording: + """Context manager that handles demo video recording. + + Recording is only started when ``args.record_steps`` is not ``None``. + On exit the window record is stopped and the framework is asked to finish + saving the video file. + + Args: + sim: The simulation manager. + args: Parsed command-line arguments. Expected to contain + ``record_steps``, ``record_fps`` and ``record_save_path``. + prefix: Prefix used for the generated video filename. + look_at: Optional camera look-at tuple for the recording. Headless + recording uses :data:`DEFAULT_DEMO_LOOK_AT` when omitted. + """ + + def __init__( + self, + sim: SimulationManager, + args: argparse.Namespace, + prefix: str = "demo", + look_at: tuple[Sequence[float], Sequence[float], Sequence[float]] | None = None, + ): + self.sim = sim + self.args = args + self.prefix = prefix + self.look_at = ( + DEFAULT_DEMO_LOOK_AT + if look_at is None and getattr(args, "headless", False) + else look_at + ) + self.is_active = False + + def __enter__(self) -> DemoRecording: + """Start recording if requested.""" + if self.args.record_steps is None: + return self + + import datetime + import warnings + from pathlib import Path + + requested_path = Path(self.args.record_save_path or "./recordings") + if requested_path.suffix.lower() == ".mp4": + requested_path.parent.mkdir(parents=True, exist_ok=True) + save_path = str(requested_path) + else: + requested_path.mkdir(parents=True, exist_ok=True) + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + save_path = str(requested_path / f"{self.prefix}_{timestamp}.mp4") + + original_width = self.sim.sim_config.width + original_height = self.sim.sim_config.height + try: + # Use a smaller resolution for recording to keep files small. + self.sim.sim_config.width = 640 + self.sim.sim_config.height = 480 + started = self.sim.start_window_record( + save_path=save_path, + fps=self.args.record_fps, + max_memory=2048, + video_prefix=self.prefix, + look_at=self.look_at, + use_sim_time=True, + ) + finally: + self.sim.sim_config.width = original_width + self.sim.sim_config.height = original_height + + if not started: + warnings.warn( + f"Failed to start recording for prefix '{self.prefix}'. Continuing without recording.", + UserWarning, + stacklevel=2, + ) + return self + + self.is_active = True + return self + + def __exit__(self, exc_type, exc, tb) -> None: + """Stop recording and wait for the file to be written.""" + if not self.is_active: + return + if self.sim.is_window_recording(): + self.sim.stop_window_record() + self.sim.wait_window_record_saves() + + +def maybe_open_window(sim: SimulationManager, args: argparse.Namespace) -> None: + """Open the viewer window unless running headless. + + Args: + sim: The simulation manager. + args: Parsed arguments containing ``headless``. + """ + if not args.headless: + sim.open_window() + + +def maybe_wait_for_user(args: argparse.Namespace, prompt: str) -> None: + """Wait for user input unless auto_play is enabled. + + Args: + args: Parsed arguments containing ``auto_play``. + prompt: Message to display when waiting. + """ + if not args.auto_play: + input(prompt) + + +def maybe_pause_for_inspection(args: argparse.Namespace) -> None: + """Pause at the end of a demo for visual inspection. + + Args: + args: Parsed arguments containing ``auto_play``. + """ + maybe_wait_for_user(args, "Demo finished. Press Enter to exit...") + + +def resolve_demo_steps( + args: argparse.Namespace, + *, + auto_play_steps: int = 300, +) -> int | None: + """Resolve the run limit for a continuous demo. + + Explicit ``--record_steps`` takes precedence. ``--auto_play`` uses a + finite default so non-interactive smoke runs terminate on their own. + Interactive runs remain open until interrupted. + + Args: + args: Parsed demo arguments. + auto_play_steps: Default update count used by ``--auto_play``. + + Returns: + Maximum number of updates, or ``None`` for an interactive run. + + Raises: + ValueError: If ``auto_play_steps`` is not positive. + """ + if auto_play_steps < 1: + raise ValueError("auto_play_steps must be at least 1") + record_steps = getattr(args, "record_steps", None) + if record_steps is not None: + return record_steps + return auto_play_steps if getattr(args, "auto_play", False) else None + + +def run_simulation_loop( + sim: SimulationManager, + *, + max_steps: int | None = None, + steps_per_update: int = 1, + sleep: float = 0.0, + log_interval: int | None = 100, + on_step: Callable[[int], None] | None = None, +) -> int: + """Run a standard simulation update loop. + + The function intentionally does not destroy ``sim``; callers should use + :class:`~embodichain.lab.sim.demo_base.DemoBase` or ``try/finally`` so + setup failures and loop failures share the same cleanup path. + + Args: + sim: Simulation manager to update. + max_steps: Optional number of update calls before returning. + steps_per_update: Physics steps advanced by each update call. + sleep: Optional wall-clock delay after each update. + log_interval: Print aggregate FPS every this many updates. Set to + ``None`` to disable progress logging. + on_step: Optional callback receiving the one-based update count. + + Returns: + Number of completed update calls. + + Raises: + ValueError: If a numeric loop option is outside its valid range. + """ + if max_steps is not None and max_steps < 1: + raise ValueError("max_steps must be at least 1") + if steps_per_update < 1: + raise ValueError("steps_per_update must be at least 1") + if sleep < 0: + raise ValueError("sleep must be non-negative") + if log_interval is not None and log_interval < 1: + raise ValueError("log_interval must be at least 1") + + started_at = time.monotonic() + last_log_at = started_at + last_log_step = 0 + step_count = 0 + + try: + while max_steps is None or step_count < max_steps: + sim.update(step=steps_per_update) + step_count += 1 + if on_step is not None: + on_step(step_count) + if sleep: + time.sleep(sleep) + + if log_interval is not None and step_count % log_interval == 0: + now = time.monotonic() + elapsed = now - last_log_at + fps = ( + sim.num_envs + * (step_count - last_log_step) + * steps_per_update + / elapsed + if elapsed > 0 + else 0.0 + ) + print(f"[INFO]: Simulation step: {step_count}, FPS: {fps:.2f}") + last_log_at = now + last_log_step = step_count + except KeyboardInterrupt: + print("\n[INFO]: Stopping simulation...") + + return step_count + + +def replay_trajectory( + sim: SimulationManager, + robot: Robot, + traj: torch.Tensor, + *, + post_steps: int = 60, + step_size: int = 4, + sleep: float = 1e-2, + arm_name: str | None = None, +) -> None: + """Replay a joint-space trajectory on a robot. + + ``traj`` may be either a 1-D tensor of shape ``(num_joints,)``, a 2-D + tensor of shape ``(num_steps, num_joints)`` or a 3-D tensor of shape + ``(batch, num_steps, num_joints)``. For 1-D input the single + configuration is held for ``post_steps``. For 2-D/3-D input each step is + applied sequentially and the final configuration is held. + + Args: + sim: The simulation manager. + robot: The robot instance. + traj: Joint position trajectory tensor. + post_steps: Number of steps to hold the final configuration. + step_size: Number of physics steps per ``sim.update()`` call. + sleep: Sleep duration between steps (seconds). + arm_name: Optional arm name passed to ``robot.set_qpos``. + """ + if traj.dim() == 1: + traj = traj.unsqueeze(0).unsqueeze(0) + elif traj.dim() == 2: + traj = traj.unsqueeze(0) + + joint_ids = robot.get_joint_ids(arm_name) if arm_name is not None else None + + for i in range(traj.shape[1]): + robot.set_qpos(qpos=traj[:, i, :], joint_ids=joint_ids) + sim.update(step=step_size) + time.sleep(sleep) + + final_qpos = traj[:, -1, :] + for _ in range(post_steps): + robot.set_qpos(qpos=final_qpos, joint_ids=joint_ids) + sim.update(step=2) + time.sleep(sleep) diff --git a/examples/data_pipeline/online_dataset_demo.py b/examples/data_pipeline/online_dataset_demo.py index 83f8af808..335893d48 100644 --- a/examples/data_pipeline/online_dataset_demo.py +++ b/examples/data_pipeline/online_dataset_demo.py @@ -28,14 +28,12 @@ Usage:: - python examples/data_pipeline/online_dataset_demo.py + python examples/data_pipeline/online_dataset_demo.py --device cpu """ from __future__ import annotations import argparse -import json -import time from pathlib import Path from torch.utils.data import DataLoader @@ -48,15 +46,52 @@ from embodichain.data_pipeline.engine.data import OnlineDataEngine, OnlineDataEngineCfg from embodichain.utils.logger import log_info +DEFAULT_CONFIG = ( + Path(__file__).resolve().parents[2] + / "embodichain_tasks/configs/gym/special/simple_task_ur10.json" +) + + +def _positive_int(value: str) -> int: + """Parse a positive integer command-line option.""" + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError("must be at least 1") + return parsed + def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="OnlineDataset demo") + parser.add_argument( + "--config", + type=Path, + default=DEFAULT_CONFIG, + help="Gym configuration file.", + ) parser.add_argument( "--device", type=str, default="cpu", help="Simulation device, e.g. 'cpu' or 'cuda:0' (default: cpu).", ) + parser.add_argument( + "--renderer", + choices=["auto", "hybrid", "fast-rt", "rt"], + default="auto", + help="Renderer backend used by the simulation worker.", + ) + parser.add_argument( + "--gpu-id", + type=int, + default=0, + help="GPU used by the simulation worker.", + ) + parser.add_argument( + "--num-batches", + type=_positive_int, + default=5, + help="Number of batches shown for each dataset mode.", + ) return parser.parse_args() @@ -66,29 +101,28 @@ def _parse_args() -> argparse.Namespace: def _build_engine(args: argparse.Namespace) -> OnlineDataEngine: - """Construct and start an OnlineDataEngine from the given CLI args.""" - config_path = Path("embodichain_tasks/configs/gym/special/simple_task_ur10.json") - if not config_path.exists(): + """Construct an OnlineDataEngine from the given CLI args.""" + if not args.config.is_file(): raise FileNotFoundError( - f"Gym config not found: {config_path}. " - "Provide a valid path via --config." + f"Gym config not found: {args.config}. Provide a valid path via --config." ) from embodichain.utils.utility import load_config - gym_config = load_config(config_path) + gym_config = load_config(args.config) gym_config["headless"] = True - gym_config.setdefault("renderer", True) - gym_config["gpu_id"] = 0 + gym_config["renderer"] = args.renderer + gym_config["gpu_id"] = args.gpu_id gym_config["device"] = args.device cfg = OnlineDataEngineCfg( - buffer_size=2, state_dim=6, gym_config=gym_config, buffer_device=args.device + buffer_size=2, + state_dim=6, + gym_config=gym_config, + # The engine shares this buffer with its simulation subprocess. + buffer_device="cpu", ) - engine = OnlineDataEngine(cfg) - engine.start() - - return engine + return OnlineDataEngine(cfg) # --------------------------------------------------------------------------- @@ -225,12 +259,13 @@ def main() -> None: engine = _build_engine(args) try: - _demo_item_mode(engine, chunk_size=32, num_batches=5) - _demo_batch_mode(engine, chunk_size=32, num_batches=5) - _demo_uniform_dynamic(engine, num_batches=5) - _demo_gmm_dynamic(engine, num_batches=5) + engine.start() + _demo_item_mode(engine, chunk_size=32, num_batches=args.num_batches) + _demo_batch_mode(engine, chunk_size=32, num_batches=args.num_batches) + _demo_uniform_dynamic(engine, num_batches=args.num_batches) + _demo_gmm_dynamic(engine, num_batches=args.num_batches) finally: - # engine.stop() + engine.stop() log_info("[Demo] Engine stopped.", color="green") diff --git a/examples/gym/pour_water.sh b/examples/gym/pour_water.sh index 2c91f31af..5d2e08a5d 100755 --- a/examples/gym/pour_water.sh +++ b/examples/gym/pour_water.sh @@ -1,3 +1,22 @@ +#!/usr/bin/env bash +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +set -euo pipefail + embodichain run-env --gym_config embodichain_tasks/configs/gym/pour_water/gym_config.json \ --action_config embodichain_tasks/configs/gym/pour_water/action_config.json \ --filter_visual_rand diff --git a/examples/gym/run_blocks_ranking_rgb.sh b/examples/gym/run_blocks_ranking_rgb.sh old mode 100644 new mode 100755 index dfbbbb817..d12350e10 --- a/examples/gym/run_blocks_ranking_rgb.sh +++ b/examples/gym/run_blocks_ranking_rgb.sh @@ -1,3 +1,22 @@ +#!/usr/bin/env bash +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +set -euo pipefail + embodichain run-env \ --gym_config embodichain_tasks/configs/gym/blocks_ranking_rgb/cobot_magic_3cam.json \ --num_envs 1 \ diff --git a/examples/gym/run_scoop_ice.sh b/examples/gym/run_scoop_ice.sh index c1c15cc07..ad066c0bd 100755 --- a/examples/gym/run_scoop_ice.sh +++ b/examples/gym/run_scoop_ice.sh @@ -1 +1,20 @@ +#!/usr/bin/env bash +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +set -euo pipefail + embodichain run-env --gym_config embodichain_tasks/configs/gym/scoop_ice/gym_config.json diff --git a/examples/gym/run_stack_blocks_two.sh b/examples/gym/run_stack_blocks_two.sh old mode 100644 new mode 100755 index 9e1558811..0b96c281d --- a/examples/gym/run_stack_blocks_two.sh +++ b/examples/gym/run_stack_blocks_two.sh @@ -1,3 +1,22 @@ +#!/usr/bin/env bash +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +set -euo pipefail + embodichain run-env \ --gym_config embodichain_tasks/configs/gym/stack_blocks_two/cobot_magic_3cam.json \ --num_envs 1 \ diff --git a/examples/sim/demo/grasp_cup_to_caffe.py b/examples/sim/demo/grasp_cup_to_caffe.py index 25a040e43..27ccb2d06 100644 --- a/examples/sim/demo/grasp_cup_to_caffe.py +++ b/examples/sim/demo/grasp_cup_to_caffe.py @@ -19,431 +19,371 @@ and performs a grasp cup to coffee machine task in a simulated environment. """ +from __future__ import annotations + import argparse + import numpy as np import torch from tqdm import tqdm -from typing import Union -from scipy.spatial.transform import Rotation as R -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.objects import Robot, RigidObject + +from embodichain.data import get_data_path from embodichain.lab.sim.cfg import ( - RenderCfg, - LightCfg, + ArticulationCfg, JointDrivePropertiesCfg, - RigidObjectCfg, RigidBodyAttributesCfg, - ArticulationCfg, + RigidObjectCfg, ) -from embodichain.lab.sim.utility.action_utils import interpolate_with_distance +from embodichain.lab.sim.demo_base import DemoBase +from embodichain.lab.sim.objects import RigidObject, Robot +from embodichain.lab.sim.robots.dexforce_w1.cfg import DexforceW1Cfg from embodichain.lab.sim.shapes import MeshCfg -from embodichain.data import get_data_path +from embodichain.lab.sim.utility.action_utils import interpolate_with_distance +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + resolve_demo_steps, + run_simulation_loop, + setup_print_options, +) from embodichain.utils import logger -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.robots.dexforce_w1.cfg import DexforceW1Cfg - -def parse_arguments(): - """ - Parse command-line arguments to configure the simulation. - Returns: - argparse.Namespace: Parsed arguments including number of environments and rendering options. - """ +class GraspCupToCaffeDemo(DemoBase): + """Grasp a cup with the dexforce_w1 right arm and place it on the caffe.""" + + def setup(self) -> None: + """Create the simulation, robot, table, caffe and cup, then settle.""" + self.sim = create_default_sim( + self.args, + arena_space=2.5, + num_envs=self.args.num_envs, + add_default_light=False, + ) + self.robot = self._create_robot() + self.table = self._create_table() + self.caffe = self._create_caffe() + self.cup = self._create_cup() + self.sim.update(step=1) + + # Apply random perturbation. + self.apply_random_xy_perturbation(self.cup, max_perturbation=0.05) + self.apply_random_xy_perturbation(self.caffe, max_perturbation=0.05) + + maybe_open_window(self.sim, self.args) + maybe_init_gpu_physics(self.sim) + + def run(self) -> None: + """Execute the grasp-and-place trajectory and keep the sim live.""" + with DemoRecording(self.sim, self.args, prefix="grasp_cup_to_caffe"): + self._run_simulation() + logger.log_info("\n Press Ctrl+C to exit simulation loop.") + run_simulation_loop( + self.sim, + max_steps=resolve_demo_steps(self.args), + steps_per_update=10, + ) + + def _create_robot(self) -> Robot: + """Create and configure the dexforce_w1 robot. + + Returns: + The configured robot instance added to the simulation. + """ + cfg = DexforceW1Cfg.from_dict( + { + "uid": "dexforce_w1", + "init_pos": [0.4, -0.5, 0.0], + } + ) + 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, + ] + return self.sim.add_robot(cfg=cfg) + + def _create_table(self) -> RigidObject: + """Create the table rigid object. + + Returns: + The table object added to the simulation. + """ + scoop_cfg = RigidObjectCfg( + uid="table", + shape=MeshCfg( + fpath=get_data_path("MultiW1Data/table_a.obj"), + ), + attrs=RigidBodyAttributesCfg( + mass=0.5, + ), + max_convex_hull_num=8, + body_type="kinematic", + init_pos=[1.1, -0.5, 0.08], + init_rot=[0.0, 0.0, 0.0], + ) + scoop = self.sim.add_rigid_object(cfg=scoop_cfg) + return scoop + + def _create_caffe(self) -> Robot: + """Create the caffe (container) articulated object. + + Returns: + The caffe object added to the simulation. + """ + container_cfg = ArticulationCfg( + uid="caffe", + fpath=get_data_path("MultiW1Data/cafe/cafe.urdf"), + init_pos=[1.05, -0.5, 0.79], + init_rot=[0, 0, -30], + attrs=RigidBodyAttributesCfg( + mass=1.0, + ), + drive_pros=JointDrivePropertiesCfg( + stiffness=1.0, damping=0.1, max_effort=100.0, drive_type="force" + ), + ) + container = self.sim.add_articulation(cfg=container_cfg) + return container + + def _create_cup(self) -> RigidObject: + """Create the cup rigid object. + + Returns: + The cup object added to the simulation. + """ + scoop_cfg = RigidObjectCfg( + uid="cup", + shape=MeshCfg( + fpath=get_data_path("MultiW1Data/paper_cup_2.obj"), + ), + attrs=RigidBodyAttributesCfg( + mass=0.3, + ), + max_convex_hull_num=1, + body_type="dynamic", + init_pos=[0.86, -0.76, 0.841], + init_rot=[0.0, 0.0, 0.0], + ) + scoop = self.sim.add_rigid_object(cfg=scoop_cfg) + return scoop + + def _create_trajectory(self) -> torch.Tensor: + """Generate the right-arm trajectory to grasp the cup and place it. + + Returns: + Interpolated trajectory of shape ``[n_envs, n_waypoint, dof]``. + """ + robot = self.robot + cup = self.cup + caffe = self.caffe + right_arm_ids = robot.get_joint_ids("right_arm") + hand_open_qpos = torch.tensor( + [0.0, 1.5, 0.0, 0.0, 0.0, 0.0], + dtype=torch.float32, + device=self.sim.device, + ) + hand_close_qpos = torch.tensor( + [0.1, 1.5, 0.3, 0.2, 0.3, 0.3], + dtype=torch.float32, + device=self.sim.device, + ) + + cup_position = cup.get_local_pose(to_matrix=True)[:, :3, 3] + + # Grasp cup waypoint generation. + rest_right_qpos = robot.get_qpos()[:, right_arm_ids] # [n_envs, dof] + right_arm_xpos = robot.compute_fk( + qpos=rest_right_qpos, name="right_arm", to_matrix=True + ) + approach_cup_relative_position = torch.tensor( + [-0.05, -0.06, 0.025], dtype=torch.float32, device=self.sim.device + ) + pick_cup_relative_position = torch.tensor( + [-0.03, -0.028, 0.021], dtype=torch.float32, device=self.sim.device + ) + + approach_xpos = right_arm_xpos.clone() + approach_xpos[:, :3, 3] = cup_position + approach_cup_relative_position + + pick_xpos = right_arm_xpos.clone() + pick_xpos[:, :3, 3] = cup_position + pick_cup_relative_position + + lift_xpos = pick_xpos.clone() + lift_xpos[:, 2, 3] += 0.07 + + # Place cup to caffe waypoint generation. + caffe_position = caffe.get_local_pose(to_matrix=True)[:, :3, 3] + place_cup_up_relative_position = torch.tensor( + [-0.14, -0.18, 0.13], dtype=torch.float32, device=self.sim.device + ) + place_cup_down_relative_position = torch.tensor( + [-0.14, -0.18, 0.09], dtype=torch.float32, device=self.sim.device + ) + + place_cup_up_pose = lift_xpos.clone() + place_cup_up_pose[:, :3, 3] = caffe_position + place_cup_up_relative_position + place_down_pose = lift_xpos.clone() + place_down_pose[:, :3, 3] = caffe_position + place_cup_down_relative_position + # Compute ik for each waypoint. + is_success, approach_qpos = robot.compute_ik( + pose=approach_xpos, joint_seed=rest_right_qpos, name="right_arm" + ) + is_success, pick_qpos = robot.compute_ik( + pose=pick_xpos, joint_seed=approach_qpos, name="right_arm" + ) + is_success, lift_qpos = robot.compute_ik( + pose=lift_xpos, joint_seed=pick_qpos, name="right_arm" + ) + is_success, place_up_qpos = robot.compute_ik( + pose=place_cup_up_pose, joint_seed=lift_qpos, name="right_arm" + ) + is_success, place_down_qpos = robot.compute_ik( + pose=place_down_pose, joint_seed=place_up_qpos, name="right_arm" + ) + + n_envs = self.sim.num_envs + + # Combine hand and arm trajectory. + arm_trajectory = torch.cat( + [ + rest_right_qpos[:, None, :], + approach_qpos[:, None, :], + pick_qpos[:, None, :], + pick_qpos[:, None, :], + lift_qpos[:, None, :], + place_up_qpos[:, None, :], + place_down_qpos[:, None, :], + place_down_qpos[:, None, :], + lift_qpos[:, None, :], + rest_right_qpos[:, None, :], + ], + dim=1, + ) + hand_trajectory = torch.cat( + [ + hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), + ], + dim=1, + ) + all_trajectory = torch.cat([arm_trajectory, hand_trajectory], dim=-1) + # Trajectory with shape [n_envs, n_waypoint, dof]. + interp_trajectory = interpolate_with_distance( + trajectory=all_trajectory, interp_num=150, device=self.sim.device + ) + return interp_trajectory + + def _run_simulation(self) -> None: + """Execute the generated trajectory to grasp the cup and place it.""" + # [n_envs, n_waypoint, dof] + interp_trajectory = self._create_trajectory() + + right_arm_ids = self.robot.get_joint_ids("right_arm") + right_hand_ids = self.robot.get_joint_ids("right_eef") + combine_ids = np.concatenate([right_arm_ids, right_hand_ids]) + n_waypoints = interp_trajectory.shape[1] + logger.log_info("Executing trajectory...") + for i in tqdm(range(n_waypoints)): + self.robot.set_qpos(interp_trajectory[:, i, :], joint_ids=combine_ids) + self.sim.update(step=10) + + @staticmethod + def apply_random_xy_perturbation( + item: RigidObject | Robot, max_perturbation: float = 0.02 + ) -> None: + """Apply random perturbation to the object's XY position. + + Args: + item: The object to perturb. + max_perturbation: Maximum perturbation magnitude. + """ + item_pose = item.get_local_pose(to_matrix=True) + item_xy = item_pose[:, :2, 3].to("cpu").numpy() + perturbation = np.random.uniform( + low=-max_perturbation, high=max_perturbation, size=item_xy.shape + ) + new_xy = item_xy + perturbation + item_pose[:, :2, 3] = torch.tensor( + new_xy, dtype=torch.float32, device=item_pose.device + ) + item.set_local_pose(item_pose) + + +def main() -> None: + """Entry point for the grasp-cup-to-caffe demo.""" + setup_print_options() parser = argparse.ArgumentParser( description="Create and simulate a robot in SimulationManager" ) - add_env_launcher_args_to_parser(parser) - return parser.parse_args() - - -def initialize_simulation(args) -> SimulationManager: - """ - Initialize the simulation environment based on the provided arguments. - - Args: - args (argparse.Namespace): Parsed command-line arguments. - - Returns: - SimulationManager: Configured simulation manager instance. - """ - config = SimulationManagerCfg( - headless=True, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), - physics_dt=1.0 / 100.0, - num_envs=args.num_envs, - arena_space=2.5, - ) - sim = SimulationManager(config) - - return sim - - -def create_robot(sim: SimulationManager) -> Robot: - """ - Create and configure a robot with an arm and a dexterous hand in the simulation. - - Args: - sim (SimulationManager): The simulation manager instance. - - Returns: - Robot: The configured robot instance added to the simulation. - """ - cfg = DexforceW1Cfg.from_dict( - { - "uid": "dexforce_w1", - "init_pos": [0.4, -0.5, 0.0], - } - ) - 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, - ] - return sim.add_robot(cfg=cfg) - - -def create_table(sim: SimulationManager) -> RigidObject: - """ - Create a table rigid object in the simulation. - - Args: - sim (SimulationManager): The simulation manager instance. - - Returns: - RigidObject: The table object added to the simulation. - """ - scoop_cfg = RigidObjectCfg( - uid="table", - shape=MeshCfg( - fpath=get_data_path("MultiW1Data/table_a.obj"), - ), - attrs=RigidBodyAttributesCfg( - mass=0.5, - ), - max_convex_hull_num=8, - body_type="kinematic", - init_pos=[1.1, -0.5, 0.08], - init_rot=[0.0, 0.0, 0.0], - ) - scoop = sim.add_rigid_object(cfg=scoop_cfg) - return scoop - - -def create_caffe(sim: SimulationManager) -> Robot: - """ - Create a caffe (container) articulated object in the simulation. - - Args: - sim (SimulationManager): The simulation manager instance. - - Returns: - Robot: The caffe object added to the simulation. - """ - container_cfg = ArticulationCfg( - uid="caffe", - fpath=get_data_path("MultiW1Data/cafe/cafe.urdf"), - init_pos=[1.05, -0.5, 0.79], - init_rot=[0, 0, -30], - attrs=RigidBodyAttributesCfg( - mass=1.0, - ), - drive_pros=JointDrivePropertiesCfg( - stiffness=1.0, damping=0.1, max_effort=100.0, drive_type="force" - ), - ) - container = sim.add_articulation(cfg=container_cfg) - return container - - -def create_cup(sim: SimulationManager) -> RigidObject: - """ - Create a cup rigid object in the simulation. - - Args: - sim (SimulationManager): The simulation manager instance. - - Returns: - RigidObject: The cup object added to the simulation. - """ - scoop_cfg = RigidObjectCfg( - uid="cup", - shape=MeshCfg( - fpath=get_data_path("MultiW1Data/paper_cup_2.obj"), - ), - attrs=RigidBodyAttributesCfg( - mass=0.3, - ), - max_convex_hull_num=1, - body_type="dynamic", - init_pos=[0.86, -0.76, 0.841], - init_rot=[0.0, 0.0, 0.0], - ) - scoop = sim.add_rigid_object(cfg=scoop_cfg) - return scoop - - -def create_trajectory( - sim: SimulationManager, robot: Robot, cup: RigidObject, caffe: Robot -) -> torch.Tensor: - """ - Generate a trajectory for the right arm to grasp the cup and move it to the caffe. - - Args: - sim (SimulationManager): The simulation manager instance. - robot (Robot): The robot instance. - cup (RigidObject): The cup object. - caffe (Robot): The caffe object. - - Returns: - torch.Tensor: Interpolated trajectory of shape [n_envs, n_waypoint, dof]. - """ - right_arm_ids = robot.get_joint_ids("right_arm") - hand_open_qpos = torch.tensor( - [0.0, 1.5, 0.0, 0.0, 0.0, 0.0], - dtype=torch.float32, - device=sim.device, - ) - hand_close_qpos = torch.tensor( - [0.1, 1.5, 0.3, 0.2, 0.3, 0.3], - dtype=torch.float32, - device=sim.device, - ) - - cup_position = cup.get_local_pose(to_matrix=True)[:, :3, 3] - - # grasp cup waypoint generation - rest_right_qpos = robot.get_qpos()[:, right_arm_ids] # [n_envs, dof] - right_arm_xpos = robot.compute_fk( - qpos=rest_right_qpos, name="right_arm", to_matrix=True - ) - approach_cup_relative_position = torch.tensor( - [-0.05, -0.06, 0.025], dtype=torch.float32, device=sim.device - ) - pick_cup_relative_position = torch.tensor( - [-0.03, -0.028, 0.021], dtype=torch.float32, device=sim.device - ) - - approach_xpos = right_arm_xpos.clone() - approach_xpos[:, :3, 3] = cup_position + approach_cup_relative_position - - pick_xpos = right_arm_xpos.clone() - pick_xpos[:, :3, 3] = cup_position + pick_cup_relative_position - - lift_xpos = pick_xpos.clone() - lift_xpos[:, 2, 3] += 0.07 - - # place cup to caffe waypoint generation - caffe_position = caffe.get_local_pose(to_matrix=True)[:, :3, 3] - place_cup_up_relative_position = torch.tensor( - [-0.14, -0.18, 0.13], dtype=torch.float32, device=sim.device - ) - place_cup_down_relative_position = torch.tensor( - [-0.14, -0.18, 0.09], dtype=torch.float32, device=sim.device - ) - - place_cup_up_pose = lift_xpos.clone() - place_cup_up_pose[:, :3, 3] = caffe_position + place_cup_up_relative_position - place_down_pose = lift_xpos.clone() - place_down_pose[:, :3, 3] = caffe_position + place_cup_down_relative_position - # compute ik for each waypoint - is_success, approach_qpos = robot.compute_ik( - pose=approach_xpos, joint_seed=rest_right_qpos, name="right_arm" - ) - is_success, pick_qpos = robot.compute_ik( - pose=pick_xpos, joint_seed=approach_qpos, name="right_arm" - ) - is_success, lift_qpos = robot.compute_ik( - pose=lift_xpos, joint_seed=pick_qpos, name="right_arm" - ) - is_success, place_up_qpos = robot.compute_ik( - pose=place_cup_up_pose, joint_seed=lift_qpos, name="right_arm" - ) - is_success, place_down_qpos = robot.compute_ik( - pose=place_down_pose, joint_seed=place_up_qpos, name="right_arm" - ) - - n_envs = sim.num_envs - - # combine hand and arm trajectory - arm_trajectory = torch.cat( - [ - rest_right_qpos[:, None, :], - approach_qpos[:, None, :], - pick_qpos[:, None, :], - pick_qpos[:, None, :], - lift_qpos[:, None, :], - place_up_qpos[:, None, :], - place_down_qpos[:, None, :], - place_down_qpos[:, None, :], - lift_qpos[:, None, :], - rest_right_qpos[:, None, :], - ], - dim=1, - ) - hand_trajectory = torch.cat( - [ - hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), - ], - dim=1, - ) - all_trajectory = torch.cat([arm_trajectory, hand_trajectory], dim=-1) - # trajetory with shape [n_envs, n_waypoint, dof] - interp_trajectory = interpolate_with_distance( - trajectory=all_trajectory, interp_num=150, device=sim.device - ) - return interp_trajectory - - -def run_simulation( - sim: SimulationManager, robot: Robot, cup: RigidObject, caffe: Robot -): - """ - Execute the generated trajectory to drive the robot to complete the grasp and place task. - - Args: - sim (SimulationManager): The simulation manager instance. - robot (Robot): The robot instance. - cup (RigidObject): The cup object. - caffe (Robot): The caffe object. - """ - # [n_envs, n_waypoint, dof] - interp_trajectory = create_trajectory(sim, robot, cup, caffe) - - right_arm_ids = robot.get_joint_ids("right_arm") - right_hand_ids = robot.get_joint_ids("right_eef") - combine_ids = np.concatenate([right_arm_ids, right_hand_ids]) - n_waypoints = interp_trajectory.shape[1] - logger.log_info(f"Executing trajectory...") - for i in tqdm(range(n_waypoints)): - robot.set_qpos(interp_trajectory[:, i, :], joint_ids=combine_ids) - sim.update(step=10) - - -def apply_random_xy_perturbation( - item: Union[RigidObject, Robot], max_perturbation: float = 0.02 -): - """ - Apply random perturbation to the object's XY position. - - Args: - item (Union[RigidObject, Robot]): The object to perturb. - max_perturbation (float): Maximum perturbation magnitude. - """ - item_pose = item.get_local_pose(to_matrix=True) - item_xy = item_pose[:, :2, 3].to("cpu").numpy() - perturbation = np.random.uniform( - low=-max_perturbation, high=max_perturbation, size=item_xy.shape - ) - new_xy = item_xy + perturbation - item_pose[:, :2, 3] = torch.tensor( - new_xy, dtype=torch.float32, device=item_pose.device - ) - item.set_local_pose(item_pose) - - -def main(): - """ - Main function to demonstrate robot simulation. - - Initializes the simulation, creates the robot and objects, and performs the grasp and place task. - """ - args = parse_arguments() - sim = initialize_simulation(args) - - robot = create_robot(sim) - table = create_table(sim) - caffe = create_caffe(sim) - cup = create_cup(sim) - sim.update(step=1) - - # apply random perturbation - apply_random_xy_perturbation(cup, max_perturbation=0.05) - apply_random_xy_perturbation(caffe, max_perturbation=0.05) - - if not args.headless: - sim.open_window() - - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - - run_simulation(sim, robot, cup, caffe) - - logger.log_info("\n Press Ctrl+C to exit simulation loop.") - try: - counter = 0 - while True: - counter += 1 - sim.update(step=10) - if counter % 10 == 0: - pass - - except KeyboardInterrupt: - logger.log_info("\n Exit") + parser = add_demo_args(parser) + args = parser.parse_args() + GraspCupToCaffeDemo(args).main() if __name__ == "__main__": diff --git a/examples/sim/demo/pick_up_cloth.py b/examples/sim/demo/pick_up_cloth.py index ae86a6e3b..b17ddf4ff 100644 --- a/examples/sim/demo/pick_up_cloth.py +++ b/examples/sim/demo/pick_up_cloth.py @@ -15,282 +15,292 @@ # ---------------------------------------------------------------------------- """ -This script demonstrates the creation and simulation of a robot with a soft object, -and performs a pressing task in a simulated environment. +This script demonstrates the creation and simulation of a robot with a cloth object, +and performs a pick-up task in a simulated environment. """ +from __future__ import annotations + import argparse +import os +import tempfile + import numpy as np -import time import open3d as o3d import torch -from dexsim.utility.path import get_resources_data_path - -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.objects import Robot, SoftObject -from embodichain.lab.sim.utility.action_utils import interpolate_with_distance -from embodichain.lab.sim.shapes import MeshCfg from embodichain.data import get_data_path -from embodichain.utils import logger from embodichain.lab.sim.cfg import ( - RenderCfg, - RigidObjectCfg, - RigidBodyAttributesCfg, - LightCfg, ClothObjectCfg, ClothPhysicalAttributesCfg, + RigidBodyAttributesCfg, + RigidObjectCfg, ) +from embodichain.lab.sim.demo_base import DemoBase +from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.robots import URRobotCfg -import os -from embodichain.lab.sim.shapes import MeshCfg, CubeCfg -import tempfile -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser - +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg +from embodichain.lab.sim.utility.action_utils import interpolate_with_distance +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, +) -def create_robot(sim: SimulationManager, position=[0.0, 0.0, 0.0]): - """ - Create and configure a robot with an arm and a dexterous hand in the simulation. - Args: - sim (SimulationManager): The simulation manager instance. +class PickUpClothDemo(DemoBase): + """Pick up a cloth with a UR10 arm and a parallel gripper.""" - Returns: - Robot: The configured robot instance added to the simulation. - """ - gripper_urdf_path = get_data_path("DH_PGC_140_50_M/DH_PGC_140_50_M.urdf") - cfg = URRobotCfg.from_dict( - { - "robot_type": "ur10", - "uid": "UR10", - "urdf_cfg": { - "components": [ - {"component_type": "hand", "urdf_path": gripper_urdf_path}, - ] - }, - "drive_pros": { - "stiffness": {"FINGER[1-2]": 1e2}, - "damping": {"FINGER[1-2]": 1e1}, - "max_effort": {"FINGER[1-2]": 1e3}, - "drive_type": "force", - }, - "control_parts": { - "hand": ["FINGER[1-2]"], - }, - "solver_cfg": { - "arm": { - "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], - ] - } - }, - "init_qpos": [ - 0.0, - -np.pi / 2, - -np.pi / 2, - np.pi / 2, - -np.pi / 2, - 0.0, - 0.0, - 0.0, + def setup(self) -> None: + """Create the simulation, robot, cloth and padding box, then settle.""" + self.sim = create_default_sim( + self.args, + arena_space=5.0, + num_envs=self.args.num_envs, + add_default_light=False, + ) + self.robot = self._create_robot() + self.cloth = self._create_cloth() + self.padding_box = self._create_padding_box() + maybe_init_gpu_physics(self.sim) + maybe_open_window(self.sim, self.args) + # Let the cloth settle before interaction. + self.sim.update(step=10) + + def run(self) -> None: + """Plan and replay the cloth pick-up trajectory.""" + grasp_xpos = torch.tensor( + [ + [ + [-1, 0, 0, 0.5], + [0, 1, 0, 0], + [0, 0, -1, 0.075], + [0, 0, 0, 1], + ], ], - "init_pos": position, - } - ) - return sim.add_robot(cfg=cfg) - - -def create_padding_box(sim: SimulationManager): - padding_box_cfg = RigidObjectCfg( - uid="padding_box", - shape=CubeCfg( - size=[0.02, 0.07, 0.05], - ), - attrs=RigidBodyAttributesCfg( - mass=1.0, - static_friction=0.01, - dynamic_friction=0.00, - restitution=0.01, - min_position_iters=32, - min_velocity_iters=8, - ), - body_type="kinematic", - init_pos=[0.5, 0.0, 0.026], - init_rot=[0.0, 0.0, 0.0], - ) - padding_box = sim.add_rigid_object(cfg=padding_box_cfg) - return padding_box - - -def create_2d_grid_mesh(width: float, height: float, nx: int = 1, ny: int = 1): - """Create a flat rectangle in the XY plane centered at `origin`. - - The rectangle is subdivided into an `nx` by `ny` grid (cells) and - triangulated. `nx=1, ny=1` yields the simple two-triangle rectangle. - - Returns an vertices and triangles. - """ - w = float(width) - h = float(height) - if nx < 1 or ny < 1: - raise ValueError("nx and ny must be >= 1") - - # Vectorized vertex positions using PyTorch - x_lin = torch.linspace(-w / 2.0, w / 2.0, steps=nx + 1, dtype=torch.float64) - y_lin = torch.linspace(-h / 2.0, h / 2.0, steps=ny + 1, dtype=torch.float64) - yy, xx = torch.meshgrid(y_lin, x_lin) # shapes: (ny+1, nx+1) - xx_flat = xx.reshape(-1) - yy_flat = yy.reshape(-1) - zz_flat = torch.full_like(xx_flat, 0, dtype=torch.float64) - verts = torch.stack([xx_flat, yy_flat, zz_flat], dim=1) # (Nverts, 3) - - # Vectorized triangle indices - idx = torch.arange((nx + 1) * (ny + 1), dtype=torch.int64).reshape(ny + 1, nx + 1) - v0 = idx[:-1, :-1].reshape(-1) - v1 = idx[:-1, 1:].reshape(-1) - v2 = idx[1:, :-1].reshape(-1) - v3 = idx[1:, 1:].reshape(-1) - tri1 = torch.stack([v0, v1, v3], dim=1) - tri2 = torch.stack([v0, v3, v2], dim=1) - faces = torch.cat([tri1, tri2], dim=0).to(torch.int32) - return verts, faces - - -def create_cloth(sim: SimulationManager): - cloth_verts, cloth_faces = create_2d_grid_mesh(width=0.3, height=0.3, nx=12, ny=12) - cloth_mesh = o3d.geometry.TriangleMesh( - vertices=o3d.utility.Vector3dVector(cloth_verts.to("cpu").numpy()), - triangles=o3d.utility.Vector3iVector(cloth_faces.to("cpu").numpy()), - ) - cloth_save_path = os.path.join(tempfile.gettempdir(), "cloth_mesh.ply") - o3d.io.write_triangle_mesh(cloth_save_path, cloth_mesh) - - cloth = sim.add_cloth_object( - cfg=ClothObjectCfg( - uid="cloth", - shape=MeshCfg(fpath=cloth_save_path), - init_pos=[0.5, 0.0, 0.3], - init_rot=[0, 0, 0], - physical_attr=ClothPhysicalAttributesCfg( - mass=0.01, - youngs=1e10, - poissons=0.4, - thickness=0.06, - bending_stiffness=0.01, - bending_damping=0.1, - dynamic_friction=0.95, - min_position_iters=30, + dtype=torch.float32, + device=self.sim.device, + ) + grasp_xpos = grasp_xpos.repeat(self.sim.num_envs, 1, 1) + grab_traj = self._get_grasp_traj(grasp_xpos) + + maybe_wait_for_user(self.args, "Press Enter to start grabbing cloth...") + + with DemoRecording(self.sim, self.args, prefix="pick_up_cloth"): + for qpos in grab_traj.unbind(dim=1): + self.robot.set_qpos(qpos) + self.sim.update(step=3) + + maybe_wait_for_user(self.args, "Press Enter to exit the simulation...") + + def _create_robot(self, position=[0.0, 0.0, 0.0]) -> Robot: + """Create and configure a UR10 robot with a parallel gripper. + + Args: + position: Initial root position of the robot. + + Returns: + The configured robot instance added to the simulation. + """ + gripper_urdf_path = get_data_path("DH_PGC_140_50_M/DH_PGC_140_50_M.urdf") + cfg = URRobotCfg.from_dict( + { + "robot_type": "ur10", + "uid": "UR10", + "urdf_cfg": { + "components": [ + {"component_type": "hand", "urdf_path": gripper_urdf_path}, + ] + }, + "drive_pros": { + "stiffness": {"FINGER[1-2]": 1e2}, + "damping": {"FINGER[1-2]": 1e1}, + "max_effort": {"FINGER[1-2]": 1e3}, + "drive_type": "force", + }, + "control_parts": { + "hand": ["FINGER[1-2]"], + }, + "solver_cfg": { + "arm": { + "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], + ] + } + }, + "init_qpos": [ + 0.0, + -np.pi / 2, + -np.pi / 2, + np.pi / 2, + -np.pi / 2, + 0.0, + 0.0, + 0.0, + ], + "init_pos": position, + } + ) + return self.sim.add_robot(cfg=cfg) + + def _create_padding_box(self): + """Create a kinematic padding box used as a grasp reference.""" + padding_box_cfg = RigidObjectCfg( + uid="padding_box", + shape=CubeCfg( + size=[0.02, 0.07, 0.05], ), + attrs=RigidBodyAttributesCfg( + mass=1.0, + static_friction=0.01, + dynamic_friction=0.00, + restitution=0.01, + min_position_iters=32, + min_velocity_iters=8, + ), + body_type="kinematic", + init_pos=[0.5, 0.0, 0.026], + init_rot=[0.0, 0.0, 0.0], ) - ) - return cloth - + padding_box = self.sim.add_rigid_object(cfg=padding_box_cfg) + return padding_box + + def _create_2d_grid_mesh( + self, width: float, height: float, nx: int = 1, ny: int = 1 + ): + """Create a flat rectangle in the XY plane centered at the origin. + + The rectangle is subdivided into an ``nx`` by ``ny`` grid (cells) and + triangulated. ``nx=1, ny=1`` yields the simple two-triangle rectangle. + + Returns vertices and triangles. + """ + w = float(width) + h = float(height) + if nx < 1 or ny < 1: + raise ValueError("nx and ny must be >= 1") + + # Vectorized vertex positions using PyTorch + x_lin = torch.linspace(-w / 2.0, w / 2.0, steps=nx + 1, dtype=torch.float64) + y_lin = torch.linspace(-h / 2.0, h / 2.0, steps=ny + 1, dtype=torch.float64) + yy, xx = torch.meshgrid(y_lin, x_lin) # shapes: (ny+1, nx+1) + xx_flat = xx.reshape(-1) + yy_flat = yy.reshape(-1) + zz_flat = torch.full_like(xx_flat, 0, dtype=torch.float64) + verts = torch.stack([xx_flat, yy_flat, zz_flat], dim=1) # (Nverts, 3) + + # Vectorized triangle indices + idx = torch.arange((nx + 1) * (ny + 1), dtype=torch.int64).reshape( + ny + 1, nx + 1 + ) + v0 = idx[:-1, :-1].reshape(-1) + v1 = idx[:-1, 1:].reshape(-1) + v2 = idx[1:, :-1].reshape(-1) + v3 = idx[1:, 1:].reshape(-1) + tri1 = torch.stack([v0, v1, v3], dim=1) + tri2 = torch.stack([v0, v3, v2], dim=1) + faces = torch.cat([tri1, tri2], dim=0).to(torch.int32) + return verts, faces + + def _create_cloth(self): + """Create the cloth object from a generated grid mesh.""" + cloth_verts, cloth_faces = self._create_2d_grid_mesh( + width=0.3, height=0.3, nx=12, ny=12 + ) + cloth_mesh = o3d.geometry.TriangleMesh( + vertices=o3d.utility.Vector3dVector(cloth_verts.to("cpu").numpy()), + triangles=o3d.utility.Vector3iVector(cloth_faces.to("cpu").numpy()), + ) + cloth_save_path = os.path.join(tempfile.gettempdir(), "cloth_mesh.ply") + o3d.io.write_triangle_mesh(cloth_save_path, cloth_mesh) + + cloth = self.sim.add_cloth_object( + cfg=ClothObjectCfg( + uid="cloth", + shape=MeshCfg(fpath=cloth_save_path), + init_pos=[0.5, 0.0, 0.3], + init_rot=[0, 0, 0], + physical_attr=ClothPhysicalAttributesCfg( + mass=0.01, + youngs=1e10, + poissons=0.4, + thickness=0.06, + bending_stiffness=0.01, + bending_damping=0.1, + dynamic_friction=0.95, + min_position_iters=30, + ), + ), + ) + return cloth -def get_grasp_traj(sim: SimulationManager, robot: Robot, grasp_xpos: torch.Tensor): - n_envs = sim.num_envs - rest_arm_qpos = robot.get_qpos("arm") + def _get_grasp_traj(self, grasp_xpos: torch.Tensor): + """Compute the interpolated arm+hand trajectory for grasping the cloth.""" + n_envs = self.sim.num_envs + rest_arm_qpos = self.robot.get_qpos("arm") - approach_xpos = grasp_xpos.clone() - approach_xpos[:, 2, 3] += 0.04 - _, qpos_approach = robot.compute_ik( - pose=approach_xpos, joint_seed=rest_arm_qpos, name="arm" - ) - _, qpos_grasp = robot.compute_ik( - pose=grasp_xpos, joint_seed=qpos_approach, name="arm" - ) - hand_open_qpos = torch.tensor([0.00, 0.00], dtype=torch.float32, device=sim.device) - hand_close_qpos = torch.tensor( - [0.025, 0.025], dtype=torch.float32, device=sim.device - ) - - arm_trajectory = torch.cat( - [ - rest_arm_qpos[:, None, :], - qpos_approach[:, None, :], - qpos_grasp[:, None, :], - qpos_grasp[:, None, :], - qpos_approach[:, None, :], - rest_arm_qpos[:, None, :], - ], - dim=1, - ) - hand_trajectory = torch.cat( - [ - hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), - hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), - ], - dim=1, - ) - all_trajectory = torch.cat([arm_trajectory, hand_trajectory], dim=-1) - interp_trajectory = interpolate_with_distance( - trajectory=all_trajectory, interp_num=220, device=sim.device - ) - return interp_trajectory + approach_xpos = grasp_xpos.clone() + approach_xpos[:, 2, 3] += 0.04 + _, qpos_approach = self.robot.compute_ik( + pose=approach_xpos, joint_seed=rest_arm_qpos, name="arm" + ) + _, qpos_grasp = self.robot.compute_ik( + pose=grasp_xpos, joint_seed=qpos_approach, name="arm" + ) + hand_open_qpos = torch.tensor( + [0.00, 0.00], dtype=torch.float32, device=self.sim.device + ) + hand_close_qpos = torch.tensor( + [0.025, 0.025], dtype=torch.float32, device=self.sim.device + ) + arm_trajectory = torch.cat( + [ + rest_arm_qpos[:, None, :], + qpos_approach[:, None, :], + qpos_grasp[:, None, :], + qpos_grasp[:, None, :], + qpos_approach[:, None, :], + rest_arm_qpos[:, None, :], + ], + dim=1, + ) + hand_trajectory = torch.cat( + [ + hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_open_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), + hand_close_qpos[None, None, :].repeat(n_envs, 1, 1), + ], + dim=1, + ) + all_trajectory = torch.cat([arm_trajectory, hand_trajectory], dim=-1) + interp_trajectory = interpolate_with_distance( + trajectory=all_trajectory, interp_num=220, device=self.sim.device + ) + return interp_trajectory -def main(): - """ - Main function to demonstrate robot simulation. - This function initializes the simulation, creates the robot and other objects, - and performs the press softbody task. - """ +def main() -> None: + """Entry point for the pick-up-cloth demo.""" + setup_print_options() parser = argparse.ArgumentParser( description="Create a simulation scene with SimulationManager" ) - add_env_launcher_args_to_parser(parser) + parser = add_demo_args(parser) + # Cloth simulation requires GPU physics; default to CUDA. + parser.set_defaults(device="cuda") args = parser.parse_args() - # Configure the simulation - sim_cfg = SimulationManagerCfg( - width=1920, - height=1080, - num_envs=args.num_envs, - headless=True, - physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device="cuda", - render_cfg=RenderCfg( - renderer=args.renderer - ), # Enable ray tracing for better visuals - ) - - # Create the simulation instance - sim = SimulationManager(sim_cfg) - - robot = create_robot(sim) - cloth = create_cloth(sim) - padding_box = create_padding_box(sim) - sim.init_gpu_physics() - sim.open_window() - sim.update(step=10) # Let the cloth settle before interaction - - grasp_xpos = torch.tensor( - [ - [ - [-1, 0, 0, 0.5], - [0, 1, 0, 0], - [0, 0, -1, 0.075], - [0, 0, 0, 1], - ], - ], - dtype=torch.float32, - device=sim.device, - ) - grasp_xpos = grasp_xpos.repeat(sim.num_envs, 1, 1) - grab_traj = get_grasp_traj(sim, robot, grasp_xpos) - input("Press Enter to start grabing cloth...") - - n_waypoint = grab_traj.shape[1] - for i in range(n_waypoint): - robot.set_qpos(grab_traj[:, i, :]) - sim.update(step=3) - input("Press Enter to exit the simulation...") + PickUpClothDemo(args).main() if __name__ == "__main__": diff --git a/examples/sim/demo/press_softbody.py b/examples/sim/demo/press_softbody.py index 944a2de43..47795b991 100644 --- a/examples/sim/demo/press_softbody.py +++ b/examples/sim/demo/press_softbody.py @@ -19,182 +19,160 @@ and performs a pressing task in a simulated environment. """ +from __future__ import annotations + import argparse + import numpy as np -import time import torch from dexsim.utility.path import get_resources_data_path -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.objects import Robot, SoftObject -from embodichain.lab.sim.utility.action_utils import interpolate_with_distance -from embodichain.lab.sim.shapes import MeshCfg -from embodichain.data import get_data_path -from embodichain.utils import logger from embodichain.lab.sim.cfg import ( - RenderCfg, - LightCfg, SoftObjectCfg, - SoftbodyVoxelAttributesCfg, SoftbodyPhysicalAttributesCfg, + SoftbodyVoxelAttributesCfg, ) -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.demo_base import DemoBase +from embodichain.lab.sim.objects import Robot, SoftObject from embodichain.lab.sim.robots import URRobotCfg +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.utility.action_utils import interpolate_with_distance +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + resolve_demo_steps, + run_simulation_loop, + setup_print_options, +) +from embodichain.utils import logger -def parse_arguments(): - """ - Parse command-line arguments to configure the simulation. - - Returns: - argparse.Namespace: Parsed arguments including number of environments, device, and rendering options. - """ +class PressSoftbodyDemo(DemoBase): + """Press a soft cow with the end link of a UR10 robot.""" + + def setup(self) -> None: + """Create the simulation, robot and soft object, then open the viewer.""" + self.sim = create_default_sim( + self.args, + arena_space=5.0, + num_envs=self.args.num_envs, + add_default_light=False, + ) + self.robot = self._create_robot() + self.soft_cow = self._create_soft_cow() + maybe_init_gpu_physics(self.sim) + maybe_open_window(self.sim, self.args) + + def run(self) -> None: + """Press the cow and keep the simulation live until interrupted.""" + with DemoRecording(self.sim, self.args, prefix="press_softbody"): + self._press_cow() + logger.log_info("\n Press Ctrl+C to exit simulation loop.") + run_simulation_loop( + self.sim, + max_steps=resolve_demo_steps(self.args), + steps_per_update=10, + ) + + def _create_robot(self) -> Robot: + """Create and configure a UR10 robot in the simulation. + + Returns: + The configured robot instance added to the simulation. + """ + cfg = URRobotCfg.from_dict( + { + "robot_type": "ur10", + "uid": "UR10", + "solver_cfg": {"arm": {"tcp": np.eye(4)}}, + "init_qpos": [ + 0.0, + -np.pi / 2, + -np.pi / 2, + np.pi / 2, + -np.pi / 2, + 0.0, + ], + } + ) + return self.sim.add_robot(cfg=cfg) + + def _create_soft_cow(self) -> SoftObject: + """Create the soft cow object in the simulation. + + Returns: + The soft cow object. + """ + cow: SoftObject = self.sim.add_soft_object( + cfg=SoftObjectCfg( + uid="cow", + shape=MeshCfg( + fpath=get_resources_data_path("Model", "cow", "cow2.obj"), + ), + init_rot=[0, 90, 0], + init_pos=[0.45, -0.1, 0.12], + voxel_attr=SoftbodyVoxelAttributesCfg( + simulation_mesh_resolution=8, + maximal_edge_length=0.5, + ), + physical_attr=SoftbodyPhysicalAttributesCfg( + youngs=5e3, + poissons=0.45, + density=100, + dynamic_friction=0.1, + ), + ), + ) + return cow + + def _press_cow(self) -> None: + """Drive the robot end link to press the soft cow.""" + start_qpos = self.robot.get_qpos() + arm_ids = self.robot.get_joint_ids("arm") + arm_start_qpos = start_qpos[:, arm_ids] + + arm_start_xpos = self.robot.compute_fk( + arm_start_qpos, name="arm", to_matrix=True + ) + press_xpos = arm_start_xpos.clone() + press_xpos[:, :3, 3] = torch.tensor( + [0.5, -0.1, 0.005], device=press_xpos.device + ) + + approach_xpos = press_xpos.clone() + approach_xpos[:, 2, 3] += 0.05 + + is_success, approach_qpos = self.robot.compute_ik( + approach_xpos, joint_seed=arm_start_qpos, name="arm" + ) + + arm_trajectory = torch.concatenate([arm_start_qpos, approach_qpos]) + interp_trajectory = interpolate_with_distance( + trajectory=arm_trajectory[None, :, :], interp_num=50, device=self.sim.device + ) + interp_trajectory = interp_trajectory[0] + for qpos in interp_trajectory: + self.robot.set_qpos( + qpos.unsqueeze(0).repeat(self.sim.num_envs, 1), joint_ids=arm_ids + ) + self.sim.update(step=5) + + +def main() -> None: + """Entry point for the press-softbody demo.""" + setup_print_options() parser = argparse.ArgumentParser( description="Create and simulate a robot in SimulationManager" ) - add_env_launcher_args_to_parser(parser) - return parser.parse_args() - - -def initialize_simulation(args): - """ - Initialize the simulation environment based on the provided arguments. - - Args: - args (argparse.Namespace): Parsed command-line arguments. - - Returns: - SimulationManager: Configured simulation manager instance. - """ - config = SimulationManagerCfg( - headless=True, - sim_device="cuda", - render_cfg=RenderCfg(renderer=args.renderer), - physics_dt=1.0 / 100.0, - num_envs=args.num_envs, - ) - sim = SimulationManager(config) - - return sim - - -def create_robot(sim: SimulationManager): - """ - Create and configure a robot with an arm and a dexterous hand in the simulation. - - Args: - sim (SimulationManager): The simulation manager instance. - - Returns: - Robot: The configured robot instance added to the simulation. - """ - cfg = URRobotCfg.from_dict( - { - "robot_type": "ur10", - "uid": "UR10", - "solver_cfg": {"arm": {"tcp": np.eye(4)}}, - "init_qpos": [ - 0.0, - -np.pi / 2, - -np.pi / 2, - np.pi / 2, - -np.pi / 2, - 0.0, - ], - } - ) - return sim.add_robot(cfg=cfg) - - -def create_soft_cow(sim: SimulationManager) -> SoftObject: - """create soft cow object in the simulation - - Args: - sim (SimulationManager): The simulation manager instance. - - Returns: - SoftObject: soft cow object - """ - cow: SoftObject = sim.add_soft_object( - cfg=SoftObjectCfg( - uid="cow", - shape=MeshCfg( - fpath=get_resources_data_path("Model", "cow", "cow2.obj"), - ), - init_rot=[0, 90, 0], - init_pos=[0.45, -0.1, 0.12], - voxel_attr=SoftbodyVoxelAttributesCfg( - simulation_mesh_resolution=8, - maximal_edge_length=0.5, - ), - physical_attr=SoftbodyPhysicalAttributesCfg( - youngs=5e3, - poissons=0.45, - density=100, - dynamic_friction=0.1, - ), - ), - ) - return cow - - -def press_cow(sim: SimulationManager, robot: Robot): - """robot press cow softbody with its end link - - Args: - sim (SimulationManager): The simulation manager instance. - robot (Robot): The robot instance to be controlled. - """ - start_qpos = robot.get_qpos() - arm_ids = robot.get_joint_ids("arm") - arm_start_qpos = start_qpos[:, arm_ids] - - arm_start_xpos = robot.compute_fk(arm_start_qpos, name="arm", to_matrix=True) - press_xpos = arm_start_xpos.clone() - press_xpos[:, :3, 3] = torch.tensor([0.5, -0.1, 0.005], device=press_xpos.device) - - approach_xpos = press_xpos.clone() - approach_xpos[:, 2, 3] += 0.05 - - is_success, approach_qpos = robot.compute_ik( - approach_xpos, joint_seed=arm_start_qpos, name="arm" - ) - - arm_trajectory = torch.concatenate([arm_start_qpos, approach_qpos]) - interp_trajectory = interpolate_with_distance( - trajectory=arm_trajectory[None, :, :], interp_num=50, device=sim.device - ) - interp_trajectory = interp_trajectory[0] - for qpos in interp_trajectory: - robot.set_qpos(qpos.unsqueeze(0).repeat(sim.num_envs, 1), joint_ids=arm_ids) - sim.update(step=5) - - -def main(): - """ - Main function to demonstrate robot simulation. - - This function initializes the simulation, creates the robot and other objects, - and performs the press softbody task. - """ - args = parse_arguments() - sim = initialize_simulation(args) - - robot = create_robot(sim) - soft_cow = create_soft_cow(sim) - sim.init_gpu_physics() - sim.open_window() - - press_cow(sim, robot) - - logger.log_info("\n Press Ctrl+C to exit simulation loop.") - try: - while True: - sim.update(step=10) - except KeyboardInterrupt: - logger.log_info("\n Exit") + parser = add_demo_args(parser) + # Soft-body simulation requires GPU physics; default to CUDA. + parser.set_defaults(device="cuda") + args = parser.parse_args() + PressSoftbodyDemo(args).main() if __name__ == "__main__": diff --git a/examples/sim/demo/scoop_ice.py b/examples/sim/demo/scoop_ice.py index 64c6654a3..3fa796c2c 100644 --- a/examples/sim/demo/scoop_ice.py +++ b/examples/sim/demo/scoop_ice.py @@ -19,543 +19,514 @@ and performs a scoop ice task in a simulated environment. """ +from __future__ import annotations + import argparse import numpy as np -import time import torch -from tqdm import tqdm from scipy.spatial.transform import Rotation as R +from tqdm import tqdm -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.objects import Robot, RigidObject, RigidObjectGroup +from embodichain.data import get_data_path from embodichain.lab.sim.cfg import ( - RenderCfg, - RigidObjectCfg, - RigidBodyAttributesCfg, ArticulationCfg, - RigidObjectGroupCfg, JointDrivePropertiesCfg, LightCfg, + RigidBodyAttributesCfg, + RigidObjectCfg, + RigidObjectGroupCfg, ) +from embodichain.lab.sim.demo_base import DemoBase from embodichain.lab.sim.material import VisualMaterialCfg +from embodichain.lab.sim.objects import RigidObject, RigidObjectGroup, Robot +from embodichain.lab.sim.robots import URRobotCfg +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg from embodichain.lab.sim.utility.action_utils import interpolate_with_distance -from embodichain.lab.sim.shapes import MeshCfg, CubeCfg -from embodichain.data import get_data_path +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_open_window, + resolve_demo_steps, + run_simulation_loop, + setup_print_options, +) from embodichain.utils import logger -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.robots import URRobotCfg -def initialize_simulation(args): - """ - Initialize the simulation environment based on the provided arguments. - - Args: - args (argparse.Namespace): Parsed command-line arguments. - - Returns: - SimulationManager: Configured simulation manager instance. - """ - config = SimulationManagerCfg( - headless=True, - render_cfg=RenderCfg(renderer=args.renderer), - physics_dt=1.0 / 100.0, - ) - sim = SimulationManager(config) - - light = sim.add_light( - cfg=LightCfg(uid="main_light", intensity=10.0, init_pos=(0, 0, 2.0)) - ) - - return sim - - -def randomize_ice_positions(sim, ice_cubes): - """ - Randomly drop ice cubes into the container within a specified range. - - Args: - sim (SimulationManager): The simulation manager instance. - ice_cubes (RigidObjectGroup): Group of ice cube objects to be randomized. - """ - num_objs = ice_cubes.num_objects - position_low = np.array([0.65, -0.45, 0.5]) - position_high = np.array([0.55, -0.35, 0.5]) - position_random = np.random.uniform( - low=position_low, high=position_high, size=(num_objs, 3) - ) - random_drop_pose_np = np.eye(4)[None, :, :].repeat(num_objs, axis=0) - random_drop_pose_np[:, :3, 3] = position_random - - # Assign random positions to each ice cube - for i in tqdm(range(num_objs), desc="Dropping ice cubes"): - ice_cubes.set_local_pose( - pose=torch.tensor( - random_drop_pose_np[i][None, None, :, :], - dtype=torch.float32, - device=sim.device, +class ScoopIceDemo(DemoBase): + """Scoop ice cubes with a UR10 arm and a BrainCo dexterous hand.""" + + def setup(self) -> None: + """Create the simulation, robot, container, scoop and ice cubes.""" + self.sim = create_default_sim( + self.args, + arena_space=5.0, + num_envs=self.args.num_envs, + add_default_light=False, + ) + self.sim.add_light( + cfg=LightCfg(uid="main_light", intensity=10.0, init_pos=(0, 0, 2.0)) + ) + + # Create simulation objects. + self.robot = self._create_robot() + self.container = self._create_container() + self.padding_box = self._create_padding_box() + self.scoop = self._create_scoop() + self.heave_ice = self._create_heave_ice() + self.ice_cubes = self._create_ice_cubes() + + maybe_open_window(self.sim, self.args) + self._randomize_ice_positions() + + def run(self) -> None: + """Grasp the scoop, perform the scoop task, then keep the sim live.""" + with DemoRecording(self.sim, self.args, prefix="scoop_ice"): + self._scoop_grasp() + self._scoop_ice() + logger.log_info("\n Press Ctrl+C to exit simulation loop.") + run_simulation_loop( + self.sim, + max_steps=resolve_demo_steps(self.args), + sleep=1e-2, + ) + + def _randomize_ice_positions(self) -> None: + """Randomly drop ice cubes into the container within a specified range.""" + ice_cubes = self.ice_cubes + num_objs = ice_cubes.num_objects + position_low = np.array([0.65, -0.45, 0.5]) + position_high = np.array([0.55, -0.35, 0.5]) + position_random = np.random.uniform( + low=position_low, high=position_high, size=(num_objs, 3) + ) + random_drop_pose_np = np.eye(4)[None, :, :].repeat(num_objs, axis=0) + random_drop_pose_np[:, :3, 3] = position_random + + # Assign random positions to each ice cube. + for i in tqdm(range(num_objs), desc="Dropping ice cubes"): + ice_cubes.set_local_pose( + pose=torch.tensor( + random_drop_pose_np[i][None, None, :, :], + dtype=torch.float32, + device=self.sim.device, + ), + obj_ids=[i], + ) + self.sim.update(step=10) + + def _create_robot(self) -> Robot: + """Create and configure a UR10 robot with a BrainCo dexterous hand. + + Returns: + The configured robot instance added to the simulation. + """ + hand_urdf_path = get_data_path( + "BrainCoHandRevo1/BrainCoLeftHand/BrainCoLeftHand.urdf" + ) + + # Define transformation for attaching the hand to the arm. + hand_attach_xpos = np.eye(4) + hand_attach_xpos[:3, :3] = R.from_rotvec([90, 0, 0], degrees=True).as_matrix() + + cfg = URRobotCfg.from_dict( + { + "robot_type": "ur10", + "uid": "ur10_with_brainco", + "urdf_cfg": { + "components": [ + { + "component_type": "hand", + "urdf_path": hand_urdf_path, + "transform": hand_attach_xpos, + } + ] + }, + "control_parts": { + "hand": [ + "LEFT_HAND_THUMB1", + "LEFT_HAND_THUMB2", + "LEFT_HAND_INDEX", + "LEFT_HAND_MIDDLE", + "LEFT_HAND_RING", + "LEFT_HAND_PINKY", + ], + }, + "drive_pros": { + "stiffness": {"LEFT_[A-Z|_]+[0-9]?": 1e2}, + "damping": {"LEFT_[A-Z|_]+[0-9]?": 1e1}, + "max_effort": {"LEFT_[A-Z|_]+[0-9]?": 1e3}, + "drive_type": "force", + }, + "solver_cfg": {"arm": {"tcp": np.eye(4)}}, + "init_qpos": [ + 0.0, + -np.pi / 2, + -np.pi / 2, + 2.5, + -np.pi / 2, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.5, + -0.00016, + -0.00010, + -0.00013, + -0.00009, + 0.0, + ], + } + ) + + return self.sim.add_robot(cfg=cfg) + + def _create_scoop(self) -> RigidObject: + """Create the scoop rigid object.""" + scoop_cfg = RigidObjectCfg( + uid="scoop", + shape=MeshCfg( + fpath=get_data_path("ScoopIceNewEnv/scoop.ply"), + ), + attrs=RigidBodyAttributesCfg( + mass=0.5, + static_friction=0.95, + dynamic_friction=0.9, + restitution=0.01, + min_position_iters=32, + min_velocity_iters=8, + ), + max_convex_hull_num=12, + body_type="dynamic", + init_pos=[0.6, 0.0, 0.09], + init_rot=[0.0, 0.0, 0.0], + ) + scoop = self.sim.add_rigid_object(cfg=scoop_cfg) + return scoop + + def _create_heave_ice(self) -> RigidObject: + """Create the heave ice rigid object.""" + heave_ice_cfg = RigidObjectCfg( + uid="heave_ice", + shape=MeshCfg( + fpath=get_data_path("ScoopIceNewEnv/ice_mesh_small/ice_000.obj"), + ), + attrs=RigidBodyAttributesCfg( + mass=0.5, + static_friction=0.95, + dynamic_friction=0.9, + restitution=0.01, + min_position_iters=32, + min_velocity_iters=8, + ), + body_type="dynamic", + init_pos=[10, 10, 0.08], + init_rot=[0.0, 0.0, 0.0], + ) + heave_ice = self.sim.add_rigid_object(cfg=heave_ice_cfg) + return heave_ice + + def _create_padding_box(self) -> RigidObject: + """Create the padding box rigid object.""" + padding_box_cfg = RigidObjectCfg( + uid="padding_box", + shape=CubeCfg( + size=[0.1, 0.16, 0.05], + ), + attrs=RigidBodyAttributesCfg( + mass=1.0, + static_friction=0.95, + dynamic_friction=0.9, + restitution=0.01, + min_position_iters=32, + min_velocity_iters=8, ), - obj_ids=[i], + body_type="kinematic", + init_pos=[0.6, 0.15, 0.025], + init_rot=[0.0, 0.0, 0.0], ) - sim.update(step=10) - - -def create_robot(sim): - """ - Create and configure a robot with an arm and a dexterous hand in the simulation. - - Args: - sim (SimulationManager): The simulation manager instance. - - Returns: - Robot: The configured robot instance added to the simulation. - """ - hand_urdf_path = get_data_path( - "BrainCoHandRevo1/BrainCoLeftHand/BrainCoLeftHand.urdf" - ) - - # Define transformation for attaching the hand to the arm - hand_attach_xpos = np.eye(4) - hand_attach_xpos[:3, :3] = R.from_rotvec([90, 0, 0], degrees=True).as_matrix() - - cfg = URRobotCfg.from_dict( - { - "robot_type": "ur10", - "uid": "ur10_with_brainco", - "urdf_cfg": { - "components": [ - { - "component_type": "hand", - "urdf_path": hand_urdf_path, - "transform": hand_attach_xpos, + heave_ice = self.sim.add_rigid_object(cfg=padding_box_cfg) + return heave_ice + + def _create_container(self) -> Robot: + """Create the container articulation.""" + container_cfg = ArticulationCfg( + uid="container", + fpath=get_data_path("ScoopIceNewEnv/IceContainer/ice_container.urdf"), + init_pos=[0.7, -0.4, 0.21], + init_rot=[0, 0, -90], + attrs=RigidBodyAttributesCfg( + mass=1.0, + static_friction=0.95, + dynamic_friction=0.9, + restitution=0.01, + min_position_iters=32, + min_velocity_iters=8, + ), + drive_pros=JointDrivePropertiesCfg( + stiffness=1.0, damping=0.1, max_effort=100.0, drive_type="force" + ), + ) + container = self.sim.add_articulation(cfg=container_cfg) + return container + + def _create_ice_cubes(self) -> RigidObjectGroup: + """Create the group of ice cube rigid objects.""" + ice_cubes_path = get_data_path("ScoopIceNewEnv/ice_mesh_small") + cfg_dict = { + "uid": "ice_cubes", + "max_num": 300, + "folder_path": ice_cubes_path, + "ext": ".obj", + "rigid_objects": { + "obj": { + "attrs": { + "mass": 0.003, + "contact_offset": 0.001, + "rest_offset": 0, + "dynamic_friction": 0.05, + "static_friction": 0.1, + "restitution": 0.01, + "min_position_iters": 32, + "min_velocity_iters": 4, + "max_depenetration_velocity": 1.0, }, - ] - }, - "control_parts": { - "hand": [ - "LEFT_HAND_THUMB1", - "LEFT_HAND_THUMB2", - "LEFT_HAND_INDEX", - "LEFT_HAND_MIDDLE", - "LEFT_HAND_RING", - "LEFT_HAND_PINKY", - ], + "shape": {"shape_type": "Mesh"}, + "init_pos": [20.0, 0, 1.0], + } }, - "drive_pros": { - "stiffness": {"LEFT_[A-Z|_]+[0-9]?": 1e2}, - "damping": {"LEFT_[A-Z|_]+[0-9]?": 1e1}, - "max_effort": {"LEFT_[A-Z|_]+[0-9]?": 1e3}, - "drive_type": "force", - }, - "solver_cfg": {"arm": {"tcp": np.eye(4)}}, - "init_qpos": [ - 0.0, - -np.pi / 2, - -np.pi / 2, - 2.5, - -np.pi / 2, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 1.5, - -0.00016, - -0.00010, - -0.00013, - -0.00009, - 0.0, - ], } - ) - - return sim.add_robot(cfg=cfg) - - -def create_scoop(sim: SimulationManager): - """Create a scoop rigid object in the simulation.""" - scoop_cfg = RigidObjectCfg( - uid="scoop", - shape=MeshCfg( - fpath=get_data_path("ScoopIceNewEnv/scoop.ply"), - ), - attrs=RigidBodyAttributesCfg( - mass=0.5, - static_friction=0.95, - dynamic_friction=0.9, - restitution=0.01, - min_position_iters=32, - min_velocity_iters=8, - ), - max_convex_hull_num=12, - body_type="dynamic", - init_pos=[0.6, 0.0, 0.09], - init_rot=[0.0, 0.0, 0.0], - ) - scoop = sim.add_rigid_object(cfg=scoop_cfg) - return scoop - - -def create_heave_ice(sim: SimulationManager): - """Create a heave ice rigid object in the simulation. Make sure that""" - heave_ice_cfg = RigidObjectCfg( - uid="heave_ice", - shape=MeshCfg( - fpath=get_data_path("ScoopIceNewEnv/ice_mesh_small/ice_000.obj"), - ), - attrs=RigidBodyAttributesCfg( - mass=0.5, - static_friction=0.95, - dynamic_friction=0.9, - restitution=0.01, - min_position_iters=32, - min_velocity_iters=8, - ), - body_type="dynamic", - init_pos=[10, 10, 0.08], - init_rot=[0.0, 0.0, 0.0], - ) - heave_ice = sim.add_rigid_object(cfg=heave_ice_cfg) - return heave_ice - - -def create_padding_box(sim: SimulationManager): - padding_box_cfg = RigidObjectCfg( - uid="padding_box", - shape=CubeCfg( - size=[0.1, 0.16, 0.05], - ), - attrs=RigidBodyAttributesCfg( - mass=1.0, - static_friction=0.95, - dynamic_friction=0.9, - restitution=0.01, - min_position_iters=32, - min_velocity_iters=8, - ), - body_type="kinematic", - init_pos=[0.6, 0.15, 0.025], - init_rot=[0.0, 0.0, 0.0], - ) - heave_ice = sim.add_rigid_object(cfg=padding_box_cfg) - return heave_ice - - -def create_container(sim: SimulationManager): - container_cfg = ArticulationCfg( - uid="container", - fpath=get_data_path("ScoopIceNewEnv/IceContainer/ice_container.urdf"), - init_pos=[0.7, -0.4, 0.21], - init_rot=[0, 0, -90], - attrs=RigidBodyAttributesCfg( - mass=1.0, - static_friction=0.95, - dynamic_friction=0.9, - restitution=0.01, - min_position_iters=32, - min_velocity_iters=8, - ), - drive_pros=JointDrivePropertiesCfg( - stiffness=1.0, damping=0.1, max_effort=100.0, drive_type="force" - ), - ) - container = sim.add_articulation(cfg=container_cfg) - return container - - -def create_ice_cubes(sim: SimulationManager): - ice_cubes_path = get_data_path("ScoopIceNewEnv/ice_mesh_small") - cfg_dict = { - "uid": "ice_cubes", - "max_num": 300, - "folder_path": ice_cubes_path, - "ext": ".obj", - "rigid_objects": { - "obj": { - "attrs": { - "mass": 0.003, - "contact_offset": 0.001, - "rest_offset": 0, - "dynamic_friction": 0.05, - "static_friction": 0.1, - "restitution": 0.01, - "min_position_iters": 32, - "min_velocity_iters": 4, - "max_depenetration_velocity": 1.0, - }, - "shape": {"shape_type": "Mesh"}, - "init_pos": [20.0, 0, 1.0], - } - }, - } - - ice_cubes_cfg = RigidObjectGroupCfg.from_dict(cfg_dict) - ice_cubes: RigidObjectGroup = sim.add_rigid_object_group(cfg=ice_cubes_cfg) - - # Set visual material for ice cubes. - # The material below only works for ray tracing backend. - # Set ior to 1.31 and material type to "BSDF" for better ice appearance. - ice_mat = sim.create_visual_material( - cfg=VisualMaterialCfg( - base_color=[1.0, 1.0, 1.0, 1.0], - ior=1.31, - roughness=0.2, - material_type="BSDF", + + ice_cubes_cfg = RigidObjectGroupCfg.from_dict(cfg_dict) + ice_cubes: RigidObjectGroup = self.sim.add_rigid_object_group(cfg=ice_cubes_cfg) + + # Set visual material for ice cubes. + # The material below only works for the ray tracing backend. + # Set ior to 1.31 and material type to "BSDF" for better ice appearance. + ice_mat = self.sim.create_visual_material( + cfg=VisualMaterialCfg( + base_color=[1.0, 1.0, 1.0, 1.0], + ior=1.31, + roughness=0.2, + material_type="BSDF", + ) + ) + ice_cubes.set_visual_material(mat=ice_mat) + + return ice_cubes + + def _scoop_grasp(self) -> None: + """Grasp the scoop and position the heave ice for scooping.""" + sim = self.sim + robot = self.robot + scoop = self.scoop + heave_ice = self.heave_ice + padding_box = self.padding_box + + rest_qpos = robot.get_qpos() + arm_ids = robot.get_joint_ids("arm") + hand_ids = robot.get_joint_ids("hand") + hand_open_qpos = torch.tensor([0.0, 1.5, 0.4, 0.4, 0.4, 0.4]) + hand_close_qpos = torch.tensor([0.4, 1.5, 1.0, 1.1, 1.1, 0.9]) + arm_rest_qpos = rest_qpos[:, arm_ids] + + # Calculate and set the drop pose for the scoop object. + padding_box_pose = padding_box.get_local_pose(to_matrix=True) + scoop_drop_relative_pose = torch.tensor( + [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.115], + [0.0, 0.0, 1.0, 0.065], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=torch.float32, + device=sim.device, + ) + scoop_drop_pose = torch.bmm( + padding_box_pose, + scoop_drop_relative_pose[None, :, :].repeat(sim.num_envs, 1, 1), + ) + scoop.set_local_pose(scoop_drop_pose) + + scoop_pose = scoop.get_local_pose(to_matrix=True) + + # Tricky implementation. + heave_ice_relative = torch.tensor( + [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, -0.13], + [0.0, 0.0, 1.0, 0.04], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=torch.float32, + device=sim.device, + )[None, :, :].repeat(sim.num_envs, 1, 1) + heave_ice_pose = torch.bmm(scoop_pose, heave_ice_relative) + heave_ice.set_local_pose(heave_ice_pose) + sim.update(step=200) + + # Move hand to grasp scoop. + scoop_pose = scoop.get_local_pose(to_matrix=True) + grasp_scoop_pose_relative = torch.tensor( + [ + [0.00522967, 0.6788424, 0.7342653, -0.05885637], + [0.99054945, 0.0971214, -0.09684561, 0.0301468], + [-0.13705578, 0.72783256, -0.6719191, 0.1040391], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=torch.float32, + device=sim.device, + )[None, :, :].repeat(sim.num_envs, 1, 1) + + grasp_scoop_pose = torch.bmm(scoop_pose, grasp_scoop_pose_relative) + pregrasp_scoop_pose = grasp_scoop_pose.clone() + pregrasp_scoop_pose[:, 2, 3] += 0.1 + is_success, pre_grasp_scoop_qpos = robot.compute_ik( + pregrasp_scoop_pose, joint_seed=arm_rest_qpos, name="arm" + ) + + is_success, grasp_scoop_qpos = robot.compute_ik( + grasp_scoop_pose, joint_seed=arm_rest_qpos, name="arm" + ) + robot.set_qpos(pre_grasp_scoop_qpos, joint_ids=arm_ids) + sim.update(step=100) + robot.set_qpos(grasp_scoop_qpos, joint_ids=arm_ids) + sim.update(step=100) + + # Close hand. + robot.set_qpos( + hand_close_qpos[None, :].repeat(sim.num_envs, 1), joint_ids=hand_ids + ) + sim.update(step=100) + + # Remove heave ice. + remove_heave_ice_pose = torch.tensor( + [ + [1.0, 0.0, 0.0, 10.0], + [0.0, 1.0, 0.0, 10.0], + [0.0, 0.0, 1.0, 0.04], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=torch.float32, + device=sim.device, + ) + heave_ice.set_local_pose(remove_heave_ice_pose[None, :, :]) + + def _scoop_ice(self) -> None: + """Lift, scoop and place the ice with the grasped scoop.""" + sim = self.sim + robot = self.robot + scoop = self.scoop + + start_qpos = robot.get_qpos() + arm_ids = robot.get_joint_ids("arm") + hand_ids = robot.get_joint_ids("hand") + hand_open_qpos = torch.tensor([0.0, 1.5, 0.4, 0.4, 0.4, 0.4]) + hand_close_qpos = torch.tensor([0.4, 1.5, 1.0, 1.1, 1.1, 0.9]) + arm_start_qpos = start_qpos[:, arm_ids] + + # Lift. + arm_start_xpos = robot.compute_fk(arm_start_qpos, name="arm", to_matrix=True) + arm_lift_xpos = arm_start_xpos.clone() + arm_lift_xpos[:, 2, 3] += 0.45 + is_success, arm_lift_qpos = robot.compute_ik( + arm_lift_xpos, joint_seed=arm_start_qpos, name="arm" + ) + + # Apply 45 degree wrist rotation. + wrist_rotation = R.from_euler("X", 45, degrees=True).as_matrix() + arm_lift_rotation = arm_lift_xpos[0, :3, :3].to("cpu").numpy() + new_rotation = wrist_rotation @ arm_lift_rotation + arm_lift_xpos_rotated = arm_lift_xpos.clone() + arm_lift_xpos_rotated[:, :3, :3] = torch.tensor( + new_rotation, dtype=torch.float32, device=sim.device + ) + arm_lift_xpos_rotated[:, :3, 3] = torch.tensor( + [0.5, -0.2, 0.55], dtype=torch.float32, device=sim.device + ) + is_success, arm_lift_qpos_rotated = robot.compute_ik( + arm_lift_xpos_rotated, joint_seed=arm_lift_qpos, name="arm" + ) + + # Into container. + scoop_dis = 0.252 + scoop_offset = scoop_dis * torch.tensor( + [0.0, -0.58123819, -0.81373347], dtype=torch.float32, device=sim.device + ) + arm_into_container_xpos = arm_lift_xpos_rotated.clone() + arm_into_container_xpos[:, :3, 3] = ( + arm_into_container_xpos[:, :3, 3] + scoop_offset + ) + is_success, arm_into_container_qpos = robot.compute_ik( + arm_into_container_xpos, joint_seed=arm_lift_qpos_rotated, name="arm" + ) + + # Apply -60 degree wrist rotation. + arm_into_container_rotation = ( + arm_into_container_xpos[0, :3, :3].to("cpu").numpy() + ) + wrist_rotation = R.from_euler("X", -60, degrees=True).as_matrix() + new_rotation = wrist_rotation @ arm_into_container_rotation + arm_scoop_xpos = arm_into_container_xpos.clone() + arm_scoop_xpos[:, :3, :3] = torch.tensor( + new_rotation, dtype=torch.float32, device=sim.device ) - ) - ice_cubes.set_visual_material(mat=ice_mat) - - return ice_cubes - - -def scoop_grasp( - sim: SimulationManager, - robot: Robot, - scoop: RigidObject, - heave_ice: RigidObject, - padding_box: RigidObject, -): - """ - Control the robot to grasp the scoop object and position the heave ice for scooping. - - Args: - sim (SimulationManager): The simulation manager instance. - robot (Robot): The robot instance to be controlled. - scoop (RigidObject): The scoop object to be grasped. - heave_ice (RigidObject): The heave ice object to be positioned. - padding_box (RigidObject): The padding box object used as a reference for positioning. - """ - rest_qpos = robot.get_qpos() - arm_ids = robot.get_joint_ids("arm") - hand_ids = robot.get_joint_ids("hand") - hand_open_qpos = torch.tensor([0.0, 1.5, 0.4, 0.4, 0.4, 0.4]) - hand_close_qpos = torch.tensor([0.4, 1.5, 1.0, 1.1, 1.1, 0.9]) - arm_rest_qpos = rest_qpos[:, arm_ids] - - # Calculate and set the drop pose for the scoop object - padding_box_pose = padding_box.get_local_pose(to_matrix=True) - scoop_drop_relative_pose = torch.tensor( - [ - [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, 0.115], - [0.0, 0.0, 1.0, 0.065], - [0.0, 0.0, 0.0, 1.0], - ], - dtype=torch.float32, - device=sim.device, - ) - scoop_drop_pose = torch.bmm( - padding_box_pose, - scoop_drop_relative_pose[None, :, :].repeat(sim.num_envs, 1, 1), - ) - scoop.set_local_pose(scoop_drop_pose) - - scoop_pose = scoop.get_local_pose(to_matrix=True) - - # tricky implementation - heave_ice_relative = torch.tensor( - [ - [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, -0.13], - [0.0, 0.0, 1.0, 0.04], - [0.0, 0.0, 0.0, 1.0], - ], - dtype=torch.float32, - device=sim.device, - )[None, :, :].repeat(sim.num_envs, 1, 1) - heave_ice_pose = torch.bmm(scoop_pose, heave_ice_relative) - heave_ice.set_local_pose(heave_ice_pose) - sim.update(step=200) - - # move hand to grasp scoop - scoop_pose = scoop.get_local_pose(to_matrix=True) - grasp_scoop_pose_relative = torch.tensor( - [ - [0.00522967, 0.6788424, 0.7342653, -0.05885637], - [0.99054945, 0.0971214, -0.09684561, 0.0301468], - [-0.13705578, 0.72783256, -0.6719191, 0.1040391], - [0.0, 0.0, 0.0, 1.0], - ], - dtype=torch.float32, - device=sim.device, - )[None, :, :].repeat(sim.num_envs, 1, 1) - - grasp_scoop_pose = torch.bmm(scoop_pose, grasp_scoop_pose_relative) - pregrasp_scoop_pose = grasp_scoop_pose.clone() - pregrasp_scoop_pose[:, 2, 3] += 0.1 - is_success, pre_grasp_scoop_qpos = robot.compute_ik( - pregrasp_scoop_pose, joint_seed=arm_rest_qpos, name="arm" - ) - - is_success, grasp_scoop_qpos = robot.compute_ik( - grasp_scoop_pose, joint_seed=arm_rest_qpos, name="arm" - ) - robot.set_qpos(pre_grasp_scoop_qpos, joint_ids=arm_ids) - sim.update(step=100) - robot.set_qpos(grasp_scoop_qpos, joint_ids=arm_ids) - sim.update(step=100) - - # close hand - robot.set_qpos(hand_close_qpos[None, :].repeat(sim.num_envs, 1), joint_ids=hand_ids) - sim.update(step=100) - - # remove heave ice - remove_heave_ice_pose = torch.tensor( - [ - [1.0, 0.0, 0.0, 10.0], - [0.0, 1.0, 0.0, 10.0], - [0.0, 0.0, 1.0, 0.04], - [0.0, 0.0, 0.0, 1.0], - ], - dtype=torch.float32, - device=sim.device, - ) - heave_ice.set_local_pose(remove_heave_ice_pose[None, :, :]) - - -def scoop_ice(sim: SimulationManager, robot: Robot, scoop: RigidObject): - """ - Control the robot to perform the scoop ice task, including lifting, scooping, - and placing the ice. - - Args: - sim (SimulationManager): The simulation manager instance. - robot (Robot): The robot instance to be controlled. - scoop (RigidObject): The scoop object used for scooping ice. - """ - start_qpos = robot.get_qpos() - arm_ids = robot.get_joint_ids("arm") - hand_ids = robot.get_joint_ids("hand") - hand_open_qpos = torch.tensor([0.0, 1.5, 0.4, 0.4, 0.4, 0.4]) - hand_close_qpos = torch.tensor([0.4, 1.5, 1.0, 1.1, 1.1, 0.9]) - arm_start_qpos = start_qpos[:, arm_ids] - - # lift - arm_start_xpos = robot.compute_fk(arm_start_qpos, name="arm", to_matrix=True) - arm_lift_xpos = arm_start_xpos.clone() - arm_lift_xpos[:, 2, 3] += 0.45 - is_success, arm_lift_qpos = robot.compute_ik( - arm_lift_xpos, joint_seed=arm_start_qpos, name="arm" - ) - - # apply 45 degree wrist rotation - wrist_rotation = R.from_euler("X", 45, degrees=True).as_matrix() - arm_lift_rotation = arm_lift_xpos[0, :3, :3].to("cpu").numpy() - new_rotation = wrist_rotation @ arm_lift_rotation - arm_lift_xpos_rotated = arm_lift_xpos.clone() - arm_lift_xpos_rotated[:, :3, :3] = torch.tensor( - new_rotation, dtype=torch.float32, device=sim.device - ) - arm_lift_xpos_rotated[:, :3, 3] = torch.tensor( - [0.5, -0.2, 0.55], dtype=torch.float32, device=sim.device - ) - is_success, arm_lift_qpos_rotated = robot.compute_ik( - arm_lift_xpos_rotated, joint_seed=arm_lift_qpos, name="arm" - ) - - # into container - scoop_dis = 0.252 - scoop_offset = scoop_dis * torch.tensor( - [0.0, -0.58123819, -0.81373347], dtype=torch.float32, device=sim.device - ) - arm_into_container_xpos = arm_lift_xpos_rotated.clone() - arm_into_container_xpos[:, :3, 3] = arm_into_container_xpos[:, :3, 3] + scoop_offset - is_success, arm_into_container_qpos = robot.compute_ik( - arm_into_container_xpos, joint_seed=arm_lift_qpos_rotated, name="arm" - ) - - # apply -60 degree wrist rotation - arm_into_container_rotation = arm_into_container_xpos[0, :3, :3].to("cpu").numpy() - wrist_rotation = R.from_euler("X", -60, degrees=True).as_matrix() - new_rotation = wrist_rotation @ arm_into_container_rotation - arm_scoop_xpos = arm_into_container_xpos.clone() - arm_scoop_xpos[:, :3, :3] = torch.tensor( - new_rotation, dtype=torch.float32, device=sim.device - ) - is_success, arm_scoop_qpos = robot.compute_ik( - arm_scoop_xpos, joint_seed=arm_into_container_qpos, name="arm" - ) - - # minor lift - arm_scoop_xpos[:, 2, 3] += 0.15 - is_success, arm_scoop_lift_qpos = robot.compute_ik( - arm_scoop_xpos, joint_seed=arm_scoop_qpos, name="arm" - ) - - # pack arm and hand trajectory - arm_trajectory = torch.concatenate( - [ - arm_start_qpos, - arm_lift_qpos, - arm_lift_qpos_rotated, - arm_into_container_qpos, - arm_scoop_qpos, - arm_scoop_lift_qpos, - ] - ) - - hand_trajectory = torch.vstack( - [ - hand_close_qpos, - hand_close_qpos, - hand_close_qpos, - hand_close_qpos, - hand_close_qpos, - hand_close_qpos, - ] - ) - - all_trajectory = torch.hstack([arm_trajectory, hand_trajectory]) - interp_trajectory = interpolate_with_distance( - trajectory=all_trajectory[None, :, :], interp_num=200, device=sim.device - ) - interp_trajectory = interp_trajectory[0] - # run trajectory - arm_ids = robot.get_joint_ids("arm") - hand_ids = robot.get_joint_ids("hand") - combine_ids = np.concatenate([arm_ids, hand_ids]) - for qpos in interp_trajectory: - robot.set_qpos(qpos.unsqueeze(0), joint_ids=combine_ids) - sim.update(step=10) - - -def main(): + is_success, arm_scoop_qpos = robot.compute_ik( + arm_scoop_xpos, joint_seed=arm_into_container_qpos, name="arm" + ) + + # Minor lift. + arm_scoop_xpos[:, 2, 3] += 0.15 + is_success, arm_scoop_lift_qpos = robot.compute_ik( + arm_scoop_xpos, joint_seed=arm_scoop_qpos, name="arm" + ) + + # Pack arm and hand trajectory. + arm_trajectory = torch.concatenate( + [ + arm_start_qpos, + arm_lift_qpos, + arm_lift_qpos_rotated, + arm_into_container_qpos, + arm_scoop_qpos, + arm_scoop_lift_qpos, + ] + ) + + hand_trajectory = torch.vstack( + [ + hand_close_qpos, + hand_close_qpos, + hand_close_qpos, + hand_close_qpos, + hand_close_qpos, + hand_close_qpos, + ] + ) + + all_trajectory = torch.hstack([arm_trajectory, hand_trajectory]) + interp_trajectory = interpolate_with_distance( + trajectory=all_trajectory[None, :, :], interp_num=200, device=sim.device + ) + interp_trajectory = interp_trajectory[0] + # Run trajectory. + arm_ids = robot.get_joint_ids("arm") + hand_ids = robot.get_joint_ids("hand") + combine_ids = np.concatenate([arm_ids, hand_ids]) + for qpos in interp_trajectory: + robot.set_qpos(qpos.unsqueeze(0), joint_ids=combine_ids) + sim.update(step=10) + + +def main() -> None: + """Entry point for the scoop-ice demo.""" + setup_print_options() parser = argparse.ArgumentParser(description="Scoop ice task simulation") - add_env_launcher_args_to_parser(parser) + parser = add_demo_args(parser) args = parser.parse_args() - - """ - Main function to demonstrate robot simulation. - - This function initializes the simulation, creates the robot and other objects, - and performs the scoop ice task. - """ - sim = initialize_simulation(args) - - # Create simulation objects - robot = create_robot(sim) - container = create_container(sim) - padding_box = create_padding_box(sim) - scoop = create_scoop(sim) - heave_ice = create_heave_ice(sim) - ice_cubes = create_ice_cubes(sim) - - sim.open_window() - - # Randomize ice positions - randomize_ice_positions(sim, ice_cubes) - - # Perform tasks - scoop_grasp(sim, robot, scoop, heave_ice, padding_box) - scoop_ice(sim, robot, scoop) - - logger.log_info("\n Press Ctrl+C to exit simulation loop.") - try: - while True: - # sim.update(step=10) - time.sleep(1e-2) - except KeyboardInterrupt: - logger.log_info("\n Exit") + ScoopIceDemo(args).main() if __name__ == "__main__": diff --git a/examples/sim/gizmo/gizmo_camera.py b/examples/sim/gizmo/gizmo_camera.py index 296c3be47..5921e937e 100644 --- a/examples/sim/gizmo/gizmo_camera.py +++ b/examples/sim/gizmo/gizmo_camera.py @@ -18,43 +18,44 @@ It shows how to create a gizmo attached to a camera for real-time pose manipulation. """ +from __future__ import annotations + import argparse -import cv2 -import numpy as np import time -import torch -torch.set_printoptions(precision=4, sci_mode=False) +import cv2 -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.sensors import Camera, CameraCfg -from embodichain.lab.sim.cfg import RigidObjectCfg, RigidBodyAttributesCfg, RenderCfg +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + resolve_demo_steps, + shutdown_sim, +) from embodichain.utils import logger -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -def main(): +def main() -> None: """Main function to demonstrate camera gizmo manipulation.""" # Parse command line arguments parser = argparse.ArgumentParser( description="Create and simulate a camera with gizmo in SimulationManager" ) - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) args = parser.parse_args() - # Configure the simulation - sim_cfg = SimulationManagerCfg( + sim = create_default_sim( + args, width=1920, height=1080, physics_dt=1.0 / 100.0, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), + add_default_light=False, ) - - # Create simulation context - sim = SimulationManager(sim_cfg) sim.set_manual_update(False) # Add some objects to the scene for camera to observe @@ -96,15 +97,13 @@ def main(): # Wait for initialization time.sleep(0.2) - # Enable gizmo for interactive camera control using the new unified API - sim.enable_gizmo(uid="gizmo_camera") - if not sim.has_gizmo("gizmo_camera"): - logger.log_error("Failed to enable gizmo for camera!") - return - - # Open simulation window (if not headless) if not args.headless: - sim.open_window() + sim.enable_gizmo(uid="gizmo_camera") + if not sim.has_gizmo("gizmo_camera"): + shutdown_sim(sim) + raise RuntimeError("Failed to enable gizmo for camera.") + + maybe_open_window(sim, args) logger.log_info("Gizmo-Camera tutorial started!") logger.log_info( @@ -116,10 +115,14 @@ def main(): logger.log_info("Press Ctrl+C to stop the simulation") # Run simulation loop - run_simulation(sim, camera) + run_simulation(sim, camera, args) -def run_simulation(sim, camera): +def run_simulation( + sim: SimulationManager, + camera: Camera, + args: argparse.Namespace, +) -> None: """Run the simulation loop with gizmo updates.""" step_count = 0 last_time = time.time() @@ -131,7 +134,8 @@ def run_simulation(sim, camera): ) try: - while True: + max_steps = resolve_demo_steps(args) + while max_steps is None or step_count < max_steps: # Update all gizmos managed by sim (including camera gizmo) sim.update_gizmos() @@ -145,7 +149,7 @@ def run_simulation(sim, camera): step_count += 1 # Display camera view in separate window - if step_count % 5 == 0: # Update display every 5 steps for performance + if not args.headless and step_count % 5 == 0: data = camera.get_data() if "color" in data: # Get RGB image and convert for OpenCV display @@ -153,25 +157,18 @@ def run_simulation(sim, camera): # 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) - - # 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") + 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("Gizmo Camera View", bgr_image) + if cv2.waitKey(1) & 0xFF == ord("h"): + sim.toggle_gizmo_visibility("gizmo_camera") # Example: Destroy gizmo after certain steps to test cleanup if step_count == 30000 and sim.has_gizmo("gizmo_camera"): @@ -209,7 +206,7 @@ def run_simulation(sim, camera): # Disable gizmo if it exists if sim.has_gizmo("gizmo_camera"): sim.disable_gizmo("gizmo_camera") - sim.destroy() + shutdown_sim(sim) logger.log_info("Simulation terminated successfully") diff --git a/examples/sim/gizmo/gizmo_object.py b/examples/sim/gizmo/gizmo_object.py index b0931f241..a265b0f1f 100644 --- a/examples/sim/gizmo/gizmo_object.py +++ b/examples/sim/gizmo/gizmo_object.py @@ -19,44 +19,46 @@ It shows the basic setup of simulation context, adding objects, and sensors. """ +from __future__ import annotations + import argparse import time -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg +from embodichain.lab.sim import SimulationManager +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg from embodichain.lab.sim.shapes import CubeCfg -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.objects import RigidObject, RigidObjectCfg +from embodichain.lab.sim.objects import RigidObjectCfg +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + resolve_demo_steps, + shutdown_sim, +) from embodichain.utils import logger -def main(): +def main() -> None: """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) + add_demo_args(parser) args = parser.parse_args() - # Configure the simulation - sim_cfg = SimulationManagerCfg( + sim = create_default_sim( + args, width=1920, height=1080, - headless=args.headless, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device=args.device, - render_cfg=RenderCfg( - renderer=args.renderer - ), # Enable ray tracing for better visuals + add_default_light=False, ) - # Create the simulation instance - sim = SimulationManager(sim_cfg) - # Add two cubes to the scene - cube1: RigidObject = sim.add_rigid_object( + sim.add_rigid_object( cfg=RigidObjectCfg( uid="cube1", shape=CubeCfg(size=[0.1, 0.1, 0.1]), @@ -70,7 +72,7 @@ def main(): init_pos=[0.0, 0.0, 1.0], ) ) - cube2: RigidObject = sim.add_rigid_object( + sim.add_rigid_object( cfg=RigidObjectCfg( uid="cube2", shape=CubeCfg(size=[0.1, 0.1, 0.1]), @@ -100,24 +102,21 @@ def main(): logger.log_info("Press Ctrl+C to stop the simulation") # Open window when the scene has been set up - if not args.headless: - sim.open_window() + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) # Run the simulation - run_simulation(sim) + run_simulation(sim, max_steps=resolve_demo_steps(args)) -def run_simulation(sim: SimulationManager): +def run_simulation(sim: SimulationManager, max_steps: int | None = None) -> None: """Run the simulation loop.""" - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - step_count = 0 gizmo_enabled = True try: last_time = time.time() last_step = 0 - while True: + while max_steps is None or step_count < max_steps: sim.update(step=1) # Update all gizmos if any are enabled @@ -128,7 +127,8 @@ def run_simulation(sim: SimulationManager): # Disable gizmo after 200000 steps (example) if step_count == 200000 and gizmo_enabled: logger.log_info("Disabling gizmo at step 200000") - sim.disable_gizmo("cube") + sim.disable_gizmo("cube1") + sim.disable_gizmo("cube2") gizmo_enabled = False # Print FPS every second @@ -146,7 +146,7 @@ def run_simulation(sim: SimulationManager): except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: - sim.destroy() + shutdown_sim(sim) logger.log_info("Simulation terminated successfully") diff --git a/examples/sim/gizmo/gizmo_robot.py b/examples/sim/gizmo/gizmo_robot.py index 6b8b9effb..2542ce1fa 100644 --- a/examples/sim/gizmo/gizmo_robot.py +++ b/examples/sim/gizmo/gizmo_robot.py @@ -17,45 +17,49 @@ Gizmo-Robot Example: Test Gizmo class on a robot (UR10) """ +from __future__ import annotations + +import argparse import time -import torch + import numpy as np -import argparse +import torch -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.solvers import PytorchSolverCfg from embodichain.lab.sim.cfg import ( - RenderCfg, + JointDrivePropertiesCfg, 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.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + resolve_demo_steps, + shutdown_sim, +) from embodichain.utils import logger -def main(): +def main() -> None: """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) + add_demo_args(parser) args = parser.parse_args() - # Configure the simulation - sim_cfg = SimulationManagerCfg( + sim = create_default_sim( + args, width=1920, height=1080, physics_dt=1.0 / 100.0, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), + add_default_light=False, ) - - sim = SimulationManager(sim_cfg) sim.set_manual_update(False) # Get UR10 URDF path @@ -102,7 +106,7 @@ def main(): 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=robot.device, ) joint_ids = robot.get_joint_ids("arm") robot.set_qpos(qpos=initial_qpos, joint_ids=joint_ids) @@ -110,26 +114,27 @@ def main(): time.sleep(0.2) # Wait for a moment to ensure everything is set up # Enable gizmo using the new API - sim.enable_gizmo(uid="ur10_gizmo_test", control_part="arm") - if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): - logger.log_error("Failed to enable gizmo!") - return + if not args.headless: + sim.enable_gizmo(uid="ur10_gizmo_test", control_part="arm") + if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): + shutdown_sim(sim) + raise RuntimeError("Failed to enable gizmo.") - sim.open_window() + maybe_open_window(sim, args) logger.log_info("Gizmo-Robot example started!") logger.log_info("Use the gizmo to drag the robot end-effector (EE)") logger.log_info("Press Ctrl+C to stop the simulation") - run_simulation(sim) + run_simulation(sim, max_steps=resolve_demo_steps(args)) -def run_simulation(sim: SimulationManager): +def run_simulation(sim: SimulationManager, max_steps: int | None = None) -> None: step_count = 0 try: last_time = time.time() last_step = 0 - while True: + while max_steps is None or step_count < max_steps: time.sleep(0.033) # 30Hz # Update all gizmos managed by sim sim.update_gizmos() @@ -149,7 +154,7 @@ def run_simulation(sim: SimulationManager): except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: - sim.destroy() + shutdown_sim(sim) logger.log_info("Simulation terminated successfully") diff --git a/examples/sim/gizmo/gizmo_scene.py b/examples/sim/gizmo/gizmo_scene.py index a37e6eb86..cac5c6e8b 100644 --- a/examples/sim/gizmo/gizmo_scene.py +++ b/examples/sim/gizmo/gizmo_scene.py @@ -22,49 +22,53 @@ Both objects can be interactively controlled through their respective gizmos. """ -import time -import torch -import numpy as np +from __future__ import annotations + import argparse +import time + import cv2 +import numpy as np +import torch -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.cfg import ( - RenderCfg, - RobotCfg, - URDFCfg, JointDrivePropertiesCfg, - RigidObjectCfg, RigidBodyAttributesCfg, + RigidObjectCfg, + RobotCfg, + URDFCfg, ) -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.shapes import CubeCfg from embodichain.lab.sim.sensors import CameraCfg from embodichain.lab.sim.solvers import PinkSolverCfg +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + resolve_demo_steps, + shutdown_sim, +) from embodichain.data import get_data_path from embodichain.utils import logger -def main(): +def main() -> None: """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) + add_demo_args(parser) args = parser.parse_args() - # Configure the simulation - sim_cfg = SimulationManagerCfg( + sim = create_default_sim( + args, width=1920, height=1080, - headless=args.headless, physics_dt=1.0 / 100.0, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), + add_default_light=False, ) - - sim = SimulationManager(sim_cfg) sim.set_manual_update(False) # Get DexForce W1 URDF path @@ -112,14 +116,14 @@ def main(): [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", + device=robot.device, ) 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", + device=robot.device, ) left_joint_ids = robot.get_joint_ids("left_arm") @@ -141,7 +145,7 @@ def main(): ), init_pos=[1.0, 0.0, 0.5], # Position to the side of the robot ) - cube = sim.add_rigid_object(cube_cfg) + sim.add_rigid_object(cube_cfg) camera_cfg = CameraCfg( uid="scene_camera", @@ -158,30 +162,27 @@ def main(): up=(0.0, 0.0, 1.0), ), ) - camera = sim.add_sensor(sensor_cfg=camera_cfg) - - # Enable gizmo for all assets after all are created and initialized - sim.enable_gizmo(uid="w1_gizmo_test", control_part="left_arm") - if not sim.has_gizmo("w1_gizmo_test", control_part="left_arm"): - logger.log_error("Failed to enable left arm gizmo!") - return - - sim.enable_gizmo(uid="w1_gizmo_test", control_part="right_arm") - if not sim.has_gizmo("w1_gizmo_test", control_part="right_arm"): - logger.log_error("Failed to enable right arm gizmo!") - return - - sim.enable_gizmo(uid="interactive_cube") - if not sim.has_gizmo("interactive_cube"): - logger.log_error("Failed to enable gizmo for cube!") - return + sim.add_sensor(sensor_cfg=camera_cfg) - sim.enable_gizmo(uid="scene_camera") - if not sim.has_gizmo("scene_camera"): - logger.log_error("Failed to enable gizmo for camera!") - return + if not args.headless: + gizmos = ( + ("w1_gizmo_test", "left_arm"), + ("w1_gizmo_test", "right_arm"), + ("interactive_cube", None), + ("scene_camera", None), + ) + for uid, control_part in gizmos: + if control_part is None: + sim.enable_gizmo(uid=uid) + enabled = sim.has_gizmo(uid) + else: + sim.enable_gizmo(uid=uid, control_part=control_part) + enabled = sim.has_gizmo(uid, control_part=control_part) + if not enabled: + shutdown_sim(sim) + raise RuntimeError(f"Failed to enable gizmo for {uid}.") - sim.open_window() + maybe_open_window(sim, args) logger.log_info("Gizmo Scene example started!") logger.log_info("Four gizmos are active in the scene:") @@ -191,23 +192,27 @@ def main(): 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) + run_simulation(sim, args) -def run_simulation(sim: SimulationManager): +def run_simulation( + sim: SimulationManager, + args: argparse.Namespace, +) -> 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: + max_steps = resolve_demo_steps(args) + while max_steps is None or step_count < max_steps: time.sleep(0.033) # 30Hz sim.update_gizmos() step_count += 1 # Display camera view in a window every 5 steps - if camera is not None and step_count % 5 == 0: + if not args.headless and camera is not None and step_count % 5 == 0: camera.update() data = camera.get_data() if "color" in data: @@ -243,7 +248,7 @@ def run_simulation(sim: SimulationManager): logger.log_info("\nStopping simulation...") finally: cv2.destroyAllWindows() - sim.destroy() + shutdown_sim(sim) logger.log_info("Simulation terminated successfully") diff --git a/examples/sim/gizmo/gizmo_w1.py b/examples/sim/gizmo/gizmo_w1.py index 6c06c4677..b54cfe63d 100644 --- a/examples/sim/gizmo/gizmo_w1.py +++ b/examples/sim/gizmo/gizmo_w1.py @@ -17,46 +17,43 @@ Gizmo-Robot Example: Test Gizmo class on a robot (UR10) """ +from __future__ import annotations + +import argparse import time -import torch + import numpy as np -import argparse +import torch -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import ( - RenderCfg, - RobotCfg, - URDFCfg, - JointDrivePropertiesCfg, +from embodichain.lab.sim import SimulationManager +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + resolve_demo_steps, + shutdown_sim, ) -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(): +def main() -> None: """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) + add_demo_args(parser) args = parser.parse_args() - # Configure the simulation - sim_cfg = SimulationManagerCfg( + sim = create_default_sim( + args, width=1920, height=1080, - headless=args.headless, physics_dt=1.0 / 100.0, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), + add_default_light=False, ) - - sim = SimulationManager(sim_cfg) sim.set_manual_update(False) cfg = DexforceW1Cfg.from_dict( @@ -132,14 +129,14 @@ def main(): [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", + device=robot.device, ) 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", + device=robot.device, ) left_joint_ids = robot.get_joint_ids("left_arm") @@ -151,31 +148,29 @@ def main(): time.sleep(0.2) # Wait for a moment to ensure everything is set up # Enable gizmo for both arms using the new API - sim.enable_gizmo(uid="w1_gizmo_test", control_part="left_arm") - if not sim.has_gizmo("w1_gizmo_test", control_part="left_arm"): - logger.log_error("Failed to enable left arm gizmo!") - return - - sim.enable_gizmo(uid="w1_gizmo_test", control_part="right_arm") - if not sim.has_gizmo("w1_gizmo_test", control_part="right_arm"): - logger.log_error("Failed to enable right arm gizmo!") - return + if not args.headless: + sim.enable_gizmo(uid="w1_gizmo_test", control_part="left_arm") + sim.enable_gizmo(uid="w1_gizmo_test", control_part="right_arm") + for arm_name in ("left_arm", "right_arm"): + if not sim.has_gizmo("w1_gizmo_test", control_part=arm_name): + shutdown_sim(sim) + raise RuntimeError(f"Failed to enable {arm_name} gizmo.") - sim.open_window() + maybe_open_window(sim, args) logger.log_info("Gizmo-DexForce W1 example started!") logger.log_info("Use the gizmos to drag both robot arms' end-effectors") logger.log_info("Press Ctrl+C to stop the simulation") - run_simulation(sim) + run_simulation(sim, max_steps=resolve_demo_steps(args)) -def run_simulation(sim: SimulationManager): +def run_simulation(sim: SimulationManager, max_steps: int | None = None) -> None: step_count = 0 try: last_time = time.time() last_step = 0 - while True: + while max_steps is None or step_count < max_steps: time.sleep(0.033) # 30Hz # Update all gizmos managed by sim sim.update_gizmos() @@ -195,7 +190,7 @@ def run_simulation(sim: SimulationManager): except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: - sim.destroy() + shutdown_sim(sim) logger.log_info("Simulation terminated successfully") diff --git a/examples/sim/planners/neural_planner.py b/examples/sim/planners/neural_planner.py index c9e0700da..508124bd9 100644 --- a/examples/sim/planners/neural_planner.py +++ b/examples/sim/planners/neural_planner.py @@ -31,9 +31,8 @@ import torch from embodichain.data.assets.planner_assets import download_neural_planner_checkpoint -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import MarkerCfg, RenderCfg +from embodichain.lab.sim import SimulationManager +from embodichain.lab.sim.cfg import MarkerCfg from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.robots.franka_panda import FrankaPandaCfg from embodichain.lab.sim.planners import ( @@ -45,12 +44,20 @@ PlanState, ) from embodichain.lab.sim.planners.neural_planner import NeuralPlanOptions +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + shutdown_sim, +) def parse_args() -> argparse.Namespace: default_device = "cuda" if torch.cuda.is_available() else "cpu" parser = argparse.ArgumentParser(description="NeuralPlanner waypoint example") - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) parser.set_defaults(device=default_device, arena_space=2.0) parser.add_argument( "--num-waypoints", @@ -203,25 +210,21 @@ def main() -> None: resolved_device.index if resolved_device.type == "cuda" else int(args.gpu_id) ) assert effective_gpu_id is not None - sim = SimulationManager( - SimulationManagerCfg( - headless=args.headless, - sim_device=sim_device, - num_envs=args.num_envs, - arena_space=args.arena_space, - gpu_id=effective_gpu_id, - render_cfg=RenderCfg(renderer=args.renderer), - ) + args.device = sim_device + args.gpu_id = effective_gpu_id + sim = create_default_sim( + args, + num_envs=args.num_envs, + arena_space=args.arena_space, + add_default_light=False, ) try: robot = create_franka(sim) arm_name = "arm" device = robot.device - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - if not args.headless: - sim.open_window() + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) start_qpos = torch.tensor( [0.0, -np.pi / 4, 0.0, -3 * np.pi / 4, 0.0, np.pi / 2, np.pi / 4], @@ -278,21 +281,22 @@ def main() -> None: f"NeuralPlanner failed for environment(s) {failed_env_ids}." ) - play_trajectory( - sim, - robot, - arm_name, - result.positions, - step_repeat=args.step_repeat, - ) - sim.update(step=args.hold_steps) + with DemoRecording(sim, args, prefix="neural_planner"): + play_trajectory( + sim, + robot, + arm_name, + result.positions, + step_repeat=args.step_repeat, + ) + sim.update(step=args.hold_steps) if args.interactive: from IPython import embed embed(header="NeuralPlanner example. Press Ctrl+D to exit.") finally: - sim.destroy() + shutdown_sim(sim) SimulationManager.flush_cleanup_queue() diff --git a/examples/sim/robot/dexforce_w1.py b/examples/sim/robot/dexforce_w1.py index beee19d9c..9ac101959 100644 --- a/examples/sim/robot/dexforce_w1.py +++ b/examples/sim/robot/dexforce_w1.py @@ -19,56 +19,68 @@ adding a pair of parallel grippers as the left and right hands. """ -import numpy as np - -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.robots import DexforceW1Cfg +from __future__ import annotations +import argparse -def main(): - np.set_printoptions(precision=5, suppress=True) +from embodichain.lab.sim.robots import DexforceW1Cfg +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, + shutdown_sim, +) - config = SimulationManagerCfg() - sim = SimulationManager(config) - cfg = DexforceW1Cfg.from_dict( - { - "uid": "dexforce_w1", - "version": "v021", - "arm_kind": "anthropomorphic", - "with_default_eef": False, - "control_parts": { - "left_eef": ["LEFT_FINGER1_JOINT", "LEFT_FINGER2_JOINT"], - "right_eef": ["RIGHT_FINGER1_JOINT", "RIGHT_FINGER2_JOINT"], - }, - "urdf_cfg": { - "components": [ - { - "component_type": "left_hand", - "urdf_path": "DH_PGC_140_50/DH_PGC_140_50.urdf", - }, - { - "component_type": "right_hand", - "urdf_path": "DH_PGC_140_50/DH_PGC_140_50.urdf", - }, - ] - }, - "drive_pros": { - "max_effort": { - "left_eef": 10.0, - "right_eef": 10.0, - } - }, - } +def main() -> None: + """Create a DexForce W1 with custom parallel-gripper end effectors.""" + parser = add_demo_args( + argparse.ArgumentParser(description="Customize DexForce W1 end effectors.") ) + args = parser.parse_args() + setup_print_options() - robot = sim.add_robot(cfg=cfg) - sim.update(step=1) - print("DexforceW1 with a user defined end-effector added to the simulation.") - - from IPython import embed - - embed() + sim = create_default_sim(args, add_default_light=False) + try: + cfg = DexforceW1Cfg.from_dict( + { + "uid": "dexforce_w1", + "version": "v021", + "arm_kind": "anthropomorphic", + "with_default_eef": False, + "control_parts": { + "left_eef": ["LEFT_FINGER1_JOINT", "LEFT_FINGER2_JOINT"], + "right_eef": ["RIGHT_FINGER1_JOINT", "RIGHT_FINGER2_JOINT"], + }, + "urdf_cfg": { + "components": [ + { + "component_type": "left_hand", + "urdf_path": "DH_PGC_140_50/DH_PGC_140_50.urdf", + }, + { + "component_type": "right_hand", + "urdf_path": "DH_PGC_140_50/DH_PGC_140_50.urdf", + }, + ] + }, + "drive_pros": { + "max_effort": { + "left_eef": 10.0, + "right_eef": 10.0, + }, + }, + } + ) + sim.add_robot(cfg=cfg) + sim.update(step=1) + print("DexforceW1 with user-defined end effectors added.") + maybe_open_window(sim, args) + maybe_wait_for_user(args, "Press Enter to exit...") + finally: + shutdown_sim(sim) if __name__ == "__main__": diff --git a/examples/sim/scene/scene_demo.py b/examples/sim/scene/scene_demo.py index 1c08af6ae..6a4bb632c 100644 --- a/examples/sim/scene/scene_demo.py +++ b/examples/sim/scene/scene_demo.py @@ -18,24 +18,32 @@ It supports loading kitchen/factory/office scenes via EmbodiChainDataset. """ +from __future__ import annotations + import argparse -import time -from pathlib import Path import math +from pathlib import Path + import embodichain.utils.logger as logger -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.cfg import ( - RenderCfg, - RigidBodyAttributesCfg, LightCfg, - RobotCfg, - URDFCfg, + RigidBodyAttributesCfg, ) from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim.objects import RigidObject, RigidObjectCfg, Robot from embodichain.data.assets.scene_assets import SceneData from embodichain.data.constants import EMBODICHAIN_DEFAULT_DATA_ROOT -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + resolve_demo_steps, + run_simulation_loop, + shutdown_sim, +) def resolve_asset_path(scene_name: str) -> str: @@ -54,7 +62,7 @@ def resolve_asset_path(scene_name: str) -> str: logger.log_info(f"Using local asset: {local_gltf}") return str(local_gltf) - scene_data = SceneData() + SceneData() extracted_dir = Path(EMBODICHAIN_DEFAULT_DATA_ROOT) / "extract" / "SceneData" glb_path = extracted_dir / f"{scene_name}.glb" @@ -72,22 +80,24 @@ def resolve_asset_path(scene_name: str) -> str: ) -def run_simulation(sim: SimulationManager): +def run_simulation( + sim: SimulationManager, + args: argparse.Namespace, +) -> None: """Run the simulation loop.""" - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - try: - while True: - time.sleep(0.01) - except KeyboardInterrupt: - logger.log_info("\n Stopping simulation...") + with DemoRecording(sim, args, prefix=f"scene_{args.scene}"): + run_simulation_loop( + sim, + max_steps=resolve_demo_steps(args), + sleep=0.01, + ) finally: - sim.destroy() + shutdown_sim(sim) logger.log_info("Simulation terminated successfully.") -def main(): +def main() -> None: parser = argparse.ArgumentParser( description="Create a simulation scene with SimulationManager" ) @@ -98,7 +108,7 @@ def main(): choices=["kitchen", "factory", "office", "local"], help="Choose which scene to load", ) - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) args = parser.parse_args() logger.log_info(f"Initializing scene '{args.scene}'") @@ -111,24 +121,20 @@ def main(): print(f"Failed to download or resolve scene asset: {e}") return - sim_cfg = SimulationManagerCfg( + sim = create_default_sim( + args, width=1920, height=1080, - headless=True, physics_dt=1.0 / 100.0, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), num_envs=args.num_envs, arena_space=10.0, + add_default_light=False, ) - sim = SimulationManager(sim_cfg) num_lights = 8 radius = 5 height = 8 intensity = 50 - lights = [] - for i in range(num_lights): angle = 2 * math.pi * i / num_lights x = radius * math.cos(angle) @@ -136,7 +142,7 @@ def main(): z = height uid = f"l{i+1}" cfg = LightCfg(uid=uid, intensity=intensity, radius=600, init_pos=[x, y, z]) - lights.append(sim.add_light(cfg)) + sim.add_light(cfg) physics_attrs = RigidBodyAttributesCfg( mass=10, @@ -147,7 +153,7 @@ def main(): try: logger.log_info(f"Loading scene asset into simulation: {asset_path}") - scene_obj: RigidObject = sim.add_rigid_object( + sim.add_rigid_object( cfg=RigidObjectCfg( uid=args.scene, shape=MeshCfg(fpath=asset_path), @@ -160,7 +166,7 @@ def main(): if args.scene == "factory": from embodichain.lab.sim.robots.dexforce_w1.cfg import DexforceW1Cfg - w1_robot: Robot = sim.add_robot( + sim.add_robot( cfg=DexforceW1Cfg.from_dict( { "uid": "dexforce_w1", @@ -174,15 +180,17 @@ def main(): except Exception as e: logger.log_info(f"Failed to load scene asset: {e}") + shutdown_sim(sim) return logger.log_info(f"Scene '{args.scene}' setup complete!") logger.log_info(f"Running simulation with {args.num_envs} environment(s)") logger.log_info("Press Ctrl+C to stop the simulation") - sim.open_window() + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) - run_simulation(sim) + run_simulation(sim, args) if __name__ == "__main__": diff --git a/examples/sim/sensors/batch_camera.py b/examples/sim/sensors/batch_camera.py index b6eb48247..6ba9f517e 100644 --- a/examples/sim/sensors/batch_camera.py +++ b/examples/sim/sensors/batch_camera.py @@ -14,113 +14,120 @@ # limitations under the License. # ---------------------------------------------------------------------------- +"""Capture a batch of mono or stereo camera images.""" + +from __future__ import annotations + +import argparse import time -import numpy as np + import matplotlib.pyplot as plt +import numpy as np +import torch -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RenderCfg, RigidObjectCfg, LightCfg +from embodichain.lab.sim.cfg import RigidObjectCfg from embodichain.lab.sim.shapes import MeshCfg -from embodichain.lab.sim.objects import RigidObject, Light from embodichain.lab.sim.sensors import ( Camera, StereoCamera, CameraCfg, StereoCameraCfg, ) -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.data import get_data_path +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + setup_print_options, + shutdown_sim, +) -def main(args): - config = SimulationManagerCfg( - headless=True, - sim_device=args.device, +def run(args: argparse.Namespace) -> None: + """Render and save or display one batch of camera frames.""" + sim = create_default_sim( + args, num_envs=args.num_envs, arena_space=2, - render_cfg=RenderCfg(renderer=args.renderer), + add_default_light=False, ) - sim = SimulationManager(config) - - rigid_obj: RigidObject = sim.add_rigid_object( - cfg=RigidObjectCfg( - uid="obj", - shape=MeshCfg(fpath=get_data_path("Chair/chair.glb")), - init_pos=(0, 0, 0.2), + try: + sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="obj", + shape=MeshCfg(fpath=get_data_path("Chair/chair.glb")), + init_pos=(0, 0, 0.2), + ) ) - ) - - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - - if args.headless is False: - sim.open_window() + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) + setup_print_options() - import torch - - torch.set_printoptions(precision=4, sci_mode=False) - - eye = (0.0, 0, 2.0) - target = (0.0, 0.0, 0.0) - if args.sensor_type == "stereo": - camera: StereoCamera = sim.add_sensor( - sensor_cfg=StereoCameraCfg( - width=640, - height=480, - extrinsics=CameraCfg.ExtrinsicsCfg(eye=eye, target=target), + eye = (0.0, 0, 2.0) + target = (0.0, 0.0, 0.0) + if args.sensor_type == "stereo": + camera: Camera | StereoCamera = sim.add_sensor( + sensor_cfg=StereoCameraCfg( + width=640, + height=480, + extrinsics=CameraCfg.ExtrinsicsCfg(eye=eye, target=target), + ) ) - ) - else: - camera: Camera = sim.add_sensor( - sensor_cfg=CameraCfg( - width=640, - height=480, - extrinsics=CameraCfg.ExtrinsicsCfg(eye=eye, target=target), + else: + camera = sim.add_sensor( + sensor_cfg=CameraCfg( + width=640, + height=480, + extrinsics=CameraCfg.ExtrinsicsCfg(eye=eye, target=target), + ) ) - ) - - # TODO: To be removed - sim.reset_objects_state() - t0 = time.time() - camera.update() - print(f"Camera update time: {time.time() - t0:.4f} seconds") + sim.reset_objects_state() - data_frame = camera.get_data() + started_at = time.perf_counter() + camera.update() + print(f"Camera update time: {time.perf_counter() - started_at:.4f} seconds") - t0 = time.time() - rgba = data_frame["color"].cpu().numpy() - if args.sensor_type == "stereo": - rgba_right = data_frame["color_right"].cpu().numpy() + data_frame = camera.get_data() + rgba = data_frame["color"].cpu().numpy() + rgba_right = ( + data_frame["color_right"].cpu().numpy() + if args.sensor_type == "stereo" + else None + ) - # plot rgba into a grid of images - grid_x = np.ceil(np.sqrt(args.num_envs)).astype(int) - grid_y = np.ceil(args.num_envs / grid_x).astype(int) - fig, axs = plt.subplots(grid_x, grid_y, figsize=(12, 6), squeeze=False) - axs = axs.flatten() - for i in range(args.num_envs): + grid_x = int(np.ceil(np.sqrt(args.num_envs))) + grid_y = int(np.ceil(args.num_envs / grid_x)) + fig, axs = plt.subplots(grid_x, grid_y, figsize=(12, 6), squeeze=False) + for env_id, axis in enumerate(axs.flatten()): + axis.axis("off") + if env_id >= args.num_envs: + continue + image = ( + np.concatenate((rgba[env_id], rgba_right[env_id]), axis=1) + if rgba_right is not None + else rgba[env_id] + ) + axis.imshow(image) + axis.set_title(f"Env {env_id}") - if args.sensor_type == "stereo": - image = np.concatenate((rgba[i], rgba_right[i]), axis=1) + if args.headless: + fig.savefig("camera_data.png") else: - image = rgba[i] - axs[i].imshow(image) - axs[i].axis("off") - axs[i].set_title(f"Env {i}") + plt.show() + plt.close(fig) + finally: + shutdown_sim(sim) - if args.headless: - plt.savefig(f"camera_data.png") - else: - plt.show() - - -if __name__ == "__main__": - import argparse +def main() -> None: + """Parse command-line arguments and capture a camera batch.""" parser = argparse.ArgumentParser(description="Run the batch robot simulation.") - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) parser.add_argument( "--sensor_type", + "--sensor-type", type=str, default="camera", choices=["stereo", "camera"], @@ -128,4 +135,8 @@ def main(args): ) args = parser.parse_args() - main(args) + run(args) + + +if __name__ == "__main__": + main() diff --git a/examples/sim/sensors/create_contact_sensor.py b/examples/sim/sensors/create_contact_sensor.py index 17c26caff..7ec85c98a 100644 --- a/examples/sim/sensors/create_contact_sensor.py +++ b/examples/sim/sensors/create_contact_sensor.py @@ -14,20 +14,17 @@ # 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. -""" +"""Create a grasping scene and report filtered robot/object contacts.""" + +from __future__ import annotations import argparse import time + import torch -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import ( - RenderCfg, - RigidBodyAttributesCfg, -) +from embodichain.lab.sim import SimulationManager +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg from embodichain.lab.sim.sensors import ( ContactSensorCfg, ArticulationContactFilterCfg, @@ -35,11 +32,22 @@ from embodichain.lab.sim.shapes import CubeCfg from embodichain.lab.sim.objects import RigidObject, RigidObjectCfg, Robot, RobotCfg from embodichain.data import get_data_path -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + resolve_demo_steps, + run_simulation_loop, + shutdown_sim, +) def create_cube( - sim: SimulationManager, uid: str, position: list = (0.0, 0.0, 0) + sim: SimulationManager, + uid: str, + position: tuple[float, float, float] = (0.0, 0.0, 0.0), ) -> RigidObject: """create cube @@ -70,7 +78,11 @@ def create_cube( return cube -def robot_grasp_pose(robot: Robot, cube: RigidObject, sim: SimulationManager): +def robot_grasp_pose( + robot: Robot, + cube: RigidObject, + sim: SimulationManager, +) -> None: sim.update(step=100) arm_ids = robot.get_joint_ids("arm") gripper_ids = robot.get_joint_ids("hand") @@ -85,12 +97,14 @@ def robot_grasp_pose(robot: Robot, cube: RigidObject, sim: SimulationManager): approach_xpos = target_xpos.clone() approach_xpos[:, 2, 3] += 0.1 - is_success, approach_qpos = robot.compute_ik( + approach_success, approach_qpos = robot.compute_ik( pose=approach_xpos, joint_seed=rest_arm_qpos, name="arm" ) - is_success, target_qpos = robot.compute_ik( + target_success, target_qpos = robot.compute_ik( pose=target_xpos, joint_seed=approach_qpos, name="arm" ) + if not (approach_success.all() and target_success.all()): + raise RuntimeError("Failed to solve the cube grasp pose.") robot.set_qpos(approach_qpos, joint_ids=arm_ids) sim.update(step=40) @@ -106,7 +120,9 @@ def robot_grasp_pose(robot: Robot, cube: RigidObject, sim: SimulationManager): def create_robot( - sim: SimulationManager, uid: str, position: list = (0.0, 0.0, 0) + sim: SimulationManager, + uid: str, + position: tuple[float, float, float] = (0.0, 0.0, 0.0), ) -> Robot: """create robot @@ -149,9 +165,9 @@ def create_robot( "init_pos": position, "init_qpos": [0.0, -1.57, 1.57, -1.57, -1.57, 0.0, 0.0, 0.0], "drive_pros": { - "stiffness": {"JOINT[1-6]": 1e4, "FINGER[1-2]_JOINT": 1e2}, - "damping": {"JOINT[1-6]": 1e3, "FINGER[1-2]_JOINT": 1e1}, - "max_effort": {"JOINT[1-6]": 1e5, "FINGER[1-2]_JOINT": 1e3}, + "stiffness": {"Joint[1-6]": 1e4, "finger[1-2]_joint": 1e2}, + "damping": {"Joint[1-6]": 1e3, "finger[1-2]_joint": 1e1}, + "max_effort": {"Joint[1-6]": 1e5, "finger[1-2]_joint": 1e3}, }, "solver_cfg": { "arm": { @@ -166,69 +182,56 @@ def create_robot( ], } }, - "control_parts": {"arm": ["JOINT[1-6]"], "hand": ["FINGER[1-2]_JOINT"]}, + "control_parts": { + "arm": ["Joint[1-6]"], + "hand": ["finger[1-2]_joint"], + }, } robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(robot_cfg_dict)) return robot -def main(): +def main() -> None: """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) + add_demo_args(parser) args = parser.parse_args() - # Configure the simulation - sim_cfg = SimulationManagerCfg( - width=1920, - height=1080, + sim = create_default_sim( + args, num_envs=args.num_envs, - headless=True, - physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device=args.device, - render_cfg=RenderCfg( - renderer=args.renderer - ), # Enable ray tracing for better visuals + add_default_light=False, ) - # Create the simulation instance - sim = SimulationManager(sim_cfg) - - # Add objects to the scene - cube0 = create_cube(sim, "cube0", position=[0.0, 0.0, 0.03]) - cube1 = create_cube(sim, "cube1", position=[0.0, 0.0, 0.06]) - cube2 = create_cube(sim, "cube2", position=[0.0, 0.0, 0.09]) - robot = create_robot(sim, "UR10_PGI", position=[0.5, 0.0, 0.0]) - - print("[INFO]: Scene setup complete!") - print(f"[INFO]: Running simulation with {args.num_envs} environment(s)") - print("[INFO]: Press Ctrl+C to stop the simulation") - - # Open window when the scene has been set up - if not args.headless: - sim.open_window() - - robot_grasp_pose(robot, cube2, sim) - # Run the simulation - run_simulation(sim) + try: + create_cube(sim, "cube0", position=(0.0, 0.0, 0.03)) + create_cube(sim, "cube1", position=(0.0, 0.0, 0.06)) + cube2 = create_cube(sim, "cube2", position=(0.0, 0.0, 0.09)) + robot = create_robot(sim, "UR10_PGI", position=(0.5, 0.0, 0.0)) + + print("[INFO]: Scene setup complete!") + print(f"[INFO]: Running simulation with {args.num_envs} environment(s)") + + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) + robot_grasp_pose(robot, cube2, sim) + run_simulation(sim, args) + finally: + shutdown_sim(sim) + print("[INFO]: Simulation terminated successfully") -def run_simulation(sim: SimulationManager): +def run_simulation(sim: SimulationManager, args: argparse.Namespace) -> None: """Run the simulation loop. Args: sim: The SimulationManager instance to run """ - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - - step_count = 0 # contact filter config contact_filter_cfg = ContactSensorCfg() contact_filter_cfg.rigid_uid_list = ["cube0", "cube1", "cube2"] @@ -240,45 +243,38 @@ def run_simulation(sim: SimulationManager): contact_sensor = sim.add_sensor(sensor_cfg=contact_filter_cfg) - try: - accmulated_cost_time = 0.0 - while True: - # Update physics simulation - sim.update(step=1) - start_time = time.time() - contact_sensor.update() - contact_report = contact_sensor.get_data() - accmulated_cost_time += time.time() - start_time - step_count += 1 - - # Print FPS every second - if step_count % 100 == 0: - average_cost_time = accmulated_cost_time / 100.0 - print( - f"[INFO]: Fetch contact cost time: {average_cost_time * 1000:.2f} ms, num_envs: {sim.num_envs}" - ) - # filter contact report for a rigid object with a articulation link - cube2_user_ids = sim.get_rigid_object("cube2").get_user_ids() - finger1_user_ids = ( - sim.get_robot("UR10_PGI").get_user_ids("finger1_link").reshape(-1) - ) - filter_user_ids = torch.cat([cube2_user_ids, finger1_user_ids]) - filter_contact_report = contact_sensor.filter_by_user_ids( - filter_user_ids - ) - # print("filter_contact_report", filter_contact_report) - # visualize contact points - contact_sensor.set_contact_point_visibility( - visible=True, rgba=(0.0, 0.0, 1.0, 1.0), point_size=6.0 - ) - accmulated_cost_time = 0.0 - - except KeyboardInterrupt: - print("\n[INFO]: Stopping simulation...") - finally: - # Clean up resources - sim.destroy() - print("[INFO]: Simulation terminated successfully") + accumulated_cost_time = 0.0 + + def update_contact(step: int) -> None: + """Update the sensor and periodically report/filter contacts.""" + nonlocal accumulated_cost_time + started_at = time.perf_counter() + contact_sensor.update() + contact_sensor.get_data() + accumulated_cost_time += time.perf_counter() - started_at + + if step % 100 != 0: + return + print( + "[INFO]: Fetch contact cost time: " + f"{accumulated_cost_time * 10:.2f} ms, num_envs: {sim.num_envs}" + ) + cube_ids = sim.get_rigid_object("cube2").get_user_ids() + finger_ids = sim.get_robot("UR10_PGI").get_user_ids("finger1_link").reshape(-1) + contact_sensor.filter_by_user_ids(torch.cat([cube_ids, finger_ids])) + contact_sensor.set_contact_point_visibility( + visible=True, + rgba=(0.0, 0.0, 1.0, 1.0), + point_size=6.0, + ) + accumulated_cost_time = 0.0 + + with DemoRecording(sim, args, prefix="contact_sensor"): + run_simulation_loop( + sim, + max_steps=resolve_demo_steps(args), + on_step=update_contact, + ) if __name__ == "__main__": diff --git a/examples/sim/solvers/differential_solver.py b/examples/sim/solvers/differential_solver.py index 11efa65d0..543ea6f7d 100644 --- a/examples/sim/solvers/differential_solver.py +++ b/examples/sim/solvers/differential_solver.py @@ -13,237 +13,174 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -import os -import time -import numpy as np -import torch -from IPython import embed - -from embodichain.data import get_data_path -from embodichain.lab.sim.cfg import RobotCfg -from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import MarkerCfg +"""Demonstrate batched differential IK on an SR5 robot.""" -def main(visualize: bool = True): - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) +from __future__ import annotations - # Set up simulation with specified device (CPU or CUDA) - sim_device = "cpu" - num_envs = 9 # Number of parallel arenas/environments - config = SimulationManagerCfg( - headless=False, sim_device=sim_device, arena_space=1.5, num_envs=num_envs - ) - sim = SimulationManager(config) - sim.set_manual_update(False) - - # Load robot URDF file - urdf = get_data_path("Rokae/SR5/SR5.urdf") - assert os.path.isfile(urdf) - - # Robot configuration - cfg_dict = { - "fpath": urdf, - "control_parts": { - "main_arm": [ - "joint1", - "joint2", - "joint3", - "joint4", - "joint5", - "joint6", - ], - }, - "solver_cfg": { - "main_arm": { - "class_type": "DifferentialSolver", - "end_link_name": "ee_link", - "root_link_name": "base_link", - }, - }, - } +import argparse +import time - robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) +import numpy as np +import torch - # Prepare initial joint positions for all environments - rad = torch.deg2rad(torch.tensor(45.0)) - arm_name = "main_arm" - fk_qpos = torch.full((num_envs, 6), rad, dtype=torch.float32, device="cpu") - # All envs start with the same qpos (can be randomized) - qpos = torch.from_numpy(np.array([0.0, 0.0, np.pi / 2, 0.0, np.pi / 2, 0.0])).to( - fk_qpos.device +from embodichain.data import get_data_path +from embodichain.lab.sim.cfg import MarkerCfg, RobotCfg +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, + shutdown_sim, +) + +TARGET_OFFSETS = ( + (0.3, 0.0, 0.0), + (0.2, -0.2, 0.0), + (0.0, 0.0, 0.2), + (0.2, 0.0, 0.2), + (-0.3, 0.0, 0.0), + (-0.2, 0.2, 0.0), + (0.0, 0.0, -0.2), + (-0.2, 0.0, -0.2), + (0.1, 0.1, -0.1), +) + + +def main() -> None: + """Solve and replay one Cartesian path per environment.""" + parser = add_demo_args( + argparse.ArgumentParser(description="Run batched differential IK.") ) - qpos = qpos.unsqueeze(0).repeat(num_envs, 1) - robot.set_qpos(qpos=qpos, joint_ids=robot.get_joint_ids(arm_name)) - - time.sleep(3.0) - fk_xpos = robot.compute_fk( - qpos=qpos, name=arm_name, to_matrix=True - ) # (num_envs, 4, 4) - - # Prepare batch start and end poses for all envs - start_pose = fk_xpos.clone() # (num_envs, 4, 4) - end_pose = fk_xpos.clone() - move_vecs = torch.tensor( - [ - [0.3, 0.0, 0.0], - [0.2, -0.2, 0.0], - [0.0, 0.0, 0.2], - [0.2, 0.0, 0.2], - [-0.3, 0.0, 0.0], - [-0.2, 0.2, 0.0], - [0.0, 0.0, -0.2], - [-0.2, 0.0, -0.2], - [0.1, 0.1, -0.1], - ], - dtype=end_pose.dtype, - device=end_pose.device, + parser.set_defaults(num_envs=9) + parser.add_argument( + "--num_steps", + "--num-steps", + type=int, + default=100, + help="Number of Cartesian interpolation steps.", ) - end_pose[ - :, :3, 3 - ] += move_vecs # Move each env's end-effector in a different direction - - num_steps = 100 - # Interpolate poses for each env - interpolated_poses = torch.stack( - [torch.lerp(start_pose, end_pose, t) for t in np.linspace(0, 1, num_steps)], - dim=1, - ) # (num_envs, num_steps, 4, 4) - - # Initial joint positions for all envs - ik_qpos = qpos.clone() - - ik_qpos_results = [] - ik_success_flags = [] - - ik_compute_begin = time.time() - # Batch IK solving for each step - for step in range(num_steps): - poses = interpolated_poses[:, step, :, :] # (num_envs, 4, 4) - if poses.shape[0] != num_envs: - poses = poses.expand(num_envs, *poses.shape[1:]) - if ik_qpos.shape[0] != num_envs: - ik_qpos = ik_qpos.expand(num_envs, *ik_qpos.shape[1:]) - assert ( - poses.shape[0] == num_envs - ), f"poses batch mismatch: {poses.shape[0]} vs {num_envs}" - assert ( - ik_qpos.shape[0] == num_envs - ), f"ik_qpos batch mismatch: {ik_qpos.shape[0]} vs {num_envs}" - - # Parallel batch IK solving - res, ik_qpos_new = robot.compute_ik( - pose=poses, joint_seed=ik_qpos, name=arm_name - ) - ik_qpos_results.append(ik_qpos_new.clone()) - ik_success_flags.append(res) - ik_qpos = ik_qpos_new # Update joint seed - ik_compute_end = time.time() - print( - f"IK compute time for {num_steps} steps and {num_envs} envs: {ik_compute_end - ik_compute_begin:.4f} seconds" + args = parser.parse_args() + setup_print_options() + + sim = create_default_sim( + args, + num_envs=args.num_envs, + arena_space=1.5, + add_default_light=False, ) - - # Collect visualization data for all steps and environments - if visualize: - draw_data = [[] for _ in range(num_envs)] - for env_id in range(num_envs): - for step in range(num_steps): - ik_qpos_new = ik_qpos_results[step] - ik_xpos = robot.compute_fk(qpos=ik_qpos_new, name=arm_name, to_matrix=True) - local_pose = robot._entities[env_id].get_world_pose() - if visualize: - fk_axis = local_pose @ end_pose[env_id].cpu().numpy() - ik_axis = local_pose @ ik_xpos[env_id].cpu().numpy() - local_axis = local_pose @ ik_xpos[env_id].cpu().numpy() - - draw_data[env_id].append( - { - "step": step, - "fk_axis": fk_axis, - "ik_axis": ik_axis, - "local_axis": local_axis, - } - ) - - if visualize: - # Batch draw all steps' data for each environment - for env_id in range(num_envs): - # Only draw fk_axis and ik_axis once per env (first step) - fk_axis = draw_data[env_id][0]["fk_axis"] - ik_axis = draw_data[env_id][0]["ik_axis"] - - sim.draw_marker( - cfg=MarkerCfg( - name=f"fk_axis_env{env_id}", - marker_type="axis", - axis_xpos=fk_axis, - axis_size=0.002, - axis_len=0.005, - arena_index=env_id, - ) + try: + robot = sim.add_robot( + cfg=RobotCfg.from_dict( + { + "fpath": get_data_path("Rokae/SR5/SR5.urdf"), + "control_parts": { + "main_arm": [f"joint{i}" for i in range(1, 7)], + }, + "solver_cfg": { + "main_arm": { + "class_type": "DifferentialSolver", + "end_link_name": "ee_link", + "root_link_name": "base_link", + }, + }, + } ) - - sim.draw_marker( - cfg=MarkerCfg( - name=f"ik_axis_env{env_id}", - marker_type="axis", - axis_xpos=ik_axis, - axis_size=0.002, - axis_len=0.005, - arena_index=env_id, + ) + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) + + arm_name = "main_arm" + qpos = torch.tensor( + [0.0, 0.0, np.pi / 2, 0.0, np.pi / 2, 0.0], + dtype=torch.float32, + device=robot.device, + ).repeat(args.num_envs, 1) + robot.set_qpos(qpos, joint_ids=robot.get_joint_ids(arm_name)) + start_pose = robot.compute_fk( + qpos=qpos, + name=arm_name, + to_matrix=True, + ) + target_pose = start_pose.clone() + offsets = torch.tensor( + [ + TARGET_OFFSETS[env_id % len(TARGET_OFFSETS)] + for env_id in range(args.num_envs) + ], + dtype=target_pose.dtype, + device=target_pose.device, + ) + target_pose[:, :3, 3] += offsets + poses = torch.stack( + [ + torch.lerp(start_pose, target_pose, t) + for t in torch.linspace( + 0.0, + 1.0, + args.num_steps, + device=robot.device, ) - ) - - # Draw the whole local_axis trajectory as a single call - local_axes = np.stack( - [item["local_axis"] for item in draw_data[env_id]], axis=0 - ) # (num_steps, 4, 4) or (num_steps, 3) + ], + dim=1, + ) - sim.draw_marker( - cfg=MarkerCfg( - name=f"local_axis_env{env_id}_trajectory", - marker_type="axis", - axis_xpos=local_axes, - axis_size=0.002, - axis_len=0.005, - arena_index=env_id, - ) + qpos_history = [] + success_history = [] + seed = qpos + started_at = time.perf_counter() + for pose in poses.unbind(dim=1): + success, solution = robot.compute_ik( + pose=pose, + joint_seed=seed, + name=arm_name, ) + seed = solution[:, 0, :] if solution.dim() == 3 else solution + qpos_history.append(seed.clone()) + success_history.append(torch.as_tensor(success, device=robot.device)) + print( + f"Solved {args.num_steps} x {args.num_envs} differential IK targets in " + f"{time.perf_counter() - started_at:.4f} seconds" + ) - # Optionally, set qpos for each step (replay or animate) - for step in range(num_steps): - ik_qpos_new = ik_qpos_results[step] - res = ik_success_flags[step] - # Only set qpos for successful IK results - if isinstance(res, (list, np.ndarray, torch.Tensor)): - for env_id, success in enumerate(res): - if success: - q = ( - ik_qpos_new[env_id] - if ik_qpos_new.dim() == 3 - else ik_qpos_new[env_id] - ) - robot.set_qpos( - qpos=q, - joint_ids=robot.get_joint_ids(arm_name), - env_ids=[env_id], + final_pose = robot.compute_fk( + qpos=qpos_history[-1], + name=arm_name, + to_matrix=True, + ) + if not args.no_vis_eef_axis: + for env_id in range(args.num_envs): + for suffix, pose in ( + ("target", target_pose[env_id]), + ("result", final_pose[env_id]), + ): + sim.draw_marker( + cfg=MarkerCfg( + name=f"differential_{suffix}_{env_id}", + marker_type="axis", + axis_xpos=pose, + axis_size=0.002, + axis_len=0.005, + arena_index=env_id, + ) ) - else: - # fallback: set all - if ik_qpos_new.dim() == 3: - robot.set_qpos( - qpos=ik_qpos_new[:, 0, :], joint_ids=robot.get_joint_ids(arm_name) - ) - else: + + joint_ids = robot.get_joint_ids(arm_name) + for solution, success in zip(qpos_history, success_history): + success_ids = success.nonzero(as_tuple=True)[0] + if success_ids.numel(): robot.set_qpos( - qpos=ik_qpos_new, joint_ids=robot.get_joint_ids(arm_name) + solution[success_ids], + joint_ids=joint_ids, + env_ids=success_ids, ) - time.sleep(0.005) - - embed(header="Test DifferentialSolver example. Press Ctrl+D to exit.") + sim.update(step=1) + maybe_wait_for_user(args, "Press Enter to exit...") + finally: + shutdown_sim(sim) if __name__ == "__main__": - main(visualize=True) + main() diff --git a/examples/sim/solvers/neural_ik_solver.py b/examples/sim/solvers/neural_ik_solver.py index 0aa781b6d..1727e2314 100644 --- a/examples/sim/solvers/neural_ik_solver.py +++ b/examples/sim/solvers/neural_ik_solver.py @@ -13,233 +13,176 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- + +"""Demonstrate batched neural IK on Franka Panda.""" + +from __future__ import annotations + import argparse -import math import time import numpy as np import torch -from IPython import embed from embodichain.data.assets.solver_assets import download_neural_ik_checkpoint -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import MarkerCfg -from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.robots.franka_panda import FrankaPandaCfg from embodichain.lab.sim.solvers import NeuralIKSolverCfg - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="NeuralIKSolver example") - parser.add_argument( - "--device", - type=str, - default="cpu", - choices=["cpu", "cuda"], - help="Compute device for tensors and the neural IK solver (default: cpu).", +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, + shutdown_sim, +) + +TARGET_OFFSETS = ( + (0.3, 0.4, -0.2), + (0.2, 0.0, 0.0), + (0.0, 0.2, 0.0), + (0.0, -0.2, -0.1), + (-0.2, 0.0, 0.0), + (0.0, -0.2, 0.0), + (0.0, 0.0, -0.15), + (-0.2, 0.2, 0.0), + (0.0, 0.2, -0.15), +) + + +def main() -> None: + """Solve and replay one neural-IK Cartesian path per environment.""" + parser = add_demo_args( + argparse.ArgumentParser(description="Run the NeuralIKSolver example.") ) parser.add_argument( - "--num-envs", + "--num_steps", + "--num-steps", type=int, - default=1, - help="Number of parallel environments to simulate. IK is solved for all " - "environments simultaneously at each step (default: 1).", + default=50, + help="Number of Cartesian interpolation steps.", ) - return parser.parse_args() - - -def _resolve_device(device: str) -> str: - if device == "cuda" and not torch.cuda.is_available(): - raise RuntimeError( - "CUDA was requested but is not available. Use --device cpu or install " - "a CUDA-enabled PyTorch build." - ) - return device - - -def _squeeze_ik_qpos(ik_qpos: torch.Tensor) -> torch.Tensor: - """Normalize IK output to (num_envs, dof).""" - if ik_qpos.dim() == 3: - return ik_qpos[:, 0, :] - return ik_qpos - - -def _pose_with_arena_offset( - pose: torch.Tensor | np.ndarray, arena_offset: torch.Tensor -) -> np.ndarray: - """Convert arena-local 4x4 pose to world frame by adding arena translation.""" - if isinstance(pose, torch.Tensor): - xpos = pose.detach().cpu().numpy() - else: - xpos = np.asarray(pose) - xpos = np.array(xpos, copy=True, dtype=np.float64) - offset = arena_offset.detach().cpu().numpy().reshape(3) - if xpos.ndim == 2: - xpos[:3, 3] += offset - elif xpos.ndim == 3: - xpos[:, :3, 3] += offset - return xpos - - -def main(): - args = parse_args() - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) - - sim_device = _resolve_device(args.device) - num_envs = args.num_envs - - config = SimulationManagerCfg( - headless=True, - sim_device=sim_device, - num_envs=num_envs, - arena_space=2.0, - ) - sim = SimulationManager(config) + args = parser.parse_args() + if args.device.startswith("cuda") and not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is not available.") + setup_print_options() + # Download before allocating simulation resources so download errors do + # not leave a partially initialized renderer behind. checkpoint_path = download_neural_ik_checkpoint() - - cfg = FrankaPandaCfg.from_dict({"robot_type": "panda"}) - cfg.solver_cfg["arm"] = NeuralIKSolverCfg( - end_link_name="fr3_hand_tcp", - root_link_name="base", - tcp=[ - [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.0], - [0.0, 0.0, 0.0, 1.0], - ], - checkpoint_path=checkpoint_path, - num_arm_joints=7, - max_steps=30, - action_scale=0.2, - hidden_dims=[256, 256], - pos_eps=0.1, - rot_eps=0.5, - ) - - robot: Robot = sim.add_robot(cfg=cfg) - - sim.open_window() - - arm_name = "arm" - device = robot.device - - seed_qpos = torch.tensor( - [0.0, -np.pi / 4, 0.0, -3 * np.pi / 4, 0.0, np.pi / 2, np.pi / 4], - dtype=torch.float32, - device=device, - ) - qpos = seed_qpos.unsqueeze(0).expand(num_envs, -1).clone() - robot.set_qpos(qpos=qpos, joint_ids=robot.get_joint_ids(arm_name)) - time.sleep(3.0) - - fk_xpos = robot.compute_fk(qpos=qpos, name=arm_name, to_matrix=True) - print(f"fk_xpos shape: {tuple(fk_xpos.shape)}") - - start_pose = fk_xpos.clone() - end_pose = fk_xpos.clone() - - # Per-environment target offsets (cycle if num_envs exceeds preset count) - move_vecs = torch.tensor( - [ - [0.3, 0.4, -0.2], - [0.2, 0.0, 0.0], - [0.0, 0.2, 0.0], - [0.0, -0.2, -0.1], - [-0.2, 0.0, 0.0], - [0.0, -0.2, 0.0], - [0.0, 0.0, -0.15], - [-0.2, 0.2, 0.0], - [0.0, 0.2, -0.15], - ], - dtype=torch.float32, - device=device, - ) - for env_id in range(num_envs): - end_pose[env_id, :3, 3] += move_vecs[env_id % move_vecs.shape[0]] - - num_steps = 50 - interpolated_poses = torch.stack( - [ - torch.lerp(start_pose, end_pose, t) - for t in torch.linspace(0.0, 1.0, num_steps, device=device) - ], - dim=1, - ) - - ik_qpos = qpos.clone() - ik_qpos_results: list[torch.Tensor] = [] - ik_success_flags: list[torch.Tensor] = [] - - print( - f"\nRunning {num_steps} batch IK steps: num_envs={num_envs}, device='{sim_device}' ..." + sim = create_default_sim( + args, + num_envs=args.num_envs, + arena_space=2.0, + add_default_light=False, ) - ik_compute_begin = time.time() - for step in range(num_steps): - poses = interpolated_poses[:, step, :, :] - res, ik_qpos_new = robot.compute_ik( - pose=poses, joint_seed=ik_qpos, name=arm_name + try: + cfg = FrankaPandaCfg.from_dict({"robot_type": "panda"}) + cfg.solver_cfg["arm"] = NeuralIKSolverCfg( + end_link_name="fr3_hand_tcp", + root_link_name="base", + tcp=np.eye(4).tolist(), + checkpoint_path=checkpoint_path, + num_arm_joints=7, + max_steps=30, + action_scale=0.2, + hidden_dims=[256, 256], + pos_eps=0.1, + rot_eps=0.5, ) - ik_qpos = _squeeze_ik_qpos(ik_qpos_new) - ik_qpos_results.append(ik_qpos.clone()) - ik_success_flags.append(res) - ik_compute_end = time.time() - print( - f"IK compute time for {num_steps} steps and {num_envs} envs: " - f"{ik_compute_end - ik_compute_begin:.4f}s" - ) - - # Draw target and achieved EE axes for each environment (final step) - final_step = num_steps - 1 - final_ik_qpos = ik_qpos_results[final_step] - final_res = ik_success_flags[final_step] - ik_xpos_all = robot.compute_fk(qpos=final_ik_qpos, name=arm_name, to_matrix=True) - arena_offsets = sim.arena_offsets - - for env_id in range(num_envs): - target_axis = _pose_with_arena_offset(end_pose[env_id], arena_offsets[env_id]) - sim.draw_marker( - cfg=MarkerCfg( - name=f"fk_target_env{env_id}", - marker_type="axis", - axis_xpos=target_axis, - axis_size=0.002, - axis_len=0.005, - arena_index=-1, - ) + robot = sim.add_robot(cfg=cfg) + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) + + arm_name = "arm" + qpos = torch.tensor( + [0.0, -np.pi / 4, 0.0, -3 * np.pi / 4, 0.0, np.pi / 2, np.pi / 4], + dtype=torch.float32, + device=robot.device, + ).repeat(args.num_envs, 1) + robot.set_qpos(qpos, joint_ids=robot.get_joint_ids(arm_name)) + + start_pose = robot.compute_fk( + qpos=qpos, + name=arm_name, + to_matrix=True, ) - - if final_res[env_id]: - ik_axis = _pose_with_arena_offset( - ik_xpos_all[env_id], arena_offsets[env_id] - ) - sim.draw_marker( - cfg=MarkerCfg( - name=f"ik_result_env{env_id}", - marker_type="axis", - axis_xpos=ik_axis, - axis_size=0.002, - axis_len=0.005, - arena_index=-1, + target_pose = start_pose.clone() + target_pose[:, :3, 3] += torch.tensor( + [ + TARGET_OFFSETS[env_id % len(TARGET_OFFSETS)] + for env_id in range(args.num_envs) + ], + dtype=target_pose.dtype, + device=target_pose.device, + ) + poses = torch.stack( + [ + torch.lerp(start_pose, target_pose, t) + for t in torch.linspace( + 0.0, + 1.0, + args.num_steps, + device=robot.device, ) - ) + ], + dim=1, + ) - # Animate: batch-apply IK qpos for successful envs, then step simulation - joint_ids = robot.get_joint_ids(arm_name) - for step in range(num_steps): - ik_qpos_step = ik_qpos_results[step] - res = ik_success_flags[step] - if res.any(): - success_ids = res.nonzero(as_tuple=True)[0] - robot.set_qpos( - qpos=ik_qpos_step[success_ids], - joint_ids=joint_ids, - env_ids=success_ids, + qpos_history = [] + success_history = [] + seed = qpos + started_at = time.perf_counter() + for pose in poses.unbind(dim=1): + success, solution = robot.compute_ik( + pose=pose, + joint_seed=seed, + name=arm_name, ) - sim.update(step=5) + seed = solution[:, 0, :] if solution.dim() == 3 else solution + qpos_history.append(seed.clone()) + success_history.append(torch.as_tensor(success, device=robot.device)) + print( + f"Solved {args.num_steps} x {args.num_envs} neural IK targets in " + f"{time.perf_counter() - started_at:.4f} seconds" + ) - embed(header="NeuralIKSolver example. Press Ctrl+D to exit.") + final_pose = robot.compute_fk( + qpos=qpos_history[-1], + name=arm_name, + to_matrix=True, + ) + if not args.no_vis_eef_axis: + for env_id in range(args.num_envs): + for suffix, pose in ( + ("target", target_pose[env_id]), + ("result", final_pose[env_id]), + ): + sim.draw_marker( + cfg=MarkerCfg( + name=f"neural_ik_{suffix}_{env_id}", + marker_type="axis", + axis_xpos=pose, + axis_size=0.002, + axis_len=0.005, + arena_index=env_id, + ) + ) + + joint_ids = robot.get_joint_ids(arm_name) + for solution, success in zip(qpos_history, success_history): + success_ids = success.nonzero(as_tuple=True)[0] + if success_ids.numel(): + robot.set_qpos( + solution[success_ids], + joint_ids=joint_ids, + env_ids=success_ids, + ) + sim.update(step=5) + maybe_wait_for_user(args, "Press Enter to exit...") + finally: + shutdown_sim(sim) if __name__ == "__main__": diff --git a/examples/sim/solvers/opw_solver.py b/examples/sim/solvers/opw_solver.py index e8ae222ca..62dce0f7d 100644 --- a/examples/sim/solvers/opw_solver.py +++ b/examples/sim/solvers/opw_solver.py @@ -13,181 +13,119 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -import os + +"""Demonstrate OPW FK/IK on both CobotMagic arms.""" + +from __future__ import annotations + +import argparse import time -import torch + import numpy as np -from IPython import embed +import torch -from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim.robots import CobotMagicCfg -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import MarkerCfg - - -def main(): - # Set print options for better readability - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) - - # Initialize simulation - sim_device = "cpu" - config = SimulationManagerCfg(headless=False, sim_device=sim_device) - sim = SimulationManager(config) - sim.set_manual_update(False) - - # Robot configuration dictionary - cfg_dict = { - "uid": "CobotMagic", - "init_pos": [0.0, 0.0, 0.7775], - "init_qpos": [ - -0.3, - 0.3, - 1.0, - 1.0, - -1.2, - -1.2, - 0.0, - 0.0, - 0.6, - 0.6, - 0.0, - 0.0, - 0.05, - 0.05, - 0.05, - 0.05, - ], - "solver_cfg": { - "left_arm": { - "class_type": "OPWSolver", - "end_link_name": "left_link6", - "root_link_name": "left_arm_base", - "tcp": [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0.143], [0, 0, 0, 1]], - }, - "right_arm": { - "class_type": "OPWSolver", - "end_link_name": "right_link6", - "root_link_name": "right_arm_base", - "tcp": [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0.143], [0, 0, 0, 1]], - }, - }, - } - - # Add robot to simulation - robot: Robot = sim.add_robot(cfg=CobotMagicCfg.from_dict(cfg_dict)) - - # Left arm control - arm_name = "left_arm" - # Set initial joint positions for left arm - qpos_seed = torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32) - qpos_fk = torch.tensor( - [[0.0, np.pi / 4, -np.pi / 4, 0.0, np.pi / 4, 0.0]], dtype=torch.float32 - ) - fk_xpos = robot.compute_fk(qpos=qpos_fk, name=arm_name, to_matrix=True) - - link_pose = robot._entities[0].get_link_pose("left_base_link") - link_pose_tensor = torch.from_numpy(link_pose).to( - fk_xpos.device, dtype=fk_xpos.dtype - ) - - # Solve IK for the left arm - res, ik_qpos = robot.compute_ik(pose=fk_xpos, name=arm_name, joint_seed=qpos_seed) - - # Measure IK computation time and visualize result - a = time.time() - if ik_qpos.dim() == 3: - ik_xpos = robot.compute_fk(qpos=ik_qpos[0][0], name=arm_name, to_matrix=True) - else: - ik_xpos = robot.compute_fk(qpos=ik_qpos, name=arm_name, to_matrix=True) - b = time.time() - print(f"Left arm IK computation time: {b-a:.6f} seconds") - - # Visualize the result in simulation - sim.draw_marker( - cfg=MarkerCfg( - name="fk_xpos_left", - marker_type="axis", - axis_xpos=np.array(fk_xpos.tolist()), - axis_size=0.002, - axis_len=0.005, - ) - ) - - sim.draw_marker( - cfg=MarkerCfg( - name="ik_xpos_left", - marker_type="axis", - axis_xpos=np.array(ik_xpos.tolist()), - axis_size=0.002, - axis_len=0.005, - ) - ) - - # Move robot to IK result joint positions - if ik_qpos.dim() == 3: - robot.set_qpos(qpos=ik_qpos[0][0], joint_ids=robot.get_joint_ids(arm_name)) - else: - robot.set_qpos(qpos=ik_qpos, joint_ids=robot.get_joint_ids(arm_name)) - - # Right arm control - arm_name_r = "right_arm" - # Set initial joint positions for right arm - qpos_seed_r = torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32) - qpos_fk_r = torch.tensor( - [[0.0, np.pi / 4, -np.pi / 4, 0.0, np.pi / 4, 0.0]], dtype=torch.float32 - ) - fk_xpos_r = robot.compute_fk(qpos=qpos_fk_r, name=arm_name_r, to_matrix=True) - - link_pose_r = robot._entities[0].get_link_pose("right_base_link") - link_pose_tensor_r = torch.from_numpy(link_pose_r).to( - fk_xpos_r.device, dtype=fk_xpos_r.dtype - ) - - # Solve IK for the right arm - res_r, ik_qpos_r = robot.compute_ik( - pose=fk_xpos_r, name=arm_name_r, joint_seed=qpos_seed_r - ) - - # Measure IK computation time and visualize result - a_r = time.time() - if ik_qpos_r.dim() == 3: - ik_xpos_r = robot.compute_fk( - qpos=ik_qpos_r[0][0], name=arm_name_r, to_matrix=True - ) - else: - ik_xpos_r = robot.compute_fk(qpos=ik_qpos_r, name=arm_name_r, to_matrix=True) - b_r = time.time() - print(f"Right arm IK computation time: {b_r-a_r:.6f} seconds") - - # Visualize the result in simulation - sim.draw_marker( - cfg=MarkerCfg( - name="fk_xpos_right", - marker_type="axis", - axis_xpos=np.array(fk_xpos_r.tolist()), - axis_size=0.002, - axis_len=0.005, - ) - ) - - sim.draw_marker( - cfg=MarkerCfg( - name="ik_xpos_right", - marker_type="axis", - axis_xpos=np.array(ik_xpos_r.tolist()), - axis_size=0.002, - axis_len=0.005, +from embodichain.lab.sim.robots import CobotMagicCfg +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, + shutdown_sim, +) + + +def _first_solution(qpos: torch.Tensor) -> torch.Tensor: + """Normalize analytic IK output to one solution per environment.""" + return qpos[:, 0, :] if qpos.dim() == 3 else qpos + + +def main() -> None: + """Solve and visualize one pose for each arm.""" + args = add_demo_args( + argparse.ArgumentParser(description="Run the OPW solver example.") + ).parse_args() + setup_print_options() + + sim = create_default_sim(args, add_default_light=False) + try: + robot = sim.add_robot( + cfg=CobotMagicCfg.from_dict( + { + "uid": "CobotMagic", + "init_pos": [0.0, 0.0, 0.7775], + "solver_cfg": { + arm: { + "class_type": "OPWSolver", + "end_link_name": f"{side}_link6", + "root_link_name": f"{side}_arm_base", + "tcp": [ + [1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 1, 0.143], + [0, 0, 0, 1], + ], + } + for arm, side in ( + ("left_arm", "left"), + ("right_arm", "right"), + ) + }, + } + ) ) - ) + maybe_open_window(sim, args) - # Move robot to IK result joint positions - if ik_qpos_r.dim() == 3: - robot.set_qpos(qpos=ik_qpos_r[0][0], joint_ids=robot.get_joint_ids(arm_name_r)) - else: - robot.set_qpos(qpos=ik_qpos_r, joint_ids=robot.get_joint_ids(arm_name_r)) - - embed(header="Test OPWSolver example. Press Ctrl+D to exit.") + target_qpos = torch.tensor( + [[0.0, np.pi / 4, -np.pi / 4, 0.0, np.pi / 4, 0.0]], + dtype=torch.float32, + device=robot.device, + ) + seed = torch.zeros_like(target_qpos) + for arm_name in ("left_arm", "right_arm"): + target_pose = robot.compute_fk( + qpos=target_qpos, + name=arm_name, + to_matrix=True, + ) + started_at = time.perf_counter() + success, solutions = robot.compute_ik( + pose=target_pose, + name=arm_name, + joint_seed=seed, + ) + elapsed = time.perf_counter() - started_at + if not torch.as_tensor(success).all(): + raise RuntimeError(f"OPW IK failed for {arm_name}.") + + solution = _first_solution(solutions) + achieved_pose = robot.compute_fk( + qpos=solution, + name=arm_name, + to_matrix=True, + ) + robot.set_qpos(solution, joint_ids=robot.get_joint_ids(arm_name)) + print(f"{arm_name} IK computation time: {elapsed:.6f} seconds") + + for suffix, pose in ( + ("target", target_pose), + ("result", achieved_pose), + ): + sim.draw_marker( + cfg=MarkerCfg( + name=f"{arm_name}_{suffix}", + marker_type="axis", + axis_xpos=pose, + axis_size=0.002, + axis_len=0.005, + ) + ) + + sim.update(step=1) + maybe_wait_for_user(args, "Press Enter to exit...") + finally: + shutdown_sim(sim) if __name__ == "__main__": diff --git a/examples/sim/solvers/pink_solver.py b/examples/sim/solvers/pink_solver.py index 6308b6128..173a9d24c 100644 --- a/examples/sim/solvers/pink_solver.py +++ b/examples/sim/solvers/pink_solver.py @@ -13,161 +13,135 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -import os -import time -import numpy as np -import torch -from IPython import embed - -from embodichain.data import get_data_path -from embodichain.lab.sim.cfg import RobotCfg -from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import MarkerCfg - -def main(): - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) +"""Demonstrate iterative Pink IK along a Cartesian path.""" - # Set up simulation with specified device (CPU or CUDA) - sim_device = "cpu" - config = SimulationManagerCfg(headless=False, sim_device=sim_device) - sim = SimulationManager(config) - sim.set_manual_update(False) +from __future__ import annotations - # Load robot URDF file - urdf = get_data_path("Rokae/SR5/SR5.urdf") - - assert os.path.isfile(urdf) - - cfg_dict = { - "fpath": urdf, - "control_parts": { - "main_arm": [ - "joint1", - "joint2", - "joint3", - "joint4", - "joint5", - "joint6", - ], - }, - "solver_cfg": { - "main_arm": { - "class_type": "PinkSolver", - "end_link_name": "ee_link", - "root_link_name": "base_link", - }, - }, - } - - robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) - - # Define a sample target pose as a 1x4x4 homogeneous matrix - rad = torch.deg2rad(torch.tensor(45.0)) +import argparse +import time - arm_name = "main_arm" - fk_qpos = torch.full((1, 6), rad, dtype=torch.float32, device="cpu") +import numpy as np +import torch - # Set initial joint positions - qpos = torch.from_numpy(np.array([0.0, 0.0, np.pi / 2, 0.0, np.pi / 2, 0.0])).to( - fk_qpos.device +from embodichain.data import get_data_path +from embodichain.lab.sim.cfg import MarkerCfg, RobotCfg +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, + shutdown_sim, +) + + +def main() -> None: + """Solve a short Cartesian path with Pink.""" + parser = add_demo_args( + argparse.ArgumentParser(description="Run the Pink solver example.") ) - qpos = qpos.unsqueeze(0) - robot.set_qpos(qpos=qpos, joint_ids=robot.get_joint_ids("main_arm")) - import time - - time.sleep(3.0) - fk_xpos = robot.compute_fk(qpos=qpos, name=arm_name, to_matrix=True) - print(f"fk_xpos: {fk_xpos}") - start_pose = fk_xpos.clone()[0] # Start pose - end_pose = fk_xpos.clone()[0] # End pose - - end_pose[:3, 3] = end_pose[:3, 3][:3] + torch.tensor( - [0.0, 0.4, 0.0], device=fk_xpos.device + parser.add_argument( + "--num_steps", + "--num-steps", + type=int, + default=100, + help="Number of Cartesian interpolation steps.", ) - - num_steps = 100 - - # Interpolate poses between start and end - interpolated_poses = [ - torch.lerp(start_pose, end_pose, t) for t in np.linspace(0, 1, num_steps) - ] - - ik_qpos = qpos - - qpos = ik_qpos - res, ik_qpos = robot.compute_ik(pose=end_pose, joint_seed=qpos, name=arm_name) - import time - - a = time.time() - if ik_qpos.dim() == 3: - ik_xpos = robot.compute_fk(qpos=ik_qpos[0][0], name=arm_name, to_matrix=True) - else: - ik_xpos = robot.compute_fk(qpos=ik_qpos, name=arm_name, to_matrix=True) - b = time.time() - print(f"ik time: {b-a}") - - ik_xpos = ik_xpos - - sim.draw_marker( - cfg=MarkerCfg( - name="fk_xpos", - marker_type="axis", - axis_xpos=np.array(end_pose.tolist()), - axis_size=0.002, - axis_len=0.005, + args = parser.parse_args() + setup_print_options() + + sim = create_default_sim(args, add_default_light=False) + try: + robot = sim.add_robot( + cfg=RobotCfg.from_dict( + { + "fpath": get_data_path("Rokae/SR5/SR5.urdf"), + "control_parts": { + "main_arm": [f"joint{i}" for i in range(1, 7)], + }, + "solver_cfg": { + "main_arm": { + "class_type": "PinkSolver", + "end_link_name": "ee_link", + "root_link_name": "base_link", + }, + }, + } + ) ) - ) + maybe_open_window(sim, args) - sim.draw_marker( - cfg=MarkerCfg( - name="ik_xpos", - marker_type="axis", - axis_xpos=np.array(ik_xpos.tolist()), - axis_size=0.002, - axis_len=0.005, + arm_name = "main_arm" + qpos = torch.tensor( + [[0.0, 0.0, np.pi / 2, 0.0, np.pi / 2, 0.0]], + dtype=torch.float32, + device=robot.device, ) - ) - - for i, pose in enumerate(interpolated_poses): - print(f"Step {i}: Moving to pose:\n{pose}") - start_time = time.time() - res, ik_qpos = robot.compute_ik(pose=pose, joint_seed=ik_qpos, name=arm_name) - end_time = time.time() - compute_time = end_time - start_time - print(f"Step {i}: IK computation time: {compute_time:.6f} seconds") - - print(f"IK result: {res}, ik_qpos: {ik_qpos}") - if not res: - print(f"Step {i}: IK failed for pose:\n{pose}") - continue - - # Set robot joint positions - if ik_qpos.dim() == 3: - robot.set_qpos(qpos=ik_qpos[0][0], joint_ids=robot.get_joint_ids(arm_name)) - else: - robot.set_qpos(qpos=ik_qpos, joint_ids=robot.get_joint_ids(arm_name)) - - # Visualize current pose - ik_xpos = robot.compute_fk(qpos=ik_qpos, name=arm_name, to_matrix=True) - ik_xpos = ik_xpos - + robot.set_qpos(qpos, joint_ids=robot.get_joint_ids(arm_name)) + start_pose = robot.compute_fk( + qpos=qpos, + name=arm_name, + to_matrix=True, + ) + target_pose = start_pose.clone() + target_pose[:, 1, 3] += 0.4 sim.draw_marker( cfg=MarkerCfg( - name=f"ik_xpos_step_{i}", + name="pink_target", marker_type="axis", - axis_xpos=np.array(ik_xpos.tolist()), + axis_xpos=target_pose, axis_size=0.002, axis_len=0.005, ) ) - # Add delay to simulate motion - time.sleep(0.005) + poses = torch.stack( + [ + torch.lerp(start_pose, target_pose, t) + for t in torch.linspace( + 0.0, + 1.0, + args.num_steps, + device=robot.device, + ) + ], + dim=1, + ) + started_at = time.perf_counter() + for pose in poses.unbind(dim=1): + success, solution = robot.compute_ik( + pose=pose, + joint_seed=qpos, + name=arm_name, + ) + if not torch.as_tensor(success).all(): + raise RuntimeError("Pink IK failed along the Cartesian path.") + qpos = solution[:, 0, :] if solution.dim() == 3 else solution + robot.set_qpos(qpos, joint_ids=robot.get_joint_ids(arm_name)) + sim.update(step=1) + print( + f"Solved {args.num_steps} Pink IK steps in " + f"{time.perf_counter() - started_at:.6f} seconds" + ) - embed(header="Test PinkSolver example. Press Ctrl+D to exit.") + achieved_pose = robot.compute_fk( + qpos=qpos, + name=arm_name, + to_matrix=True, + ) + sim.draw_marker( + cfg=MarkerCfg( + name="pink_result", + marker_type="axis", + axis_xpos=achieved_pose, + axis_size=0.002, + axis_len=0.005, + ) + ) + maybe_wait_for_user(args, "Press Enter to exit...") + finally: + shutdown_sim(sim) if __name__ == "__main__": diff --git a/examples/sim/solvers/pinocchio_solver.py b/examples/sim/solvers/pinocchio_solver.py index 6d70305e5..33f512da2 100644 --- a/examples/sim/solvers/pinocchio_solver.py +++ b/examples/sim/solvers/pinocchio_solver.py @@ -13,113 +13,106 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -import os -import time -import numpy as np -import torch -from IPython import embed - -from embodichain.data import get_data_path -from embodichain.lab.sim.cfg import RobotCfg -from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import MarkerCfg - - -def main(): - # Set print options for better readability - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) - - # Initialize simulation - sim_device = "cpu" - config = SimulationManagerCfg(headless=False, sim_device=sim_device) - sim = SimulationManager(config) - sim.set_manual_update(False) - - # Load robot URDF file - urdf = get_data_path("DexforceW1V021/DexforceW1_v02_1.urdf") - assert os.path.isfile(urdf) - # Robot configuration dictionary - cfg_dict = { - "fpath": urdf, - "control_parts": { - "left_arm": [f"LEFT_J{i+1}" for i in range(7)], - "right_arm": [f"RIGHT_J{i+1}" for i in range(7)], - }, - "solver_cfg": { - "left_arm": { - "class_type": "PinocchioSolver", - "end_link_name": "left_ee", - "root_link_name": "left_arm_base", - }, - "right_arm": { - "class_type": "PinocchioSolver", - "end_link_name": "right_ee", - "root_link_name": "right_arm_base", - }, - }, - } +"""Demonstrate Pinocchio FK/IK on a DexForce W1 arm.""" - robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) - arm_name = "left_arm" - # Set initial joint positions for left arm - qpos_seed = torch.tensor( - [[0.0, 0.1, 0.0, -np.pi / 4, 0.0, 0.0, 0.0]], dtype=torch.float32 - ) - qpos_fk = torch.tensor( - [[0.0, 0.0, 0.0, -np.pi / 4, 0.0, 0.0, 0.0]], dtype=torch.float32 - ) - fk_xpos = robot.compute_fk(qpos=qpos_fk, name=arm_name, to_matrix=True) - link_pose = robot._entities[0].get_link_pose("left_arm_base") - link_pose_tensor = torch.from_numpy(link_pose).to( - fk_xpos.device, dtype=fk_xpos.dtype - ) +from __future__ import annotations - # Solve IK for the left arm - res, ik_qpos = robot.compute_ik(pose=fk_xpos, name=arm_name, joint_seed=qpos_seed) - - # Measure IK computation time and visualize result - a = time.time() - if ik_qpos.dim() == 3: - ik_xpos = robot.compute_fk(qpos=ik_qpos[0][0], name=arm_name, to_matrix=True) - else: - ik_xpos = robot.compute_fk(qpos=ik_qpos, name=arm_name, to_matrix=True) - b = time.time() - print(f"IK computation time: {b-a:.6f} seconds") +import argparse +import time - fk_xpos = link_pose_tensor @ fk_xpos - ik_xpos = link_pose_tensor @ ik_xpos +import numpy as np +import torch - # Visualize the result in simulation - sim.draw_marker( - cfg=MarkerCfg( - name="fk_xpos", - marker_type="axis", - axis_xpos=np.array(fk_xpos.tolist()), - axis_size=0.002, - axis_len=0.005, +from embodichain.data import get_data_path +from embodichain.lab.sim.cfg import MarkerCfg, RobotCfg +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, + shutdown_sim, +) + + +def main() -> None: + """Solve one left-arm pose and compare FK with the IK result.""" + args = add_demo_args( + argparse.ArgumentParser(description="Run the Pinocchio solver example.") + ).parse_args() + setup_print_options() + + sim = create_default_sim(args, add_default_light=False) + try: + robot = sim.add_robot( + cfg=RobotCfg.from_dict( + { + "fpath": get_data_path("DexforceW1V021/DexforceW1_v02_1.urdf"), + "control_parts": { + "left_arm": [f"LEFT_J{i + 1}" for i in range(7)], + }, + "solver_cfg": { + "left_arm": { + "class_type": "PinocchioSolver", + "end_link_name": "left_ee", + "root_link_name": "left_arm_base", + }, + }, + } + ) ) - ) + maybe_open_window(sim, args) - sim.draw_marker( - cfg=MarkerCfg( - name="ik_xpos", - marker_type="axis", - axis_xpos=np.array(ik_xpos.tolist()), - axis_size=0.002, - axis_len=0.005, + arm_name = "left_arm" + target_qpos = torch.tensor( + [[0.0, 0.0, 0.0, -np.pi / 4, 0.0, 0.0, 0.0]], + dtype=torch.float32, + device=robot.device, + ) + seed = target_qpos.clone() + seed[:, 1] = 0.1 + target_pose = robot.compute_fk( + qpos=target_qpos, + name=arm_name, + to_matrix=True, ) - ) - - # Move robot to IK result joint positions - if ik_qpos.dim() == 3: - robot.set_qpos(qpos=ik_qpos[0][0], joint_ids=robot.get_joint_ids(arm_name)) - else: - robot.set_qpos(qpos=ik_qpos, joint_ids=robot.get_joint_ids(arm_name)) - embed(header="Test PinocchioSolver example. Press Ctrl+D to exit.") + started_at = time.perf_counter() + success, solution = robot.compute_ik( + pose=target_pose, + name=arm_name, + joint_seed=seed, + ) + elapsed = time.perf_counter() - started_at + if not torch.as_tensor(success).all(): + raise RuntimeError("Pinocchio IK failed.") + if solution.dim() == 3: + solution = solution[:, 0, :] + + achieved_pose = robot.compute_fk( + qpos=solution, + name=arm_name, + to_matrix=True, + ) + robot.set_qpos(solution, joint_ids=robot.get_joint_ids(arm_name)) + print(f"IK computation time: {elapsed:.6f} seconds") + + for suffix, pose in (("target", target_pose), ("result", achieved_pose)): + sim.draw_marker( + cfg=MarkerCfg( + name=f"pinocchio_{suffix}", + marker_type="axis", + axis_xpos=pose, + axis_size=0.002, + axis_len=0.005, + ) + ) + + sim.update(step=1) + maybe_wait_for_user(args, "Press Enter to exit...") + finally: + shutdown_sim(sim) if __name__ == "__main__": diff --git a/examples/sim/solvers/pytorch_solver.py b/examples/sim/solvers/pytorch_solver.py index 5d954ff61..77141e008 100644 --- a/examples/sim/solvers/pytorch_solver.py +++ b/examples/sim/solvers/pytorch_solver.py @@ -1,224 +1,184 @@ -import os +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Demonstrate batched PyTorch IK on a DexForce W1 arm.""" + +from __future__ import annotations + +import argparse import time + import numpy as np import torch -from IPython import embed from embodichain.data import get_data_path -from embodichain.lab.sim.cfg import RobotCfg -from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import MarkerCfg - - -def main(): - # Set numpy and torch print options for better readability - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) - - # Initialize simulation environment (CPU or CUDA) - sim_device = "cpu" - num_envs = 9 # Number of parallel environments - config = SimulationManagerCfg( - headless=False, sim_device=sim_device, arena_space=2.0, num_envs=num_envs - ) - sim = SimulationManager(config) - sim.set_manual_update(False) - - # Load robot URDF file - urdf = get_data_path("DexforceW1V021/DexforceW1_v02_1.urdf") - assert os.path.isfile(urdf) - - # Robot configuration dictionary - cfg_dict = { - "fpath": urdf, - "control_parts": { - "left_arm": [ - "LEFT_J1", - "LEFT_J2", - "LEFT_J3", - "LEFT_J4", - "LEFT_J5", - "LEFT_J6", - "LEFT_J7", - ], - }, - "solver_cfg": { - "left_arm": { - "class_type": "PytorchSolver", - "end_link_name": "left_ee", - "root_link_name": "left_arm_base", - }, - }, - } - - # Add robot to simulation - robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) - - # Prepare initial joint positions for all environments - arm_name = "left_arm" - qpos = ( - torch.tensor([0.0, 0.0, 0.0, -np.pi / 2, 0.0, 0.0, 0.0], dtype=torch.float32) - .unsqueeze(0) - .repeat(num_envs, 1) +from embodichain.lab.sim.cfg import MarkerCfg, RobotCfg +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, + shutdown_sim, +) + +TARGET_OFFSETS = ( + (0.2, 0.0, 0.0), + (0.0, 0.2, 0.0), + (0.0, -0.2, -0.5), + (-0.2, 0.0, 0.0), + (-0.2, 0.0, 0.0), + (0.0, -0.2, 0.0), + (0.0, 0.0, -0.5), + (-0.2, 0.2, 0.0), + (0.0, 0.2, -0.5), +) + + +def main() -> None: + """Solve and replay one batched Cartesian path.""" + parser = add_demo_args( + argparse.ArgumentParser(description="Run batched PyTorch IK.") ) - robot.set_qpos(qpos=qpos, joint_ids=robot.get_joint_ids(arm_name)) - - time.sleep(2.0) - fk_xpos = robot.compute_fk( - qpos=qpos, name=arm_name, to_matrix=True - ) # (num_envs, 4, 4) - - # Prepare batch start and end poses for all envs - start_pose = fk_xpos.clone() - end_pose = fk_xpos.clone() - move_vecs = torch.tensor( - [ - [0.2, 0.0, 0.0], - [0.0, 0.2, 0.0], - [0.0, -0.2, -0.5], - [-0.2, 0.0, 0.0], - [-0.2, 0.0, 0.0], - [0.0, -0.2, 0.0], - [0.0, 0.0, -0.5], - [-0.2, 0.2, 0.0], - [0.0, 0.2, -0.5], - ], - dtype=end_pose.dtype, - device=end_pose.device, + parser.set_defaults(num_envs=9) + parser.add_argument( + "--num_steps", + "--num-steps", + type=int, + default=50, + help="Number of Cartesian interpolation steps.", ) - end_pose[:, :3, 3] += move_vecs - - num_steps = 50 - # Interpolate poses for each env - interpolated_poses = torch.stack( - [torch.lerp(start_pose, end_pose, t) for t in np.linspace(0, 1, num_steps)], - dim=1, - ) # (num_envs, num_steps, 4, 4) - - # Initial joint positions for all envs - ik_qpos = qpos.clone() - ik_qpos_results = [] - ik_success_flags = [] - - # Batch IK solving for each step - ik_compute_begin = time.time() - for step in range(num_steps): - poses = interpolated_poses[:, step, :, :] # (num_envs, 4, 4) - if poses.shape[0] != num_envs: - poses = poses.expand(num_envs, *poses.shape[1:]) - if ik_qpos.shape[0] != num_envs: - ik_qpos = ik_qpos.expand(num_envs, *ik_qpos.shape[1:]) - assert poses.shape[0] == num_envs - assert ik_qpos.shape[0] == num_envs - - # Parallel batch IK solving - res, ik_qpos_new = robot.compute_ik( - pose=poses, joint_seed=ik_qpos, name=arm_name - ) - ik_qpos_results.append(ik_qpos_new.clone()) - ik_success_flags.append(res) - ik_qpos = ik_qpos_new # Update joint seed - ik_compute_end = time.time() - print( - f"IK compute time for {num_steps} steps and {num_envs} envs: {ik_compute_end - ik_compute_begin:.4f} seconds" + args = parser.parse_args() + setup_print_options() + + sim = create_default_sim( + args, + num_envs=args.num_envs, + arena_space=2.0, + add_default_light=False, ) - - # Collect visualization data for all steps and environments - draw_data = [[] for _ in range(num_envs)] - for env_id in range(num_envs): - for step in range(num_steps): - ik_qpos_new = ik_qpos_results[step] - ik_xpos = robot.compute_fk(qpos=ik_qpos_new, name=arm_name, to_matrix=True) - local_pose = robot._entities[env_id].get_link_pose("left_arm_base") - if isinstance(local_pose, np.ndarray): - local_pose = torch.from_numpy(local_pose).to( - ik_xpos.device, dtype=ik_xpos.dtype - ) - fk_axis = (local_pose @ end_pose[env_id]).cpu().numpy() - ik_axis = (local_pose @ ik_xpos[env_id]).cpu().numpy() - local_axis = (local_pose @ ik_xpos[env_id]).cpu().numpy() - draw_data[env_id].append( + try: + robot = sim.add_robot( + cfg=RobotCfg.from_dict( { - "step": step, - "fk_axis": fk_axis, - "ik_axis": ik_axis, - "local_axis": local_axis, + "fpath": get_data_path("DexforceW1V021/DexforceW1_v02_1.urdf"), + "control_parts": { + "left_arm": [f"LEFT_J{i}" for i in range(1, 8)], + }, + "solver_cfg": { + "left_arm": { + "class_type": "PytorchSolver", + "end_link_name": "left_ee", + "root_link_name": "left_arm_base", + }, + }, } ) - - # Batch draw: only draw fk_axis and ik_axis once per env, draw local_axis trajectory for all steps - for env_id in range(num_envs): - fk_axis = draw_data[env_id][0]["fk_axis"] - ik_axis = draw_data[env_id][0]["ik_axis"] - - sim.draw_marker( - cfg=MarkerCfg( - name=f"fk_axis_env{env_id}", - marker_type="axis", - axis_xpos=fk_axis, - axis_size=0.002, - axis_len=0.005, - arena_index=env_id, - ) ) - - sim.draw_marker( - cfg=MarkerCfg( - name=f"ik_axis_env{env_id}", - marker_type="axis", - axis_xpos=ik_axis, - axis_size=0.002, - axis_len=0.005, - arena_index=env_id, - ) + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) + + arm_name = "left_arm" + qpos = torch.tensor( + [0.0, 0.0, 0.0, -np.pi / 2, 0.0, 0.0, 0.0], + dtype=torch.float32, + device=robot.device, + ).repeat(args.num_envs, 1) + robot.set_qpos(qpos, joint_ids=robot.get_joint_ids(arm_name)) + start_pose = robot.compute_fk( + qpos=qpos, + name=arm_name, + to_matrix=True, ) - - # Draw the whole local_axis trajectory as a single call (if supported) - local_axes = np.stack( - [item["local_axis"] for item in draw_data[env_id]], axis=0 + target_pose = start_pose.clone() + target_pose[:, :3, 3] += torch.tensor( + [ + TARGET_OFFSETS[env_id % len(TARGET_OFFSETS)] + for env_id in range(args.num_envs) + ], + dtype=target_pose.dtype, + device=target_pose.device, + ) + poses = torch.stack( + [ + torch.lerp(start_pose, target_pose, t) + for t in torch.linspace( + 0.0, + 1.0, + args.num_steps, + device=robot.device, + ) + ], + dim=1, ) - sim.draw_marker( - cfg=MarkerCfg( - name=f"local_axis_env{env_id}_trajectory", - marker_type="axis", - axis_xpos=local_axes, - axis_size=0.002, - axis_len=0.005, - arena_index=env_id, + qpos_history = [] + success_history = [] + seed = qpos + started_at = time.perf_counter() + for pose in poses.unbind(dim=1): + success, solution = robot.compute_ik( + pose=pose, + joint_seed=seed, + name=arm_name, ) + seed = solution[:, 0, :] if solution.dim() == 3 else solution + qpos_history.append(seed.clone()) + success_history.append(torch.as_tensor(success, device=robot.device)) + print( + f"Solved {args.num_steps} x {args.num_envs} PyTorch IK targets in " + f"{time.perf_counter() - started_at:.4f} seconds" ) - # Optionally, set qpos for each step (replay or animate) - for step in range(num_steps): - ik_qpos_new = ik_qpos_results[step] - res = ik_success_flags[step] - if isinstance(res, (list, np.ndarray, torch.Tensor)): - for env_id, success in enumerate(res): - if success: - q = ( - ik_qpos_new[env_id] - if ik_qpos_new.dim() == 3 - else ik_qpos_new[env_id] - ) - robot.set_qpos( - qpos=q, - joint_ids=robot.get_joint_ids(arm_name), - env_ids=[env_id], + final_pose = robot.compute_fk( + qpos=qpos_history[-1], + name=arm_name, + to_matrix=True, + ) + if not args.no_vis_eef_axis: + for env_id in range(args.num_envs): + for suffix, pose in ( + ("target", target_pose[env_id]), + ("result", final_pose[env_id]), + ): + sim.draw_marker( + cfg=MarkerCfg( + name=f"pytorch_{suffix}_{env_id}", + marker_type="axis", + axis_xpos=pose, + axis_size=0.002, + axis_len=0.005, + arena_index=env_id, + ) ) - else: - if ik_qpos_new.dim() == 3: - robot.set_qpos( - qpos=ik_qpos_new[:, 0, :], joint_ids=robot.get_joint_ids(arm_name) - ) - else: + + joint_ids = robot.get_joint_ids(arm_name) + for solution, success in zip(qpos_history, success_history): + success_ids = success.nonzero(as_tuple=True)[0] + if success_ids.numel(): robot.set_qpos( - qpos=ik_qpos_new, joint_ids=robot.get_joint_ids(arm_name) + solution[success_ids], + joint_ids=joint_ids, + env_ids=success_ids, ) - time.sleep(0.005) - - embed(header="Test PytorchSolver batch example. Press Ctrl+D to exit.") + sim.update(step=1) + maybe_wait_for_user(args, "Press Enter to exit...") + finally: + shutdown_sim(sim) if __name__ == "__main__": diff --git a/examples/sim/solvers/srs_solver.py b/examples/sim/solvers/srs_solver.py index 502726dee..407509be5 100644 --- a/examples/sim/solvers/srs_solver.py +++ b/examples/sim/solvers/srs_solver.py @@ -14,70 +14,81 @@ # limitations under the License. # ---------------------------------------------------------------------------- +"""Demonstrate SRS FK/IK on a DexForce W1 arm.""" + +from __future__ import annotations + +import argparse import time + import numpy as np import torch -from IPython import embed - -from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.robots import DexforceW1Cfg - - -def main(): - # Set print options for better readability - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) - - # Initialize simulation - sim_device = "cpu" - sim = SimulationManager( - SimulationManagerCfg( - headless=False, sim_device=sim_device, width=2200, height=1200 - ) - ) - - sim.set_manual_update(False) - - robot: Robot = sim.add_robot(cfg=DexforceW1Cfg.from_dict({"uid": "dexforce_w1"})) - arm_name = "left_arm" - # Set initial joint positions for left arm - qpos_fk_list = [ - torch.tensor([[0.0, 0.0, 0.0, -np.pi / 2, 0.0, 0.0, 0.0]], dtype=torch.float32), - ] - robot.set_qpos(qpos_fk_list[0], joint_ids=robot.get_joint_ids(arm_name)) - - time.sleep(0.5) - - fk_xpos_batch = torch.cat(qpos_fk_list, dim=0) - - fk_xpos_list = robot.compute_fk(qpos=fk_xpos_batch, name=arm_name, to_matrix=True) - - start_time = time.time() - res, ik_qpos = robot.compute_ik( - pose=fk_xpos_list, - name=arm_name, - # joint_seed=qpos_fk_list[0], - return_all_solutions=True, - ) - end_time = time.time() - print( - f"Batch IK computation time for {len(fk_xpos_list)} poses: {end_time - start_time:.6f} seconds" +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, + shutdown_sim, +) + + +def main() -> None: + """Solve one left-arm pose with all analytic solutions enabled.""" + args = add_demo_args( + argparse.ArgumentParser(description="Run the SRS solver example.") + ).parse_args() + setup_print_options() + + sim = create_default_sim( + args, + width=2200, + height=1200, + add_default_light=False, ) + try: + robot = sim.add_robot(cfg=DexforceW1Cfg.from_dict({"uid": "dexforce_w1"})) + maybe_open_window(sim, args) + + arm_name = "left_arm" + qpos = torch.tensor( + [[0.0, 0.0, 0.0, -np.pi / 2, 0.0, 0.0, 0.0]], + dtype=torch.float32, + device=robot.device, + ) + target_pose = robot.compute_fk( + qpos=qpos, + name=arm_name, + to_matrix=True, + ) - if ik_qpos.dim() == 3: - first_solutions = ik_qpos[:, 0, :] - else: - first_solutions = ik_qpos - robot.set_qpos(first_solutions, joint_ids=robot.get_joint_ids(arm_name)) - - ik_xpos_list = robot.compute_fk(qpos=first_solutions, name=arm_name, to_matrix=True) - - print("fk_xpos_list: ", fk_xpos_list) - print("ik_xpos_list: ", ik_xpos_list) + started_at = time.perf_counter() + success, solutions = robot.compute_ik( + pose=target_pose, + name=arm_name, + return_all_solutions=True, + ) + elapsed = time.perf_counter() - started_at + if not torch.as_tensor(success).all(): + raise RuntimeError("SRS IK failed.") + + solution = solutions[:, 0, :] if solutions.dim() == 3 else solutions + robot.set_qpos(solution, joint_ids=robot.get_joint_ids(arm_name)) + achieved_pose = robot.compute_fk( + qpos=solution, + name=arm_name, + to_matrix=True, + ) - embed(header="Test SRSSolver example. Press Ctrl+D to exit.") + print(f"IK computation time: {elapsed:.6f} seconds") + print("Target FK pose:", target_pose) + print("IK result pose:", achieved_pose) + sim.update(step=1) + maybe_wait_for_user(args, "Press Enter to exit...") + finally: + shutdown_sim(sim) if __name__ == "__main__": diff --git a/examples/sim/utility/workspace_analyzer/analyze_cartesian_workspace.py b/examples/sim/utility/workspace_analyzer/analyze_cartesian_workspace.py index 8d2b5b9c0..b4ccabdb6 100644 --- a/examples/sim/utility/workspace_analyzer/analyze_cartesian_workspace.py +++ b/examples/sim/utility/workspace_analyzer/analyze_cartesian_workspace.py @@ -14,91 +14,117 @@ # limitations under the License. # ---------------------------------------------------------------------------- -import torch +"""Analyze the sampled Cartesian workspace of DexForce W1.""" + +from __future__ import annotations + +import argparse + import numpy as np -from IPython import embed +import torch -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import MarkerCfg from embodichain.lab.sim.robots import DexforceW1Cfg -from embodichain.lab.sim.cfg import MarkerCfg, RenderCfg -from embodichain.lab.sim.utility.workspace_analyzer.workspace_analyzer import ( - WorkspaceAnalyzer, - WorkspaceAnalyzerConfig, - AnalysisMode, +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, + shutdown_sim, ) from embodichain.lab.sim.utility.workspace_analyzer.configs.visualization_config import ( VisualizationConfig, ) +from embodichain.lab.sim.utility.workspace_analyzer.workspace_analyzer import ( + AnalysisMode, + WorkspaceAnalyzer, + WorkspaceAnalyzerConfig, +) -if __name__ == "__main__": - # Example usage - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) - config = SimulationManagerCfg( - headless=False, - sim_device="cuda", - width=1080, - height=1080, +def main() -> None: + """Run Cartesian workspace analysis.""" + parser = add_demo_args( + argparse.ArgumentParser(description="Analyze the W1 Cartesian workspace.") ) - sim = SimulationManager(config) - - cfg = DexforceW1Cfg.from_dict( - {"uid": "dexforce_w1", "version": "v021", "arm_kind": "industrial"} + parser.add_argument( + "--num_samples", + "--num-samples", + type=int, + default=50000, + help="Number of Cartesian poses to sample.", ) - robot = sim.add_robot(cfg=cfg) - print("DexforceW1 robot added to the simulation.") + parser.set_defaults(device="cuda") + args = parser.parse_args() + setup_print_options() - # Set left arm joint positions (mirrored) - left_qpos = torch.tensor( - [0, -np.pi / 4, 0.0, -np.pi / 2, -np.pi / 4, 0.0, 0.0], - dtype=torch.float32, - device=robot.device, - ) - right_qpos = -left_qpos - robot.set_qpos( - qpos=left_qpos, - joint_ids=robot.get_joint_ids("left_arm"), - ) - # Set right arm joint positions (mirrored) - robot.set_qpos( - qpos=right_qpos, - joint_ids=robot.get_joint_ids("right_arm"), + sim = create_default_sim( + args, + width=1080, + height=1080, + add_default_light=False, ) + try: + robot = sim.add_robot( + cfg=DexforceW1Cfg.from_dict( + { + "uid": "dexforce_w1", + "version": "v021", + "arm_kind": "industrial", + } + ) + ) + maybe_open_window(sim, args) - left_arm_pose = robot.compute_fk( - qpos=left_qpos, - name="left_arm", - to_matrix=True, - ) + left_qpos = torch.tensor( + [0, -np.pi / 4, 0.0, -np.pi / 2, -np.pi / 4, 0.0, 0.0], + dtype=torch.float32, + device=robot.device, + ) + robot.set_qpos(left_qpos, joint_ids=robot.get_joint_ids("left_arm")) + robot.set_qpos(-left_qpos, joint_ids=robot.get_joint_ids("right_arm")) + left_arm_pose = robot.compute_fk( + qpos=left_qpos, + name="left_arm", + to_matrix=True, + ) + sim.draw_marker( + cfg=MarkerCfg( + name="left_arm_pose_axis", + marker_type="axis", + axis_xpos=left_arm_pose, + axis_size=0.005, + axis_len=0.15, + arena_index=0, + ) + ) - sim.draw_marker( - cfg=MarkerCfg( - name=f"left_arm_pose_axis", - marker_type="axis", - axis_xpos=left_arm_pose, - axis_size=0.005, - axis_len=0.15, - arena_index=0, + results = WorkspaceAnalyzer( + robot=robot, + config=WorkspaceAnalyzerConfig( + mode=AnalysisMode.CARTESIAN_SPACE, + visualization=VisualizationConfig( + show_unreachable_points=False, + point_size=8.0, + ), + control_part_name="left_arm", + ), + sim_manager=sim, + ).analyze( + num_samples=args.num_samples, + visualize=not args.headless, ) - ) + print("\nCartesian Space Results:") + print( + f" Reachable points: {results['num_reachable']} / {results['num_samples']}" + ) + print(f" Analysis time: {results['analysis_time']:.2f}s") + print(f" Metrics: {results['metrics']}") + maybe_wait_for_user(args, "Press Enter to exit...") + finally: + shutdown_sim(sim) - cartesian_config = WorkspaceAnalyzerConfig( - mode=AnalysisMode.CARTESIAN_SPACE, - visualization=VisualizationConfig( - show_unreachable_points=False, point_size=8.0 - ), - control_part_name="left_arm", - ) - wa_cartesian = WorkspaceAnalyzer( - robot=robot, config=cartesian_config, sim_manager=sim - ) - results_cartesian = wa_cartesian.analyze(num_samples=50000, visualize=True) - print(f"\nCartesian Space Results:") - print( - f" Reachable points: {results_cartesian['num_reachable']} / {results_cartesian['num_samples']}" - ) - print(f" Analysis time: {results_cartesian['analysis_time']:.2f}s") - print(f" Metrics: {results_cartesian['metrics']}") - embed(header="Workspace Analyzer Test Environment") +if __name__ == "__main__": + main() diff --git a/examples/sim/utility/workspace_analyzer/analyze_joint_workspace.py b/examples/sim/utility/workspace_analyzer/analyze_joint_workspace.py index 5c658fa98..3fbb0b576 100644 --- a/examples/sim/utility/workspace_analyzer/analyze_joint_workspace.py +++ b/examples/sim/utility/workspace_analyzer/analyze_joint_workspace.py @@ -14,41 +14,68 @@ # limitations under the License. # ---------------------------------------------------------------------------- -import torch -import numpy as np -from IPython import embed +"""Analyze the sampled joint-space workspace of DexForce W1.""" + +from __future__ import annotations + +import argparse -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.robots import DexforceW1Cfg +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, + shutdown_sim, +) from embodichain.lab.sim.utility.workspace_analyzer.workspace_analyzer import ( WorkspaceAnalyzer, ) -if __name__ == "__main__": - # Example usage - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) - - config = SimulationManagerCfg(headless=False, sim_device="cpu") - sim_manager = SimulationManager(config) - sim_manager.set_manual_update(False) - cfg = DexforceW1Cfg.from_dict( - {"uid": "dexforce_w1", "version": "v021", "arm_kind": "industrial"} +def main() -> None: + """Run joint-space workspace analysis.""" + parser = add_demo_args( + argparse.ArgumentParser(description="Analyze the W1 joint workspace.") + ) + parser.add_argument( + "--num_samples", + "--num-samples", + type=int, + default=30000, + help="Number of joint configurations to sample.", ) - robot = sim_manager.add_robot(cfg=cfg) - print("DexforceW1 robot added to the simulation.") + args = parser.parse_args() + setup_print_options() - print("Example: Joint Space Analysis") + sim = create_default_sim(args, add_default_light=False) + try: + sim.set_manual_update(False) + robot = sim.add_robot( + cfg=DexforceW1Cfg.from_dict( + { + "uid": "dexforce_w1", + "version": "v021", + "arm_kind": "industrial", + } + ) + ) + maybe_open_window(sim, args) + print("DexforceW1 robot added to the simulation.") - wa_joint = WorkspaceAnalyzer(robot=robot, sim_manager=sim_manager) - results_joint = wa_joint.analyze(num_samples=30000, visualize=True) + results = WorkspaceAnalyzer(robot=robot, sim_manager=sim).analyze( + num_samples=args.num_samples, + visualize=not args.headless, + ) + print("\nJoint Space Results:") + print(f" Valid points: {results['num_valid']} / {results['num_samples']}") + print(f" Analysis time: {results['analysis_time']:.2f}s") + print(f" Metrics: {results['metrics']}") + maybe_wait_for_user(args, "Press Enter to exit...") + finally: + shutdown_sim(sim) - print(f"\nJoint Space Results:") - print( - f" Valid points: {results_joint['num_valid']} / {results_joint['num_samples']}" - ) - print(f" Analysis time: {results_joint['analysis_time']:.2f}s") - print(f" Metrics: {results_joint['metrics']}") - embed(header="End of Joint Space Analysis Example") +if __name__ == "__main__": + main() diff --git a/examples/sim/utility/workspace_analyzer/analyze_plane_workspace.py b/examples/sim/utility/workspace_analyzer/analyze_plane_workspace.py index 8bd1b4ce1..e1d5fe68b 100644 --- a/examples/sim/utility/workspace_analyzer/analyze_plane_workspace.py +++ b/examples/sim/utility/workspace_analyzer/analyze_plane_workspace.py @@ -14,94 +14,120 @@ # limitations under the License. # ---------------------------------------------------------------------------- -import torch +"""Analyze a planar slice of the DexForce W1 workspace.""" + +from __future__ import annotations + +import argparse + import numpy as np -from IPython import embed +import torch -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import MarkerCfg from embodichain.lab.sim.robots import DexforceW1Cfg -from embodichain.lab.sim.utility.workspace_analyzer.workspace_analyzer import ( - WorkspaceAnalyzer, - WorkspaceAnalyzerConfig, - AnalysisMode, +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, + shutdown_sim, ) -from embodichain.lab.sim.cfg import MarkerCfg, RenderCfg from embodichain.lab.sim.utility.workspace_analyzer.configs.visualization_config import ( VisualizationConfig, ) +from embodichain.lab.sim.utility.workspace_analyzer.workspace_analyzer import ( + AnalysisMode, + WorkspaceAnalyzer, + WorkspaceAnalyzerConfig, +) -if __name__ == "__main__": - # Example usage - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) - config = SimulationManagerCfg( - headless=False, - sim_device="cpu", - width=1080, - height=1080, +def main() -> None: + """Run planar workspace analysis.""" + parser = add_demo_args( + argparse.ArgumentParser(description="Analyze a planar W1 workspace slice.") ) - sim = SimulationManager(config) - sim.set_manual_update(False) - - cfg = DexforceW1Cfg.from_dict( - {"uid": "dexforce_w1", "version": "v021", "arm_kind": "industrial"} + parser.add_argument( + "--num_samples", + "--num-samples", + type=int, + default=1500, + help="Number of planar poses to sample.", ) - robot = sim.add_robot(cfg=cfg) - print("DexforceW1 robot added to the simulation.") + args = parser.parse_args() + setup_print_options() - # Set left arm joint positions (mirrored) - left_qpos = torch.tensor([0, -np.pi / 4, 0.0, -np.pi / 2, -np.pi / 4, 0.0, 0.0]) - right_qpos = -left_qpos - robot.set_qpos( - qpos=left_qpos, - joint_ids=robot.get_joint_ids("left_arm"), - ) - # Set right arm joint positions (mirrored) - robot.set_qpos( - qpos=right_qpos, - joint_ids=robot.get_joint_ids("right_arm"), + sim = create_default_sim( + args, + width=1080, + height=1080, + add_default_light=False, ) + try: + sim.set_manual_update(False) + robot = sim.add_robot( + cfg=DexforceW1Cfg.from_dict( + { + "uid": "dexforce_w1", + "version": "v021", + "arm_kind": "industrial", + } + ) + ) + maybe_open_window(sim, args) - left_arm_pose = robot.compute_fk( - qpos=left_qpos, - name="left_arm", - to_matrix=True, - ) + left_qpos = torch.tensor( + [0, -np.pi / 4, 0.0, -np.pi / 2, -np.pi / 4, 0.0, 0.0], + dtype=torch.float32, + device=robot.device, + ) + robot.set_qpos(left_qpos, joint_ids=robot.get_joint_ids("left_arm")) + robot.set_qpos(-left_qpos, joint_ids=robot.get_joint_ids("right_arm")) + left_arm_pose = robot.compute_fk( + qpos=left_qpos, + name="left_arm", + to_matrix=True, + ) + sim.draw_marker( + cfg=MarkerCfg( + name="left_arm_pose_axis", + marker_type="axis", + axis_xpos=left_arm_pose, + axis_size=0.005, + axis_len=0.15, + arena_index=0, + ) + ) - sim.draw_marker( - cfg=MarkerCfg( - name=f"left_arm_pose_axis", - marker_type="axis", - axis_xpos=left_arm_pose, - axis_size=0.005, - axis_len=0.15, - arena_index=0, + results = WorkspaceAnalyzer( + robot=robot, + config=WorkspaceAnalyzerConfig( + mode=AnalysisMode.PLANE_SAMPLING, + plane_normal=torch.tensor([0.0, 0.0, 1.0]), + plane_point=torch.tensor([0.0, 0.0, 1.2]), + reference_pose=left_arm_pose[0].cpu().numpy(), + visualization=VisualizationConfig( + show_unreachable_points=False, + vis_type="axis", + ), + control_part_name="left_arm", + ), + sim_manager=sim, + ).analyze( + num_samples=args.num_samples, + visualize=not args.headless, ) - ) + print("\nPlanar Workspace Results:") + print( + f" Reachable points: {results['num_reachable']} / {results['num_samples']}" + ) + print(f" Analysis time: {results['analysis_time']:.2f}s") + print(f" Metrics: {results['metrics']}") + maybe_wait_for_user(args, "Press Enter to exit...") + finally: + shutdown_sim(sim) - cartesian_config = WorkspaceAnalyzerConfig( - mode=AnalysisMode.PLANE_SAMPLING, - plane_normal=torch.tensor([0.0, 0.0, 1.0]), - plane_point=torch.tensor([0.0, 0.0, 1.2]), - # plane_bounds=torch.tensor([[-0.5, 0.5], [-0.5, 0.5]]), - reference_pose=left_arm_pose[0] - .cpu() - .numpy(), # Use computed left arm pose as reference - visualization=VisualizationConfig( - show_unreachable_points=False, vis_type="axis" - ), - control_part_name="left_arm", - ) - wa_cartesian = WorkspaceAnalyzer( - robot=robot, config=cartesian_config, sim_manager=sim - ) - results_cartesian = wa_cartesian.analyze(num_samples=1500, visualize=True) - print(f"\nCartesian Space Results:") - print( - f" Reachable points: {results_cartesian['num_reachable']} / {results_cartesian['num_samples']}" - ) - print(f" Analysis time: {results_cartesian['analysis_time']:.2f}s") - print(f" Metrics: {results_cartesian['metrics']}") - embed(header="Workspace Analyzer Test Environment") +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/atomic_action/coordinated_pickment.py b/scripts/tutorials/atomic_action/coordinated_pickment.py index cb00e9fe3..cf9403e5a 100644 --- a/scripts/tutorials/atomic_action/coordinated_pickment.py +++ b/scripts/tutorials/atomic_action/coordinated_pickment.py @@ -37,7 +37,6 @@ import numpy as np import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.atomic_actions import ( Affordance, @@ -57,6 +56,7 @@ from embodichain.lab.sim.objects import RigidObject, Robot from embodichain.lab.sim.shapes import CubeCfg, MeshCfg from embodichain.lab.sim.solvers import PytorchSolverCfg +from embodichain.lab.sim.utility.demo_utils import add_demo_args, shutdown_sim from embodichain.utils import logger from embodichain.utils.math import matrix_from_euler from scripts.tutorials.atomic_action.tutorial_utils import ( @@ -151,25 +151,23 @@ class PickmentObjectPreset: def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the demo.""" parser = argparse.ArgumentParser(description="Dual-arm coordinated pickment demo") - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) parser.set_defaults(device="cpu", renderer="hybrid") parser.add_argument( "--diagnose_plan", + "--diagnose-plan", action="store_true", help="Plan and print diagnostics without playing the trajectory.", ) parser.add_argument( "--debug_state", + "--debug-state", action="store_true", help="Log hand targets and object poses during execution.", ) - parser.add_argument( - "--auto_play", - action="store_true", - help="Run the viewer demo without waiting for keyboard input.", - ) parser.add_argument( "--headless_play", + "--headless-play", action="store_true", help="Execute planned trajectories without opening the viewer window.", ) @@ -179,11 +177,6 @@ def parse_arguments() -> argparse.Namespace: default="pencil", help="Object mesh to grasp in the coordinated pickment demo.", ) - parser.add_argument( - "--no_vis_eef_axis", - action="store_true", - help="Do not draw the pickment target/grasp coordinate frames before planning.", - ) return parser.parse_args() @@ -787,8 +780,11 @@ def main() -> None: arena_space=3.0, light_pos=(0.0, -0.4, 3.0), ) - robot = create_dual_ur5_robot(sim) - run_coordinated_pickment_demo(args, sim, robot) + try: + robot = create_dual_ur5_robot(sim) + run_coordinated_pickment_demo(args, sim, robot) + finally: + shutdown_sim(sim) if __name__ == "__main__": diff --git a/scripts/tutorials/atomic_action/coordinated_placement.py b/scripts/tutorials/atomic_action/coordinated_placement.py index f6b89073a..1653fad09 100644 --- a/scripts/tutorials/atomic_action/coordinated_placement.py +++ b/scripts/tutorials/atomic_action/coordinated_placement.py @@ -38,7 +38,6 @@ import torch from scipy.spatial.transform import Rotation as SciRotation -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.atomic_actions import ( Affordance, @@ -63,6 +62,7 @@ from embodichain.lab.sim.objects import RigidObject, Robot from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim.solvers import URSolverCfg +from embodichain.lab.sim.utility.demo_utils import add_demo_args, shutdown_sim from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( broadcast_pose_batch, @@ -192,30 +192,23 @@ def transform_baseline_pose( def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the demo.""" parser = argparse.ArgumentParser(description="Dual-arm coordinated placement demo") - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) parser.set_defaults(device="cuda", renderer="hybrid") parser.add_argument( "--diagnose_plan", + "--diagnose-plan", action="store_true", help="Plan and print diagnostics without playing the trajectory.", ) parser.add_argument( "--debug_state", + "--debug-state", action="store_true", help="Log hand targets and object poses during execution.", ) - parser.add_argument( - "--auto_play", - action="store_true", - help="Run the viewer demo without waiting for keyboard input.", - ) - parser.add_argument( - "--no_vis_eef_axis", - action="store_true", - help="Do not draw coordinated placement target coordinate frames.", - ) parser.add_argument( "--headless_play", + "--headless-play", action="store_true", help="Execute planned trajectories without opening the viewer window.", ) @@ -1065,8 +1058,11 @@ def main() -> None: arena_space=3.0, light_pos=(0.0, -0.4, 3.0), ) - robot = create_dual_ur5_robot(sim) - run_coordinated_placement_demo(args, sim, robot) + try: + robot = create_dual_ur5_robot(sim) + run_coordinated_placement_demo(args, sim, robot) + finally: + shutdown_sim(sim) if __name__ == "__main__": diff --git a/scripts/tutorials/atomic_action/move_end_effector.py b/scripts/tutorials/atomic_action/move_end_effector.py index 9ef606ad3..6731cb0cd 100644 --- a/scripts/tutorials/atomic_action/move_end_effector.py +++ b/scripts/tutorials/atomic_action/move_end_effector.py @@ -19,109 +19,164 @@ from __future__ import annotations import argparse -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, EndEffectorPoseTarget, MoveEndEffector, MoveEndEffectorCfg, ) +from embodichain.lab.sim.demo_base import DemoBase +from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + format_tensor, + maybe_open_window, + maybe_wait_for_user, + replay_trajectory, + setup_print_options, +) from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, - broadcast_pose_batch, - broadcast_waypoint_pose_batch, - create_toppra_motion_generator, - create_tutorial_simulation, + create_ur5_gripper_robot_cfg, draw_axis_marker, - make_top_down_eef_pose, - prepare_tutorial_scene, - replay_trajectory, + get_tutorial_window_size, ) MOVE_SAMPLE_INTERVAL = 80 POST_TRAJECTORY_STEPS = 120 -def parse_arguments() -> argparse.Namespace: - """Parse command-line arguments for the MoveEndEffector tutorial.""" - parser = argparse.ArgumentParser( - description="Demonstrate MoveEndEffector with a multi-waypoint pose trajectory." - ) - add_env_launcher_args_to_parser(parser) - parser.add_argument("--auto_play", action="store_true") - parser.add_argument("--no_vis_eef_axis", action="store_true") - return parser.parse_args() +class MoveEndEffectorDemo(DemoBase): + """Demo that moves a UR5 end effector through a multi-waypoint trajectory.""" + def setup(self) -> None: + """Create simulation, robot, motion generator and atomic action engine.""" + width, height = get_tutorial_window_size(self.args) + self.sim = create_default_sim( + self.args, + width=width, + height=height, + physics_dt=1.0 / 100.0, + arena_space=2.5, + ) + self.robot = self.sim.add_robot( + cfg=create_ur5_gripper_robot_cfg(init_pos=(0.0, 0.0, 0.0)) + ) + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=self.robot.uid)) + ) -def main() -> None: - """Move the robot end effector through two pose waypoints.""" - args = parse_arguments() - sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot(sim) - motion_gen = create_toppra_motion_generator(robot) - - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register( - MoveEndEffector( - motion_gen, - cfg=MoveEndEffectorCfg(sample_interval=MOVE_SAMPLE_INTERVAL), + move_cfg = MoveEndEffectorCfg( + control_part="arm", + sample_interval=MOVE_SAMPLE_INTERVAL, ) - ) - poses = torch.stack( - [ - make_top_down_eef_pose( - torch.tensor([0.30, -0.20, 0.36], device=sim.device) - ), - make_top_down_eef_pose(torch.tensor([0.45, 0.10, 0.30], device=sim.device)), - ] - ) - n_envs = robot.get_qpos().shape[0] - if not args.no_vis_eef_axis: - for name, pose in zip(("target", "side"), poses, strict=True): + self.atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + self.atomic_engine.register(MoveEndEffector(motion_gen, cfg=move_cfg)) + + self.target_pose = self._make_top_down_eef_pose() + self.side_pose = self._make_side_eef_pose() + + maybe_open_window(self.sim, self.args) + if not self.args.no_vis_eef_axis: draw_axis_marker( - sim, - f"move_end_effector_{name}_axis", - broadcast_pose_batch(pose, num_envs=n_envs), + self.sim, "move_end_effector_target_axis", self.target_pose ) - wait_for_user = prepare_tutorial_scene( - sim, args, "Inspect the robot, then press Enter to plan MoveEndEffector..." - ) + draw_axis_marker(self.sim, "move_end_effector_side_axis", self.side_pose) + + def run(self) -> None: + """Plan and replay the MoveEndEffector trajectory.""" + maybe_wait_for_user( + self.args, "Inspect the robot, then press Enter to plan MoveEndEffector..." + ) - success, trajectory, _ = engine.run( - [ - ( - "move_end_effector", - EndEffectorPoseTarget(broadcast_waypoint_pose_batch(poses, n_envs)), + n_envs = self.robot.get_qpos().shape[0] + multi_waypoint_xpos = ( + torch.stack([self.target_pose, self.side_pose], dim=0) + .unsqueeze(0) + .repeat(n_envs, 1, 1, 1) + ) + logger.log_info( + "Planning MoveEndEffector through multi-waypoint trajectory: " + f"xpos0={format_tensor(self.target_pose[:3, 3])} -> " + f"xpos1={format_tensor(self.side_pose[:3, 3])}" + ) + success, traj, _ = self.atomic_engine.run( + steps=[ + ( + "move_end_effector", + EndEffectorPoseTarget(xpos=multi_waypoint_xpos), + ) + ] + ) + if not success.all(): + logger.log_warning("Failed to plan MoveEndEffector demo trajectory.") + return + + maybe_wait_for_user( + self.args, "Press Enter to replay the MoveEndEffector demo..." + ) + + with DemoRecording(self.sim, self.args, prefix="move_end_effector_auto_play"): + replay_trajectory( + self.sim, + self.robot, + traj, + post_steps=POST_TRAJECTORY_STEPS, + step_size=4, + sleep=1e-2, ) - ] - ) - if not success.all(): - logger.log_warning("Failed to plan MoveEndEffector demo trajectory.") - return - - if wait_for_user: - input("Press Enter to replay the MoveEndEffector demo...") - replay_trajectory( - sim, - robot, - trajectory, - args, - video_prefix="move_end_effector_auto_play", - hold_steps=POST_TRAJECTORY_STEPS, + + maybe_wait_for_user(self.args, "Press Enter to exit the simulation...") + + def _make_top_down_eef_pose(self) -> torch.Tensor: + pose = torch.eye(4, dtype=torch.float32, device=self.sim.device) + pose[:3, :3] = torch.tensor( + [ + [-0.0539, -0.9985, -0.0022], + [-0.9977, 0.0540, -0.0401], + [0.0401, 0.0000, -0.9992], + ], + dtype=torch.float32, + device=self.sim.device, + ) + pose[:3, 3] = torch.tensor( + [0.30, -0.20, 0.36], dtype=torch.float32, device=self.sim.device + ) + return pose + + def _make_side_eef_pose(self) -> torch.Tensor: + """A second waypoint offset from the top-down pose for the multi-waypoint demo.""" + pose = torch.eye(4, dtype=torch.float32, device=self.sim.device) + pose[:3, :3] = torch.tensor( + [ + [-0.0539, -0.9985, -0.0022], + [-0.9977, 0.0540, -0.0401], + [0.0401, 0.0000, -0.9992], + ], + dtype=torch.float32, + device=self.sim.device, + ) + pose[:3, 3] = torch.tensor( + [0.45, 0.10, 0.30], dtype=torch.float32, device=self.sim.device + ) + return pose + + +def main() -> None: + """Entry point for the MoveEndEffector demo.""" + setup_print_options() + parser = argparse.ArgumentParser( + description="Demonstrate MoveEndEffector with a multi-waypoint pose trajectory." ) - if wait_for_user: - input("Press Enter to exit the simulation...") + parser = add_demo_args(parser) + args = parser.parse_args() + MoveEndEffectorDemo(args).main() if __name__ == "__main__": diff --git a/scripts/tutorials/atomic_action/move_held_object.py b/scripts/tutorials/atomic_action/move_held_object.py index d0acc5afd..3b6c8e980 100644 --- a/scripts/tutorials/atomic_action/move_held_object.py +++ b/scripts/tutorials/atomic_action/move_held_object.py @@ -19,18 +19,13 @@ from __future__ import annotations import argparse -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) +import time import torch from embodichain.data import get_data_path -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.atomic_actions import ( + AntipodalAffordance, AtomicActionEngine, EndEffectorPoseTarget, GraspTarget, @@ -39,29 +34,52 @@ MoveEndEffectorCfg, MoveHeldObject, MoveHeldObjectCfg, + ObjectSemantics, PickUp, PickUpCfg, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.demo_base import DemoBase from embodichain.lab.sim.objects import RigidObject +from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, +) +from embodichain.toolkits.graspkit.pg_grasp.antipodal_generator import ( + AntipodalSamplerCfg, + GraspGeneratorCfg, +) +from embodichain.toolkits.graspkit.pg_grasp.gripper_collision_checker import ( + GripperCollisionCfg, +) from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, - broadcast_pose_batch, - clone_local_pose_from_first_env, - create_antipodal_semantics, - create_toppra_motion_generator, - create_tutorial_simulation, + create_ur5_gripper_robot_cfg, draw_axis_marker, - get_hand_open_close_qpos, - make_eef_pose_at, - prepare_tutorial_scene, - replay_trajectory, + get_tutorial_window_size, ) -OBJECT_MESH_PATH = "PaperCup/paper_cup.ply" +GRIPPER_MAX_OPEN_WIDTH = 0.080 +GRIPPER_FINGER_LENGTH = 0.088 +GRIPPER_ROOT_Z_WIDTH = 0.096 +GRIPPER_Y_THICKNESS = 0.040 + +OBJECT_LABEL = "sugar_box" +OBJECT_MESH_PATH = "SugarBox/sugar_box_usd/sugar_box.usda" OBJECT_XY = (-0.42, -0.08) +OBJECT_APPROACH_DIRECTION = (0.0, 0.0, -1.0) +OBJECT_MIN_HAND_CLOSE_QPOS = 0.024 +OBJECT_INIT_ROT = (0.0, 0.0, 0.0) +OBJECT_BODY_SCALE = (0.8, 0.8, 0.8) +OBJECT_MASS = 0.05 +OBJECT_USE_USD_PROPERTIES = False + MOVE_SAMPLE_INTERVAL = 60 PICK_SAMPLE_INTERVAL = 120 MOVE_HELD_OBJECT_SAMPLE_INTERVAL = 120 @@ -69,145 +87,254 @@ POST_TRAJECTORY_STEPS = 240 -def parse_arguments() -> argparse.Namespace: - """Parse command-line arguments for the MoveHeldObject tutorial.""" - parser = argparse.ArgumentParser( - description="Pick up a paper cup and move it by object pose." - ) - add_env_launcher_args_to_parser(parser) - parser.add_argument("--n_sample", type=int, default=10000) - parser.add_argument("--force_reannotate", action="store_true") - parser.add_argument("--auto_play", action="store_true") - parser.add_argument("--no_vis_eef_axis", action="store_true") - return parser.parse_args() - - -def create_pick_object(sim) -> RigidObject: - """Create and settle the paper cup used by the tutorial.""" - obj = sim.add_rigid_object( - cfg=RigidObjectCfg( - uid="paper_cup", +class MoveHeldObjectDemo(DemoBase): + """Demo that picks up an object and moves it to an object-frame target pose.""" + + def setup(self) -> None: + """Create simulation, robot, object, motion generator and action engine.""" + width, height = get_tutorial_window_size(self.args) + self.sim = create_default_sim( + self.args, + width=width, + height=height, + physics_dt=1.0 / 100.0, + arena_space=2.5, + ) + self.robot = self.sim.add_robot( + cfg=create_ur5_gripper_robot_cfg(init_pos=(0.0, 0.0, 0.0)) + ) + self.obj = self._create_pick_object() + + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=self.robot.uid)) + ) + + hand_open, hand_close = self._get_hand_open_close_qpos() + move_cfg = MoveEndEffectorCfg( + control_part="arm", + sample_interval=MOVE_SAMPLE_INTERVAL, + ) + pickup_cfg = PickUpCfg( + control_part="arm", + hand_control_part="hand", + hand_open_qpos=hand_open, + hand_close_qpos=hand_close, + approach_direction=torch.tensor( + OBJECT_APPROACH_DIRECTION, + dtype=torch.float32, + device=self.sim.device, + ), + pre_grasp_distance=0.15, + lift_height=0.16, + sample_interval=PICK_SAMPLE_INTERVAL, + hand_interp_steps=HAND_INTERP_STEPS, + ) + move_held_object_cfg = MoveHeldObjectCfg( + control_part="arm", + hand_control_part="hand", + hand_close_qpos=hand_close, + sample_interval=MOVE_HELD_OBJECT_SAMPLE_INTERVAL, + ) + + self.atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + self.atomic_engine.register(MoveEndEffector(motion_gen, cfg=move_cfg)) + self.atomic_engine.register(PickUp(motion_gen, cfg=pickup_cfg)) + self.atomic_engine.register( + MoveHeldObject(motion_gen, cfg=move_held_object_cfg) + ) + + self.semantics = self._create_object_semantics() + obj_pose = self.obj.get_local_pose(to_matrix=True) + move_position = obj_pose[0, :3, 3].clone() + move_position[2] = 0.36 + self.move_target = self._make_pre_pick_eef_pose(move_position) + self.object_target_pose = self._make_object_target_pose() + + maybe_open_window(self.sim, self.args) + if not self.args.no_vis_eef_axis: + draw_axis_marker( + self.sim, "move_held_object_target_axis", self.object_target_pose + ) + + def run(self) -> None: + """Plan and replay the move_end_effector -> pick_up -> move_held_object trajectory.""" + maybe_wait_for_user( + self.args, "Inspect the sugar box, then press Enter to plan..." + ) + + logger.log_info("Planning move_end_effector -> pick_up -> move_held_object") + move_held_object_target = HeldObjectPoseTarget( + object_target_pose=self.object_target_pose + ) + start_time = time.time() + success, traj, _ = self.atomic_engine.run( + steps=[ + ("move_end_effector", EndEffectorPoseTarget(xpos=self.move_target)), + ("pick_up", GraspTarget(semantics=self.semantics)), + ("move_held_object", move_held_object_target), + ] + ) + cost_time = time.time() - start_time + logger.log_info(f"Plan trajectory cost time: {cost_time:.2f} seconds") + if not success.all(): + logger.log_warning("Failed to plan move_held_object demo trajectory.") + return + + maybe_wait_for_user( + self.args, "Press Enter to replay the move_held_object demo..." + ) + + with DemoRecording(self.sim, self.args, prefix="move_held_object_auto_play"): + self._replay_move_held_object_trajectory(traj) + + maybe_wait_for_user(self.args, "Press Enter to exit the simulation...") + + def _create_pick_object(self) -> RigidObject: + cfg = RigidObjectCfg( + uid=OBJECT_LABEL, shape=MeshCfg(fpath=get_data_path(OBJECT_MESH_PATH)), attrs=RigidBodyAttributesCfg( - mass=0.01, + mass=OBJECT_MASS, dynamic_friction=0.97, static_friction=0.99, ), max_convex_hull_num=16, - init_pos=[*OBJECT_XY, 0.0], - body_scale=(0.75, 0.75, 1.0), + init_pos=[OBJECT_XY[0], OBJECT_XY[1], 0.0], + init_rot=OBJECT_INIT_ROT, + body_scale=OBJECT_BODY_SCALE, + use_usd_properties=OBJECT_USE_USD_PROPERTIES, ) - ) - sim.update(step=10) - clone_local_pose_from_first_env(obj) - obj.clear_dynamics() - return obj - - -def make_object_target_pose(device: torch.device) -> torch.Tensor: - """Build the desired final paper-cup pose in the object frame.""" - pose = torch.eye(4, dtype=torch.float32, device=device) - pose[:3, :3] = torch.tensor( - [[0.0, 0.0, -1.0], [0.0, 1.0, 0.0], [1.0, 0.0, 0.0]], - dtype=torch.float32, - device=device, - ) - pose[:3, 3] = torch.tensor([-0.3, -0.3, 0.5], device=device) - return pose + obj = self.sim.add_rigid_object(cfg=cfg) + # Settle the object to ensure it is resting on the ground before planning + self.sim.update(step=10) + return obj -def main() -> None: - """Plan MoveEndEffector -> PickUp -> MoveHeldObject.""" - args = parse_arguments() - sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot(sim) - obj = create_pick_object(sim) - motion_gen = create_toppra_motion_generator(robot) - hand_open, hand_close = get_hand_open_close_qpos(robot) - - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register( - MoveEndEffector( - motion_gen, MoveEndEffectorCfg(sample_interval=MOVE_SAMPLE_INTERVAL) + def _get_hand_open_close_qpos(self) -> tuple[torch.Tensor, torch.Tensor]: + hand_limits = self.robot.get_qpos_limits(name="hand")[0].to( + device=self.sim.device, dtype=torch.float32 ) - ) - engine.register( - PickUp( - motion_gen, - PickUpCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - pre_grasp_distance=0.15, - lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, + hand_open = hand_limits[:, 0] + hand_close_limit = hand_limits[:, 1] + hand_close = torch.minimum( + hand_close_limit, + torch.full_like(hand_close_limit, OBJECT_MIN_HAND_CLOSE_QPOS), + ) + return hand_open, hand_close + + def _build_grasp_generator_cfg(self) -> GraspGeneratorCfg: + return GraspGeneratorCfg( + viser_port=11801, + antipodal_sampler_cfg=AntipodalSamplerCfg( + n_sample=self.args.n_sample, + max_length=GRIPPER_MAX_OPEN_WIDTH, + min_length=0.003, ), + is_partial_annotate=False, + is_filter_ground_collision=False, ) - ) - engine.register( - MoveHeldObject( - motion_gen, - MoveHeldObjectCfg( - hand_close_qpos=hand_close, - sample_interval=MOVE_HELD_OBJECT_SAMPLE_INTERVAL, + + def _build_gripper_collision_cfg(self) -> GripperCollisionCfg: + return GripperCollisionCfg( + max_open_length=GRIPPER_MAX_OPEN_WIDTH, + finger_length=GRIPPER_FINGER_LENGTH, + y_thickness=GRIPPER_Y_THICKNESS, + root_z_width=GRIPPER_ROOT_Z_WIDTH, + open_check_margin=0.002, + point_sample_dense=0.012, + ) + + def _create_object_semantics(self) -> ObjectSemantics: + return ObjectSemantics( + label=OBJECT_LABEL, + geometry={ + "mesh_vertices": self.obj.get_vertices(env_ids=[0], scale=True)[0], + "mesh_triangles": self.obj.get_triangles(env_ids=[0])[0], + }, + affordance=AntipodalAffordance( + mesh_vertices=self.obj.get_vertices(env_ids=[0], scale=True)[0], + mesh_triangles=self.obj.get_triangles(env_ids=[0])[0], + gripper_collision_cfg=self._build_gripper_collision_cfg(), + generator_cfg=self._build_grasp_generator_cfg(), + force_reannotate=self.args.force_reannotate, ), + entity=self.obj, ) - ) - semantics = create_antipodal_semantics( - obj, - label="paper_cup", - n_sample=args.n_sample, - force_reannotate=args.force_reannotate, - ) - move_position = obj.get_local_pose(to_matrix=True)[0, :3, 3].clone() - move_position[2] = 0.36 - n_envs = robot.get_qpos().shape[0] - move_target = broadcast_pose_batch(make_eef_pose_at(robot, move_position), n_envs) - object_target = broadcast_pose_batch(make_object_target_pose(sim.device), n_envs) - if not args.no_vis_eef_axis: - draw_axis_marker(sim, "move_held_object_target_axis", object_target) - wait_for_user = prepare_tutorial_scene( - sim, args, "Inspect the paper cup, then press Enter to plan..." - ) + def _make_pre_pick_eef_pose(self, position: torch.Tensor) -> torch.Tensor: + pose = self.robot.compute_fk( + qpos=self.robot.get_qpos(name="arm"), + name="arm", + to_matrix=True, + )[0].clone() + pose[:3, 3] = position + return pose + + def _make_object_target_pose(self) -> torch.Tensor: + pose = torch.eye(4, dtype=torch.float32, device=self.sim.device) + pose[:3, :3] = torch.tensor( + [ + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0], + [1.0, 0.0, 0.0], + ], + dtype=torch.float32, + device=self.sim.device, + ) + pose[:3, 3] = torch.tensor( + [-0.3, -0.3, 0.5], dtype=torch.float32, device=self.sim.device + ) + return pose + + def _replay_move_held_object_trajectory(self, traj: torch.Tensor) -> None: + post_grasp_clear_step = self._compute_pick_close_end_step() + should_clear_object_dynamics = True + for i in range(traj.shape[1]): + self.robot.set_qpos(traj[:, i, :]) + self.sim.update(step=4) + if should_clear_object_dynamics and i + 1 >= post_grasp_clear_step: + self.obj.clear_dynamics() + should_clear_object_dynamics = False + logger.log_info(f"Object dynamics cleared after grasp at step={i}") + time.sleep(1e-2) - success, trajectory, _ = engine.run( - [ - ("move_end_effector", EndEffectorPoseTarget(move_target)), - ("pick_up", GraspTarget(semantics)), - ("move_held_object", HeldObjectPoseTarget(object_target)), - ] + logger.log_info("MoveHeldObject keeps the sugar box suspended in the gripper.") + + final_qpos = traj[:, -1, :] + for _ in range(POST_TRAJECTORY_STEPS): + self.robot.set_qpos(final_qpos) + self.sim.update(step=2) + time.sleep(1e-2) + + @staticmethod + def _compute_pick_close_end_step() -> int: + motion_waypoints = PICK_SAMPLE_INTERVAL - HAND_INTERP_STEPS + n_approach = int(round(motion_waypoints) * 0.6) + return MOVE_SAMPLE_INTERVAL + n_approach + HAND_INTERP_STEPS + + +def main() -> None: + """Entry point for the MoveHeldObject demo.""" + setup_print_options() + parser = argparse.ArgumentParser( + description="Demonstrate MoveHeldObject holding a sugar box in the gripper." ) - if not success.all(): - logger.log_warning("Failed to plan MoveHeldObject demo trajectory.") - return - - if wait_for_user: - input("Press Enter to replay the MoveHeldObject demo...") - clear_after_step = ( - MOVE_SAMPLE_INTERVAL - + round((PICK_SAMPLE_INTERVAL - HAND_INTERP_STEPS) * 0.6) - + HAND_INTERP_STEPS + parser = add_demo_args(parser) + parser.add_argument( + "--n_sample", + "--n-sample", + type=int, + default=10000, + help="Number of samples for antipodal grasp generation.", ) - dynamics_cleared = False - - def clear_object_dynamics(step_idx: int, _: int) -> None: - nonlocal dynamics_cleared - if not dynamics_cleared and step_idx + 1 >= clear_after_step: - obj.clear_dynamics() - dynamics_cleared = True - - replay_trajectory( - sim, - robot, - trajectory, - args, - video_prefix="move_held_object_auto_play", - hold_steps=POST_TRAJECTORY_STEPS, - on_trajectory_step=clear_object_dynamics, + parser.add_argument( + "--force_reannotate", + "--force-reannotate", + action="store_true", + help="Force grasp region re-annotation instead of using cached data.", ) - if wait_for_user: - input("Press Enter to exit the simulation...") + args = parser.parse_args() + MoveHeldObjectDemo(args).main() if __name__ == "__main__": diff --git a/scripts/tutorials/atomic_action/move_joints.py b/scripts/tutorials/atomic_action/move_joints.py index e4dd06fd4..235652f0b 100644 --- a/scripts/tutorials/atomic_action/move_joints.py +++ b/scripts/tutorials/atomic_action/move_joints.py @@ -19,16 +19,9 @@ from __future__ import annotations import argparse -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, JointPositionTarget, @@ -36,92 +29,128 @@ MoveJointsCfg, NamedJointPositionTarget, ) +from embodichain.lab.sim.demo_base import DemoBase +from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + replay_trajectory, + setup_print_options, +) from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, - create_toppra_motion_generator, - create_tutorial_simulation, + create_ur5_gripper_robot_cfg, draw_axis_marker, - prepare_tutorial_scene, - replay_trajectory, + get_tutorial_window_size, ) MOVE_JOINTS_SAMPLE_INTERVAL = 80 POST_TRAJECTORY_STEPS = 120 -def parse_arguments() -> argparse.Namespace: - """Parse command-line arguments for the MoveJoints tutorial.""" - parser = argparse.ArgumentParser( - description="Demonstrate MoveJoints with named and explicit qpos targets." - ) - add_env_launcher_args_to_parser(parser) - parser.add_argument("--auto_play", action="store_true") - parser.add_argument("--no_vis_eef_axis", action="store_true") - return parser.parse_args() +class MoveJointsDemo(DemoBase): + """Demo that moves a UR5 arm through named and explicit joint targets.""" + def setup(self) -> None: + """Create simulation, robot, motion generator and atomic action engine.""" + width, height = get_tutorial_window_size(self.args) + self.sim = create_default_sim( + self.args, + width=width, + height=height, + physics_dt=1.0 / 100.0, + arena_space=2.5, + ) + self.robot = self.sim.add_robot( + cfg=create_ur5_gripper_robot_cfg(init_pos=(0.0, 0.0, 0.0)) + ) + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=self.robot.uid)) + ) -def main() -> None: - """Move the robot arm through a named target and two explicit waypoints.""" - args = parse_arguments() - sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot(sim) - motion_gen = create_toppra_motion_generator(robot) - - ready, mid, home = ( - torch.tensor(qpos, dtype=torch.float32, device=sim.device) - for qpos in ( - [0.35, -1.20, 1.30, -1.65, -1.57, 0.20], - [0.15, -1.40, 1.45, -1.60, -1.57, 0.10], - [0.0, -1.57, 1.57, -1.57, -1.57, 0.0], + ready_qpos = self._make_arm_qpos([0.35, -1.20, 1.30, -1.65, -1.57, 0.20]) + mid_qpos = self._make_arm_qpos([0.15, -1.40, 1.45, -1.60, -1.57, 0.10]) + home_qpos = self._make_arm_qpos([0.0, -1.57, 1.57, -1.57, -1.57, 0.0]) + move_joints_cfg = MoveJointsCfg( + control_part="arm", + sample_interval=MOVE_JOINTS_SAMPLE_INTERVAL, + named_joint_positions={"ready": ready_qpos}, ) - ) - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register( - MoveJoints( - motion_gen, - cfg=MoveJointsCfg( - sample_interval=MOVE_JOINTS_SAMPLE_INTERVAL, - named_joint_positions={"ready": ready}, - ), + + self.atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + self.atomic_engine.register(MoveJoints(motion_gen, cfg=move_joints_cfg)) + + self.mid_qpos = mid_qpos + self.home_qpos = home_qpos + + maybe_open_window(self.sim, self.args) + if not self.args.no_vis_eef_axis: + self._draw_start_eef_axis() + + def run(self) -> None: + """Plan and replay the MoveJoints trajectory.""" + maybe_wait_for_user( + self.args, "Inspect the robot, then press Enter to plan MoveJoints..." ) - ) - if not args.no_vis_eef_axis: - draw_axis_marker( - sim, - "move_joints_start_eef_axis", - robot.compute_fk(robot.get_qpos(name="arm"), name="arm", to_matrix=True), + n_envs = self.robot.get_qpos().shape[0] + multi_waypoint_qpos = ( + torch.stack([self.mid_qpos, self.home_qpos], dim=0) + .unsqueeze(0) + .repeat(n_envs, 1, 1) ) - wait_for_user = prepare_tutorial_scene( - sim, args, "Inspect the robot, then press Enter to plan MoveJoints..." - ) + logger.log_info( + "Planning MoveJoints: NamedJointPositionTarget('ready') -> " + "multi-waypoint trajectory (mid -> home)" + ) + success, traj, _ = self.atomic_engine.run( + steps=[ + ("move_joints", NamedJointPositionTarget(name="ready")), + ("move_joints", JointPositionTarget(qpos=multi_waypoint_qpos)), + ] + ) + if not success.all(): + logger.log_warning("Failed to plan MoveJoints demo trajectory.") + return + + maybe_wait_for_user(self.args, "Press Enter to replay the MoveJoints demo...") + + with DemoRecording(self.sim, self.args, prefix="move_joints_auto_play"): + replay_trajectory( + self.sim, + self.robot, + traj, + post_steps=POST_TRAJECTORY_STEPS, + step_size=4, + sleep=1e-2, + ) + + maybe_wait_for_user(self.args, "Press Enter to exit the simulation...") + + def _make_arm_qpos(self, values: list[float]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.float32, device=self.sim.device) + + def _draw_start_eef_axis(self) -> None: + eef_pose = self.robot.compute_fk( + qpos=self.robot.get_qpos(name="arm"), + name="arm", + to_matrix=True, + ) + draw_axis_marker(self.sim, "move_joints_start_eef_axis", eef_pose) - waypoints = ( - torch.stack([mid, home]).unsqueeze(0).repeat(robot.get_qpos().shape[0], 1, 1) - ) - success, trajectory, _ = engine.run( - [ - ("move_joints", NamedJointPositionTarget("ready")), - ("move_joints", JointPositionTarget(waypoints)), - ] - ) - if not success.all(): - logger.log_warning("Failed to plan MoveJoints demo trajectory.") - return - - if wait_for_user: - input("Press Enter to replay the MoveJoints demo...") - replay_trajectory( - sim, - robot, - trajectory, - args, - video_prefix="move_joints_auto_play", - hold_steps=POST_TRAJECTORY_STEPS, + +def main() -> None: + """Entry point for the MoveJoints demo.""" + setup_print_options() + parser = argparse.ArgumentParser( + description="Demonstrate MoveJoints with named and explicit qpos targets." ) - if wait_for_user: - input("Press Enter to exit the simulation...") + parser = add_demo_args(parser) + args = parser.parse_args() + MoveJointsDemo(args).main() if __name__ == "__main__": diff --git a/scripts/tutorials/atomic_action/pickup.py b/scripts/tutorials/atomic_action/pickup.py index 2aa7b3082..87217e466 100644 --- a/scripts/tutorials/atomic_action/pickup.py +++ b/scripts/tutorials/atomic_action/pickup.py @@ -14,49 +14,75 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Demonstrate PickUp on a cube with a configurable approach direction.""" +"""Demonstrate PickUp on an upright object with configurable approach.""" from __future__ import annotations import argparse -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) +import time import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.data import get_data_path from embodichain.lab.sim.atomic_actions import ( + AntipodalAffordance, AtomicActionEngine, GraspTarget, + ObjectSemantics, PickUp, PickUpCfg, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.demo_base import DemoBase from embodichain.lab.sim.objects import RigidObject -from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + format_tensor, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, +) +from embodichain.toolkits.graspkit.pg_grasp.antipodal_generator import ( + AntipodalSamplerCfg, + GraspGeneratorCfg, +) +from embodichain.toolkits.graspkit.pg_grasp.gripper_collision_checker import ( + GripperCollisionCfg, +) from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, - clone_local_pose_from_first_env, - create_antipodal_semantics, - create_toppra_motion_generator, - create_tutorial_simulation, + create_ur5_gripper_robot_cfg, draw_axis_marker, - get_hand_open_close_qpos, - initialize_pre_pick_robot_pose, - prepare_tutorial_scene, - replay_trajectory, + get_tutorial_window_size, ) -OBJECT_SIZE = (0.05, 0.05, 0.05) +GRIPPER_MAX_OPEN_WIDTH = 0.080 +GRIPPER_FINGER_LENGTH = 0.088 +GRIPPER_ROOT_Z_WIDTH = 0.096 +GRIPPER_Y_THICKNESS = 0.040 + +OBJECT_MIN_HAND_CLOSE_QPOS = 0.024 OBJECT_XY = (-0.42, -0.08) + +OBJECT_PRESETS = { + "sugar_box": { + "label": "sugar_box", + "mesh_path": "SugarBox/sugar_box_usd/sugar_box.usda", + "init_rot": (0.0, 0.0, 0.0), + "body_scale": (0.8, 0.8, 0.8), + "mass": 0.05, + "use_usd_properties": False, + }, +} + PICK_SAMPLE_INTERVAL = 120 HAND_INTERP_STEPS = 12 POST_TRAJECTORY_STEPS = 240 + APPROACH_DIRECTIONS = { "top": (0.0, 0.0, -1.0), "side": (0.0, 1.0, 0.0), @@ -64,127 +90,283 @@ } -def parse_arguments() -> argparse.Namespace: - """Parse command-line arguments for the PickUp tutorial.""" - parser = argparse.ArgumentParser(description="Demonstrate PickUp on a cube.") - add_env_launcher_args_to_parser(parser) - parser.add_argument("--n_sample", type=int, default=10000) - parser.add_argument("--force_reannotate", action="store_true") - parser.add_argument("--auto_play", action="store_true") - parser.add_argument( - "--approach", choices=[*APPROACH_DIRECTIONS, "custom"], default="top" - ) - parser.add_argument("--custom_approach_direction", type=float, nargs=3) - parser.add_argument("--no_vis_eef_axis", action="store_true") - return parser.parse_args() +class PickUpDemo(DemoBase): + """Demo that picks up an object using an antipodal grasp affordance.""" + + def setup(self) -> None: + """Create simulation, robot, object, motion generator and action engine.""" + width, height = get_tutorial_window_size(self.args) + self.sim = create_default_sim( + self.args, + width=width, + height=height, + physics_dt=1.0 / 100.0, + arena_space=2.5, + ) + self.robot = self.sim.add_robot( + cfg=create_ur5_gripper_robot_cfg(init_pos=(0.0, 0.0, 0.0)) + ) + self.obj = self._create_pick_object(self.args.object) + + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=self.robot.uid)) + ) + hand_open, hand_close = self._get_hand_open_close_qpos() + approach_direction = self._resolve_approach_direction() + self._initialize_pre_pick_robot_pose(hand_open) + pickup_cfg = PickUpCfg( + control_part="arm", + hand_control_part="hand", + hand_open_qpos=hand_open, + hand_close_qpos=hand_close, + approach_direction=approach_direction, + pre_grasp_distance=0.15, + lift_height=0.16, + sample_interval=PICK_SAMPLE_INTERVAL, + hand_interp_steps=HAND_INTERP_STEPS, + ) + + self.atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + self.atomic_engine.register(PickUp(motion_gen, cfg=pickup_cfg)) + + self.semantics = self._create_object_semantics() + self.approach_direction = approach_direction + + maybe_open_window(self.sim, self.args) + if not self.args.no_vis_eef_axis: + self._draw_pick_object_axis() + + def run(self) -> None: + """Plan and replay the PickUp trajectory.""" + maybe_wait_for_user( + self.args, + f"Inspect the upright {self.args.object}, then press Enter to plan...", + ) + + logger.log_info( + f"Planning pick_up for {self.args.object} with " + f"approach_direction={format_tensor(self.approach_direction)}" + ) + start_time = time.time() + success, traj, _ = self.atomic_engine.run( + steps=[("pick_up", GraspTarget(semantics=self.semantics))] + ) + cost_time = time.time() - start_time + logger.log_info(f"Plan trajectory cost time: {cost_time:.2f} seconds") + if not success.all(): + logger.log_warning("Failed to plan pickup demo trajectory.") + return + + maybe_wait_for_user(self.args, "Press Enter to replay the pickup demo...") + + with DemoRecording( + self.sim, self.args, prefix=f"pickup_{self.args.object}_auto_play" + ): + self._replay_pickup_trajectory(traj) + + maybe_wait_for_user(self.args, "Press Enter to exit the simulation...") -def create_pick_object(sim) -> RigidObject: - """Create a settled cube for antipodal grasp planning.""" - obj = sim.add_rigid_object( - cfg=RigidObjectCfg( - uid="cube", - shape=CubeCfg(size=list(OBJECT_SIZE)), + def _create_pick_object(self, object_name: str) -> RigidObject: + preset = OBJECT_PRESETS[object_name] + cfg = RigidObjectCfg( + uid=preset["label"], + shape=MeshCfg(fpath=get_data_path(preset["mesh_path"])), attrs=RigidBodyAttributesCfg( - mass=0.05, + mass=preset["mass"], dynamic_friction=0.97, static_friction=0.99, ), max_convex_hull_num=16, - init_pos=[*OBJECT_XY, OBJECT_SIZE[2]], + init_pos=[OBJECT_XY[0], OBJECT_XY[1], 0.0], + init_rot=preset["init_rot"], + body_scale=preset["body_scale"], + use_usd_properties=preset["use_usd_properties"], ) - ) - sim.update(step=10) - clone_local_pose_from_first_env(obj) - obj.clear_dynamics() - return obj - - -def resolve_approach_direction( - args: argparse.Namespace, device: torch.device -) -> torch.Tensor: - """Resolve and validate a normalized approach direction.""" - direction = ( - args.custom_approach_direction - if args.approach == "custom" - else APPROACH_DIRECTIONS[args.approach] - ) - if direction is None: - raise ValueError( - "--custom_approach_direction is required for --approach custom." + obj = self.sim.add_rigid_object(cfg=cfg) + + # Settle the object to ensure it is resting on the ground before planning + self.sim.update(step=10) + return obj + + def _build_grasp_generator_cfg(self) -> GraspGeneratorCfg: + return GraspGeneratorCfg( + viser_port=11801, + antipodal_sampler_cfg=AntipodalSamplerCfg( + n_sample=self.args.n_sample, + max_length=GRIPPER_MAX_OPEN_WIDTH, + min_length=0.003, + ), + is_partial_annotate=False, + is_filter_ground_collision=False, ) - approach = torch.tensor(direction, dtype=torch.float32, device=device) - if torch.linalg.norm(approach) < 1e-6: - raise ValueError("approach_direction must be non-zero.") - return torch.nn.functional.normalize(approach, dim=0) + def _build_gripper_collision_cfg(self) -> GripperCollisionCfg: + return GripperCollisionCfg( + max_open_length=GRIPPER_MAX_OPEN_WIDTH, + finger_length=GRIPPER_FINGER_LENGTH, + y_thickness=GRIPPER_Y_THICKNESS, + root_z_width=GRIPPER_ROOT_Z_WIDTH, + open_check_margin=0.002, + point_sample_dense=0.012, + ) -def main() -> None: - """Plan and replay a sampled antipodal PickUp trajectory.""" - args = parse_arguments() - sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot(sim) - obj = create_pick_object(sim) - hand_open, hand_close = get_hand_open_close_qpos(robot) - initialize_pre_pick_robot_pose(robot, obj, hand_open) - motion_gen = create_toppra_motion_generator(robot) - - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register( - PickUp( - motion_gen, - cfg=PickUpCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - approach_direction=resolve_approach_direction(args, sim.device), - pre_grasp_distance=0.15, - lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, + def _create_object_semantics(self) -> ObjectSemantics: + label = OBJECT_PRESETS[self.args.object]["label"] + return ObjectSemantics( + label=label, + geometry={ + "mesh_vertices": self.obj.get_vertices(env_ids=[0], scale=True)[0], + "mesh_triangles": self.obj.get_triangles(env_ids=[0])[0], + }, + affordance=AntipodalAffordance( + mesh_vertices=self.obj.get_vertices(env_ids=[0], scale=True)[0], + mesh_triangles=self.obj.get_triangles(env_ids=[0])[0], + gripper_collision_cfg=self._build_gripper_collision_cfg(), + generator_cfg=self._build_grasp_generator_cfg(), + force_reannotate=self.args.force_reannotate, ), + entity=self.obj, + ) + + def _get_hand_open_close_qpos(self) -> tuple[torch.Tensor, torch.Tensor]: + hand_limits = self.robot.get_qpos_limits(name="hand")[0].to( + device=self.sim.device, dtype=torch.float32 + ) + hand_open = hand_limits[:, 0] + hand_close_limit = hand_limits[:, 1] + hand_close = torch.minimum( + hand_close_limit, + torch.full_like(hand_close_limit, OBJECT_MIN_HAND_CLOSE_QPOS), + ) + return hand_open, hand_close + + def _resolve_approach_direction(self) -> torch.Tensor: + if self.args.approach == "custom": + if self.args.custom_approach_direction is None: + raise ValueError( + "--custom_approach_direction is required when --approach custom." + ) + direction = self.args.custom_approach_direction + else: + direction = APPROACH_DIRECTIONS[self.args.approach] + + approach_direction = torch.tensor( + direction, dtype=torch.float32, device=self.sim.device + ) + norm = torch.linalg.norm(approach_direction) + if norm < 1e-6: + raise ValueError("approach_direction must be non-zero.") + return approach_direction / norm + + def _make_pre_pick_eef_pose(self, position: torch.Tensor) -> torch.Tensor: + pose = self.robot.compute_fk( + qpos=self.robot.get_qpos(name="arm"), + name="arm", + to_matrix=True, + ).clone() + pose[:, :3, 3] = position + return pose + + def _initialize_pre_pick_robot_pose(self, hand_open: torch.Tensor) -> None: + obj_pose = self.obj.get_local_pose(to_matrix=True) + move_position = obj_pose[:, :3, 3].clone() + move_position[:, 2] = 0.36 + pre_pick_pose = self._make_pre_pick_eef_pose(move_position) + ik_success, arm_qpos = self.robot.compute_ik( + pose=pre_pick_pose, + joint_seed=self.robot.get_qpos(name="arm"), + name="arm", + ) + if not torch.all(ik_success): + raise RuntimeError("Failed to initialize the robot at the pre-pick pose.") + + n_envs = self.robot.get_qpos().shape[0] + hand_qpos = hand_open.unsqueeze(0).repeat(n_envs, 1) + for target in (False, True): + self.robot.set_qpos(arm_qpos, name="arm", target=target) + self.robot.set_qpos(hand_qpos, name="hand", target=target) + self.robot.clear_dynamics() + + def _replay_pickup_trajectory(self, traj: torch.Tensor) -> None: + post_grasp_clear_step = self._compute_pick_close_end_step() + should_clear_object_dynamics = True + for i in range(traj.shape[1]): + self.robot.set_qpos(traj[:, i, :]) + self.sim.update(step=4) + if should_clear_object_dynamics and i + 1 >= post_grasp_clear_step: + self.obj.clear_dynamics() + should_clear_object_dynamics = False + logger.log_info(f"Object dynamics cleared after grasp at step={i}") + time.sleep(1e-2) + + logger.log_info( + f"PickUp keeps the upright {self.args.object} suspended in the gripper." ) + + final_qpos = traj[:, -1, :] + for _ in range(POST_TRAJECTORY_STEPS): + self.robot.set_qpos(final_qpos) + self.sim.update(step=2) + time.sleep(1e-2) + + @staticmethod + def _compute_pick_close_end_step() -> int: + motion_waypoints = PICK_SAMPLE_INTERVAL - HAND_INTERP_STEPS + n_approach = int(round(motion_waypoints) * 0.6) + return n_approach + HAND_INTERP_STEPS + + def _draw_pick_object_axis(self) -> None: + draw_axis_marker( + self.sim, + "pickup_object_axis", + self.obj.get_local_pose(to_matrix=True), + ) + + +def main() -> None: + """Entry point for the PickUp demo.""" + setup_print_options() + parser = argparse.ArgumentParser( + description="Demonstrate PickUp on an upright object." ) - semantics = create_antipodal_semantics( - obj, - label="cube", - n_sample=args.n_sample, - force_reannotate=args.force_reannotate, + parser = add_demo_args(parser) + parser.add_argument( + "--object", + choices=sorted(OBJECT_PRESETS.keys()), + default="sugar_box", + help="Object preset to pick.", ) - if not args.no_vis_eef_axis: - draw_axis_marker(sim, "pickup_object_axis", obj.get_local_pose(to_matrix=True)) - wait_for_user = prepare_tutorial_scene( - sim, args, "Inspect the cube, then press Enter to plan PickUp..." + parser.add_argument( + "--n_sample", + "--n-sample", + type=int, + default=10000, + help="Number of samples for antipodal grasp generation.", ) - - success, trajectory, _ = engine.run([("pick_up", GraspTarget(semantics))]) - if not success.all(): - logger.log_warning("Failed to plan PickUp demo trajectory.") - return - - if wait_for_user: - input("Press Enter to replay the PickUp demo...") - clear_after_step = ( - round((PICK_SAMPLE_INTERVAL - HAND_INTERP_STEPS) * 0.6) + HAND_INTERP_STEPS + parser.add_argument( + "--force_reannotate", + "--force-reannotate", + action="store_true", + help="Force grasp region re-annotation instead of using cached data.", ) - dynamics_cleared = False - - def clear_object_dynamics(step_idx: int, _: int) -> None: - nonlocal dynamics_cleared - if not dynamics_cleared and step_idx + 1 >= clear_after_step: - obj.clear_dynamics() - dynamics_cleared = True - - replay_trajectory( - sim, - robot, - trajectory, - args, - video_prefix="pickup_cube_auto_play", - hold_steps=POST_TRAJECTORY_STEPS, - on_trajectory_step=clear_object_dynamics, + parser.add_argument( + "--approach", + choices=["top", "side", "side_y", "custom"], + default="top", + help="Pick approach direction preset.", + ) + parser.add_argument( + "--custom_approach_direction", + "--custom-approach-direction", + type=float, + nargs=3, + default=None, + metavar=("X", "Y", "Z"), + help="World-frame approach direction used when --approach custom.", ) - if wait_for_user: - input("Press Enter to exit the simulation...") + args = parser.parse_args() + PickUpDemo(args).main() if __name__ == "__main__": diff --git a/scripts/tutorials/atomic_action/place.py b/scripts/tutorials/atomic_action/place.py index 6a8752d5a..99fc46e77 100644 --- a/scripts/tutorials/atomic_action/place.py +++ b/scripts/tutorials/atomic_action/place.py @@ -19,46 +19,64 @@ from __future__ import annotations import argparse -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) +import time import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.data import get_data_path from embodichain.lab.sim.atomic_actions import ( + AntipodalAffordance, AtomicActionEngine, EndEffectorPoseTarget, GraspTarget, + ObjectSemantics, PickUp, PickUpCfg, Place, PlaceCfg, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.demo_base import DemoBase from embodichain.lab.sim.objects import RigidObject -from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, +) +from embodichain.toolkits.graspkit.pg_grasp.antipodal_generator import ( + AntipodalSamplerCfg, + GraspGeneratorCfg, +) +from embodichain.toolkits.graspkit.pg_grasp.gripper_collision_checker import ( + GripperCollisionCfg, +) from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, - broadcast_pose_batch, - broadcast_waypoint_pose_batch, - clone_local_pose_from_first_env, - create_antipodal_semantics, - create_toppra_motion_generator, - create_tutorial_simulation, + create_ur5_gripper_robot_cfg, draw_axis_marker, - get_hand_open_close_qpos, - initialize_pre_pick_robot_pose, - prepare_tutorial_scene, - replay_trajectory, + get_tutorial_window_size, ) -OBJECT_SIZE = (0.05, 0.05, 0.05) +GRIPPER_MAX_OPEN_WIDTH = 0.080 +GRIPPER_FINGER_LENGTH = 0.088 +GRIPPER_ROOT_Z_WIDTH = 0.096 +GRIPPER_Y_THICKNESS = 0.040 + +OBJECT_LABEL = "sugar_box" +OBJECT_MESH_PATH = "SugarBox/sugar_box_usd/sugar_box.usda" OBJECT_XY = (-0.42, -0.08) +OBJECT_MIN_HAND_CLOSE_QPOS = 0.024 +OBJECT_APPROACH_DIRECTION = (0.0, 0.0, -1.0) +OBJECT_INIT_ROT = (0.0, 0.0, 0.0) +OBJECT_BODY_SCALE = (0.8, 0.8, 0.8) +OBJECT_MASS = 0.05 +OBJECT_USE_USD_PROPERTIES = False + PICK_SAMPLE_INTERVAL = 120 PLACE_SAMPLE_INTERVAL = 120 HAND_INTERP_STEPS = 12 @@ -66,155 +84,273 @@ PLACE_LIFT_HEIGHT = 0.14 -def parse_arguments() -> argparse.Namespace: - """Parse command-line arguments for the Place tutorial.""" - parser = argparse.ArgumentParser( - description="Pick up a cube and place it at a target pose." - ) - add_env_launcher_args_to_parser(parser) - parser.add_argument("--n_sample", type=int, default=10000) - parser.add_argument("--force_reannotate", action="store_true") - parser.add_argument("--auto_play", action="store_true") - parser.add_argument("--no_vis_eef_axis", action="store_true") - return parser.parse_args() - - -def create_pick_object(sim) -> RigidObject: - """Create a settled cube for the PickUp and Place sequence.""" - obj = sim.add_rigid_object( - cfg=RigidObjectCfg( - uid="cube", - shape=CubeCfg(size=list(OBJECT_SIZE)), +class PlaceDemo(DemoBase): + """Demo that picks up an object and places it at a target end-effector pose.""" + + def setup(self) -> None: + """Create simulation, robot, object, motion generator and action engine.""" + width, height = get_tutorial_window_size(self.args) + self.sim = create_default_sim( + self.args, + width=width, + height=height, + physics_dt=1.0 / 100.0, + arena_space=2.5, + ) + self.robot = self.sim.add_robot( + cfg=create_ur5_gripper_robot_cfg(init_pos=(0.0, 0.0, 0.0)) + ) + self.obj = self._create_pick_object() + + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=self.robot.uid)) + ) + + hand_open, hand_close = self._get_hand_open_close_qpos() + self._initialize_pre_pick_robot_pose(hand_open) + pickup_cfg = PickUpCfg( + control_part="arm", + hand_control_part="hand", + hand_open_qpos=hand_open, + hand_close_qpos=hand_close, + approach_direction=torch.tensor( + OBJECT_APPROACH_DIRECTION, + dtype=torch.float32, + device=self.sim.device, + ), + pre_grasp_distance=0.15, + lift_height=0.16, + sample_interval=PICK_SAMPLE_INTERVAL, + hand_interp_steps=HAND_INTERP_STEPS, + ) + place_cfg = PlaceCfg( + control_part="arm", + hand_control_part="hand", + hand_open_qpos=hand_open, + hand_close_qpos=hand_close, + lift_height=PLACE_LIFT_HEIGHT, + sample_interval=PLACE_SAMPLE_INTERVAL, + hand_interp_steps=HAND_INTERP_STEPS, + ) + + self.atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + self.atomic_engine.register(PickUp(motion_gen, cfg=pickup_cfg)) + self.atomic_engine.register(Place(motion_gen, cfg=place_cfg)) + + self.semantics = self._create_object_semantics() + self.place_eef_poses = self._make_place_eef_poses() + + maybe_open_window(self.sim, self.args) + if not self.args.no_vis_eef_axis: + draw_axis_marker(self.sim, "place_target_axis", self.place_eef_poses[-1]) + + def run(self) -> None: + """Plan and replay the PickUp -> Place trajectory.""" + maybe_wait_for_user( + self.args, + "Inspect the object, then press Enter to plan PickUp -> Place...", + ) + + n_envs = self.robot.get_qpos().shape[0] + multi_waypoint_xpos = self.place_eef_poses.unsqueeze(0).repeat(n_envs, 1, 1, 1) + place_target = EndEffectorPoseTarget(xpos=multi_waypoint_xpos) + logger.log_info("Planning PickUp precondition -> Place release trajectory") + success, traj, _ = self.atomic_engine.run( + steps=[ + ("pick_up", GraspTarget(semantics=self.semantics)), + ("place", place_target), + ] + ) + if not success.all(): + logger.log_warning("Failed to plan Place demo trajectory.") + return + + maybe_wait_for_user(self.args, "Press Enter to replay the Place demo...") + + with DemoRecording(self.sim, self.args, prefix="place_auto_play"): + self._replay_place_trajectory(traj) + + maybe_wait_for_user(self.args, "Press Enter to exit the simulation...") + + def _create_pick_object(self) -> RigidObject: + cfg = RigidObjectCfg( + uid=OBJECT_LABEL, + shape=MeshCfg(fpath=get_data_path(OBJECT_MESH_PATH)), attrs=RigidBodyAttributesCfg( - mass=0.05, + mass=OBJECT_MASS, dynamic_friction=0.97, static_friction=0.99, enable_ccd=True, ), max_convex_hull_num=16, - init_pos=[*OBJECT_XY, 0.5 * OBJECT_SIZE[2]], + init_pos=[OBJECT_XY[0], OBJECT_XY[1], 0.05], + init_rot=OBJECT_INIT_ROT, + body_scale=OBJECT_BODY_SCALE, + use_usd_properties=OBJECT_USE_USD_PROPERTIES, ) - ) - sim.update(step=10) - clone_local_pose_from_first_env(obj) - obj.clear_dynamics() - return obj - - -def make_place_eef_poses(device: torch.device) -> torch.Tensor: - """Build hover and release waypoints for the multi-waypoint Place target.""" - rotation = torch.tensor( - [ - [0.0539, 0.9985, -0.0022], - [0.9977, -0.0540, -0.0401], - [-0.0401, 0.0, -0.9992], - ], - dtype=torch.float32, - device=device, - ) - poses = [] - for position in ((-0.40, 0.48, 0.20), (-0.40, 0.48, 0.10)): - pose = torch.eye(4, dtype=torch.float32, device=device) - pose[:3, :3], pose[:3, 3] = rotation, torch.tensor(position, device=device) - poses.append(pose) - return torch.stack(poses) + obj = self.sim.add_rigid_object(cfg=cfg) + # Set the object to a stable pose on the ground by simulating a few steps. + self.sim.update(step=10) + return obj -def main() -> None: - """Plan and replay PickUp followed by a multi-waypoint Place.""" - args = parse_arguments() - sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot(sim) - obj = create_pick_object(sim) - motion_gen = create_toppra_motion_generator(robot) - hand_open, hand_close = get_hand_open_close_qpos(robot) - initialize_pre_pick_robot_pose(robot, obj, hand_open) - - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register( - PickUp( - motion_gen, - cfg=PickUpCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - pre_grasp_distance=0.15, - lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, - ), + def _get_hand_open_close_qpos(self) -> tuple[torch.Tensor, torch.Tensor]: + hand_limits = self.robot.get_qpos_limits(name="hand")[0].to( + device=self.sim.device, dtype=torch.float32 ) - ) - engine.register( - Place( - motion_gen, - cfg=PlaceCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - lift_height=PLACE_LIFT_HEIGHT, - sample_interval=PLACE_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, + hand_open = hand_limits[:, 0] + hand_close_limit = hand_limits[:, 1] + hand_close = torch.minimum( + hand_close_limit, + torch.full_like(hand_close_limit, OBJECT_MIN_HAND_CLOSE_QPOS), + ) + return hand_open, hand_close + + def _build_grasp_generator_cfg(self) -> GraspGeneratorCfg: + return GraspGeneratorCfg( + viser_port=11801, + antipodal_sampler_cfg=AntipodalSamplerCfg( + n_sample=self.args.n_sample, + max_length=GRIPPER_MAX_OPEN_WIDTH, + min_length=0.003, ), + is_partial_annotate=False, + is_filter_ground_collision=False, ) - ) - semantics = create_antipodal_semantics( - obj, - label="cube", - n_sample=args.n_sample, - force_reannotate=args.force_reannotate, - ) - place_poses = make_place_eef_poses(sim.device) - if not args.no_vis_eef_axis: - draw_axis_marker( - sim, - "place_target_axis", - broadcast_pose_batch(place_poses[-1], robot.get_qpos().shape[0]), - ) - wait_for_user = prepare_tutorial_scene( - sim, args, "Inspect the cube, then press Enter to plan PickUp -> Place..." - ) + def _build_gripper_collision_cfg(self) -> GripperCollisionCfg: + return GripperCollisionCfg( + max_open_length=GRIPPER_MAX_OPEN_WIDTH, + finger_length=GRIPPER_FINGER_LENGTH, + y_thickness=GRIPPER_Y_THICKNESS, + root_z_width=GRIPPER_ROOT_Z_WIDTH, + open_check_margin=0.002, + point_sample_dense=0.012, + ) - success, trajectory, _ = engine.run( - [ - ("pick_up", GraspTarget(semantics)), - ( - "place", - EndEffectorPoseTarget( - broadcast_waypoint_pose_batch( - place_poses, robot.get_qpos().shape[0] - ) - ), + def _create_object_semantics(self) -> ObjectSemantics: + return ObjectSemantics( + label=OBJECT_LABEL, + geometry={ + "mesh_vertices": self.obj.get_vertices(env_ids=[0], scale=True)[0], + "mesh_triangles": self.obj.get_triangles(env_ids=[0])[0], + }, + affordance=AntipodalAffordance( + mesh_vertices=self.obj.get_vertices(env_ids=[0], scale=True)[0], + mesh_triangles=self.obj.get_triangles(env_ids=[0])[0], + gripper_collision_cfg=self._build_gripper_collision_cfg(), + generator_cfg=self._build_grasp_generator_cfg(), + force_reannotate=self.args.force_reannotate, ), - ] + entity=self.obj, + ) + + def _make_pre_pick_eef_pose(self, position: torch.Tensor) -> torch.Tensor: + pose = self.robot.compute_fk( + qpos=self.robot.get_qpos(name="arm"), + name="arm", + to_matrix=True, + ).clone() + pose[:, :3, 3] = position + return pose + + def _initialize_pre_pick_robot_pose(self, hand_open: torch.Tensor) -> None: + obj_pose = self.obj.get_local_pose(to_matrix=True) + move_position = obj_pose[:, :3, 3].clone() + move_position[:, 2] = 0.36 + pre_pick_pose = self._make_pre_pick_eef_pose(move_position) + ik_success, arm_qpos = self.robot.compute_ik( + pose=pre_pick_pose, + joint_seed=self.robot.get_qpos(name="arm"), + name="arm", + ) + if not torch.all(ik_success): + raise RuntimeError("Failed to initialize the robot at the pre-pick pose.") + + n_envs = self.robot.get_qpos().shape[0] + hand_qpos = hand_open.unsqueeze(0).repeat(n_envs, 1) + for target in (False, True): + self.robot.set_qpos(arm_qpos, name="arm", target=target) + self.robot.set_qpos(hand_qpos, name="hand", target=target) + self.robot.clear_dynamics() + + def _make_place_eef_poses(self) -> torch.Tensor: + """Build a multi-waypoint place trajectory ``(n_waypoint, 4, 4)``. + + Two waypoints are returned: a higher hover pose and the final release pose. + ``Place`` approaches from above the first waypoint, descends through each + waypoint in order, opens the gripper at the last, and retracts — so this + exercises the multi-waypoint descent path. + """ + rotation = torch.tensor( + [ + [0.0539, 0.9985, -0.0022], + [0.9977, -0.0540, -0.0401], + [-0.0401, -0.0000, -0.9992], + ], + dtype=torch.float32, + device=self.sim.device, + ) + hover_pose = torch.eye(4, dtype=torch.float32, device=self.sim.device) + hover_pose[:3, :3] = rotation + hover_pose[:3, 3] = torch.tensor( + [-0.40, 0.48, 0.20], dtype=torch.float32, device=self.sim.device + ) + place_pose = torch.eye(4, dtype=torch.float32, device=self.sim.device) + place_pose[:3, :3] = rotation + place_pose[:3, 3] = torch.tensor( + [-0.40, 0.48, 0.10], dtype=torch.float32, device=self.sim.device + ) + return torch.stack([hover_pose, place_pose], dim=0) + + def _replay_place_trajectory(self, traj: torch.Tensor) -> None: + post_grasp_clear_step = self._compute_pick_close_end_step() + should_clear_object_dynamics = True + for i in range(traj.shape[1]): + self.robot.set_qpos(traj[:, i, :]) + self.sim.update(step=4) + if should_clear_object_dynamics and i + 1 >= post_grasp_clear_step: + self.obj.clear_dynamics() + should_clear_object_dynamics = False + logger.log_info(f"Object dynamics cleared after grasp at step={i}") + time.sleep(1e-2) + + logger.log_info("Place opens the gripper and clears WorldState.held_object.") + final_qpos = traj[:, -1, :] + for _ in range(POST_TRAJECTORY_STEPS): + self.robot.set_qpos(final_qpos) + self.sim.update(step=2) + time.sleep(1e-2) + + @staticmethod + def _compute_pick_close_end_step() -> int: + motion_waypoints = PICK_SAMPLE_INTERVAL - HAND_INTERP_STEPS + n_approach = int(round(motion_waypoints) * 0.6) + return n_approach + HAND_INTERP_STEPS + + +def main() -> None: + """Entry point for the Place demo.""" + setup_print_options() + parser = argparse.ArgumentParser( + description="Demonstrate Place by picking an object and releasing it at a target pose." ) - if not success.all(): - logger.log_warning("Failed to plan Place demo trajectory.") - return - - if wait_for_user: - input("Press Enter to replay the Place demo...") - clear_after_step = ( - round((PICK_SAMPLE_INTERVAL - HAND_INTERP_STEPS) * 0.6) + HAND_INTERP_STEPS + parser = add_demo_args(parser) + parser.add_argument( + "--n_sample", + "--n-sample", + type=int, + default=10000, + help="Number of samples for antipodal grasp generation.", ) - dynamics_cleared = False - - def clear_object_dynamics(step_idx: int, _: int) -> None: - nonlocal dynamics_cleared - if not dynamics_cleared and step_idx + 1 >= clear_after_step: - obj.clear_dynamics() - dynamics_cleared = True - - replay_trajectory( - sim, - robot, - trajectory, - args, - video_prefix="place_auto_play", - hold_steps=POST_TRAJECTORY_STEPS, - on_trajectory_step=clear_object_dynamics, + parser.add_argument( + "--force_reannotate", + "--force-reannotate", + action="store_true", + help="Force grasp region re-annotation instead of using cached data.", ) - if wait_for_user: - input("Press Enter to exit the simulation...") + args = parser.parse_args() + PlaceDemo(args).main() if __name__ == "__main__": diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py index 10deaad3a..2ca2e96e4 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -19,16 +19,11 @@ from __future__ import annotations import argparse -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) +import time import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.data import get_data_path from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, EndEffectorPoseTarget, @@ -38,57 +33,207 @@ PressCfg, ) from embodichain.lab.sim.cfg import ( + JointDrivePropertiesCfg, RigidBodyAttributesCfg, RigidObjectCfg, + RobotCfg, + URDFCfg, ) +from embodichain.lab.sim.demo_base import DemoBase from embodichain.lab.sim.material import VisualMaterialCfg -from embodichain.lab.sim.objects import RigidObject +from embodichain.lab.sim.objects import RigidObject, Robot +from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + format_tensor, + maybe_init_gpu_physics, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, +) from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, - broadcast_pose_batch, - create_toppra_motion_generator, - create_tutorial_simulation, draw_axis_marker, - format_tensor, - get_hand_open_close_qpos, - make_top_down_eef_pose, - prepare_tutorial_scene, - replay_trajectory, + get_tutorial_window_size, + make_ur5_solver_cfg, ) +GRIPPER_URDF_PATH = "DH_PGI_140_80/DH_PGI_140_80.urdf" +GRIPPER_HAND_JOINT_PATTERN = "gripper_finger1_joint_1" +ARM_JOINT_PATTERN = "joint[0-9]" +GRIPPER_TCP_Z = 0.15 + MOVE_SAMPLE_INTERVAL = 60 PRESS_SAMPLE_INTERVAL = 90 HAND_INTERP_STEPS = 12 POST_TRAJECTORY_STEPS = 180 -BLOCK_SIZE = (0.12, 0.12, 0.06) +TABLE_SIZE = [1.0, 1.4, 0.05] +TABLE_TOP_Z = -0.045 +BLOCK_SIZE = [0.12, 0.12, 0.06] +BLOCK_CENTER = [-0.30, -0.12, TABLE_TOP_Z + 0.5 * BLOCK_SIZE[2]] PRESS_CLEARANCE = 0.13 PRESS_SURFACE_OFFSET = 0.003 DEFAULT_PRESS_TOLERANCE = 0.01 -def parse_arguments() -> argparse.Namespace: - """Parse command-line arguments for the Press tutorial.""" - parser = argparse.ArgumentParser(description="Demonstrate Press on a wooden block.") - add_env_launcher_args_to_parser(parser) - parser.add_argument("--auto_play", action="store_true") - parser.add_argument("--debug_state", action="store_true") - parser.add_argument( - "--press_tolerance", type=float, default=DEFAULT_PRESS_TOLERANCE - ) - parser.add_argument("--block_pos", type=float, nargs=2, default=(-0.30, -0.12)) - parser.add_argument("--no_vis_eef_axis", action="store_true") - return parser.parse_args() +class PressDemo(DemoBase): + """Demo that presses the center of a regular wooden block.""" + + def setup(self) -> None: + """Create simulation, custom robot, table, block and action engine.""" + width, height = get_tutorial_window_size(self.args) + self.sim = create_default_sim( + self.args, + width=width, + height=height, + physics_dt=1.0 / 100.0, + arena_space=2.5, + ) + self.robot = self._create_robot() + self._create_table() + block_center = [ + self.args.block_pos[0], + self.args.block_pos[1], + TABLE_TOP_Z + 0.5 * BLOCK_SIZE[2], + ] + self.block = self._create_wooden_block(block_center) + + self._settle_object(self.block, step=5) + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=self.robot.uid)) + ) + hand_close = self._get_hand_close_qpos() + + self.atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + self.atomic_engine.register( + MoveEndEffector( + motion_gen, + cfg=MoveEndEffectorCfg( + control_part="arm", + sample_interval=MOVE_SAMPLE_INTERVAL, + ), + ) + ) + self.atomic_engine.register( + Press( + motion_gen, + cfg=PressCfg( + control_part="arm", + hand_control_part="hand", + hand_close_qpos=hand_close, + sample_interval=PRESS_SAMPLE_INTERVAL, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ) + ) + + block_pose = self.block.get_local_pose(to_matrix=True) + block_top_z = block_pose[0, 2, 3] + 0.5 * BLOCK_SIZE[2] + press_position = block_pose[0, :3, 3].clone() + press_position[2] = block_top_z + PRESS_SURFACE_OFFSET + move_position = press_position.clone() + move_position[2] = block_top_z + PRESS_CLEARANCE + + self.move_target = self._make_top_down_eef_pose(move_position) + self.press_target = self._make_top_down_eef_pose(press_position) + + maybe_open_window(self.sim, self.args) + if not self.args.no_vis_eef_axis: + self._draw_press_target_axis(self.press_target) + + def run(self) -> None: + """Plan and replay the MoveEndEffector -> Press trajectory.""" + maybe_wait_for_user( + self.args, "Inspect the wooden block, then press Enter to plan..." + ) + + logger.log_info("Planning MoveEndEffector -> Press") + start_time = time.time() + success, traj, _ = self.atomic_engine.run( + steps=[ + ("move_end_effector", EndEffectorPoseTarget(xpos=self.move_target)), + ("press", EndEffectorPoseTarget(xpos=self.press_target)), + ] + ) + cost_time = time.time() - start_time + logger.log_info(f"Plan trajectory cost time: {cost_time:.2f} seconds") + if not success.all(): + logger.log_warning("Failed to plan Press demo trajectory.") + return + + is_center_hit, center_error, hit_step, hit_pos, expected_pos = ( + self._compute_press_center_check(traj) + ) + logger.log_info( + "Press center check: " + f"success={is_center_hit}, " + f"xy_error={center_error:.4f} m, " + f"hit_step={hit_step}, " + f"hit_pos={format_tensor(hit_pos)}, " + f"expected={format_tensor(expected_pos)}" + ) + if not is_center_hit: + logger.log_warning( + "Press planned trajectory did not reach the block center within " + f"{self.args.press_tolerance:.4f} m." + ) + return + + maybe_wait_for_user(self.args, "Press Enter to replay the Press demo...") + + with DemoRecording(self.sim, self.args, prefix="press_auto_play"): + self._replay_press_trajectory(traj) + + maybe_wait_for_user(self.args, "Press Enter to exit the simulation...") + def _create_robot(self) -> Robot: + ur5_urdf_path = get_data_path("UniversalRobots/UR5/UR5.urdf") + gripper_urdf_path = get_data_path(GRIPPER_URDF_PATH) + cfg = RobotCfg( + uid="UR5", + urdf_cfg=URDFCfg( + components=[ + {"component_type": "arm", "urdf_path": ur5_urdf_path}, + {"component_type": "hand", "urdf_path": gripper_urdf_path}, + ] + ), + drive_pros=JointDrivePropertiesCfg( + stiffness={ARM_JOINT_PATTERN: 1e4, GRIPPER_HAND_JOINT_PATTERN: 1e3}, + damping={ARM_JOINT_PATTERN: 1e3, GRIPPER_HAND_JOINT_PATTERN: 1e2}, + max_effort={ARM_JOINT_PATTERN: 1e5, GRIPPER_HAND_JOINT_PATTERN: 1e4}, + drive_type="force", + ), + control_parts={ + "arm": [ARM_JOINT_PATTERN], + "hand": [GRIPPER_HAND_JOINT_PATTERN], + }, + solver_cfg={"arm": make_ur5_solver_cfg(GRIPPER_TCP_Z)}, + init_qpos=[0.0, -1.57, 1.57, -1.57, -1.57, 0.0, 0.0, 0.0], + init_pos=(0.0, 0.0, 0.0), + ) + return self.sim.add_robot(cfg=cfg) -def create_wooden_block(sim, center: list[float]) -> RigidObject: - """Create the static block used as a press target.""" - return sim.add_rigid_object( - cfg=RigidObjectCfg( + def _create_table(self) -> RigidObject: + cfg = RigidObjectCfg( + uid="table", + shape=CubeCfg(size=TABLE_SIZE), + body_type="static", + attrs=RigidBodyAttributesCfg( + dynamic_friction=0.8, + static_friction=0.9, + ), + init_pos=[-0.30, 0.10, TABLE_TOP_Z - 0.5 * TABLE_SIZE[2]], + ) + return self.sim.add_rigid_object(cfg=cfg) + + def _create_wooden_block(self, block_center: list[float]) -> RigidObject: + cfg = RigidObjectCfg( uid="wooden_block", shape=CubeCfg( - size=list(BLOCK_SIZE), + size=BLOCK_SIZE, visual_material=VisualMaterialCfg( uid="wooden_block_mat", base_color=[0.58, 0.32, 0.14, 1.0], @@ -96,145 +241,147 @@ def create_wooden_block(sim, center: list[float]) -> RigidObject: ), ), body_type="static", - attrs=RigidBodyAttributesCfg(dynamic_friction=0.8, static_friction=0.9), - init_pos=center, + attrs=RigidBodyAttributesCfg( + dynamic_friction=0.8, + static_friction=0.9, + ), + init_pos=block_center, ) - ) + return self.sim.add_rigid_object(cfg=cfg) + def _settle_object(self, obj: RigidObject, step: int = 5) -> None: + maybe_init_gpu_physics(self.sim) + obj.reset() + self.sim.update(step=step) + obj.clear_dynamics() -def compute_press_center_check( - robot, - trajectory: torch.Tensor, - block: RigidObject, - tolerance: float, -) -> tuple[bool, float, int, torch.Tensor, torch.Tensor]: - """Return whether the press trajectory reaches the block center tolerance.""" - arm_joint_ids = robot.get_joint_ids(name="arm") - start = MOVE_SAMPLE_INTERVAL + HAND_INTERP_STEPS - arm_traj = trajectory[ - :, start : MOVE_SAMPLE_INTERVAL + PRESS_SAMPLE_INTERVAL, arm_joint_ids - ] - fk_pose = torch.stack( - [ - robot.compute_fk(qpos=qpos, name="arm", to_matrix=True) - for qpos in arm_traj.unbind(dim=1) - ], - dim=1, - ) - block_center = block.get_local_pose(to_matrix=True)[:, :3, 3] - target_z = block_center[:, 2] + 0.5 * BLOCK_SIZE[2] + PRESS_SURFACE_OFFSET - xy_error = torch.linalg.norm( - fk_pose[:, :, :2, 3] - block_center[:, None, :2], dim=2 - ) - z_error = torch.abs(fk_pose[:, :, 2, 3] - target_z[:, None]) - best_idx = (xy_error + z_error).argmin(dim=1) - env_idx = torch.arange(trajectory.shape[0], device=trajectory.device) - best_pos = fk_pose[env_idx, best_idx, :3, 3] - center_error = torch.linalg.norm(best_pos[:, :2] - block_center[:, :2], dim=1) - worst_env = int(center_error.argmax().item()) - expected = torch.stack( - [block_center[worst_env, 0], block_center[worst_env, 1], target_z[worst_env]] - ) - return ( - bool(torch.all(center_error <= tolerance)), - float(center_error[worst_env].item()), - start + int(best_idx[worst_env].item()), - best_pos[worst_env], - expected, - ) + def _get_hand_close_qpos(self) -> torch.Tensor: + hand_limits = self.robot.get_qpos_limits(name="hand")[0].to( + device=self.sim.device, dtype=torch.float32 + ) + return hand_limits[:, 1] + + def _make_top_down_eef_pose(self, position: torch.Tensor) -> torch.Tensor: + pose = torch.eye(4, dtype=torch.float32, device=position.device) + pose[:3, :3] = torch.tensor( + [ + [-0.0539, -0.9985, -0.0022], + [-0.9977, 0.0540, -0.0401], + [0.0401, 0.0000, -0.9992], + ], + dtype=torch.float32, + device=position.device, + ) + pose[:3, 3] = position + return pose + def _draw_press_target_axis(self, press_target: torch.Tensor) -> None: + if press_target.dim() == 2: + press_target = press_target.unsqueeze(0) + draw_axis_marker(self.sim, "press_target_axis", press_target) -def main() -> None: - """Plan, verify, and replay MoveEndEffector followed by Press.""" - args = parse_arguments() - sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot(sim) - block = create_wooden_block(sim, [*args.block_pos, 0.5 * BLOCK_SIZE[2]]) - if sim.device.type == "cuda": - sim.init_gpu_physics() - block.reset() - sim.update(step=5) - block.clear_dynamics() - - motion_gen = create_toppra_motion_generator(robot) - hand_close = get_hand_open_close_qpos(robot)[1] - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register( - MoveEndEffector( - motion_gen, MoveEndEffectorCfg(sample_interval=MOVE_SAMPLE_INTERVAL) + def _compute_press_center_check( + self, traj: torch.Tensor + ) -> tuple[bool, float, int, torch.Tensor, torch.Tensor]: + arm_joint_ids = self.robot.get_joint_ids(name="arm") + press_segment_start = MOVE_SAMPLE_INTERVAL + HAND_INTERP_STEPS + press_segment_end = MOVE_SAMPLE_INTERVAL + PRESS_SAMPLE_INTERVAL + arm_traj = traj[:, press_segment_start:press_segment_end, arm_joint_ids] + fk_pose = torch.stack( + [ + self.robot.compute_fk( + qpos=waypoint.unsqueeze(0), + name="arm", + to_matrix=True, + )[0] + for waypoint in arm_traj[0] + ], + dim=0, ) - ) - engine.register( - Press( - motion_gen, - PressCfg( - hand_close_qpos=hand_close, - sample_interval=PRESS_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, + + block_pose = self.block.get_local_pose(to_matrix=True) + block_center = block_pose[0, :3, 3] + block_top_z = block_center[2] + 0.5 * BLOCK_SIZE[2] + target_xy = block_center[:2] + target_z = block_top_z + PRESS_SURFACE_OFFSET + + xy_error = torch.linalg.norm(fk_pose[:, :2, 3] - target_xy, dim=1) + z_error = torch.abs(fk_pose[:, 2, 3] - target_z) + combined_error = xy_error + z_error + best_idx = int(torch.argmin(combined_error).item()) + best_pos = fk_pose[best_idx, :3, 3] + center_error = float(torch.linalg.norm(best_pos[:2] - target_xy).item()) + return ( + center_error <= self.args.press_tolerance, + center_error, + press_segment_start + best_idx, + best_pos, + torch.tensor( + [target_xy[0], target_xy[1], target_z], + dtype=torch.float32, + device=traj.device, ), ) - ) - block_center = block.get_local_pose(to_matrix=True)[0, :3, 3] - press_position = block_center.clone() - press_position[2] += 0.5 * BLOCK_SIZE[2] + PRESS_SURFACE_OFFSET - move_position = press_position.clone() - move_position[2] += PRESS_CLEARANCE - PRESS_SURFACE_OFFSET - n_envs = robot.get_qpos().shape[0] - move_target = broadcast_pose_batch(make_top_down_eef_pose(move_position), n_envs) - press_target = broadcast_pose_batch(make_top_down_eef_pose(press_position), n_envs) - if not args.no_vis_eef_axis: - draw_axis_marker(sim, "press_target_axis", press_target) - wait_for_user = prepare_tutorial_scene( - sim, args, "Inspect the wooden block, then press Enter to plan..." - ) + def _replay_press_trajectory(self, traj: torch.Tensor) -> None: + log_stride = max(1, traj.shape[1] // 10) + for i in range(traj.shape[1]): + self.robot.set_qpos(traj[:, i, :]) + self.sim.update(step=4) + if self.args.debug_state and ( + i % log_stride == 0 or i == traj.shape[1] - 1 + ): + self._log_block_state(f"replay step {i}/{traj.shape[1] - 1}") + time.sleep(1e-2) - success, trajectory, _ = engine.run( - [ - ("move_end_effector", EndEffectorPoseTarget(move_target)), - ("press", EndEffectorPoseTarget(press_target)), - ] + logger.log_info("Press returned the end-effector to the starting pose.") + final_qpos = traj[:, -1, :] + for i in range(POST_TRAJECTORY_STEPS): + self.robot.set_qpos(final_qpos) + self.sim.update(step=2) + if self.args.debug_state and i % max(1, POST_TRAJECTORY_STEPS // 5) == 0: + self._log_block_state(f"post step {i}/{POST_TRAJECTORY_STEPS - 1}") + time.sleep(1e-2) + + def _log_block_state(self, label: str) -> None: + block_pose = self.block.get_local_pose(to_matrix=True) + logger.log_info( + f"{label}: pos={format_tensor(block_pose[0, :3, 3])}, " + f"z_axis={format_tensor(block_pose[0, :3, 2])}" + ) + + +def main() -> None: + """Entry point for the Press demo.""" + setup_print_options() + parser = argparse.ArgumentParser( + description="Demonstrate Press on the center of a regular wooden block." ) - if not success.all(): - logger.log_warning("Failed to plan Press demo trajectory.") - return - is_center_hit, center_error, hit_step, hit_pos, expected_pos = ( - compute_press_center_check(robot, trajectory, block, args.press_tolerance) + parser = add_demo_args(parser) + parser.add_argument( + "--debug_state", + "--debug-state", + action="store_true", + help="Log the block pose during replay.", ) - logger.log_info( - "Press center check: " - f"success={is_center_hit}, xy_error={center_error:.4f} m, hit_step={hit_step}, " - f"hit_pos={format_tensor(hit_pos)}, expected={format_tensor(expected_pos)}" + parser.add_argument( + "--press_tolerance", + "--press-tolerance", + type=float, + default=DEFAULT_PRESS_TOLERANCE, + help="XY tolerance in meters for checking whether press reaches block center.", ) - if not is_center_hit: - logger.log_warning( - "Press trajectory did not reach the block center within tolerance." - ) - return - - if wait_for_user: - input("Press Enter to replay the Press demo...") - - def log_state(step_idx: int, total_steps: int) -> None: - if args.debug_state and ( - step_idx % max(1, total_steps // 10) == 0 or step_idx == total_steps - 1 - ): - logger.log_info( - f"replay step {step_idx}/{total_steps - 1}: " - f"pos={format_tensor(block.get_local_pose(to_matrix=True)[0, :3, 3])}" - ) - - replay_trajectory( - sim, - robot, - trajectory, - args, - video_prefix="press_auto_play", - hold_steps=POST_TRAJECTORY_STEPS, - on_trajectory_step=log_state, + parser.add_argument( + "--block_pos", + "--block-pos", + type=float, + nargs=2, + default=BLOCK_CENTER[:2], + metavar=("X", "Y"), + help="Initial XY position of the wooden block center.", ) - if wait_for_user: - input("Press Enter to exit the simulation...") + args = parser.parse_args() + PressDemo(args).main() if __name__ == "__main__": diff --git a/scripts/tutorials/grasp/grasp_generator.py b/scripts/tutorials/grasp/grasp_generator.py index c20596142..e8b2e765e 100644 --- a/scripts/tutorials/grasp/grasp_generator.py +++ b/scripts/tutorials/grasp/grasp_generator.py @@ -19,25 +19,25 @@ in a simulated environment using the SimulationManager and grasp planning utilities. """ +from __future__ import annotations + import argparse -import numpy as np import time + +import numpy as np import torch -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.objects import Robot, RigidObject from embodichain.lab.sim.utility.action_utils import interpolate_with_distance from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim.solvers import URSolverCfg from embodichain.data import get_data_path -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from dexsim.utility.path import get_resources_data_path from embodichain.utils import logger from embodichain.lab.sim.cfg import ( - RenderCfg, JointDrivePropertiesCfg, RobotCfg, - LightCfg, RigidBodyAttributesCfg, RigidObjectCfg, URDFCfg, @@ -50,9 +50,18 @@ from embodichain.toolkits.graspkit.pg_grasp.gripper_collision_checker import ( GripperCollisionCfg, ) +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + replay_trajectory, + shutdown_sim, +) -def parse_arguments(): +def parse_arguments() -> argparse.Namespace: """ Parse command-line arguments to configure the simulation. @@ -62,7 +71,14 @@ def parse_arguments(): parser = argparse.ArgumentParser( description="Create and simulate a robot in SimulationManager" ) - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) + parser.add_argument( + "--n_sample", + "--n-sample", + type=int, + default=10000, + help="Number of antipodal grasp samples.", + ) return parser.parse_args() @@ -76,28 +92,17 @@ def initialize_simulation(args) -> SimulationManager: Returns: SimulationManager: Configured simulation manager instance. """ - config = SimulationManagerCfg( - headless=True, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), - physics_dt=1.0 / 100.0, + return create_default_sim( + args, + num_envs=args.num_envs, arena_space=2.5, ) - sim = SimulationManager(config) - light = sim.add_light( - cfg=LightCfg( - uid="main_light", - color=(0.6, 0.6, 0.6), - intensity=30.0, - init_pos=(1.0, 0, 3.0), - ) - ) - return sim - - -def create_robot(sim: SimulationManager, position=[0.0, 0.0, 0.0]) -> Robot: +def create_robot( + sim: SimulationManager, + position: tuple[float, float, float] = (0.0, 0.0, 0.0), +) -> Robot: """ Create and configure a robot with an arm and a dexterous hand in the simulation. @@ -146,9 +151,9 @@ def create_robot(sim: SimulationManager, position=[0.0, 0.0, 0.0]) -> Robot: return sim.add_robot(cfg=cfg) -def create_obj(sim: SimulationManager): +def create_obj(sim: SimulationManager) -> RigidObject: mug_cfg = RigidObjectCfg( - uid="table", + uid="mug", shape=MeshCfg( fpath=get_resources_data_path("Model", "BakeTexture", "hdr_color_mesh.ply"), ), @@ -162,11 +167,14 @@ def create_obj(sim: SimulationManager): init_pos=[0.55, 0.0, 0.08], init_rot=[0.0, 0.0, 0.0], ) - mug = sim.add_rigid_object(cfg=mug_cfg) - return mug + return sim.add_rigid_object(cfg=mug_cfg) -def get_grasp_traj(sim: SimulationManager, robot: Robot, grasp_xpos: torch.Tensor): +def get_grasp_traj( + sim: SimulationManager, + robot: Robot, + grasp_xpos: torch.Tensor, +) -> torch.Tensor: n_envs = sim.num_envs rest_arm_qpos = robot.get_qpos("arm") @@ -213,82 +221,84 @@ def get_grasp_traj(sim: SimulationManager, robot: Robot, grasp_xpos: torch.Tenso return interp_trajectory -if __name__ == "__main__": - import time - +def main() -> None: + """Plan and replay a mug grasp.""" args = parse_arguments() sim = initialize_simulation(args) - robot = create_robot(sim, position=[0.0, 0.0, 0.0]) - obj = create_obj(sim) - - # get mug grasp pose - grasp_cfg = GraspGeneratorCfg( - viser_port=11801, - antipodal_sampler_cfg=AntipodalSamplerCfg( - n_sample=10000, max_length=0.088, min_length=0.003 - ), - is_partial_annotate=False, - is_filter_ground_collision=True, - n_top_grasps=30, - ) - sim.open_window() - - # Annotate part of the mug to be grasped by following the instructions in the visualization window: - # 1. View grasp object in browser (e.g http://localhost:11801) - # 2. press 'Rect Select Region', select grasp region - # 3. press 'Confirm Selection' to finish grasp region selection. - - start_time = time.time() - - gripper_collision_cfg = GripperCollisionCfg( - max_open_length=0.088, finger_length=0.078, point_sample_dense=0.012 - ) - - # Extract mesh data from the mug and create grasp generator - vertices = obj.get_vertices(env_ids=[0], scale=True)[0] - triangles = obj.get_triangles(env_ids=[0])[0] - grasp_generator = GraspGenerator( - vertices=vertices, - triangles=triangles, - cfg=grasp_cfg, - gripper_collision_cfg=gripper_collision_cfg, - ) - - # Annotate grasp region (populates internal antipodal point pairs) - grasp_generator.annotate() + try: + robot = create_robot(sim) + obj = create_obj(sim) + maybe_open_window(sim, args) + + grasp_cfg = GraspGeneratorCfg( + viser_port=11801, + antipodal_sampler_cfg=AntipodalSamplerCfg( + n_sample=args.n_sample, + max_length=0.088, + min_length=0.003, + ), + is_partial_annotate=False, + is_filter_ground_collision=True, + n_top_grasps=30, + ) + started_at = time.perf_counter() + grasp_generator = GraspGenerator( + vertices=obj.get_vertices(env_ids=[0], scale=True)[0], + triangles=obj.get_triangles(env_ids=[0])[0], + cfg=grasp_cfg, + gripper_collision_cfg=GripperCollisionCfg( + max_open_length=0.088, + finger_length=0.078, + point_sample_dense=0.012, + ), + ) - # Compute grasp poses per environment - approach_direction = torch.tensor( - [0, 0, -1], dtype=torch.float32, device=sim.device - ) - obj_poses = obj.get_local_pose(to_matrix=True) - grasp_xpos_list = [] + # The first run opens Viser for selecting the mug's graspable region; + # later runs reuse the cached annotation. + grasp_generator.annotate() - rest_xpos = robot.compute_fk( - qpos=robot.get_qpos("arm"), name="arm", to_matrix=True - )[0] - for i, obj_pose in enumerate(obj_poses): - is_success, grasp_pose, open_length = grasp_generator.get_grasp_poses( - obj_pose, - approach_direction, - visualize_collision=False, - visualize_pose=True, + approach_direction = torch.tensor( + [0, 0, -1], + dtype=torch.float32, + device=sim.device, ) - if is_success: - grasp_xpos_list.append(grasp_pose.unsqueeze(0)) - else: - logger.log_warning(f"No valid grasp pose found for {i}-th object.") - grasp_xpos_list.append(rest_xpos.unsqueeze(0)) + rest_pose = robot.compute_fk( + qpos=robot.get_qpos("arm"), + name="arm", + to_matrix=True, + )[0] + grasp_poses = [] + for env_id, obj_pose in enumerate(obj.get_local_pose(to_matrix=True)): + success, grasp_pose, _ = grasp_generator.get_grasp_poses( + obj_pose, + approach_direction, + visualize_collision=False, + visualize_pose=True, + ) + if not success: + logger.log_warning( + f"No valid grasp pose found for environment {env_id}." + ) + grasp_pose = rest_pose + grasp_poses.append(grasp_pose.unsqueeze(0)) + + logger.log_info( + f"Grasp pose generation took {time.perf_counter() - started_at:.2f} seconds" + ) + trajectory = get_grasp_traj(sim, robot, torch.cat(grasp_poses, dim=0)) + maybe_wait_for_user(args, "Press Enter to start the mug grasp...") + with DemoRecording(sim, args, prefix="grasp_generator"): + replay_trajectory( + sim, + robot, + trajectory, + post_steps=0, + step_size=4, + ) + maybe_wait_for_user(args, "Press Enter to exit...") + finally: + shutdown_sim(sim) - grasp_xpos = torch.cat(grasp_xpos_list, dim=0) - cost_time = time.time() - start_time - logger.log_info(f"Get grasp pose cost time: {cost_time:.2f} seconds") - grab_traj = get_grasp_traj(sim, robot, grasp_xpos) - input("Press Enter to start the grab mug demo...") - n_waypoint = grab_traj.shape[1] - for i in range(n_waypoint): - robot.set_qpos(grab_traj[:, i, :]) - sim.update(step=4) - time.sleep(1e-2) - input("Press Enter to exit the simulation...") +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/gym/modular_env.py b/scripts/tutorials/gym/modular_env.py index dbe05ec26..699994476 100644 --- a/scripts/tutorials/gym/modular_env.py +++ b/scripts/tutorials/gym/modular_env.py @@ -14,9 +14,14 @@ # limitations under the License. # ---------------------------------------------------------------------------- -import torch +"""Configure and run a manager-based modular environment.""" + +from __future__ import annotations -from typing import List, Dict, Any +import argparse + +import gymnasium as gym +import torch import embodichain.lab.gym.envs.managers.randomization as rand import embodichain.lab.gym.envs.managers.events as events @@ -109,7 +114,7 @@ class ExampleCfg(EmbodiedEnvCfg): ) # Define the sensor configuration using StereoCameraCfg - sensor: List[SensorCfg] = [ + sensor: list[SensorCfg] = [ StereoCameraCfg( uid="eye_in_head", width=960, @@ -125,7 +130,7 @@ class ExampleCfg(EmbodiedEnvCfg): ) ] - background: List[RigidObjectCfg] = [ + background: list[RigidObjectCfg] = [ RigidObjectCfg( uid="table", shape=MeshCfg( @@ -144,7 +149,7 @@ class ExampleCfg(EmbodiedEnvCfg): ), ] - rigid_object: List[RigidObjectCfg] = [ + rigid_object: list[RigidObjectCfg] = [ RigidObjectCfg( uid="fork", shape=MeshCfg( @@ -155,7 +160,7 @@ class ExampleCfg(EmbodiedEnvCfg): ), ] - articulation_cfg: List[ArticulationCfg] = [ + articulation_cfg: list[ArticulationCfg] = [ ArticulationCfg( uid="drawer", fpath="SlidingBoxDrawer/SlidingBoxDrawer.urdf", @@ -175,14 +180,12 @@ class ModularEnv(EmbodiedEnv): and uses custom event and observation managers. """ - def __init__(self, cfg: EmbodiedEnvCfg, **kwargs): + def __init__(self, cfg: EmbodiedEnvCfg, **kwargs) -> None: super().__init__(cfg, **kwargs) -if __name__ == "__main__": - import gymnasium as gym - import argparse - +def main() -> None: + """Run five short episodes of the modular environment.""" from embodichain.lab.sim import SimulationManagerCfg from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser @@ -202,9 +205,20 @@ def __init__(self, cfg: EmbodiedEnvCfg, **kwargs): # Create the Gym environment env = gym.make("ModularEnv-v1", cfg=env_cfg) - for i in range(5): - obs, info = env.reset() + try: + device = env.get_wrapper_attr("device") + for _ in range(5): + env.reset() + for _ in range(100): + action = torch.zeros( + env.action_space.shape, + dtype=torch.float32, + device=device, + ) + env.step(action) + finally: + env.close() - for i in range(100): - action = torch.zeros(env.action_space.shape, dtype=torch.float32) - obs, reward, done, truncated, info = env.step(action) + +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/gym/random_reach.py b/scripts/tutorials/gym/random_reach.py index ddea8b777..f20f160d8 100644 --- a/scripts/tutorials/gym/random_reach.py +++ b/scripts/tutorials/gym/random_reach.py @@ -14,9 +14,16 @@ # limitations under the License. # ---------------------------------------------------------------------------- -import torch -import numpy as np +"""Minimal Gym environment with random joint-space reach commands.""" + +from __future__ import annotations + +import argparse +import time + import gymnasium as gym +import numpy as np +import torch from embodichain.lab.gym.envs import BaseEnv, EnvCfg from embodichain.lab.sim import SimulationManagerCfg @@ -34,6 +41,7 @@ @register_env("RandomReach-v1", override=True) class RandomReachEnv(BaseEnv): + """Environment that moves a UR10 toward random joint targets.""" robot_init_qpos = np.array( [1.57079, -1.57079, 1.57079, -1.57079, -1.57079, -3.14159] @@ -41,18 +49,21 @@ class RandomReachEnv(BaseEnv): def __init__( self, - num_envs=1, - headless=False, - device="cpu", - renderer="hybrid", + num_envs: int = 1, + headless: bool = False, + device: str = "cpu", + renderer: str = "hybrid", + gpu_id: int = 0, + arena_space: float = 2.0, **kwargs, - ): + ) -> None: env_cfg = EnvCfg( sim_cfg=SimulationManagerCfg( headless=headless, - arena_space=2.0, + arena_space=arena_space, sim_device=device, render_cfg=RenderCfg(renderer=renderer), + gpu_id=gpu_id, ), num_envs=num_envs, ) @@ -113,10 +124,8 @@ def _extend_obs(self, obs: EnvObs, **kwargs) -> EnvObs: return obs -if __name__ == "__main__": - import argparse - import time - +def main() -> None: + """Run a short random-reach throughput demo.""" from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser parser = argparse.ArgumentParser( @@ -131,50 +140,41 @@ def _extend_obs(self, obs: EnvObs, **kwargs) -> EnvObs: headless=args.headless, device=args.device, renderer=args.renderer, + gpu_id=args.gpu_id, + arena_space=args.arena_space, ) - for episode in range(10): - print("Episode:", episode) - env.reset() - start_time = time.time() - total_steps = 0 - - for i in range(100): - action = env.action_space.sample() - action = torch.as_tensor( - action, dtype=torch.float32, device=env.get_wrapper_attr("device") + try: + device = env.get_wrapper_attr("device") + num_envs = env.get_wrapper_attr("num_envs") + init_pose = ( + torch.as_tensor( + env.unwrapped.robot_init_qpos, + dtype=torch.float32, + device=device, ) + .unsqueeze(0) + .repeat(num_envs, 1) + ) - init_pose = env.unwrapped.robot_init_qpos - init_pose = ( - torch.as_tensor( - init_pose, - dtype=torch.float32, - device=env.get_wrapper_attr("device"), - ) - .unsqueeze_(0) - .repeat(env.get_wrapper_attr("num_envs"), 1) - ) - action = ( - init_pose - + torch.rand_like( - action, dtype=torch.float32, device=env.get_wrapper_attr("device") - ) - * 0.2 - - 0.1 - ) + for episode in range(10): + print("Episode:", episode) + env.reset() + started_at = time.perf_counter() + + for _ in range(100): + action = init_pose + torch.rand_like(init_pose) * 0.2 - 0.1 + env.step(action) - obs, reward, done, truncated, info = env.step(action) - total_steps += env.get_wrapper_attr("num_envs") + elapsed = time.perf_counter() - started_at + total_steps = 100 * num_envs + if elapsed > 0: + print(f"Total steps: {total_steps}") + print(f"Elapsed time: {elapsed:.2f} seconds") + print(f"FPS: {total_steps / elapsed:.2f}") + finally: + env.close() - end_time = time.time() - elapsed_time = end_time - start_time - if elapsed_time > 0: - fps = total_steps / elapsed_time - print(f"Total steps: {total_steps}") - print(f"Elapsed time: {elapsed_time:.2f} seconds") - print(f"FPS: {fps:.2f}") - else: - print("Elapsed time is too short to calculate FPS.") - env.close() +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/sim/create_cloth.py b/scripts/tutorials/sim/create_cloth.py index 1f0d883cc..c8dbf208d 100644 --- a/scripts/tutorials/sim/create_cloth.py +++ b/scripts/tutorials/sim/create_cloth.py @@ -19,24 +19,34 @@ It shows the basic setup of simulation context, adding objects, lighting, and sensors. """ +from __future__ import annotations + import argparse import os import tempfile -import time -import torch + import open3d as o3d -from dexsim.utility.path import get_resources_data_path -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +import torch + +from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.cfg import ( - RenderCfg, - RigidObjectCfg, RigidBodyAttributesCfg, ClothObjectCfg, ClothPhysicalAttributesCfg, + RigidObjectCfg, ) -from embodichain.lab.sim.shapes import MeshCfg, CubeCfg from embodichain.lab.sim.objects import ClothObject +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + resolve_demo_steps, + run_simulation_loop, + shutdown_sim, +) def create_2d_grid_mesh(width: float, height: float, nx: int = 1, ny: int = 1): @@ -80,23 +90,17 @@ def main(): parser = argparse.ArgumentParser( description="Create a simulation scene with SimulationManager" ) - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) + # Cloth simulation requires GPU physics. + parser.set_defaults(device="cuda") args = parser.parse_args() - # Configure the simulation - sim_cfg = SimulationManagerCfg( - width=1920, - height=1080, - headless=True, + sim = create_default_sim( + args, num_envs=args.num_envs, - physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device="cuda", # soft simulation only supports cuda device - render_cfg=RenderCfg(renderer=args.renderer), + add_default_light=False, ) - # Create the simulation instance - sim = SimulationManager(sim_cfg) - print("[INFO]: Scene setup complete!") cloth_verts, cloth_faces = create_2d_grid_mesh(width=0.3, height=0.3, nx=12, ny=12) @@ -142,61 +146,42 @@ def main(): init_pos=[0.5, 0.0, 0.04], init_rot=[0.0, 0.0, 0.0], ) - padding_box = sim.add_rigid_object(cfg=padding_box_cfg) + sim.add_rigid_object(cfg=padding_box_cfg) print("[INFO]: Add soft object complete!") # Open window when the scene has been set up - if not args.headless: - sim.open_window() + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) print(f"[INFO]: Running simulation with {args.num_envs} environment(s)") print("[INFO]: Press Ctrl+C to stop the simulation") # Run the simulation - run_simulation(sim, cloth) + run_simulation(sim, cloth, args) -def run_simulation(sim: SimulationManager, cloth: ClothObject) -> None: +def run_simulation( + sim: SimulationManager, + cloth: ClothObject, + args: argparse.Namespace, +) -> None: """Run the simulation loop. Args: sim: The SimulationManager instance to run - soft_obj: soft object + cloth: Cloth object to reset periodically. + args: Parsed demo arguments. """ - # Initialize GPU physics - sim.init_gpu_physics() - - step_count = 0 - try: - last_time = time.time() - last_step = 0 - while True: - # Update physics simulation - sim.update(step=1) - step_count += 1 - - # Print FPS every second - 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 - ) - print(f"[INFO]: Simulation step: {step_count}, FPS: {fps:.2f}") - last_time = current_time - last_step = step_count - if step_count % 500 == 0: - cloth.reset() - - except KeyboardInterrupt: - print("\n[INFO]: Stopping simulation...") + with DemoRecording(sim, args, prefix="create_cloth"): + run_simulation_loop( + sim, + max_steps=resolve_demo_steps(args), + on_step=lambda step: cloth.reset() if step % 500 == 0 else None, + ) finally: - # Clean up resources - sim.destroy() + shutdown_sim(sim) print("[INFO]: Simulation terminated successfully") diff --git a/scripts/tutorials/sim/create_rigid_constraint.py b/scripts/tutorials/sim/create_rigid_constraint.py index 5944dd597..752bbec83 100644 --- a/scripts/tutorials/sim/create_rigid_constraint.py +++ b/scripts/tutorials/sim/create_rigid_constraint.py @@ -23,16 +23,23 @@ import argparse import sys +import time -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.cfg import ( - RigidObjectCfg, - RigidConstraintCfg, RigidBodyAttributesCfg, - RenderCfg, + RigidConstraintCfg, + RigidObjectCfg, ) from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + shutdown_sim, +) # Number of physics sub-steps per update call. STEPS_PER_UPDATE = 1 @@ -49,7 +56,7 @@ def main(): parser = argparse.ArgumentParser( description="Attach and detach two cubes via a fixed rigid constraint" ) - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) args = parser.parse_args() # The simulation teardown (``SimulationManager.destroy``) calls ``os._exit``, @@ -57,20 +64,13 @@ def main(): # ``print`` below is visible even when the script is piped to a file. sys.stdout.reconfigure(line_buffering=True) - # Configure the simulation. - sim_cfg = SimulationManagerCfg( - width=1920, - height=1080, - headless=args.headless, - physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), + sim = create_default_sim( + args, num_envs=args.num_envs, arena_space=3.0, + add_default_light=False, ) - sim = SimulationManager(sim_cfg) - # Shared physics attributes for the two cubes. physics_attrs = RigidBodyAttributesCfg( mass=0.2, @@ -99,46 +99,45 @@ def main(): ) ) - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + maybe_init_gpu_physics(sim) print("[INFO]: Scene setup complete with two cubes (cube_a, cube_b).") - # --- Phase 1: attach the two cubes with a fixed constraint --------------- - # With default (None) local frames the constraint welds the cubes at their - # *current* relative pose: local_frame_a defaults to identity and - # local_frame_b is computed as inv(pose_B) @ pose_A, so the offset is - # preserved rather than the two origins being pulled together. - constraint = sim.create_rigid_constraint( - cfg=RigidConstraintCfg( - name="cube_weld", - rigid_object_a_uid="cube_a", - rigid_object_b_uid="cube_b", + try: + # With default local frames the constraint preserves the cubes' + # current relative pose instead of pulling their origins together. + sim.create_rigid_constraint( + cfg=RigidConstraintCfg( + name="cube_weld", + rigid_object_a_uid="cube_a", + rigid_object_b_uid="cube_b", + ) ) - ) - print("[INFO]: Created constraint 'cube_weld' between cube_a and cube_b.") - - # Open the viewer (unless --headless) so the welded motion is visible. - if not args.headless: - sim.open_window() + print("[INFO]: Created constraint 'cube_weld' between cube_a and cube_b.") - print("[INFO]: Stepping physics while ATTACHED (relative pose held):") - _run_phase(sim, cube_a, cube_b, attached=True) + maybe_open_window(sim, args) - # --- Phase 2: remove the constraint ------------------------------------ - sim.remove_rigid_constraint("cube_weld") - assert "cube_weld" not in sim.get_rigid_constraint_uid_list() - print("\n[INFO]: Removed constraint 'cube_weld'. cube_a and cube_b are now free.") + with DemoRecording(sim, args, prefix="create_rigid_constraint"): + print("[INFO]: Stepping physics while ATTACHED (relative pose held):") + _run_phase(sim, cube_a, cube_b) - import time + sim.remove_rigid_constraint("cube_weld") + assert "cube_weld" not in sim.get_rigid_constraint_uid_list() + print( + "\n[INFO]: Removed constraint 'cube_weld'. " + "cube_a and cube_b are now free." + ) - time.sleep(2.0) # Wait a moment so the viewer can show the constraint removal. + # Give an interactive viewer a moment to show the removal. + if not args.headless: + time.sleep(2.0) - print("[INFO]: Stepping physics while DETACHED (relative pose may drift):") - _run_phase(sim, cube_a, cube_b, attached=False) + print("[INFO]: Stepping physics while DETACHED (relative pose may drift):") + _run_phase(sim, cube_a, cube_b) - print("\n[INFO]: Tutorial complete.") - sim.destroy() + print("\n[INFO]: Tutorial complete.") + finally: + shutdown_sim(sim) def _relative_z(cube_a, cube_b) -> float: @@ -160,14 +159,13 @@ def _relative_z(cube_a, cube_b) -> float: return float(pose_b[0, 2, 3] - pose_a[0, 2, 3]) -def _run_phase(sim, cube_a, cube_b, attached: bool) -> None: +def _run_phase(sim, cube_a, cube_b) -> None: """Step the simulation for one phase and print the bodies' relative z. Args: sim: The :class:`SimulationManager`. cube_a: The first :class:`RigidObject`. cube_b: The second :class:`RigidObject`. - attached: True while the constraint is active, False after removal. """ rel_z = _relative_z(cube_a, cube_b) print(f" step {0:4d}: relative z (cube_b - cube_a) = {rel_z:.4f} m") diff --git a/scripts/tutorials/sim/create_rigid_object_group.py b/scripts/tutorials/sim/create_rigid_object_group.py index d681dc919..4275af5a2 100644 --- a/scripts/tutorials/sim/create_rigid_object_group.py +++ b/scripts/tutorials/sim/create_rigid_object_group.py @@ -18,18 +18,27 @@ This script demonstrates how to create a rigid object group using SimulationManager. """ +from __future__ import annotations + import argparse -import time -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg -from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.sim import SimulationManager +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg from embodichain.lab.sim.objects import ( - RigidObjectGroup, RigidObjectGroupCfg, RigidObjectCfg, ) +from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + resolve_demo_steps, + run_simulation_loop, + shutdown_sim, +) def main(): @@ -39,26 +48,16 @@ def main(): parser = argparse.ArgumentParser( description="Create a simulation scene with SimulationManager" ) - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) args = parser.parse_args() - # Configure the simulation - sim_cfg = SimulationManagerCfg( - width=1920, - height=1080, - headless=True, - physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device=args.device, - render_cfg=RenderCfg( - renderer=args.renderer - ), # Enable ray tracing for better visuals + sim = create_default_sim( + args, num_envs=args.num_envs, arena_space=3.0, + add_default_light=False, ) - # Create the simulation instance - sim = SimulationManager(sim_cfg) - physics_attrs = RigidBodyAttributesCfg( mass=1.0, dynamic_friction=0.5, @@ -67,7 +66,7 @@ def main(): ) # Add objects to the scene - obj_group: RigidObjectGroup = sim.add_rigid_object_group( + sim.add_rigid_object_group( cfg=RigidObjectGroupCfg( uid="obj_group", rigid_objects={ @@ -98,52 +97,25 @@ def main(): print("[INFO]: Press Ctrl+C to stop the simulation") # Open window when the scene has been set up - if not args.headless: - sim.open_window() + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) # Run the simulation - run_simulation(sim) + run_simulation(sim, args) -def run_simulation(sim: SimulationManager): +def run_simulation(sim: SimulationManager, args: argparse.Namespace) -> None: """Run the simulation loop. Args: sim: The SimulationManager instance to run """ - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - - step_count = 0 - try: - last_time = time.time() - last_step = 0 - while True: - # Update physics simulation - sim.update(step=1) - step_count += 1 - - # Print FPS every second - 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 - ) - print(f"[INFO]: Simulation step: {step_count}, FPS: {fps:.2f}") - last_time = current_time - last_step = step_count - - except KeyboardInterrupt: - print("\n[INFO]: Stopping simulation...") + with DemoRecording(sim, args, prefix="create_rigid_object_group"): + run_simulation_loop(sim, max_steps=resolve_demo_steps(args)) finally: - # Clean up resources - sim.destroy() + shutdown_sim(sim) print("[INFO]: Simulation terminated successfully") diff --git a/scripts/tutorials/sim/create_robot.py b/scripts/tutorials/sim/create_robot.py index 6a0b8ee59..e368aae67 100644 --- a/scripts/tutorials/sim/create_robot.py +++ b/scripts/tutorials/sim/create_robot.py @@ -19,25 +19,33 @@ It shows how to load a robot from URDF, set up control parts, and run basic simulation. """ +from __future__ import annotations + import argparse import numpy as np -import time import torch -torch.set_printoptions(precision=4, sci_mode=False) - from scipy.spatial.transform import Rotation as R -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.cfg import ( - RenderCfg, JointDrivePropertiesCfg, RobotCfg, URDFCfg, ) from embodichain.data import get_data_path -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + resolve_demo_steps, + run_simulation_loop, + setup_print_options, + shutdown_sim, +) ACTION_SWITCH_INTERVAL = 100 ACTION_CYCLE_STEPS = 4 * ACTION_SWITCH_INTERVAL @@ -50,34 +58,30 @@ def main(): parser = argparse.ArgumentParser( description="Create and simulate a robot in SimulationManager" ) - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) args = parser.parse_args() + setup_print_options() # Initialize simulation print("Creating simulation...") - config = SimulationManagerCfg( - headless=True, - sim_device=args.device, + sim = create_default_sim( + args, arena_space=3.0, - render_cfg=RenderCfg(renderer=args.renderer), - physics_dt=1.0 / 100.0, num_envs=args.num_envs, + add_default_light=False, ) - sim = SimulationManager(config) # Create robot configuration robot = create_robot(sim) # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + maybe_init_gpu_physics(sim) # Open visualization window if not headless - if not args.headless: - sim.open_window() + maybe_open_window(sim, args) # Run simulation loop - run_simulation(sim, robot) + run_simulation(sim, robot, args) def create_robot(sim): @@ -137,15 +141,17 @@ def create_robot(sim): return robot -def run_simulation(sim: SimulationManager, robot: Robot): +def run_simulation( + sim: SimulationManager, + robot: Robot, + args: argparse.Namespace, +) -> None: """Run the simulation loop with robot control.""" print("Starting simulation...") print("Robot will move through different poses") print("Press Ctrl+C to stop") - step_count = 0 - arm_joint_ids = robot.get_joint_ids("arm") # Define some target joint positions for demonstration arm_position1 = ( @@ -170,35 +176,36 @@ def run_simulation(sim: SimulationManager, robot: Robot): hand_position_open = robot.body_data.qpos_limits[:, hand_joint_ids, 1] hand_position_close = robot.body_data.qpos_limits[:, hand_joint_ids, 0] - try: - while True: - # Update physics - sim.update(step=1) - cycle_step = step_count % ACTION_CYCLE_STEPS - - if cycle_step == 0: - robot.set_qpos(qpos=arm_position1, joint_ids=arm_joint_ids) - print(f"Moving to arm position 1") + def update_target(step: int) -> None: + """Switch arm and hand targets at fixed simulation intervals.""" + cycle_step = (step - 1) % ACTION_CYCLE_STEPS - if cycle_step == ACTION_SWITCH_INTERVAL: - robot.set_qpos(qpos=arm_position2, joint_ids=arm_joint_ids) - print(f"Moving to arm position 2") + if cycle_step == 0: + robot.set_qpos(qpos=arm_position1, joint_ids=arm_joint_ids) + print("Moving to arm position 1") - if cycle_step == 2 * ACTION_SWITCH_INTERVAL: - robot.set_qpos(qpos=hand_position_close, joint_ids=hand_joint_ids) - print(f"Closing hand") + if cycle_step == ACTION_SWITCH_INTERVAL: + robot.set_qpos(qpos=arm_position2, joint_ids=arm_joint_ids) + print("Moving to arm position 2") - if cycle_step == 3 * ACTION_SWITCH_INTERVAL: - robot.set_qpos(qpos=hand_position_open, joint_ids=hand_joint_ids) - print(f"Opening hand") + if cycle_step == 2 * ACTION_SWITCH_INTERVAL: + robot.set_qpos(qpos=hand_position_close, joint_ids=hand_joint_ids) + print("Closing hand") - step_count += 1 + if cycle_step == 3 * ACTION_SWITCH_INTERVAL: + robot.set_qpos(qpos=hand_position_open, joint_ids=hand_joint_ids) + print("Opening hand") - except KeyboardInterrupt: - print("Stopping simulation...") + try: + with DemoRecording(sim, args, prefix="create_robot"): + run_simulation_loop( + sim, + max_steps=resolve_demo_steps(args), + on_step=update_target, + ) finally: print("Cleaning up...") - sim.destroy() + shutdown_sim(sim) if __name__ == "__main__": diff --git a/scripts/tutorials/sim/create_scene.py b/scripts/tutorials/sim/create_scene.py index b6c1df4b1..f5a8400fc 100644 --- a/scripts/tutorials/sim/create_scene.py +++ b/scripts/tutorials/sim/create_scene.py @@ -22,14 +22,24 @@ """ import argparse -import time -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg +from embodichain.lab.sim import SimulationManager +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg from embodichain.lab.sim.shapes import CubeCfg, MeshCfg from embodichain.lab.sim.objects import RigidObject, RigidObjectCfg -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.data import get_data_path +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + run_simulation_loop, + shutdown_sim, +) + +DEFAULT_HEADLESS_STEPS = 1000 +RECORD_LOOK_AT = ((2.6, -2.2, 1.6), (0.0, 0.0, 0.45), (0.0, 0.0, 1.0)) def main(): @@ -39,46 +49,21 @@ def main(): parser = argparse.ArgumentParser( description="Create a simulation scene with SimulationManager" ) - add_env_launcher_args_to_parser(parser) - parser.add_argument( - "--record-steps", - type=int, - default=1000, - help="Number of simulation steps to record before exiting in headless mode.", - ) - parser.add_argument( - "--record-fps", - type=int, - default=20, - help="Output video FPS for headless recording.", - ) - parser.add_argument( - "--record-save-path", - type=str, - default=None, - help="Optional mp4 output path for headless recording.", - ) + add_demo_args(parser) + parser.set_defaults(record_fps=20) args = parser.parse_args() + if args.headless and args.record_steps is None: + args.record_steps = DEFAULT_HEADLESS_STEPS - # Configure the simulation - sim_cfg = SimulationManagerCfg( - width=1920, - height=1080, - headless=True, - physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device=args.device, - render_cfg=RenderCfg( - renderer=args.renderer, - ), + sim = create_default_sim( + args, num_envs=args.num_envs, arena_space=3.0, + add_default_light=False, ) - # Create the simulation instance - sim = SimulationManager(sim_cfg) - # Add cube object to the scene - cube: RigidObject = sim.add_rigid_object( + sim.add_rigid_object( cfg=RigidObjectCfg( uid="cube", shape=CubeCfg(size=[0.1, 0.1, 0.1]), @@ -95,7 +80,7 @@ def main(): # Add chair object to the scene path = get_data_path("Chair/chair.glb") - chair: RigidObject = sim.add_rigid_object( + sim.add_rigid_object( cfg=RigidObjectCfg( uid="chair", shape=MeshCfg(fpath=path), @@ -114,84 +99,33 @@ def main(): print("[INFO]: Press Ctrl+C to stop the simulation") # Open window when the scene has been set up - if not args.headless: - sim.open_window() - - if args.headless: - if not sim.start_window_record( - save_path=args.record_save_path, - fps=args.record_fps, - max_memory=2048, - video_prefix="create_scene_headless", - look_at=((2.6, -2.2, 1.6), (0.0, 0.0, 0.45), (0.0, 0.0, 1.0)), - ): - raise RuntimeError("Failed to start headless recording") - print("[INFO]: Headless recording enabled.") - print( - "[INFO]: The output path is reported by `SimulationManager.start_window_record()`." - ) - print(f"[INFO]: Running {args.record_steps} steps before exporting the video") + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) - # Run the simulation - run_simulation( - sim, - max_steps=args.record_steps if args.headless else None, - ) + run_simulation(sim, args) def run_simulation( sim: SimulationManager, - max_steps: int | None = None, -): + args: argparse.Namespace, +) -> None: """Run the simulation loop. Args: sim: The SimulationManager instance to run. - max_steps: Optional maximum number of simulation steps to execute. + args: Parsed demo arguments. """ - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - - step_count = 0 - try: - last_time = time.time() - last_step = 0 - while True: - # Update physics simulation - sim.update(step=1) - step_count += 1 - - # Print FPS every second - 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 - ) - print(f"[INFO]: Simulation step: {step_count}, FPS: {fps:.2f}") - last_time = current_time - last_step = step_count - - if max_steps is not None and step_count >= max_steps: - print( - f"[INFO]: Reached {max_steps} steps. Exporting headless recording..." - ) - break - - except KeyboardInterrupt: - print("\n[INFO]: Stopping simulation...") + with DemoRecording( + sim, + args, + prefix="create_scene", + look_at=RECORD_LOOK_AT, + ): + run_simulation_loop(sim, max_steps=args.record_steps) finally: - if sim.is_window_recording(): - sim.stop_window_record() - sim.wait_window_record_saves() - - # Clean up resources - sim.destroy() + shutdown_sim(sim) if __name__ == "__main__": diff --git a/scripts/tutorials/sim/create_sensor.py b/scripts/tutorials/sim/create_sensor.py index 5600c2cea..fb67fe181 100644 --- a/scripts/tutorials/sim/create_sensor.py +++ b/scripts/tutorials/sim/create_sensor.py @@ -19,21 +19,20 @@ It shows how to configure a camera sensor, attach it to the robot's end-effector, and visualize the sensor's output during simulation. """ +from __future__ import annotations + import argparse + +import cv2 import numpy as np import torch -import cv2 - -torch.set_printoptions(precision=4, sci_mode=False) from scipy.spatial.transform import Rotation as R -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.sensors import Camera, CameraCfg from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.cfg import ( - RenderCfg, JointDrivePropertiesCfg, RobotCfg, URDFCfg, @@ -41,12 +40,23 @@ ) from embodichain.lab.sim.shapes import CubeCfg from embodichain.data import get_data_path +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + resolve_demo_steps, + run_simulation_loop, + setup_print_options, + shutdown_sim, +) ACTION_SWITCH_INTERVAL = 100 ACTION_CYCLE_STEPS = 2 * ACTION_SWITCH_INTERVAL -def mask_to_color_map(mask, user_ids, fix_seed=True): +def mask_to_color_map(mask, user_ids): """ Convert instance mask into color map. :param mask: Instance mask map. @@ -78,25 +88,24 @@ def main(): parser = argparse.ArgumentParser( description="Create and simulate a robot in SimulationManager" ) - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) parser.add_argument( "--attach_sensor", + "--attach-sensor", action="store_true", help="Attach sensor to robot end-effector", ) args = parser.parse_args() + setup_print_options() # Initialize simulation print("Creating simulation...") - config = SimulationManagerCfg( - headless=True, - sim_device=args.device, + sim = create_default_sim( + args, arena_space=3.0, - render_cfg=RenderCfg(renderer=args.renderer), - physics_dt=1.0 / 100.0, num_envs=args.num_envs, + add_default_light=False, ) - sim = SimulationManager(config) # Create robot configuration robot = create_robot(sim) @@ -113,15 +122,13 @@ def main(): sim.add_rigid_object(cfg=cube_cfg) # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + maybe_init_gpu_physics(sim) # Open visualization window if not headless - if not args.headless: - sim.open_window() + maybe_open_window(sim, args) # Run simulation loop - run_simulation(sim, robot, sensor) + run_simulation(sim, robot, sensor, args) def create_sensor(sim: SimulationManager, args): @@ -132,7 +139,7 @@ def create_sensor(sim: SimulationManager, args): # extrinsics params pos = [0.09, 0.05, 0.04] - quat = R.from_euler("xyz", [-35, 135, 0], degrees=True).as_quat().tolist() + quat_xyzw = R.from_euler("xyz", [-35, 135, 0], degrees=True).as_quat() # If attach_sensor is True, attach to robot end-effector; otherwise, place it in the scene if args.attach_sensor: @@ -140,8 +147,10 @@ def create_sensor(sim: SimulationManager, args): else: parent = None pos = [1.2, -0.2, 1.5] - quat = R.from_euler("xyz", [0, 180, 0], degrees=True).as_quat().tolist() - quat = [quat[3], quat[0], quat[1], quat[2]] # Convert to (w, x, y, z) + quat_xyzw = R.from_euler("xyz", [0, 180, 0], degrees=True).as_quat() + + # CameraCfg uses (w, x, y, z), while SciPy returns (x, y, z, w). + quat = [quat_xyzw[3], *quat_xyzw[:3]] # create camera sensor and attach to robot end-effector camera: Camera = sim.add_sensor( @@ -271,15 +280,18 @@ def get_sensor_image(camera: Camera, headless=False, step_count=0): plt.close(fig) -def run_simulation(sim: SimulationManager, robot: Robot, camera: Camera): +def run_simulation( + sim: SimulationManager, + robot: Robot, + camera: Camera, + args: argparse.Namespace, +) -> None: """Run the simulation loop with robot and camera sensor control.""" print("Starting simulation...") print("Robot will move through different poses") print("Press Ctrl+C to stop") - step_count = 0 - arm_joint_ids = robot.get_joint_ids("arm") # Define some target joint positions for demonstration @@ -299,33 +311,32 @@ def run_simulation(sim: SimulationManager, robot: Robot, camera: Camera): .repeat(sim.num_envs, 1) ) - try: - while True: - # Update physics - sim.update(step=1) - cycle_step = step_count % ACTION_CYCLE_STEPS + def update_target(step: int) -> None: + """Move the robot and capture images at each target switch.""" + cycle_step = (step - 1) % ACTION_CYCLE_STEPS - if cycle_step == 0: - robot.set_qpos(qpos=arm_position1, joint_ids=arm_joint_ids) - print(f"Moving to arm position 1") + if cycle_step == 0: + robot.set_qpos(qpos=arm_position1, joint_ids=arm_joint_ids) + print("Moving to arm position 1") - # Refresh and get image from sensor - get_sensor_image(camera) + get_sensor_image(camera, headless=args.headless, step_count=step) - if cycle_step == ACTION_SWITCH_INTERVAL: - robot.set_qpos(qpos=arm_position2, joint_ids=arm_joint_ids) - print(f"Moving to arm position 2") + if cycle_step == ACTION_SWITCH_INTERVAL: + robot.set_qpos(qpos=arm_position2, joint_ids=arm_joint_ids) + print("Moving to arm position 2") - # Refresh and get image from sensor - get_sensor_image(camera) + get_sensor_image(camera, headless=args.headless, step_count=step) - step_count += 1 - - except KeyboardInterrupt: - print("Stopping simulation...") + try: + with DemoRecording(sim, args, prefix="create_sensor"): + run_simulation_loop( + sim, + max_steps=resolve_demo_steps(args), + on_step=update_target, + ) finally: print("Cleaning up...") - sim.destroy() + shutdown_sim(sim) if __name__ == "__main__": diff --git a/scripts/tutorials/sim/create_softbody.py b/scripts/tutorials/sim/create_softbody.py index 3b8973ef7..413f5da32 100644 --- a/scripts/tutorials/sim/create_softbody.py +++ b/scripts/tutorials/sim/create_softbody.py @@ -19,21 +19,32 @@ It shows the basic setup of simulation context, adding objects, lighting, and sensors. """ +from __future__ import annotations + import argparse -import time + from dexsim.utility.path import get_resources_data_path -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser + +from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.cfg import ( - RenderCfg, - SoftbodyVoxelAttributesCfg, SoftbodyPhysicalAttributesCfg, + SoftbodyVoxelAttributesCfg, ) -from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim.objects import ( SoftObject, SoftObjectCfg, ) +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + resolve_demo_steps, + run_simulation_loop, + shutdown_sim, +) def main(): @@ -43,25 +54,17 @@ def main(): parser = argparse.ArgumentParser( description="Create a simulation scene with SimulationManager" ) - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) + # Soft-body simulation requires GPU physics. + parser.set_defaults(device="cuda") args = parser.parse_args() - # Configure the simulation - sim_cfg = SimulationManagerCfg( - width=1920, - height=1080, - headless=True, + sim = create_default_sim( + args, num_envs=args.num_envs, - physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device="cuda", # soft simulation only supports cuda device - render_cfg=RenderCfg( - renderer=args.renderer - ), # Enable ray tracing for better visuals + add_default_light=False, ) - # Create the simulation instance - sim = SimulationManager(sim_cfg) - print("[INFO]: Scene setup complete!") # add softbody to the scene @@ -88,17 +91,21 @@ def main(): print("[INFO]: Add soft object complete!") # Open window when the scene has been set up - if not args.headless: - sim.open_window() + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) print(f"[INFO]: Running simulation with {args.num_envs} environment(s)") print("[INFO]: Press Ctrl+C to stop the simulation") # Run the simulation - run_simulation(sim, cow) + run_simulation(sim, cow, args) -def run_simulation(sim: SimulationManager, soft_obj: SoftObject) -> None: +def run_simulation( + sim: SimulationManager, + soft_obj: SoftObject, + args: argparse.Namespace, +) -> None: """Run the simulation loop. Args: @@ -106,39 +113,15 @@ def run_simulation(sim: SimulationManager, soft_obj: SoftObject) -> None: soft_obj: soft object """ - # Initialize GPU physics - sim.init_gpu_physics() - - step_count = 0 - try: - last_time = time.time() - last_step = 0 - while True: - # Update physics simulation - sim.update(step=1) - step_count += 1 - - # Print FPS every second - 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 - ) - print(f"[INFO]: Simulation step: {step_count}, FPS: {fps:.2f}") - last_time = current_time - last_step = step_count - if step_count % 500 == 0: - soft_obj.reset() - - except KeyboardInterrupt: - print("\n[INFO]: Stopping simulation...") + with DemoRecording(sim, args, prefix="create_softbody"): + run_simulation_loop( + sim, + max_steps=resolve_demo_steps(args), + on_step=lambda step: soft_obj.reset() if step % 500 == 0 else None, + ) finally: - # Clean up resources - sim.destroy() + shutdown_sim(sim) print("[INFO]: Simulation terminated successfully") diff --git a/scripts/tutorials/sim/export_usd.py b/scripts/tutorials/sim/export_usd.py index c6cb91c74..00f742932 100644 --- a/scripts/tutorials/sim/export_usd.py +++ b/scripts/tutorials/sim/export_usd.py @@ -18,21 +18,30 @@ This script demonstrates how to export a simulation scene to a usd file using the SimulationManager. """ +from __future__ import annotations + import argparse + import numpy as np -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.objects import Robot, RigidObject + +from embodichain.lab.sim import SimulationManager +from embodichain.lab.sim.objects import Articulation, Robot, RigidObject from embodichain.lab.sim.cfg import ( - RenderCfg, - LightCfg, + ArticulationCfg, JointDrivePropertiesCfg, - RigidObjectCfg, RigidBodyAttributesCfg, - ArticulationCfg, + RigidObjectCfg, ) from embodichain.lab.sim.shapes import MeshCfg from embodichain.data import get_data_path +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + resolve_demo_steps, + run_simulation_loop, + shutdown_sim, +) from embodichain.utils import logger from embodichain.lab.sim.robots.dexforce_w1.cfg import DexforceW1Cfg @@ -48,7 +57,7 @@ def parse_arguments(): parser = argparse.ArgumentParser( description="Create and simulate a robot in SimulationManager" ) - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) return parser.parse_args() @@ -62,26 +71,11 @@ def initialize_simulation(args) -> SimulationManager: Returns: SimulationManager: Configured simulation manager instance. """ - config = SimulationManagerCfg( - headless=True, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), - physics_dt=1.0 / 100.0, + return create_default_sim( + args, num_envs=1, arena_space=2.5, ) - sim = SimulationManager(config) - - light = sim.add_light( - cfg=LightCfg( - uid="main_light", - color=(0.6, 0.6, 0.6), - intensity=30.0, - init_pos=(1.0, 0, 3.0), - ) - ) - - return sim def create_robot(sim: SimulationManager) -> Robot: @@ -189,7 +183,7 @@ def create_table(sim: SimulationManager) -> RigidObject: return scoop -def create_caffe(sim: SimulationManager) -> Robot: +def create_caffe(sim: SimulationManager) -> Articulation: """ Create a caffe (container) articulated object in the simulation. @@ -197,7 +191,7 @@ def create_caffe(sim: SimulationManager) -> Robot: sim (SimulationManager): The simulation manager instance. Returns: - Robot: The caffe object added to the simulation. + Articulation: The caffe object added to the simulation. """ container_cfg = ArticulationCfg( uid="caffe", @@ -212,8 +206,7 @@ def create_caffe(sim: SimulationManager) -> Robot: ), ) print(f"Loading URDF file from: {container_cfg.fpath}") - container = sim.add_articulation(cfg=container_cfg) - return container + return sim.add_articulation(cfg=container_cfg) def create_cup(sim: SimulationManager) -> RigidObject: @@ -252,23 +245,21 @@ def main(): args = parse_arguments() sim = initialize_simulation(args) - robot = create_robot(sim) - table = create_table(sim) - caffe = create_caffe(sim) - cup = create_cup(sim) - - sim.export_usd("w1_coffee_scene.usda") - - logger.log_info("Scene exported successfully.") - - if not args.headless: - sim.open_window() - logger.log_info("Press Ctrl+C to exit.") - try: - while True: - sim.update(step=1) - except KeyboardInterrupt: - logger.log_info("\nExit") + try: + create_robot(sim) + create_table(sim) + create_caffe(sim) + create_cup(sim) + + sim.export_usd("w1_coffee_scene.usda") + logger.log_info("Scene exported successfully.") + + maybe_open_window(sim, args) + if not args.headless or args.auto_play: + logger.log_info("Press Ctrl+C to exit.") + run_simulation_loop(sim, max_steps=resolve_demo_steps(args)) + finally: + shutdown_sim(sim) if __name__ == "__main__": diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index 6d6613f9a..847578a2a 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -17,22 +17,30 @@ Gizmo-Robot Example: Test Gizmo class on a robot (UR10) """ +from __future__ import annotations + +import argparse import time -import torch + import numpy as np -import argparse +import torch -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.cfg import ( - RenderCfg, + JointDrivePropertiesCfg, RobotCfg, URDFCfg, - JointDrivePropertiesCfg, ) from embodichain.lab.sim.solvers import PinkSolverCfg from embodichain.data import get_data_path +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + resolve_demo_steps, + shutdown_sim, +) from embodichain.utils import logger @@ -43,19 +51,16 @@ def main(): parser = argparse.ArgumentParser( description="Create a simulation scene with SimulationManager" ) - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) args = parser.parse_args() - # Configure the simulation - sim_cfg = SimulationManagerCfg( + sim = create_default_sim( + args, width=1920, height=1080, physics_dt=1.0 / 100.0, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), + add_default_light=False, ) - - sim = SimulationManager(sim_cfg) sim.set_manual_update(False) # Get UR10 URDF path @@ -101,23 +106,24 @@ def main(): sim.enable_gizmo(uid="ur10_gizmo_test", control_part="arm") if not sim.has_gizmo("ur10_gizmo_test", control_part="arm"): logger.log_error("Failed to enable gizmo!") + shutdown_sim(sim) return - sim.open_window() + maybe_open_window(sim, args) logger.log_info("Gizmo-Robot example started!") logger.log_info("Use the gizmo to drag the robot end-effector (EE)") logger.log_info("Press Ctrl+C to stop the simulation") - run_simulation(sim) + run_simulation(sim, max_steps=resolve_demo_steps(args)) -def run_simulation(sim: SimulationManager): +def run_simulation(sim: SimulationManager, max_steps: int | None = None) -> None: step_count = 0 try: last_time = time.time() last_step = 0 - while True: + while max_steps is None or step_count < max_steps: time.sleep(0.033) # 30Hz # Update all gizmos managed by sim sim.update_gizmos() @@ -137,7 +143,7 @@ def run_simulation(sim: SimulationManager): except KeyboardInterrupt: logger.log_info("\nStopping simulation...") finally: - sim.destroy() + shutdown_sim(sim) logger.log_info("Simulation terminated successfully") diff --git a/scripts/tutorials/sim/import_usd.py b/scripts/tutorials/sim/import_usd.py index ada74edf9..dc7f0253e 100644 --- a/scripts/tutorials/sim/import_usd.py +++ b/scripts/tutorials/sim/import_usd.py @@ -1,5 +1,5 @@ # ---------------------------------------------------------------------------- -# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# 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. @@ -20,12 +20,12 @@ Multiple arenas are not supported when importing USD files. """ +from __future__ import annotations + import argparse -import time -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg +from embodichain.lab.sim import SimulationManager +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg from embodichain.lab.sim.shapes import CubeCfg, MeshCfg from embodichain.lab.sim.objects import ( RigidObject, @@ -34,6 +34,16 @@ Robot, ) from embodichain.data import get_data_path +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + resolve_demo_steps, + run_simulation_loop, + shutdown_sim, +) def main(): @@ -43,26 +53,16 @@ def main(): parser = argparse.ArgumentParser( description="Create a simulation scene with SimulationManager" ) - add_env_launcher_args_to_parser(parser) + add_demo_args(parser) args = parser.parse_args() - # Configure the simulation - sim_cfg = SimulationManagerCfg( - width=1920, - height=1080, - headless=True, - physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device=args.device, - render_cfg=RenderCfg( - renderer=args.renderer, - ), # Enable ray tracing for better visuals + sim = create_default_sim( + args, num_envs=1, arena_space=3.0, + add_default_light=False, ) - # Create the simulation instance - sim = SimulationManager(sim_cfg) - cube: RigidObject = sim.add_rigid_object( cfg=RigidObjectCfg( uid="cube", @@ -104,56 +104,32 @@ def main(): ) # Open window when the scene has been set up - if not args.headless: - sim.open_window() + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) print("[INFO]: Scene setup complete!") print("[INFO]: Press Ctrl+C to stop the simulation") # Run the simulation - run_simulation(sim) + run_simulation(sim, args) -def run_simulation(sim: SimulationManager): +def run_simulation(sim: SimulationManager, args: argparse.Namespace) -> None: """Run the simulation loop. Args: sim: The SimulationManager instance to run """ - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - - step_count = 0 - try: - last_time = time.time() - last_step = 0 - while True: - # Update physics simulation - sim.update(step=1) - time.sleep(0.03) # Sleep to limit update rate (optional) - step_count += 1 - - # Print FPS every second - 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 - ) - # print(f"[INFO]: Simulation step: {step_count}, FPS: {fps:.2f}") - last_time = current_time - last_step = step_count - - except KeyboardInterrupt: - print("\n[INFO]: Stopping simulation...") + with DemoRecording(sim, args, prefix="import_usd"): + run_simulation_loop( + sim, + max_steps=resolve_demo_steps(args), + sleep=0.03, + ) finally: - # Clean up resources - sim.destroy() + shutdown_sim(sim) print("[INFO]: Simulation terminated successfully") diff --git a/scripts/tutorials/sim/motion_generator.py b/scripts/tutorials/sim/motion_generator.py index 7d1988103..0050307fe 100644 --- a/scripts/tutorials/sim/motion_generator.py +++ b/scripts/tutorials/sim/motion_generator.py @@ -17,15 +17,12 @@ from __future__ import annotations import argparse -import time from collections.abc import Sequence import numpy as np import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RenderCfg +from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.planners import ( MotionGenCfg, @@ -37,12 +34,20 @@ ) from embodichain.lab.sim.planners.utils import TrajectorySampleMethod from embodichain.lab.sim.robots import CobotMagicCfg +from embodichain.lab.sim.utility.demo_utils import ( + DemoRecording, + add_demo_args, + create_default_sim, + maybe_init_gpu_physics, + maybe_open_window, + setup_print_options, + shutdown_sim, +) RECORD_WIDTH = 1920 RECORD_HEIGHT = 1080 DEFAULT_ARENA_SPACE = 3.0 DEFAULT_RECORD_TARGET_Z = 0.95 -DEFAULT_RECORD_MAX_MEMORY = 2048 def parse_args() -> argparse.Namespace: @@ -50,35 +55,10 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Generate and replay MotionGenerator trajectories for one or more environments." ) - add_env_launcher_args_to_parser(parser) - parser.add_argument( - "--arena-space", - type=float, - default=DEFAULT_ARENA_SPACE, - help="Spacing between replicated tutorial environments.", - ) - parser.add_argument( - "--step-delay", - type=float, - default=0.1, - help="Seconds to wait between trajectory waypoints during playback.", - ) - parser.add_argument( - "--record-fps", - type=int, - default=20, - help="Output video FPS for headless recording.", - ) - parser.add_argument( - "--record-save-path", - type=str, - default=None, - help="Optional mp4 output path for headless recording.", - ) - parser.add_argument( - "--disable-record", - action="store_true", - help="Disable automatic whole-scene recording in headless mode.", + add_demo_args(parser) + parser.set_defaults( + arena_space=DEFAULT_ARENA_SPACE, + record_fps=20, ) return parser.parse_args() @@ -183,133 +163,101 @@ def create_demo_trajectory( return [qpos_begin, qpos_mid, qpos_final], [xpos_begin, xpos_mid, xpos_final] -def start_headless_recording( - sim: SimulationManager, - args: argparse.Namespace, -) -> bool: - """Start headless viewer recording with a whole-scene camera.""" - if not args.headless or args.disable_record: - return False - - look_at = compute_record_look_at( - num_envs=sim.num_envs, - arena_space=sim.sim_config.arena_space, - ) - if not sim.start_window_record( - save_path=args.record_save_path, - fps=args.record_fps, - max_memory=DEFAULT_RECORD_MAX_MEMORY, - video_prefix="motion_generator_headless", - look_at=look_at, - use_sim_time=False, - ): - raise RuntimeError("Failed to start headless recording") - - print("[INFO]: Headless recording enabled.") - print( - "[INFO]: The output path is reported by `SimulationManager.start_window_record()`." - ) - return True - - def main() -> None: """Run the motion-generator tutorial.""" args = parse_args() - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) - - sim = SimulationManager( - SimulationManagerCfg( - width=RECORD_WIDTH, - height=RECORD_HEIGHT, - headless=True, - physics_dt=1.0 / 100.0, - sim_device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), - num_envs=args.num_envs, - arena_space=args.arena_space, - ) - ) - - robot: Robot = sim.add_robot(cfg=CobotMagicCfg.from_dict({"uid": "CobotMagic"})) - arm_name = "left_arm" - - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + setup_print_options() - if not args.headless: - sim.open_window() - - print( - f"[INFO]: Running motion generator tutorial with {sim.num_envs} environment(s)" + sim = create_default_sim( + args, + width=RECORD_WIDTH, + height=RECORD_HEIGHT, + num_envs=args.num_envs, + arena_space=args.arena_space, + add_default_light=False, ) - - recording_started = start_headless_recording(sim, args) try: - qpos_list, xpos_list = create_demo_trajectory( - robot=robot, - arm_name=arm_name, + robot: Robot = sim.add_robot(cfg=CobotMagicCfg.from_dict({"uid": "CobotMagic"})) + arm_name = "left_arm" + maybe_init_gpu_physics(sim) + maybe_open_window(sim, args) + print( + f"[INFO]: Running motion generator tutorial with {sim.num_envs} environment(s)" + ) + + look_at = compute_record_look_at( num_envs=sim.num_envs, + arena_space=sim.sim_config.arena_space, ) + with DemoRecording( + sim, + args, + prefix="motion_generator", + look_at=look_at, + ): + qpos_list, xpos_list = create_demo_trajectory( + robot=robot, + arm_name=arm_name, + num_envs=sim.num_envs, + ) - motion_generator = MotionGenerator( - cfg=MotionGenCfg( - planner_cfg=ToppraPlannerCfg( - robot_uid=robot.uid, + motion_generator = MotionGenerator( + cfg=MotionGenCfg( + planner_cfg=ToppraPlannerCfg( + robot_uid=robot.uid, + ) ) ) - ) - options = MotionGenOptions( - control_part=arm_name, - start_qpos=qpos_list[0], - is_interpolate=True, - is_linear=False, - plan_opts=ToppraPlanOptions( - constraints={ - "velocity": 0.2, - "acceleration": 0.5, - }, - sample_method=TrajectorySampleMethod.QUANTITY, - sample_interval=20, - ), - ) + options = MotionGenOptions( + control_part=arm_name, + start_qpos=qpos_list[0], + is_interpolate=True, + is_linear=False, + plan_opts=ToppraPlanOptions( + constraints={ + "velocity": 0.2, + "acceleration": 0.5, + }, + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=20, + ), + ) - joint_plan = motion_generator.generate( - target_states=[PlanState.from_qpos(qpos) for qpos in qpos_list], - options=options, - ) - if joint_plan.positions is None: - raise RuntimeError("Joint-space planning did not produce any positions.") - move_robot_along_trajectory( - sim=sim, - robot=robot, - arm_name=arm_name, - qpos_trajectory=joint_plan.positions, - ) + joint_plan = motion_generator.generate( + target_states=[PlanState.from_qpos(qpos) for qpos in qpos_list], + options=options, + ) + if joint_plan.positions is None: + raise RuntimeError( + "Joint-space planning did not produce any positions." + ) + move_robot_along_trajectory( + sim=sim, + robot=robot, + arm_name=arm_name, + qpos_trajectory=joint_plan.positions, + ) - options.is_linear = True - cartesian_plan = motion_generator.generate( - target_states=[PlanState.from_xpos(xpos) for xpos in xpos_list], - options=options, - ) - if cartesian_plan.positions is None: - raise RuntimeError( - "Cartesian-space planning did not produce any positions." + options.is_linear = True + cartesian_plan = motion_generator.generate( + target_states=[PlanState.from_xpos(xpos) for xpos in xpos_list], + options=options, + ) + if cartesian_plan.positions is None: + raise RuntimeError( + "Cartesian-space planning did not produce any positions." + ) + sim.reset() + move_robot_along_trajectory( + sim=sim, + robot=robot, + arm_name=arm_name, + qpos_trajectory=cartesian_plan.positions, ) - sim.reset() - move_robot_along_trajectory( - sim=sim, - robot=robot, - arm_name=arm_name, - qpos_trajectory=cartesian_plan.positions, - ) finally: - if sim.is_window_recording(): - sim.stop_window_record() - sim.wait_window_record_saves() - sim.destroy() + shutdown_sim(sim) if __name__ == "__main__": diff --git a/scripts/tutorials/sim/srs_solver.py b/scripts/tutorials/sim/srs_solver.py index 2fb8eddae..2176c7087 100644 --- a/scripts/tutorials/sim/srs_solver.py +++ b/scripts/tutorials/sim/srs_solver.py @@ -14,70 +14,87 @@ # limitations under the License. # ---------------------------------------------------------------------------- +"""Demonstrate batched FK and IK with the SRS solver.""" + +from __future__ import annotations + +import argparse import time + import numpy as np import torch -from IPython import embed - from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.robots import DexforceW1Cfg - - -def main(): - # Set print options for better readability - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) - - # Initialize simulation - sim_device = "cpu" - sim = SimulationManager( - SimulationManagerCfg( - headless=False, sim_device=sim_device, width=2200, height=1200 - ) +from embodichain.lab.sim.utility.demo_utils import ( + add_demo_args, + create_default_sim, + maybe_open_window, + maybe_wait_for_user, + setup_print_options, + shutdown_sim, +) + + +def main() -> None: + """Run the SRS solver tutorial.""" + parser = add_demo_args( + argparse.ArgumentParser(description="Run the SRS FK/IK tutorial.") + ) + args = parser.parse_args() + setup_print_options() + + sim = create_default_sim( + args, + width=2200, + height=1200, + add_default_light=False, ) - sim.set_manual_update(False) - - robot: Robot = sim.add_robot(cfg=DexforceW1Cfg.from_dict({"uid": "dexforce_w1"})) - arm_name = "left_arm" - # Set initial joint positions for left arm - qpos_fk_list = [ - torch.tensor([[0.0, 0.0, 0.0, -np.pi / 2, 0.0, 0.0, 0.0]], dtype=torch.float32), - ] - robot.set_qpos(qpos_fk_list[0], joint_ids=robot.get_joint_ids(arm_name)) - - time.sleep(0.5) - - fk_xpos_batch = torch.cat(qpos_fk_list, dim=0) - - fk_xpos_list = robot.compute_fk(qpos=fk_xpos_batch, name=arm_name, to_matrix=True) + try: + sim.set_manual_update(False) - start_time = time.time() - res, ik_qpos = robot.compute_ik( - pose=fk_xpos_list, - name=arm_name, - # joint_seed=qpos_fk_list[0], - return_all_solutions=True, - ) - end_time = time.time() - print( - f"Batch IK computation time for {len(fk_xpos_list)} poses: {end_time - start_time:.6f} seconds" - ) + robot: Robot = sim.add_robot( + cfg=DexforceW1Cfg.from_dict({"uid": "dexforce_w1"}) + ) + arm_name = "left_arm" + qpos = torch.tensor( + [[0.0, 0.0, 0.0, -np.pi / 2, 0.0, 0.0, 0.0]], + dtype=torch.float32, + ) + robot.set_qpos(qpos, joint_ids=robot.get_joint_ids(arm_name)) - if ik_qpos.dim() == 3: - first_solutions = ik_qpos[:, 0, :] - else: - first_solutions = ik_qpos - robot.set_qpos(first_solutions, joint_ids=robot.get_joint_ids(arm_name)) + if not args.auto_play: + time.sleep(0.5) - ik_xpos_list = robot.compute_fk(qpos=first_solutions, name=arm_name, to_matrix=True) + fk_pose = robot.compute_fk(qpos=qpos, name=arm_name, to_matrix=True) - print("fk_xpos_list: ", fk_xpos_list) - print("ik_xpos_list: ", ik_xpos_list) + started_at = time.perf_counter() + success, ik_qpos = robot.compute_ik( + pose=fk_pose, + name=arm_name, + return_all_solutions=True, + ) + elapsed = time.perf_counter() - started_at + print( + f"Batch IK computation time for {len(fk_pose)} poses: {elapsed:.6f} seconds" + ) + print("IK success:", success) + + first_solution = ik_qpos[:, 0, :] if ik_qpos.dim() == 3 else ik_qpos + robot.set_qpos(first_solution, joint_ids=robot.get_joint_ids(arm_name)) + ik_pose = robot.compute_fk( + qpos=first_solution, + name=arm_name, + to_matrix=True, + ) - embed(header="Test SRSSolver example. Press Ctrl-D to exit.") + print("FK poses:", fk_pose) + print("IK poses:", ik_pose) + maybe_open_window(sim, args) + maybe_wait_for_user(args, "Press Enter to exit...") + finally: + shutdown_sim(sim) if __name__ == "__main__": diff --git a/tests/data_pipeline/test_online_data.py b/tests/data_pipeline/test_online_data.py index 4b931e8a6..10d50bdae 100644 --- a/tests/data_pipeline/test_online_data.py +++ b/tests/data_pipeline/test_online_data.py @@ -31,6 +31,8 @@ import multiprocessing as mp import unittest +from unittest.mock import Mock + import pytest import torch @@ -138,6 +140,29 @@ def setup_method(self) -> None: # ----------------------------------------------------------------------- + def test_start_raises_when_worker_exits_before_initialization(self) -> None: + """A failed simulation worker must not leave start() waiting forever.""" + engine = object.__new__(OnlineDataEngine) + engine.cfg = Mock() + engine.shared_buffer = Mock() + engine._lock_index = Mock() + engine._fill_signal = Mock() + engine._init_signal = Mock() + engine._init_signal.is_set.return_value = False + engine._close_signal = Mock() + + process = Mock() + process.pid = 1234 + process.is_alive.return_value = False + engine._mp_ctx = Mock() + engine._mp_ctx.Process.return_value = process + + with pytest.raises(RuntimeError, match="exited before initializing"): + engine.start() + + process.start.assert_called_once() + engine._fill_signal.set.assert_called_once() + def test_sample_batch_shape(self) -> None: """sample_batch returns TensorDict with shape [batch_size, chunk_size].""" BATCH = 3 diff --git a/tests/gym/utils/test_gym_utils.py b/tests/gym/utils/test_gym_utils.py index d03272f8a..20fcaf1e0 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -314,6 +314,20 @@ def test_launcher_preserves_gym_renderer_when_cli_omits_override(): assert merged_config["render_cfg"]["renderer"] == "rt" +def test_launcher_accepts_hyphenated_common_options(): + """Hyphenated aliases should map to the existing underscore destinations.""" + parser = argparse.ArgumentParser() + add_env_launcher_args_to_parser(parser) + + args = parser.parse_args( + ["--num-envs", "4", "--arena-space", "3.5", "--gpu-id", "2"] + ) + + assert args.num_envs == 4 + assert args.arena_space == 3.5 + assert args.gpu_id == 2 + + def test_sensor_and_extra_obs_together(): """Test that both sensors and extra observations work together.""" config = { diff --git a/tests/sim/test_demo_base.py b/tests/sim/test_demo_base.py new file mode 100644 index 000000000..2385c2cdf --- /dev/null +++ b/tests/sim/test_demo_base.py @@ -0,0 +1,67 @@ +# ---------------------------------------------------------------------------- +# 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 pytest + +from embodichain.lab.sim.demo_base import DemoBase + + +def test_demo_base_runs_setup_run_cleanup(): + class SimpleDemo(DemoBase): + def setup(self): + self.sim = Mock(spec=["destroy"]) + + def run(self): + self.ran = True + + demo = SimpleDemo(SimpleNamespace()) + demo.main() + assert demo.ran is True + demo.sim.destroy.assert_called_once() + + +def test_demo_base_cleanup_runs_even_if_run_raises(): + class BrokenDemo(DemoBase): + def setup(self): + self.sim = Mock(spec=["destroy"]) + + def run(self): + raise RuntimeError("boom") + + demo = BrokenDemo(SimpleNamespace()) + with pytest.raises(RuntimeError, match="boom"): + demo.main() + demo.sim.destroy.assert_called_once() + + +def test_demo_base_cleanup_runs_even_if_setup_raises_after_creating_sim(): + class BrokenSetupDemo(DemoBase): + def setup(self): + self.sim = Mock(spec=["destroy"]) + raise RuntimeError("setup failed") + + def run(self): + raise AssertionError("run must not be called") + + demo = BrokenSetupDemo(SimpleNamespace()) + with pytest.raises(RuntimeError, match="setup failed"): + demo.main() + demo.sim.destroy.assert_called_once() diff --git a/tests/sim/utility/test_demo_utils.py b/tests/sim/utility/test_demo_utils.py new file mode 100644 index 000000000..bad847e5b --- /dev/null +++ b/tests/sim/utility/test_demo_utils.py @@ -0,0 +1,335 @@ +# ---------------------------------------------------------------------------- +# 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 argparse +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from unittest.mock import Mock, call, patch + +from embodichain.lab.sim.utility.demo_utils import ( + DEFAULT_DEMO_LOOK_AT, + DemoRecording, + add_demo_args, + create_default_sim, + format_tensor, + maybe_init_gpu_physics, + maybe_open_window, + maybe_wait_for_user, + replay_trajectory, + resolve_demo_steps, + run_simulation_loop, + setup_print_options, + shutdown_sim, +) + + +def test_add_demo_args_adds_expected_flags(): + parser = argparse.ArgumentParser() + parser = add_demo_args(parser) + args = parser.parse_args( + [ + "--headless", + "--auto-play", + "--record-fps", + "60", + "--num-envs", + "2", + "--arena-space", + "3.0", + "--gpu-id", + "1", + ] + ) + assert args.headless is True + assert args.auto_play is True + assert args.record_fps == 60 + assert args.record_steps is None + assert args.no_vis_eef_axis is False + assert args.num_envs == 2 + assert args.arena_space == 3.0 + assert args.gpu_id == 1 + + +def test_add_demo_args_rejects_non_positive_record_steps(): + parser = add_demo_args(argparse.ArgumentParser()) + with pytest.raises(SystemExit): + parser.parse_args(["--record_steps", "0"]) + + +def test_format_tensor_rounds_and_moves_to_cpu(): + tensor = torch.tensor([1.23456789, 2.34567891]) + result = format_tensor(tensor) + assert result == "[1.2346, 2.3457]" + + +def test_setup_print_options_sets_numpy_and_torch(): + setup_print_options() + assert np.get_printoptions()["precision"] == 5 + assert np.get_printoptions()["suppress"] is True + assert torch._tensor_str.PRINT_OPTS.precision == 5 + assert torch._tensor_str.PRINT_OPTS.sci_mode is False + + +def test_shutdown_sim_calls_destroy(): + sim = Mock(spec=["destroy"]) + shutdown_sim(sim) + sim.destroy.assert_called_once() + + +def test_shutdown_sim_finishes_active_recording_before_destroy(): + sim = Mock( + spec=[ + "destroy", + "is_window_recording", + "stop_window_record", + "wait_window_record_saves", + ] + ) + sim.is_window_recording.return_value = True + + shutdown_sim(sim) + + assert sim.method_calls == [ + call.is_window_recording(), + call.stop_window_record(), + call.wait_window_record_saves(), + call.destroy(), + ] + + +def test_create_default_sim_forwards_num_envs_and_headless(): + args = SimpleNamespace(headless=True, device="cpu", renderer="auto", gpu_id=2) + with ( + patch("embodichain.lab.sim.SimulationManager") as mock_sm, + patch("embodichain.lab.sim.SimulationManagerCfg") as mock_cfg_cls, + ): + create_default_sim(args, num_envs=4, add_default_light=False) + cfg_kwargs = mock_cfg_cls.call_args.kwargs + assert cfg_kwargs["num_envs"] == 4 + assert cfg_kwargs["headless"] is True + assert cfg_kwargs["gpu_id"] == 2 + mock_sm.assert_called_once_with(mock_cfg_cls.return_value) + + +def test_create_default_sim_adds_light_when_requested(): + args = SimpleNamespace(headless=True, device="cpu", renderer="auto") + with ( + patch("embodichain.lab.sim.SimulationManager") as mock_sm, + patch("embodichain.lab.sim.SimulationManagerCfg"), + ): + sim = create_default_sim(args, num_envs=1, add_default_light=True) + sim.add_light.assert_called_once() + mock_sm.return_value.add_light.assert_called_once() + + +def test_maybe_init_physics_inits_when_enabled(): + sim = Mock(spec=["is_use_gpu_physics", "init_gpu_physics"]) + sim.is_use_gpu_physics = True + maybe_init_gpu_physics(sim) + sim.init_gpu_physics.assert_called_once() + + +def test_maybe_init_physics_skips_when_disabled(): + sim = Mock(spec=["is_use_gpu_physics", "init_gpu_physics"]) + sim.is_use_gpu_physics = False + maybe_init_gpu_physics(sim) + sim.init_gpu_physics.assert_not_called() + + +def _make_recording_sim(): + sim = Mock( + spec=[ + "start_window_record", + "stop_window_record", + "wait_window_record_saves", + "is_window_recording", + "sim_config", + ] + ) + sim.sim_config = SimpleNamespace(width=1920, height=1080) + sim.start_window_record.return_value = True + sim.is_window_recording.return_value = False + return sim + + +def test_demo_recording_does_nothing_when_record_steps_is_none(): + sim = _make_recording_sim() + args = SimpleNamespace( + record_steps=None, + record_fps=30, + record_save_path="/tmp", + auto_play=False, + headless=True, + ) + with DemoRecording(sim, args, prefix="demo"): + pass + sim.start_window_record.assert_not_called() + + +def test_demo_recording_starts_and_stops_window_record(): + sim = _make_recording_sim() + sim.is_window_recording.return_value = True + args = SimpleNamespace( + record_steps=10, + record_fps=30, + record_save_path="/tmp/recordings", + auto_play=False, + headless=True, + ) + with DemoRecording(sim, args, prefix="demo") as rec: + assert rec.is_active is True + sim.start_window_record.assert_called_once() + call_kwargs = sim.start_window_record.call_args.kwargs + assert call_kwargs["fps"] == 30 + assert call_kwargs["video_prefix"] == "demo" + assert "/tmp/recordings" in call_kwargs["save_path"] + assert call_kwargs["save_path"].endswith(".mp4") + assert call_kwargs["look_at"] == DEFAULT_DEMO_LOOK_AT + sim.stop_window_record.assert_called_once() + sim.wait_window_record_saves.assert_called_once() + + +def test_demo_recording_passes_look_at(): + sim = _make_recording_sim() + args = SimpleNamespace( + record_steps=10, + record_fps=30, + record_save_path="/tmp", + auto_play=False, + headless=True, + ) + look_at = ((1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0)) + with DemoRecording(sim, args, prefix="demo", look_at=look_at): + pass + call_kwargs = sim.start_window_record.call_args.kwargs + assert call_kwargs["look_at"] == look_at + + +def test_demo_recording_uses_exact_mp4_path(): + sim = _make_recording_sim() + args = SimpleNamespace( + record_steps=10, + record_fps=30, + record_save_path="/tmp/custom_demo.mp4", + auto_play=False, + headless=False, + ) + with DemoRecording(sim, args, prefix="ignored"): + pass + call_kwargs = sim.start_window_record.call_args.kwargs + assert call_kwargs["save_path"] == "/tmp/custom_demo.mp4" + assert call_kwargs["look_at"] is None + + +def test_demo_recording_warns_and_skips_on_start_failure(): + sim = _make_recording_sim() + sim.start_window_record.return_value = False + args = SimpleNamespace( + record_steps=10, + record_fps=30, + record_save_path="/tmp", + auto_play=False, + headless=True, + ) + with pytest.warns(UserWarning, match="Failed to start recording"): + with DemoRecording(sim, args, prefix="demo"): + pass + sim.stop_window_record.assert_not_called() + + +def test_maybe_open_window_opens_when_not_headless(): + sim = Mock(spec=["open_window"]) + args = SimpleNamespace(headless=False) + maybe_open_window(sim, args) + sim.open_window.assert_called_once() + + +def test_maybe_open_window_does_nothing_when_headless(): + sim = Mock(spec=["open_window"]) + args = SimpleNamespace(headless=True) + maybe_open_window(sim, args) + sim.open_window.assert_not_called() + + +def test_maybe_wait_for_user_prompts_when_not_auto_play(): + args = SimpleNamespace(auto_play=False) + with patch("builtins.input", return_value="") as mock_input: + maybe_wait_for_user(args, "Press enter") + mock_input.assert_called_once_with("Press enter") + + +def test_maybe_wait_for_user_skips_when_auto_play(): + args = SimpleNamespace(auto_play=True) + with patch("builtins.input") as mock_input: + maybe_wait_for_user(args, "Press enter") + mock_input.assert_not_called() + + +def test_resolve_demo_steps_prefers_explicit_record_steps(): + args = SimpleNamespace(auto_play=True, record_steps=12) + assert resolve_demo_steps(args, auto_play_steps=5) == 12 + + +def test_resolve_demo_steps_makes_auto_play_finite(): + assert ( + resolve_demo_steps( + SimpleNamespace(auto_play=True, record_steps=None), + auto_play_steps=5, + ) + == 5 + ) + assert ( + resolve_demo_steps(SimpleNamespace(auto_play=False, record_steps=None)) is None + ) + + +def test_run_simulation_loop_updates_until_limit_and_calls_hook(): + sim = Mock(spec=["update", "num_envs"]) + sim.num_envs = 2 + on_step = Mock() + + completed = run_simulation_loop( + sim, + max_steps=3, + steps_per_update=2, + log_interval=None, + on_step=on_step, + ) + + assert completed == 3 + assert sim.update.call_args_list == [call(step=2), call(step=2), call(step=2)] + assert on_step.call_args_list == [call(1), call(2), call(3)] + + +def test_replay_trajectory_sets_qpos_and_updates_sim(): + robot = Mock(spec=["set_qpos", "get_joint_ids"]) + robot.get_joint_ids.return_value = None + sim = Mock(spec=["update"]) + # Shape: (batch=1, num_steps=2, num_joints=3) + traj = torch.tensor( + [ + [[0.0, 0.1, 0.2], [0.3, 0.4, 0.5]], + ] + ) + replay_trajectory(sim, robot, traj, post_steps=1, step_size=4, sleep=0.0) + assert robot.set_qpos.call_count == 3 # 2 traj + 1 post + assert sim.update.call_count == 3 + sim.update.assert_has_calls([call(step=4), call(step=4), call(step=2)])