diff --git a/.agents/skills/add-atomic-action/SKILL.md b/.agents/skills/add-atomic-action/SKILL.md index a7de806ce..7ac1fd8bb 100644 --- a/.agents/skills/add-atomic-action/SKILL.md +++ b/.agents/skills/add-atomic-action/SKILL.md @@ -22,10 +22,11 @@ full-DoF trajectory. | Base classes (`ActionCfg`, `AtomicAction`, `WorldState`, `ActionResult`, typed targets, `ObjectSemantics`) | `embodichain/lab/sim/atomic_actions/core.py` | | Affordance types (`Affordance`, `AntipodalAffordance`, `InteractionPoints`) | `embodichain/lab/sim/atomic_actions/affordance.py` | | Stateless trajectory helpers (`TrajectoryBuilder`) | `embodichain/lab/sim/atomic_actions/trajectory.py` | -| Built-in actions (reference implementations) | `embodichain/lab/sim/atomic_actions/actions.py` | +| Built-in action primitives (reference implementations) | `embodichain/lab/sim/atomic_actions/primitives/` | +| Backward-compatible action re-export module | `embodichain/lab/sim/atomic_actions/actions.py` | | Engine + global registry (`register_action`, `AtomicActionEngine.register` / `run`) | `embodichain/lab/sim/atomic_actions/engine.py` | | Public API exports | `embodichain/lab/sim/atomic_actions/__init__.py` | -| Reference docs | `docs/source/overview/sim/atomic_actions.md` | +| Reference docs | `docs/source/overview/sim/atomic_actions/index.md`, `docs/source/overview/sim/atomic_actions/builtin_actions.md` | ## The Contract (read first) @@ -43,19 +44,29 @@ inherit from `MoveEndEffector` or any other action. Each action: full-DoF shaped `(n_envs, n_waypoints, robot.dof)` and `next_state` is the successor `WorldState` (advance `last_qpos` to the trajectory's final row; set/clear/preserve `held_object` per the action's semantics). + - `success` is a per-environment boolean tensor of shape `(n_envs,)` (or a + scalar bool). Use `ActionResult.success_all` (or `.success.all()`) when you + need a single aggregate boolean. There is **no** `validate` method, **no** `**kwargs`, **no** `start_qpos` parameter, **no** `updates_held_object_state` flag, and **no** `get_held_object_state`. The `WorldState` is the single channel for inter-action state. +`ActionCfg` (and therefore every action cfg) carries two motion-source fields used +by `TrajectoryBuilder.plan_arm_traj`: +- `motion_source: str = "ik_interp"` — `"ik_interp"` (batched IK + interpolation) + or `"motion_gen"` (delegates to the batched `MotionGenerator`). +- `planner_type: str | None = None` — `"toppra"` or `"neural"`; required when + `motion_source="motion_gen"`. + ## Steps ### 1. Define the config Add a `@configclass`-decorated class that extends `ActionCfg` **directly** (the cfg -hierarchy is flat — do not inherit from another action's cfg). Place it in -`embodichain/lab/sim/atomic_actions/actions.py` alongside the existing configs, or in -a new file if the action is large. +hierarchy is flat — do not inherit from another action's cfg). For a built-in +primitive, place the config beside the action class in +`embodichain/lab/sim/atomic_actions/primitives/.py`. ```python from __future__ import annotations @@ -147,20 +158,20 @@ class Push(AtomicAction): control_part=self.cfg.control_part, ) - # 3. Plan the arm trajectory via the builder (uses IK + interpolation). + # 3. Plan the arm trajectory via the builder (uses IK + interpolation by default; + # set cfg.motion_source="motion_gen" to use the MotionGenerator instead). target_states = [ [PlanState(xpos=contact_xpos[i], move_type=MoveType.EEF_MOVE)] for i in range(self.n_envs) ] - ok, arm_traj = self.builder.plan_arm_traj( + success, arm_traj = self.builder.plan_arm_traj( target_states, start_arm_qpos, self.cfg.sample_interval, control_part=self.cfg.control_part, arm_dof=self.arm_dof, + cfg=self.cfg, ) - if not ok: - return self._fail(state) # 4. Embed the arm slice into a full-DoF trajectory (n_envs, n_wp, robot.dof). full = torch.empty( @@ -172,7 +183,7 @@ class Push(AtomicAction): full[:, :, self.arm_joint_ids] = arm_traj return ActionResult( - success=True, + success=success, trajectory=full, next_state=WorldState( last_qpos=full[:, -1, :].clone(), @@ -182,7 +193,7 @@ class Push(AtomicAction): def _fail(self, state: WorldState) -> ActionResult: return ActionResult( - success=False, + success=torch.zeros(self.n_envs, dtype=torch.bool, device=self.device), trajectory=torch.empty( (self.n_envs, 0, self.robot_dof), dtype=torch.float32, @@ -195,6 +206,8 @@ class Push(AtomicAction): **Rules:** - `execute()` returns an `ActionResult` — never a bare tuple. - `trajectory` shape is always `(n_envs, n_waypoints, robot.dof)` (full robot DoF). +- Pass `cfg=self.cfg` to every `self.builder.plan_arm_traj(...)` call so the builder + reads `motion_source` / `planner_type` from the action config. - Use `self.builder.` for all trajectory math (`resolve_pose_target`, `resolve_joint_target`, `resolve_start_qpos`, `apply_local_offset`, `plan_arm_traj`, `plan_joint_traj`, `split_three_phase`, `interpolate_hand_qpos`). Do not reimplement @@ -228,11 +241,25 @@ register_action("push", Push) ### 5. Export from the public API -Add the config, action class, and any new target to +Add the config, action class, and any new target to the package exports. For a +built-in primitive, first export it from +`embodichain/lab/sim/atomic_actions/primitives/__init__.py`: + +```python +from .push import Push, PushCfg + +__all__ = [ + ..., + "Push", + "PushCfg", +] +``` + +Then export it from the public API in `embodichain/lab/sim/atomic_actions/__init__.py`: ```python -from .actions import Push, PushCfg +from .primitives import Push, PushCfg # (and from .core import PushTarget if you defined one) __all__ = [ @@ -242,12 +269,16 @@ __all__ = [ ] ``` +Keep `embodichain/lab/sim/atomic_actions/actions.py` as a compatibility facade; +update it only if the new built-in should also be available from the legacy +`embodichain.lab.sim.atomic_actions.actions` import path. + ### 6. Update the supported actions table -Add a row to the table in `docs/source/overview/sim/atomic_actions.md`: +Add a row to the table in `docs/source/overview/sim/atomic_actions/builtin_actions.md`: ```markdown -| `Push` | `PushCfg` | `PushTarget` — contact pose | Approach → push forward | +| `Push` | Single | `PushTarget` — contact pose | Approach → push forward | Add a demo asset or `N/A` | ``` ### 7. Write a test @@ -273,7 +304,7 @@ def test_push_action_returns_full_dof_trajectory(): ): result = action.execute(PushTarget(contact_pose=torch.eye(4)), state) assert isinstance(result, ActionResult) - assert result.success is True + assert result.success_all is True assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) # push preserves held_object assert result.next_state.held_object is state.held_object @@ -289,6 +320,8 @@ def test_push_action_returns_full_dof_trajectory(): | `execute(target, start_qpos=None, **kwargs)` | Signature is `execute(self, target, state: WorldState) -> ActionResult`. No `**kwargs`, no `start_qpos`. | | Reimplementing IK / interpolation inline | Use `self.builder.plan_arm_traj(...)`, `self.builder.plan_joint_traj(...)`, and friends. | | Returning arm-only or arm+hand trajectory | Always embed into full `robot.dof` before returning. | +| Forgetting `cfg=self.cfg` in `plan_arm_traj` | The builder defaults to `motion_source="ik_interp"`; pass `cfg=self.cfg` to opt into `motion_gen` / `planner_type`. | +| Treating `ActionResult.success` as a scalar | It is `(n_envs,)` for batched actions; use `.success_all` or `.success.all()` for a single bool. | | `name` not matching the engine registration key | Keep `cfg.name` identical to the key passed to `engine.register(...)` / `register_action(...)`. | | Forgetting to export from `__init__.py` | Users import from the public API — missing exports cause `ImportError`. | | Inheriting another action's cfg | Cfgs are flat; extend `ActionCfg` directly and declare the fields you need. | @@ -299,8 +332,8 @@ def test_push_action_returns_full_dof_trajectory(): |------|--------| | 1 | Define a flat `@configclass` extending `ActionCfg` with a unique `name` | | 2 | Define a typed target (or reuse `EndEffectorPoseTarget` / `JointPositionTarget` / `NamedJointPositionTarget` / `GraspTarget` / `HeldObjectPoseTarget`) | -| 3 | Subclass `AtomicAction` directly, set `TargetType`, compose `TrajectoryBuilder`, implement `execute(target, state) -> ActionResult` | +| 3 | Subclass `AtomicAction` directly, set `TargetType`, compose `TrajectoryBuilder`, implement `execute(target, state) -> ActionResult` (pass `cfg=self.cfg` to `plan_arm_traj` and return per-env `success`) | | 4 | Register: `engine.register(Push(mg, cfg=...))` (instance) or `register_action("push", Push)` (class) | -| 5 | Export config + action (+ target) from `__init__.py` | -| 6 | Add a row to the supported-actions table in the overview docs | +| 5 | Export config + action (+ target) from `primitives/__init__.py` and `atomic_actions/__init__.py` | +| 6 | Add a row to the supported-actions table in `builtin_actions.md` and update API reference docs | | 7 | Write behavioural tests (target type, full-DoF shape, `WorldState` contract) | diff --git a/.agents/skills/add-test/SKILL.md b/.agents/skills/add-test/SKILL.md index e23c0cd0c..979229305 100644 --- a/.agents/skills/add-test/SKILL.md +++ b/.agents/skills/add-test/SKILL.md @@ -89,6 +89,7 @@ class TestMySimComponent: def teardown_method(self): self.sim.destroy() + SimulationManager.flush_cleanup_queue() def test_basic_behavior(self): result = self.sim.do_something() @@ -99,6 +100,32 @@ class TestMySimComponent: self.sim.do_something(bad_input) ``` +## Resource-Aware Test Classification + +Use the narrowest test type that proves the behavior: + +- Prefer mocks or CPU-only inputs for pure logic and shape validation. +- `tests/conftest.py` automatically classifies conventional CUDA, renderer, and + real-simulation tests. Add `@pytest.mark.gpu`, `@pytest.mark.slow`, or + `@pytest.mark.requires_sim` explicitly only when the test's node id/source + cannot reveal that requirement (for example, a hidden CUDA helper or an + end-to-end toolkit test). + +GPU tests are skipped by default to keep normal test runs within the shared VRAM budget. Run them explicitly and serially: + +```bash +# Default suite: GPU-marked tests are skipped. +pytest tests/ + +# Dedicated GPU suite. Do not add -n unless pytest-xdist is installed. +pytest tests/ --run-gpu -m gpu +``` + +For backend/device matrices, run the complete contract on one representative +configuration and use small (one environment, low-resolution) smoke tests for +the remaining configurations. Always destroy a real `SimulationManager` and +flush its cleanup queue in teardown. + ## Mocking Patterns for Functor Tests Most functor tests don't need a live simulation. Use mock objects following the pattern in `tests/gym/envs/managers/test_reward_functors.py`: @@ -199,6 +226,9 @@ pytest tests//test_.py -v # Single test function pytest tests//test_.py::test_expected_output -v +# GPU-specific test +pytest tests//test_.py --run-gpu -m gpu -v + # Single test class method pytest tests//test_.py::TestMyClass::test_basic_behavior -v ``` @@ -219,6 +249,8 @@ black tests//test_.py | `from __future__` | Required after header | | Magic numbers | Define as named constants with explanatory comments | | Simulation tests | Initialize/teardown in `setup_method`/`teardown_method` | +| CUDA coverage | Use `@pytest.mark.gpu`; run with `--run-gpu -m gpu` | +| Long integration | Use `@pytest.mark.slow`; keep it out of normal PR runs | | Pure-logic tests | Use mock objects, no real sim | | `SceneEntityCfg` | Use `MagicMock(uid="...")` in tests | | Assertions | `assert`, `pytest.approx`, `torch.allclose`, `pytest.raises` | @@ -232,14 +264,16 @@ black tests//test_.py | Using real `SimulationManager` for functor tests | Use `MockEnv`/`MockSim` — much faster, no GPU needed | | Hardcoded numbers without explanation | Define as `EXPECTED_DISTANCE = 0.5 # cube at origin, target at (0.5, 0, 0)` | | Testing multiple concepts in one function | Split into separate `test_` functions | -| Forgetting `teardown_method` | Always call `self.sim.destroy()` in teardown | +| Forgetting cleanup | Call `self.sim.destroy()` and `SimulationManager.flush_cleanup_queue()` in teardown | +| Using full matrices | Use one full representative case and low-resource smoke coverage elsewhere | | Not running `black` on test file | CI checks all files including tests | ## Quick Reference | Action | Command | |--------|---------| -| Run all tests | `pytest tests/` | +| Run default tests | `pytest tests/` | +| Run GPU tests | `pytest tests/ --run-gpu -m gpu` | | Run single file | `pytest tests//test_.py -v` | | Run single test | `pytest tests/::test_ -v` | | Run with print output | `pytest -s tests//test_.py` | diff --git a/.agents/skills/project-dev-context/SKILL.md b/.agents/skills/project-dev-context/SKILL.md index 85423ab0b..28a9e952d 100644 --- a/.agents/skills/project-dev-context/SKILL.md +++ b/.agents/skills/project-dev-context/SKILL.md @@ -16,7 +16,7 @@ Use this skill when: - the request says `刷新项目上下文` - the request says `更新项目上下文` - the request says `写项目上下文` -- the request names a known project topic such as `env-framework`, `manager-functor`, or `ik-solvers` +- the request names a known project topic such as `env-framework`, `manager-functor`, `ik-solvers`, or `atomic-actions` ## Start here diff --git a/.github/workflows/docs-pages.yml b/.github/workflows/docs-pages.yml index 82ae738b5..d1fbeeefb 100644 --- a/.github/workflows/docs-pages.yml +++ b/.github/workflows/docs-pages.yml @@ -14,6 +14,8 @@ on: description: "Existing CI/CD Pipeline run ID to deploy from instead of rebuilding" required: false type: string + release: + types: [published] permissions: actions: read @@ -30,7 +32,8 @@ jobs: if: > (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success' && - github.event.workflow_run.event == 'push') || + github.event.workflow_run.event == 'push' && + github.event.workflow_run.head_branch == 'main') || (github.event_name == 'workflow_dispatch' && inputs.artifact_run_id != '') runs-on: ubuntu-latest @@ -56,7 +59,11 @@ jobs: uses: actions/deploy-pages@v4 build-and-deploy: - if: github.event_name == 'workflow_dispatch' && inputs.artifact_run_id == '' + if: > + (github.event_name == 'workflow_dispatch' && inputs.artifact_run_id == '') || + (github.event_name == 'release' && + github.event.action == 'published' && + startsWith(github.event.release.tag_name, 'v')) runs-on: Linux env: NVIDIA_DRIVER_CAPABILITIES: all @@ -77,7 +84,7 @@ jobs: steps: - uses: actions/checkout@v4 with: - ref: ${{ inputs.ref != '' && inputs.ref || 'main' }} + ref: ${{ github.event_name == 'release' && github.event.release.tag_name || inputs.ref != '' && inputs.ref || 'main' }} - uses: ./.github/actions/activate-conda - name: Restore full multi-version docs site @@ -92,7 +99,7 @@ jobs: shell: bash run: | SITE_URL="https://${{ github.repository_owner }}.github.io/${{ github.event.repository.name }}" - TARGET_REF="${{ inputs.ref != '' && inputs.ref || 'main' }}" + TARGET_REF="${{ github.event_name == 'release' && github.event.release.tag_name || inputs.ref != '' && inputs.ref || 'main' }}" SKIP_VERSION="main" if [[ "${TARGET_REF}" == v* ]]; then SKIP_VERSION="${TARGET_REF}" @@ -105,7 +112,7 @@ jobs: - name: Build docs site shell: bash run: | - TARGET_REF="${{ inputs.ref != '' && inputs.ref || 'main' }}" + TARGET_REF="${{ github.event_name == 'release' && github.event.release.tag_name || inputs.ref != '' && inputs.ref || 'main' }}" pip install -e ".[gensim]" \ --extra-index-url http://pyp.open3dv.site:2345/simple/ \ --trusted-host pyp.open3dv.site \ @@ -144,7 +151,11 @@ jobs: path: ${{ github.workspace }}/docs/build/html deploy-manual-build: - if: github.event_name == 'workflow_dispatch' && inputs.artifact_run_id == '' + if: > + (github.event_name == 'workflow_dispatch' && inputs.artifact_run_id == '') || + (github.event_name == 'release' && + github.event.action == 'published' && + startsWith(github.event.release.tag_name, 'v')) needs: build-and-deploy runs-on: ubuntu-latest environment: diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 53e7dc324..38d474e72 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -156,17 +156,27 @@ jobs: steps: - uses: actions/checkout@v4 - uses: ./.github/actions/activate-conda - - name: Run tests + - name: Install test package run: | pip install -e ".[gensim]" \ --extra-index-url http://pyp.open3dv.site:2345/simple/ \ --trusted-host pyp.open3dv.site \ --extra-index-url https://download.blender.org/pypi/ - echo "Unit test Start" + pip install pytest-xdist + + - name: Run default tests + run: | + echo "Default test suite (GPU-marked tests are skipped)" export HF_ENDPOINT=https://hf-mirror.com pytest tests/docs -q --confcutdir=tests/docs pytest tests + - name: Run GPU tests serially + run: | + echo "Dedicated GPU test suite" + export HF_ENDPOINT=https://hf-mirror.com + pytest tests --run-gpu -m gpu + release-build: if: startsWith(github.ref, 'refs/tags/v') needs: lint @@ -204,24 +214,3 @@ jobs: with: name: python-distributions path: dist/ - - release-publish: - if: startsWith(github.ref, 'refs/tags/v') - needs: release-build - runs-on: ubuntu-latest - environment: - name: pypi - url: https://pypi.org/p/embodichain - permissions: - contents: read - id-token: write # PyPI Trusted Publishing - - steps: - - name: (Release) Download distributions - uses: actions/download-artifact@v4 - with: - name: python-distributions - path: dist/ - - - name: (Release) Publish to PyPI - uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.github/workflows/release-publish.yml b/.github/workflows/release-publish.yml new file mode 100644 index 000000000..7af7ec1e6 --- /dev/null +++ b/.github/workflows/release-publish.yml @@ -0,0 +1,37 @@ +name: Publish Release + +on: + release: + types: [published] + +permissions: + contents: read + +jobs: + publish: + if: startsWith(github.event.release.tag_name, 'v') + runs-on: ubuntu-latest + environment: + name: pypi + url: https://pypi.org/p/embodichain + permissions: + contents: read + id-token: write # PyPI Trusted Publishing + + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ github.event.release.tag_name }} + + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Build distributions from the published tag + run: | + python -m pip install --upgrade pip + python -m pip install build + python -m build + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.gitignore b/.gitignore index 7405b2797..f84991adf 100644 --- a/.gitignore +++ b/.gitignore @@ -190,6 +190,8 @@ embodichain/agents/policy/runs/* *.pth outputs test_configs/* +embodichain_data/ +scripts/tutorials/atomic_action/*.glb wandb/ *.mp4 @@ -200,4 +202,4 @@ wandb/ embodichain/VERSION # benchmark results -scripts/benchmark/rl/reports/* \ No newline at end of file +scripts/benchmark/rl/reports/* diff --git a/AGENTS.md b/AGENTS.md index 0fd66ab38..4cbd9f2d4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -21,7 +21,7 @@ the agent should: 3. Load only the matched Markdown files under `agent_context/` 4. Avoid reading `docs/source/` unless the user explicitly asks for the Sphinx documentation -Available topics: `env-framework`, `manager-functor`, `ik-solvers`, `robot-system`, `sensor-system`, `motion-planning`, `rl-learning`, `configclass-pattern`, `randomization`. +Available topics: `env-framework`, `manager-functor`, `ik-solvers`, `robot-system`, `sensor-system`, `motion-planning`, `atomic-actions`, `rl-learning`, `configclass-pattern`, `randomization`. --- diff --git a/VERSION b/VERSION index ee1372d33..717903969 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.2.2 +0.2.3 diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 31441b253..27ae31649 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -355,3 +355,41 @@ topics: - env-framework - rl-training status: active + + - id: atomic-actions + title: Atomic Actions + aliases: + - atomic action + - atomic actions + - motion primitive + - action primitive + - AtomicAction + - AtomicActionEngine + - TrajectoryBuilder + - 原子动作 + keywords: + - atomic + - action + - AtomicAction + - AtomicActionEngine + - TrajectoryBuilder + - ActionResult + - WorldState + - ActionCfg + - motion_source + - plan_arm_traj + - register_action + paths: + - topics/atomic-actions/atomic-actions.md + source_of_truth: + - embodichain/lab/sim/atomic_actions/core.py + - embodichain/lab/sim/atomic_actions/engine.py + - embodichain/lab/sim/atomic_actions/trajectory.py + - embodichain/lab/sim/atomic_actions/primitives/ + - embodichain/lab/sim/atomic_actions/actions.py + - embodichain/lab/sim/atomic_actions/__init__.py + related_topics: + - motion-planning + - robot-system + - ik-solvers + status: active diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md new file mode 100644 index 000000000..1113208b0 --- /dev/null +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -0,0 +1,145 @@ +# Atomic Actions + +## Entry Points + +| What | Path | +|---|---| +| Base classes, typed targets, configs | `embodichain/lab/sim/atomic_actions/core.py` | +| Engine and global registry | `embodichain/lab/sim/atomic_actions/engine.py` | +| Trajectory helpers | `embodichain/lab/sim/atomic_actions/trajectory.py` | +| Built-in primitives | `embodichain/lab/sim/atomic_actions/primitives/` | +| Legacy re-export facade | `embodichain/lab/sim/atomic_actions/actions.py` | +| Public API | `embodichain/lab/sim/atomic_actions/__init__.py` | + +## Overview + +Atomic actions are env-batched motion primitives chained by `AtomicActionEngine`. Each action receives a typed target and a `WorldState`, plans a full-DoF trajectory for all environments, and returns an `ActionResult`. The engine threads `WorldState` from one action to the next and concatenates trajectories along the time axis. + +``` +AtomicActionEngine + ├─ AtomicAction(s) ← one primitive per class, e.g. MoveEndEffector, PickUp + │ │ + │ └── TrajectoryBuilder ← IK/interpolation and MotionGenerator dispatch + │ + └── WorldState ← last_qpos + held_object/coordinated_held_object +``` + +All tensor shapes carry a leading batch dim `B = n_envs`. + +## Core Types + +### Typed Targets + +Frozen dataclasses accepted by actions via their `TargetType` class variable. + +| Target | Holds | Used by | +|---|---|---| +| `EndEffectorPoseTarget(xpos)` | `(4,4)`, `(B,4,4)` or `(B,n_waypoint,4,4)` EEF pose | `MoveEndEffector`, `Place`, `Press` | +| `JointPositionTarget(qpos)` | `(dof,)`, `(B,dof)` or `(B,n_waypoint,dof)` joint positions | `MoveJoints` | +| `NamedJointPositionTarget(name)` | Name resolved from `MoveJointsCfg.named_joint_positions` | `MoveJoints` | +| `GraspTarget(semantics)` | `ObjectSemantics` describing the object to grasp | `PickUp` | +| `HeldObjectPoseTarget(pose)` | `(4,4)` or `(B,4,4)` target pose for the held object | `MoveHeldObject` | +| `CoordinatedPickmentTarget(...)` | Shared object + left/right object-to-EEF transforms | `CoordinatedPickment` | +| `CoordinatedPlacementTarget(...)` | Two held-object states + target poses | `CoordinatedPlacement` | + +### WorldState + +Threaded between actions: +- `last_qpos: torch.Tensor` — shape `(B, robot.dof)`, robot joint positions at the start of the next action. +- `held_object: HeldObjectState | None` — object held by one gripper. +- `coordinated_held_object: CoordinatedHeldObjectState | None` — object jointly held by two grippers. + +`HeldObjectState` stores the object's semantics plus the object-to-EEF transform and grasp pose (both `(B, 4, 4)`). + +### ActionResult + +Every `execute()` returns: +- `success: bool | torch.Tensor` — per-env boolean tensor of shape `(B,)` for batched actions. +- `trajectory: torch.Tensor` — full-robot trajectory `(B, n_waypoints, robot.dof)`. +- `next_state: WorldState` — state to feed into the next action. + +Helpers: +- `ActionResult.success_all` — `True` only when every env succeeded. +- `bool(action_result)` — deprecated; delegates to `success_all` and emits a `DeprecationWarning`. + +## Action Configuration + +`ActionCfg` (base for all action configs): + +| Field | Type | Default | Notes | +|---|---|---|---| +| `name` | `str` | `"default"` | Engine registration key | +| `control_part` | `str` | `"arm"` | Robot control part to move | +| `interpolation_type` | `str` | `"linear"` | Interpolation flavor | +| `velocity_limit` | `float \| None` | `None` | Used on the `motion_gen` path | +| `acceleration_limit` | `float \| None` | `None` | Used on the `motion_gen` path | +| `motion_source` | `str` | `"ik_interp"` | `"ik_interp"` (batched IK + interpolation) or `"motion_gen"` (batched `MotionGenerator`) | +| `planner_type` | `str \| None` | `None` | `"toppra"` or `"neural"; required when `motion_source="motion_gen"` | + +The base config is flat: every action cfg extends `ActionCfg` directly, even if it also carries hand open/close fields (see `PickUpCfg` / `PlaceCfg`). + +## TrajectoryBuilder + +Stateless helper owned by each action. Key methods: + +| Method | Purpose | +|---|---| +| `resolve_pose_target(target, n_envs)` | Broadcast EEF target to `(B,4,4)` or `(B,n,4,4)` | +| `resolve_joint_target(target, n_envs, joint_dof, control_part)` | Broadcast joint target to `(B,dof)` or `(B,n,dof)` | +| `resolve_start_qpos(start_qpos, n_envs, arm_dof, control_part)` | Broadcast start qpos to `(B, arm_dof)` | +| `plan_arm_traj(target_states_list, start_qpos, n_waypoints, control_part, arm_dof, cfg=None)` | Returns `(success:(B,), trajectory:(B,n_waypoints,arm_dof))`. Selects `ik_interp` or `motion_gen` from `cfg.motion_source`. | +| `plan_joint_traj(start_qpos, target_qpos, n_waypoints)` | Joint-space interpolation; always succeeds. | +| `split_three_phase(...)` | Split sample interval into motion / hand-interp / motion phases. | +| `interpolate_hand_qpos(...)` | Interpolate gripper qpos between two states. | + +`plan_arm_traj` input contract for actions: `target_states_list` is `list[list[PlanState]]` where the outer list is per-env and the inner list is per-waypoint. The builder internally converts to a batched `list[PlanState]` (each carrying `(B, ...)` tensors) when dispatching to `MotionGenerator`. + +## AtomicActionEngine + +```python +engine = AtomicActionEngine(motion_generator) +engine.register(MoveEndEffector(motion_generator, cfg=MoveEndEffectorCfg())) +success, traj, final_state = engine.run(steps=[("move_end_effector", target)]) +``` + +`run(steps, state=None) -> (success, full_traj, final_state)`: +- `success` is a `(B,)` bool tensor indicating which environments completed every step. +- Failed environments hold their last successful joint position in both `full_traj` and `final_state.last_qpos` for the remainder of the sequence. +- If all envs fail, the loop stops early. +- `state` defaults to `WorldState(last_qpos=robot.get_qpos().clone())`. + +## Built-in Primitives + +| Action | Target | Notes | +|---|---|---| +| `MoveEndEffector` | `EndEffectorPoseTarget` | EEF pose move | +| `MoveJoints` | `JointPositionTarget` / `NamedJointPositionTarget` | Joint-space interpolation | +| `PickUp` | `GraspTarget` | Approach → close gripper → lift; populates `held_object` | +| `MoveHeldObject` | `HeldObjectPoseTarget` | Move held object; preserves `held_object` | +| `Place` | `EndEffectorPoseTarget` | Lower → open gripper → retract; clears `held_object` | +| `Press` | `EndEffectorPoseTarget` | Close gripper → press down → return | +| `CoordinatedPickment` | `CoordinatedPickmentTarget` | Dual-arm shared-object pick | +| `CoordinatedPlacement` | `CoordinatedPlacementTarget` | Dual-arm placement | + +## Implementing a New Action + +1. Create a flat `@configclass` extending `ActionCfg` with a unique `name`. +2. Reuse an existing target or define a new frozen dataclass in `core.py`. +3. Subclass `AtomicAction` directly (do not inherit from another action). Set `TargetType` and compose a `TrajectoryBuilder`. +4. Implement `execute(self, target, state: WorldState) -> ActionResult`: + - Resolve batched targets and start qpos via `self.builder`. + - Call `self.builder.plan_arm_traj(..., cfg=self.cfg)` if using arm motion. + - Return per-env `success` (a `(B,)` tensor if any env can fail, or `torch.ones(...)` for always-succeeding paths). + - Embed the arm trajectory into full-DoF shape `(B, n_wp, robot.dof)`. + - Advance `last_qpos` to the final row and preserve/update/clear `held_object`. +5. Register an instance with the engine or globally via `register_action(name, ActionClass)`. +6. Export from `primitives/__init__.py` and `atomic_actions/__init__.py`. + +## Common Failure Modes + +- **Forgetting `cfg=self.cfg` in `plan_arm_traj`** — without it, `motion_source` defaults to `"ik_interp"` and `planner_type` is ignored. +- **Treating `success` as scalar** — `ActionResult.success` is `(B,)` for all built-ins; use `success_all` or `success.all()` for a single bool. +- **Using `bool(action_result)` in new code** — still works but emits a `DeprecationWarning`; prefer `.success_all`. +- **Returning arm-only trajectory** — actions must embed into `(B, n_wp, robot.dof)` before returning. +- **`motion_source="motion_gen"` without a MotionGenerator** — the engine passes its own `motion_generator` to each action's `TrajectoryBuilder`; if it is `None`, the action raises `ValueError` at execute time. +- **Wrong planner_type** — `planner_type` must be registered in `MotionGenerator._support_planner_dict` (currently `"toppra"` or `"neural"`). diff --git a/agent_context/topics/motion-planning/motion-planning.md b/agent_context/topics/motion-planning/motion-planning.md index 4398013be..56fd40454 100644 --- a/agent_context/topics/motion-planning/motion-planning.md +++ b/agent_context/topics/motion-planning/motion-planning.md @@ -5,12 +5,12 @@ | What | Path | |---|---| | Planner registry | `embodichain/lab/sim/planners/__init__.py` | -| Base planner class & config | `embodichain/lab/sim/planners/base_planner.py` → `BasePlanner`, `BasePlannerCfg`, `PlanOptions` | +| Base planner class & config | `embodichain/lab/sim/planners/base_planner.py` → `BasePlanner`, `BasePlannerCfg`, `PlanOptions`, `validate_plan_options` | | TOPPRA planner | `embodichain/lab/sim/planners/toppra_planner.py` → `ToppraPlanner`, `ToppraPlannerCfg`, `ToppraPlanOptions` | | Neural planner | `embodichain/lab/sim/planners/neural_planner.py` → `NeuralPlanner`, `NeuralPlannerCfg`, `NeuralPlanOptions` | | Planner assets | `embodichain/data/assets/planner_assets.py` → `download_neural_planner_checkpoint()` | | Motion generator | `embodichain/lab/sim/planners/motion_generator.py` → `MotionGenerator`, `MotionGenCfg`, `MotionGenOptions` | -| Planner utilities & data types | `embodichain/lab/sim/planners/utils.py` → `PlanState`, `PlanResult`, `MoveType`, `MovePart`, `TrajectorySampleMethod` | +| Planner utilities & data types | `embodichain/lab/sim/planners/utils.py` → `PlanState`, `PlanResult`, `MoveType`, `MovePart`, `TrajectorySampleMethod`, `interpolate_xpos_batched` | ## Overview @@ -20,12 +20,14 @@ The planning stack has two layers: All planners resolve their robot at init via `SimulationManager.get_instance().get_robot(cfg.robot_uid)`. +The entire stack is **env-batched** (`B = num_envs`). `PlanState` / `PlanResult` tensors carry a leading `B` dimension; `BasePlanner.plan()` and `MotionGenerator.generate()` operate on `B` environments in one call. + ## Planner Hierarchy ``` BasePlanner (ABC) - ├─ ToppraPlanner Time-optimal path parameterization - └─ NeuralPlanner (experimental) APG waypoint rollout + ├─ ToppraPlanner Time-optimal path parameterization (fork-pool fan-out) + └─ NeuralPlanner (experimental) APG waypoint rollout (native batching) MotionGenerator Wraps any BasePlanner; adds interpolation and multi-part support ``` @@ -33,7 +35,7 @@ MotionGenerator Wraps any BasePlanner; adds interpolation and multi-pa Config hierarchy: ``` BasePlannerCfg robot_uid (MISSING), planner_type - ├─ ToppraPlannerCfg planner_type = "toppra" + ├─ ToppraPlannerCfg planner_type = "toppra", max_workers, mp_context └─ NeuralPlannerCfg planner_type = "neural", checkpoint_path (MISSING) MotionGenCfg planner_cfg (MISSING — must be a BasePlannerCfg subclass) @@ -42,7 +44,7 @@ PlanOptions (empty base) ├─ ToppraPlanOptions constraints, sample_method, sample_interval └─ NeuralPlanOptions control_part, start_qpos, max_steps -MotionGenOptions start_qpos, control_part, plan_opts, is_interpolate, +MotionGenOptions start_qpos (B, DOF), control_part, plan_opts, is_interpolate, interpolate_nums, is_linear, interpolate_position_step, interpolate_angle_step ``` @@ -54,8 +56,15 @@ MotionGenOptions start_qpos, control_part, plan_opts, is_interpolate, Time-optimal path parameterization using the [toppra](https://github.com/hungpham2511/toppra) library. - **Dependency**: `pip install toppra==0.6.3` (import-time error if missing). -- **Single-instance only**: raises `NotImplementedError` if `robot.num_instances > 1`. -- **Method**: `plan(target_states, options=ToppraPlanOptions()) → PlanResult` +- **Batched**: accepts `target_states` whose tensor fields carry leading batch dim `B`. Internally fans out `B` independent single-env TOPPRA solves across a `ProcessPoolExecutor`. +- **Method**: `plan(target_states, options=ToppraPlanOptions()) -> PlanResult` + +`ToppraPlannerCfg` fields: + +| Field | Type | Default | Notes | +|---|---|---|---| +| `max_workers` | `int \| None` | `None` | Worker process count. `None` → `min(cpu_count() // 2, B)`. | +| `mp_context` | `str \| None` | `None` | Multiprocessing start method. `None` auto-selects `fork` on CPU and `spawn` on GPU; can be set to `"fork"` or `"spawn"`. | `ToppraPlanOptions` fields: @@ -65,14 +74,21 @@ Time-optimal path parameterization using the [toppra](https://github.com/hungpha | `sample_method` | `TrajectorySampleMethod` | `QUANTITY` | `TIME`, `QUANTITY`, or `DISTANCE` | | `sample_interval` | `float \| int` | `0.01` | Time interval (seconds) or sample count depending on method | +Worker details: +- The pure-numpy module-level worker `_toppra_solve_one_env` is picklable and never touches CUDA/Warp/sim state. +- `B == 1` or `max_workers == 1` uses an inline fallback (no IPC). +- `TIME` sampling can produce per-env waypoint counts; shorter trajectories are tail-padded by repeating the final waypoint and `duration` records the real endpoint per env. +- Per-env failures set `success[b] = False` and fill the env's trajectory with its start qpos; other envs continue. `BrokenProcessPool` tears the pool down and rebuilds it on the next call. + ### NeuralPlanner (experimental) Learning-based EEF waypoint planner. Franka Panda only. - Checkpoint: `download_neural_planner_checkpoint()` from HuggingFace (gated, needs `HF_TOKEN`) - Use via `MotionGenerator` with `planner_type="neural"` and `plan_opts=NeuralPlanOptions(...)` -- Input: `EEF_MOVE` `PlanState` list with 4×4 `xpos` +- Input: `EEF_MOVE` `PlanState` list with batched `xpos:(B, 4, 4)` - Key cfg: `checkpoint_path` (from download), `control_part` +- Natively batched: transformer forward, reach checks, and convergence holds all operate on `(B, ...)`. ### MotionGenerator @@ -81,12 +97,13 @@ Unified interface for trajectory planning with optional pre-interpolation. - Wraps a `BasePlanner` instance (resolved from `planner_cfg.planner_type`). - Supported planner types: `{"toppra": (ToppraPlanner, ToppraPlannerCfg), "neural": (NeuralPlanner, NeuralPlannerCfg)}`. - `MotionGenCfg.planner_cfg` is **MISSING** — must be provided. +- `generate()` and `interpolate_trajectory()` are env-batched (`B, N, DOF`). `MotionGenOptions` fields: | Field | Type | Default | Notes | |---|---|---|---| -| `start_qpos` | `torch.Tensor \| None` | `None` | Override starting joint config; `None` = use current robot state | +| `start_qpos` | `torch.Tensor \| None` | `None` | Override starting joint config, shape `(B, DOF)`; `None` = use current robot state | | `control_part` | `str \| None` | `None` | Robot control part name (must match `RobotCfg.control_parts` key) | | `plan_opts` | `PlanOptions \| None` | `None` | Passed to the underlying planner | | `is_interpolate` | `bool` | `False` | Pre-interpolate waypoints before planning | @@ -99,35 +116,43 @@ Unified interface for trajectory planning with optional pre-interpolation. ### PlanState (input) -Describes one waypoint or action: +Describes one waypoint or action. Tensor fields carry a leading batch dim `B`; enum/scalar fields are shared across `B`. | Field | Type | Notes | |---|---|---| | `move_type` | `MoveType` | `TOOL`, `EEF_MOVE`, `JOINT_MOVE`, `SYNC`, `PAUSE` | | `move_part` | `MovePart` | `LEFT`, `RIGHT`, `BOTH`, `TORSO`, `ALL` | -| `xpos` | `torch.Tensor \| None` | 4×4 target TCP pose (for `EEF_MOVE`) | -| `qpos` | `torch.Tensor \| None` | Target joint angles `(DOF,)` (for `JOINT_MOVE`) | +| `xpos` | `torch.Tensor \| None` | Target TCP pose `(B, 4, 4)` for `EEF_MOVE` | +| `qpos` | `torch.Tensor \| None` | Target joint angles `(B, DOF)` for `JOINT_MOVE` | +| `qvel` / `qacc` | `torch.Tensor \| None` | Target joint velocities / accelerations `(B, DOF)` | | `is_open` | `bool` | Tool open/close (for `TOOL`) | | `is_world_coordinate` | `bool` | `True` = world frame; `False` = relative | | `pause_seconds` | `float` | Duration for `PAUSE` move type | +Convenience constructors: +- `PlanState.from_qpos(qpos:(B,DOF), move_type=JOINT_MOVE, ...) -> PlanState` +- `PlanState.from_xpos(xpos:(B,4,4), move_type=EEF_MOVE, ...) -> PlanState` +- `PlanState.single(qpos=(DOF,)\|None, xpos=(4,4)\|None, ...) -> PlanState` — unsqueezes single-env tensors to `B=1` (idempotent on already-batched tensors). + ### PlanResult (output) | Field | Type | Notes | |---|---|---| -| `success` | `bool \| torch.Tensor` | Whether planning succeeded | -| `xpos_list` | `torch.Tensor \| None` | EEF poses `(N, 4, 4)` | -| `positions` | `torch.Tensor \| None` | Joint positions `(N, DOF)` | -| `velocities` | `torch.Tensor \| None` | Joint velocities `(N, DOF)` | -| `accelerations` | `torch.Tensor \| None` | Joint accelerations `(N, DOF)` | -| `dt` | `torch.Tensor \| None` | Per-step time durations `(N,)` | -| `duration` | `float \| torch.Tensor` | Total trajectory time (seconds) | +| `success` | `bool \| torch.Tensor` | Per-env success `(B,)` bool tensor (or scalar bool) | +| `xpos_list` | `torch.Tensor \| None` | EEF poses `(B, N, 4, 4)` | +| `positions` | `torch.Tensor \| None` | Joint positions `(B, N, DOF)` | +| `velocities` | `torch.Tensor \| None` | Joint velocities `(B, N, DOF)` | +| `accelerations` | `torch.Tensor \| None` | Joint accelerations `(B, N, DOF)` | +| `dt` | `torch.Tensor \| None` | Per-step time durations `(B, N)` | +| `duration` | `float \| torch.Tensor` | Total trajectory time per env `(B,)` | + +Helper: `PlanResult.is_all_success() -> bool` returns `True` only when every env succeeded. ### MoveType enum | Value | Meaning | |---|---| -| `TOOL` | Tool open/close command | +| `TOOL` | Tool open or close command | | `EEF_MOVE` | End-effector Cartesian move (IK + trajectory) | | `JOINT_MOVE` | Joint-space move (trajectory planning only) | | `SYNC` | Synchronized dual-arm movement | @@ -161,11 +186,13 @@ Describes one waypoint or action: ### validate_plan_options decorator -Applied to `plan()` methods to type-check the `options` argument at runtime. Supports three styles: +Applied to `plan()` methods to type-check the `options` argument at runtime and enforce batch consistency. Supports three styles: - `@validate_plan_options` — bare; validates against base `PlanOptions`. - `@validate_plan_options()` — called with no args; same as above. - `@validate_plan_options(options_cls=MyPlanOptions)` — custom options class. +The decorator checks that every `PlanState` in `target_states` shares the same leading batch dim `B` and that `B` matches `robot.num_instances` (or is `1`). + ### Constraint checking `BasePlanner.is_satisfied_constraint(vels, accs, constraints)` verifies trajectory outputs stay within limits. Tolerance: 10% for velocity, 25% for acceleration. Supports batch dimensions `(B, N, DOF)`. @@ -175,8 +202,9 @@ Applied to `plan()` methods to type-check the `options` argument at runtime. Sup - **`robot_uid` is MISSING** — `BasePlannerCfg.robot_uid` defaults to `MISSING`. Forgetting to set it raises `ValueError` at planner init. - **Robot not found** — planner init calls `SimulationManager.get_instance().get_robot(uid)`. If the robot hasn't been added to the sim yet, this returns `None` and raises `ValueError`. - **toppra not installed** — `ToppraPlanner` import fails with `ImportError` at module load time if `toppra==0.6.3` is not installed. -- **Multi-instance robot with ToppraPlanner** — `ToppraPlanner.__init__` raises `NotImplementedError` if `robot.num_instances > 1`. Use a batch-capable planner or single-instance setup. -- **Wrong PlanOptions subclass** — `@validate_plan_options(options_cls=ToppraPlanOptions)` rejects non-matching options types. Passing a base `PlanOptions()` to `ToppraPlanner.plan()` will error. +- **Batch dim mismatch** — `@validate_plan_options` raises `ValueError` if `PlanState` entries have inconsistent `B` or if `B` does not equal `robot.num_instances`. +- **Single-env caller shape mismatch** — legacy callers passing `(DOF,)` qpos or `(4,4)` xpos must wrap with `PlanState.single(...)` or call `from_qpos`/`from_xpos` with a leading `B=1` dim. - **MotionGenerator planner_type not registered** — if `planner_cfg.planner_type` is not in `_support_planner_dict`, `MotionGenerator.__init__` fails. Register new planners there first. - **Interpolation with unsupported MoveType** — pre-interpolation in `MotionGenOptions` only works for `EEF_MOVE` and `JOINT_MOVE`. Using it with `TOOL`, `SYNC`, or `PAUSE` is ignored or produces unexpected results. - **Constraint tolerance** — `is_satisfied_constraint` allows 10% velocity / 25% acceleration overshoot. Dense waypoint trajectories may appear to violate constraints but pass validation. +- **Fork safety with GPU sim** — `ToppraPlannerCfg.mp_context=None` defaults to `spawn` on GPU to avoid fork-after-CUDA-init hazards. Force `fork` only when the sim device is CPU or you have verified it is safe. diff --git a/agent_context/topics/randomization/randomization.md b/agent_context/topics/randomization/randomization.md index b7f894a65..057206007 100644 --- a/agent_context/topics/randomization/randomization.md +++ b/agent_context/topics/randomization/randomization.md @@ -42,13 +42,16 @@ The `__init__.py` of the randomization package re-exports everything via `from . | `randomize_visual_material` | Random material properties | (varies) | | `randomize_camera_extrinsics` | Camera pose | `pos_range`, `euler_range` (attach mode) or `eye_range`, `target_range`, `up_range` (look-at mode) | | `randomize_camera_intrinsics` | Camera intrinsic params | (varies) | -| `randomize_light` | Light pos/color/intensity | `position_range`, `color_range`, `intensity_range` | +| `randomize_light` | Light pos/color/intensity/direction | `position_range`, `color_range`, `intensity_range`, `direction_range` | | `randomize_emission_light` | Emission light props | (varies) | | `randomize_indirect_lighting` | Indirect lighting | (varies) | - Camera extrinsics auto-detect mode: if `extrinsics.parent` is set → attach mode (pos/euler perturbation via `set_local_pose`); if `extrinsics.eye` is set → look-at mode (eye/target/up perturbation via `look_at`). - `set_rigid_object_visual_material` is deterministic (not random) but uses the same functor mechanism for fixed material assignment at reset. - Light randomization applies the **same values across all envs** (documented limitation). +- ``position_range`` is ignored for global scene lights (``"sun"``, ``"direction"``). Use ``direction_range`` instead. +- ``direction_range`` is only applicable for directional light types (``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, ``"mesh"``). +- Global lights (``"sun"``, ``"direction"``) have ``num_instances == 1``; all other types are batched per environment. ### Spatial (`spatial.py`) diff --git a/docs/source/_static/atomic_actions/coordinated_pickment.gif b/docs/source/_static/atomic_actions/coordinated_pickment.gif new file mode 100644 index 000000000..cbacf8494 Binary files /dev/null and b/docs/source/_static/atomic_actions/coordinated_pickment.gif differ diff --git a/docs/source/_static/atomic_actions/coordinated_placement.gif b/docs/source/_static/atomic_actions/coordinated_placement.gif new file mode 100644 index 000000000..a10785652 Binary files /dev/null and b/docs/source/_static/atomic_actions/coordinated_placement.gif differ diff --git a/docs/source/_static/atomic_actions/press.gif b/docs/source/_static/atomic_actions/press.gif new file mode 100644 index 000000000..1eafdc319 Binary files /dev/null and b/docs/source/_static/atomic_actions/press.gif differ diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.primitives.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.primitives.rst new file mode 100644 index 000000000..9b7defb95 --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.primitives.rst @@ -0,0 +1,91 @@ +embodichain.lab.sim.atomic_actions.primitives +============================================ + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives + + .. rubric:: Built-in Primitive Actions + + .. autosummary:: + + MoveEndEffectorCfg + MoveEndEffector + MoveJointsCfg + MoveJoints + PickUpCfg + PickUp + MoveHeldObjectCfg + MoveHeldObject + PlaceCfg + Place + PressCfg + Press + CoordinatedPickmentCfg + CoordinatedPickment + CoordinatedPlacementCfg + CoordinatedPlacement + +.. currentmodule:: embodichain.lab.sim.atomic_actions.primitives + +MoveEndEffector +--------------- + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives.move_end_effector + :members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict + +MoveJoints +---------- + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives.move_joints + :members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict + +PickUp +------ + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives.pick_up + :members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict + +MoveHeldObject +-------------- + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives.move_held_object + :members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict + +Place +----- + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives.place + :members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict + +Press +----- + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives.press + :members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict + +CoordinatedPickment +------------------- + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives.coordinated_pickment + :members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict + +CoordinatedPlacement +-------------------- + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives.coordinated_placement + :members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst index c2982d04e..fb8360b77 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst @@ -16,8 +16,11 @@ embodichain.lab.sim.atomic_actions NamedJointPositionTarget GraspTarget HeldObjectPoseTarget + CoordinatedPickmentTarget + CoordinatedPlacementTarget Target HeldObjectState + CoordinatedHeldObjectState WorldState ActionResult ActionCfg @@ -33,10 +36,31 @@ embodichain.lab.sim.atomic_actions MoveHeldObject PlaceCfg Place + PressCfg + Press + CoordinatedPickmentCfg + CoordinatedPickment + CoordinatedPlacementCfg + CoordinatedPlacement AtomicActionEngine +.. toctree:: + :maxdepth: 1 + :hidden: + + embodichain.lab.sim.atomic_actions.primitives + .. currentmodule:: embodichain.lab.sim.atomic_actions +Layout +------ + +The public API is exported from ``embodichain.lab.sim.atomic_actions``. Built-in +primitive implementations live under +``embodichain.lab.sim.atomic_actions.primitives`` and +``embodichain.lab.sim.atomic_actions.actions`` remains a compatibility re-export +for existing imports. + Core ---- @@ -76,12 +100,24 @@ Core :members: :show-inheritance: +.. autoclass:: CoordinatedPickmentTarget + :members: + :show-inheritance: + +.. autoclass:: CoordinatedPlacementTarget + :members: + :show-inheritance: + .. autodata:: Target .. autoclass:: HeldObjectState :members: :show-inheritance: +.. autoclass:: CoordinatedHeldObjectState + :members: + :show-inheritance: + .. autoclass:: WorldState :members: :show-inheritance: @@ -153,6 +189,33 @@ Actions :members: :show-inheritance: +.. autoclass:: PressCfg + :members: + :exclude-members: __init__, copy, replace, to_dict + :show-inheritance: + +.. autoclass:: Press + :members: + :show-inheritance: + +.. autoclass:: CoordinatedPickmentCfg + :members: + :exclude-members: __init__, copy, replace, to_dict + :show-inheritance: + +.. autoclass:: CoordinatedPickment + :members: + :show-inheritance: + +.. autoclass:: CoordinatedPlacementCfg + :members: + :exclude-members: __init__, copy, replace, to_dict + :show-inheritance: + +.. autoclass:: CoordinatedPlacement + :members: + :show-inheritance: + Engine & Registry ----------------- diff --git a/docs/source/features/interaction/window.md b/docs/source/features/interaction/window.md index a83e5ccf2..10faf2110 100644 --- a/docs/source/features/interaction/window.md +++ b/docs/source/features/interaction/window.md @@ -38,9 +38,12 @@ In ORBIT mode, plain `W/A/S/D/Q/E` does not move the view. Hold **Left Ctrl** wh | Input | Operation | |-------|-----------| | **Viewer recording (toggle)** | Press **`r`** to **start** recording what the interactive viewer shows, and press **`r`** again to **stop** and save as MP4 videos. Recording uses a hidden camera that follows the live viewer camera pose, so the exported videos match the on-screen view. Useful for debugging and recording demos. | +| **Print camera pose** | Press **`p`** to print the current viewer pose as an executable `window.set_look_at(...)` call. | Recording hotkey registration is controlled by `SimConfig.window_record.enable_hotkey` (enabled by default). You can also call `SimulationManager.start_window_record()`, `stop_window_record()`, or `toggle_window_record()` programmatically. +The camera-pose hotkey is controlled by `SimulationManagerCfg.window_camera_pose.enable_hotkey` and prints look-at form by default. Set `SimulationManagerCfg.window_camera_pose.convert_to_look_at=False` to print the raw 4x4 pose matrix instead. The same output can be requested programmatically with `SimulationManager.print_window_camera_pose()`. + ## Customizing Window Events Users can create their own custom window interaction controls by subclassing the `ObjectManipulator` class (provided by `dexsim`). This allows for the implementation of specific behaviors and responses to user inputs. diff --git a/docs/source/overview/gym/action_functors.md b/docs/source/overview/gym/action_functors.md index 225424da0..375c0a3c9 100644 --- a/docs/source/overview/gym/action_functors.md +++ b/docs/source/overview/gym/action_functors.md @@ -5,6 +5,13 @@ This page lists all available action terms that can be used with the Action Manager. Action terms are configured using {class}`~cfg.ActionTermCfg` and are responsible for processing raw actions from the policy and converting them to the format expected by the robot (e.g., qpos, qvel, qf). +## Quick Reference + +- Use this page when defining the policy-facing action space for RL or closed-loop control tasks. +- Action terms are configured with {class}`~cfg.ActionTermCfg`. +- Action terms transform one batched policy action into robot control commands for all environments. +- Check each term's ``action_dim`` to determine how many values the policy must output. + ````{tip} **Using an AI coding agent?** Use the **`/add-functor`** skill to scaffold a new action term with the correct class structure, `ActionTermCfg` registration, and module placement in `actions.py`. ```` diff --git a/docs/source/overview/gym/dataset_functors.md b/docs/source/overview/gym/dataset_functors.md index 92fc71cbd..c8fe58224 100644 --- a/docs/source/overview/gym/dataset_functors.md +++ b/docs/source/overview/gym/dataset_functors.md @@ -5,6 +5,10 @@ This page lists all available dataset functors that can be used with the Dataset Manager. Dataset functors are configured using {class}`~cfg.DatasetFunctorCfg` and are responsible for collecting and saving episode data during environment interaction. +```{note} +This page covers structured dataset export. If you only need human-viewable debug or demo videos from a fixed camera, use {class}`~record.record_camera_data` on {doc}`event_functors`. +``` + ````{tip} **Using an AI coding agent?** Use the **`/add-functor`** skill to scaffold a new dataset functor with the correct signature, `DatasetFunctorCfg` registration, and module placement in `datasets.py`. ```` @@ -78,6 +82,23 @@ The LeRobotRecorder saves the following data for each frame: - ``observation.images.{sensor_name}``: Camera images (if sensors present) - ``observation.images.{sensor_name}_right``: Right camera images (for stereo cameras) +### Dataset Recording vs Video Recording + +```{list-table} Recording Options +:header-rows: 1 +:widths: 30 35 35 + +* - Need + - Use + - Why +* - Training or imitation-learning data + - {class}`~datasets.LeRobotRecorder` + - Saves structured observation, action, and metadata for downstream pipelines. +* - Quick qualitative inspection or demos + - {class}`~record.record_camera_data` + - Saves MP4 videos from a dedicated camera without creating a training dataset. +``` + ## Usage Example ```python diff --git a/docs/source/overview/gym/env.md b/docs/source/overview/gym/env.md index 19835b01b..95c067f95 100644 --- a/docs/source/overview/gym/env.md +++ b/docs/source/overview/gym/env.md @@ -109,32 +109,127 @@ The {class}`~envs.EmbodiedEnvCfg` class exposes the following additional paramet ```python from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.envs.managers import ( + ActionTermCfg, + DatasetFunctorCfg, + EventCfg, + ObservationCfg, + RewardCfg, + SceneEntityCfg, +) from embodichain.utils import configclass +@configclass +class MyEventCfg: + randomize_object_pose: EventCfg = EventCfg( + func="randomize_rigid_object_pose", + mode="reset", + params={ + "entity_cfg": SceneEntityCfg(uid="cube"), + "position_range": [[-0.08, -0.08, 0.0], [0.08, 0.08, 0.0]], + "relative_position": True, + }, + ) + record_debug_video: EventCfg = EventCfg( + func="record_camera_data", + mode="interval", + interval_step=1, + params={ + "name": "overview_cam", + "eye": (1.0, 0.0, 1.2), + "target": (0.0, 0.0, 0.4), + "save_path": "./outputs/videos", + }, + ) + + +@configclass +class MyObservationCfg: + object_pose: ObservationCfg = ObservationCfg( + func="get_object_pose", + mode="add", + name="object/cube/pose", + params={ + "entity_cfg": SceneEntityCfg(uid="cube"), + "to_matrix": False, + }, + ) + normalized_qpos: ObservationCfg = ObservationCfg( + func="normalize_robot_joint_data", + mode="modify", + name="robot/qpos", + params={}, + ) + + +@configclass +class MyRewardCfg: + approach_target: RewardCfg = RewardCfg( + func="distance_to_target", + weight=1.0, + params={ + "entity_cfg": SceneEntityCfg(uid="cube"), + "target_pose_key": "goal_pose", + "exponential": True, + "sigma": 0.2, + }, + ) + success_bonus: RewardCfg = RewardCfg( + func="success_reward", + weight=10.0, + params={}, + ) + + +@configclass +class MyActionCfg: + delta_qpos: ActionTermCfg = ActionTermCfg( + func="DeltaQposTerm", + params={"scale": 0.1}, + ) + + +@configclass +class MyDatasetCfg: + lerobot: DatasetFunctorCfg = DatasetFunctorCfg( + func="LeRobotRecorder", + params={ + "save_path": "./outputs/datasets/my_task", + "robot_meta": {"robot_type": "my_robot", "control_freq": 25}, + "instruction": {"lang": "move the cube to the goal"}, + "use_videos": False, + }, + ) + + @configclass class MyTaskEnvCfg(EmbodiedEnvCfg): - # 1. Define Scene Components - robot = ... # Robot configuration - sensor = [...] # List of sensors (e.g., Cameras) - light = ... # Lighting configuration - - # 2. Define Objects - rigid_object = [...] # Dynamic objects (e.g., tools, debris) - rigid_object_group = [...] # Object groups (efficient for many similar objects) - articulation = [...] # Articulated objects (e.g., cabinets) - - # 3. Define Managers - events = ... # Event settings (Randomization, etc.) - observations = ... # Custom observation spec - dataset = ... # Data collection settings - - # 4. Action Manager (for RL tasks) - actions = ... # Action preprocessing (e.g., DeltaQposTerm with scale) - extensions = { # Task-specific parameters (e.g., success_threshold) + # Scene assets are task-specific and usually come from existing robot/object cfgs. + robot = ... + sensor = [...] + rigid_object = [...] + + # Manager configs plug into the shared environment lifecycle. + events = MyEventCfg() + observations = MyObservationCfg() + rewards = MyRewardCfg() + actions = MyActionCfg() + dataset = MyDatasetCfg() + + init_rollout_buffer = True + extensions = { "success_threshold": 0.1, } ``` +This example shows the typical division of responsibilities: + +- ``events`` mutate or record the scene during ``startup``, ``reset``, or ``interval`` phases. +- ``observations`` expose task state to policies or data pipelines. +- ``rewards`` shape RL behavior. +- ``actions`` define how policy outputs map to robot control commands. +- ``dataset`` controls structured episode export, independent from debug-video recording. + ## Manager Systems The manager systems in {class}`~envs.EmbodiedEnv` provide modular, configuration-driven functionality for handling complex simulation behaviors. Each manager uses a **functor-based** architecture, allowing you to compose behaviors through configuration without modifying environment code. Functors are reusable functions or classes (inheriting from {class}`~envs.managers.Functor`) that operate on the environment state, configured through {class}`~envs.managers.cfg.FunctorCfg`. @@ -185,6 +280,10 @@ The manager operates in a single mode ``"save"`` which handles both recording an * ``instruction``: Task instruction dictionary. * ``use_videos``: Whether to save video recordings of episodes. +```{note} +The Dataset Manager handles structured training data. If you want debug or demo videos from a dedicated camera, use {class}`~envs.managers.record.record_camera_data` documented in {doc}`event_functors`. +``` + The dataset manager is called automatically during {meth}`~envs.Env.step()`, ensuring all observation-action pairs are recorded without additional user code. ## Reinforcement Learning Environment diff --git a/docs/source/overview/gym/event_functors.md b/docs/source/overview/gym/event_functors.md index fb110ef3f..fb844da6d 100644 --- a/docs/source/overview/gym/event_functors.md +++ b/docs/source/overview/gym/event_functors.md @@ -64,7 +64,15 @@ This page lists all available event functors that can be used with the Event Man "base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]]}} ``` * - {func}`~randomization.visual.randomize_light` - - Vary light position, color, and intensity within specified ranges. + - Vary light position, color, intensity, and direction within specified ranges. + + .. note:: + ``position_range`` is ignored for global scene lights (``"sun"``, ``"direction"``). + Use ``direction_range`` instead for these types. + + .. note:: + ``direction_range`` is only applicable for directional light types + (``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, ``"mesh"``). ```json {"func": "randomize_light", @@ -72,7 +80,8 @@ This page lists all available event functors that can be used with the Event Man "params": {"entity_cfg": {"uid": "light_1"}, "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], - "intensity_range": [50.0, 100.0]}} + "intensity_range": [50.0, 100.0], + "direction_range": [[-0.1, -0.1, -0.1], [0.1, 0.1, 0.1]]}} ``` * - {func}`~randomization.visual.randomize_emission_light` - Randomize global emission light color and intensity. Applies the same emission light properties across all environments. @@ -195,6 +204,15 @@ This page lists all available event functors that can be used with the Event Man "params": {"grid_origin": [0.5, -0.3], "grid_size": [0.4, 0.6], "grid_res": [4, 6], "entity_cfg": {"uid": "objects"}}} ``` +* - {class}`~randomization.spatial.randomize_anchor_height` + - Randomize the height of an anchor object and shift other objects by the same delta. Implemented as a Functor class. Samples a per-environment height delta (uniform range or discrete candidates), moves the anchor object relative to its configured initial position, and adds the same delta to the Z component of every other included object while preserving XY and rotation. + + ```json + {"func": "randomize_anchor_height", "mode": "reset", + "params": {"anchor_uid": "table", + "height_delta_range": [[-0.05], [0.05]], + "exclude_uids": ["floor"]}} + ``` ``` ## Geometry Randomization @@ -294,6 +312,119 @@ This page lists all available event functors that can be used with the Event Man ``` ``` +## Recording + +```{list-table} Recording Functors +:header-rows: 1 +:widths: 25 75 + +* - Functor Name + - Description +* - {class}`~record.record_camera_data` + - Record RGB frames from a dedicated camera and save them as an MP4 video when an episode resets. This is useful for debugging, qualitative evaluation, and demo capture. The functor creates its own camera from the configured pose and intrinsics, captures frames during ``interval`` execution, and writes one video per episode. + + ```json + {"func": "record_camera_data", "mode": "interval", "interval_step": 1, + "params": {"name": "debug_cam", + "resolution": [640, 480], + "eye": [1.0, 0.0, 1.2], + "target": [0.0, 0.0, 0.5], + "up": [0.0, 0.0, 1.0], + "save_path": "./outputs/videos"}} + ``` +* - {class}`~record.record_camera_data_async` + - Record RGB frames for several environments independently, then merge them into a single tiled MP4 once each tracked environment finishes an episode. This variant is useful when you want side-by-side qualitative comparison across vectorized environments. + + ```json + {"func": "record_camera_data_async", "mode": "interval", "interval_step": 1, + "params": {"name": "overview_async_cam", + "resolution": [640, 480], + "eye": [1.0, 0.0, 1.2], + "target": [0.0, 0.0, 0.5], + "up": [0.0, 0.0, 1.0], + "save_path": "./outputs/videos"}} + ``` +``` + +### record_camera_data + +The ``record_camera_data`` functor lives under the record manager module, but it is configured through the event pipeline because it runs periodically during stepping and flushes video output at episode reset. + +```{note} +This is separate from {doc}`dataset_functors`. Use dataset functors when you need training data such as observations and actions. Use ``record_camera_data`` when you need human-viewable videos for debugging or demos. +``` + +#### Typical Usage + +- Set ``mode="interval"`` so the camera captures frames during stepping. +- Use ``interval_step=1`` to record every environment step, or increase it to reduce file size. +- Choose a fixed third-person camera pose with ``eye``, ``target``, and ``up``. +- Save output videos with ``save_path``. + +#### Parameters + +```{list-table} record_camera_data Parameters +:header-rows: 1 +:widths: 30 70 + +* - Parameter + - Description +* - ``name`` + - Camera name used for the internally created sensor and output file naming. Default: ``"default"``. +* - ``resolution`` + - Output image size as ``[width, height]``. Default: ``[640, 480]``. +* - ``eye`` + - Camera position in world coordinates. Default: ``[0, 0, 2]``. +* - ``target`` + - Look-at target in world coordinates. Default: ``[0, 0, 0]``. +* - ``up`` + - Up vector for the camera pose. Default: ``[0, 0, 1]``. +* - ``intrinsics`` + - Camera intrinsics as ``[fx, fy, cx, cy]``. Defaults to values derived from the configured resolution. +* - ``max_env_num`` + - Maximum number of environments to tile into one frame when recording vectorized environments. +* - ``save_path`` + - Output directory for generated MP4 files. Default: ``./outputs/videos``. +``` + +#### Behavior Notes + +- In vectorized environments, frames from multiple environments are tiled into a single composite image before video encoding. +- Videos are flushed during episode initialization for environments that are resetting, not by the Dataset Manager. +- If the process exits without another reset, the final episode may not be flushed to disk. +- This functor captures RGB imagery for visualization. It does not record actions, rewards, or structured training data. + +### record_camera_data_async + +The ``record_camera_data_async`` functor is the multi-environment variant of ``record_camera_data``. It tracks a small set of environments independently, waits until each one has completed an episode, and then writes a single tiled video that combines them. + +#### Typical Usage + +- Use this variant when you want one comparison video spanning multiple vectorized environments. +- Configure it with ``mode="interval"`` just like ``record_camera_data``. +- Keep the number of tracked environments modest; the current implementation records up to four environments. +- Use the same camera pose parameters as the single-environment recorder. + +#### Parameters + +The async recorder accepts the same parameters as ``record_camera_data``: + +- ``name`` +- ``resolution`` +- ``eye`` +- ``target`` +- ``up`` +- ``intrinsics`` +- ``max_env_num`` +- ``save_path`` + +#### Behavior Notes + +- Frames are buffered separately for each tracked environment and merged later into one tiled video. +- The current implementation tracks at most four environments, even if ``num_envs`` is larger. +- Output is written only after all tracked environments have completed an episode, so video generation may lag behind individual resets. +- This variant is meant for qualitative inspection and comparison, not structured dataset export. + ## Usage Example ```python diff --git a/docs/source/overview/gym/index.rst b/docs/source/overview/gym/index.rst index 61deac36f..439445cfd 100644 --- a/docs/source/overview/gym/index.rst +++ b/docs/source/overview/gym/index.rst @@ -165,17 +165,37 @@ For most new tasks, start from Choosing Where to Start ----------------------- +.. list-table:: + :header-rows: 1 + :widths: 34 66 + + * - If you want to... + - Start here + * - Understand the full ``EmbodiedEnv`` configuration surface + - :doc:`env` + * - Connect policy outputs to robot control terms + - :doc:`action_functors` + * - Randomize scenes, place objects, or capture debug/demo videos + - :doc:`event_functors` + * - Add task observations or transform existing observation entries + - :doc:`observation_functors` + * - Compose RL reward terms + - :doc:`reward_functors` + * - Save structured demonstrations or offline training datasets + - :doc:`dataset_functors` + - Start with :doc:`env` for the full :class:`~embodichain.lab.gym.envs.embodied_env.EmbodiedEnv` configuration and custom task guide. - Use :doc:`action_functors` when connecting policy outputs to robot control. - Use :doc:`event_functors` for reset randomization, visual randomization, and - scene perturbations. + scene perturbations. This page also covers runtime camera-video capture via + ``record_camera_data``. - Use :doc:`observation_functors` to add task observations without changing base environment code. - Use :doc:`reward_functors` when composing RL reward terms. - Use :doc:`dataset_functors` when recording demonstrations or exporting - datasets. + datasets. Use this page for structured episode data, not debug video capture. Documentation Quality Notes --------------------------- diff --git a/docs/source/overview/gym/observation_functors.md b/docs/source/overview/gym/observation_functors.md index bb67cce6e..3b6ead061 100644 --- a/docs/source/overview/gym/observation_functors.md +++ b/docs/source/overview/gym/observation_functors.md @@ -5,6 +5,13 @@ This page lists all available observation functors that can be used with the Observation Manager. Observation functors are configured using {class}`~cfg.ObservationCfg` and can operate in two modes: ``modify`` (update existing observations) or ``add`` (add new observations). +## Quick Reference + +- Use ``mode="add"`` to create a new observation entry in the observation dictionary. +- Use ``mode="modify"`` to update an existing observation entry in place. +- Most returned tensors are batched over ``num_envs`` as the leading dimension. +- Use hierarchical names such as ``"object/cup/pose"`` to keep custom observations organized. + ````{tip} **Using an AI coding agent?** Use the **`/add-functor`** skill to scaffold a new observation functor with the correct signature (`env, obs, entity_cfg, ...`), module placement in `observations.py`, and `__all__` export. Use **`/add-test`** to generate mock-based tests. ```` diff --git a/docs/source/overview/gym/reward_functors.md b/docs/source/overview/gym/reward_functors.md index fb91cbf0f..45fc1081a 100644 --- a/docs/source/overview/gym/reward_functors.md +++ b/docs/source/overview/gym/reward_functors.md @@ -5,6 +5,13 @@ This page lists all available reward functors that can be used with the Reward Manager. Reward functors are configured using {class}`~cfg.RewardCfg` and return scalar reward tensors that are weighted and summed to form the total environment reward. +## Quick Reference + +- Use this page when composing scalar reward terms for RL tasks. +- Reward functors are configured with {class}`~cfg.RewardCfg`. +- Each reward functor returns a tensor of shape ``(num_envs,)`` before weighting. +- Final rewards are the weighted sum of all configured reward terms. + ````{tip} **Using an AI coding agent?** Use the **`/add-functor`** skill to scaffold a new reward functor with the correct signature (`env, obs, action, info, ...`), module placement in `rewards.py`, and `__all__` export. Use **`/add-test`** to generate mock-based tests. ```` diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index 01320709e..cd11dd5d7 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -7,6 +7,10 @@ The following actions are available out of the box: +```{note} +The built-in atomic actions currently support gripper-based manipulation only. Dexterous-hand manipulation is not supported yet. +``` + | Action | Arm | Target type | Motion phases | Demo | |---|---|---|---|---| | `MoveEndEffector` | Single | `EndEffectorPoseTarget` — EEF pose | Move end-effector to pose | MoveEndEffector | @@ -14,6 +18,9 @@ The following actions are available out of the box: | `PickUp` | Single | `GraspTarget` — object semantics | Approach → close gripper → lift | PickUp | | `MoveHeldObject` | Single | `HeldObjectPoseTarget` — held-object pose | Move held object while keeping gripper closed | MoveHeldObject | | `Place` | Single | `EndEffectorPoseTarget` — EEF release pose | Lower → open gripper → retract | Place | +| `Press` | Single | `EndEffectorPoseTarget` — EEF press pose | Close gripper → press down → return | Press | +| `CoordinatedPickment` | Dual | `CoordinatedPickmentTarget` — shared-object pose | Approach both ends → close both grippers → lift → move object | CoordinatedPickment | +| `CoordinatedPlacement` | Dual | `CoordinatedPlacementTarget` — two held-object poses | Move support object → align placing object → release placing hand → retreat | CoordinatedPlacement | --- @@ -117,7 +124,102 @@ down to the target pose. On success, the returned `WorldState` clears `held_obje | `hand_interp_steps` | `5` | Waypoints for the gripper open phase | | `sample_interval` | `80` | Total waypoints across all three phases | -**Target:** `EndEffectorPoseTarget(xpos=...)` — the EEF pose at release, a `torch.Tensor` of shape -`(4, 4)`, `(n_envs, 4, 4)` or `(n_envs, n_waypoint, 4, 4)`. +**Target:** `EndEffectorPoseTarget(xpos=..., tcp_symmetry="none")` — the EEF pose at +release, a `torch.Tensor` of shape `(4, 4)`, `(n_envs, 4, 4)` or +`(n_envs, n_waypoint, 4, 4)`. Keep the default +`tcp_symmetry="none"` when the TCP orientation is strict. Use +`tcp_symmetry="z_roll_180"` only when releasing with TCP x/y flipped is physically +equivalent; `Place` then chooses the closer TCP z-roll 180 variant from +`WorldState.last_qpos` and applies that same variant across all release waypoints. ![Place demo](../../../_static/atomic_actions/place.gif) + +--- + +## `Press` + +Three-phase contact motion: *close gripper → press down → return*. This is useful +for button-like or contact-based interactions where the end-effector should reach a +target pose and then return to the pre-press arm pose. + +`Press` does not create or clear `WorldState.held_object`; it preserves the state +threaded into it. + +| Config field | Default | Description | +|---|---|---| +| `hand_close_qpos` | `None` | **Required.** Gripper closed joint positions | +| `hand_control_part` | `"hand"` | Robot control part for the gripper | +| `hand_interp_steps` | `5` | Waypoints for the gripper close phase | +| `sample_interval` | `80` | Total waypoints across all three phases | + +**Target:** `EndEffectorPoseTarget(xpos=...)` — the EEF pose to press, a `torch.Tensor` +of shape `(4, 4)` or `(n_envs, 4, 4)`. + +![Press demo](../../../_static/atomic_actions/press.gif) + +--- + +## `CoordinatedPickment` + +Dual-arm grasp motion for one shared object. Both arms move to object-relative +grasp poses, close both grippers, lift the object, and move it to an object pose +while keeping both grippers closed. On success, the returned `WorldState` carries +`coordinated_held_object` (`CoordinatedHeldObjectState`) and leaves +`held_object` as `None`. + +| Config field | Default | Description | +|---|---|---| +| `control_part` | `"dual_arm"` | Combined arm control part | +| `left_arm_control_part` / `right_arm_control_part` | `"left_arm"` / `"right_arm"` | Arm control parts for each grasp | +| `left_hand_control_part` / `right_hand_control_part` | `"left_hand"` / `"right_hand"` | Hand control parts for each gripper | +| `pre_grasp_distance` | `0.10` | Distance to back away from each grasp TCP | +| `lift_height` | `0.08` | World-Z lift distance before moving to the target pose | +| `object_motion_keyframes` | `6` | Sparse object-pose IK keyframes for synchronized motion | +| `sample_interval` | `120` | Total waypoints across all phases | + +**Target:** `CoordinatedPickmentTarget(...)` with a target object pose, object +semantics, and left/right object-to-EEF transforms. + +**Tutorial:** `scripts/tutorials/atomic_action/coordinated_pickment.py` + +![CoordinatedPickment demo](../../../_static/atomic_actions/coordinated_pickment.gif) + +--- + +## `CoordinatedPlacement` + +Dual-arm object-centric placement. The support arm moves its held object to a lower +target pose and keeps its gripper closed. The placing arm moves its held object to +the aligned upper target pose, optionally opens the placing hand, then lifts away. + +`CoordinatedPlacement` is intentionally explicit about dual-arm state: the target +contains both `placing_held_object` and `support_held_object`. This avoids relying +on the engine's single `WorldState.held_object` slot to infer two simultaneously +held objects. + +| Config field | Default | Description | +|---|---|---| +| `control_part` | `"dual_arm"` | Robot control part containing both arms | +| `placing_arm_control_part` | `"left_arm"` | Arm that releases the placed object | +| `support_arm_control_part` | `"right_arm"` | Arm that keeps holding the support object | +| `placing_hand_control_part` | `"left_hand"` | Placing gripper control part | +| `support_hand_control_part` | `"right_hand"` | Support gripper control part | +| `placing_hand_open_qpos` | `None` | **Required.** Placing gripper open joint positions | +| `placing_hand_close_qpos` | `None` | **Required.** Placing gripper closed joint positions | +| `support_hand_close_qpos` | `None` | **Required.** Support gripper closed joint positions | +| `release` | `True` | Whether to open the placing gripper | +| `placing_height_offset` | `0.0` | World-Z offset applied to the placing object target pose | +| `support_height_offset` | `0.0` | World-Z offset applied to the support object target pose | +| `lift_height` | `0.08` | Placing-arm lift distance after release (m) | +| `hand_interp_steps` | `10` | Waypoints for placing-hand release | +| `hold_steps` | `4` | Alignment hold waypoints before release | +| `retreat_steps` | `16` | Placing-arm retreat waypoints | +| `sample_interval` | `100` | Total waypoints across all phases | + +**Target:** `CoordinatedPlacementTarget(...)` with placing/support object target +poses plus the corresponding `HeldObjectState` values. On success, the returned +`WorldState.held_object` is the support object's held state. + +**Tutorial:** `scripts/tutorials/atomic_action/coordinated_placement.py` + +![CoordinatedPlacement demo](../../../_static/atomic_actions/coordinated_placement.gif) diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index 119021bda..f1f120a71 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -5,6 +5,10 @@ Atomic actions are the building blocks for automated robot motion generation. Each action encapsulates a complete, self-contained motion primitive — such as picking up an object or moving to a pose — that can be chained together to form complex manipulation workflows. +```{note} +Atomic actions currently support gripper-based manipulation only. Dexterous-hand manipulation is not supported yet. +``` + ## Design Overview The module is organized into three layers: @@ -16,7 +20,8 @@ AtomicActionEngine ← orchestrates a sequence of (name, typed_target) │ │ │ └── MotionGenerator ← low-level trajectory planner (IK + trajectory optimization) │ - └── WorldState ← threaded action-to-action (last_qpos + held_object) + └── WorldState ← threaded action-to-action + (last_qpos + held_object/coordinated_held_object) ``` Each action receives a typed target and a `WorldState`, runs its planning pipeline, and @@ -25,11 +30,13 @@ the `next_state` of each action as the input state of the next, then concatenate trajectories into one contiguous sequence: ``` -GraspTarget(semantics) ──► AtomicAction.execute(target, state) +GraspTarget(semantics, grasp_xpos=None) ──► AtomicAction.execute(target, state) EndEffectorPoseTarget(xpos) │ JointPositionTarget(qpos) ├─ IK solve when pose-based NamedJointPositionTarget(name) ├─ Motion plan / interpolation HeldObjectPoseTarget(pose) └─ Gripper interpolation when needed +CoordinatedPickmentTarget(...) │ +CoordinatedPlacementTarget(...) │ │ ActionResult (success, full-DoF traj, next_state) @@ -57,11 +64,13 @@ and every action declares the target type, or tuple of target types, it accepts | Target | Constructor | Used by | |---|---|---| -| `EndEffectorPoseTarget` | `EndEffectorPoseTarget(xpos)` | `MoveEndEffector`, `Place` | +| `EndEffectorPoseTarget` | `EndEffectorPoseTarget(xpos, tcp_symmetry="none")` | `MoveEndEffector`, `Place`, `Press` | | `JointPositionTarget` | `JointPositionTarget(qpos)` | `MoveJoints` | | `NamedJointPositionTarget` | `NamedJointPositionTarget(name)` | `MoveJoints` | -| `GraspTarget` | `GraspTarget(semantics)` | `PickUp` | +| `GraspTarget` | `GraspTarget(semantics, grasp_xpos=None)` | `PickUp` | | `HeldObjectPoseTarget` | `HeldObjectPoseTarget(object_target_pose)` | `MoveHeldObject` | +| `CoordinatedPickmentTarget` | `CoordinatedPickmentTarget(...)` | `CoordinatedPickment` | +| `CoordinatedPlacementTarget` | `CoordinatedPlacementTarget(...)` | `CoordinatedPlacement` | `Target` is the union of these typed target dataclasses. @@ -94,22 +103,28 @@ action's `TargetType` before calling `execute`: | Target | Holds | Accepted by | |---|---|---| -| `EndEffectorPoseTarget(xpos)` | EEF pose tensor `(4,4)`, `(n_envs,4,4)` or `(n_envs, n_waypoint, 4, 4)` | `MoveEndEffector`, `Place` | +| `EndEffectorPoseTarget(xpos, tcp_symmetry="none")` | EEF pose tensor `(4,4)`, `(n_envs,4,4)` or `(n_envs, n_waypoint, 4, 4)`; `Place` may opt into TCP z-roll 180 equivalence | `MoveEndEffector`, `Place`, `Press` | | `JointPositionTarget(qpos)` | Control-part qpos tensor `(control_dof,)`, `(n_envs, control_dof)` or `(n_envs, n_waypoint, control_dof)` | `MoveJoints` | | `NamedJointPositionTarget(name)` | Name resolved from `MoveJointsCfg.named_joint_positions` | `MoveJoints` | -| `GraspTarget(semantics)` | `ObjectSemantics` (affordance + entity) | `PickUp` | +| `GraspTarget(semantics, grasp_xpos=None)` | `ObjectSemantics` plus an optional preselected TCP grasp pose | `PickUp` | | `HeldObjectPoseTarget(object_target_pose)` | Desired held-object pose tensor | `MoveHeldObject` | +| `CoordinatedPickmentTarget(...)` | Shared object semantics plus left/right grasp transforms and target object pose | `CoordinatedPickment` | +| `CoordinatedPlacementTarget(...)` | Two held-object states plus object-centric placing/support target poses | `CoordinatedPlacement` | -`WorldState` is threaded between actions and carries the robot's `last_qpos` plus an optional -`held_object: HeldObjectState`. The built-in actions update it as follows: +`WorldState` is threaded between actions and carries the robot's `last_qpos` plus optional +`held_object: HeldObjectState` and `coordinated_held_object: CoordinatedHeldObjectState`. +The built-in actions update it as follows: | Action | Effect on `held_object` | |---|---| | `PickUp` | Populates it (computed object-to-EEF transform) | | `MoveHeldObject` | Requires it; preserves it unchanged | | `Place` | Clears it to `None` | +| `CoordinatedPickment` | Leaves `held_object` as `None` and populates `coordinated_held_object` | +| `CoordinatedPlacement` | Returns the support arm's `HeldObjectState`; the placing object is released | | `MoveEndEffector` | Leaves it unchanged | | `MoveJoints` | Leaves it unchanged | +| `Press` | Leaves it unchanged | If a step fails, `run()` returns `success=False` with the partial trajectory concatenated up to (but not including) the failed step, and the `WorldState` going into that step. @@ -136,6 +151,8 @@ from embodichain.lab.sim.atomic_actions import ( MoveHeldObjectCfg, Place, PlaceCfg, + Press, + PressCfg, MoveEndEffector, MoveEndEffectorCfg, ) @@ -159,6 +176,11 @@ move_held_object_cfg = MoveHeldObjectCfg( hand_close_qpos=torch.tensor([0.025, 0.025]), ) move_cfg = MoveEndEffectorCfg(control_part="arm") +press_cfg = PressCfg( + control_part="arm", + hand_control_part="hand", + hand_close_qpos=torch.tensor([0.025, 0.025]), +) move_joints_cfg = MoveJointsCfg( control_part="arm", named_joint_positions={"home": torch.zeros(6)}, @@ -170,6 +192,7 @@ engine.register(PickUp(motion_gen, cfg=pickup_cfg)) engine.register(MoveHeldObject(motion_gen, cfg=move_held_object_cfg)) engine.register(Place(motion_gen, cfg=place_cfg)) engine.register(MoveEndEffector(motion_gen, cfg=move_cfg)) +engine.register(Press(motion_gen, cfg=press_cfg)) engine.register(MoveJoints(motion_gen, cfg=move_joints_cfg)) # 3. Describe the object to pick @@ -203,6 +226,10 @@ is_success, traj, final_state = engine.run( You can add any motion primitive by subclassing `AtomicAction`, composing a `TrajectoryBuilder` for the shared planning math, and registering an instance with the engine. +Built-in primitives live one action per module under +`embodichain/lab/sim/atomic_actions/primitives/`, while +`embodichain.lab.sim.atomic_actions` remains the public import surface and +`embodichain.lab.sim.atomic_actions.actions` stays as a compatibility re-export. ### Step 1 — Define the config @@ -297,6 +324,9 @@ builtin_actions - `scripts/tutorials/atomic_action/pickup.py` - `scripts/tutorials/atomic_action/move_held_object.py` - `scripts/tutorials/atomic_action/place.py` + - `scripts/tutorials/atomic_action/press.py` + - `scripts/tutorials/atomic_action/coordinated_pickment.py` + - `scripts/tutorials/atomic_action/coordinated_placement.py` Run a demo in headless CPU mode with `--auto_play --headless --device cpu` to record an MP4 under `outputs/videos`. For example: diff --git a/docs/source/overview/sim/sim_articulation.md b/docs/source/overview/sim/sim_articulation.md index 0950a6fba..812de8b45 100644 --- a/docs/source/overview/sim/sim_articulation.md +++ b/docs/source/overview/sim/sim_articulation.md @@ -16,6 +16,7 @@ Articulations are configured using the {class}`~cfg.ArticulationCfg` dataclass. | `fix_base` | `bool` | `True` | Whether to fix the base of the articulation. | | `use_usd_properties` | `bool` | `False` | If True, use physical properties from USD file; if False, override with config values. Only effective for usd files. | | `init_qpos` | `List[float]` | `None` | Initial joint positions. | +| `qpos_limits` | `Tensor` / `Dict[str, List[float]]` | `None` | Override joint position limits. Replaces asset limits and may either tighten or expand the range. | | `body_scale` | `List[float]` | `[1.0, 1.0, 1.0]` | Scaling factors for the articulation links. | | `disable_self_collisions` | `bool` | `True` | Whether to disable self-collisions. | | `drive_props` | `JointDrivePropertiesCfg` | `...` | Default drive properties. | @@ -68,6 +69,49 @@ The `drive_props` parameter controls the joint physics behavior. It is defined u | `armature` | `float` / `Dict` | `0.0` | Joint armature added to joint-space inertia ($kg$ for prismatic, $kg \cdot m^2$ for revolute). | | `drive_type` | `str` | `"none"` | Drive mode: `"force"`(driven by a force), `"acceleration"`(driven by an acceleration) or `none`(no force). | +### Joint Position Limits + +Use `qpos_limits` to override the limits defined in the asset file. This is the +articulation's effective physical limit in simulation, so it is also the range +used when `set_qpos(...)` clamps requested joint positions. + +```python +from embodichain.lab.sim.cfg import ArticulationCfg +import torch + +# Override specific joints by name or regex at initialization. +art_cfg = ArticulationCfg( + fpath="assets/robots/franka/franka.urdf", + qpos_limits={ + "panda_joint1": [-2.5, 2.5], + "panda_joint[2-4]": [-1.5, 1.5], + }, +) + +# Or provide a (dof, 2) tensor/array applied to all environments. +dof = 7 +art_cfg = ArticulationCfg( + fpath="assets/robots/franka/franka.urdf", + qpos_limits=torch.tensor([[-2.0, 2.0]] * dof), +) +``` + +You can also change the active limits at runtime: + +```python +# Tighten limits for joints 0 and 1 on all environments. +new_limits = torch.tensor([[-0.1, 0.1], [-0.2, 0.2]], device=device) +articulation.set_qpos_limits(new_limits, joint_ids=[0, 1]) + +# Query the effective limits. +effective_limits = articulation.get_qpos_limits(joint_ids=[0, 1]) +``` + +If you need solver-only planning limits that are tighter than the physical +articulation limits, configure them on the solver with +`SolverCfg.user_qpos_limits`; that behavior lives in the solver layer, not the +articulation layer. + ### Setup & Initialization ```python diff --git a/docs/source/overview/sim/sim_assets.md b/docs/source/overview/sim/sim_assets.md index ab49dfbbb..fca7351f4 100644 --- a/docs/source/overview/sim/sim_assets.md +++ b/docs/source/overview/sim/sim_assets.md @@ -112,14 +112,27 @@ For Rigid Object tutorial, please refer to the [Create Scene](https://dexforce.g ### Lights -Configured via `LightCfg`. +Configured via `LightCfg`. Supports six light types matching the dexsim rendering backend. | Parameter | Type | Default | Description | | :--- | :--- | :--- | :--- | -| `light_type` | `Literal` | `"point"` | Type of light (currently only "point"). | -| `color` | `tuple` | `(1.0, 1.0, 1.0)` | RGB color. | -| `intensity` | `float` | `50.0` | Intensity in watts/m^2. | -| `radius` | `float` | `1e2` | Falloff radius. | +| `light_type` | `Literal` | `"point"` | Light type: ``"point"``, ``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, or ``"mesh"``. | +| `color` | `tuple` | `(1.0, 1.0, 1.0)` | RGB color of the light source. | +| `intensity` | `float` | `30.0` | Intensity of the light source in watts/m^2. | +| `enable_shadow` | `bool` | `True` | Whether the light casts shadows. | +| `radius` | `float` | `10.0` | Falloff radius (only for ``"point"``). | +| `direction` | `tuple` | `(0.0, 0.0, -1.0)` | Direction vector for directional types (``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, ``"mesh"``). | +| `spot_angle_inner` | `float` | `30.0` | Inner cone angle in degrees (only for ``"spot"``). | +| `spot_angle_outer` | `float` | `45.0` | Outer cone angle in degrees (only for ``"spot"``). | +| `rect_width` | `float` | `1.0` | Width of rectangular area light (only for ``"rect"``). | +| `rect_height` | `float` | `1.0` | Height of rectangular area light (only for ``"rect"``). | +| `mesh_path` | `str` | `""` | Asset path for mesh-based emissive lights (only for ``"mesh"``). | + +.. attention:: + ``"sun"`` and ``"direction"`` are **global scene lights** (infinite-distance directional light + sources). They are created as a single instance on the root environment, not batched per + environment. Use :meth:`Light.set_direction` instead of :meth:`Light.set_local_pose` for + these types. ```{toctree} diff --git a/docs/source/overview/sim/sim_robot.md b/docs/source/overview/sim/sim_robot.md index e184c28cf..a683a3cab 100644 --- a/docs/source/overview/sim/sim_robot.md +++ b/docs/source/overview/sim/sim_robot.md @@ -104,5 +104,19 @@ The Robot class overrides standard Articulation methods to support the name argu | `set_qf(..., name="part")` | Set joint efforts for a specific part. | | `get_qpos(name="part")` | Get joint positions of a specific part. | | `get_qvel(name="part")` | Get joint velocities of a specific part. | +| `set_qpos_limits(..., name="part")` | Set physical joint position limits for a specific part. | +| `get_qpos_limits(name="part")` | Get physical joint position limits of a specific part. | + +```python +# Tighten the left arm physical workspace +left_arm_limits = robot.get_qpos_limits(name="left_arm") +left_arm_limits[..., 0] += 0.05 # tighten lower bound +left_arm_limits[..., 1] -= 0.05 # tighten upper bound +robot.set_qpos_limits(left_arm_limits, name="left_arm") +``` + +For solver-only planning limits, configure `SolverCfg.user_qpos_limits` on the +robot's solver config. Robot solver synchronization will intersect those solver +limits with the robot's current effective `qpos_limits`. For more API details, refer to the {class}`~objects.Robot` documentation. diff --git a/docs/source/overview/sim/solvers/srs_solver.md b/docs/source/overview/sim/solvers/srs_solver.md index 2b26ee6d0..0cd7350dc 100644 --- a/docs/source/overview/sim/solvers/srs_solver.md +++ b/docs/source/overview/sim/solvers/srs_solver.md @@ -51,7 +51,7 @@ cfg = SRSSolverCfg( end_link_name="left_ee", root_link_name="left_arm_base", dh_params=arm_params.dh_params, - user_qpos_limit=arm_params.qpos_limits, + user_qpos_limits=arm_params.qpos_limits, T_e_oe=arm_params.T_e_oe, T_b_ob=arm_params.T_b_ob, link_lengths=arm_params.link_lengths, diff --git a/docs/source/tutorial/atomic_actions.rst b/docs/source/tutorial/atomic_actions.rst index b65f56792..5644792b8 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -5,7 +5,7 @@ Atomic Actions EmbodiChain's **atomic action** layer provides a high-level, composable interface for common manipulation primitives such as *move end-effector*, *move joints*, *pick up*, -*move held object*, and *place*. Each action +*move held object*, *place*, and *press*. Each action encapsulates the full planning pipeline — grasp-pose estimation, IK, trajectory generation, and gripper interpolation — behind a single ``execute(target, state)`` call, making it straightforward to chain multiple actions together into complex robot behaviours. @@ -18,13 +18,14 @@ Key Features ``GraspTarget`` (wrapping an ``ObjectSemantics``), or ``HeldObjectPoseTarget``. The engine checks each step's target against the action's declared ``TargetType`` before running. - **Built-in primitives** — ``MoveEndEffector``, ``MoveJoints``, ``PickUp``, ``MoveHeldObject``, - and ``Place`` + ``Place``, ``Press``, ``CoordinatedPickment``, and ``CoordinatedPlacement`` cover the most common tabletop manipulation workflows out of the box. See :doc:`/overview/sim/atomic_actions/index` for configs and target types. - **Extensible registry** — custom action *classes* can be registered globally with ``register_action``; action *instances* are registered per-engine under a name. - **Engine orchestration** — ``AtomicActionEngine.run(steps, state)`` sequences named - ``(name, typed_target)`` steps, threads a ``WorldState`` (``last_qpos`` + ``held_object``) + ``(name, typed_target)`` steps, threads a ``WorldState`` (``last_qpos`` + + ``held_object`` / ``coordinated_held_object``) from one action into the next, and returns a single concatenated full-DOF trajectory ready to replay in the simulator. @@ -34,7 +35,7 @@ For the full design overview, architecture diagram, and extension guide see The Code -------- -Focused demo scripts are available for the five built-in primitives in the +Focused demo scripts are available for the built-in primitives in the ``scripts/tutorials/atomic_action`` directory: - ``move_end_effector.py`` @@ -42,6 +43,9 @@ Focused demo scripts are available for the five built-in primitives in the - ``pickup.py`` - ``move_held_object.py`` - ``place.py`` +- ``press.py`` +- ``coordinated_pickment.py`` +- ``coordinated_placement.py`` Each script supports interactive inspection by default. Add ``--auto_play`` to skip keyboard prompts, and combine it with ``--headless --device cpu`` to record an MP4 under @@ -54,6 +58,15 @@ keyboard prompts, and combine it with ``--headless --device cpu`` to record an M python scripts/tutorials/atomic_action/pickup.py --headless --auto_play --device cpu python scripts/tutorials/atomic_action/move_held_object.py --headless --auto_play --device cpu python scripts/tutorials/atomic_action/place.py --headless --auto_play --device cpu + python scripts/tutorials/atomic_action/press.py --headless --auto_play --device cpu + python scripts/tutorials/atomic_action/coordinated_pickment.py --headless --auto_play --device cpu + python scripts/tutorials/atomic_action/coordinated_placement.py --headless --auto_play --device cpu + +The concrete implementations are organized one primitive per module under +``embodichain/lab/sim/atomic_actions/primitives``. Public imports from +``embodichain.lab.sim.atomic_actions`` remain the recommended API, and +``embodichain.lab.sim.atomic_actions.actions`` stays as a compatibility +re-export surface. Typical Usage ------------- diff --git a/docs/source/tutorial/modular_env.rst b/docs/source/tutorial/modular_env.rst index eef801c39..af4d14992 100644 --- a/docs/source/tutorial/modular_env.rst +++ b/docs/source/tutorial/modular_env.rst @@ -129,11 +129,12 @@ Configures a stereo camera system using :class:`StereoCameraCfg`: :language: python :lines: 120-130 -Defines scene illumination with controllable point lights: +Defines scene illumination with configurable lights: -- **Type**: Point light for realistic shadows -- **Properties**: Configurable color, intensity, and position -- **UID**: Named reference for event system manipulation +- **Types**: Supports ``"point"``, ``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, and ``"mesh"``. +- **Global lights**: ``"sun"`` and ``"direction"`` are global scene lights (single instance, infinite distance). +- **Properties**: Configurable color, intensity, position, and direction. +- **UID**: Named reference for event system manipulation. **Rigid Objects** diff --git a/docs/superpowers/plans/2026-07-13-newton-runtime-contracts-stage-1.md b/docs/superpowers/plans/2026-07-13-newton-runtime-contracts-stage-1.md new file mode 100644 index 000000000..049f6cc96 --- /dev/null +++ b/docs/superpowers/plans/2026-07-13-newton-runtime-contracts-stage-1.md @@ -0,0 +1,2617 @@ +# Newton Runtime Contracts Stage 1 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:** Replace the current private, stale-ID Newton integration with a +public DexSim contract that supports exact prepare/step semantics, +generation-aware rigid/articulation bindings, transactional runtime rebuild, +full multi-arena frames, and independent simultaneous simulation managers. + +**Architecture:** DexSim remains the sole owner of Newton model/state/control +truth and publishes stable entity references, immutable generation-tagged +bindings, prepare results, and rebuild events. EmbodiChain owns environment +semantics through an explicit `BackendSceneContext`, rebinds views after a +generation change, initializes only newly added objects, and keeps its existing +public object and manager calls as compatibility wrappers. + +**Tech Stack:** Python 3.11+, DexSim, NVIDIA Newton, NVIDIA Warp, PyTorch, +Gymnasium, pytest, Sphinx Markdown, Black 26.3.1. + +## Global Constraints + +- Source of truth: `docs/superpowers/specs/2026-07-13-newton-runtime-contracts-design.md`. +- EmbodiChain branch: `feature/newton-physics-backend` at or after `1d4b9eb`. +- DexSim branch: `feature/embodichain-newton-contracts` from `dev@5281bce`. +- DexSim package version is exactly `0.4.4`; `NEWTON_INTEGRATION_API_VERSION` is exactly `2`. +- `prepare()` creates a ready runtime and never advances simulation time; a requested update performs exactly the requested number of physics steps after prepare. +- Model generation starts at `0`, first successful finalize commits `1`, and only successful model replacement increments it. +- Runtime topology mutation is supported for rigid bodies and articulations; soft-body and cloth mutation fails explicitly. +- Runtime IDs are generation-scoped. Stable public identity is `(world_token, entity_handle, entity_kind)`. +- Public poses use `float32`, `xyzw`, and explicit world/arena-local frames; arena conversion uses the full SE(3) transform. +- Existing `SimulationManager`, `RigidObject`, and `Articulation` public call surfaces remain source compatible. +- Core Newton paths must not use `dexsim.default_world()`, global `get_physics_scene()`, or the default `SimulationManager` instance. +- Default-backend behavior remains unchanged. +- Stage 2 differentiable execution is not implemented by this plan; its plan is written only after this stage's merge gate passes. + +--- + +## Repository Baselines and Reference Files + +Run before Task 1: + +```bash +git -C /root/sources/dexsim branch --show-current +git -C /root/sources/dexsim rev-parse HEAD +git -C /root/sources/EmbodiChain branch --show-current +git -C /root/sources/EmbodiChain rev-parse HEAD +``` + +Expected branch names are `feature/embodichain-newton-contracts` and +`feature/newton-physics-backend`. Record the actual starting SHAs in the +execution log; do not reset either repository to the SHAs above. + +Use these implementations as focused references, without copying their +singleton/Omniverse assumptions: + +- DexSim runtime: `/root/sources/dexsim/python/dexsim/engine/newton_physics/newton_manager.py` +- DexSim rebuild: `/root/sources/dexsim/python/dexsim/engine/newton_physics/rebuild.py` +- DexSim articulation spans: `/root/sources/dexsim/python/dexsim/engine/newton_physics/articulation/articulation.py` +- IsaacLab clone mapping: `/root/sources/IsaacLab/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py` +- IsaacLab rebinding: `/root/sources/IsaacLab/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py` +- IsaacLab FK invalidation: `/root/sources/IsaacLab/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py` + +## File Map + +### DexSim files created + +- `python/dexsim/engine/newton_physics/contracts.py` — API version, stable references, prepare/rebuild results, subscriptions, model leases, and public errors. +- `python/dexsim/engine/newton_physics/bindings.py` — immutable rigid/articulation bindings and explicit q/qd spans. +- `python/dexsim/engine/newton_physics/runtime_snapshot.py` — generation-independent rigid/articulation snapshot and restore data. +- `python/test/engine/newton_physics/newton_contract_test_utils.py` — shared world, rigid, articulation, and state-array test helpers. +- `python/test/engine/newton_physics/test_newton_public_contract.py` — public API/version/prepare/attachment tests. +- `python/test/engine/newton_physics/test_newton_bindings.py` — generation, cross-world, static-body, and joint-span tests. +- `python/test/engine/newton_physics/test_newton_transactional_rebuild.py` — rigid/articulation preservation and rollback tests. +- `python/test/engine/newton_physics/test_newton_multi_world_runtime.py` — same-device isolation and cleanup tests. + +### DexSim files modified + +- `version.txt` — package patch version `0.4.4`. +- `python/dexsim/engine/newton_physics/__init__.py` — public exports. +- `python/dexsim/engine/newton_physics/newton_manager.py` — generation, prepare, public attachment/binding delegation, candidate commit, close, events. +- `python/dexsim/engine/newton_physics/rebuild.py` — candidate build/restore/validate/atomic commit. +- `python/dexsim/engine/newton_physics/registry.py` — world-token ownership and public owner lookup. +- `python/dexsim/engine/newton_physics/rigid_body/add_body.py` — legacy patch delegates to public attachment. +- `python/dexsim/engine/newton_physics/rigid_body/registration.py` — canonical descriptor replay into a chosen build target. +- `python/dexsim/engine/newton_physics/articulation/articulation.py` — stable ref and explicit q/qd span export. +- `python/dexsim/engine/newton_physics/articulation/skeleton_bridge.py` — canonical articulation replay and removal delta. +- `python/dexsim/engine/newton_physics/world.py` — prepare-then-step behavior and closed-world checks. +- `python/dexsim/engine/newton_physics/integration.py` — deterministic per-world teardown. +- `python/dexsim/engine/newton_physics/capture_coordinator.py` — weak, device-scoped capture coordination. + +### EmbodiChain files created + +- `embodichain/lab/sim/physics/context.py` — explicit backend/world/scene/arena ownership and full transform table. +- `tests/sim/newton_contract_test_utils.py` — deterministic Newton manager and asset-config fixtures used by the Stage 1 tests. +- `tests/sim/test_newton_scene_context.py` — ownership and full-frame conversion tests. +- `tests/sim/test_newton_rebuild_bindings.py` — pending initialization and generation refresh tests. +- `tests/sim/test_newton_multi_manager.py` — two-manager isolation and cleanup tests. + +### EmbodiChain files modified + +- `pyproject.toml` — exact `dexsim_engine==0.4.4` dependency. +- `embodichain/lab/sim/cfg.py` — strict Newton configuration validation. +- `embodichain/lab/sim/common.py` — remove base-constructor virtual reset. +- `embodichain/lab/sim/physics/__init__.py` — export the context/capability types. +- `embodichain/lab/sim/physics/base.py` — structured capability and prepare-result contracts. +- `embodichain/lab/sim/physics/default.py` — default-backend compatibility implementation. +- `embodichain/lab/sim/physics/newton.py` — API handshake, event subscription, generation, pending initialization, close. +- `embodichain/lab/sim/sim_manager.py` — context construction, exact update lifecycle, remove invalidation, idempotent close/reset. +- `embodichain/lab/sim/utility/sim_utils.py` — public `attach_rigid_body` use; remove private registry/meta writes. +- `embodichain/lab/sim/objects/rigid_object.py` — explicit context and deferred initialization. +- `embodichain/lab/sim/objects/articulation.py` — explicit context, separate position/velocity widths, deferred initialization. +- `embodichain/lab/sim/objects/robot.py` — pass owner context and defer reset until fully constructed. +- `embodichain/lab/sim/objects/backends/newton.py` — generation-aware bindings, full frames, FK, explicit unsupported errors. +- `embodichain/lab/sim/objects/backends/default.py` — accept explicit context without behavior changes. +- `tests/sim/test_backend_parity.py` — structured capability matrix. +- `tests/sim/test_newton_finalize_lifecycle.py` — exact prepare and initialization semantics. +- `tests/sim/objects/test_rigid_object.py` — multi-arena/rebuild compatibility cases. +- `tests/sim/objects/test_articulation.py` — q/qd span, FK, link/root frame, rebuild cases. +- `docs/source/overview/sim/sim_manager.md` — public lifecycle, capabilities, topology mutation, and cleanup contract. +- `design/newton-backend-design.md` — replace obsolete Target 4 claims with verified Stage 1 status. + +--- + +### Task 1: Publish the DexSim integration contract and generation lifecycle + +**Files:** + +- Create: `/root/sources/dexsim/python/dexsim/engine/newton_physics/contracts.py` +- Create: `/root/sources/dexsim/python/test/engine/newton_physics/newton_contract_test_utils.py` +- Create: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_public_contract.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/newton_manager.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/__init__.py` +- Modify: `/root/sources/dexsim/version.txt` + +**Interfaces:** + +- Produces: `NEWTON_INTEGRATION_API_VERSION = 2`. +- Produces: `NewtonEntityRef`, `NewtonPrepareResult`, `NewtonPrepareFailure`, `NewtonModelRebuiltEvent`, `NewtonRuntimeStatus`, `NewtonSubscription`, `NewtonModelLease`. +- Produces: `NewtonIntegrationError` subclasses for stale generation, cross-world use, unsupported operation, closed runtime, rebuild failure, and active model lease. +- Produces: `NewtonManager.world_token`, `NewtonManager.model_generation`, `NewtonManager.runtime_status`, and `NewtonManager.prepare()`. +- Produces: `NewtonManager.acquire_model_lease() -> NewtonModelLease`; Stage 2 sessions consume this primitive without changing its rebuild-safety semantics. + +- [ ] **Step 1: Add shared deterministic test helpers** + +Create `newton_contract_test_utils.py` and import its symbols explicitly from +each new DexSim test file: + +```python +DT = 0.01 + + +def make_world(device: str = "cpu"): + config = dexsim.WorldConfig() + config.open_windows = False + config.use_default_physics = False + config.backend = dexsim.types.Backend.OPENGL + config.renderer = dexsim.types.Renderer.HYBRID + world = dexsim.World(config) + world.set_physics_backend("Newton", cfg=NewtonCfg(device=device)) + manager = get_newton_manager(world) + manager._dexsim_renderer = config.renderer + manager._visualizer_disabled = True + return world, world.get_env(), manager + + +@pytest.fixture +def newton_world(): + world, env, manager = make_world() + try: + yield world, env + finally: + world.quit() + + +@pytest.fixture +def two_newton_worlds(): + first = make_world() + second = make_world() + try: + yield first, second + finally: + first[0].quit() + second[0].quit() + + +@pytest.fixture +def two_cuda_worlds(): + if not wp.is_cuda_available(): + pytest.skip("CUDA is required for same-device capture coordination.") + first = make_world("cuda:0") + second = make_world("cuda:0") + try: + yield first, second + finally: + first[0].quit() + second[0].quit() + + +def dynamic_box(arena, name: str, z: float = 1.0): + obj = arena.create_cube(0.1, 0.1, 0.1) + obj.set_name(name) + obj.set_location(0.0, 0.0, z) + attr = PhysicalAttr() + attr.mass = 1.0 + obj.add_rigidbody(ActorType.DYNAMIC, RigidBodyShape.BOX, attr) + return obj + + +def static_plane(arena, name: str): + obj = arena.create_plane(0.0, 10.0) + obj.set_name(name) + obj.add_rigidbody(ActorType.STATIC, RigidBodyShape.PLANE, PhysicalAttr()) + return obj + + +def test_articulation(arena, name: str): + path = get_resources_data_path("Robot", "UR5GPI", "UR5_pgi.urdf") + articulation = arena.load_urdf(path) + articulation.set_name(name) + return articulation + + +def assign_body_state(state, body_id: int, pose, velocity, acceleration) -> None: + body_q = state.body_q.numpy() + body_qd = state.body_qd.numpy() + body_qdd = state.body_qdd.numpy() + body_q[body_id] = np.asarray(pose, dtype=np.float32) + body_qd[body_id] = np.asarray(velocity, dtype=np.float32) + body_qdd[body_id] = np.asarray(acceleration, dtype=np.float32) + state.body_q.assign(body_q) + state.body_qd.assign(body_qd) + state.body_qdd.assign(body_qdd) +``` + +Each test file imports `dynamic_box as _dynamic_box`, +`static_plane as _static_plane`, `test_articulation as +_test_urdf_articulation`, and `assign_body_state as _assign_body_state`, so all +helper names in the following snippets are defined. + +- [ ] **Step 2: Write contract and prepare tests that fail on the current API** + +Add tests with these assertions: + +```python +def test_public_contract_version_and_initial_generation(newton_world): + world, _ = newton_world + mgr = get_newton_manager(world) + assert Version(dexsim.__version__).base_version == "0.4.4" + assert NEWTON_INTEGRATION_API_VERSION == 2 + assert mgr.model_generation == 0 + assert mgr.runtime_status.model_finalized is False + + +def test_prepare_finalizes_without_advancing_time(newton_world): + world, env = newton_world + _dynamic_box(env, "box") + mgr = get_newton_manager(world) + before = mgr._sim_time + result = mgr.prepare() + assert result.generation == 1 + assert result.did_build is True + assert result.did_rebuild is False + assert len(result.added_entities) == 1 + assert result.removed_entities == () + assert mgr._sim_time == before + assert mgr.model_generation == 1 + assert mgr.runtime_status.solver_ready is True + + +def test_second_prepare_is_idempotent(newton_world): + world, env = newton_world + _dynamic_box(env, "box") + mgr = get_newton_manager(world) + first = mgr.prepare() + second = mgr.prepare() + assert first.generation == second.generation == 1 + assert second.did_build is False + assert second.did_rebuild is False +``` + +- [ ] **Step 3: Run the focused test and confirm the missing-contract failure** + +Run: + +```bash +cd /root/sources/dexsim +pytest -q python/test/engine/newton_physics/test_newton_public_contract.py +``` + +Expected: collection fails because the new public symbols and `prepare()` do +not exist. + +- [ ] **Step 4: Add the public value types and errors** + +Implement the contract with immutable, typed values: + +```python +from __future__ import annotations + + +NEWTON_INTEGRATION_API_VERSION = 2 + + +class NewtonIntegrationError(RuntimeError): + """Base error for the public Newton integration contract.""" + + +class NewtonStaleBindingError(NewtonIntegrationError): + pass + + +class NewtonCrossWorldError(NewtonIntegrationError): + pass + + +class NewtonUnsupportedOperationError(NewtonIntegrationError): + pass + + +class NewtonClosedError(NewtonIntegrationError): + pass + + +class NewtonRebuildError(NewtonIntegrationError): + def __init__(self, failure: NewtonPrepareFailure) -> None: + self.failure = failure + super().__init__(failure.message) + + +class NewtonActiveLeaseError(NewtonIntegrationError): + pass + + +@dataclass(frozen=True, slots=True) +class NewtonEntityRef: + world_token: int + entity_handle: int + entity_kind: Literal["rigid", "articulation"] + + +@dataclass(frozen=True, slots=True) +class NewtonPrepareResult: + generation: int + did_build: bool + did_rebuild: bool + added_entities: tuple[NewtonEntityRef, ...] = () + removed_entities: tuple[NewtonEntityRef, ...] = () + + +@dataclass(frozen=True, slots=True) +class NewtonModelRebuiltEvent: + old_generation: int + new_generation: int + added_entities: tuple[NewtonEntityRef, ...] + removed_entities: tuple[NewtonEntityRef, ...] + + +@dataclass(frozen=True, slots=True) +class NewtonRuntimeStatus: + model_finalized: bool + solver_ready: bool + running: bool + stale: bool + closed: bool + + +@dataclass(frozen=True, slots=True) +class NewtonPrepareFailure: + generation: int + operation: Literal["build", "rebuild"] + message: str + + +class NewtonSubscription: + def __init__(self, unsubscribe: Callable[[], None]) -> None: + self._unsubscribe = unsubscribe + self._closed = False + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._unsubscribe() + + +class NewtonModelLease: + def __init__( + self, + *, + world_token: int, + generation: int, + model: object, + release: Callable[[], None], + ) -> None: + self.world_token = world_token + self.generation = generation + self.model = model + self._release = release + self._closed = False + + def __enter__(self) -> NewtonModelLease: + return self + + def __exit__(self, *exc_info: object) -> None: + self.close() + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._release() +``` + +Keep `NewtonModelState.BUILDER` and `NewtonModelState.READY` working for +existing callers. Add `_model_generation`, `_has_stepped`, `_closed`, and a +monotonic per-manager `world_token`. `prepare()` delegates build/rebuild to the +later transactional helper, returns an idempotent result, and never calls +`simulate()` or increments `_sim_time`. +Add +`subscribe_model_rebuilt(callback: Callable[[NewtonModelRebuiltEvent], None]) -> NewtonSubscription`; +callbacks are stored per manager and invoked only after an atomic successful +commit. A failed build raises `NewtonRebuildError` carrying a +`NewtonPrepareFailure` and emits no rebuilt event. +`acquire_model_lease()` requires a finalized, non-stale model, increments a +per-manager active-lease counter, and returns a lease that strongly owns that +exact model and generation. Lease release is idempotent and decrements the +counter exactly once. This is a lifecycle primitive only; no differentiable +session or tape behavior is added in Stage 1. + +- [ ] **Step 5: Export the contract and bump the package patch version** + +Export the new symbols from `newton_physics/__init__.py` and change only: + +```text +DEXSIM_VERSION_MAJOR 0 +DEXSIM_VERSION_MINOR 4 +DEXSIM_VERSION_PATCH 4 +``` + +- [ ] **Step 6: Refresh the editable DexSim package metadata** + +The development install points at `build_Release/lib/python_package`; its +Python modules are source symlinks, but its `dexsim/version.txt` is a generated +copy. Refresh it through the repository-supported setup script after changing +the root version: + +```bash +cd /root/sources/dexsim +PYTHON_BIN="$(command -v python)" ./setup_dev_python.sh -j12 +python - <<'PY' +from packaging.version import Version +import dexsim +assert Version(dexsim.__version__).base_version == "0.4.4" +print(dexsim.__version__, dexsim.__file__) +PY +``` + +Expected: the editable install reports base version `0.4.4` and imports from +the local development package. Generated `build_Release` contents are never +staged or committed. + +- [ ] **Step 7: Run the focused tests** + +Run: + +```bash +pytest -q python/test/engine/newton_physics/test_newton_public_contract.py +``` + +Expected: all tests in the file pass; no simulation-time change occurs during +prepare. + +- [ ] **Step 8: Commit the public contract** + +```bash +git -C /root/sources/dexsim add version.txt python/dexsim/engine/newton_physics/contracts.py python/dexsim/engine/newton_physics/newton_manager.py python/dexsim/engine/newton_physics/__init__.py python/test/engine/newton_physics/newton_contract_test_utils.py python/test/engine/newton_physics/test_newton_public_contract.py +git -C /root/sources/dexsim commit -m "feat(newton): publish runtime integration contract" +``` + +--- + +### Task 2: Replace private rigid registration with stable public attachment + +**Files:** + +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/contracts.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/newton_manager.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/rigid_body/add_body.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/rigid_body/registration.py` +- Modify: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_public_contract.py` + +**Interfaces:** + +- Produces: `NewtonManager.attach_rigid_body(...) -> NewtonEntityRef`. +- Produces: immutable `NewtonRigidDescriptor` replay records owned by DexSim. +- Produces: `NewtonManager.entity_ref(entity)`, `descriptor_for(ref)`, and read-only `descriptors`. +- Consumes: `NewtonEntityRef` and generation lifecycle from Task 1. + +- [ ] **Step 1: Add failing public attachment, clone-world, and descriptor-isolation tests** + +```python +def test_attach_rigid_body_returns_stable_ref(newton_world): + world, env = newton_world + cube = env.create_cube(0.1, 0.2, 0.3) + mgr = get_newton_manager(world) + ref = mgr.attach_rigid_body( + cube, + actor_type=ActorType.DYNAMIC, + shape_type=RigidBodyShape.BOX, + physical_attr=PhysicalAttr(), + ) + assert ref.world_token == mgr.world_token + assert ref.entity_handle == cube.get_native_handle() + assert ref.entity_kind == "rigid" + assert mgr.entity_ref(cube) == ref + + +def test_child_arena_attachment_uses_child_world(newton_world): + world, env = newton_world + arena = env.add_arena("arena_a") + cube = arena.create_cube(0.1, 0.1, 0.1) + mgr = get_newton_manager(world) + ref = mgr.attach_rigid_body( + cube, + actor_type=ActorType.DYNAMIC, + shape_type=RigidBodyShape.BOX, + physical_attr=PhysicalAttr(), + ) + mgr.prepare() + binding = mgr.bind_rigid_entities((ref,)) + assert binding.world_ids_host.tolist() == [0] + assert mgr._model.body_world.numpy()[binding.body_ids_host[0]] == 0 + + +def test_clone_descriptors_do_not_alias(newton_world): + world, env = newton_world + arena_a = env.add_arena("arena_a") + arena_b = env.add_arena("arena_b") + prototype = arena_a.create_sphere(0.1) + prototype.set_name("prototype") + prototype.add_rigidbody( + ActorType.DYNAMIC, RigidBodyShape.SPHERE, PhysicalAttr() + ) + clone = arena_a.clone_actor_to( + "prototype", arena_b, "clone", ObjectCloneOptions() + ) + mgr = get_newton_manager(world) + source = mgr.descriptor_for(mgr.entity_ref(prototype)) + target = mgr.descriptor_for(mgr.entity_ref(clone)) + assert source is not target + assert source.world_id == 0 + assert target.world_id == 1 + + +def test_mesh_geometry_descriptor_is_owned_by_dexsim(newton_world): + world, env = newton_world + entity = env.create_cube(0.1, 0.1, 0.1) + vertices = np.array( + [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float32 + ) + triangles = np.array( + [[0, 1, 2], [0, 1, 3], [0, 2, 3], [1, 2, 3]], dtype=np.int32 + ) + geometry = GeometryDesc.mesh(vertices=vertices, triangles=triangles) + mgr = get_newton_manager(world) + ref = mgr.attach_rigid_body( + entity, + actor_type=ActorType.DYNAMIC, + shape_type=RigidBodyShape.MESH, + physical_attr=PhysicalAttr(), + geometry_desc=geometry, + ) + vertices[0] = 99.0 + owned = mgr.descriptor_for(ref).geometry_desc + assert np.allclose(owned.vertices[0], [0.0, 0.0, 0.0]) + + +def test_desc_native_box_prepares_without_shape_parameter_fallback(newton_world): + world, env = newton_world + entity = env.create_cube(0.1, 0.2, 0.3) + mgr = get_newton_manager(world) + ref = mgr.attach_rigid_body( + entity, + actor_type=ActorType.DYNAMIC, + shape_type=RigidBodyShape.BOX, + body_desc=RigidBodyPhysicsDesc.dynamic(mass=1.0), + shape_desc=NewtonCollisionDesc(ke=1000.0, kd=50.0, margin=0.01), + geometry_desc=GeometryDesc.cube((0.1, 0.2, 0.3)), + ) + result = mgr.prepare() + assert result.generation == 1 + assert mgr.bind_rigid_entities((ref,)).body_ids_host[0] >= 0 + + +def test_desc_native_sphere_prepares_from_owned_geometry(newton_world): + world, env = newton_world + entity = env.create_sphere(0.2) + mgr = get_newton_manager(world) + ref = mgr.attach_rigid_body( + entity, + actor_type=ActorType.DYNAMIC, + shape_type=RigidBodyShape.SPHERE, + body_desc=RigidBodyPhysicsDesc.dynamic(mass=1.0), + shape_desc=NewtonCollisionDesc(ke=1000.0, kd=50.0), + geometry_desc=GeometryDesc.sphere(0.2), + ) + mgr.prepare() + descriptor = mgr.descriptor_for(ref) + assert descriptor.geometry_desc.radius == pytest.approx(0.2) + assert mgr.bind_rigid_entities((ref,)).body_ids_host[0] >= 0 +``` + +- [ ] **Step 2: Run the tests and confirm public attachment is absent** + +Run the attachment tests above with `pytest -q`. Expected: failures report +missing `attach_rigid_body`, `entity_ref`, or `descriptor_for`. + +- [ ] **Step 3: Add canonical descriptor ownership and the public method** + +Add an immutable descriptor that deep-copies mutable desc-native inputs: + +```python +@dataclass(frozen=True, slots=True) +class NewtonRigidDescriptor: + entity_ref: NewtonEntityRef + arena_handle: int + world_id: int + actor_type: ActorType + shape_type: RigidBodyShape + node_scale: tuple[float, float, float] + body_scale: tuple[float, float, float] + physical_attr: PhysicalAttr | None + body_desc: object | None + shape_desc: object | None + geometry_desc: object | None +``` + +The `object` annotations above stand for the existing concrete DexSim spawn +descriptor types, not borrowed arbitrary objects. At attachment time, normalize +legacy `PhysicalAttr` into owned body/collision values and recursively copy all +desc-native data. Copy NumPy arrays with canonical `float32`/`int32` dtypes and +mark them read-only. Resolve file-backed mesh data, convex/ACD hulls, and SDF +mesh/config inputs into descriptor-owned replay data so a rebuild neither reads +mutable entity metadata nor depends on a later file change. `descriptor_for()` +returns this immutable snapshot; it never returns `dexsim_meta`. + +Implement the public method with this exact surface: + +```python +def attach_rigid_body( + self, + entity, + *, + actor_type, + shape_type, + physical_attr=None, + body_desc=None, + shape_desc=None, + geometry_desc=None, +) -> NewtonEntityRef: + self._assert_open() + ref = self._ref_for_entity(entity, "rigid") + descriptor = make_rigid_descriptor( + manager=self, + entity=entity, + entity_ref=ref, + actor_type=actor_type, + shape_type=shape_type, + physical_attr=physical_attr, + body_desc=body_desc, + shape_desc=shape_desc, + geometry_desc=geometry_desc, + ) + self._rigid_descriptors[ref] = descriptor + replay_rigid_descriptor(self, descriptor, entity) + self._record_added(ref) + self.mark_runtime_model_stale() + return ref +``` + +The legacy `MeshObject.add_rigidbody` patches call this method and return its +reference. `register_mesh_object_to_newton_patch` remains temporarily importable +for DexSim compatibility but delegates through descriptor replay and is no +longer called by EmbodiChain. +Legacy `add_sdf_rigidbody` and `add_acd_rigidbody` project their SDF config or +resolved convex hulls into the same `geometry_desc` contract before delegation; +the existing SDF/ACD regression tests must keep passing. + +- [ ] **Step 4: Run attachment and existing scene mutation tests** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_public_contract.py python/test/engine/newton_physics/test_newton_scene_mutations.py +``` + +Expected: both files pass, including global world `-1` and child world IDs. + +- [ ] **Step 5: Commit public rigid attachment** + +```bash +git -C /root/sources/dexsim add python/dexsim/engine/newton_physics/contracts.py python/dexsim/engine/newton_physics/newton_manager.py python/dexsim/engine/newton_physics/rigid_body/add_body.py python/dexsim/engine/newton_physics/rigid_body/registration.py python/test/engine/newton_physics/test_newton_public_contract.py +git -C /root/sources/dexsim commit -m "feat(newton): add stable rigid attachment API" +``` + +--- + +### Task 3: Add immutable generation-aware rigid and articulation bindings + +**Files:** + +- Create: `/root/sources/dexsim/python/dexsim/engine/newton_physics/bindings.py` +- Create: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_bindings.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/newton_manager.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/articulation/articulation.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/__init__.py` + +**Interfaces:** + +- Produces: `NewtonJointSpan`, `NewtonArticulationSpan`, `RigidEntityBinding`, `ArticulationBinding`. +- Produces: `joint_span_from_builder(builder, joint_id) -> NewtonJointSpan`. +- Produces: `NewtonManager.bind_rigid_entities(refs, device=None)`. +- Produces: `NewtonManager.bind_articulations(refs, device=None)`. +- Consumes: stable references from Task 2. + +- [ ] **Step 1: Add failing binding contract tests** + +```python +def test_rigid_binding_is_int32_and_generation_scoped(newton_world): + world, env = newton_world + dynamic = _dynamic_box(env, "dynamic") + static = _static_plane(env, "static") + mgr = get_newton_manager(world) + result = mgr.prepare() + binding = mgr.bind_rigid_entities( + (mgr.entity_ref(dynamic), mgr.entity_ref(static)) + ) + assert binding.generation == result.generation + assert binding.body_ids_host.dtype == np.int32 + assert binding.shape_ids_host.dtype == np.int32 + assert binding.body_ids_host[1] == -1 + binding.assert_current(mgr) + + +def test_binding_rejects_other_world(two_newton_worlds): + (world_a, env_a, mgr_a), (_, _, mgr_b) = two_newton_worlds + box = _dynamic_box(env_a, "box") + mgr_a.prepare() + with pytest.raises(NewtonCrossWorldError): + mgr_b.bind_rigid_entities((mgr_a.entity_ref(box),)) + + +def test_articulation_binding_exposes_distinct_q_and_qd_spans(newton_world): + world, env = newton_world + art = _test_urdf_articulation(env, "arm") + mgr = get_newton_manager(world) + mgr.prepare() + binding = mgr.bind_articulations((mgr.entity_ref(art),)) + assert binding.qpos_width > 0 + assert binding.qvel_width > 0 + assert all( + span.q_width >= 0 for spans in binding.joint_spans for span in spans + ) + assert all( + span.qd_width >= 0 for spans in binding.joint_spans for span in spans + ) + assert binding.joint_spans_wp.dtype == wp.int32 +``` + +Add builder-only unit cases for revolute, spherical, and free-joint widths: + +```python +@pytest.mark.parametrize( + "q_width, qd_width", + [(1, 1), (4, 3), (7, 6)], +) +def test_joint_span_uses_distinct_position_and_velocity_widths( + q_width, qd_width +): + builder = SimpleNamespace( + joint_q_start=[0, q_width], + joint_qd_start=[0, qd_width], + joint_q=[0.0] * q_width, + joint_qd=[0.0] * qd_width, + joint_target_pos=[0.0] * qd_width, + joint_target_vel=[0.0] * qd_width, + ) + span = joint_span_from_builder(builder, joint_id=0) + assert span.q_start == span.qd_start == 0 + assert span.q_width == q_width + assert span.qd_width == qd_width + assert span.target_q_start == span.qd_start + assert span.target_q_width == qd_width + assert span.target_qd_start == span.qd_start + assert span.target_qd_width == qd_width +``` + +- [ ] **Step 2: Run the binding file and confirm collection/API failures** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_bindings.py +``` + +Expected: missing binding classes/methods. + +- [ ] **Step 3: Implement immutable bindings with O(1) validation** + +```python +@dataclass(frozen=True, slots=True) +class NewtonJointSpan: + joint_id: int + q_start: int + q_width: int + qd_start: int + qd_width: int + target_q_start: int + target_q_width: int + target_qd_start: int + target_qd_width: int + + +@dataclass(frozen=True, slots=True) +class NewtonArticulationSpan: + ref: NewtonEntityRef + q_start: int + q_width: int + qd_start: int + qd_width: int + target_q_start: int + target_q_width: int + target_qd_start: int + target_qd_width: int + control_start: int + control_width: int + + +@dataclass(frozen=True, slots=True) +class RigidEntityBinding: + world_token: int + generation: int + refs: tuple[NewtonEntityRef, ...] + body_ids_host: np.ndarray + shape_ids_host: np.ndarray + world_ids_host: np.ndarray + body_ids_wp: wp.array + shape_ids_wp: wp.array + world_ids_wp: wp.array + + def assert_current(self, manager: NewtonManager) -> None: + assert_binding_owner_and_generation(self, manager) + + +@dataclass(frozen=True, slots=True) +class ArticulationBinding: + world_token: int + generation: int + refs: tuple[NewtonEntityRef, ...] + articulation_ids_host: np.ndarray + root_body_ids_host: np.ndarray + link_body_ids_host: np.ndarray + world_ids_host: np.ndarray + articulation_spans: tuple[NewtonArticulationSpan, ...] + joint_spans: tuple[tuple[NewtonJointSpan, ...], ...] + joint_spans_wp: wp.array + qpos_width: int + qvel_width: int + target_qpos_width: int + target_qvel_width: int +``` + +Resolve all host arrays once, upload `int32` device arrays once, make NumPy +arrays read-only, and validate a binding by comparing only `world_token` and +`generation`. Do not resolve entity IDs in steady-state fetch/apply calls. +`joint_span_from_builder` computes each end from the next start entry, or from +the corresponding flat array length for the final joint. Current position uses +`joint_q_start/q_width`; current velocity uses `joint_qd_start/qd_width`. +Newton `joint_target_pos`, `joint_target_vel`, and `joint_f` are per-DOF, so +their starts and widths use the qd span even for spherical/free joints. Do not +reuse the coordinate-width q span for position targets. + +- [ ] **Step 4: Run binding and simulation-index regressions** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_bindings.py python/test/engine/newton_physics/test_newton_sim_index.py +``` + +Expected: all tests pass; old `get_sim_index()` remains a generation-local +compatibility view, not a stable handle. + +- [ ] **Step 5: Commit bindings** + +```bash +git -C /root/sources/dexsim add python/dexsim/engine/newton_physics/bindings.py python/dexsim/engine/newton_physics/newton_manager.py python/dexsim/engine/newton_physics/articulation/articulation.py python/dexsim/engine/newton_physics/__init__.py python/test/engine/newton_physics/test_newton_bindings.py +git -C /root/sources/dexsim commit -m "feat(newton): add generation-aware entity bindings" +``` + +--- + +### Task 4: Make rigid rebuild transactional and preserve both state buffers + +**Files:** + +- Create: `/root/sources/dexsim/python/dexsim/engine/newton_physics/runtime_snapshot.py` +- Create: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_transactional_rebuild.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/rebuild.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/newton_manager.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/world.py` + +**Interfaces:** + +- Produces: `NewtonRuntimeSnapshot`, `RigidRuntimeSnapshot`, and candidate-runtime commit. +- Produces: successful rebuild event after commit only. +- Produces: rebuild rejection while a model-generation lease is active. +- Consumes: descriptors and bindings from Tasks 2–3. + +- [ ] **Step 1: Add failing state-preservation, rollback, and exact-step tests** + +```python +def test_rigid_rebuild_preserves_both_states_and_wrench(newton_world): + world, env = newton_world + box = _dynamic_box(env, "box") + mgr = get_newton_manager(world) + mgr.prepare() + ref = mgr.entity_ref(box) + body_id = mgr.bind_rigid_entities((ref,)).body_ids_host[0] + pose0 = np.array([0.1, 0.2, 1.3, 0.0, 0.0, 0.0, 1.0], np.float32) + pose1 = np.array([0.2, 0.3, 1.4, 0.0, 0.0, 0.0, 1.0], np.float32) + _assign_body_state( + mgr._state_0, + body_id, + pose0, + [1, 2, 3, 4, 5, 6], + [0.1, 0.2, 0.3, 0.4, 0.5, 0.6], + ) + _assign_body_state( + mgr._state_1, + body_id, + pose1, + [6, 5, 4, 3, 2, 1], + [0.6, 0.5, 0.4, 0.3, 0.2, 0.1], + ) + external_forces = mgr._external_forces.numpy() + external_forces[body_id] = [1, 2, 3, 4, 5, 6] + mgr._external_forces.assign(external_forces) + _dynamic_box(env, "new_box") + result = mgr.prepare() + new_id = mgr.bind_rigid_entities((ref,)).body_ids_host[0] + assert result.did_rebuild is True + assert np.allclose(mgr._state_0.body_q.numpy()[new_id], pose0) + assert np.allclose(mgr._state_1.body_q.numpy()[new_id], pose1) + assert np.allclose( + mgr._state_0.body_qdd.numpy()[new_id], + [0.1, 0.2, 0.3, 0.4, 0.5, 0.6], + ) + assert np.allclose( + mgr._state_1.body_qdd.numpy()[new_id], + [0.6, 0.5, 0.4, 0.3, 0.2, 0.1], + ) + assert np.allclose(mgr._external_forces.numpy()[new_id], [1, 2, 3, 4, 5, 6]) + + +def test_failed_candidate_keeps_old_runtime_and_generation(newton_world, monkeypatch): + world, env = newton_world + box = _dynamic_box(env, "box") + mgr = get_newton_manager(world) + mgr.prepare() + old_model = mgr._model + old_generation = mgr.model_generation + _dynamic_box(env, "new_box") + monkeypatch.setattr(rebuild, "_validate_candidate", lambda candidate: (_ for _ in ()).throw(ValueError("injected"))) + with pytest.raises(NewtonRebuildError, match="injected"): + mgr.prepare() + assert mgr._model is old_model + assert mgr.model_generation == old_generation + assert mgr.lifecycle_state is NewtonModelState.STALE + + +def test_world_update_prepares_then_executes_one_step(newton_world): + world, env = newton_world + _dynamic_box(env, "box", z=1.0) + mgr = get_newton_manager(world) + before = mgr._sim_time + world.update(0.01) + assert mgr.model_generation == 1 + assert mgr._sim_time == pytest.approx(before + 0.01) + + +def test_active_model_lease_blocks_rebuild_until_released(newton_world): + world, env = newton_world + _dynamic_box(env, "box") + mgr = get_newton_manager(world) + mgr.prepare() + lease = mgr.acquire_model_lease() + assert lease.world_token == mgr.world_token + assert lease.generation == mgr.model_generation + assert lease.model is mgr._model + _dynamic_box(env, "new_box") + with pytest.raises(NewtonActiveLeaseError, match="generation 1"): + mgr.prepare() + assert mgr.model_generation == 1 + lease.close() + lease.close() + assert mgr.prepare().generation == 2 +``` + +- [ ] **Step 2: Run the new file and confirm current rebuild destroys/aliases state** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_transactional_rebuild.py +``` + +Expected: failures show missing snapshot types, non-transactional clearing, or +the current first-update skip; the active-lease case currently rebuilds or has +no lease surface. + +- [ ] **Step 3: Add complete rigid snapshots keyed by stable reference** + +```python +@dataclass(frozen=True, slots=True) +class RigidRuntimeSnapshot: + ref: NewtonEntityRef + state_0_pose: np.ndarray + state_0_velocity: np.ndarray + state_0_acceleration: np.ndarray + state_1_pose: np.ndarray + state_1_velocity: np.ndarray + state_1_acceleration: np.ndarray + external_wrench: np.ndarray + + +@dataclass(frozen=True, slots=True) +class NewtonRuntimeSnapshot: + generation: int + rigid: dict[NewtonEntityRef, RigidRuntimeSnapshot] + articulations: dict[NewtonEntityRef, ArticulationRuntimeSnapshot] +``` + +Capture only references present in the old finalized model. Restore surviving +references after candidate finalization using the candidate binding. Copy +arrays into both candidate states; do not make both buffers identical when the +old buffers differed. + +- [ ] **Step 4: Build and validate a candidate before mutating the live manager** + +Use an unregistered candidate manager that shares only the stable world token: + +```python +candidate = NewtonManager( + cfg=copy.deepcopy(manager.cfg), + world_token=manager.world_token, + register_live=False, +) +candidate.set_dexsim_world(world) +candidate.replay_descriptors(manager.descriptors) +candidate.start_simulation() +restore_runtime_snapshot(candidate, snapshot) +_validate_candidate(candidate) +manager._commit_candidate(candidate, delta) +``` + +Before constructing a candidate, `prepare()` checks the active-lease count. A +topology-changing prepare raises `NewtonActiveLeaseError` with the world token +and leased generation while the count is non-zero; idempotent prepare of an +unchanged READY model remains allowed. The live model, generation, topology +delta, and STALE status remain unchanged after rejection. Stage 2 must acquire +this lease before recording a tape and close it only after backward or explicit +session detach. + +`_commit_candidate` swaps builder/model/states/control/contacts/pipeline/solver, +entity mappings, descriptors, articulation runtime bindings, caches, and graph +as one non-raising assignment block. It increments generation once, sets READY, +detaches the transferred resources from the candidate, then emits one rebuilt +event. Retire the old runtime only after the swap; cleanup failure is logged +without rolling back a runtime already published to subscribers. Candidate +failure before the swap closes candidate resources, +leaves the live runtime fields unchanged, keeps STALE, and raises +`NewtonRebuildError` chained from the original exception. + +- [ ] **Step 5: Make World.update call prepare and still step** + +Replace skip-on-build behavior with: + +```python +prepare_result = mgr.prepare() +mgr.step(dt) +if mgr.should_sync_to_dexsim(step_override=sync_to_dexsim): + _push_newton_state_to_dexsim(mgr) +``` + +The result is available to integrations but never consumes the requested +physics step. + +- [ ] **Step 6: Run lifecycle and mutation regressions** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_transactional_rebuild.py python/test/engine/newton_physics/test_newton_scene_lifecycle.py python/test/engine/newton_physics/test_newton_scene_mutations.py python/test/engine/newton_physics/test_newton_body_dynamics.py +``` + +Expected: all pass; update time increments on the first call and after rebuild. + +- [ ] **Step 7: Commit transactional rigid rebuild** + +```bash +git -C /root/sources/dexsim add python/dexsim/engine/newton_physics/runtime_snapshot.py python/dexsim/engine/newton_physics/rebuild.py python/dexsim/engine/newton_physics/newton_manager.py python/dexsim/engine/newton_physics/world.py python/test/engine/newton_physics/test_newton_transactional_rebuild.py +git -C /root/sources/dexsim commit -m "refactor(newton): rebuild rigid runtime transactionally" +``` + +--- + +### Task 5: Preserve articulation runtime state and rebind explicit spans + +**Files:** + +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/runtime_snapshot.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/rebuild.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/articulation/articulation.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/articulation/skeleton_bridge.py` +- Modify: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_transactional_rebuild.py` +- Modify: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_bindings.py` + +**Interfaces:** + +- Produces: complete `ArticulationRuntimeSnapshot` and canonical articulation replay. +- Produces: `NewtonManager.forward_kinematics(articulation_mask=None)` as the public FK synchronization entry point; it updates both ping-pong states. +- Consumes: `ArticulationBinding` from Task 3 and candidate transaction from Task 4. + +- [ ] **Step 1: Add failing articulation preservation and removal tests** + +Define these local test helpers above the tests: + +```python +def _assign_slice(owner, name: str, start: int, values: np.ndarray) -> None: + source = getattr(owner, name) + data = source.numpy() + data[start : start + len(values)] = values + source.assign(data) + + +def _read_slice(owner, name: str, start: int, width: int) -> np.ndarray: + return getattr(owner, name).numpy()[start : start + width].copy() + + +def _assign_existing_owners( + manager, owner_names: tuple[str, ...], name: str, start: int, values +) -> None: + assigned = False + for owner_name in owner_names: + owner = getattr(manager, owner_name, None) + if owner is None or getattr(owner, name, None) is None: + continue + _assign_slice(owner, name, start, values) + assigned = True + assert assigned, f"{name} is unavailable on {owner_names}" + + +def _read_first_owner( + manager, owner_names: tuple[str, ...], name: str, start: int, width: int +) -> np.ndarray: + for owner_name in owner_names: + owner = getattr(manager, owner_name, None) + if owner is not None and getattr(owner, name, None) is not None: + return _read_slice(owner, name, start, width) + raise AssertionError(f"{name} is unavailable on {owner_names}") + + +def _write_articulation_arrays( + manager, + binding, + state_0_q, + state_0_qd, + state_1_q, + state_1_qd, + model_target_q, + model_target_qd, + control_target_q, + control_target_qd, + generalized_force, + active_control, +) -> None: + span = binding.articulation_spans[0] + _assign_slice(manager._state_0, "joint_q", span.q_start, state_0_q) + _assign_slice(manager._state_0, "joint_qd", span.qd_start, state_0_qd) + _assign_slice(manager._state_1, "joint_q", span.q_start, state_1_q) + _assign_slice(manager._state_1, "joint_qd", span.qd_start, state_1_qd) + _assign_existing_owners( + manager, + ("_model",), + "joint_target_pos", + span.target_q_start, + model_target_q, + ) + _assign_existing_owners( + manager, + ("_model",), + "joint_target_vel", + span.target_qd_start, + model_target_qd, + ) + _assign_existing_owners( + manager, + ("_control",), + "joint_target_pos", + span.target_q_start, + control_target_q, + ) + _assign_existing_owners( + manager, + ("_control",), + "joint_target_vel", + span.target_qd_start, + control_target_qd, + ) + _assign_existing_owners( + manager, + ("_model",), + "joint_f", + span.control_start, + generalized_force, + ) + _assign_existing_owners( + manager, + ("_control",), + "joint_f", + span.control_start, + active_control, + ) + + +def _read_q(manager, binding, state_name="_state_0"): + span = binding.articulation_spans[0] + return _read_slice( + getattr(manager, state_name), "joint_q", span.q_start, span.q_width + ) + + +def _read_qd(manager, binding, state_name="_state_0"): + span = binding.articulation_spans[0] + return _read_slice( + getattr(manager, state_name), "joint_qd", span.qd_start, span.qd_width + ) + + +def _read_target_q(manager, binding, owner_name): + span = binding.articulation_spans[0] + return _read_first_owner( + manager, + (owner_name,), + "joint_target_pos", + span.target_q_start, + span.target_q_width, + ) + + +def _read_target_qd(manager, binding, owner_name): + span = binding.articulation_spans[0] + return _read_first_owner( + manager, + (owner_name,), + "joint_target_vel", + span.target_qd_start, + span.target_qd_width, + ) + + +def _read_generalized_force(manager, binding): + span = binding.articulation_spans[0] + return _read_first_owner( + manager, + ("_model",), + "joint_f", + span.control_start, + span.control_width, + ) + + +def _read_active_control(manager, binding): + span = binding.articulation_spans[0] + return _read_first_owner( + manager, + ("_control",), + "joint_f", + span.control_start, + span.control_width, + ) +``` + +```python +def test_articulation_rebuild_preserves_current_target_and_control(newton_world): + world, env = newton_world + art = _test_urdf_articulation(env, "arm") + mgr = get_newton_manager(world) + mgr.prepare() + ref = mgr.entity_ref(art) + binding = mgr.bind_articulations((ref,)) + q0 = np.linspace(0.01, 0.01 * binding.qpos_width, binding.qpos_width, dtype=np.float32) + qd0 = np.linspace(0.02, 0.02 * binding.qvel_width, binding.qvel_width, dtype=np.float32) + q1 = q0 + 0.4 + qd1 = qd0 + 0.5 + model_target_q = np.linspace( + -0.03, + -0.03 * binding.target_qpos_width, + binding.target_qpos_width, + dtype=np.float32, + ) + model_target_qd = np.linspace( + -0.04, + -0.04 * binding.target_qvel_width, + binding.target_qvel_width, + dtype=np.float32, + ) + control_target_q = model_target_q - 0.7 + control_target_qd = model_target_qd - 0.8 + generalized_force = np.full(binding.qvel_width, 0.3, dtype=np.float32) + active_control = np.full(binding.qvel_width, -0.6, dtype=np.float32) + _write_articulation_arrays( + mgr, + binding, + q0, + qd0, + q1, + qd1, + model_target_q, + model_target_qd, + control_target_q, + control_target_qd, + generalized_force, + active_control, + ) + _dynamic_box(env, "topology_change") + mgr.prepare() + rebound = mgr.bind_articulations((ref,)) + assert np.allclose(_read_q(mgr, rebound), q0) + assert np.allclose(_read_qd(mgr, rebound), qd0) + assert np.allclose(_read_q(mgr, rebound, "_state_1"), q1) + assert np.allclose(_read_qd(mgr, rebound, "_state_1"), qd1) + assert np.allclose( + _read_target_q(mgr, rebound, "_model"), model_target_q + ) + assert np.allclose( + _read_target_qd(mgr, rebound, "_model"), model_target_qd + ) + assert np.allclose( + _read_target_q(mgr, rebound, "_control"), control_target_q + ) + assert np.allclose( + _read_target_qd(mgr, rebound, "_control"), control_target_qd + ) + assert np.allclose( + _read_generalized_force(mgr, rebound), generalized_force + ) + assert np.allclose(_read_active_control(mgr, rebound), active_control) + + +def test_removed_articulation_ref_cannot_rebind(newton_world): + world, env = newton_world + art = _test_urdf_articulation(env, "arm") + mgr = get_newton_manager(world) + mgr.prepare() + ref = mgr.entity_ref(art) + env.remove_skeleton("arm") + mgr.prepare() + with pytest.raises(NewtonStaleBindingError, match="removed"): + mgr.bind_articulations((ref,)) + + +def test_articulation_rebuild_preserves_drive_limits_and_feedforward(newton_world): + world, env = newton_world + art = _test_urdf_articulation(env, "arm") + mgr = get_newton_manager(world) + mgr.prepare() + ref = mgr.entity_ref(art) + binding = mgr.bind_articulations((ref,)) + span = binding.articulation_spans[0] + width = span.control_width + expected = { + "joint_target_ke": np.full(width, 11.0, np.float32), + "joint_target_kd": np.full(width, 1.2, np.float32), + "joint_friction": np.full(width, 0.13, np.float32), + "joint_armature": np.full(width, 0.07, np.float32), + "joint_target_mode": np.full(width, 1, np.int32), + "joint_effort_limit": np.full(width, 9.0, np.float32), + "joint_velocity_limit": np.full(width, 4.0, np.float32), + "joint_limit_lower": np.full(width, -0.9, np.float32), + "joint_limit_upper": np.full(width, 0.9, np.float32), + } + for name, values in expected.items(): + _assign_slice(mgr._model, name, span.control_start, values) + feedforward = np.full(width, 0.23, np.float32) + _assign_slice(mgr._control, "joint_act", span.control_start, feedforward) + _dynamic_box(env, "topology_change") + mgr.prepare() + rebound = mgr.bind_articulations((ref,)).articulation_spans[0] + for name, values in expected.items(): + actual = _read_slice( + mgr._model, name, rebound.control_start, rebound.control_width + ) + assert np.allclose(actual, values) + assert np.allclose( + _read_slice( + mgr._control, + "joint_act", + rebound.control_start, + rebound.control_width, + ), + feedforward, + ) +``` + +- [ ] **Step 2: Run the articulation tests and confirm runtime data is lost** + +Expected: current rebuild either omits the articulation or loses current/target +state and control. + +- [ ] **Step 3: Implement complete articulation snapshots** + +```python +@dataclass(frozen=True, slots=True) +class ArticulationRuntimeSnapshot: + ref: NewtonEntityRef + state_0_joint_q: np.ndarray + state_0_joint_qd: np.ndarray + state_1_joint_q: np.ndarray + state_1_joint_qd: np.ndarray + model_target_joint_q: np.ndarray | None + model_target_joint_qd: np.ndarray | None + control_target_joint_q: np.ndarray | None + control_target_joint_qd: np.ndarray | None + model_joint_f: np.ndarray | None + control_joint_f: np.ndarray | None + control_joint_act: np.ndarray | None + drive_stiffness: np.ndarray + drive_damping: np.ndarray + drive_friction: np.ndarray + drive_armature: np.ndarray + drive_target_mode: np.ndarray + drive_effort_limit: np.ndarray + drive_velocity_limit: np.ndarray + joint_limit_lower: np.ndarray + joint_limit_upper: np.ndarray + root_state_0: np.ndarray + root_state_1: np.ndarray +``` + +Read/write each field through explicit spans from `ArticulationBinding`. +Model defaults and active `Control` targets/forces are captured separately; +never collapse them just because the normal setter currently writes both. +An owner field is `None` only when that Newton model/control array is genuinely +absent, and restore preserves that absence. +Capture active feed-forward control from `Control.joint_act`. Capture drive and +limit arrays from `joint_target_ke`, `joint_target_kd`, `joint_friction`, +`joint_armature`, `joint_target_mode`, `joint_effort_limit`, +`joint_velocity_limit`, `joint_limit_lower`, and `joint_limit_upper` on their +live owner (`Control` when exposed, otherwise `Model`) and restore them before +candidate validation. Each `root_state_*` is exactly 13 `float32` values: world pose in +`xyz+xyzw` followed by linear and angular velocity. Contacts are not captured; +the candidate collision pipeline regenerates them. +Canonical articulation replay reconstructs links/joints/drives into the +candidate and then refreshes existing `NewtonArticulation` wrapper metadata at +commit. Never interpret an active-joint ordinal as a flattened q/qd index. + +- [ ] **Step 4: Invalidate FK after current q writes** + +When current q is restored or written, build a boolean articulation mask with +shape `(model.articulation_count,)`. Evaluate FK independently against both +ping-pong states so their distinct q/qd histories remain distinct: + +```python +for state in (manager._state_0, manager._state_1): + eval_fk( + manager._model, + state.joint_q, + state.joint_qd, + state, + articulation_mask, + ) +``` + +Run visual synchronization only after FK is current. + +- [ ] **Step 5: Run articulation, binding, and rebuild tests** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_transactional_rebuild.py python/test/engine/newton_physics/test_newton_bindings.py python/test/engine/newton_physics/test_newton_physics_scene.py -k articulation +``` + +Expected: all selected tests pass, including spherical/free-joint span cases. + +- [ ] **Step 6: Commit articulation preservation** + +```bash +git -C /root/sources/dexsim add python/dexsim/engine/newton_physics/runtime_snapshot.py python/dexsim/engine/newton_physics/rebuild.py python/dexsim/engine/newton_physics/articulation/articulation.py python/dexsim/engine/newton_physics/articulation/skeleton_bridge.py python/test/engine/newton_physics/test_newton_transactional_rebuild.py python/test/engine/newton_physics/test_newton_bindings.py +git -C /root/sources/dexsim commit -m "feat(newton): preserve articulation state across rebuild" +``` + +--- + +### Task 6: Isolate same-device worlds and make DexSim cleanup deterministic + +**Files:** + +- Create: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_multi_world_runtime.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/newton_manager.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/registry.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/integration.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/capture_coordinator.py` + +**Interfaces:** + +- Produces: idempotent `NewtonManager.close()` and public `manager_for_entity(entity)`. +- Produces: weak device-level CUDA capture coordination without shared physics state. +- Consumes: per-world tokens and generations from Tasks 1–5. + +- [ ] **Step 1: Add failing two-world isolation and close tests** + +```python +def test_two_worlds_build_step_rebuild_and_close_independently(two_newton_worlds): + (world_a, env_a, mgr_a), (world_b, env_b, mgr_b) = two_newton_worlds + box_a = _dynamic_box(env_a, "a") + box_b = _dynamic_box(env_b, "b") + world_a.update(0.01) + world_b.update(0.02) + assert mgr_a.world_token != mgr_b.world_token + assert mgr_a.model_generation == mgr_b.model_generation == 1 + assert mgr_a._model is not mgr_b._model + _dynamic_box(env_a, "a2") + world_a.update(0.01) + assert mgr_a.model_generation == 2 + assert mgr_b.model_generation == 1 + mgr_a.close() + mgr_a.close() + with pytest.raises(NewtonClosedError): + mgr_a.bind_rigid_entities((mgr_a.entity_ref(box_a),)) + assert manager_for_entity(box_b) is mgr_b + world_b.update(0.02) + + +def test_capture_coordinator_holds_only_weak_manager_refs(two_cuda_worlds): + (_, _, mgr_a), (_, _, mgr_b) = two_cuda_worlds + coordinator = capture_coordinator_for_device(mgr_a.device) + assert coordinator.manager_count == 2 + mgr_a.close() + assert mgr_a not in tuple(coordinator.managers) + assert mgr_b in tuple(coordinator.managers) + assert coordinator.manager_count == 1 +``` + +- [ ] **Step 2: Run the file; confirm leakage/cross-world assumptions fail** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_multi_world_runtime.py +``` + +Expected: failures expose absent close/owner lookup or strong global state. + +- [ ] **Step 3: Implement deterministic per-world teardown** + +`NewtonManager.close()` sets closed once, invalidates graphs/caches, clears +callbacks and subscriptions, releases model/state/control/contact/solver and +renderer resources, unregisters arena/entity ownership, and removes only this +manager from weak coordination. Every public method calls `_assert_open()`. + +Expose owner lookup without leaking private registries: + +```python +def manager_for_entity(entity) -> NewtonManager | None: + arena = entity.get_arena() + return manager_for_arena(arena) +``` + +The CUDA coordinator stores `weakref.WeakSet[NewtonManager]`, serializes only +capture operations for a device, and has a finite diagnostic timeout. It owns +no builder/model/state/control/solver arrays. Expose a read-only `managers` +tuple and derived `manager_count` for diagnostics/tests. + +- [ ] **Step 4: Run the DexSim Stage 1 suite** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_public_contract.py python/test/engine/newton_physics/test_newton_bindings.py python/test/engine/newton_physics/test_newton_transactional_rebuild.py python/test/engine/newton_physics/test_newton_multi_world_runtime.py python/test/engine/newton_physics/test_newton_scene_lifecycle.py python/test/engine/newton_physics/test_newton_scene_mutations.py python/test/engine/newton_physics/test_newton_sim_index.py python/test/engine/newton_physics/test_newton_physics_scene.py +``` + +Expected: zero failures and no surviving per-world registrations after fixture +teardown. + +- [ ] **Step 5: Commit world isolation and cleanup** + +```bash +git -C /root/sources/dexsim add python/dexsim/engine/newton_physics/newton_manager.py python/dexsim/engine/newton_physics/registry.py python/dexsim/engine/newton_physics/integration.py python/dexsim/engine/newton_physics/capture_coordinator.py python/test/engine/newton_physics/test_newton_multi_world_runtime.py +git -C /root/sources/dexsim commit -m "fix(newton): isolate and close per-world runtimes" +``` + +--- + +### Task 7: Add EmbodiChain API handshake, structured capabilities, and scene context + +**Files:** + +- Create: `embodichain/lab/sim/physics/context.py` +- Create: `tests/sim/newton_contract_test_utils.py` +- Create: `tests/sim/test_newton_scene_context.py` +- Modify: `pyproject.toml` +- Modify: `embodichain/lab/sim/cfg.py` +- Modify: `embodichain/lab/sim/physics/base.py` +- Modify: `embodichain/lab/sim/physics/default.py` +- Modify: `embodichain/lab/sim/physics/newton.py` +- Modify: `embodichain/lab/sim/physics/__init__.py` +- Modify: `tests/sim/test_backend_parity.py` + +**Interfaces:** + +- Produces: `PhysicsCapabilities`, `PhysicsPrepareResult`, `BackendSceneContext`. +- Produces: exact DexSim package/API validation during Newton activation. +- Consumes: DexSim public contract from Tasks 1–6. + +- [ ] **Step 1: Add shared EmbodiChain contract test fixtures** + +Create `tests/sim/newton_contract_test_utils.py`; each new Stage 1 test file +imports the fixtures and factories it uses: + +```python +ART_PATH = "SlidingBoxDrawer/SlidingBoxDrawer.urdf" + + +def box_cfg(uid: str, z: float = 1.0) -> RigidObjectCfg: + return RigidObjectCfg.from_dict( + { + "uid": uid, + "shape": {"shape_type": "Cube", "size": [0.1, 0.1, 0.1]}, + "attrs": {"mass": 1.0}, + "body_type": "dynamic", + "init_pos": (0.0, 0.0, z), + } + ) + + +def arm_cfg(uid: str) -> ArticulationCfg: + return ArticulationCfg.from_dict( + { + "uid": uid, + "fpath": get_data_path(ART_PATH), + "drive_pros": {"drive_type": "force"}, + } + ) + + +@pytest.fixture +def newton_sim(): + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + device="cpu", + num_envs=2, + physics_cfg=NewtonPhysicsCfg( + device="cpu", num_substeps=2, use_cuda_graph=False + ), + ) + ) + try: + yield sim + finally: + close = getattr(sim, "close", None) + if close is None: + sim.destroy(exit_process=False) + else: + close() + + +@pytest.fixture +def fake_manager(): + identity = np.eye(4, dtype=np.float32) + rotated = np.eye(4, dtype=np.float32) + rotated[:3, :3] = Rotation.from_euler("z", 90, degrees=True).as_matrix() + rotated[:3, 3] = [2.0, 3.0, 0.0] + + class Root: + def __init__(self, pose): + self.pose = pose + def get_world_pose(self): + return self.pose.copy() + + class Arena: + def __init__(self, pose): + self.root = Root(pose) + def get_root_node(self): + return self.root + + class Manager: + def __init__(self): + self._arenas = [Arena(identity), Arena(rotated)] + self.device = torch.device("cpu") + self.is_closed = False + + return Manager() +``` + +- [ ] **Step 2: Add failing version, capability, and context tests** + +```python +def test_newton_backend_requires_exact_contract(monkeypatch): + monkeypatch.setattr(dexsim, "__version__", "0.4.3") + backend = NewtonPhysicsBackend(SimpleNamespace()) + with pytest.raises(RuntimeError, match="0.4.4"): + backend._validate_integration_contract() + + +def test_capabilities_are_structured(): + caps = NewtonPhysicsBackend(SimpleNamespace()).capabilities + assert caps.runtime_topology_mutation == frozenset({"rigid", "articulation"}) + assert caps.multi_world is True + assert caps.articulation_acceleration is False + assert "soft_body" not in caps.asset_kinds + assert "cloth" not in caps.asset_kinds + + +def test_scene_context_uses_full_arena_transform(fake_manager): + context = BackendSceneContext(fake_manager) + transforms = context.arena_transforms + assert transforms.shape == (2, 4, 4) + local = torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]]) + world = context.local_to_world(local, torch.tensor([1])) + roundtrip = context.world_to_local(world, torch.tensor([1])) + assert torch.allclose(roundtrip, local, atol=1e-6) + assert not torch.allclose(world[:, :3], local[:, :3] + transforms[1, :3, 3]) +``` + +The fake second arena must contain both non-zero translation and a 90-degree Z +rotation so the last assertion detects translation-only conversion. + +- [ ] **Step 3: Run the pure-Python tests and confirm missing types/validation** + +```bash +pytest -q tests/sim/test_backend_parity.py tests/sim/test_newton_scene_context.py -m "not requires_sim" +``` + +Expected: missing capability/context contracts. + +- [ ] **Step 4: Add structured backend contracts** + +```python +@dataclass(frozen=True, slots=True) +class PhysicsCapabilities: + asset_kinds: frozenset[str] + runtime_topology_mutation: frozenset[str] + solver_gradients: frozenset[str] + cuda_graph: bool + partial_reset: bool + forward_kinematics: bool + heterogeneous_joint_spans: bool + runtime_collision_filter: bool + contact_sensor: bool + articulation_acceleration: bool + multi_world: bool + + +@dataclass(frozen=True, slots=True) +class PhysicsPrepareResult: + generation: int | None + did_build: bool + did_rebuild: bool + added_entities: tuple[object, ...] = () + removed_entities: tuple[object, ...] = () +``` + +Keep every existing `supports_*` property as a wrapper over `capabilities`. +The default backend returns `generation=None` and preserves existing behavior. +Add abstract/default implementations for `model_generation`, +`queue_initialization(obj)`, `prepare() -> PhysicsPrepareResult`, and +idempotent `close()` so later tasks do not branch on backend names. + +- [ ] **Step 5: Add strict config and dependency validation** + +Change the dependency to: + +```toml +"dexsim_engine==0.4.4", +``` + +Validate positive `physics_dt`, positive `num_substeps`, normalized device, +recognized solver parameters, gradient/solver compatibility, broad phase, and +CUDA graph combinations before world construction. On backend activation, +require both `dexsim.__version__ == "0.4.4"` and +`NEWTON_INTEGRATION_API_VERSION == 2`. Accept local source builds whose public +version is `0.4.4+` by comparing +`packaging.version.Version(dexsim.__version__).base_version` to `"0.4.4"`; +reject every other base version. + +- [ ] **Step 6: Implement explicit owner context** + +```python +class BackendSceneContext: + def __init__(self, manager: SimulationManager) -> None: + self._manager_ref = weakref.ref(manager) + + @property + def manager(self) -> SimulationManager: + manager = self._manager_ref() + if manager is None or manager.is_closed: + raise RuntimeError("BackendSceneContext owner is closed.") + return manager + + @property + def world(self): + return self.manager.get_world() + + @property + def scene(self): + return self.manager.physics.get_scene() + + @property + def physics(self): + return self.manager.physics + + @property + def generation(self) -> int | None: + return self.manager.physics.model_generation +``` + +Build device-resident `(N, 4, 4)` world transforms and inverses from each +arena root node's full world pose. Provide batched `local_to_world()` and +`world_to_local()` for pose tensors in `xyzw`. Maintain weak maps from arena +and entity native handles to contexts. `register_entities()` and +`for_entities()` provide the source-compatible constructor fallback without +consulting a default world or default SimulationManager; unknown external +entities raise an ownership error that tells the caller to pass `context=`. +Expose `register_entity(entity, ref=None)`, `register_entities(entities)`, and +`entity_ref(entity)`; the last method returns the DexSim stable reference +recorded during Newton attachment. + +Use homogeneous composition for both conversion directions: + +```python +def _xyzw_pose_to_matrix(pose: torch.Tensor) -> torch.Tensor: + matrix = torch.eye(4, dtype=pose.dtype, device=pose.device).repeat( + pose.shape[0], 1, 1 + ) + matrix[:, :3, 3] = pose[:, :3] + matrix[:, :3, :3] = matrix_from_quat( + convert_quat(pose[:, 3:7], to="wxyz") + ) + return matrix + + +def _matrix_to_xyzw_pose(matrix: torch.Tensor) -> torch.Tensor: + quat = convert_quat(quat_from_matrix(matrix[:, :3, :3]), to="xyzw") + return torch.cat((matrix[:, :3, 3], quat), dim=-1) + + +def local_to_world(self, pose: torch.Tensor, env_ids: torch.Tensor) -> torch.Tensor: + local = _xyzw_pose_to_matrix(pose) + world = torch.bmm(self.arena_transforms[env_ids.long()], local) + return _matrix_to_xyzw_pose(world) + + +def world_to_local(self, pose: torch.Tensor, env_ids: torch.Tensor) -> torch.Tensor: + world = _xyzw_pose_to_matrix(pose) + local = torch.bmm(self.inverse_arena_transforms[env_ids.long()], world) + return _matrix_to_xyzw_pose(local) +``` + +- [ ] **Step 7: Run config/context/capability tests** + +```bash +pytest -q tests/sim/test_backend_parity.py tests/sim/test_newton_scene_context.py tests/sim/test_physics_attrs.py +``` + +Expected: all pass without creating a real simulation for pure contract cases. + +- [ ] **Step 8: Commit the EmbodiChain contract foundation** + +```bash +git add pyproject.toml embodichain/lab/sim/cfg.py embodichain/lab/sim/physics/context.py embodichain/lab/sim/physics/base.py embodichain/lab/sim/physics/default.py embodichain/lab/sim/physics/newton.py embodichain/lab/sim/physics/__init__.py tests/sim/newton_contract_test_utils.py tests/sim/test_backend_parity.py tests/sim/test_newton_scene_context.py +git commit -m "refactor(sim): add explicit physics scene contracts" +``` + +--- + +### Task 8: Route EmbodiChain spawning and objects through explicit ownership + +**Files:** + +- Modify: `embodichain/lab/sim/common.py` +- Modify: `embodichain/lab/sim/sim_manager.py` +- Modify: `embodichain/lab/sim/utility/sim_utils.py` +- Modify: `embodichain/lab/sim/objects/rigid_object.py` +- Modify: `embodichain/lab/sim/objects/articulation.py` +- Modify: `embodichain/lab/sim/objects/robot.py` +- Modify: `embodichain/lab/sim/objects/backends/default.py` +- Modify: `tests/sim/test_newton_scene_context.py` +- Modify: `tests/sim/test_newton_finalize_lifecycle.py` + +**Interfaces:** + +- Produces: manager-owned context passed into all physical objects/views. +- Produces: explicit pending-initialization queue; constructors never invoke overridable `reset()`. +- Consumes: `BackendSceneContext` and DexSim `attach_rigid_body`. + +- [ ] **Step 1: Add failing no-global and no-constructor-reset tests** + +```python +def test_rigid_and_articulation_construction_do_not_use_default_world( + monkeypatch, newton_sim +): + monkeypatch.setattr(dexsim, "default_world", lambda: (_ for _ in ()).throw(AssertionError("global"))) + rigid = newton_sim.add_rigid_object(box_cfg("box")) + art = newton_sim.add_articulation(arm_cfg("arm")) + assert rigid.context is newton_sim.scene_context + assert art.context is newton_sim.scene_context + + +def test_physical_object_context_keyword_is_source_compatible(): + assert inspect.signature(RigidObject).parameters["context"].default is None + assert inspect.signature(Articulation).parameters["context"].default is None + + +def test_batch_entity_constructor_never_calls_virtual_reset(): + class Probe(BatchEntity): + def reset(self, env_ids=None): + raise AssertionError("virtual reset from base constructor") + def set_local_pose(self, pose, env_ids=None): + return None + def get_local_pose(self, to_matrix=False): + return torch.zeros(1, 7) + Probe( + cfg=ObjectBaseCfg(uid="probe"), + entities=[object()], + device=torch.device("cpu"), + ) +``` + +- [ ] **Step 2: Run focused tests and confirm global lookup/base reset failures** + +```bash +pytest -q tests/sim/test_newton_scene_context.py tests/sim/test_newton_finalize_lifecycle.py +``` + +- [ ] **Step 3: Remove base virtual reset and pass context from the manager** + +`BatchEntity.__init__` stores fields only. Keep the `auto_reset` keyword for +source compatibility, but ignore it and emit a one-time deprecation warning +when `True`; no virtual method is called. + +Create `self.scene_context = BackendSceneContext(self)` immediately after +backend activation and arena construction. Pass it explicitly: + +```python +rigid_obj = RigidObject( + cfg=cfg, + entities=obj_list, + device=self.device, + context=self.scene_context, +) +self.physics.queue_initialization(rigid_obj) +``` + +Use the same pattern for articulations and robots. Default-only light and rigid +group constructors call their own `reset()` explicitly after all subclass +fields are initialized, preserving current behavior. + +Add `context: BackendSceneContext | None = None` as the final keyword to rigid, +articulation, and robot constructors. Resolve `None` with +`BackendSceneContext.for_entities(entities)`. The manager registers spawned +entities against its context before constructing their wrapper, so existing +positional constructor calls retain their signature and no global owner lookup +is needed. + +- [ ] **Step 4: Replace EmbodiChain private Newton attachment** + +In `_attach_newton_rigidbody_desc`, retain EmbodiChain descriptor resolution +and warnings, then call only: + +```python +manager = context.physics.newton_manager +entity_ref = manager.attach_rigid_body( + obj, + actor_type=body_type, + shape_type=shape_type, + body_desc=body, + shape_desc=shape, +) +context.register_entity(obj, entity_ref) +``` + +Delete imports of `register_mesh_object_to_newton_patch`, +`_get_entity_native_handle`, writes to `mgr.dexsim_meta`, and hard-coded world +`-1`. The standard legacy `add_rigidbody` route stores the returned reference +through the same context registration method. + +Thread `context` through `load_mesh_objects_from_cfg`, +`spawn_rigid_object_entities`, `spawn_articulation_entities`, and +`spawn_usd_articulation_entities`. Compatibility defaults resolve from the +provided entities; the SimulationManager core path always passes its context. +Replace `_is_newton_backend_active()`, `_newton_solver_type()`, and +`get_dexsim_arenas()` use in physical spawn/object paths with context +properties. Object `destroy()` methods remove entities from their owning +context arenas, never from a default world. + +- [ ] **Step 5: Run spawn and default-backend regressions** + +```bash +pytest -q tests/sim/test_newton_scene_context.py tests/sim/test_newton_finalize_lifecycle.py +pytest -q tests/sim/objects/test_rigid_object.py tests/sim/objects/test_articulation.py -k "constructor or spawn or desc_native" +``` + +Expected: selected tests pass; no core object construction resolves a default +world. + +- [ ] **Step 6: Commit explicit ownership and spawning** + +```bash +git add embodichain/lab/sim/common.py embodichain/lab/sim/sim_manager.py embodichain/lab/sim/utility/sim_utils.py embodichain/lab/sim/objects/rigid_object.py embodichain/lab/sim/objects/articulation.py embodichain/lab/sim/objects/robot.py embodichain/lab/sim/objects/backends/default.py tests/sim/test_newton_scene_context.py tests/sim/test_newton_finalize_lifecycle.py +git commit -m "refactor(sim): make physical object ownership explicit" +``` + +--- + +### Task 9: Rebind rigid views by generation and initialize only new objects + +**Files:** + +- Create: `tests/sim/test_newton_rebuild_bindings.py` +- Modify: `embodichain/lab/sim/physics/newton.py` +- Modify: `embodichain/lab/sim/objects/backends/newton.py` +- Modify: `embodichain/lab/sim/objects/rigid_object.py` +- Modify: `tests/sim/objects/test_rigid_object.py` +- Modify: `tests/sim/test_newton_finalize_lifecycle.py` + +**Interfaces:** + +- Produces: `NewtonRigidBodyView._ensure_binding()` with O(1) generation check. +- Produces: event-driven cache invalidation and exact pending initialization. +- Consumes: DexSim `RigidEntityBinding`, prepare result, and rebuild event. + +- [ ] **Step 1: Add failing rebind/state-preservation/initialization tests** + +```python +def test_rigid_view_refreshes_once_after_generation_change(newton_sim): + old = newton_sim.add_rigid_object(box_cfg("old", z=1.0)) + first = newton_sim.physics.prepare() + old_generation = old._data.body_view.binding.generation + old_pose = old.get_local_pose().clone() + new = newton_sim.add_rigid_object(box_cfg("new", z=2.0)) + second = newton_sim.physics.prepare() + assert second.generation == first.generation + 1 + assert old_generation == first.generation + assert old._data.body_view.binding.generation == second.generation + assert new._data.body_view.binding.generation == second.generation + assert torch.allclose(old.get_local_pose(), old_pose, atol=1e-5) + assert torch.allclose(new.get_local_pose()[:, 2], torch.tensor([2.0] * new.num_instances)) + + +def test_existing_object_is_not_reset_on_rebuild(newton_sim, mocker): + old = newton_sim.add_rigid_object(box_cfg("old")) + newton_sim.physics.prepare() + reset = mocker.spy(old, "reset") + newton_sim.add_rigid_object(box_cfg("new")) + newton_sim.physics.prepare() + reset.assert_not_called() +``` + +- [ ] **Step 2: Run the new file and observe stale IDs/global reset** + +```bash +pytest -q tests/sim/test_newton_rebuild_bindings.py +``` + +Expected: stale binding or `_reset_entities_after_finalize` resets existing +objects. + +- [ ] **Step 3: Replace permanent IDs with one binding per view** + +```python +def _ensure_binding(self) -> RigidEntityBinding: + generation = self._context.generation + if self._binding is None or self._binding.generation != generation: + refs = tuple(self._context.entity_ref(entity) for entity in self._entities) + self._binding = self._manager.bind_rigid_entities(refs, device=self._device) + self._invalidate_derived_caches() + self._binding.assert_current(self._manager) + return self._binding + +@property +def binding(self) -> RigidEntityBinding: + return self._ensure_binding() +``` + +Every fetch/apply operation calls this once and passes its batched device IDs +to `NewtonPhysicsScene`. Remove permanent `_body_ids`, sorted-ID, and XY-offset +caches; rebuild derived caches only inside `_invalidate_derived_caches()`. + +- [ ] **Step 4: Subscribe backend lifecycle and initialize only prepare additions** + +`NewtonPhysicsBackend.activate()` subscribes to model rebuilt events and marks +registered views dirty. `queue_initialization(obj)` stores object identity and +its stable refs. After `prepare()` and view rebind, initialize only queued +objects whose refs appear in `result.added_entities`, then remove them from the +queue. First build follows the same path. Existing objects are never reset by a +rebuild. + +- [ ] **Step 5: Convert rigid poses with full context transforms** + +`fetch_pose()` converts DexSim world poses to arena-local with +`context.world_to_local`; `apply_pose()` converts local to world with +`context.local_to_world`. Remove all `[:2, 3]`, XY-only, and +`get_all_arenas()` conversion logic from the Newton view. + +- [ ] **Step 6: Run rigid and lifecycle tests** + +```bash +pytest -q tests/sim/test_newton_rebuild_bindings.py tests/sim/test_newton_finalize_lifecycle.py tests/sim/objects/test_rigid_object.py -k "Newton or newton or local_pose or reset" +``` + +Expected: selected tests pass for 1, 2, and 8 arenas, including a rotated +arena fixture. + +- [ ] **Step 7: Commit rigid rebinding** + +```bash +git add embodichain/lab/sim/physics/newton.py embodichain/lab/sim/objects/backends/newton.py embodichain/lab/sim/objects/rigid_object.py tests/sim/test_newton_rebuild_bindings.py tests/sim/objects/test_rigid_object.py tests/sim/test_newton_finalize_lifecycle.py +git commit -m "fix(newton): rebind rigid views after runtime rebuild" +``` + +--- + +### Task 10: Rebind articulation views, separate q/qd widths, and enforce FK/capabilities + +**Files:** + +- Modify: `embodichain/lab/sim/objects/backends/newton.py` +- Modify: `embodichain/lab/sim/objects/articulation.py` +- Modify: `embodichain/lab/sim/objects/robot.py` +- Modify: `tests/sim/test_newton_rebuild_bindings.py` +- Modify: `tests/sim/objects/test_articulation.py` +- Modify: `tests/sim/objects/test_rigid_object.py` +- Modify: `tests/sim/objects/test_robot.py` + +**Interfaces:** + +- Produces: `NewtonArticulationView._ensure_binding()` and real `compute_kinematics()`. +- Produces: `_articulation_buffer_shapes(num_instances, qpos_width, qvel_width, target_qpos_width, target_qvel_width)` used by `ArticulationData`. +- Produces: explicit unsupported-operation errors for q acceleration and other absent capabilities. +- Consumes: DexSim `ArticulationBinding` and full context transforms. + +- [ ] **Step 1: Add failing articulation rebind, frame, width, and FK tests** + +```python +def test_articulation_view_rebinds_and_preserves_targets(newton_sim): + arm = newton_sim.add_articulation(arm_cfg("arm")) + newton_sim.physics.prepare() + q = torch.full((arm.num_instances, arm.qpos_width), 0.1, device=arm.device) + qd = torch.full((arm.num_instances, arm.qvel_width), 0.2, device=arm.device) + target_q = torch.full( + (arm.num_instances, arm.target_qpos_width), + -0.1, + device=arm.device, + ) + target_qd = torch.full( + (arm.num_instances, arm.target_qvel_width), + -0.2, + device=arm.device, + ) + arm.set_qpos(q, target=False) + arm.set_qvel(qd, target=False) + arm.set_qpos(target_q, target=True) + arm.set_qvel(target_qd, target=True) + before_link = arm.get_link_pose(arm.link_names[-1]).clone() + newton_sim.add_rigid_object(box_cfg("topology_change")) + newton_sim.physics.prepare() + assert torch.allclose(arm.get_qpos(), q) + assert torch.allclose(arm.get_qvel(), qd) + assert torch.allclose(arm.get_qpos(target=True), target_q) + assert torch.allclose(arm.get_qvel(target=True), target_qd) + assert torch.allclose(arm.get_link_pose(arm.link_names[-1]), before_link, atol=1e-5) + + +def test_qpos_write_updates_link_pose_without_physics_step(newton_sim): + arm = newton_sim.add_articulation(arm_cfg("arm")) + newton_sim.physics.prepare() + before = arm.get_link_pose(arm.link_names[-1]).clone() + q = arm.get_qpos().clone() + q[:, 0] += 0.2 + arm.set_qpos(q, target=False) + after = arm.get_link_pose(arm.link_names[-1]) + assert not torch.allclose(after, before) + + +def test_newton_qacc_is_explicitly_unsupported(newton_sim): + arm = newton_sim.add_articulation(arm_cfg("arm")) + newton_sim.physics.prepare() + with pytest.raises(NotImplementedError, match="acceleration"): + _ = arm.body_data.qacc +``` + +Add pure allocation coverage for distinct widths: + +```python +def test_articulation_buffer_shapes_keep_q_and_qd_widths_distinct(): + shapes = _articulation_buffer_shapes( + num_instances=3, + qpos_width=7, + qvel_width=6, + target_qpos_width=6, + target_qvel_width=6, + ) + assert shapes["qpos"] == (3, 7) + assert shapes["target_qpos"] == (3, 6) + assert shapes["qvel"] == (3, 6) + assert shapes["target_qvel"] == (3, 6) + assert shapes["qf"] == (3, 6) +``` + +- [ ] **Step 2: Run the selected articulation tests and confirm stale/FK/zero-qacc failures** + +```bash +pytest -q tests/sim/test_newton_rebuild_bindings.py tests/sim/objects/test_articulation.py -k "Newton or newton or qpos or qacc or link_pose" +``` + +- [ ] **Step 3: Bind articulation IDs, link IDs, and spans as one unit** + +Implement the same O(1) generation guard as the rigid view. Resolve root/link +IDs and joint spans only through `manager.bind_articulations`; delete direct +reads of `dexsim_meta_links`, `get_gpu_index()`, and permanent articulation ID +lists from the Newton view. + +Expose and use separate widths: + +```python +@property +def qpos_width(self) -> int: + return self._ensure_binding().qpos_width + +@property +def qvel_width(self) -> int: + return self._ensure_binding().qvel_width + +@property +def target_qpos_width(self) -> int: + return self._ensure_binding().target_qpos_width + +@property +def target_qvel_width(self) -> int: + return self._ensure_binding().target_qvel_width +``` + +Allocate current `qpos` with `qpos_width`, current `qvel/qf` with +`qvel_width`, and targets with their explicit binding widths. For DexSim 0.4.4 +both target position and target velocity are per-DOF and therefore their +widths equal `qvel_width`, but callers consume the explicit properties rather +than inferring that relationship. Existing `dof` remains a compatibility alias +for all-1-DOF assets and raises a clear error when its old ambiguous meaning +would truncate a non-scalar joint. +Expose matching read-only `Articulation.qpos_width` and +`Articulation.qvel_width` properties plus `target_qpos_width` and +`target_qvel_width` that delegate to the view. + +```python +def _articulation_buffer_shapes( + num_instances: int, + qpos_width: int, + qvel_width: int, + target_qpos_width: int, + target_qvel_width: int, +) -> dict[str, tuple[int, int]]: + return { + "qpos": (num_instances, qpos_width), + "target_qpos": (num_instances, target_qpos_width), + "qvel": (num_instances, qvel_width), + "target_qvel": (num_instances, target_qvel_width), + "qf": (num_instances, qvel_width), + } +``` + +- [ ] **Step 4: Implement frame-correct root/link reads and FK** + +Convert root and link world poses with the complete arena transform. On current +q writes, call the DexSim public FK invalidation/evaluation path for the +affected articulation IDs. Implement `compute_kinematics(env_ids)` by mapping +the selected environment rows through the binding's `articulation_ids_host`, +constructing a boolean mask of length `model.articulation_count`, and calling +the public manager FK method; it must not be a no-op. + +Replace fabricated q acceleration with: + +```python +raise NotImplementedError( + "Newton articulation joint acceleration is not exposed by DexSim 0.4.4." +) +``` + +Use the same explicit pattern for unsupported runtime collision-filter or +sensor operations; do not return plausible zeros or success. + +Add `@pytest.mark.gpu` to the Newton-backed rigid, articulation, and robot test +classes because they configure `device="cuda"` even though their node IDs do +not contain `cuda`. This keeps them out of the CPU job and includes them in the +serial `--run-gpu -m gpu` merge gate. + +- [ ] **Step 5: Run articulation and robot regressions** + +```bash +pytest -q tests/sim/test_newton_rebuild_bindings.py tests/sim/objects/test_articulation.py tests/sim/objects/test_robot.py -k "Newton or newton" +``` + +Expected: Newton articulation/robot selections pass; skips only correspond to +capabilities explicitly outside Stage 1. + +- [ ] **Step 6: Commit articulation rebinding** + +```bash +git add embodichain/lab/sim/objects/backends/newton.py embodichain/lab/sim/objects/articulation.py embodichain/lab/sim/objects/robot.py tests/sim/test_newton_rebuild_bindings.py tests/sim/objects/test_articulation.py tests/sim/objects/test_rigid_object.py tests/sim/objects/test_robot.py +git commit -m "fix(newton): bind articulation state by model generation" +``` + +--- + +### Task 11: Complete SimulationManager mutation, exact update, multi-instance, and close lifecycle + +**Files:** + +- Create: `tests/sim/test_newton_multi_manager.py` +- Modify: `embodichain/lab/sim/sim_manager.py` +- Modify: `embodichain/lab/sim/physics/base.py` +- Modify: `embodichain/lab/sim/physics/default.py` +- Modify: `embodichain/lab/sim/physics/newton.py` +- Modify: `tests/sim/test_newton_finalize_lifecycle.py` +- Modify: `tests/sim/test_newton_rebuild_bindings.py` + +**Interfaces:** + +- Produces: idempotent `SimulationManager.close()` and close-before-remove `reset()`. +- Produces: exact requested-step update after prepare. +- Consumes: pending initialization and per-world DexSim close. + +- [ ] **Step 1: Add failing update/remove/close/two-manager tests** + +```python +def test_update_runs_exact_requested_steps_after_rebuild(newton_sim): + box = newton_sim.add_rigid_object(box_cfg("box", z=1.0)) + manager = newton_sim.newton_manager + before = manager._sim_time + newton_sim.update(physics_dt=0.01, step=3) + assert manager._sim_time == pytest.approx(before + 0.03) + newton_sim.add_rigid_object(box_cfg("new", z=2.0)) + before = manager._sim_time + newton_sim.update(physics_dt=0.01, step=2) + assert manager._sim_time == pytest.approx(before + 0.02) + + +def test_remove_invalidates_and_survivor_rebinds(newton_sim): + keep = newton_sim.add_rigid_object(box_cfg("keep")) + remove = newton_sim.add_rigid_object(box_cfg("remove")) + newton_sim.physics.prepare() + keep_pose = keep.get_local_pose().clone() + assert newton_sim.remove_asset("remove") is True + result = newton_sim.physics.prepare() + assert result.did_rebuild is True + assert torch.allclose(keep.get_local_pose(), keep_pose, atol=1e-5) + with pytest.raises(Exception, match="removed|closed|stale"): + remove.get_local_pose() + + +def test_close_and_reset_are_idempotent(newton_sim): + instance_id = newton_sim.instance_id + world = newton_sim.get_world() + newton_sim.close() + newton_sim.close() + assert newton_sim.is_closed is True + assert SimulationManager.is_instantiated(instance_id) is False + SimulationManager.reset(instance_id) + assert dexsim.engine.newton_physics.get_newton_manager(world) is None + + +@pytest.mark.gpu +def test_two_same_device_managers_are_isolated(): + def make_sim(): + return SimulationManager( + SimulationManagerCfg( + headless=True, + device="cuda:0", + num_envs=2, + physics_cfg=NewtonPhysicsCfg( + device="cuda:0", use_cuda_graph=False + ), + ) + ) + + first = make_sim() + second = make_sim() + try: + first_box = first.add_rigid_object(box_cfg("first")) + second_box = second.add_rigid_object(box_cfg("second")) + first.update(physics_dt=0.01, step=1) + second.update(physics_dt=0.01, step=1) + assert first.get_world() is not second.get_world() + assert first.get_physics_scene() is not second.get_physics_scene() + assert first.scene_context is not second.scene_context + assert first.newton_manager.world_token != second.newton_manager.world_token + first.add_rigid_object(box_cfg("first_new")) + first.update(physics_dt=0.01, step=1) + assert first.newton_manager.model_generation == 2 + assert second.newton_manager.model_generation == 1 + first.close() + second.update(physics_dt=0.01, step=1) + second_pose = second_box.get_local_pose().clone() + assert torch.isfinite(second_pose).all() + assert first_box.context is not second_box.context + finally: + first.close() + second.close() +``` + +- [ ] **Step 2: Run the lifecycle files and confirm current reset/leak behavior** + +```bash +pytest -q tests/sim/test_newton_finalize_lifecycle.py tests/sim/test_newton_rebuild_bindings.py tests/sim/test_newton_multi_manager.py --run-gpu +``` + +- [ ] **Step 3: Return prepare results and preserve exact update count** + +Change `PhysicsBackend.prepare()` and `ensure_initialized()` to return +`PhysicsPrepareResult`. `SimulationManager.update()` calls prepare once, then +executes the existing world update loop exactly `step` times. It never adds a +warmup step and never drops the first requested step. + +- [ ] **Step 4: Make every topology mutation invalidate and every removal close its view** + +After successful rigid/articulation/robot add or remove, call +`physics.invalidate()`. Removal destroys the wrapper, unregisters its context +refs, and leaves surviving refs queued for rebind but not reset. Soft/cloth +add/remove on Newton raises `NotImplementedError` before mutating registries. + +- [ ] **Step 5: Add idempotent close and safe registry allocation** + +```python +def close(self) -> None: + if self._is_closed: + return + self._is_closed = True + first_error = None + try: + self.wait_window_record_saves() + self.physics.close() + except Exception as exc: + first_error = exc + try: + if self._world is not None: + self._world.quit() + except Exception as exc: + if first_error is None: + first_error = exc + finally: + self._instances.pop(self.instance_id, None) + if first_error is not None: + raise first_error +``` + +`reset(instance_id)` calls `instance.close()` before removing it. Preserve +`destroy(exit_process=...)` as a wrapper around close plus its documented +process-exit policy. Allocate instance IDs monotonically rather than from +`len(_instances)`, so closing a non-last manager cannot overwrite a live entry. +DexSim manager close is idempotent, so the backend-close/world-quit sequence is +safe even when `World.quit()` invokes the same integration teardown again. + +- [ ] **Step 6: Run lifecycle, multi-manager, and default-backend tests** + +```bash +pytest -q tests/sim/test_newton_finalize_lifecycle.py tests/sim/test_newton_rebuild_bindings.py tests/sim/test_newton_multi_manager.py tests/sim/test_backend_parity.py tests/sim/objects/test_rigid_object.py tests/sim/objects/test_articulation.py --run-gpu +``` + +Expected: zero failures; fixture teardown finds no stale world or manager +registration. + +- [ ] **Step 7: Commit manager lifecycle completion** + +```bash +git add embodichain/lab/sim/sim_manager.py embodichain/lab/sim/physics/base.py embodichain/lab/sim/physics/default.py embodichain/lab/sim/physics/newton.py tests/sim/test_newton_finalize_lifecycle.py tests/sim/test_newton_rebuild_bindings.py tests/sim/test_newton_multi_manager.py +git commit -m "fix(sim): close and rebuild Newton managers safely" +``` + +--- + +### Task 12: Document the contract and run the Stage 1 merge gate + +**Files:** + +- Modify: `docs/source/overview/sim/sim_manager.md` +- Modify: `design/newton-backend-design.md` +- Verify: both repositories' Stage 1 diffs. + +**Interfaces:** + +- Consumes: all Stage 1 interfaces and tests. +- Produces: a verified foundation for the later Stage 2 differentiable plan. + +- [ ] **Step 1: Update public documentation with executable examples** + +Document: + +```python +sim = SimulationManager( + SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg(device="cuda:0"), + num_envs=4, + headless=True, + ) +) +try: + cube = sim.add_rigid_object(cube_cfg) + sim.finalize_newton_physics() + sim.update(step=1) +finally: + sim.close() +``` + +State that prepare/finalize does not advance time, add/remove of rigid bodies +and articulations rebuilds transactionally, old runtime IDs must not be cached, +soft/cloth topology mutation is unsupported, q acceleration is unavailable, +and two managers are isolated. Mark the old Target 4 implementation claims as +superseded by the 2026-07-13 design and record only tests actually passing. + +- [ ] **Step 2: Run DexSim formatting and the full Newton test directory** + +```bash +cd /root/sources/dexsim +black --check --diff python/dexsim/engine/newton_physics python/test/engine/newton_physics +pytest -q python/test/engine/newton_physics +``` + +Expected: Black exits zero and pytest reports zero failures. + +- [ ] **Step 3: Run EmbodiChain focused CPU/headless tests** + +```bash +cd /root/sources/EmbodiChain +pytest -q tests/sim/test_backend_parity.py tests/sim/test_physics_attrs.py tests/sim/test_newton_scene_context.py tests/sim/test_newton_finalize_lifecycle.py +``` + +Expected: zero failures. + +- [ ] **Step 4: Run EmbodiChain serial GPU Newton integration tests** + +```bash +pytest -q tests/sim/test_newton_rebuild_bindings.py tests/sim/test_newton_multi_manager.py tests/sim/objects/test_rigid_object.py tests/sim/objects/test_articulation.py tests/sim/objects/test_robot.py --run-gpu -m gpu +``` + +Expected: zero failures; skips are listed and checked against structured +capabilities. + +- [ ] **Step 5: Run the complete EmbodiChain regression and docs build** + +```bash +pytest -q tests +black --check --diff --color ./ +LC_ALL=C.UTF-8 LANG=C.UTF-8 make -C docs html +``` + +Expected: zero pytest failures, Black leaves all files unchanged, and Sphinx +builds without new warnings/errors. + +- [ ] **Step 6: Inspect both diffs and dependency/API versions** + +```bash +if rg -n "dexsim\.default_world\(\)|get_physics_scene\(\)|SimulationManager\.get_instance\(" embodichain/lab/sim/objects/rigid_object.py embodichain/lab/sim/objects/articulation.py embodichain/lab/sim/objects/backends/newton.py; then + echo "core Newton object/view path still contains global owner lookup" >&2 + exit 1 +fi +if rg -n "register_mesh_object_to_newton_patch|_get_entity_native_handle|dexsim_meta" embodichain/lab/sim/utility/sim_utils.py embodichain/lab/sim/objects/backends/newton.py; then + echo "EmbodiChain still consumes private DexSim Newton integration state" >&2 + exit 1 +fi +git -C /root/sources/dexsim diff --check dev...HEAD +git -C /root/sources/dexsim log --oneline --decorate dev..HEAD +git -C /root/sources/EmbodiChain diff --check main...HEAD +git -C /root/sources/EmbodiChain log --oneline --decorate main..HEAD +python - <<'PY' +import dexsim +from dexsim.engine.newton_physics import NEWTON_INTEGRATION_API_VERSION +assert dexsim.__version__.split("+")[0] == "0.4.4" +assert NEWTON_INTEGRATION_API_VERSION == 2 +print(dexsim.__version__, NEWTON_INTEGRATION_API_VERSION) +PY +``` + +Expected: no whitespace errors, reviewable commit series in each repository, +package base version `0.4.4`, API version `2`. + +- [ ] **Step 7: Commit Stage 1 documentation** + +```bash +git add docs/source/overview/sim/sim_manager.md design/newton-backend-design.md +git commit -m "docs: describe Newton runtime lifecycle contracts" +``` + +- [ ] **Step 8: Stop at the Stage 1 review gate** + +Report exact command outputs, failures/skips, both branch heads, and remaining +known upstream limitations. Do not start Stage 2. After Stage 1 is accepted, +use `superpowers:brainstorming` only if Stage 1 changed the approved design; +otherwise use `superpowers:writing-plans` to create the dependent +differentiable dynamics/kinematics implementation plan. + +--- + +## Spec Coverage Checklist + +| Specification requirement | Implemented by | +|---|---| +| DexSim public API/version/generation/prepare | Tasks 1–3 | +| Public rigid attachment and canonical descriptors | Task 2 | +| Generation-aware rigid/articulation bindings | Tasks 3, 9, 10 | +| Transactional rebuild and rigid state preservation | Task 4 | +| Active model-generation lease blocks rebuild | Tasks 1, 4 | +| Articulation state/control preservation and q/qd spans | Task 5 | +| Full arena/world frame conversion | Tasks 7, 9, 10 | +| No global/default-world core ownership | Tasks 7–8 | +| Pending initialization without constructor virtual reset | Tasks 8–9 | +| Two same-device worlds/managers and deterministic cleanup | Tasks 6, 11 | +| Structured capabilities and explicit unsupported errors | Tasks 7, 10 | +| Exact prepare/step count | Tasks 1, 4, 11 | +| Existing public API and default backend compatibility | Tasks 7–12 | +| Documentation and complete merge gate | Task 12 | diff --git a/docs/superpowers/specs/2026-07-13-newton-runtime-contracts-design.md b/docs/superpowers/specs/2026-07-13-newton-runtime-contracts-design.md new file mode 100644 index 000000000..133189299 --- /dev/null +++ b/docs/superpowers/specs/2026-07-13-newton-runtime-contracts-design.md @@ -0,0 +1,749 @@ +# Newton Runtime Contracts and Differentiable Execution Design + +**Status:** Approved in design discussion; awaiting written-spec review + +**Date:** 2026-07-13 + +**EmbodiChain branch:** `feature/newton-physics-backend` + +**DexSim implementation branch:** `feature/embodichain-newton-contracts` + +**Target DexSim package version:** `0.4.4` + +## 1. Purpose and supersession + +This specification replaces the multi-environment and differentiable-runtime +portions of: + +- `docs/superpowers/specs/2026-06-21-newton-backend-pr-design.md` +- `docs/superpowers/plans/2026-06-22-newton-backend-pr.md` +- the corresponding Target 4 and Target 5 status claims in + `design/newton-backend-design.md` + +The previous documents describe mutually incompatible clone-at-finalize and +spawn-time clone designs, treat rebuild-time runtime IDs as permanent, and +conflate a forward-kinematics demonstration with differentiable Newton +dynamics. They remain historical records but are not implementation sources +after this specification is accepted. + +The work is delivered in two sequential stages: + +1. A coordinated DexSim and EmbodiChain refactor covering the Newton public + integration contract, lifecycle, multi-world isolation, runtime topology + mutation, rigid bodies, and articulations. +2. A differentiable execution layer built on the resulting authoritative + state, generation, binding, and lifecycle contracts. It supports both real + solver dynamics and pure kinematics. + +Stage 2 starts only after Stage 1 correctness and lifecycle tests pass. + +## 2. Goals + +### 2.1 Stage 1 goals + +- Make DexSim the sole authority for Newton model, state, control, contacts, + entity metadata, runtime mappings, and model generation. +- Remove EmbodiChain use of private DexSim registry, registration, and + `dexsim_meta` details. +- Preserve the existing public `SimulationManager`, `RigidObject`, and + `Articulation` call surfaces. +- Support initial build and post-finalize add/remove for rigid bodies and + articulations. +- Preserve surviving rigid and articulation state across rebuilds. +- Rebind all runtime IDs after every successful model rebuild. +- Correctly isolate global and child-arena Newton worlds. +- Support two or more simultaneous `SimulationManager`/DexSim `World` + instances in one process, including same-GPU operation and deterministic + teardown. +- Make arena-local and world-frame pose semantics explicit and consistent. +- Support articulation topologies whose position and velocity widths differ, + including spherical and free joints at the binding-contract level. +- Preserve the default physics backend behavior. + +### 2.2 Stage 2 goals + +- Make the default differentiable path execute + `DifferentiableStepper` and the configured Newton solver. +- Match normal simulation time and substep semantics exactly. +- Support pure-kinematics environment steps through `newton.eval_fk` without + misrepresenting them as dynamics. +- Support non-zero action gradients, finite-difference validation, and + continuous multi-step differentiation. +- Preserve the normal environment lifecycle while making only a minimal, + explicit differentiable-output addition to the functor surface. +- Make state-buffer ownership safe across forward, backward, reset, rebuild, + and multiple worlds. + +## 3. Non-goals + +- A general functor-system rewrite. +- Runtime topology mutation for soft bodies or cloth. Such attempts fail + explicitly in this iteration. +- Replacing EmbodiChain's environment framework with IsaacLab's architecture. +- Copying IsaacLab's class-level physics singleton, USD/Fabric coupling, or + backend discovery by class-name convention. +- Broad renderer, sensor, or solver performance optimization unrelated to the + new contracts. +- Heterogeneous articulations inside one `Articulation` batch. Separate + batched assets may have different topologies, while instances within one + batch retain the same topology. + +## 4. Architectural boundary + +The ownership rule is: + +> DexSim owns physical truth; EmbodiChain owns environment semantics. + +```text +SimulationManager + -> PhysicsBackend + -> BackendSceneContext + -> DexSim World / NewtonManager + -> arena transforms + -> model generation + -> entity and view registries +``` + +DexSim owns: + +- `ModelBuilder`, `Model`, both runtime `State` buffers, `Control`, contacts, + collision pipeline, solver, and CUDA graph; +- stable entity references and canonical replay descriptors; +- world assignment, runtime body/shape/articulation/link/joint mappings; +- build, rebuild, snapshot, restore, commit, and generation transitions; +- the public tensor and binding contracts used by integrations. + +EmbodiChain owns: + +- backend selection and environment orchestration; +- public object APIs and backend-independent views; +- arena-local frame semantics and arena transform tables; +- pending initialization of newly added objects; +- environment reset, step count, observations, rewards, hooks, and datasets; +- differentiable environment policy and task-specific action/output kernels. + +Core object and utility code must not resolve its owner through +`dexsim.default_world()` or the default `SimulationManager` instance. + +## 5. Lifecycle and generation + +### 5.1 Lifecycle phases + +The integration exposes the following conceptual phases: + +```text +BUILDING + -> MODEL_FINALIZED + -> VIEWS_BOUND + -> SOLVER_READY + -> RUNNING + -> STALE + -> rebuild + -> VIEWS_BOUND +``` + +DexSim may retain its internal state enum, but public results must distinguish +successful model finalization, successful binding readiness, solver readiness, +staleness, and closure. `READY` alone must not ambiguously mean both “model +exists” and “all external consumers are rebound.” + +### 5.2 Model generation + +Each `NewtonManager` has a public, read-only `model_generation`: + +- it starts at `0` before the first finalized model; +- the first successful finalize commits generation `1`; +- every successful model-replacing rebuild increments it once; +- live writes that do not replace or re-index model arrays do not increment it; +- failed candidate builds do not increment it; +- separate worlds maintain independent generations. + +All runtime bindings carry the generation against which they were resolved. +Using a stale binding raises an explicit generation error or causes the owning +view to rebind before access; it never silently uses old IDs. + +### 5.3 Prepare result and rebuild events + +DexSim exposes a public prepare result equivalent to: + +```python +@dataclass(frozen=True) +class NewtonPrepareResult: + generation: int + did_build: bool + did_rebuild: bool + added_entities: tuple[NewtonEntityRef, ...] + removed_entities: tuple[NewtonEntityRef, ...] +``` + +After an atomic runtime commit, DexSim publishes a `MODEL_REBUILT` event with +the old and new generations and the topology delta. Failed builds publish a +failure result but never a success event. EmbodiChain subscribes when its +backend activates and unsubscribes during close. + +`PhysicsBackend.prepare()` returns an EmbodiChain-level result carrying the +same generation and rebuild facts. It establishes a ready-to-step runtime but +does not itself advance simulation time. + +## 6. DexSim public integration contract + +### 6.1 API version + +DexSim exports: + +```python +NEWTON_INTEGRATION_API_VERSION = 2 +``` + +The package patch version advances to `0.4.4`. EmbodiChain pins the exact +package version and validates the integration API version at backend +activation. This prevents two materially different Newton integrations from +sharing an indistinguishable dependency version. + +### 6.2 Stable entity references + +Registration returns an opaque, stable `NewtonEntityRef` containing enough +identity to reject cross-world use. Runtime integer IDs are not exposed as +stable handles. + +The identity is derived from the owning world and DexSim entity, not from a +model body index. Removal invalidates the reference for new bindings while +allowing rebuild snapshots to identify that the entity should be omitted. + +### 6.3 Public rigid-body attachment + +DexSim provides a supported attachment API equivalent to: + +```python +attach_rigid_body( + entity, + *, + actor_type, + shape_type, + physical_attr=None, + body_desc=None, + shape_desc=None, + geometry_desc=None, +) -> NewtonEntityRef +``` + +This is the only integration entry point EmbodiChain uses for Newton rigid +bodies. It: + +- resolves the owning manager from the entity's actual arena/world; +- derives global world `-1` or the correct child-world index; +- captures mesh, box, sphere, and other supported geometry parameters; +- stores a canonical descriptor that can be replayed during rebuild; +- supports the legacy `PhysicalAttr` projection and Newton-native body/shape + descriptors; +- makes descriptor ownership per entity so clone mutation cannot alias the + prototype; +- marks a finalized runtime stale when topology changes. + +Clone operations recompute target world metadata from the target arena rather +than copying the prototype's world index. + +### 6.4 Public generation-aware bindings + +DexSim provides binding operations equivalent to: + +```python +bind_rigid_entities(refs) -> RigidEntityBinding +bind_articulations(refs) -> ArticulationBinding +``` + +`RigidEntityBinding` includes generation, body IDs where applicable, shape +IDs, and world IDs. Static entities may have shape IDs without body IDs. + +`ArticulationBinding` includes generation, articulation and link body IDs, +world IDs, and explicit per-active-joint spans for: + +- current q position; +- target q position; +- q velocity; +- target q velocity; +- generalized force/control. + +It separately reports `qpos_width` and `dof`/`qvel_width`. An active-joint +index is never interpreted as a flattened DOF index. + +Bindings use `int32` indices and declare device, dtype, shape, ownership, +mutability, and lifetime. Public state data uses `float32`; public quaternions +use `xyzw`. + +## 7. Transactional rebuild and state restoration + +### 7.1 Rebuild sequence + +Runtime topology mutation follows: + +```text +add/remove entity + -> manager STALE + -> snapshot by stable entity reference + -> build candidate builder/model/state/control/solver + -> restore surviving entities into candidate runtime + -> validate candidate mappings and resources + -> atomically commit candidate runtime + -> generation + 1 + -> publish MODEL_REBUILT + -> EmbodiChain rebinds views + -> initialize only newly added entities + -> FK and required DexSim visual synchronization +``` + +The old runtime is not cleared before the candidate is validated. Candidate +construction may temporarily use additional memory; correctness and rollback +take precedence over rebuild-time peak memory. + +If candidate construction or restoration fails: + +- the manager remains `STALE` and stepping is prohibited; +- the previous runtime remains available for diagnostics but is not presented + as current physical truth for the mutated scene; +- generation does not change; +- no rebuilt event is emitted; +- callers may correct the scene/configuration and retry prepare. + +### 7.2 Snapshot coverage + +Snapshots are keyed by stable entity reference rather than runtime body ID. + +Rigid state coverage: + +- pose; +- linear and angular velocity; +- linear and angular acceleration; +- pending external force and torque; +- both ping-pong state buffers where fields exist. + +Articulation coverage: + +- root pose and velocity; +- current and target q position; +- current and target q velocity; +- generalized forces and active controls; +- relevant drive/control state; +- both ping-pong state buffers. + +Contacts are regenerated and are not restored. Removed entities are omitted. +New entities receive descriptor/default state and are reported in the prepare +result for owner-side initialization. + +### 7.3 Differentiable model leases + +A live differentiable session holds a lease on its model generation. A +topology-changing rebuild is rejected while an outstanding tape depends on +that generation. The user or environment must finish backward or explicitly +close/detach the session before rebuilding. This prevents model arrays from +being freed while Warp autograd still references them. + +## 8. EmbodiChain backend and scene context + +### 8.1 BackendSceneContext + +Every simulated object receives its owner explicitly through a context that +contains: + +- `SimulationManager` identity; +- DexSim `World` and physics scene; +- active `PhysicsBackend`; +- arena list and full world transforms; +- current backend/model generation; +- entity/view registration helpers. + +Objects no longer call `dexsim.default_world()`, global +`get_physics_scene()`, or default-instance arena utilities in core paths. + +### 8.2 View rebinding + +Rigid and articulation views store a binding rather than permanent IDs. Before +each batch operation they perform an O(1) generation comparison. A mismatch +causes one binding refresh for the complete batch, invalidating dependent +sorted-ID and arena-transform caches. + +READY steady state does not re-resolve IDs, allocate bindings, or loop over +entities in Python. + +### 8.3 Pending initialization + +EmbodiChain records newly added objects in `pending_initialization`. After a +successful prepare and view rebind, only those objects receive their initial +state/reset. Existing objects retain the state restored by DexSim. + +Base entity constructors do not call overridable `reset()` methods. Object +initialization is an explicit manager/lifecycle phase. + +### 8.4 Frame contract + +The public API distinguishes world and arena-local frames. Existing +`set_local_pose()` and `get_local_pose()` remain arena-local before and after +finalization. + +Conversion uses the complete arena rigid transform, including rotation, not +only XY translation. Root and link pose data returned by DexSim global APIs is +converted in the view. Quaternion convention is consistently `xyzw`. + +Velocity and wrench APIs retain their documented frames; any API whose frame +is currently ambiguous is documented and validated as part of this refactor +rather than inferred from method names. + +### 8.5 Articulation data contract + +EmbodiChain stores separate q-position and velocity/force widths and delegates +active-joint span resolution to `ArticulationBinding`. Current all-1-DOF robot +calls remain source compatible. Spherical and free joints use their actual q +and qd widths. + +Writing current q position triggers required FK invalidation/evaluation before +link pose or visual state is reported. Unsupported data such as articulation +q acceleration is represented through capabilities and raises an explicit +unsupported-operation error rather than returning plausible zeros. + +## 9. Multi-world isolation and cleanup + +Registries, generation counters, entity mappings, solvers, state buffers, +CUDA graphs, and callbacks are keyed by owning world. A reference or binding +from one world cannot be used with another. + +Same-GPU CUDA capture may use a device-level coordinator for capture safety, +but that coordinator does not own simulation state and stores only weak +manager references. Capture timeout and peer diagnostics use existing public +or implemented helpers and cannot wait indefinitely by default. + +EmbodiChain adds an idempotent `SimulationManager.close()` that never exits the +process. It releases backend subscriptions, bindings, DexSim world resources, +CUDA graphs, and instance registry entries. + +Existing cleanup surfaces remain compatible: + +- `destroy()` remains available and preserves its documented exit-process + compatibility behavior; +- `SimulationManager.reset(instance_id)` closes the selected live instance + before removing it, so a new instance cannot inherit its world state; +- repeated close/reset is safe. + +## 10. Capabilities and validation + +Backend capabilities become structured and cover operations in addition to +asset categories. The Newton capability description includes at least: + +- supported asset kinds; +- supported solver and gradient combinations; +- CUDA graph support and invalidation rules; +- partial reset and FK support; +- runtime topology mutation by asset kind; +- heterogeneous q/qd span support; +- runtime collision-filter support; +- contact sensor and acceleration-field support; +- multi-world support. + +Configuration validates positive dt/substeps, device normalization, solver +parameters, gradient requirements, collision pipeline compatibility, and CUDA +graph combinations before finalization. Unconsumed solver parameters are +errors, not silently ignored fields. + +Unsupported operations fail at configuration or API boundaries. In +particular, this iteration rejects runtime topology mutation for soft bodies +and cloth, and reports upstream Newton limitations instead of returning fake +data. + +## 11. Differentiable execution architecture + +### 11.1 Functional core and stateful environment + +The differentiable layer has two levels: + +1. `DifferentiableSession`, a generation-bound functional rollout owner. +2. `DifferentiableEmbodiedEnv`, a stateful Gym/EmbodiChain wrapper. + +The session owns independent state, control, contact, and tape buffers. It +never records directly into buffers that a later environment step will +overwrite before backward. Each forward retains its required buffers in the +autograd context until backward or explicit release. + +The environment maintains a functional session state across steps and mirrors +the resulting state into the normal runtime for non-differentiable consumers, +rendering, and existing object APIs. Mirror writes do not replace tape-owned +buffers. + +Any generation change invalidates the session. Reset detaches the reset +environments from prior episode history. + +### 11.2 Explicit execution modes + +Configuration selects: + +```python +DifferentiableStepCfg( + mode="dynamics" | "kinematics", + bptt_horizon_steps=None, +) +``` + +The existing `truncate_backward_at` input remains accepted as a deprecated +alias for `bptt_horizon_steps`. Its former ambiguous solver-substep meaning is +not retained. Truncation occurs only at environment-step boundaries. + +### 11.3 Dynamics mode + +Dynamics mode must execute `DifferentiableStepper` and the configured solver. +For one environment step: + +```text +write action/control + -> repeat sim_steps_per_control physics steps + -> repeat Newton num_substeps solver steps + -> clear forces + -> apply pending external forces + -> collide + -> DifferentiableStepper.step + -> swap state + -> clear one-shot external inputs +``` + +The total solver step count is: + +```text +sim_steps_per_control * NewtonPhysicsCfg.num_substeps +``` + +The solver dt is `physics_dt / num_substeps`. Control remains applied with the +same cadence as normal simulation. State ownership and final-buffer selection +are independent of odd/even substep count. + +No FK-only fallback is permitted when a task is configured for dynamics. A +zero gradient caused by an unsupported control path fails validation/tests +and must be corrected at the control/solver contract. + +### 11.4 Kinematics mode + +Kinematics mode executes: + +```text +action + -> task-defined q-position update + -> newton.eval_fk + -> body/link state + -> differentiable observations and reward +``` + +It does not run collision or a solver and does not advance physical simulation +time. It does advance the environment episode step. Runtime q position and +DexSim visual state are synchronized after the functional result when enabled +by the environment. + +Kinematics is a first-class, explicitly named mode, not evidence that dynamics +differentiation works. + +### 11.5 Autograd output contract + +The PyTorch/Warp bridge uses explicit outputs equivalent to: + +```python +DifferentiableOutput( + name="reward", + tensor=reward_torch, + source=reward_warp_array, + requires_grad=True, +) +``` + +Every differentiable output has a Warp source whose gradient is seeded by the +custom backward. Observation and reward are handled independently, allowing +both `loss(obs)` and `loss(reward)` to propagate to action. Shape, dtype, +device, contiguity, and finite-value checks occur at the bridge boundary. + +Terminated, truncated, info, and other non-differentiable outputs explicitly +declare that they do not receive a gradient. + +## 12. Environment and minimal functor integration + +`DifferentiableEmbodiedEnv.step()` preserves the normal lifecycle: + +```text +action preprocessing + -> differentiable action mapping + -> dynamics or kinematics execution + -> differentiable observation/reward output + -> ordinary info and termination + -> episode counters + -> hooks and dataset handling + -> reset completed environments +``` + +This iteration does not redesign functors. Tasks provide thin differentiable +action and output adapters, normally implemented with Warp kernels. Existing +ordinary functors continue to run and are detached unless explicitly backed by +a `DifferentiableOutput` source. + +If configuration claims an observation or reward term is differentiable but +no valid Warp source is supplied, construction or the first validated step +raises an error. The system does not silently sever the graph. + +Non-differentiable side effects such as logging, dataset recording, and most +hooks remain outside the tape. + +## 13. Franka reference environments + +The Franka reach task exposes two explicit configurations: + +- **Dynamics:** action maps to a differentiable Newton effort/control path and + must pass through `DifferentiableStepper` and the semi-implicit solver. +- **Kinematics:** action updates joint q position and uses `newton.eval_fk`. + +Both modes use the same documented frame convention. Arena-local targets are +converted consistently against world-frame Newton body state, or body state is +converted to arena-local before reward evaluation. + +The reference task registers through the normal task import path, uses a +deterministic local/fixture asset for required tests, closes its environment in +all test outcomes, and does not depend on a network download for required CI. + +The dynamics acceptance path uses a control mode that is expected to produce a +real action-to-state gradient. The kinematics task remains useful as a faster +smoke test but is reported separately. + +## 14. Error handling + +Errors identify the owning world, entity reference, operation, and expected +versus actual generation where relevant. The design distinguishes: + +- invalid configuration; +- unsupported backend capability; +- stale/removed entity binding; +- closed world/session; +- candidate rebuild failure; +- active differentiable lease blocking rebuild; +- cross-world reference use; +- tensor contract mismatch. + +Runtime collision-filter changes and other setup-only fields either trigger a +documented rebuild or fail explicitly; an API must not report success after +updating metadata that the live model does not consume. + +## 15. Testing and acceptance + +### 15.1 DexSim tests + +- Public rigid attachment for mesh, box, and sphere descriptors. +- Correct global and child-world IDs for one, two, and eight child arenas. +- Clone descriptors are independent and use the target arena's world ID. +- Initial finalize increments generation once. +- Add/remove rebuild increments generation and produces new runtime mappings. +- Rigid pose/velocity/acceleration/external-wrench state survives rebuild. +- Articulation root/current-target q/qd/qf/control state survives rebuild. +- Both state buffers remain valid after restore. +- Candidate build/restore failure leaves a diagnosable non-half-initialized + manager and does not increment generation. +- Revolute, spherical, and free-joint q/qd spans bind correctly. +- Two same-GPU worlds can build, step, rebuild, capture where enabled, and + close independently. +- Repeated close and failed-construction cleanup are safe. + +### 15.2 EmbodiChain tests + +- A view's cached binding changes from the old to the new generation after + rebuild and addresses the correct entity in every environment. +- Adding an entity preserves old-object state and initializes only the new + entity. +- Removing an entity invalidates its binding without changing surviving + object identity or state. +- Rigid, articulation root, and link local poses are correct in every arena, + including non-zero rotations. +- No core object/view lookup resolves through `default_world()`. +- Two `SimulationManager` instances do not share world, scene, arena, + generation, bindings, or cleanup state. +- Existing public object and manager calls remain source compatible. +- Default-backend behavior and tests remain unchanged. +- Capability errors replace silent q-acceleration, collision-filter, or sensor + no-ops covered by the new surface. + +### 15.3 Differentiable tests + +- Dynamics uses the real `DifferentiableStepper` and expected solver-step + count. +- Action gradient is finite, non-zero, and has the expected shape. +- Central finite difference agrees in direction and reasonable tolerance with + autograd for a deterministic small scene. +- At least three consecutive environment steps advance state and backpropagate + safely. +- Odd and even Newton substep counts select the correct final state. +- `loss(obs)` and `loss(reward)` independently propagate to action. +- Kinematics FK pose and gradient match a direct `eval_fk` reference. +- Runtime mirror updates do not corrupt a still-live backward pass. +- BPTT truncation detaches at the requested environment-step boundary. +- Reset prevents cross-episode gradient leakage. +- A generation change invalidates an old session; a live model lease blocks + rebuild until released. +- Dynamics and kinematics handle arena-local targets consistently in multiple + environments. + +### 15.4 Verification order + +The merge gate runs in this order: + +```text +DexSim unit tests + -> EmbodiChain CPU/headless contract tests + -> serial GPU Newton integration tests + -> multi-world lifecycle tests + -> differentiable finite-difference tests + -> complete EmbodiChain regression suite + -> formatting and project pre-commit checks +``` + +GPU and external-simulation tests use the repository's registered markers and +deterministic teardown. Required tests do not silently skip because an asset +download failed. + +## 16. Performance constraints + +The correctness refactor must preserve an efficient steady state: + +- generation checks are O(1); +- READY views do not rebind without a generation change; +- binding refresh is batched; +- per-step pose/state access does not loop over entities in Python; +- arena transform tables are device-resident and rebuilt only when their + owning context changes; +- differentiable rollout buffers are deliberately owned and reused only when + doing so cannot invalidate an outstanding tape. + +Broad solver/render benchmarking is outside this specification. Focused +benchmarks may be added to demonstrate that generation-aware binding does not +regress steady-state batch access. + +## 17. Implementation and repository sequencing + +After this written specification is approved: + +1. Create DexSim branch `feature/embodichain-newton-contracts` from its current + `dev` branch. +2. Write a Stage 1 implementation plan spanning DexSim and EmbodiChain, with + tests preceding implementation changes. +3. Implement and verify the DexSim public contract and lifecycle first. +4. Update EmbodiChain to consume that contract and complete multi-env, + rebuild, articulation, and multi-manager parity. +5. Run the Stage 1 merge gate. +6. Write the dependent Stage 2 implementation plan. +7. Implement dynamics and kinematics sessions, bridge, environment, and + reference tasks. +8. Run the complete differentiable and repository merge gates. + +Stage 1 and Stage 2 remain reviewable as separate commit series even though +the first two previously proposed PR scopes are now one coordinated refactor. + +## 18. Accepted trade-offs + +- Transactional rebuild temporarily consumes more memory than clearing the + old runtime first. +- Explicit bindings and generation checks add types and lifecycle plumbing but + remove unsafe permanent IDs. +- Runtime articulation mutation requires upstream snapshot work rather than an + EmbodiChain-only workaround. +- Differentiable tape-owned buffers use more memory than mutating manager state + in place; this is required for correct backward ownership. +- Functor integration remains deliberately narrow in this iteration. +- Soft-body and cloth runtime mutation is explicitly postponed rather than + approximated with incomplete state preservation. diff --git a/embodichain/data/assets/demo_assets.py b/embodichain/data/assets/demo_assets.py index 9b00a654a..255002ab6 100644 --- a/embodichain/data/assets/demo_assets.py +++ b/embodichain/data/assets/demo_assets.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import open3d as o3d import os @@ -49,3 +51,20 @@ def __init__(self, data_root: str = None): prefix = type(self).__name__ path = EMBODICHAIN_DEFAULT_DATA_ROOT if data_root is None else data_root super().__init__(prefix, data_descriptor, path) + + +class CoordinatedPlacementAndPickment(EmbodiChainDataset): + """Dataset class for coordinated placement and pickment tutorial meshes.""" + + def __init__(self, data_root: str = None): + data_descriptor = o3d.data.DataDescriptor( + os.path.join( + EMBODICHAIN_DOWNLOAD_PREFIX, + demo_assets, + "coordinated_placement_and_pickment.zip", + ), + "297c10b386a4d7a8ccb68926d69425e9", + ) + prefix = type(self).__name__ + path = EMBODICHAIN_DEFAULT_DATA_ROOT if data_root is None else data_root + super().__init__(prefix, data_descriptor, path) diff --git a/embodichain/data/assets/robot_assets.py b/embodichain/data/assets/robot_assets.py index f37cfd3a1..b5f6dfb05 100644 --- a/embodichain/data/assets/robot_assets.py +++ b/embodichain/data/assets/robot_assets.py @@ -357,7 +357,6 @@ class Franka(EmbodiChainDataset): Franka/ Panda/ Panda.urdf - PandaHand.urdf PandaWithHand.urdf FR3/ fr3.urdf @@ -374,8 +373,8 @@ class Franka(EmbodiChainDataset): def __init__(self, data_root: str = None): data_descriptor = o3d.data.DataDescriptor( - os.path.join(EMBODICHAIN_DOWNLOAD_PREFIX, robot_assets, "FrankaV2.zip"), - "f0675b9da98126bc3d4e18c98ef5e06c", + os.path.join(EMBODICHAIN_DOWNLOAD_PREFIX, robot_assets, "FrankaV3.zip"), + "ec53536615ddc1c44660f2d66fdf82ae", ) prefix = type(self).__name__ path = EMBODICHAIN_DEFAULT_DATA_ROOT if data_root is None else data_root diff --git a/embodichain/lab/gym/envs/differentiable_env.py b/embodichain/lab/gym/envs/differentiable_env.py index 829b15739..95dfbdad0 100644 --- a/embodichain/lab/gym/envs/differentiable_env.py +++ b/embodichain/lab/gym/envs/differentiable_env.py @@ -25,13 +25,13 @@ Usage: class MyTask(DifferentiableEmbodiedEnv): - def _apply_action_kernel(self, action_wp, tape): ... + def _apply_dynamics_action_kernel(self, action_wp, control, tape): ... def _read_outputs(self, final_state) -> dict: ... """ from __future__ import annotations -from typing import Any, Callable +from typing import Any, Callable, Literal import torch @@ -46,11 +46,19 @@ def _read_outputs(self, final_state) -> dict: ... class DifferentiableEmbodiedEnv(EmbodiedEnv): """EmbodiedEnv variant that exposes APG-ready :py:meth:`step`. - Subclasses must implement :meth:`_apply_action_kernel` and - :meth:`_read_outputs`; the rest of the EmbodiedEnv contract (reset, - observation managers, reward functors) carries over. + Dynamics subclasses must implement :meth:`_apply_dynamics_action_kernel` + and :meth:`_read_outputs`; the rest of the EmbodiedEnv contract (reset, + observation managers, reward functors) carries over. The default + ``dynamics`` route invokes the Newton solver through + :class:`NewtonStepFunc` using a detached trajectory-local control buffer. + Subclasses that intentionally use FK-only stepping must explicitly select + ``kinematics`` and implement :meth:`_make_kinematic_step_fn` together with + the legacy :meth:`_apply_action_kernel` hook. """ + differentiable_step_mode: Literal["dynamics", "kinematics"] = "dynamics" + """Stepping route used by :meth:`_build_sim_state_dict`.""" + def __init__(self, cfg: EmbodiedEnvCfg, *args, **kwargs) -> None: self._validate_diff_cfg(cfg) super().__init__(cfg, *args, **kwargs) @@ -74,17 +82,38 @@ def _validate_diff_cfg(cfg: EmbodiedEnvCfg) -> None: # -- subclass contract ------------------------------------------------ # - def _apply_action_kernel(self, action_wp: Any, tape: Any) -> None: - """Inside the open Warp tape, write the action into Newton control. + def _apply_dynamics_action_kernel( + self, + action_wp: Any, + control: Any, + tape: Any, + ) -> None: + """Write an action into a detached dynamics trajectory control buffer. Implementations launch a Warp kernel that reads ``action_wp`` (a ``wp.array(dtype=wp.float32, requires_grad=True)`` of shape - ``[num_envs * action_dim]``) and writes into - ``self.sim.physics.newton_manager._control`` so the next stepper - call uses the new control. + ``[num_envs * action_dim]``) and writes into the supplied ``control``. + It is the isolated control owned by the active manager trajectory; do + not write ``self.sim.physics.newton_manager._control`` while the tape + is active. ``tape`` is the caller-owned active Warp tape for this + callback only; the bridge clears the per-step binding after tape exit. """ raise NotImplementedError( - "Subclasses of DifferentiableEmbodiedEnv must implement " + "Dynamics subclasses of DifferentiableEmbodiedEnv must migrate " + "their legacy _apply_action_kernel(action_wp, tape) hook to " + "_apply_dynamics_action_kernel(action_wp, control, tape)." + ) + + def _apply_action_kernel(self, action_wp: Any, tape: Any) -> None: + """Write an action for the explicitly selected kinematics route. + + This legacy hook is deliberately reserved for + ``differentiable_step_mode = 'kinematics'``. It receives no detached + solver control because FK-only environments do not invoke Newton + solver dynamics. + """ + raise NotImplementedError( + "Kinematics subclasses of DifferentiableEmbodiedEnv must implement " "_apply_action_kernel(action_wp, tape)." ) @@ -102,44 +131,36 @@ def _read_outputs(self, final_state: Any) -> dict: "_read_outputs(final_state)." ) - def _make_step_fn(self) -> Callable[[], Any]: - """Return a callable that advances the sim inside the open tape. - - The returned callable takes no arguments and returns the final - Newton :class:`State` after stepping. It is invoked by - :class:`NewtonStepFunc` inside the ``with tape:`` block, so any - Warp kernel launches (or differentiable Newton calls like - ``eval_fk``) are recorded on the tape. - - The default implementation runs the differentiable - :class:`DifferentiableStepper` for ``sim_steps_per_control`` - substeps. Subclasses can override this to swap in an FK-only - differentiable path (bypassing the dynamics solver when it does - not propagate grad through control inputs) or any other - tape-tracked stepping strategy. + def _make_kinematic_step_fn(self) -> Callable[[], Any]: + """Return the explicitly selected FK-only stepping callback. + + Subclasses must override this hook only when they set + :attr:`differentiable_step_mode` to ``"kinematics"``. This keeps + kinematics distinct from the default solver-dynamics route. + + Raises: + NotImplementedError: If kinematics mode has no named FK hook. """ - manager = self.sim - substeps = self.cfg.sim_steps_per_control - nm = manager.physics.newton_manager - stepper = manager.create_differentiable_stepper() - state_in = nm._state_0 - state_out = nm._model.state() - contacts = stepper.create_contacts() - dt_val = nm.solver_dt - - def _step(): - for _ in range(substeps): - stepper.step(state_in, state_out, contacts=contacts, dt=dt_val) - state_in, state_out = state_out, state_in - return state_in - - return _step + raise NotImplementedError( + "DifferentiableEmbodiedEnv in kinematics mode requires " + "_make_kinematic_step_fn()." + ) # -- gym surface ------------------------------------------------------ # def step(self, action: torch.Tensor): + """Advance one differentiable control step. + + Terminal environments are auto-reset only when the call cannot retain + a Warp tape for backward. A grad-tracked step returns terminal + observations unchanged and records ``deferred_reset_ids`` in ``info``; + callers must run backward before resetting those environments. + """ if not isinstance(action, torch.Tensor): action = torch.as_tensor(action, dtype=torch.float32) + retains_tape_for_backward = bool( + torch.is_grad_enabled() and action.requires_grad + ) sim_state = self._build_sim_state_dict(action) outputs = NewtonStepFunc.apply(action, sim_state) obs, reward, terminated, truncated = outputs[:4] @@ -148,29 +169,82 @@ def step(self, action: torch.Tensor): done_mask = terminated | truncated if done_mask.any(): reset_ids = done_mask.nonzero(as_tuple=False).squeeze(-1) - fresh_obs, _ = self.reset(options={"reset_ids": reset_ids}) - obs = torch.where( - done_mask.unsqueeze(-1).expand_as(obs), - fresh_obs.detach(), - obs, - ) + if retains_tape_for_backward: + info["requires_reset_after_backward"] = True + info["deferred_reset_ids"] = reset_ids.detach().clone() + else: + fresh_obs, _ = self.reset(options={"reset_ids": reset_ids}) + obs = torch.where( + done_mask.unsqueeze(-1).expand_as(obs), + fresh_obs.detach(), + obs, + ) return obs, reward, terminated, truncated, info def _build_sim_state_dict(self, action: torch.Tensor) -> dict: - return { + mode = self.differentiable_step_mode + if mode not in {"dynamics", "kinematics"}: + raise ValueError( + "differentiable_step_mode must be 'dynamics' or 'kinematics', " + f"got {mode!r}." + ) + + action_kernel, tape_binder = self._action_kernel_for_mode(mode) + sim_state = { "manager": self.sim, + "step_mode": mode, "substeps": self.cfg.sim_steps_per_control, - "action_to_control_kernel": self._wrap_action_kernel(), + "action_to_control_kernel": action_kernel, "kernel_args": (), "obs_reward_fn": self._read_outputs, - "step_fn": self._make_step_fn(), "last_info": {}, } + if tape_binder is not None: + sim_state["_bind_dynamics_tape"] = tape_binder + if mode == "kinematics": + sim_state["step_fn"] = self._make_kinematic_step_fn() + return sim_state + + def _action_kernel_for_mode( + self, + mode: str, + ) -> tuple[Callable[..., None], Callable[[Any | None], None] | None]: + """Build the mode-specific action callback consumed by NewtonStepFunc.""" + if mode == "dynamics": + dynamics_hook = getattr(self, "_apply_dynamics_action_kernel", None) + if ( + not callable(dynamics_hook) + or getattr(dynamics_hook, "__func__", None) + is DifferentiableEmbodiedEnv._apply_dynamics_action_kernel + ): + raise NotImplementedError( + "Dynamics environments using the legacy " + "_apply_action_kernel(action_wp, tape) must migrate to " + "_apply_dynamics_action_kernel(action_wp, control, tape)." + ) + return self._wrap_dynamics_action_kernel(dynamics_hook) + return self._wrap_kinematic_action_kernel(), None + + @staticmethod + def _wrap_dynamics_action_kernel( + dynamics_hook: Callable[..., None], + ) -> tuple[Callable[..., None], Callable[[Any | None], None]]: + """Expose a local-control hook with tape ownership scoped per step.""" + active_tape: list[Any | None] = [None] + + def _bind_tape(tape: Any | None) -> None: + active_tape[0] = tape + + def _inner(action_wp: Any, control: Any, *_: Any) -> None: + dynamics_hook(action_wp, control, tape=active_tape[0]) + + return _inner, _bind_tape - def _wrap_action_kernel(self): + def _wrap_kinematic_action_kernel(self): + """Expose the strict legacy action hook only for kinematics mode.""" env = self - def _inner(action_wp, *_): - env._apply_action_kernel(action_wp, tape=None) + def _inner(action_wp: Any, tape: Any, *_: Any) -> None: + env._apply_action_kernel(action_wp, tape=tape) return _inner diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index 18137d87a..be9a0076c 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -566,6 +566,9 @@ def _initialize_episode( if self.cfg.rewards: self.reward_manager.reset(env_ids=env_ids) + if self.cfg.dataset: + self.dataset_manager.reset(env_ids=env_ids) + def _infer_rollout_buffer_mode(self, rollout_buffer: TensorDict) -> str: """Infer whether the rollout buffer is expert recording or RL training data.""" if { diff --git a/embodichain/lab/gym/envs/managers/randomization/spatial.py b/embodichain/lab/gym/envs/managers/randomization/spatial.py index 1af1c09f1..d97c52d08 100644 --- a/embodichain/lab/gym/envs/managers/randomization/spatial.py +++ b/embodichain/lab/gym/envs/managers/randomization/spatial.py @@ -17,9 +17,14 @@ from __future__ import annotations import torch -from typing import TYPE_CHECKING, Union, List - -from embodichain.lab.sim.objects import RigidObject, Robot, Articulation +from typing import TYPE_CHECKING, List + +from embodichain.lab.sim.objects import ( + RigidObject, + Robot, + Articulation, + RigidObjectGroup, +) from embodichain.lab.gym.envs.managers.cfg import SceneEntityCfg from embodichain.lab.gym.envs.managers import Functor, FunctorCfg from embodichain.utils.math import sample_uniform, matrix_from_euler, matrix_from_quat @@ -28,6 +33,17 @@ if TYPE_CHECKING: from embodichain.lab.gym.envs import EmbodiedEnv +__all__ = [ + "get_random_pose", + "randomize_rigid_object_pose", + "randomize_robot_eef_pose", + "randomize_robot_qpos", + "randomize_articulation_root_pose", + "randomize_target_pose", + "planner_grid_cell_sampler", + "randomize_anchor_height", +] + def get_random_pose( init_pos: torch.Tensor, @@ -583,3 +599,227 @@ def __call__( if physics_update_step > 0: env.sim.update(step=physics_update_step) + + +class randomize_anchor_height(Functor): + """Randomize the height of an anchor object and shift other objects by the same delta. + + This functor samples a per-environment height delta, moves the anchor object + relative to its configured initial position, and adds the same delta to the + Z component of every other included object while preserving XY and rotation. + + The functor is configured through :class:`FunctorCfg` parameters, following the + same pattern as :class:`planner_grid_cell_sampler`. + """ + + _VALID_GROUPS = {"background", "rigid_object", "rigid_object_group", "articulation"} + + def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv): + """Initialize the functor. + + Args: + cfg: The functor configuration. + env: The environment instance. + """ + super().__init__(cfg, env) + + def _get_object( + self, env: EmbodiedEnv, uid: str + ) -> RigidObject | Articulation | RigidObjectGroup | None: + """Get a rigid object, articulation, or rigid object group by UID.""" + if uid in env.sim.get_rigid_object_uid_list(): + return env.sim.get_rigid_object(uid) + if uid in env.sim.get_articulation_uid_list(): + return env.sim.get_articulation(uid) + if ( + hasattr(env.sim, "get_rigid_object_group_uid_list") + and uid in env.sim.get_rigid_object_group_uid_list() + ): + return env.sim.get_rigid_object_group(uid) + return None + + def _resolve_affected_uids( + self, + env: EmbodiedEnv, + anchor_uid: str, + include_groups: list[str] | None, + exclude_uids: list[str], + ) -> list[str]: + """Resolve the list of UIDs that should be shifted.""" + if include_groups is None: + include_groups = [ + "background", + "rigid_object", + "rigid_object_group", + "articulation", + ] + + invalid_groups = set(include_groups) - self._VALID_GROUPS + if invalid_groups: + raise ValueError( + f"Invalid include_groups: {sorted(invalid_groups)}. " + f"Valid options are: {sorted(self._VALID_GROUPS)}." + ) + + uids: set[str] = set() + if any(g in include_groups for g in ("background", "rigid_object")): + uids.update(env.sim.get_rigid_object_uid_list()) + if "rigid_object_group" in include_groups: + if hasattr(env.sim, "get_rigid_object_group_uid_list"): + uids.update(env.sim.get_rigid_object_group_uid_list()) + if "articulation" in include_groups: + uids.update(env.sim.get_articulation_uid_list()) + + exclude = set(exclude_uids) | {anchor_uid} + return sorted(uids - exclude) + + def _sample_delta( + self, + num_envs: int, + height_delta_range: tuple[list[float], list[float]] | None, + height_delta_candidates: list[float] | None, + ) -> torch.Tensor: + """Sample a height delta for each environment.""" + device = self._env.device + + if height_delta_range is not None: + low = torch.tensor(height_delta_range[0], device=device) + high = torch.tensor(height_delta_range[1], device=device) + return sample_uniform( + lower=low, upper=high, size=(num_envs, 1), device=device + ).squeeze(-1) + + # Discrete sampling + candidates = torch.tensor(height_delta_candidates, device=device) + indices = torch.randint(0, len(candidates), (num_envs,), device=device) + return candidates[indices] + + def _move_object_z( + self, + obj: RigidObject | Articulation | RigidObjectGroup, + delta_z: torch.Tensor, + env_ids: torch.Tensor, + absolute: bool = False, + ) -> None: + """Move an object in Z by delta_z. + + Args: + obj: The object to move (RigidObject, Articulation, or RigidObjectGroup). + delta_z: Per-environment Z offset. + env_ids: Target environment IDs. + absolute: If True, set Z to obj.cfg.init_pos[2] + delta_z. + If False, add delta_z to the current Z. + """ + if isinstance(obj, RigidObjectGroup): + if absolute: + logger.log_warning( + "absolute=True is not supported for RigidObjectGroup; using relative shift." + ) + pose = obj.get_local_pose(to_matrix=True) # (N, M, 4, 4) + pose[env_ids, :, 2, 3] += delta_z.unsqueeze(-1) + obj.set_local_pose(pose[env_ids], env_ids=env_ids) + obj.clear_dynamics(env_ids=env_ids) + return + + # Both RigidObject and Articulation return (N, 7) by default: + # (x, y, z, qw, qx, qy, qz) + pose = obj.get_local_pose() # (N, 7) + current_z = pose[env_ids, 2] + if absolute: + init_z = torch.tensor( + obj.cfg.init_pos[2], dtype=torch.float32, device=obj.device + ) + new_z = init_z + delta_z + else: + new_z = current_z + delta_z + pose[env_ids, 2] = new_z + + obj.set_local_pose(pose[env_ids], env_ids=env_ids) + obj.clear_dynamics(env_ids=env_ids) + + def __call__( + self, + env: EmbodiedEnv, + env_ids: torch.Tensor | None, + anchor_uid: str, + height_delta_range: tuple[list[float], list[float]] | None = None, + height_delta_candidates: list[float] | None = None, + include_groups: list[str] | None = None, + exclude_uids: list[str] | None = None, + physics_update_step: int = 0, + store_key: str = "anchor_height_delta", + ) -> None: + """Apply the height randomization. + + Args: + env: The environment instance. + env_ids: Target environment IDs. If None, all environments are targeted. + anchor_uid: Exact UID of the anchor object whose height is randomized. + height_delta_range: Uniform sampling range for the height delta. + height_delta_candidates: Discrete set of allowed height delta values. + include_groups: Object groups to shift. ``None`` means all groups. + exclude_uids: Additional UIDs to skip beyond the anchor object. + physics_update_step: Number of physics update steps to apply after moving objects. + store_key: Attribute name on ``env`` where the sampled delta is stored. + """ + if env_ids is None: + env_ids = torch.arange(env.num_envs, device=env.device) + elif not isinstance(env_ids, torch.Tensor): + env_ids = torch.tensor(env_ids, device=env.device) + + if len(env_ids) == 0: + return + + # Validate sampling configuration + if height_delta_range is None and height_delta_candidates is None: + raise ValueError( + "Either 'height_delta_range' or 'height_delta_candidates' must be provided." + ) + if height_delta_candidates is not None and len(height_delta_candidates) == 0: + raise ValueError("'height_delta_candidates' must not be empty.") + if height_delta_range is not None and height_delta_candidates is not None: + logger.log_warning( + "Both 'height_delta_range' and 'height_delta_candidates' provided; " + "using 'height_delta_range'." + ) + + if exclude_uids is None: + exclude_uids = [] + + # Confirm anchor exists + anchor = self._get_object(env, anchor_uid) + if anchor is None: + raise ValueError( + f"Anchor object with uid '{anchor_uid}' not found in the scene." + ) + + # Build affected UID list + affected_uids = self._resolve_affected_uids( + env, anchor_uid, include_groups, exclude_uids + ) + + num_envs = len(env_ids) + delta_z = self._sample_delta( + num_envs, height_delta_range, height_delta_candidates + ) + + # Move anchor relative to its initial pose + self._move_object_z(anchor, delta_z, env_ids, absolute=True) + + # Move affected objects relative to their current pose + for uid in affected_uids: + obj = self._get_object(env, uid) + if obj is None: + logger.log_warning( + f"Affected object '{uid}' no longer exists; skipping height shift." + ) + continue + self._move_object_z(obj, delta_z, env_ids, absolute=False) + + # Physics settle + if physics_update_step > 0: + env.sim.update(step=physics_update_step) + + # Store delta for downstream use + if store_key: + setattr(env, store_key, delta_z) diff --git a/embodichain/lab/gym/envs/managers/randomization/visual.py b/embodichain/lab/gym/envs/managers/randomization/visual.py index c49707bd5..bb4a08ac9 100644 --- a/embodichain/lab/gym/envs/managers/randomization/visual.py +++ b/embodichain/lab/gym/envs/managers/randomization/visual.py @@ -272,6 +272,7 @@ def randomize_light( position_range: tuple[list[float], list[float]] | None = None, color_range: tuple[list[float], list[float]] | None = None, intensity_range: tuple[float, float] | None = None, + direction_range: tuple[list[float], list[float]] | None = None, ) -> None: """Randomize light properties by adding, scaling, or setting random values. @@ -289,10 +290,17 @@ def randomize_light( position_range is the x, y, z value added into light's cfg.init_pos. color_range is the absolute r, g, b value set to the light object. intensity_range is the value added into light's cfg.intensity. + direction_range is the x, y, z value added into light's cfg.direction. + (Only applicable for ``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, and ``"mesh"`` light types.) .. tip:: This function uses CPU tensors to assign light properties. + .. warning:: + ``position_range`` is ignored for global scene lights (``"sun"``, ``"direction"``) + because they are infinite-distance lights with no meaningful position. + Use ``direction_range`` instead for these light types. + Args: env (EmbodiedEnv): The environment instance. env_ids (Union[torch.Tensor, None]): The environment IDs to apply the randomization. @@ -300,25 +308,45 @@ def randomize_light( position_range (tuple[list[float], list[float]] | None): The range for the position randomization. color_range (tuple[list[float], list[float]] | None): The range for the color randomization. intensity_range (tuple[float, float] | None): The range for the intensity randomization. + direction_range (tuple[list[float], list[float]] | None): The range for the direction randomization. + Only applicable for directional light types (``"sun"``, ``"direction"``, ``"spot"``, + ``"rect"``, ``"mesh"``). """ light: Light = env.sim.get_light(entity_cfg.uid) - num_instance = len(env_ids) + if light is None: + return + + is_global = light.is_global if hasattr(light, "is_global") else False + + # For global lights, normalize env_ids to avoid index-out-of-range + if is_global or light.num_instances == 1: + num_instance = 1 + effective_env_ids = None + else: + num_instance = len(env_ids) + effective_env_ids = env_ids if position_range: - init_pos = light.cfg.init_pos - new_pos = ( - torch.tensor(init_pos, dtype=torch.float32) - .unsqueeze_(0) - .repeat(num_instance, 1) - ) - random_value = sample_uniform( - lower=torch.tensor(position_range[0]), - upper=torch.tensor(position_range[1]), - size=new_pos.shape, - ) - new_pos += random_value - light.set_local_pose(new_pos, env_ids=env_ids) + if is_global: + logger.log_warning( + f"position_range ignored for global light '{entity_cfg.uid}' " + f"(type='{light.cfg.light_type}'). Use direction_range instead." + ) + else: + init_pos = light.cfg.init_pos + new_pos = ( + torch.tensor(init_pos, dtype=torch.float32) + .unsqueeze_(0) + .repeat(num_instance, 1) + ) + random_value = sample_uniform( + lower=torch.tensor(position_range[0]), + upper=torch.tensor(position_range[1]), + size=new_pos.shape, + ) + new_pos += random_value + light.set_local_pose(new_pos, env_ids=effective_env_ids) if color_range: color = torch.zeros((num_instance, 3), dtype=torch.float32) @@ -328,7 +356,7 @@ def randomize_light( size=color.shape, ) color += random_value - light.set_color(color, env_ids=env_ids) + light.set_color(color, env_ids=effective_env_ids) if intensity_range: init_intensity = light.cfg.intensity @@ -344,7 +372,31 @@ def randomize_light( ) new_intensity += random_value new_intensity.squeeze_(1) - light.set_intensity(new_intensity, env_ids=env_ids) + light.set_intensity(new_intensity, env_ids=effective_env_ids) + + if direction_range: + light_type = light.cfg.light_type + if light_type not in ("sun", "direction", "spot", "rect", "mesh"): + logger.log_warning( + f"direction_range ignored for light '{entity_cfg.uid}' " + f"(type='{light_type}'). Direction only applicable to " + f"'sun', 'direction', 'spot', 'rect', and 'mesh' types." + ) + return + + init_dir = light.cfg.direction + new_dir = ( + torch.tensor(init_dir, dtype=torch.float32) + .unsqueeze_(0) + .repeat(num_instance, 1) + ) + random_value = sample_uniform( + lower=torch.tensor(direction_range[0]), + upper=torch.tensor(direction_range[1]), + size=new_dir.shape, + ) + new_dir += random_value + light.set_direction(new_dir, env_ids=effective_env_ids) def randomize_emission_light( diff --git a/embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py b/embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py index 71220213d..b80bba940 100644 --- a/embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py +++ b/embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py @@ -19,8 +19,8 @@ produces ``action.grad`` that flows back through a differentiable forward-kinematics path (``newton.eval_fk``). The semi_implicit solver does not propagate grad through ``joint_target_pos`` to -``body_q`` (the grad path is zero), so we bypass the dynamics -solver and run FK directly, matching the reference APG +``body_q`` (the grad path is zero), so this task explicitly selects +the kinematics route and runs FK directly, matching the reference APG implementation in ``/root/sources/analytic_policy_gradients/envs/franka_reach_env.py``. """ @@ -115,13 +115,16 @@ class FrankaReachApgEnv(DifferentiableEmbodiedEnv): action -> new_joint_q (action kernel) -> eval_fk -> body_q -> reward kernel -> reward_wp -> tape.backward -> action.grad - The dynamics solver (semi_implicit) is bypassed because it does not - propagate gradient through ``joint_target_pos`` to ``body_q`` (the - stiffness-driven grad path evaluates to zero in practice). This - matches the reference APG env's workaround. + This task explicitly uses the ``kinematics`` route because the + semi_implicit dynamics solver does not propagate gradient through + ``joint_target_pos`` to ``body_q`` (the stiffness-driven grad path + evaluates to zero in practice). This matches the reference APG env's + FK-only workaround without changing the default route for other + differentiable environments. """ metadata = {"render_modes": ["human"], "default_num_envs": 4} + differentiable_step_mode = "kinematics" def __init__( self, @@ -208,7 +211,9 @@ def _cache_franka_buffers(self) -> None: self._limit_lo_wp = wp.array(lo, dtype=wp.float32, device=self._wp_device) self._limit_hi_wp = wp.array(hi, dtype=wp.float32, device=self._wp_device) self._n_joints_per_env = int(len(model.joint_q) // self.sim.num_envs) - # Fresh FK state; reused across step calls (eval_fk overwrites it). + # Every taped forward replaces these with private primal buffers before + # the bridge opens its tape. They must never alias manager live state. + self._current_joint_q_snapshot: wp.array | None = None self._fk_state = model.state() self._new_joint_q: wp.array | None = None # Per-env global EE body indices into the flat body_q array. @@ -275,13 +280,20 @@ def _sample_new_targets(self, env_ids: torch.Tensor) -> None: # -- DifferentiableEmbodiedEnv contract ------------------------------ # - def _make_step_fn(self) -> Callable[[], Any]: - """FK bypass: compute body_q from new_joint_q via newton.eval_fk. + def _build_sim_state_dict(self, action: torch.Tensor) -> dict: + """Detach FK primal buffers before the parent opens a Warp tape.""" + nm = self.sim.physics.newton_manager + self._current_joint_q_snapshot = wp.clone(nm._state_0.joint_q) + self._fk_state = nm._model.state() + return super()._build_sim_state_dict(action) + + def _make_kinematic_step_fn(self) -> Callable[[], Any]: + """Explicit FK hook: compute body_q from new_joint_q via ``eval_fk``. The semi_implicit solver does not propagate grad through ``joint_target_pos`` to ``body_q`` (the grad path is zero), so - we bypass the dynamics solver and run forward kinematics - directly inside the tape. ``self._new_joint_q`` is populated by + this kinematics-mode task runs forward kinematics directly + inside the tape. ``self._new_joint_q`` is populated by :meth:`_apply_action_kernel` before this callable runs. """ env = self @@ -303,11 +315,15 @@ def _apply_action_kernel(self, action_wp: Any, tape: Any) -> None: Writes ``new_joint_q = clamp(current_q + action * scale, lo, hi)`` into a freshly allocated ``self._new_joint_q`` Warp array. The - FK step function then consumes this array via ``newton.eval_fk``. + explicit kinematic hook then consumes this array via ``newton.eval_fk``. """ - nm = self.sim.physics.newton_manager n = self.sim.num_envs total = n * FRANKA_NUM_ARM_JOINTS + if self._current_joint_q_snapshot is None: + raise RuntimeError( + "Franka kinematics requires a detached joint_q snapshot " + "before opening its Warp tape." + ) # Allocate a fresh new_joint_q each call so each forward pass # has its own grad graph (the tape records the kernel writes). self._new_joint_q = wp.zeros( @@ -321,7 +337,7 @@ def _apply_action_kernel(self, action_wp: Any, tape: Any) -> None: dim=total, inputs=[ action_wp, - nm._state_0.joint_q, + self._current_joint_q_snapshot, self._new_joint_q, self._limit_lo_wp, self._limit_hi_wp, @@ -399,8 +415,9 @@ def step(self, action: torch.Tensor): The parent :meth:`DifferentiableEmbodiedEnv.step` runs the differentiable bridge. After it returns, we update - ``nm._state_0.joint_q`` (detached) for envs that were not - auto-reset so the next step starts from the new configuration. + ``nm._state_0.joint_q`` for non-terminal envs so the next step starts + from the new configuration. The tape reads a per-forward detached + snapshot, so this live continuation cannot overwrite its primal input. """ if not isinstance(action, torch.Tensor): action = torch.as_tensor(action, dtype=torch.float32) @@ -437,8 +454,9 @@ def reset( Args: seed: Optional RNG seed for deterministic resets. options: Optional dict; supports ``{"reset_ids": }`` - for partial resets (used by auto-reset in - :meth:`DifferentiableEmbodiedEnv.step`). + for partial resets. No-grad terminal steps use it for + auto-reset; grad-tracked terminal steps expose the IDs in + ``info`` for an explicit reset after backward. Returns: Tuple of ``(obs, info)``. diff --git a/embodichain/lab/gym/envs/tasks/tableware/blocks_ranking_rgb.py b/embodichain/lab/gym/envs/tasks/tableware/blocks_ranking_rgb.py index 33bcc52a9..b93ffd8d7 100644 --- a/embodichain/lab/gym/envs/tasks/tableware/blocks_ranking_rgb.py +++ b/embodichain/lab/gym/envs/tasks/tableware/blocks_ranking_rgb.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import torch import numpy as np @@ -165,8 +167,8 @@ def _append_move_for_arm( ) # Joint space trajectory target_states = [ - PlanState(qpos=qpos_start[0], move_type=MoveType.JOINT_MOVE), - PlanState( + PlanState.single(qpos=qpos_start[0], move_type=MoveType.JOINT_MOVE), + PlanState.single( qpos=qpos_end[0], move_type=MoveType.JOINT_MOVE, ), @@ -175,7 +177,7 @@ def _append_move_for_arm( target_states=target_states, options=options ) - for qpos_item in plan_result.positions: + for qpos_item in plan_result.positions[0]: qpos = torch.as_tensor( qpos_item, dtype=torch.float32, device=self.device ) diff --git a/embodichain/lab/gym/envs/tasks/tableware/pour_water/action_bank.py b/embodichain/lab/gym/envs/tasks/tableware/pour_water/action_bank.py index c97e97a4a..e68fbf787 100644 --- a/embodichain/lab/gym/envs/tasks/tableware/pour_water/action_bank.py +++ b/embodichain/lab/gym/envs/tasks/tableware/pour_water/action_bank.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import torch import numpy as np from copy import deepcopy @@ -186,7 +188,9 @@ def plan_trajectory( ) plan_state = [ - PlanState(qpos=torch.as_tensor(qpos), move_type=MoveType.JOINT_MOVE) + PlanState.single( + qpos=torch.as_tensor(qpos), move_type=MoveType.JOINT_MOVE + ) for qpos in keyposes ] @@ -200,7 +204,7 @@ def plan_trajectory( ), ) - return ret.positions.detach().cpu().numpy().T + return ret.positions[0].detach().cpu().numpy().T @staticmethod @tag_edge diff --git a/embodichain/lab/gym/envs/tasks/tableware/stack_blocks_two.py b/embodichain/lab/gym/envs/tasks/tableware/stack_blocks_two.py index efabfe244..ca1fc1fbb 100644 --- a/embodichain/lab/gym/envs/tasks/tableware/stack_blocks_two.py +++ b/embodichain/lab/gym/envs/tasks/tableware/stack_blocks_two.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import torch import numpy as np @@ -206,11 +208,11 @@ def create_demo_action_list(self, *args, **kwargs): ) # Joint space trajectory target_states = [ - PlanState( + PlanState.single( qpos=torch.as_tensor(qpos_waypoints_np[start_idx]), move_type=MoveType.JOINT_MOVE, ), - PlanState( + PlanState.single( qpos=torch.as_tensor(qpos_waypoints_np[end_idx]), move_type=MoveType.JOINT_MOVE, ), @@ -220,7 +222,7 @@ def create_demo_action_list(self, *args, **kwargs): ) # Convert to torch and add to action list - for qpos_item in plan_result.positions: + for qpos_item in plan_result.positions[0]: qpos = torch.as_tensor( qpos_item, dtype=torch.float32, device=self.device ) diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index ba1333598..1f09d5509 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -827,6 +827,12 @@ def add_env_launcher_args_to_parser(parser: argparse.ArgumentParser) -> None: default=False, action="store_true", ) + parser.add_argument( + "--max_episodes", + help="Override the max_episodes value from the gym config.", + default=None, + type=int, + ) def merge_args_with_gym_config(args: argparse.Namespace, gym_config: dict) -> dict: @@ -849,6 +855,8 @@ def merge_args_with_gym_config(args: argparse.Namespace, gym_config: dict) -> di merged_config["physics"] = args.physics merged_config["gpu_id"] = args.gpu_id merged_config["arena_space"] = args.arena_space + if args.max_episodes is not None: + merged_config["max_episodes"] = args.max_episodes return merged_config diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index bd96d7bea..381bfd8ce 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -18,10 +18,13 @@ This module provides a unified interface for the atomic motion primitives (``move_end_effector``, ``move_joints``, ``pick_up``, ``move_held_object``, -``place``), with typed targets, a ``WorldState`` threaded across sequenced -actions, and extensible custom action registration. +``place``, ``press``, ``coordinated_pickment``, ``coordinated_placement``), +with typed targets, a ``WorldState`` threaded across sequenced actions, and +extensible custom action registration. """ +from __future__ import annotations + from .affordance import ( Affordance, AntipodalAffordance, @@ -31,6 +34,9 @@ ActionCfg, ActionResult, AtomicAction, + CoordinatedHeldObjectState, + CoordinatedPickmentTarget, + CoordinatedPlacementTarget, GraspTarget, HeldObjectState, HeldObjectPoseTarget, @@ -41,24 +47,30 @@ Target, WorldState, ) -from .actions import ( - MoveEndEffector, - MoveJoints, - MoveHeldObject, - PickUp, - Place, - MoveEndEffectorCfg, - MoveJointsCfg, - MoveHeldObjectCfg, - PickUpCfg, - PlaceCfg, -) from .engine import ( AtomicActionEngine, register_action, unregister_action, get_registered_actions, ) +from .primitives import ( + CoordinatedPickment, + CoordinatedPickmentCfg, + CoordinatedPlacement, + CoordinatedPlacementCfg, + MoveEndEffector, + MoveEndEffectorCfg, + MoveHeldObject, + MoveHeldObjectCfg, + MoveJoints, + MoveJointsCfg, + PickUp, + PickUpCfg, + Place, + PlaceCfg, + Press, + PressCfg, +) from .trajectory import TrajectoryBuilder __all__ = [ @@ -68,10 +80,13 @@ "InteractionPoints", "ObjectSemantics", "HeldObjectState", + "CoordinatedHeldObjectState", "HeldObjectPoseTarget", "JointPositionTarget", "NamedJointPositionTarget", "EndEffectorPoseTarget", + "CoordinatedPickmentTarget", + "CoordinatedPlacementTarget", "GraspTarget", "Target", "WorldState", @@ -79,16 +94,22 @@ "ActionCfg", "AtomicAction", # Action implementations + "CoordinatedPickment", + "CoordinatedPlacement", "MoveEndEffector", "MoveJoints", "MoveHeldObject", "PickUp", "Place", + "Press", + "CoordinatedPickmentCfg", + "CoordinatedPlacementCfg", "MoveEndEffectorCfg", "MoveJointsCfg", "MoveHeldObjectCfg", "PickUpCfg", "PlaceCfg", + "PressCfg", # Engine "AtomicActionEngine", "register_action", diff --git a/embodichain/lab/sim/atomic_actions/actions.py b/embodichain/lab/sim/atomic_actions/actions.py index ada502945..37911831b 100644 --- a/embodichain/lab/sim/atomic_actions/actions.py +++ b/embodichain/lab/sim/atomic_actions/actions.py @@ -14,794 +14,48 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Concrete atomic actions built on :class:`AtomicAction` and :class:`TrajectoryBuilder`. +"""Compatibility re-exports for built-in atomic actions. -Five sibling actions live here: :class:`MoveEndEffector`, :class:`MoveJoints`, -:class:`PickUp`, :class:`MoveHeldObject`, and :class:`Place`. Each inherits -:class:`AtomicAction` directly and composes a :class:`TrajectoryBuilder` for -shared trajectory math. ``execute`` takes a typed target plus a -:class:`WorldState` and returns an :class:`ActionResult` whose trajectory is -full-robot DoF shaped ``(n_envs, n_waypoints, robot.dof)``. +Concrete action implementations live in ``atomic_actions.primitives``. Importing +from ``atomic_actions.actions`` remains supported for existing callers. """ from __future__ import annotations -import torch -from typing import ClassVar - -from embodichain.lab.sim.planners import PlanState, MoveType -from embodichain.utils import configclass, logger -from embodichain.utils.math import pose_inv - -from .affordance import AntipodalAffordance -from .core import ( - ActionCfg, - ActionResult, - AtomicAction, - GraspTarget, - HeldObjectState, - HeldObjectPoseTarget, - JointPositionTarget, - NamedJointPositionTarget, - ObjectSemantics, - EndEffectorPoseTarget, - WorldState, +from .primitives import ( + CoordinatedPickment, + CoordinatedPickmentCfg, + CoordinatedPlacement, + CoordinatedPlacementCfg, + MoveEndEffector, + MoveEndEffectorCfg, + MoveHeldObject, + MoveHeldObjectCfg, + MoveJoints, + MoveJointsCfg, + PickUp, + PickUpCfg, + Place, + PlaceCfg, + Press, + PressCfg, ) -from .trajectory import TrajectoryBuilder - -# ============================================================================= -# Cfg classes (flat — no inheritance among action configs) -# ============================================================================= - - -@configclass -class MoveEndEffectorCfg(ActionCfg): - name: str = "move_end_effector" - """Name of the action, used for identification and logging.""" - - sample_interval: int = 50 - """Number of waypoints in the planned trajectory.""" - - -@configclass -class MoveJointsCfg(ActionCfg): - name: str = "move_joints" - """Name of the action, used for identification and logging.""" - - sample_interval: int = 50 - """Number of waypoints in the interpolated joint-space trajectory.""" - - named_joint_positions: dict[str, torch.Tensor] | None = None - """Optional named joint targets resolved by ``NamedJointPositionTarget``.""" - - -@configclass -class PickUpCfg(ActionCfg): - name: str = "pick_up" - """Name of the action, used for identification and logging.""" - - sample_interval: int = 80 - """Number of waypoints for the full trajectory (approach + hand + lift).""" - - hand_interp_steps: int = 5 - """Number of waypoints for the gripper close interpolation phase.""" - - hand_control_part: str = "hand" - """Name of the robot part that controls the hand joints.""" - - hand_open_qpos: torch.Tensor | None = None - """Joint positions for the open hand state, shape ``[hand_dof,]``.""" - - hand_close_qpos: torch.Tensor | None = None - """Joint positions for the closed hand state, shape ``[hand_dof,]``.""" - - lift_height: float = 0.1 - """Height (m) to lift the end-effector after closing the gripper.""" - - pre_grasp_distance: float = 0.15 - """Distance to offset back from the grasp pose along the approach direction.""" - - approach_direction: torch.Tensor = torch.tensor([0, 0, -1], dtype=torch.float32) - """Approach direction in the object local frame.""" - - -@configclass -class MoveHeldObjectCfg(ActionCfg): - name: str = "move_held_object" - """Name of the action, used for identification and logging.""" - - sample_interval: int = 50 - """Number of waypoints in the planned trajectory.""" - - hand_control_part: str = "hand" - """Name of the robot part that controls the hand joints.""" - - hand_close_qpos: torch.Tensor | None = None - """Joint positions for the closed hand state, shape ``[hand_dof,]``.""" - - -@configclass -class PlaceCfg(ActionCfg): - name: str = "place" - """Name of the action, used for identification and logging.""" - - sample_interval: int = 80 - """Number of waypoints for the full trajectory (down + hand + back).""" - - hand_interp_steps: int = 5 - """Number of waypoints for the gripper open interpolation phase.""" - - hand_control_part: str = "hand" - """Name of the robot part that controls the hand joints.""" - - hand_open_qpos: torch.Tensor | None = None - """Joint positions for the open hand state, shape ``[hand_dof,]``.""" - - hand_close_qpos: torch.Tensor | None = None - """Joint positions for the closed hand state, shape ``[hand_dof,]``.""" - - lift_height: float = 0.1 - """Height (m) to retract the end-effector after opening the gripper.""" - - -# ============================================================================= -# Shared helpers private to this module -# ============================================================================= - - -def _resolve_object_target( - target: torch.Tensor, *, n_envs: int, device: torch.device -) -> torch.Tensor: - """Broadcast an object target pose to ``(n_envs, 4, 4)`` or validate it.""" - target = target.to(device=device, dtype=torch.float32) - if target.shape == (4, 4): - target = target.unsqueeze(0).repeat(n_envs, 1, 1) - if target.shape != (n_envs, 4, 4): - logger.log_error( - f"object_target_pose must be (4, 4) or ({n_envs}, 4, 4), but got {target.shape}", - ValueError, - ) - return target - - -def _arm_qpos_from_state( - state: WorldState, arm_joint_ids, robot_dof: int -) -> torch.Tensor: - """Extract the arm slice of the full-DoF last_qpos carried in WorldState.""" - return state.last_qpos[:, arm_joint_ids] - - -# ============================================================================= -# MoveEndEffector -# ============================================================================= - - -class MoveEndEffector(AtomicAction): - """Plan a free-space end-effector move to a target pose. - - The :class:`EndEffectorPoseTarget` may carry either a single waypoint - ``(n_envs, 4, 4)`` (or a broadcastable ``(4, 4)``) or a multi-waypoint - trajectory ``(n_envs, n_waypoint, 4, 4)``. In the multi-waypoint case the - action plans a single trajectory that visits every waypoint in order, - starting from the inherited ``WorldState.last_qpos`` — IK is solved for each - waypoint with the previous waypoint's solution as the seed. - """ - - TargetType: ClassVar[type] = EndEffectorPoseTarget - - def __init__( - self, - motion_generator, - cfg: MoveEndEffectorCfg | None = None, - ) -> None: - super().__init__(motion_generator, cfg or MoveEndEffectorCfg()) - self.builder = TrajectoryBuilder(motion_generator) - self.n_envs = self.robot.get_qpos().shape[0] - self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.arm_dof = len(self.arm_joint_ids) - self.robot_dof = self.robot.dof - - def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionResult: - move_xpos = self.builder.resolve_pose_target(target.xpos, n_envs=self.n_envs) - start_qpos = self.builder.resolve_start_qpos( - _arm_qpos_from_state(state, self.arm_joint_ids, self.robot_dof), - n_envs=self.n_envs, - arm_dof=self.arm_dof, - control_part=self.cfg.control_part, - ) - target_states_list = self._build_target_states(move_xpos) - ok, arm_traj = self.builder.plan_arm_traj( - target_states_list, - start_qpos, - self.cfg.sample_interval, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - ) - if not ok: - return self._fail(state) - full = self._embed(arm_traj, state.last_qpos) - return ActionResult( - success=True, - trajectory=full, - next_state=WorldState( - last_qpos=full[:, -1, :].clone(), held_object=state.held_object - ), - ) - - def _build_target_states(self, move_xpos: torch.Tensor) -> list[list[PlanState]]: - """Build per-env PlanState lists from a single- or multi-waypoint target. - - ``move_xpos`` is the resolved target: 3D ``(n_envs, 4, 4)`` for a single - waypoint or 4D ``(n_envs, n_waypoint, 4, 4)`` for a trajectory. - """ - if move_xpos.dim() == 3: - move_xpos = move_xpos.unsqueeze(1) - n_waypoint = move_xpos.shape[1] - return [ - [ - PlanState(xpos=move_xpos[i, j], move_type=MoveType.EEF_MOVE) - for j in range(n_waypoint) - ] - for i in range(self.n_envs) - ] - - def _embed( - self, arm_traj: torch.Tensor, last_full_qpos: torch.Tensor - ) -> torch.Tensor: - n_wp = arm_traj.shape[1] - full = torch.empty( - (self.n_envs, n_wp, self.robot_dof), - dtype=torch.float32, - device=self.device, - ) - full[:, :, :] = last_full_qpos.unsqueeze(1) - full[:, :, self.arm_joint_ids] = arm_traj - return full - - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=False, - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, - ) - - -# ============================================================================= -# MoveJoints -# ============================================================================= - - -class MoveJoints(AtomicAction): - """Plan a joint-space move for the configured control part. - - The :class:`JointPositionTarget` may carry either a single waypoint - ``(n_envs, control_dof)`` or a multi-waypoint trajectory - ``(n_envs, n_waypoint, control_dof)``. In the multi-waypoint case the - action plans a single trajectory that visits every waypoint in order, - starting from the inherited ``WorldState.last_qpos``. - """ - - TargetType: ClassVar[tuple[type, ...]] = ( - JointPositionTarget, - NamedJointPositionTarget, - ) - - def __init__( - self, - motion_generator, - cfg: MoveJointsCfg | None = None, - ) -> None: - super().__init__(motion_generator, cfg or MoveJointsCfg()) - self.builder = TrajectoryBuilder(motion_generator) - self.n_envs = self.robot.get_qpos().shape[0] - self.joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.joint_dof = len(self.joint_ids) - self.robot_dof = self.robot.dof - self.named_joint_positions = self.cfg.named_joint_positions or {} - - def execute( - self, - target: JointPositionTarget | NamedJointPositionTarget, - state: WorldState, - ) -> ActionResult: - target_qpos = self.builder.resolve_joint_target( - self._resolve_target_qpos(target), - n_envs=self.n_envs, - joint_dof=self.joint_dof, - control_part=self.cfg.control_part, - ) - start_qpos = self.builder.resolve_start_qpos( - state.last_qpos[:, self.joint_ids], - n_envs=self.n_envs, - arm_dof=self.joint_dof, - control_part=self.cfg.control_part, - ) - joint_traj = self.builder.plan_joint_traj( - start_qpos, target_qpos, self.cfg.sample_interval - ) - full = self._embed(joint_traj, state.last_qpos) - return ActionResult( - success=True, - trajectory=full, - next_state=WorldState( - last_qpos=full[:, -1, :].clone(), held_object=state.held_object - ), - ) - - def _resolve_target_qpos( - self, target: JointPositionTarget | NamedJointPositionTarget - ) -> torch.Tensor: - if isinstance(target, JointPositionTarget): - return target.qpos - if target.name not in self.named_joint_positions: - logger.log_error( - f"Unknown named joint-position target '{target.name}' for " - f"MoveJoints. Available targets: {sorted(self.named_joint_positions)}", - KeyError, - ) - return self.named_joint_positions[target.name] - - def _embed( - self, joint_traj: torch.Tensor, last_full_qpos: torch.Tensor - ) -> torch.Tensor: - n_wp = joint_traj.shape[1] - full = torch.empty( - (self.n_envs, n_wp, self.robot_dof), - dtype=torch.float32, - device=self.device, - ) - full[:, :, :] = last_full_qpos.unsqueeze(1) - full[:, :, self.joint_ids] = joint_traj - return full - - -# ============================================================================= -# PickUp -# ============================================================================= - - -class PickUp(AtomicAction): - """Approach a grasp pose, close the gripper, lift.""" - - TargetType: ClassVar[type] = GraspTarget - - def __init__( - self, - motion_generator, - cfg: PickUpCfg | None = None, - ) -> None: - super().__init__(motion_generator, cfg or PickUpCfg()) - self.builder = TrajectoryBuilder(motion_generator) - self.n_envs = self.robot.get_qpos().shape[0] - self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part) - self.arm_dof = len(self.arm_joint_ids) - self.robot_dof = self.robot.dof - - if self.cfg.hand_open_qpos is None: - logger.log_error( - "hand_open_qpos must be specified in PickUpCfg", ValueError - ) - if self.cfg.hand_close_qpos is None: - logger.log_error( - "hand_close_qpos must be specified in PickUpCfg", ValueError - ) - self.hand_open_qpos = self.cfg.hand_open_qpos.to(self.device) - self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) - self.approach_direction = self.cfg.approach_direction.to(self.device) - - def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: - sem = target.semantics - if not isinstance(sem.affordance, AntipodalAffordance): - logger.log_error( - "PickUp requires an AntipodalAffordance on the target semantics.", - ValueError, - ) - if sem.entity is None: - logger.log_error( - "PickUp requires an entity on the target semantics.", ValueError - ) - - is_success, grasp_xpos = self._resolve_grasp_pose(sem) - if not self.builder.all_envs_success(is_success): - logger.log_warning("PickUp failed to resolve a grasp pose.") - return self._fail(state) - - # Pre-grasp by offsetting backwards along grasp z. - grasp_z = grasp_xpos[:, :3, 2] - pre_grasp_xpos = self.builder.apply_local_offset( - grasp_xpos, -grasp_z * self.cfg.pre_grasp_distance - ) - - start_arm_qpos = self.builder.resolve_start_qpos( - _arm_qpos_from_state(state, self.arm_joint_ids, self.robot_dof), - n_envs=self.n_envs, - arm_dof=self.arm_dof, - control_part=self.cfg.control_part, - ) - - n_approach, n_close, n_lift = self.builder.split_three_phase( - self.cfg.sample_interval, - self.cfg.hand_interp_steps, - first_phase_name="approach", - third_phase_name="lift", - ) - - # Phase 1: approach - target_states_list = [ - [ - PlanState(xpos=pre_grasp_xpos[i], move_type=MoveType.EEF_MOVE), - PlanState(xpos=grasp_xpos[i], move_type=MoveType.EEF_MOVE), - ] - for i in range(self.n_envs) - ] - ok, approach_arm = self.builder.plan_arm_traj( - target_states_list, - start_arm_qpos, - n_approach, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - ) - if not ok: - logger.log_warning("PickUp failed to plan the approach trajectory.") - return self._fail(state) - - # Phase 3: lift (planned from grasp qpos) - grasp_arm_qpos = approach_arm[:, -1, :] - lift_xpos = self.builder.apply_local_offset( - grasp_xpos, - torch.tensor([0, 0, 1], device=self.device) * self.cfg.lift_height, - ) - target_states_list = [ - [PlanState(xpos=lift_xpos[i], move_type=MoveType.EEF_MOVE)] - for i in range(self.n_envs) - ] - ok, lift_arm = self.builder.plan_arm_traj( - target_states_list, - grasp_arm_qpos, - n_lift, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - ) - if not ok: - logger.log_warning("PickUp failed to plan the lift trajectory.") - return self._fail(state) - - # Phase 2: hand close (arm held at grasp qpos) - hand_close_path = self.builder.interpolate_hand_qpos( - self.hand_open_qpos, self.hand_close_qpos, n_waypoints=n_close - ) - - full = torch.empty( - (self.n_envs, n_approach + n_close + n_lift, self.robot_dof), - dtype=torch.float32, - device=self.device, - ) - full[:, :, :] = state.last_qpos.unsqueeze(1) - full[:, :n_approach, self.arm_joint_ids] = approach_arm - full[:, :n_approach, self.hand_joint_ids] = self.hand_open_qpos - full[:, n_approach : n_approach + n_close, self.arm_joint_ids] = ( - grasp_arm_qpos.unsqueeze(1) - ) - full[:, n_approach : n_approach + n_close, self.hand_joint_ids] = ( - hand_close_path - ) - full[:, n_approach + n_close :, self.arm_joint_ids] = lift_arm - full[:, n_approach + n_close :, self.hand_joint_ids] = self.hand_close_qpos - - obj_poses = sem.entity.get_local_pose(to_matrix=True) - object_to_eef = torch.bmm(pose_inv(obj_poses), grasp_xpos) - held = HeldObjectState( - semantics=sem, object_to_eef=object_to_eef, grasp_xpos=grasp_xpos - ) - return ActionResult( - success=True, - trajectory=full, - next_state=WorldState(last_qpos=full[:, -1, :].clone(), held_object=held), - ) - - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=False, - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, - ) - - def _resolve_grasp_pose( - self, semantics: ObjectSemantics - ) -> tuple[torch.Tensor, torch.Tensor]: - obj_poses = semantics.entity.get_local_pose(to_matrix=True) - grasp_poses_result = semantics.affordance.get_valid_grasp_poses( - obj_poses=obj_poses, approach_direction=self.approach_direction - ) - n_envs = obj_poses.shape[0] - init_qpos = self.robot.get_qpos(name=self.cfg.control_part) - n_max_pose = max(r[0].shape[0] for r in grasp_poses_result) - grasp_xpos_padding = torch.zeros( - (n_envs, n_max_pose, 4, 4), dtype=torch.float32, device=self.device - ) - grasp_cost_padding = torch.full( - (n_envs, n_max_pose), - float("inf"), - dtype=torch.float32, - device=self.device, - ) - for i in range(n_envs): - n_pose = grasp_poses_result[i][0].shape[0] - grasp_xpos_padding[i, :n_pose] = grasp_poses_result[i][0] - grasp_cost_padding[i, :n_pose] = grasp_poses_result[i][1] - grasp_xpos_padding[i, n_pose:] = grasp_poses_result[i][0][0] - grasp_cost_padding[i, n_pose:] = grasp_poses_result[i][1][0] - init_qpos_repeat = init_qpos[:, None, :].repeat(1, n_max_pose, 1) - ik_success, _ = self.robot.compute_batch_ik( - pose=grasp_xpos_padding, - name=self.cfg.control_part, - joint_seed=init_qpos_repeat, - ) - grasp_cost_masked = torch.where(ik_success, grasp_cost_padding, 10000.0) - best_cost, best_idx = grasp_cost_masked.min(dim=1) - is_success = best_cost < 9999.0 - best_grasp_xpos = grasp_xpos_padding[ - torch.arange(n_envs, device=self.device), best_idx - ] - return is_success, best_grasp_xpos - - -# ============================================================================= -# MoveHeldObject -# ============================================================================= - - -class MoveHeldObject(AtomicAction): - """Move the held object to a target object pose; keep the gripper closed.""" - - TargetType: ClassVar[type] = HeldObjectPoseTarget - - def __init__( - self, - motion_generator, - cfg: MoveHeldObjectCfg | None = None, - ) -> None: - super().__init__(motion_generator, cfg or MoveHeldObjectCfg()) - self.builder = TrajectoryBuilder(motion_generator) - self.n_envs = self.robot.get_qpos().shape[0] - self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part) - self.arm_dof = len(self.arm_joint_ids) - self.robot_dof = self.robot.dof - - if self.cfg.hand_close_qpos is None: - logger.log_error( - "hand_close_qpos must be specified in MoveHeldObjectCfg", ValueError - ) - self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) - - def execute(self, target: HeldObjectPoseTarget, state: WorldState) -> ActionResult: - if state.held_object is None: - logger.log_error( - "MoveHeldObject requires WorldState.held_object — run PickUp first.", - ValueError, - ) - object_target_pose = _resolve_object_target( - target.object_target_pose, n_envs=self.n_envs, device=self.device - ) - object_to_eef = state.held_object.object_to_eef.to( - device=self.device, dtype=torch.float32 - ) - if object_to_eef.shape == (4, 4): - object_to_eef = object_to_eef.unsqueeze(0).repeat(self.n_envs, 1, 1) - move_eef_xpos = torch.bmm(object_target_pose, object_to_eef) - - start_arm_qpos = self.builder.resolve_start_qpos( - _arm_qpos_from_state(state, self.arm_joint_ids, self.robot_dof), - n_envs=self.n_envs, - arm_dof=self.arm_dof, - control_part=self.cfg.control_part, - ) - - target_states_list = [ - [PlanState(xpos=move_eef_xpos[i], move_type=MoveType.EEF_MOVE)] - for i in range(self.n_envs) - ] - ok, arm_traj = self.builder.plan_arm_traj( - target_states_list, - start_arm_qpos, - self.cfg.sample_interval, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - ) - if not ok: - logger.log_warning("MoveHeldObject failed to plan trajectory.") - return self._fail(state) - - full = torch.empty( - (self.n_envs, arm_traj.shape[1], self.robot_dof), - dtype=torch.float32, - device=self.device, - ) - full[:, :, :] = state.last_qpos.unsqueeze(1) - full[:, :, self.arm_joint_ids] = arm_traj - full[:, :, self.hand_joint_ids] = self.hand_close_qpos - - return ActionResult( - success=True, - trajectory=full, - next_state=WorldState( - last_qpos=full[:, -1, :].clone(), - held_object=state.held_object, - ), - ) - - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=False, - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, - ) - - -# ============================================================================= -# Place -# ============================================================================= - - -class Place(AtomicAction): - """Lower the held object to a place pose, open the gripper, retract. - - The :class:`EndEffectorPoseTarget` may carry either a single waypoint - ``(n_envs, 4, 4)`` (or a broadcastable ``(4, 4)``) or a multi-waypoint - trajectory ``(n_envs, n_waypoint, 4, 4)``. In the multi-waypoint case the - down phase visits every waypoint in order — approaching from above the - first waypoint, descending through each waypoint, then opening the gripper - at the final waypoint and retracting to above the last waypoint. Starting - joint positions are inherited from ``WorldState.last_qpos``. - """ - - TargetType: ClassVar[type] = EndEffectorPoseTarget - - def __init__( - self, - motion_generator, - cfg: PlaceCfg | None = None, - ) -> None: - super().__init__(motion_generator, cfg or PlaceCfg()) - self.builder = TrajectoryBuilder(motion_generator) - self.n_envs = self.robot.get_qpos().shape[0] - self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part) - self.arm_dof = len(self.arm_joint_ids) - self.robot_dof = self.robot.dof - - if self.cfg.hand_open_qpos is None: - logger.log_error("hand_open_qpos must be specified in PlaceCfg", ValueError) - if self.cfg.hand_close_qpos is None: - logger.log_error( - "hand_close_qpos must be specified in PlaceCfg", ValueError - ) - self.hand_open_qpos = self.cfg.hand_open_qpos.to(self.device) - self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) - - def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionResult: - place_xpos = self.builder.resolve_pose_target(target.xpos, n_envs=self.n_envs) - # Normalize a single-waypoint (n_envs, 4, 4) target to (n_envs, 1, 4, 4) - # so the multi-waypoint descent path below is uniform. - if place_xpos.dim() == 3: - place_xpos = place_xpos.unsqueeze(1) - n_waypoint = place_xpos.shape[1] - - start_arm_qpos = self.builder.resolve_start_qpos( - _arm_qpos_from_state(state, self.arm_joint_ids, self.robot_dof), - n_envs=self.n_envs, - arm_dof=self.arm_dof, - control_part=self.cfg.control_part, - ) - n_down, n_open, n_back = self.builder.split_three_phase( - self.cfg.sample_interval, - self.cfg.hand_interp_steps, - first_phase_name="approach", - third_phase_name="back", - ) - - lift_offset = torch.tensor([0, 0, 1], device=self.device) * self.cfg.lift_height - # Approach from above the first waypoint; retract to above the last. - # For a single waypoint these coincide, matching the legacy behavior. - approach_xpos = self.builder.apply_local_offset(place_xpos[:, 0], lift_offset) - retract_xpos = self.builder.apply_local_offset(place_xpos[:, -1], lift_offset) - - # Phase 1: down (approach → every place waypoint in order) - target_states_list = [ - [PlanState(xpos=approach_xpos[i], move_type=MoveType.EEF_MOVE)] - + [ - PlanState(xpos=place_xpos[i, j], move_type=MoveType.EEF_MOVE) - for j in range(n_waypoint) - ] - for i in range(self.n_envs) - ] - ok, down_arm = self.builder.plan_arm_traj( - target_states_list, - start_arm_qpos, - n_down, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - ) - if not ok: - return self._fail(state) - reach_arm_qpos = down_arm[:, -1, :] - - # Phase 3: back (retract to above the last waypoint) - target_states_list = [ - [PlanState(xpos=retract_xpos[i], move_type=MoveType.EEF_MOVE)] - for i in range(self.n_envs) - ] - ok, back_arm = self.builder.plan_arm_traj( - target_states_list, - reach_arm_qpos, - n_back, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - ) - if not ok: - return self._fail(state) - - # Phase 2: hand open (arm held at reach qpos) - hand_open_path = self.builder.interpolate_hand_qpos( - self.hand_close_qpos, self.hand_open_qpos, n_waypoints=n_open - ) - - full = torch.empty( - (self.n_envs, n_down + n_open + n_back, self.robot_dof), - dtype=torch.float32, - device=self.device, - ) - full[:, :, :] = state.last_qpos.unsqueeze(1) - full[:, :n_down, self.arm_joint_ids] = down_arm - full[:, :n_down, self.hand_joint_ids] = self.hand_close_qpos - full[:, n_down : n_down + n_open, self.arm_joint_ids] = ( - reach_arm_qpos.unsqueeze(1) - ) - full[:, n_down : n_down + n_open, self.hand_joint_ids] = hand_open_path - full[:, n_down + n_open :, self.arm_joint_ids] = back_arm - full[:, n_down + n_open :, self.hand_joint_ids] = self.hand_open_qpos - - return ActionResult( - success=True, - trajectory=full, - next_state=WorldState(last_qpos=full[:, -1, :].clone(), held_object=None), - ) - - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=False, - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, - ) - __all__ = [ + "CoordinatedPickment", + "CoordinatedPickmentCfg", + "CoordinatedPlacement", + "CoordinatedPlacementCfg", "MoveEndEffector", "MoveEndEffectorCfg", - "MoveJoints", - "MoveJointsCfg", "MoveHeldObject", "MoveHeldObjectCfg", + "MoveJoints", + "MoveJointsCfg", "PickUp", "PickUpCfg", "Place", "PlaceCfg", + "Press", + "PressCfg", ] diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index 0f9f29094..4e80a4fa8 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -19,7 +19,7 @@ import torch from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Any, ClassVar, TYPE_CHECKING +from typing import Any, ClassVar, Literal, TYPE_CHECKING from embodichain.lab.sim.common import BatchEntity from embodichain.utils import configclass @@ -68,7 +68,7 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class EndEffectorPoseTarget: - """End-effector pose target. Used by MoveEndEffector and Place.""" + """End-effector pose target. Used by MoveEndEffector, Place, and Press.""" xpos: torch.Tensor """Target end-effector homogeneous transform. @@ -77,9 +77,25 @@ class EndEffectorPoseTarget: - ``(4, 4)`` or ``(n_envs, 4, 4)`` — a single waypoint. - ``(n_envs, n_waypoint, 4, 4)`` — a multi-waypoint trajectory; waypoints - are visited in order. (Only consumed as multi-waypoint by MoveEndEffector.) + are visited in order. (Consumed as multi-waypoint by MoveEndEffector and + Place.) """ + tcp_symmetry: Literal["none", "z_roll_180"] = "none" + """Optional TCP-frame symmetry allowed by the target semantics. + + ``"none"`` preserves the pose exactly. ``"z_roll_180"`` lets supporting + actions choose between the pose and its TCP z-roll 180 equivalent, which + flips TCP x/y while preserving TCP z and translation. + """ + + def __post_init__(self) -> None: + if self.tcp_symmetry not in ("none", "z_roll_180"): + raise ValueError( + "tcp_symmetry must be one of 'none' or 'z_roll_180', " + f"but got {self.tcp_symmetry!r}" + ) + @dataclass(frozen=True) class JointPositionTarget: @@ -106,10 +122,19 @@ class NamedJointPositionTarget: @dataclass(frozen=True) class GraspTarget: - """Pickup target. The grasp pose is solved from the affordance + entity at execute time.""" + """Pickup target with an affordance-selected or explicitly supplied grasp pose.""" semantics: ObjectSemantics + grasp_xpos: torch.Tensor | None = None + """Optional end-effector grasp pose. + + When omitted, :class:`PickUp` selects a grasp from the target affordance. + Supplying a pose with shape ``(4, 4)`` or ``(n_envs, 4, 4)`` skips grasp + sampling, which is useful when perception or task geometry has already + selected a grasp. + """ + @dataclass(frozen=True) class HeldObjectPoseTarget: @@ -119,12 +144,60 @@ class HeldObjectPoseTarget: """(4, 4) or (n_envs, 4, 4) target pose for the held object.""" +@dataclass(frozen=True) +class CoordinatedPickmentTarget: + """Object-centric target for picking and moving one object with two hands.""" + + object_target_pose: torch.Tensor + """Target pose for the shared object, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" + + object_semantics: ObjectSemantics + """Semantic description of the shared object.""" + + left_object_to_eef: torch.Tensor + """Transform from object frame to left end-effector frame.""" + + right_object_to_eef: torch.Tensor + """Transform from object frame to right end-effector frame.""" + + object_initial_pose: torch.Tensor | None = None + """Optional initial object pose. Defaults to ``object_semantics.entity`` pose.""" + + +@dataclass(frozen=True) +class CoordinatedPlacementTarget: + """Object-centric target for dual-arm coordinated placement.""" + + placing_object_target_pose: torch.Tensor + """Target pose for the object released by the placing arm.""" + + support_object_target_pose: torch.Tensor + """Target pose for the object held by the support arm.""" + + placing_held_object: HeldObjectState + """Held-object state for the placing arm.""" + + support_held_object: HeldObjectState + """Held-object state for the support arm.""" + + placing_height_offset: float | None = None + """World-Z offset above the placing object target pose.""" + + support_height_offset: float | None = None + """World-Z offset above the support object target pose.""" + + release: bool | None = None + """Whether the placing hand releases. ``None`` uses the action config.""" + + Target = ( EndEffectorPoseTarget | JointPositionTarget | NamedJointPositionTarget | GraspTarget | HeldObjectPoseTarget + | CoordinatedPickmentTarget + | CoordinatedPlacementTarget ) @@ -147,6 +220,26 @@ class HeldObjectState: """Batched end-effector pose used to grasp the object, shape [n_envs, 4, 4].""" +@dataclass +class CoordinatedHeldObjectState: + """State of a single object jointly held by two robot hands.""" + + semantics: ObjectSemantics + """Semantic object currently held by the two grippers.""" + + left_object_to_eef: torch.Tensor + """Transform from object frame to left end-effector frame, shape ``[n_envs, 4, 4]``.""" + + right_object_to_eef: torch.Tensor + """Transform from object frame to right end-effector frame, shape ``[n_envs, 4, 4]``.""" + + left_grasp_xpos: torch.Tensor + """Left end-effector grasp pose for the shared object, shape ``[n_envs, 4, 4]``.""" + + right_grasp_xpos: torch.Tensor + """Right end-effector grasp pose for the shared object, shape ``[n_envs, 4, 4]``.""" + + @dataclass class WorldState: """State the engine threads through a sequence of actions.""" @@ -157,13 +250,17 @@ class WorldState: held_object: HeldObjectState | None = None """Object currently held by the gripper, or None.""" + coordinated_held_object: CoordinatedHeldObjectState | None = None + """Object currently held by two grippers, or None.""" + @dataclass class ActionResult: """Return value of every AtomicAction.execute call.""" - success: bool - """Whether the action produced a valid full-DoF trajectory.""" + success: bool | torch.Tensor + """Whether the action produced a valid full-DoF trajectory. + Can be a bool or a per-environment boolean tensor of shape (n_envs,).""" trajectory: torch.Tensor """Full-robot trajectory, shape (n_envs, n_waypoints, robot.dof).""" @@ -171,6 +268,23 @@ class ActionResult: next_state: WorldState """World state to feed into the next action.""" + @property + def success_all(self) -> bool: + """True only if all environments succeeded.""" + if isinstance(self.success, torch.Tensor): + return bool(torch.all(self.success).item()) + return bool(self.success) + + def __bool__(self) -> bool: + import warnings as _w + + _w.warn( + "ActionResult bool() is deprecated; use .success_all", + DeprecationWarning, + stacklevel=2, + ) + return self.success_all + # ============================================================================= # Configuration base @@ -186,6 +300,12 @@ class ActionCfg: interpolation_type: str = "linear" velocity_limit: float | None = None acceleration_limit: float | None = None + motion_source: str = "ik_interp" + """Trajectory source: 'ik_interp' (default, batched IK + linear interp) + or 'motion_gen' (batched MotionGenerator).""" + planner_type: str | None = None + """Planner type for motion_source='motion_gen': 'toppra' | 'neural'. + Required when motion_source='motion_gen'.""" # ============================================================================= @@ -232,6 +352,9 @@ def execute(self, target: Target, state: WorldState) -> ActionResult: "ActionCfg", "ActionResult", "AtomicAction", + "CoordinatedHeldObjectState", + "CoordinatedPickmentTarget", + "CoordinatedPlacementTarget", "GraspTarget", "HeldObjectState", "HeldObjectPoseTarget", diff --git a/embodichain/lab/sim/atomic_actions/engine.py b/embodichain/lab/sim/atomic_actions/engine.py index 4ab3e306d..59574af4b 100644 --- a/embodichain/lab/sim/atomic_actions/engine.py +++ b/embodichain/lab/sim/atomic_actions/engine.py @@ -90,7 +90,7 @@ def run( self, steps: Iterable[tuple[str, Target]], state: WorldState | None = None, - ) -> tuple[bool, torch.Tensor, WorldState]: + ) -> tuple[torch.Tensor, torch.Tensor, WorldState]: """Run a sequence of named actions, threading WorldState through. Args: @@ -100,10 +100,10 @@ def run( Returns: ``(success, concatenated_full_dof_trajectory, final_state)``. - On failure ``success`` is False and the trajectory is the - concatenation of steps that completed before the failure; the - failing step contributes no waypoints. ``final_state`` is the - state going into the failed step. + ``success`` is a ``(B,)`` boolean tensor indicating which + environments completed every step. Failed environments hold their + last successful joint position in both ``full_traj`` and + ``final_state.last_qpos`` for the remainder of the sequence. An empty ``steps`` iterable is a successful no-op returning an empty trajectory and the seed state. @@ -111,11 +111,13 @@ def run( if state is None: state = WorldState(last_qpos=self.robot.get_qpos().clone()) + b = state.last_qpos.shape[0] full_traj = torch.empty( - (state.last_qpos.shape[0], 0, self.robot.dof), + (b, 0, self.robot.dof), dtype=torch.float32, device=self.device, ) + alive = torch.ones(b, dtype=torch.bool, device=self.device) for name, target in steps: if name not in self._actions: @@ -127,13 +129,29 @@ def run( f"{_target_type_name(action.TargetType)}, got {type(target).__name__}", TypeError, ) + if not alive.any(): + # All envs dead: no further motion to plan. + break + prev_last_qpos = state.last_qpos.clone() result: ActionResult = action.execute(target, state) - if not result.success: - return False, full_traj, state - full_traj = torch.cat([full_traj, result.trajectory], dim=1) + step_success = ( + result.success + if isinstance(result.success, torch.Tensor) + else torch.tensor(bool(result.success), device=self.device) + ) + step_success = step_success.to(self.device) + alive = alive & step_success + # Failed envs freeze at their last successful qpos for this step's trajectory. + traj = result.trajectory + held_rows = prev_last_qpos.unsqueeze(1).repeat(1, traj.shape[1], 1) + traj = torch.where(alive[:, None, None], traj, held_rows) + full_traj = torch.cat([full_traj, traj], dim=1) state = result.next_state + state.last_qpos = torch.where( + alive[:, None], state.last_qpos, prev_last_qpos + ) - return True, full_traj, state + return alive, full_traj, state __all__ = [ diff --git a/embodichain/lab/sim/atomic_actions/primitives/__init__.py b/embodichain/lab/sim/atomic_actions/primitives/__init__.py new file mode 100644 index 000000000..0998f4d18 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/__init__.py @@ -0,0 +1,47 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Built-in atomic action primitive implementations.""" + +from __future__ import annotations + +from .coordinated_pickment import CoordinatedPickment, CoordinatedPickmentCfg +from .coordinated_placement import CoordinatedPlacement, CoordinatedPlacementCfg +from .move_end_effector import MoveEndEffector, MoveEndEffectorCfg +from .move_held_object import MoveHeldObject, MoveHeldObjectCfg +from .move_joints import MoveJoints, MoveJointsCfg +from .pick_up import PickUp, PickUpCfg +from .place import Place, PlaceCfg +from .press import Press, PressCfg + +__all__ = [ + "CoordinatedPickment", + "CoordinatedPickmentCfg", + "CoordinatedPlacement", + "CoordinatedPlacementCfg", + "MoveEndEffector", + "MoveEndEffectorCfg", + "MoveHeldObject", + "MoveHeldObjectCfg", + "MoveJoints", + "MoveJointsCfg", + "PickUp", + "PickUpCfg", + "Place", + "PlaceCfg", + "Press", + "PressCfg", +] diff --git a/embodichain/lab/sim/atomic_actions/primitives/_helpers.py b/embodichain/lab/sim/atomic_actions/primitives/_helpers.py new file mode 100644 index 000000000..2aecb8cad --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/_helpers.py @@ -0,0 +1,52 @@ +# ---------------------------------------------------------------------------- +# 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 internal helpers for concrete atomic actions.""" + +from __future__ import annotations + +import torch + +from embodichain.utils import logger + +from ..core import WorldState + + +def resolve_object_target( + target: torch.Tensor, + *, + n_envs: int, + device: torch.device, + name: str = "object_target_pose", +) -> torch.Tensor: + """Broadcast an object target pose to ``(n_envs, 4, 4)`` or validate it.""" + target = target.to(device=device, dtype=torch.float32) + if target.shape == (4, 4): + target = target.unsqueeze(0).repeat(n_envs, 1, 1) + if target.shape != (n_envs, 4, 4): + logger.log_error( + f"{name} must be (4, 4) or ({n_envs}, 4, 4), but got {target.shape}", + ValueError, + ) + return target + + +def arm_qpos_from_state(state: WorldState, arm_joint_ids: list[int]) -> torch.Tensor: + """Extract the arm slice of the full-DoF ``last_qpos`` carried in state.""" + return state.last_qpos[:, arm_joint_ids] + + +__all__ = ["arm_qpos_from_state", "resolve_object_target"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py new file mode 100644 index 000000000..d8c2b9c72 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py @@ -0,0 +1,782 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""CoordinatedPickment atomic action implementation.""" + +from __future__ import annotations + +from typing import ClassVar + +import torch + +from embodichain.utils import configclass, logger +from embodichain.utils.math import matrix_from_quat, quat_from_matrix + +from ..core import ( + ActionCfg, + ActionResult, + AtomicAction, + CoordinatedHeldObjectState, + CoordinatedPickmentTarget, + WorldState, +) +from ..trajectory import TrajectoryBuilder + + +@configclass +class CoordinatedPickmentCfg(ActionCfg): + name: str = "coordinated_pickment" + """Name of the action, used for identification and logging.""" + + control_part: str = "dual_arm" + """Combined control part containing left and right arm joints.""" + + left_arm_control_part: str = "left_arm" + """Left arm control part used to grasp one end of the object.""" + + right_arm_control_part: str = "right_arm" + """Right arm control part used to grasp the other end of the object.""" + + left_hand_control_part: str = "left_hand" + """Hand attached to the left arm.""" + + right_hand_control_part: str = "right_hand" + """Hand attached to the right arm.""" + + left_hand_open_qpos: torch.Tensor | None = None + """Left hand qpos for the open state.""" + + left_hand_close_qpos: torch.Tensor | None = None + """Left hand qpos for the closed state.""" + + right_hand_open_qpos: torch.Tensor | None = None + """Right hand qpos for the open state.""" + + right_hand_close_qpos: torch.Tensor | None = None + """Right hand qpos for the closed state.""" + + object_motion_keyframes: int = 6 + """Number of object-pose keyframes solved by IK before joint-space interpolation.""" + + pre_grasp_distance: float = 0.10 + """World distance to retreat from each grasp pose along negative TCP z.""" + + lift_height: float = 0.08 + """World-Z lift distance before moving to the object target pose.""" + + sample_interval: int = 120 + """Number of waypoints for the full coordinated pickment trajectory.""" + + hand_interp_steps: int = 10 + """Number of waypoints used for the simultaneous hand close phase.""" + + hold_steps: int = 4 + """Number of waypoints to hold the final object target pose.""" + + +class _DualArmHelpers: + """Shared trajectory helpers for dual-arm coordinated actions.""" + + def _init_dual_arm_parts( + self, + *, + first_arm_control_part: str, + second_arm_control_part: str, + first_hand_control_part: str, + second_hand_control_part: str, + ) -> None: + self.builder = TrajectoryBuilder(self.motion_generator) + self.n_envs = self.robot.get_qpos().shape[0] + self.robot_dof = self.robot.dof + self.dual_arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) + self.first_arm_joint_ids = self.robot.get_joint_ids(name=first_arm_control_part) + self.second_arm_joint_ids = self.robot.get_joint_ids( + name=second_arm_control_part + ) + self.first_hand_joint_ids = self.robot.get_joint_ids( + name=first_hand_control_part + ) + self.second_hand_joint_ids = self.robot.get_joint_ids( + name=second_hand_control_part + ) + self.first_arm_dof = len(self.first_arm_joint_ids) + self.second_arm_dof = len(self.second_arm_joint_ids) + self.dual_arm_dof = len(self.dual_arm_joint_ids) + self.first_hand_dof = len(self.first_hand_joint_ids) + self.second_hand_dof = len(self.second_hand_joint_ids) + self._dual_id_to_col = { + joint_id: col for col, joint_id in enumerate(self.dual_arm_joint_ids) + } + self._first_arm_cols = self._lookup_joint_columns( + self.first_arm_joint_ids, + self._dual_id_to_col, + first_arm_control_part, + ) + self._second_arm_cols = self._lookup_joint_columns( + self.second_arm_joint_ids, + self._dual_id_to_col, + second_arm_control_part, + ) + + @staticmethod + def _lookup_joint_columns( + joint_ids: list[int], + joint_id_to_col: dict[int, int], + control_part: str, + ) -> list[int]: + missing = [ + joint_id for joint_id in joint_ids if joint_id not in joint_id_to_col + ] + if missing: + logger.log_error( + f"Joints {missing} from '{control_part}' are not included in " + "the configured dual-arm control part.", + ValueError, + ) + return [joint_id_to_col[joint_id] for joint_id in joint_ids] + + def _fail(self, state: WorldState) -> ActionResult: + return ActionResult( + success=False, + trajectory=torch.empty( + (self.n_envs, 0, self.robot_dof), + dtype=torch.float32, + device=self.device, + ), + next_state=state, + ) + + def _expand_qpos(self, qpos: torch.Tensor, dof: int, name: str) -> torch.Tensor: + qpos = qpos.to(device=self.device, dtype=torch.float32) + if qpos.shape == (dof,): + return qpos.unsqueeze(0).repeat(self.n_envs, 1) + if qpos.shape == (self.n_envs, dof): + return qpos + logger.log_error( + f"{name} must have shape ({dof},) or " + f"({self.n_envs}, {dof}), but got {qpos.shape}", + ValueError, + ) + raise AssertionError("unreachable") + + def _resolve_pose(self, pose: torch.Tensor, name: str) -> torch.Tensor: + pose = pose.to(device=self.device, dtype=torch.float32) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).repeat(self.n_envs, 1, 1) + if pose.shape != (self.n_envs, 4, 4): + logger.log_error( + f"{name} must have shape (4, 4) or " + f"({self.n_envs}, 4, 4), but got {pose.shape}", + ValueError, + ) + return pose + + def _resolve_dual_arm_start( + self, + state: WorldState, + ) -> tuple[torch.Tensor, torch.Tensor]: + dual_start = state.last_qpos[:, self.dual_arm_joint_ids].to( + device=self.device, dtype=torch.float32 + ) + return ( + dual_start[:, self._first_arm_cols], + dual_start[:, self._second_arm_cols], + ) + + def _plan_named_arm_trajectory( + self, + control_part: str, + start_qpos: torch.Tensor, + target_poses: torch.Tensor, + n_waypoints: int, + ) -> tuple[bool, torch.Tensor]: + n_state = target_poses.shape[1] + arm_dof = start_qpos.shape[-1] + trajectory = torch.zeros( + (self.n_envs, n_state, arm_dof), + dtype=torch.float32, + device=self.device, + ) + qpos_seed = start_qpos + for i in range(n_state): + is_success, qpos = self.robot.compute_ik( + pose=target_poses[:, i], + name=control_part, + joint_seed=qpos_seed, + ) + if not self.builder.all_envs_success(is_success): + logger.log_warning( + f"Failed to compute IK for {control_part} target state {i}." + ) + return False, trajectory + trajectory[:, i] = qpos + qpos_seed = qpos + + trajectory = torch.cat([start_qpos.unsqueeze(1), trajectory], dim=1) + return True, ( + self.builder.plan_joint_traj( + trajectory[:, 0], + trajectory[:, -1], + n_waypoints, + ) + if n_state == 1 + else self._interpolate_keyframe_qpos(trajectory, n_waypoints) + ) + + def _compose_dual_arm_trajectory( + self, + first_arm_traj: torch.Tensor, + second_arm_traj: torch.Tensor, + ) -> torch.Tensor: + n_waypoints = first_arm_traj.shape[1] + dual_arm_traj = torch.zeros( + (self.n_envs, n_waypoints, self.dual_arm_dof), + dtype=torch.float32, + device=self.device, + ) + dual_arm_traj[:, :, self._first_arm_cols] = first_arm_traj + dual_arm_traj[:, :, self._second_arm_cols] = second_arm_traj + return dual_arm_traj + + def _assemble_phase( + self, + state: WorldState, + first_arm_traj: torch.Tensor, + second_arm_traj: torch.Tensor, + first_hand_traj: torch.Tensor, + second_hand_traj: torch.Tensor, + ) -> torch.Tensor: + n_waypoints = first_arm_traj.shape[1] + full = torch.empty( + (self.n_envs, n_waypoints, self.robot_dof), + dtype=torch.float32, + device=self.device, + ) + full[:, :, :] = state.last_qpos.to(self.device).unsqueeze(1) + full[:, :, self.dual_arm_joint_ids] = self._compose_dual_arm_trajectory( + first_arm_traj, second_arm_traj + ) + full[:, :, self.first_hand_joint_ids] = first_hand_traj + full[:, :, self.second_hand_joint_ids] = second_hand_traj + return full + + @staticmethod + def _repeat_qpos(qpos: torch.Tensor, n_waypoints: int) -> torch.Tensor: + return qpos.unsqueeze(1).repeat(1, n_waypoints, 1) + + def _interpolate_qpos( + self, + start_qpos: torch.Tensor, + end_qpos: torch.Tensor, + n_waypoints: int, + ) -> torch.Tensor: + weights = torch.linspace( + 0.0, + 1.0, + steps=n_waypoints, + device=self.device, + dtype=start_qpos.dtype, + ) + return torch.lerp( + start_qpos.unsqueeze(1), + end_qpos.unsqueeze(1), + weights[None, :, None], + ) + + def _interpolate_keyframe_qpos( + self, keyframe_qpos: torch.Tensor, n_waypoints: int + ) -> torch.Tensor: + n_keyframes = keyframe_qpos.shape[1] + keyframe_indices = ( + torch.linspace( + 0, + n_waypoints - 1, + steps=n_keyframes, + device=self.device, + ) + .round() + .to(dtype=torch.long) + ) + return self._interpolate_qpos_keyframes( + keyframe_qpos, keyframe_indices, n_waypoints + ) + + def _interpolate_qpos_keyframes( + self, + keyframe_qpos: torch.Tensor, + keyframe_indices: torch.Tensor, + n_waypoints: int, + ) -> torch.Tensor: + trajectory = torch.zeros( + (self.n_envs, n_waypoints, keyframe_qpos.shape[-1]), + dtype=torch.float32, + device=self.device, + ) + for segment_idx in range(len(keyframe_indices) - 1): + start_idx = int(keyframe_indices[segment_idx].item()) + end_idx = int(keyframe_indices[segment_idx + 1].item()) + n_segment = end_idx - start_idx + 1 + weights = torch.linspace( + 0.0, + 1.0, + steps=n_segment, + dtype=keyframe_qpos.dtype, + device=self.device, + ) + segment = torch.lerp( + keyframe_qpos[:, segment_idx : segment_idx + 1], + keyframe_qpos[:, segment_idx + 1 : segment_idx + 2], + weights[None, :, None], + ) + trajectory[:, start_idx : end_idx + 1] = segment + return trajectory + + def _interpolate_object_pose( + self, + start_pose: torch.Tensor, + end_pose: torch.Tensor, + n_waypoints: int, + *, + include_orientation: bool, + ) -> torch.Tensor: + weights = torch.linspace( + 0.0, + 1.0, + steps=n_waypoints, + device=self.device, + dtype=start_pose.dtype, + ) + poses = start_pose.unsqueeze(1).repeat(1, n_waypoints, 1, 1) + poses[:, :, :3, 3] = torch.lerp( + start_pose[:, None, :3, 3], + end_pose[:, None, :3, 3], + weights[None, :, None], + ) + if not include_orientation: + return poses + + start_quat = quat_from_matrix(start_pose[:, :3, :3]) + end_quat = quat_from_matrix(end_pose[:, :3, :3]) + quat_dot = torch.sum(start_quat * end_quat, dim=-1, keepdim=True) + end_quat = torch.where(quat_dot < 0.0, -end_quat, end_quat) + quat = torch.lerp( + start_quat.unsqueeze(1), + end_quat.unsqueeze(1), + weights[None, :, None], + ) + quat = quat / torch.linalg.norm(quat, dim=-1, keepdim=True).clamp_min(1e-8) + poses[:, :, :3, :3] = matrix_from_quat(quat.reshape(-1, 4)).reshape( + self.n_envs, n_waypoints, 3, 3 + ) + return poses + + +class CoordinatedPickment(AtomicAction): + """Pick and move a single object pinched by two hands.""" + + TargetType: ClassVar[type] = CoordinatedPickmentTarget + + _assemble_phase = _DualArmHelpers._assemble_phase + _compose_dual_arm_trajectory = _DualArmHelpers._compose_dual_arm_trajectory + _expand_qpos = _DualArmHelpers._expand_qpos + _fail = _DualArmHelpers._fail + _init_dual_arm_parts = _DualArmHelpers._init_dual_arm_parts + _interpolate_keyframe_qpos = _DualArmHelpers._interpolate_keyframe_qpos + _interpolate_object_pose = _DualArmHelpers._interpolate_object_pose + _interpolate_qpos = _DualArmHelpers._interpolate_qpos + _interpolate_qpos_keyframes = _DualArmHelpers._interpolate_qpos_keyframes + _lookup_joint_columns = staticmethod(_DualArmHelpers._lookup_joint_columns) + _plan_named_arm_trajectory = _DualArmHelpers._plan_named_arm_trajectory + _repeat_qpos = staticmethod(_DualArmHelpers._repeat_qpos) + _resolve_dual_arm_start = _DualArmHelpers._resolve_dual_arm_start + _resolve_pose = _DualArmHelpers._resolve_pose + + def __init__( + self, + motion_generator, + cfg: CoordinatedPickmentCfg | None = None, + ) -> None: + super().__init__(motion_generator, cfg or CoordinatedPickmentCfg()) + self._init_dual_arm_parts( + first_arm_control_part=self.cfg.left_arm_control_part, + second_arm_control_part=self.cfg.right_arm_control_part, + first_hand_control_part=self.cfg.left_hand_control_part, + second_hand_control_part=self.cfg.right_hand_control_part, + ) + self.left_arm_joint_ids = self.first_arm_joint_ids + self.right_arm_joint_ids = self.second_arm_joint_ids + self.left_hand_joint_ids = self.first_hand_joint_ids + self.right_hand_joint_ids = self.second_hand_joint_ids + self.left_arm_dof = self.first_arm_dof + self.right_arm_dof = self.second_arm_dof + self.left_hand_dof = self.first_hand_dof + self.right_hand_dof = self.second_hand_dof + + self._validate_hand_qpos_cfg() + self.left_hand_open_qpos = self._expand_qpos( + self.cfg.left_hand_open_qpos, self.left_hand_dof, "left_hand_open_qpos" + ) + self.left_hand_close_qpos = self._expand_qpos( + self.cfg.left_hand_close_qpos, self.left_hand_dof, "left_hand_close_qpos" + ) + self.right_hand_open_qpos = self._expand_qpos( + self.cfg.right_hand_open_qpos, self.right_hand_dof, "right_hand_open_qpos" + ) + self.right_hand_close_qpos = self._expand_qpos( + self.cfg.right_hand_close_qpos, + self.right_hand_dof, + "right_hand_close_qpos", + ) + + def _validate_hand_qpos_cfg(self) -> None: + for name in ( + "left_hand_open_qpos", + "left_hand_close_qpos", + "right_hand_open_qpos", + "right_hand_close_qpos", + ): + if getattr(self.cfg, name) is None: + logger.log_error( + f"{name} must be specified in CoordinatedPickmentCfg", + ValueError, + ) + + def _resolve_object_initial_pose( + self, target: CoordinatedPickmentTarget + ) -> torch.Tensor: + if target.object_initial_pose is not None: + return self._resolve_pose(target.object_initial_pose, "object_initial_pose") + if target.object_semantics.entity is None: + logger.log_error( + "CoordinatedPickmentTarget requires object_initial_pose when " + "object_semantics.entity is not provided.", + ValueError, + ) + return self._resolve_pose( + target.object_semantics.entity.get_local_pose(to_matrix=True), + "object_initial_pose", + ) + + def _resolve_target( + self, + target: CoordinatedPickmentTarget, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + CoordinatedHeldObjectState, + ]: + object_initial_pose = self._resolve_object_initial_pose(target) + object_target_pose = self._resolve_pose( + target.object_target_pose, "object_target_pose" + ) + left_object_to_eef = self._resolve_pose( + target.left_object_to_eef, "left_object_to_eef" + ) + right_object_to_eef = self._resolve_pose( + target.right_object_to_eef, "right_object_to_eef" + ) + + left_grasp_xpos = torch.bmm(object_initial_pose, left_object_to_eef) + right_grasp_xpos = torch.bmm(object_initial_pose, right_object_to_eef) + left_target_xpos = torch.bmm(object_target_pose, left_object_to_eef) + right_target_xpos = torch.bmm(object_target_pose, right_object_to_eef) + held_state = CoordinatedHeldObjectState( + semantics=target.object_semantics, + left_object_to_eef=left_object_to_eef, + right_object_to_eef=right_object_to_eef, + left_grasp_xpos=left_grasp_xpos, + right_grasp_xpos=right_grasp_xpos, + ) + return ( + object_initial_pose, + object_target_pose, + left_grasp_xpos, + right_grasp_xpos, + left_target_xpos, + right_target_xpos, + held_state, + ) + + def _compute_segment_lengths(self) -> dict[str, int]: + n_close = max(2, self.cfg.hand_interp_steps) + n_hold = max(0, self.cfg.hold_steps) + n_motion = self.cfg.sample_interval - n_close - n_hold + n_approach = n_motion // 3 + n_lift = n_motion // 3 + n_move = n_motion - n_approach - n_lift + if min(n_approach, n_lift, n_move) < 2: + logger.log_error( + "Not enough waypoints for coordinated pickment. Please increase " + "sample_interval or decrease hand_interp_steps/hold_steps.", + ValueError, + ) + return { + "approach": n_approach, + "close": n_close, + "lift": n_lift, + "move": n_move, + "hold": n_hold, + } + + def get_segment_lengths(self) -> dict[str, int]: + return self._compute_segment_lengths() + + def _compute_pre_grasp_xpos(self, grasp_xpos: torch.Tensor) -> torch.Tensor: + grasp_z = grasp_xpos[:, :3, 2] + return self.builder.apply_local_offset( + grasp_xpos, -grasp_z * self.cfg.pre_grasp_distance + ) + + def _select_motion_keyframe_indices(self, n_waypoints: int) -> torch.Tensor: + n_keyframes = min(max(2, int(self.cfg.object_motion_keyframes)), n_waypoints) + return ( + torch.linspace( + 0, + n_waypoints - 1, + steps=n_keyframes, + device=self.device, + ) + .round() + .to(dtype=torch.long) + ) + + def _plan_synchronized_object_motion( + self, + left_start_qpos: torch.Tensor, + right_start_qpos: torch.Tensor, + object_pose_traj: torch.Tensor, + left_object_to_eef: torch.Tensor, + right_object_to_eef: torch.Tensor, + ) -> tuple[bool, torch.Tensor, torch.Tensor]: + n_waypoints = object_pose_traj.shape[1] + keyframe_indices = self._select_motion_keyframe_indices(n_waypoints) + left_traj = torch.zeros( + (self.n_envs, len(keyframe_indices), left_start_qpos.shape[-1]), + dtype=torch.float32, + device=self.device, + ) + right_traj = torch.zeros( + (self.n_envs, len(keyframe_indices), right_start_qpos.shape[-1]), + dtype=torch.float32, + device=self.device, + ) + left_qpos_seed = left_start_qpos + right_qpos_seed = right_start_qpos + for keyframe_col, waypoint_idx in enumerate(keyframe_indices.tolist()): + left_xpos = torch.bmm(object_pose_traj[:, waypoint_idx], left_object_to_eef) + right_xpos = torch.bmm( + object_pose_traj[:, waypoint_idx], right_object_to_eef + ) + left_success, left_qpos = self.robot.compute_ik( + pose=left_xpos, + name=self.cfg.left_arm_control_part, + joint_seed=left_qpos_seed, + ) + right_success, right_qpos = self.robot.compute_ik( + pose=right_xpos, + name=self.cfg.right_arm_control_part, + joint_seed=right_qpos_seed, + ) + if not self.builder.all_envs_success(left_success): + logger.log_warning( + f"Failed to compute IK for {self.cfg.left_arm_control_part} " + f"object waypoint {waypoint_idx}." + ) + return False, left_traj, right_traj + if not self.builder.all_envs_success(right_success): + logger.log_warning( + f"Failed to compute IK for {self.cfg.right_arm_control_part} " + f"object waypoint {waypoint_idx}." + ) + return False, left_traj, right_traj + left_traj[:, keyframe_col] = left_qpos + right_traj[:, keyframe_col] = right_qpos + left_qpos_seed = left_qpos + right_qpos_seed = right_qpos + + return ( + True, + self._interpolate_qpos_keyframes(left_traj, keyframe_indices, n_waypoints), + self._interpolate_qpos_keyframes(right_traj, keyframe_indices, n_waypoints), + ) + + def execute( + self, target: CoordinatedPickmentTarget, state: WorldState + ) -> ActionResult: + ( + object_initial_pose, + object_target_pose, + left_grasp_xpos, + right_grasp_xpos, + left_target_xpos, + right_target_xpos, + held_state, + ) = self._resolve_target(target) + left_start_qpos, right_start_qpos = self._resolve_dual_arm_start(state) + segments = self._compute_segment_lengths() + + left_pre_grasp_xpos = self._compute_pre_grasp_xpos(left_grasp_xpos) + right_pre_grasp_xpos = self._compute_pre_grasp_xpos(right_grasp_xpos) + left_approach_targets = torch.stack( + [left_pre_grasp_xpos, left_grasp_xpos], dim=1 + ) + right_approach_targets = torch.stack( + [right_pre_grasp_xpos, right_grasp_xpos], dim=1 + ) + ok, left_approach_traj = self._plan_named_arm_trajectory( + self.cfg.left_arm_control_part, + left_start_qpos, + left_approach_targets, + segments["approach"], + ) + if not ok: + return self._fail(state) + ok, right_approach_traj = self._plan_named_arm_trajectory( + self.cfg.right_arm_control_part, + right_start_qpos, + right_approach_targets, + segments["approach"], + ) + if not ok: + return self._fail(state) + + left_grasp_qpos = left_approach_traj[:, -1] + right_grasp_qpos = right_approach_traj[:, -1] + approach_trajectory = self._assemble_phase( + state, + left_approach_traj, + right_approach_traj, + self._repeat_qpos(self.left_hand_open_qpos, segments["approach"]), + self._repeat_qpos(self.right_hand_open_qpos, segments["approach"]), + ) + + close_trajectory = self._assemble_phase( + state, + self._repeat_qpos(left_grasp_qpos, segments["close"]), + self._repeat_qpos(right_grasp_qpos, segments["close"]), + self._interpolate_qpos( + self.left_hand_open_qpos, + self.left_hand_close_qpos, + segments["close"], + ), + self._interpolate_qpos( + self.right_hand_open_qpos, + self.right_hand_close_qpos, + segments["close"], + ), + ) + + lift_object_pose = self.builder.apply_local_offset( + object_initial_pose, + torch.tensor([0.0, 0.0, self.cfg.lift_height], device=self.device), + ) + lift_object_traj = self._interpolate_object_pose( + object_initial_pose, + lift_object_pose, + segments["lift"], + include_orientation=False, + ) + ok, left_lift_traj, right_lift_traj = self._plan_synchronized_object_motion( + left_grasp_qpos, + right_grasp_qpos, + lift_object_traj, + held_state.left_object_to_eef, + held_state.right_object_to_eef, + ) + if not ok: + return self._fail(state) + + left_lift_qpos = left_lift_traj[:, -1] + right_lift_qpos = right_lift_traj[:, -1] + lift_trajectory = self._assemble_phase( + state, + left_lift_traj, + right_lift_traj, + self._repeat_qpos(self.left_hand_close_qpos, segments["lift"]), + self._repeat_qpos(self.right_hand_close_qpos, segments["lift"]), + ) + + move_object_traj = self._interpolate_object_pose( + lift_object_pose, + object_target_pose, + segments["move"], + include_orientation=True, + ) + ok, left_move_traj, right_move_traj = self._plan_synchronized_object_motion( + left_lift_qpos, + right_lift_qpos, + move_object_traj, + held_state.left_object_to_eef, + held_state.right_object_to_eef, + ) + if not ok: + return self._fail(state) + + left_target_qpos = left_move_traj[:, -1] + right_target_qpos = right_move_traj[:, -1] + move_trajectory = self._assemble_phase( + state, + left_move_traj, + right_move_traj, + self._repeat_qpos(self.left_hand_close_qpos, segments["move"]), + self._repeat_qpos(self.right_hand_close_qpos, segments["move"]), + ) + + hold_trajectory = torch.empty( + (self.n_envs, 0, self.robot_dof), dtype=torch.float32, device=self.device + ) + if segments["hold"] > 0: + hold_trajectory = self._assemble_phase( + state, + self._repeat_qpos(left_target_qpos, segments["hold"]), + self._repeat_qpos(right_target_qpos, segments["hold"]), + self._repeat_qpos(self.left_hand_close_qpos, segments["hold"]), + self._repeat_qpos(self.right_hand_close_qpos, segments["hold"]), + ) + + full = torch.cat( + [ + approach_trajectory, + close_trajectory, + lift_trajectory, + move_trajectory, + hold_trajectory, + ], + dim=1, + ) + coordinated_held_object = CoordinatedHeldObjectState( + semantics=held_state.semantics, + left_object_to_eef=held_state.left_object_to_eef, + right_object_to_eef=held_state.right_object_to_eef, + left_grasp_xpos=left_target_xpos, + right_grasp_xpos=right_target_xpos, + ) + return ActionResult( + success=True, + trajectory=full, + next_state=WorldState( + last_qpos=full[:, -1, :].clone(), + held_object=None, + coordinated_held_object=coordinated_held_object, + ), + ) + + +__all__ = ["CoordinatedPickment", "CoordinatedPickmentCfg"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py new file mode 100644 index 000000000..f8c2f88d7 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py @@ -0,0 +1,473 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""CoordinatedPlacement atomic action implementation.""" + +from __future__ import annotations + +from typing import ClassVar + +import torch + +from embodichain.lab.sim.planners import MoveType, PlanState +from embodichain.utils import configclass, logger + +from ._helpers import resolve_object_target +from ..core import ( + ActionCfg, + ActionResult, + AtomicAction, + CoordinatedPlacementTarget, + HeldObjectState, + WorldState, +) +from ..trajectory import TrajectoryBuilder + + +@configclass +class CoordinatedPlacementCfg(ActionCfg): + name: str = "coordinated_placement" + """Name of the action, used for identification and logging.""" + + control_part: str = "dual_arm" + """Robot control part containing both placing and support arms.""" + + placing_arm_control_part: str = "left_arm" + """Arm that places and releases its held object.""" + + support_arm_control_part: str = "right_arm" + """Arm that moves the support object and keeps holding it.""" + + placing_hand_control_part: str = "left_hand" + """Hand attached to the placing arm.""" + + support_hand_control_part: str = "right_hand" + """Hand attached to the support arm.""" + + placing_hand_open_qpos: torch.Tensor | None = None + """Placing-hand qpos for the open state, shape ``[hand_dof,]``.""" + + placing_hand_close_qpos: torch.Tensor | None = None + """Placing-hand qpos for the closed state, shape ``[hand_dof,]``.""" + + support_hand_close_qpos: torch.Tensor | None = None + """Support-hand qpos for the closed state, shape ``[hand_dof,]``.""" + + release: bool = True + """Whether to open the placing hand at the aligned placement pose.""" + + placing_height_offset: float = 0.0 + """Default World-Z offset above the placing object target pose.""" + + support_height_offset: float = 0.0 + """Default World-Z offset above the support object target pose.""" + + lift_height: float = 0.08 + """World-Z lift distance for the placing arm after release.""" + + sample_interval: int = 100 + """Number of waypoints for the full coordinated placement trajectory.""" + + hand_interp_steps: int = 10 + """Number of waypoints for the placing-hand release interpolation.""" + + hold_steps: int = 4 + """Number of waypoints to hold alignment before releasing.""" + + retreat_steps: int = 16 + """Number of waypoints used for the placing-arm lift retreat.""" + + +class CoordinatedPlacement(AtomicAction): + """Coordinate two held objects: support object below, placing object above.""" + + TargetType: ClassVar[type] = CoordinatedPlacementTarget + + def __init__( + self, + motion_generator, + cfg: CoordinatedPlacementCfg | None = None, + ) -> None: + super().__init__(motion_generator, cfg or CoordinatedPlacementCfg()) + self.builder = TrajectoryBuilder(motion_generator) + self.n_envs = self.robot.get_qpos().shape[0] + self.robot_dof = self.robot.dof + + self.dual_arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) + self.placing_arm_joint_ids = self.robot.get_joint_ids( + name=self.cfg.placing_arm_control_part + ) + self.support_arm_joint_ids = self.robot.get_joint_ids( + name=self.cfg.support_arm_control_part + ) + self.placing_hand_joint_ids = self.robot.get_joint_ids( + name=self.cfg.placing_hand_control_part + ) + self.support_hand_joint_ids = self.robot.get_joint_ids( + name=self.cfg.support_hand_control_part + ) + self.joint_ids = ( + self.dual_arm_joint_ids + + self.placing_hand_joint_ids + + self.support_hand_joint_ids + ) + self.placing_arm_dof = len(self.placing_arm_joint_ids) + self.support_arm_dof = len(self.support_arm_joint_ids) + self.placing_hand_dof = len(self.placing_hand_joint_ids) + self.support_hand_dof = len(self.support_hand_joint_ids) + + self._validate_hand_qpos_cfg() + self.placing_hand_open_qpos = self.builder.expand_hand_qpos( + self.cfg.placing_hand_open_qpos, + n_envs=self.n_envs, + hand_dof=self.placing_hand_dof, + ) + self.placing_hand_close_qpos = self.builder.expand_hand_qpos( + self.cfg.placing_hand_close_qpos, + n_envs=self.n_envs, + hand_dof=self.placing_hand_dof, + ) + self.support_hand_close_qpos = self.builder.expand_hand_qpos( + self.cfg.support_hand_close_qpos, + n_envs=self.n_envs, + hand_dof=self.support_hand_dof, + ) + + def execute( + self, target: CoordinatedPlacementTarget, state: WorldState + ) -> ActionResult: + placing_xpos, support_xpos, release, support_held_object = self._resolve_target( + target + ) + placing_start_qpos, support_start_qpos = self._resolve_start_qpos(state) + segments = self._compute_segment_lengths(release) + + placing_lift_xpos = self.builder.apply_local_offset( + placing_xpos, + torch.tensor( + [0.0, 0.0, self.cfg.lift_height], + dtype=torch.float32, + device=self.device, + ), + ) + + ok, placing_approach_traj = self._plan_named_arm_trajectory( + self.cfg.placing_arm_control_part, + placing_start_qpos, + torch.stack([placing_lift_xpos, placing_xpos], dim=1), + segments["approach"], + ) + if not ok: + logger.log_warning("CoordinatedPlacement failed to plan placing approach.") + return self._fail(state) + + ok, support_approach_traj = self._plan_named_arm_trajectory( + self.cfg.support_arm_control_part, + support_start_qpos, + support_xpos.unsqueeze(1), + segments["approach"], + ) + if not ok: + logger.log_warning("CoordinatedPlacement failed to plan support approach.") + return self._fail(state) + + placing_place_qpos = placing_approach_traj[:, -1] + support_place_qpos = support_approach_traj[:, -1] + approach_trajectory = self._assemble_phase( + state.last_qpos, + placing_approach_traj, + support_approach_traj, + self._repeat_qpos(self.placing_hand_close_qpos, segments["approach"]), + self._repeat_qpos(self.support_hand_close_qpos, segments["approach"]), + ) + + hold_trajectory = self._empty_phase() + if segments["hold"] > 0: + hold_trajectory = self._assemble_phase( + state.last_qpos, + self._repeat_qpos(placing_place_qpos, segments["hold"]), + self._repeat_qpos(support_place_qpos, segments["hold"]), + self._repeat_qpos(self.placing_hand_close_qpos, segments["hold"]), + self._repeat_qpos(self.support_hand_close_qpos, segments["hold"]), + ) + + release_trajectory = self._empty_phase() + if release: + release_trajectory = self._assemble_phase( + state.last_qpos, + self._repeat_qpos(placing_place_qpos, segments["release"]), + self._repeat_qpos(support_place_qpos, segments["release"]), + self.builder.interpolate_hand_qpos( + self.placing_hand_close_qpos, + self.placing_hand_open_qpos, + n_waypoints=segments["release"], + ), + self._repeat_qpos(self.support_hand_close_qpos, segments["release"]), + ) + + ok, placing_retreat_traj = self._plan_named_arm_trajectory( + self.cfg.placing_arm_control_part, + placing_place_qpos, + placing_lift_xpos.unsqueeze(1), + segments["retreat"], + ) + if not ok: + logger.log_warning("CoordinatedPlacement failed to plan placing retreat.") + return self._fail(state) + + placing_hand_retreat_qpos = ( + self.placing_hand_open_qpos if release else self.placing_hand_close_qpos + ) + retreat_trajectory = self._assemble_phase( + state.last_qpos, + placing_retreat_traj, + self._repeat_qpos(support_place_qpos, segments["retreat"]), + self._repeat_qpos(placing_hand_retreat_qpos, segments["retreat"]), + self._repeat_qpos(self.support_hand_close_qpos, segments["retreat"]), + ) + + full = torch.cat( + [ + approach_trajectory, + hold_trajectory, + release_trajectory, + retreat_trajectory, + ], + dim=1, + ) + return ActionResult( + success=True, + trajectory=full, + next_state=WorldState( + last_qpos=full[:, -1, :].clone(), + held_object=support_held_object, + ), + ) + + def _validate_hand_qpos_cfg(self) -> None: + required_names = ( + "placing_hand_open_qpos", + "placing_hand_close_qpos", + "support_hand_close_qpos", + ) + for name in required_names: + if getattr(self.cfg, name) is None: + logger.log_error( + f"{name} must be specified in CoordinatedPlacementCfg", + ValueError, + ) + + def _resolve_object_pose( + self, + pose: torch.Tensor, + height_offset: float, + name: str, + ) -> torch.Tensor: + object_pose = resolve_object_target( + pose, + n_envs=self.n_envs, + device=self.device, + name=name, + ) + return self.builder.apply_local_offset( + object_pose, + torch.tensor( + [0.0, 0.0, height_offset], + dtype=torch.float32, + device=self.device, + ), + ) + + def _resolve_object_to_eef( + self, + held_state: HeldObjectState, + name: str, + ) -> torch.Tensor: + return self._resolve_held_matrix( + held_state.object_to_eef, + f"{name}.object_to_eef", + ) + + def _resolve_held_matrix(self, matrix: torch.Tensor, name: str) -> torch.Tensor: + matrix = matrix.to(device=self.device, dtype=torch.float32) + if matrix.shape == (4, 4): + matrix = matrix.unsqueeze(0).repeat(self.n_envs, 1, 1) + if matrix.shape != (self.n_envs, 4, 4): + logger.log_error( + f"{name} must have shape (4, 4) or ({self.n_envs}, 4, 4), " + f"but got {matrix.shape}", + ValueError, + ) + return matrix + + def _resolve_held_state( + self, + held_state: HeldObjectState, + name: str, + object_to_eef: torch.Tensor, + ) -> HeldObjectState: + return HeldObjectState( + semantics=held_state.semantics, + object_to_eef=object_to_eef, + grasp_xpos=self._resolve_held_matrix( + held_state.grasp_xpos, + f"{name}.grasp_xpos", + ), + ) + + def _resolve_target( + self, + target: CoordinatedPlacementTarget, + ) -> tuple[torch.Tensor, torch.Tensor, bool, HeldObjectState]: + placing_height_offset = ( + self.cfg.placing_height_offset + if target.placing_height_offset is None + else target.placing_height_offset + ) + support_height_offset = ( + self.cfg.support_height_offset + if target.support_height_offset is None + else target.support_height_offset + ) + placing_object_pose = self._resolve_object_pose( + target.placing_object_target_pose, + placing_height_offset, + "placing_object_target_pose", + ) + support_object_pose = self._resolve_object_pose( + target.support_object_target_pose, + support_height_offset, + "support_object_target_pose", + ) + placing_object_to_eef = self._resolve_object_to_eef( + target.placing_held_object, + "placing_held_object", + ) + support_object_to_eef = self._resolve_object_to_eef( + target.support_held_object, + "support_held_object", + ) + placing_xpos = torch.bmm(placing_object_pose, placing_object_to_eef) + support_xpos = torch.bmm(support_object_pose, support_object_to_eef) + release = self.cfg.release if target.release is None else target.release + return ( + placing_xpos, + support_xpos, + release, + self._resolve_held_state( + target.support_held_object, + "support_held_object", + support_object_to_eef, + ), + ) + + def _resolve_start_qpos( + self, state: WorldState + ) -> tuple[torch.Tensor, torch.Tensor]: + if state.last_qpos.shape != (self.n_envs, self.robot_dof): + logger.log_error( + f"WorldState.last_qpos must have shape ({self.n_envs}, {self.robot_dof}), " + f"but got {state.last_qpos.shape}", + ValueError, + ) + start_qpos = state.last_qpos.to(device=self.device, dtype=torch.float32) + return ( + start_qpos[:, self.placing_arm_joint_ids], + start_qpos[:, self.support_arm_joint_ids], + ) + + def _compute_segment_lengths(self, release: bool) -> dict[str, int]: + n_release = max(2, self.cfg.hand_interp_steps) if release else 0 + n_hold = max(0, self.cfg.hold_steps) + n_retreat = max(2, self.cfg.retreat_steps) + n_approach = self.cfg.sample_interval - n_hold - n_release - n_retreat + if n_approach < 2: + logger.log_error( + "Not enough waypoints for coordinated placement. Increase " + "sample_interval or decrease hold/release/retreat steps.", + ValueError, + ) + return { + "approach": n_approach, + "hold": n_hold, + "release": n_release, + "retreat": n_retreat, + } + + def _plan_named_arm_trajectory( + self, + control_part: str, + start_qpos: torch.Tensor, + target_poses: torch.Tensor, + n_waypoints: int, + ) -> tuple[bool, torch.Tensor]: + target_states_list = [ + [ + PlanState(xpos=target_poses[i, j], move_type=MoveType.EEF_MOVE) + for j in range(target_poses.shape[1]) + ] + for i in range(self.n_envs) + ] + success, trajectory = self.builder.plan_arm_traj( + target_states_list, + start_qpos, + n_waypoints, + control_part=control_part, + arm_dof=start_qpos.shape[-1], + ) + return self.builder.all_envs_success(success), trajectory + + @staticmethod + def _repeat_qpos(qpos: torch.Tensor, n_waypoints: int) -> torch.Tensor: + return qpos.unsqueeze(1).repeat(1, n_waypoints, 1) + + def _empty_phase(self) -> torch.Tensor: + return torch.empty( + (self.n_envs, 0, self.robot_dof), + dtype=torch.float32, + device=self.device, + ) + + def _assemble_phase( + self, + base_full_qpos: torch.Tensor, + placing_arm_traj: torch.Tensor, + support_arm_traj: torch.Tensor, + placing_hand_traj: torch.Tensor, + support_hand_traj: torch.Tensor, + ) -> torch.Tensor: + n_waypoints = placing_arm_traj.shape[1] + full = base_full_qpos.to(device=self.device, dtype=torch.float32) + full = full.unsqueeze(1).repeat(1, n_waypoints, 1).clone() + full[:, :, self.placing_arm_joint_ids] = placing_arm_traj + full[:, :, self.support_arm_joint_ids] = support_arm_traj + full[:, :, self.placing_hand_joint_ids] = placing_hand_traj + full[:, :, self.support_hand_joint_ids] = support_hand_traj + return full + + def _fail(self, state: WorldState) -> ActionResult: + return ActionResult( + success=False, + trajectory=torch.empty( + (self.n_envs, 0, self.robot_dof), + dtype=torch.float32, + device=self.device, + ), + next_state=state, + ) + + +__all__ = ["CoordinatedPlacement", "CoordinatedPlacementCfg"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py b/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py new file mode 100644 index 000000000..d42d7b13e --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py @@ -0,0 +1,139 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""MoveEndEffector atomic action implementation.""" + +from __future__ import annotations + +from typing import ClassVar + +import torch + +from embodichain.lab.sim.planners import MoveType, PlanState +from embodichain.utils import configclass + +from ._helpers import arm_qpos_from_state +from ..core import ( + ActionCfg, + ActionResult, + AtomicAction, + EndEffectorPoseTarget, + WorldState, +) +from ..trajectory import TrajectoryBuilder + + +@configclass +class MoveEndEffectorCfg(ActionCfg): + name: str = "move_end_effector" + """Name of the action, used for identification and logging.""" + + sample_interval: int = 50 + """Number of waypoints in the planned trajectory.""" + + +class MoveEndEffector(AtomicAction): + """Plan a free-space end-effector move to a target pose. + + The :class:`EndEffectorPoseTarget` may carry either a single waypoint + ``(n_envs, 4, 4)`` (or a broadcastable ``(4, 4)``) or a multi-waypoint + trajectory ``(n_envs, n_waypoint, 4, 4)``. In the multi-waypoint case the + action plans a single trajectory that visits every waypoint in order, + starting from the inherited ``WorldState.last_qpos``; IK is solved for each + waypoint with the previous waypoint's solution as the seed. + """ + + TargetType: ClassVar[type] = EndEffectorPoseTarget + + def __init__( + self, + motion_generator, + cfg: MoveEndEffectorCfg | None = None, + ) -> None: + super().__init__(motion_generator, cfg or MoveEndEffectorCfg()) + self.builder = TrajectoryBuilder(motion_generator) + self.n_envs = self.robot.get_qpos().shape[0] + self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) + self.arm_dof = len(self.arm_joint_ids) + self.robot_dof = self.robot.dof + + def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionResult: + move_xpos = self.builder.resolve_pose_target(target.xpos, n_envs=self.n_envs) + start_qpos = self.builder.resolve_start_qpos( + arm_qpos_from_state(state, self.arm_joint_ids), + n_envs=self.n_envs, + arm_dof=self.arm_dof, + control_part=self.cfg.control_part, + ) + target_states_list = self._build_target_states(move_xpos) + success, arm_traj = self.builder.plan_arm_traj( + target_states_list, + start_qpos, + self.cfg.sample_interval, + control_part=self.cfg.control_part, + arm_dof=self.arm_dof, + cfg=self.cfg, + ) + full = self._embed(arm_traj, state.last_qpos) + return ActionResult( + success=success, + trajectory=full, + next_state=WorldState( + last_qpos=full[:, -1, :].clone(), + held_object=state.held_object, + coordinated_held_object=state.coordinated_held_object, + ), + ) + + def _build_target_states(self, move_xpos: torch.Tensor) -> list[list[PlanState]]: + """Build per-env PlanState lists from a single- or multi-waypoint target.""" + if move_xpos.dim() == 3: + move_xpos = move_xpos.unsqueeze(1) + n_waypoint = move_xpos.shape[1] + return [ + [ + PlanState(xpos=move_xpos[i, j], move_type=MoveType.EEF_MOVE) + for j in range(n_waypoint) + ] + for i in range(self.n_envs) + ] + + def _embed( + self, arm_traj: torch.Tensor, last_full_qpos: torch.Tensor + ) -> torch.Tensor: + n_wp = arm_traj.shape[1] + full = torch.empty( + (self.n_envs, n_wp, self.robot_dof), + dtype=torch.float32, + device=self.device, + ) + full[:, :, :] = last_full_qpos.unsqueeze(1) + full[:, :, self.arm_joint_ids] = arm_traj + return full + + def _fail(self, state: WorldState) -> ActionResult: + return ActionResult( + success=torch.zeros(self.n_envs, dtype=torch.bool, device=self.device), + trajectory=torch.empty( + (self.n_envs, 0, self.robot_dof), + dtype=torch.float32, + device=self.device, + ), + next_state=state, + ) + + +__all__ = ["MoveEndEffector", "MoveEndEffectorCfg"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py b/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py new file mode 100644 index 000000000..343414e3c --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py @@ -0,0 +1,178 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""MoveHeldObject atomic action implementation.""" + +from __future__ import annotations + +from typing import ClassVar + +import torch + +from embodichain.lab.sim.planners import MoveType, PlanState +from embodichain.utils import configclass, logger + +from ._helpers import arm_qpos_from_state, resolve_object_target +from ..core import ( + ActionCfg, + ActionResult, + AtomicAction, + HeldObjectPoseTarget, + WorldState, +) +from embodichain.utils.math import ( + axis_angle_to_rotation_matrix, +) +from ..trajectory import TrajectoryBuilder + + +@configclass +class MoveHeldObjectCfg(ActionCfg): + name: str = "move_held_object" + """Name of the action, used for identification and logging.""" + + sample_interval: int = 50 + """Number of waypoints in the planned trajectory.""" + + hand_control_part: str = "hand" + """Name of the robot part that controls the hand joints.""" + + hand_close_qpos: torch.Tensor | None = None + """Joint positions for the closed hand state, shape ``[hand_dof,]``.""" + + obj_upright_direction: torch.Tensor | None = None + """Optional object local direction to align with world up after grasping. By dafault we will use (0, 0, 1).""" + + pick_rotate_upright: float | None = None + """Optional rotation (radians) about the grasp y-axis to apply to the grasp pose""" + + +class MoveHeldObject(AtomicAction): + """Move the held object to a target object pose; keep the gripper closed.""" + + TargetType: ClassVar[type] = HeldObjectPoseTarget + + def __init__( + self, + motion_generator, + cfg: MoveHeldObjectCfg | None = None, + ) -> None: + super().__init__(motion_generator, cfg or MoveHeldObjectCfg()) + self.builder = TrajectoryBuilder(motion_generator) + self.n_envs = self.robot.get_qpos().shape[0] + self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) + self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part) + self.arm_dof = len(self.arm_joint_ids) + self.robot_dof = self.robot.dof + + if self.cfg.hand_close_qpos is None: + logger.log_error( + "hand_close_qpos must be specified in MoveHeldObjectCfg", ValueError + ) + self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) + + def execute(self, target: HeldObjectPoseTarget, state: WorldState) -> ActionResult: + if state.held_object is None: + logger.log_error( + "MoveHeldObject requires WorldState.held_object - run PickUp first.", + ValueError, + ) + object_target_pose = resolve_object_target( + target.object_target_pose, n_envs=self.n_envs, device=self.device + ) + start_arm_qpos = self.builder.resolve_start_qpos( + arm_qpos_from_state(state, self.arm_joint_ids), + n_envs=self.n_envs, + arm_dof=self.arm_dof, + control_part=self.cfg.control_part, + ) + if self.cfg.pick_rotate_upright is not None: + held_eef_xpos = self.robot.compute_fk( + qpos=start_arm_qpos, name=self.cfg.control_part, to_matrix=True + ) + held_obj_xpos = state.held_object.semantics.entity.get_local_pose( + to_matrix=True + ) + if self.cfg.obj_upright_direction is None: + upright_direction = torch.tensor([0, 0, 1], device=self.device) + else: + upright_direction = self.cfg.obj_upright_direction.to(self.device) + obj_upright = (upright_direction * held_obj_xpos[:, :3, :3]).sum(axis=2) + + grasp_ry = held_eef_xpos[:, :3, 1] + dot_result = (grasp_ry * obj_upright).sum(axis=1) + # revert flag is -1 if the dot product is negative, 1 if positive + revert_flag = torch.where(dot_result < 0, 1.0, -1.0) + grasp_rx = held_eef_xpos[:, :3, 0] + # rotate util upright + rota_axis_angle = -0.5 * torch.pi * revert_flag * grasp_rx + gripper_rotate_offset = axis_angle_to_rotation_matrix(rota_axis_angle) + # modified target xpos rotation + object_target_pose[:, :3, :3] = torch.bmm( + gripper_rotate_offset, held_obj_xpos[:, :3, :3] + ) + object_to_eef = state.held_object.object_to_eef.to( + device=self.device, dtype=torch.float32 + ) + if object_to_eef.shape == (4, 4): + object_to_eef = object_to_eef.unsqueeze(0).repeat(self.n_envs, 1, 1) + move_eef_xpos = torch.bmm(object_target_pose, object_to_eef) + + target_states_list = [ + [PlanState(xpos=move_eef_xpos[i], move_type=MoveType.EEF_MOVE)] + for i in range(self.n_envs) + ] + success, arm_traj = self.builder.plan_arm_traj( + target_states_list, + start_arm_qpos, + self.cfg.sample_interval, + control_part=self.cfg.control_part, + arm_dof=self.arm_dof, + cfg=self.cfg, + ) + + full = torch.empty( + (self.n_envs, arm_traj.shape[1], self.robot_dof), + dtype=torch.float32, + device=self.device, + ) + full[:, :, :] = state.last_qpos.unsqueeze(1) + full[:, :, self.arm_joint_ids] = arm_traj + full[:, :, self.hand_joint_ids] = self.hand_close_qpos + + return ActionResult( + success=success, + trajectory=full, + next_state=WorldState( + last_qpos=full[:, -1, :].clone(), + held_object=state.held_object, + coordinated_held_object=state.coordinated_held_object, + ), + ) + + def _fail(self, state: WorldState) -> ActionResult: + return ActionResult( + success=torch.zeros(self.n_envs, dtype=torch.bool, device=self.device), + trajectory=torch.empty( + (self.n_envs, 0, self.robot_dof), + dtype=torch.float32, + device=self.device, + ), + next_state=state, + ) + + +__all__ = ["MoveHeldObject", "MoveHeldObjectCfg"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py new file mode 100644 index 000000000..bf2b29681 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py @@ -0,0 +1,136 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""MoveJoints atomic action implementation.""" + +from __future__ import annotations + +from typing import ClassVar + +import torch + +from embodichain.utils import configclass, logger + +from ..core import ( + ActionCfg, + ActionResult, + AtomicAction, + JointPositionTarget, + NamedJointPositionTarget, + WorldState, +) +from ..trajectory import TrajectoryBuilder + + +@configclass +class MoveJointsCfg(ActionCfg): + name: str = "move_joints" + """Name of the action, used for identification and logging.""" + + sample_interval: int = 50 + """Number of waypoints in the interpolated joint-space trajectory.""" + + named_joint_positions: dict[str, torch.Tensor] | None = None + """Optional named joint targets resolved by ``NamedJointPositionTarget``.""" + + +class MoveJoints(AtomicAction): + """Plan a joint-space move for the configured control part. + + The :class:`JointPositionTarget` may carry either a single waypoint + ``(n_envs, control_dof)`` or a multi-waypoint trajectory + ``(n_envs, n_waypoint, control_dof)``. In the multi-waypoint case the + action plans a single trajectory that visits every waypoint in order, + starting from the inherited ``WorldState.last_qpos``. + """ + + TargetType: ClassVar[tuple[type, ...]] = ( + JointPositionTarget, + NamedJointPositionTarget, + ) + + def __init__( + self, + motion_generator, + cfg: MoveJointsCfg | None = None, + ) -> None: + super().__init__(motion_generator, cfg or MoveJointsCfg()) + self.builder = TrajectoryBuilder(motion_generator) + self.n_envs = self.robot.get_qpos().shape[0] + self.joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) + self.joint_dof = len(self.joint_ids) + self.robot_dof = self.robot.dof + self.named_joint_positions = self.cfg.named_joint_positions or {} + + def execute( + self, + target: JointPositionTarget | NamedJointPositionTarget, + state: WorldState, + ) -> ActionResult: + target_qpos = self.builder.resolve_joint_target( + self._resolve_target_qpos(target), + n_envs=self.n_envs, + joint_dof=self.joint_dof, + control_part=self.cfg.control_part, + ) + start_qpos = self.builder.resolve_start_qpos( + state.last_qpos[:, self.joint_ids], + n_envs=self.n_envs, + arm_dof=self.joint_dof, + control_part=self.cfg.control_part, + ) + joint_traj = self.builder.plan_joint_traj( + start_qpos, target_qpos, self.cfg.sample_interval + ) + full = self._embed(joint_traj, state.last_qpos) + return ActionResult( + success=torch.ones(self.n_envs, dtype=torch.bool, device=self.device), + trajectory=full, + next_state=WorldState( + last_qpos=full[:, -1, :].clone(), + held_object=state.held_object, + coordinated_held_object=state.coordinated_held_object, + ), + ) + + def _resolve_target_qpos( + self, target: JointPositionTarget | NamedJointPositionTarget + ) -> torch.Tensor: + if isinstance(target, JointPositionTarget): + return target.qpos + if target.name not in self.named_joint_positions: + logger.log_error( + f"Unknown named joint-position target '{target.name}' for " + f"MoveJoints. Available targets: {sorted(self.named_joint_positions)}", + KeyError, + ) + return self.named_joint_positions[target.name] + + def _embed( + self, joint_traj: torch.Tensor, last_full_qpos: torch.Tensor + ) -> torch.Tensor: + n_wp = joint_traj.shape[1] + full = torch.empty( + (self.n_envs, n_wp, self.robot_dof), + dtype=torch.float32, + device=self.device, + ) + full[:, :, :] = last_full_qpos.unsqueeze(1) + full[:, :, self.joint_ids] = joint_traj + return full + + +__all__ = ["MoveJoints", "MoveJointsCfg"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py new file mode 100644 index 000000000..bf290842a --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -0,0 +1,334 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""PickUp atomic action implementation.""" + +from __future__ import annotations + +from typing import ClassVar + +import torch + +from embodichain.lab.sim.planners import MoveType, PlanState +from embodichain.utils import configclass, logger +from embodichain.utils.math import ( + pose_inv, + quat_error_magnitude, + quat_from_matrix, + axis_angle_to_rotation_matrix, +) + +from ._helpers import arm_qpos_from_state +from ..affordance import AntipodalAffordance +from ..core import ( + ActionCfg, + ActionResult, + AtomicAction, + GraspTarget, + HeldObjectState, + ObjectSemantics, + WorldState, +) +from ..trajectory import TrajectoryBuilder + + +@configclass +class PickUpCfg(ActionCfg): + name: str = "pick_up" + """Name of the action, used for identification and logging.""" + + sample_interval: int = 80 + """Number of waypoints for the full trajectory (approach + hand + lift).""" + + hand_interp_steps: int = 5 + """Number of waypoints for the gripper close interpolation phase.""" + + hand_control_part: str = "hand" + """Name of the robot part that controls the hand joints.""" + + hand_open_qpos: torch.Tensor | None = None + """Joint positions for the open hand state, shape ``[hand_dof,]``.""" + + hand_close_qpos: torch.Tensor | None = None + """Joint positions for the closed hand state, shape ``[hand_dof,]``.""" + + lift_height: float = 0.1 + """Height (m) to lift the end-effector after closing the gripper.""" + + pre_grasp_distance: float = 0.15 + """Distance to offset back from the grasp pose along the approach direction.""" + + approach_direction: torch.Tensor = torch.tensor([0, 0, -1], dtype=torch.float32) + """Approach direction in the object local frame.""" + + obj_upright_direction: torch.Tensor | None = None + """Optional object local direction to align with world up after grasping. By dafault we will use (0, 0, 1).""" + + rotate_upright: float | None = None + """Optional rotation (radians) about the grasp y-axis to apply to the grasp pose""" + + +class PickUp(AtomicAction): + """Approach a grasp pose, close the gripper, lift.""" + + TargetType: ClassVar[type] = GraspTarget + + def __init__( + self, + motion_generator, + cfg: PickUpCfg | None = None, + ) -> None: + super().__init__(motion_generator, cfg or PickUpCfg()) + self.builder = TrajectoryBuilder(motion_generator) + self.n_envs = self.robot.get_qpos().shape[0] + self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) + self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part) + self.arm_dof = len(self.arm_joint_ids) + self.robot_dof = self.robot.dof + + if self.cfg.hand_open_qpos is None: + logger.log_error( + "hand_open_qpos must be specified in PickUpCfg", ValueError + ) + if self.cfg.hand_close_qpos is None: + logger.log_error( + "hand_close_qpos must be specified in PickUpCfg", ValueError + ) + self.hand_open_qpos = self.cfg.hand_open_qpos.to(self.device) + self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) + self.approach_direction = self.cfg.approach_direction.to(self.device) + + def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: + sem = target.semantics + if target.grasp_xpos is None and not isinstance( + sem.affordance, AntipodalAffordance + ): + logger.log_error( + "PickUp requires an AntipodalAffordance when grasp_xpos is not set.", + ValueError, + ) + if sem.entity is None: + logger.log_error( + "PickUp requires an entity on the target semantics.", ValueError + ) + start_arm_qpos = self.builder.resolve_start_qpos( + arm_qpos_from_state(state, self.arm_joint_ids), + n_envs=self.n_envs, + arm_dof=self.arm_dof, + control_part=self.cfg.control_part, + ) + if target.grasp_xpos is None: + is_success, grasp_xpos = self._resolve_grasp_pose(sem, start_arm_qpos) + else: + grasp_xpos = self.builder.resolve_pose_target( + target.grasp_xpos, n_envs=self.n_envs + ) + is_success = torch.ones(self.n_envs, dtype=torch.bool, device=self.device) + + # apply grasp yr rotation offset if specified + if self.cfg.rotate_upright is not None: + if self.cfg.obj_upright_direction is None: + upright_direction = torch.tensor([0, 0, 1], device=self.device) + else: + upright_direction = self.cfg.obj_upright_direction.to(self.device) + obj_pose = sem.entity.get_local_pose(to_matrix=True) + obj_upright = (upright_direction * obj_pose[:, :3, :3]).sum(axis=2) + grasp_ry = grasp_xpos[:, :3, 1] + dot_result = (grasp_ry * obj_upright).sum(axis=1) + # revert flag is -1 if the dot product is negative, 1 if positive + revert_flag = torch.where(dot_result < 0, 1.0, -1.0) + grasp_rx = grasp_xpos[:, :3, 0] + rota_axis_angle = self.cfg.rotate_upright * revert_flag * grasp_rx + rota_offset = axis_angle_to_rotation_matrix(rota_axis_angle) + upright_grasp_rota = torch.bmm(rota_offset, grasp_xpos[:, :3, :3]) + grasp_xpos[:, :3, :3] = upright_grasp_rota + + if not self.builder.all_envs_success(is_success): + logger.log_warning("PickUp failed to resolve a grasp pose.") + return self._fail(state) + + grasp_z = grasp_xpos[:, :3, 2] + pre_grasp_xpos = self.builder.apply_local_offset( + grasp_xpos, -grasp_z * self.cfg.pre_grasp_distance + ) + + n_approach, n_close, n_lift = self.builder.split_three_phase( + self.cfg.sample_interval, + self.cfg.hand_interp_steps, + first_phase_name="approach", + third_phase_name="lift", + ) + + target_states_list = [ + [ + PlanState(xpos=pre_grasp_xpos[i], move_type=MoveType.EEF_MOVE), + PlanState(xpos=grasp_xpos[i], move_type=MoveType.EEF_MOVE), + ] + for i in range(self.n_envs) + ] + approach_success, approach_arm = self.builder.plan_arm_traj( + target_states_list, + start_arm_qpos, + n_approach, + control_part=self.cfg.control_part, + arm_dof=self.arm_dof, + cfg=self.cfg, + ) + + grasp_arm_qpos = approach_arm[:, -1, :] + lift_xpos = self.builder.apply_local_offset( + grasp_xpos, + torch.tensor([0, 0, 1], device=self.device) * self.cfg.lift_height, + ) + target_states_list = [ + [PlanState(xpos=lift_xpos[i], move_type=MoveType.EEF_MOVE)] + for i in range(self.n_envs) + ] + lift_success, lift_arm = self.builder.plan_arm_traj( + target_states_list, + grasp_arm_qpos, + n_lift, + control_part=self.cfg.control_part, + arm_dof=self.arm_dof, + cfg=self.cfg, + ) + success = approach_success & lift_success + + hand_close_path = self.builder.interpolate_hand_qpos( + self.hand_open_qpos, self.hand_close_qpos, n_waypoints=n_close + ) + + full = torch.empty( + (self.n_envs, n_approach + n_close + n_lift, self.robot_dof), + dtype=torch.float32, + device=self.device, + ) + full[:, :, :] = state.last_qpos.unsqueeze(1) + full[:, :n_approach, self.arm_joint_ids] = approach_arm + full[:, :n_approach, self.hand_joint_ids] = self.hand_open_qpos + full[:, n_approach : n_approach + n_close, self.arm_joint_ids] = ( + grasp_arm_qpos.unsqueeze(1) + ) + full[:, n_approach : n_approach + n_close, self.hand_joint_ids] = ( + hand_close_path + ) + full[:, n_approach + n_close :, self.arm_joint_ids] = lift_arm + full[:, n_approach + n_close :, self.hand_joint_ids] = self.hand_close_qpos + + obj_poses = sem.entity.get_local_pose(to_matrix=True) + object_to_eef = torch.bmm(pose_inv(obj_poses), grasp_xpos) + held = HeldObjectState( + semantics=sem, object_to_eef=object_to_eef, grasp_xpos=grasp_xpos + ) + return ActionResult( + success=success, + trajectory=full, + next_state=WorldState( + last_qpos=full[:, -1, :].clone(), + held_object=held, + coordinated_held_object=state.coordinated_held_object, + ), + ) + + def _fail(self, state: WorldState) -> ActionResult: + return ActionResult( + success=torch.zeros(self.n_envs, dtype=torch.bool, device=self.device), + trajectory=torch.empty( + (self.n_envs, 0, self.robot_dof), + dtype=torch.float32, + device=self.device, + ), + next_state=state, + ) + + def _resolve_grasp_pose( + self, semantics: ObjectSemantics, start_qpos: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + obj_poses = semantics.entity.get_local_pose(to_matrix=True) + grasp_poses_result = semantics.affordance.get_valid_grasp_poses( + obj_poses=obj_poses, approach_direction=self.approach_direction + ) + n_envs = obj_poses.shape[0] + n_max_pose = max(r[0].shape[0] for r in grasp_poses_result) + grasp_xpos_padding = torch.zeros( + (n_envs, n_max_pose, 4, 4), dtype=torch.float32, device=self.device + ) + grasp_cost_padding = torch.full( + (n_envs, n_max_pose), + float("inf"), + dtype=torch.float32, + device=self.device, + ) + for i in range(n_envs): + n_pose = grasp_poses_result[i][0].shape[0] + grasp_poses = grasp_poses_result[i][0].to( + device=self.device, dtype=torch.float32 + ) + grasp_costs = grasp_poses_result[i][1].to( + device=self.device, dtype=torch.float32 + ) + grasp_xpos_padding[i, :n_pose] = grasp_poses + grasp_cost_padding[i, :n_pose] = grasp_costs + grasp_xpos_padding[i, n_pose:] = grasp_poses[0] + grasp_cost_padding[i, n_pose:] = grasp_costs[0] + grasp_xpos_padding, ik_success = self._select_symmetric_grasp_variants( + grasp_xpos_padding, start_qpos + ) + grasp_cost_masked = torch.where(ik_success, grasp_cost_padding, 10000.0) + best_cost, best_idx = grasp_cost_masked.min(dim=1) + is_success = best_cost < 9999.0 + best_grasp_xpos = grasp_xpos_padding[ + torch.arange(n_envs, device=self.device), best_idx + ] + return is_success, best_grasp_xpos + + def _select_symmetric_grasp_variants( + self, grasp_xpos: torch.Tensor, start_qpos: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Choose the closest TCP z-roll variant, then validate reachability.""" + n_envs, n_pose = grasp_xpos.shape[:2] + mirrored_grasp_xpos = grasp_xpos.clone() + mirrored_grasp_xpos[..., :3, 0] = -mirrored_grasp_xpos[..., :3, 0] + mirrored_grasp_xpos[..., :3, 1] = -mirrored_grasp_xpos[..., :3, 1] + grasp_variants = torch.stack([grasp_xpos, mirrored_grasp_xpos], dim=2) + + start_xpos = self.robot.compute_fk( + qpos=start_qpos, + name=self.cfg.control_part, + to_matrix=True, + ) + start_quat = quat_from_matrix(start_xpos[:, :3, :3]) + variant_quat = quat_from_matrix(grasp_variants[..., :3, :3]) + start_quat = start_quat[:, None, None, :].expand_as(variant_quat) + rotation_error = quat_error_magnitude( + variant_quat.reshape(-1, 4), + start_quat.reshape(-1, 4), + ).reshape(n_envs, n_pose, 2) + best_variant_idx = rotation_error.argmin(dim=2) + + env_idx = torch.arange(n_envs, device=self.device)[:, None] + pose_idx = torch.arange(n_pose, device=self.device)[None, :] + selected_grasp_xpos = grasp_variants[env_idx, pose_idx, best_variant_idx] + start_qpos_repeat = start_qpos[:, None, :].repeat(1, n_pose, 1) + ik_success, _ = self.robot.compute_batch_ik( + pose=selected_grasp_xpos, + name=self.cfg.control_part, + joint_seed=start_qpos_repeat, + ) + return selected_grasp_xpos, ik_success + + +__all__ = ["PickUp", "PickUpCfg"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/place.py b/embodichain/lab/sim/atomic_actions/primitives/place.py new file mode 100644 index 000000000..942cdf3b5 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/place.py @@ -0,0 +1,231 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Place atomic action implementation.""" + +from __future__ import annotations + +from typing import ClassVar + +import torch + +from embodichain.lab.sim.planners import MoveType, PlanState +from embodichain.utils import configclass, logger +from embodichain.utils.math import quat_error_magnitude, quat_from_matrix + +from ._helpers import arm_qpos_from_state +from ..core import ( + ActionCfg, + ActionResult, + AtomicAction, + EndEffectorPoseTarget, + WorldState, +) +from ..trajectory import TrajectoryBuilder + + +@configclass +class PlaceCfg(ActionCfg): + name: str = "place" + """Name of the action, used for identification and logging.""" + + sample_interval: int = 80 + """Number of waypoints for the full trajectory (down + hand + back).""" + + hand_interp_steps: int = 5 + """Number of waypoints for the gripper open interpolation phase.""" + + hand_control_part: str = "hand" + """Name of the robot part that controls the hand joints.""" + + hand_open_qpos: torch.Tensor | None = None + """Joint positions for the open hand state, shape ``[hand_dof,]``.""" + + hand_close_qpos: torch.Tensor | None = None + """Joint positions for the closed hand state, shape ``[hand_dof,]``.""" + + lift_height: float = 0.1 + """Height (m) to retract the end-effector after opening the gripper.""" + + +class Place(AtomicAction): + """Lower the held object to a place pose, open the gripper, retract. + + The :class:`EndEffectorPoseTarget` may carry either a single waypoint + ``(n_envs, 4, 4)`` (or a broadcastable ``(4, 4)``) or a multi-waypoint + trajectory ``(n_envs, n_waypoint, 4, 4)``. In the multi-waypoint case the + down phase visits every waypoint in order; approaching from above the + first waypoint, descending through each waypoint, then opening the gripper + at the final waypoint and retracting to above the last waypoint. Starting + joint positions are inherited from ``WorldState.last_qpos``. + """ + + TargetType: ClassVar[type] = EndEffectorPoseTarget + + def __init__( + self, + motion_generator, + cfg: PlaceCfg | None = None, + ) -> None: + super().__init__(motion_generator, cfg or PlaceCfg()) + self.builder = TrajectoryBuilder(motion_generator) + self.n_envs = self.robot.get_qpos().shape[0] + self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) + self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part) + self.arm_dof = len(self.arm_joint_ids) + self.robot_dof = self.robot.dof + + if self.cfg.hand_open_qpos is None: + logger.log_error("hand_open_qpos must be specified in PlaceCfg", ValueError) + if self.cfg.hand_close_qpos is None: + logger.log_error( + "hand_close_qpos must be specified in PlaceCfg", ValueError + ) + self.hand_open_qpos = self.cfg.hand_open_qpos.to(self.device) + self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) + + def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionResult: + place_xpos = self.builder.resolve_pose_target(target.xpos, n_envs=self.n_envs) + if place_xpos.dim() == 3: + place_xpos = place_xpos.unsqueeze(1) + + start_arm_qpos = self.builder.resolve_start_qpos( + arm_qpos_from_state(state, self.arm_joint_ids), + n_envs=self.n_envs, + arm_dof=self.arm_dof, + control_part=self.cfg.control_part, + ) + if target.tcp_symmetry == "z_roll_180": + place_xpos = self._select_tcp_symmetric_place_variant( + place_xpos, start_arm_qpos + ) + n_waypoint = place_xpos.shape[1] + n_down, n_open, n_back = self.builder.split_three_phase( + self.cfg.sample_interval, + self.cfg.hand_interp_steps, + first_phase_name="approach", + third_phase_name="back", + ) + + lift_offset = torch.tensor([0, 0, 1], device=self.device) * self.cfg.lift_height + approach_xpos = self.builder.apply_local_offset(place_xpos[:, 0], lift_offset) + retract_xpos = self.builder.apply_local_offset(place_xpos[:, -1], lift_offset) + + target_states_list = [ + [PlanState(xpos=approach_xpos[i], move_type=MoveType.EEF_MOVE)] + + [ + PlanState(xpos=place_xpos[i, j], move_type=MoveType.EEF_MOVE) + for j in range(n_waypoint) + ] + for i in range(self.n_envs) + ] + down_success, down_arm = self.builder.plan_arm_traj( + target_states_list, + start_arm_qpos, + n_down, + control_part=self.cfg.control_part, + arm_dof=self.arm_dof, + cfg=self.cfg, + ) + reach_arm_qpos = down_arm[:, -1, :] + + target_states_list = [ + [PlanState(xpos=retract_xpos[i], move_type=MoveType.EEF_MOVE)] + for i in range(self.n_envs) + ] + back_success, back_arm = self.builder.plan_arm_traj( + target_states_list, + reach_arm_qpos, + n_back, + control_part=self.cfg.control_part, + arm_dof=self.arm_dof, + cfg=self.cfg, + ) + success = down_success & back_success + + hand_open_path = self.builder.interpolate_hand_qpos( + self.hand_close_qpos, self.hand_open_qpos, n_waypoints=n_open + ) + + full = torch.empty( + (self.n_envs, n_down + n_open + n_back, self.robot_dof), + dtype=torch.float32, + device=self.device, + ) + full[:, :, :] = state.last_qpos.unsqueeze(1) + full[:, :n_down, self.arm_joint_ids] = down_arm + full[:, :n_down, self.hand_joint_ids] = self.hand_close_qpos + full[:, n_down : n_down + n_open, self.arm_joint_ids] = ( + reach_arm_qpos.unsqueeze(1) + ) + full[:, n_down : n_down + n_open, self.hand_joint_ids] = hand_open_path + full[:, n_down + n_open :, self.arm_joint_ids] = back_arm + full[:, n_down + n_open :, self.hand_joint_ids] = self.hand_open_qpos + + return ActionResult( + success=success, + trajectory=full, + next_state=WorldState( + last_qpos=full[:, -1, :].clone(), + held_object=None, + coordinated_held_object=state.coordinated_held_object, + ), + ) + + def _fail(self, state: WorldState) -> ActionResult: + return ActionResult( + success=torch.zeros(self.n_envs, dtype=torch.bool, device=self.device), + trajectory=torch.empty( + (self.n_envs, 0, self.robot_dof), + dtype=torch.float32, + device=self.device, + ), + next_state=state, + ) + + def _select_tcp_symmetric_place_variant( + self, place_xpos: torch.Tensor, start_qpos: torch.Tensor + ) -> torch.Tensor: + """Choose the closest TCP z-roll variant for an opt-in place target.""" + mirrored_place_xpos = place_xpos.clone() + mirrored_place_xpos[..., :3, 0] = -mirrored_place_xpos[..., :3, 0] + mirrored_place_xpos[..., :3, 1] = -mirrored_place_xpos[..., :3, 1] + place_variants = torch.stack([place_xpos, mirrored_place_xpos], dim=2) + + start_xpos = self.robot.compute_fk( + qpos=start_qpos, + name=self.cfg.control_part, + to_matrix=True, + ) + start_quat = quat_from_matrix(start_xpos[:, :3, :3]) + first_waypoint_quat = quat_from_matrix(place_variants[:, 0, :, :3, :3]) + start_quat = start_quat[:, None, :].expand_as(first_waypoint_quat) + rotation_error = quat_error_magnitude( + first_waypoint_quat.reshape(-1, 4), + start_quat.reshape(-1, 4), + ).reshape(self.n_envs, 2) + best_variant_idx = rotation_error.argmin(dim=1) + + env_idx = torch.arange(self.n_envs, device=self.device)[:, None] + waypoint_idx = torch.arange(place_xpos.shape[1], device=self.device)[None, :] + return place_variants[ + env_idx, + waypoint_idx, + best_variant_idx[:, None], + ] + + +__all__ = ["Place", "PlaceCfg"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/press.py b/embodichain/lab/sim/atomic_actions/primitives/press.py new file mode 100644 index 000000000..c30edbba6 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/press.py @@ -0,0 +1,178 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Press atomic action implementation.""" + +from __future__ import annotations + +from typing import ClassVar + +import torch + +from embodichain.lab.sim.planners import MoveType, PlanState +from embodichain.utils import configclass, logger + +from ._helpers import arm_qpos_from_state +from ..core import ( + ActionCfg, + ActionResult, + AtomicAction, + EndEffectorPoseTarget, + WorldState, +) +from ..trajectory import TrajectoryBuilder + + +@configclass +class PressCfg(ActionCfg): + name: str = "press" + """Name of the action, used for identification and logging.""" + + sample_interval: int = 80 + """Number of waypoints for the full trajectory (hand close + down + back).""" + + hand_interp_steps: int = 5 + """Number of waypoints for closing the gripper before pressing.""" + + hand_control_part: str = "hand" + """Name of the robot part that controls the hand joints.""" + + hand_close_qpos: torch.Tensor | None = None + """Joint positions for the closed hand state, shape ``[hand_dof,]``.""" + + +class Press(AtomicAction): + """Close the gripper, press down to a target pose, then return.""" + + TargetType: ClassVar[type] = EndEffectorPoseTarget + + def __init__( + self, + motion_generator, + cfg: PressCfg | None = None, + ) -> None: + super().__init__(motion_generator, cfg or PressCfg()) + self.builder = TrajectoryBuilder(motion_generator) + self.n_envs = self.robot.get_qpos().shape[0] + self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) + self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part) + self.arm_dof = len(self.arm_joint_ids) + self.hand_dof = len(self.hand_joint_ids) + self.robot_dof = self.robot.dof + + if self.cfg.hand_close_qpos is None: + logger.log_error( + "hand_close_qpos must be specified in PressCfg", ValueError + ) + self.hand_close_qpos = self.builder.expand_hand_qpos( + self.cfg.hand_close_qpos, + n_envs=self.n_envs, + hand_dof=self.hand_dof, + ) + + def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionResult: + press_xpos = self.builder.resolve_pose_target(target.xpos, n_envs=self.n_envs) + start_arm_qpos = self.builder.resolve_start_qpos( + arm_qpos_from_state(state, self.arm_joint_ids), + n_envs=self.n_envs, + arm_dof=self.arm_dof, + control_part=self.cfg.control_part, + ) + start_hand_qpos = state.last_qpos[:, self.hand_joint_ids] + + n_close, n_down, n_back = self._compute_phase_waypoints() + + hand_close_path = self.builder.interpolate_hand_qpos( + start_hand_qpos, + self.hand_close_qpos, + n_waypoints=n_close, + ) + + target_states_list = [ + [PlanState(xpos=press_xpos[i], move_type=MoveType.EEF_MOVE)] + for i in range(self.n_envs) + ] + down_success, down_arm = self.builder.plan_arm_traj( + target_states_list, + start_arm_qpos, + n_down, + control_part=self.cfg.control_part, + arm_dof=self.arm_dof, + cfg=self.cfg, + ) + + press_arm_qpos = down_arm[:, -1, :] + back_arm = self.builder.plan_joint_traj(press_arm_qpos, start_arm_qpos, n_back) + success = down_success + + full = torch.empty( + (self.n_envs, n_close + n_down + n_back, self.robot_dof), + dtype=torch.float32, + device=self.device, + ) + full[:, :, :] = state.last_qpos.unsqueeze(1) + full[:, :n_close, self.arm_joint_ids] = start_arm_qpos.unsqueeze(1) + full[:, :n_close, self.hand_joint_ids] = hand_close_path + full[:, n_close : n_close + n_down, self.arm_joint_ids] = down_arm + full[:, n_close : n_close + n_down, self.hand_joint_ids] = ( + self.hand_close_qpos.unsqueeze(1) + ) + full[:, n_close + n_down :, self.arm_joint_ids] = back_arm + full[:, n_close + n_down :, self.hand_joint_ids] = ( + self.hand_close_qpos.unsqueeze(1) + ) + + return ActionResult( + success=success, + trajectory=full, + next_state=WorldState( + last_qpos=full[:, -1, :].clone(), + held_object=state.held_object, + coordinated_held_object=state.coordinated_held_object, + ), + ) + + def _compute_phase_waypoints(self) -> tuple[int, int, int]: + n_close = self.cfg.hand_interp_steps + if n_close < 1: + logger.log_error( + "hand_interp_steps must be at least 1 for PressCfg.", ValueError + ) + + motion_waypoints = self.cfg.sample_interval - n_close + n_down = motion_waypoints // 2 + n_back = motion_waypoints - n_down + if n_down < 2 or n_back < 2: + logger.log_error( + "Not enough waypoints for press trajectory. Increase " + "sample_interval or decrease hand_interp_steps.", + ValueError, + ) + return n_close, n_down, n_back + + def _fail(self, state: WorldState) -> ActionResult: + return ActionResult( + success=torch.zeros(self.n_envs, dtype=torch.bool, device=self.device), + trajectory=torch.empty( + (self.n_envs, 0, self.robot_dof), + dtype=torch.float32, + device=self.device, + ), + next_state=state, + ) + + +__all__ = ["Press", "PressCfg"] diff --git a/embodichain/lab/sim/atomic_actions/trajectory.py b/embodichain/lab/sim/atomic_actions/trajectory.py index 58c7a8bfd..e2fc7938e 100644 --- a/embodichain/lab/sim/atomic_actions/trajectory.py +++ b/embodichain/lab/sim/atomic_actions/trajectory.py @@ -23,9 +23,10 @@ import numpy as np import torch -from embodichain.lab.sim.planners import PlanState +from embodichain.lab.sim.planners import PlanState, PlanResult, MoveType from embodichain.lab.sim.planners.motion_generator import MotionGenOptions from embodichain.lab.sim.planners.toppra_planner import ToppraPlanOptions +from embodichain.lab.sim.planners.utils import TrajectorySampleMethod from embodichain.lab.sim.utility.action_utils import interpolate_with_distance from embodichain.utils import logger @@ -286,31 +287,45 @@ def split_three_phase( return first, second, third # ------------------------------------------------------------------ - # MotionGen options + # Arm trajectory planning # ------------------------------------------------------------------ - def build_motion_gen_options( + def plan_arm_traj( self, + target_states_list: list[list[PlanState]], start_qpos: torch.Tensor, + n_waypoints: int, *, - sample_interval: int, control_part: str, - ) -> MotionGenOptions: - """Build planner options. Reads ``start_qpos[0]`` because the planner shares options across envs; pass the full batched tensor for type uniformity with other helpers.""" - return MotionGenOptions( - start_qpos=start_qpos[0], + arm_dof: int, + cfg: "ActionCfg | None" = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Plan batched arm trajectories for all environments. + + Returns ``(success:(B,), trajectory:(B, n_waypoints, arm_dof))``. + ``cfg.motion_source`` selects 'ik_interp' (default) or 'motion_gen'. + """ + motion_source = ( + getattr(cfg, "motion_source", "ik_interp") if cfg else "ik_interp" + ) + if motion_source == "motion_gen": + return self._plan_motion_gen( + target_states_list, + start_qpos, + n_waypoints, + control_part=control_part, + arm_dof=arm_dof, + cfg=cfg, + ) + return self._plan_ik_interp( + target_states_list, + start_qpos, + n_waypoints, control_part=control_part, - is_interpolate=True, - is_linear=False, - interpolate_position_step=0.001, - plan_opts=ToppraPlanOptions(sample_interval=sample_interval), + arm_dof=arm_dof, ) - # ------------------------------------------------------------------ - # Arm trajectory planning - # ------------------------------------------------------------------ - - def plan_arm_traj( + def _plan_ik_interp( self, target_states_list: list[list[PlanState]], start_qpos: torch.Tensor, @@ -318,8 +333,8 @@ def plan_arm_traj( *, control_part: str, arm_dof: int, - ) -> tuple[bool, torch.Tensor]: - """Plan batched arm trajectories for all environments.""" + ) -> tuple[torch.Tensor, torch.Tensor]: + """Batched IK + interpolation fallback trajectory source.""" n_envs = start_qpos.shape[0] n_state = len(target_states_list[0]) xpos_traj = torch.zeros( @@ -332,6 +347,7 @@ def plan_arm_traj( trajectory = torch.zeros( (n_envs, n_state, arm_dof), dtype=torch.float32, device=self.device ) + success = torch.ones(n_envs, dtype=torch.bool, device=self.device) qpos_seed = start_qpos for j in range(n_state): is_success, qpos = self.robot.compute_ik( @@ -341,14 +357,120 @@ def plan_arm_traj( logger.log_warning( f"Failed to compute IK for target state {j} in some environments." ) - return False, trajectory + success = success & is_success trajectory[:, j] = qpos qpos_seed = qpos trajectory = torch.concatenate([start_qpos.unsqueeze(1), trajectory], dim=1) + # Failed envs: hold start qpos across all waypoints + if not success.all(): + held = start_qpos.unsqueeze(1).repeat(1, trajectory.shape[1], 1) + trajectory = torch.where(success[:, None, None], trajectory, held) interp = interpolate_with_distance( trajectory=trajectory, interp_num=n_waypoints, device=self.device ) - return True, interp + return success, interp + + def _plan_motion_gen( + self, + target_states_list: list[list[PlanState]], + start_qpos: torch.Tensor, + n_waypoints: int, + *, + control_part: str, + arm_dof: int, + cfg: "ActionCfg | None", + ) -> tuple[torch.Tensor, torch.Tensor]: + """Motion-generator trajectory source.""" + if self.motion_generator is None: + logger.log_error( + "motion_source='motion_gen' requires a MotionGenerator on the engine", + ValueError, + ) + n_envs = start_qpos.shape[0] + plan_states = self._to_batched_plan_states(target_states_list, n_envs) + plan_opts = self._build_plan_opts(cfg, n_waypoints) + planner_type = getattr(cfg, "planner_type", None) + is_interpolate = planner_type != "neural" + result: PlanResult = self.motion_generator.generate( + plan_states, + MotionGenOptions( + start_qpos=start_qpos, + control_part=control_part, + plan_opts=plan_opts, + is_interpolate=is_interpolate, + ), + ) + success = ( + result.success + if isinstance(result.success, torch.Tensor) + else torch.tensor(result.success, device=self.device) + ) + positions = result.positions + # Resample to n_waypoints if the planner returned a different count + if positions.shape[1] != n_waypoints: + positions = interpolate_with_distance( + trajectory=positions, interp_num=n_waypoints, device=self.device + ) + # Failed envs hold start qpos + if not success.all(): + held = start_qpos.unsqueeze(1).repeat(1, positions.shape[1], 1) + positions = torch.where(success[:, None, None], positions, held) + return success, positions + + def _to_batched_plan_states( + self, target_states_list: list[list[PlanState]], n_envs: int + ) -> list[PlanState]: + """Convert per-env PlanState lists into a batched list[PlanState]. + + Each output PlanState carries a leading batch dim ``B`` so it matches + the planner contract. + """ + n_state = len(target_states_list[0]) + batched: list[PlanState] = [] + for j in range(n_state): + sample = target_states_list[0][j] + if sample.xpos is not None: + xpos = torch.stack( + [target_states_list[i][j].xpos for i in range(n_envs)] + ) # (B, 4, 4) + batched.append( + PlanState( + xpos=xpos, + move_type=MoveType.EEF_MOVE, + move_part=sample.move_part, + ) + ) + else: + qpos = torch.stack( + [target_states_list[i][j].qpos for i in range(n_envs)] + ) # (B, DOF) + batched.append( + PlanState( + qpos=qpos, + move_type=MoveType.JOINT_MOVE, + move_part=sample.move_part, + ) + ) + return batched + + def _build_plan_opts(self, cfg: "ActionCfg | None", n_waypoints: int): + """Build planner options from action configuration.""" + planner_type = getattr(cfg, "planner_type", None) + if planner_type in (None, "toppra"): + constraints = {} + vl = getattr(cfg, "velocity_limit", None) + al = getattr(cfg, "acceleration_limit", None) + constraints["velocity"] = vl if vl is not None else 0.2 + constraints["acceleration"] = al if al is not None else 0.5 + return ToppraPlanOptions( + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=n_waypoints, + constraints=constraints, + ) + # neural: planner reads its own cfg; pass minimal options + from embodichain.lab.sim.planners.neural_planner import NeuralPlanOptions + + return NeuralPlanOptions() def plan_joint_traj( self, diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index 3643fea9b..acb3305a5 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -22,7 +22,7 @@ import numpy as np import torch -from typing import Sequence, Union, Dict, Literal, List, Any, Optional, TYPE_CHECKING +from typing import Sequence, Dict, Literal, List, Any, Optional, TYPE_CHECKING from dataclasses import field, MISSING from dexsim.types import ( @@ -357,7 +357,7 @@ class MarkerCfg: marker_type: Literal["axis", "line", "point"] = "axis" """Type of marker to display. Can be 'axis' (3D coordinate frame), 'line', or 'point'. (only axis supported now)""" - axis_xpos: List[np.ndarray] = None + axis_xpos: torch.Tensor | None = None """List of 4x4 transformation matrices defining the position and orientation of each axis marker.""" axis_size: float = 0.002 @@ -427,6 +427,17 @@ def validate_physics_cfg(physics_cfg: PhysicsCfg) -> None: physics_backend_from_cfg(physics_cfg) +@configclass +class WindowCameraPoseCfg: + """Configuration for printing the interactive viewer camera pose.""" + + enable_hotkey: bool = True + """Whether to register the ``p`` hotkey when the window opens.""" + + convert_to_look_at: bool = True + """Whether the hotkey prints a ``set_look_at`` call instead of a matrix.""" + + @configclass class NewtonCollisionAttributesCfg: """Newton-specific per-shape collision/contact attributes. @@ -653,6 +664,8 @@ def attr(self) -> PhysicalAttr: attr.sleep_threshold = self.sleep_threshold attr.restitution = self.restitution attr.enable_ccd = self.enable_ccd + attr.max_linear_velocity = self.max_linear_velocity + attr.max_angular_velocity = self.max_angular_velocity attr.max_depenetration_velocity = self.max_depenetration_velocity attr.min_position_iters = self.min_position_iters attr.min_velocity_iters = self.min_velocity_iters @@ -660,7 +673,7 @@ def attr(self) -> PhysicalAttr: @classmethod def from_dict( - cls, init_dict: Dict[str, Union[str, float, int]] + cls, init_dict: Dict[str, str | float | int] ) -> RigidBodyAttributesCfg: """Initialize the configuration from a dictionary.""" cfg = cls() @@ -737,7 +750,7 @@ def merged_cfg(self, base: RigidBodyAttributesCfg) -> RigidBodyAttributesCfg: @classmethod def from_dict( - cls, init_dict: Dict[str, Union[str, float, int, bool]] + cls, init_dict: Dict[str, str | float | int | bool] ) -> RigidBodyAttributesOverrideCfg: """Initialize the configuration from a dictionary.""" cfg = cls() @@ -1057,7 +1070,7 @@ class JointDrivePropertiesCfg: If the drive type is "none", then no force will be applied to joint. """ - stiffness: Union[Dict[str, float], float] = 1e4 + stiffness: Dict[str, float] | float = 1e4 """Stiffness of the joint drive. The unit depends on the joint model: @@ -1066,7 +1079,7 @@ class JointDrivePropertiesCfg: * For angular joints, the unit is kg-m^2/s^2/rad (N-m/rad). """ - damping: Union[Dict[str, float], float] = 1e3 + damping: Dict[str, float] | float = 1e3 """Damping of the joint drive. The unit depends on the joint model: @@ -1075,20 +1088,20 @@ class JointDrivePropertiesCfg: * For angular joints, the unit is kg-m^2/s/rad (N-m-s/rad). """ - max_effort: Union[Dict[str, float], float] = 1e10 + max_effort: Dict[str, float] | float = 1e10 """Maximum effort that can be applied to the joint (in kg-m^2/s^2).""" - max_velocity: Union[Dict[str, float], float] = 1e10 + max_velocity: Dict[str, float] | float = 1e10 """Maximum velocity that the joint can reach (in rad/s or m/s). For linear joints, this is the maximum linear velocity with unit m/s. For angular joints, this is the maximum angular velocity with unit rad/s. """ - friction: Union[Dict[str, float], float] = 0.0 + friction: Dict[str, float] | float = 0.0 """Friction coefficient of the joint""" - armature: Union[Dict[str, float], float] = 0.0 + armature: Dict[str, float] | float = 0.0 """Joint armature added to joint-space spatial inertia. Units depend on the joint model: @@ -1099,7 +1112,7 @@ class JointDrivePropertiesCfg: @classmethod def from_dict( - cls, init_dict: Dict[str, Union[str, float, int]] + cls, init_dict: Dict[str, str | float | int] ) -> JointDrivePropertiesCfg: """Initialize the configuration from a dictionary.""" cfg = cls() @@ -1121,7 +1134,7 @@ class ObjectBaseCfg: It is used as a base class for specific asset configurations. """ - uid: Union[str, None] = None + uid: str | None = None init_pos: tuple[float, float, float] = (0.0, 0.0, 0.0) """Position of the root in simulation world frame. Defaults to (0.0, 0.0, 0.0).""" @@ -1133,7 +1146,7 @@ class ObjectBaseCfg: """4x4 transformation matrix of the root in local frame. If specified, it will override init_pos and init_rot.""" @classmethod - def from_dict(cls, init_dict: Dict[str, Union[str, float, tuple]]) -> ObjectBaseCfg: + def from_dict(cls, init_dict: Dict[str, str | float | tuple]) -> ObjectBaseCfg: """Initialize the configuration from a dictionary.""" cfg = cls() # Create a new instance of the class (cls) for key, value in init_dict.items(): @@ -1174,19 +1187,108 @@ def from_dict(cls, init_dict: Dict[str, Union[str, float, tuple]]) -> ObjectBase class LightCfg(ObjectBaseCfg): """Configuration for a light asset in the simulation. - This class extends the base asset configuration to include specific properties for lights, + Supports six light types matching the dexsim rendering backend: + + - ``"point"``: Per-environment omnidirectional point light with position + and falloff radius. Created as a batched light (one per environment). + - ``"sun"``: Global directional sun light (infinite distance). Created as + a single scene-level instance. Uses direction only; position is ignored. + Sun-specific fields (``angular_radius``, ``halo_size``, ``halo_falloff``) + are reserved for future backend support. + - ``"direction"``: Global pure directional light at infinite distance. + Created as a single scene-level instance. Direction only; no position. + - ``"spot"``: Per-environment spotlight with position, direction, and + inner/outer cone angles. Created as a batched light. + - ``"rect"``: Per-environment rectangular area light with position, + direction, width, and height. Created as a batched light. + - ``"mesh"``: Per-environment mesh-based emissive light. Requires a + :class:`~dexsim.models.MeshObject` via + :meth:`embodichain.lab.sim.objects.light.Light.set_mesh` + (not tensor-batched). Created as a batched light. + + .. attention:: + The ``angular_radius``, ``halo_size``, and ``halo_falloff`` fields are + reserved for future use. The dexsim Python bindings do not yet expose + setters for these sun-specific properties. """ - # TODO: to be added more light type, such as spot, sun, etc. - light_type: Literal["point"] = "point" + light_type: Literal["point", "sun", "direction", "spot", "rect", "mesh"] = "point" + """Light type. Supported: ``"point"``, ``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, ``"mesh"``.""" + + # ------------------------------------------------------------------ + # Universal properties (apply to all light types) + # ------------------------------------------------------------------ color: tuple[float, float, float] = (1.0, 1.0, 1.0) + """RGB color of the light source. Defaults to white ``(1.0, 1.0, 1.0)``.""" + + intensity: float = 30.0 + """Intensity of the light source in watts/m^2. Defaults to ``30.0``.""" + + enable_shadow: bool = True + """Whether the light casts shadows. Defaults to ``True``.""" + + # ------------------------------------------------------------------ + # Point light + # ------------------------------------------------------------------ + + radius: float = 10.0 + """Falloff radius for point lights. Only used when ``light_type="point"``. Defaults to ``10.0``.""" + + # ------------------------------------------------------------------ + # Directional properties (sun, direction, spot, rect, mesh) + # ------------------------------------------------------------------ + + direction: tuple[float, float, float] = (0.0, 0.0, -1.0) + """Direction vector for directional, spot, rect, and mesh lights. + Defaults to ``(0.0, 0.0, -1.0)`` (pointing down along -Z).""" + + # ------------------------------------------------------------------ + # Sun light (reserved — Python bindings not yet available) + # ------------------------------------------------------------------ + + angular_radius: float = 0.5 + """Angular radius of the sun disc in degrees. Reserved for future use.""" + + halo_size: float = 10.0 + """Halo size for sun light. Reserved for future use.""" + + halo_falloff: float = 3.0 + """Halo falloff for sun light. Reserved for future use.""" + + # ------------------------------------------------------------------ + # Spot light + # ------------------------------------------------------------------ - intensity: float = 50.0 - """Intensity of the light source with unit of watts/m^2.""" + spot_angle_inner: float = 30.0 + """Inner cone angle of the spotlight in degrees. Only used when ``light_type="spot"``. + Defaults to ``30.0``.""" - radius: float = 1e2 - """Falloff of the light, only used for point light.""" + spot_angle_outer: float = 45.0 + """Outer cone angle of the spotlight in degrees. Only used when ``light_type="spot"``. + Defaults to ``45.0``.""" + + # ------------------------------------------------------------------ + # Rect light + # ------------------------------------------------------------------ + + rect_width: float = 1.0 + """Width of the rectangular area light. Only used when ``light_type="rect"``. + Defaults to ``1.0``.""" + + rect_height: float = 1.0 + """Height of the rectangular area light. Only used when ``light_type="rect"``. + Defaults to ``1.0``.""" + + # ------------------------------------------------------------------ + # Mesh light + # ------------------------------------------------------------------ + + mesh_path: str = "" + """Asset path for mesh-based emissive lights. Only used when ``light_type="mesh"``. + The actual mesh assignment is done via + :meth:`embodichain.lab.sim.objects.light.Light.set_mesh` which accepts a + :class:`dexsim.models.MeshObject`. This field stores the path for reference.""" @configclass @@ -1206,20 +1308,43 @@ class RigidObjectCfg(ObjectBaseCfg): body_type: Literal["dynamic", "kinematic", "static"] = "dynamic" - max_convex_hull_num: int = 1 + max_convex_hull_num: int = MISSING """The maximum number of convex hulls that will be created for the rigid body. - If `max_convex_hull_num` is set to larger than 1, the rigid body will be decomposed into multiple convex hulls using coacd alogorithm. + .. deprecated:: + Use :attr:`MeshCfg.max_convex_hull_num` instead. This field is kept for + backward compatibility and overrides the shape-level value when explicitly set. + + If set to larger than 1, the rigid body will be decomposed into multiple convex hulls + using the approximate convex decomposition method specified by :attr:`acd_method`. Reference: https://github.com/SarahWeiii/CoACD """ - sdf_resolution: int = 0 + acd_method: str = MISSING + """The method used for approximate convex decomposition (ACD) of the mesh. + + .. deprecated:: + Use :attr:`MeshCfg.acd_method` instead. This field is kept for + backward compatibility and overrides the shape-level value when explicitly set. + + Currently, ``"coacd"`` and ``"vhacd"`` are supported. Only used when + :attr:`max_convex_hull_num` is set to larger than 1. + """ + + sdf_resolution: int = MISSING """Resolution for the signed distance field (SDF) of the rigid body. - The spacing of the uniformly sampled SDF is equal to the largest AABB extent of the mesh, divided by the resolution. - if `sdf_resolution` is set to larger than 0, a SDF will be generated for collision detection. - SDF will increase the accuracy of collision, but also take more time to initialize and simulate.""" - body_scale: Union[tuple, list] = (1.0, 1.0, 1.0) + .. deprecated:: + Use :attr:`MeshCfg.sdf_resolution` instead. This field is kept for + backward compatibility and overrides the shape-level value when explicitly set. + + The spacing of the uniformly sampled SDF is equal to the largest AABB extent + of the mesh, divided by the resolution. If ``sdf_resolution`` is set to larger + than 0, an SDF will be generated for collision detection. SDF will increase the + accuracy of collision, but also takes more time to initialize and simulate. + """ + + body_scale: tuple | list = (1.0, 1.0, 1.0) """Scale of the rigid body in the simulation world frame.""" use_usd_properties: bool = False @@ -1301,7 +1426,7 @@ class RigidObjectGroupCfg: ) """ - uid: Union[str, None] = None + uid: str | None = None rigid_objects: Dict[str, RigidObjectCfg] = MISSING """Configuration for the rigid objects in the group.""" @@ -1432,12 +1557,12 @@ class RigidConstraintCfg: class URDFCfg: """Standalone configuration class for URDF assembly.""" - components: Dict[str, Dict[str, Union[str, Dict, np.ndarray]]] = field( + components: Dict[str, Dict[str, str | Dict | np.ndarray]] = field( default_factory=dict ) """Dictionary of robot components to be assembled.""" - sensors: Dict[str, Dict[str, Union[str, np.ndarray]]] = field(default_factory=dict) + sensors: Dict[str, Dict[str, str | np.ndarray]] = field(default_factory=dict) """Dictionary of sensors to be attached to the robot.""" use_signature_check: bool = True @@ -1455,7 +1580,7 @@ class URDFCfg: fpath_prefix: str = EMBODICHAIN_DEFAULT_DATA_ROOT + "/assembled" """Output directory prefix for the assembled URDF file.""" - component_prefix: List[tuple[str, Union[str, None]]] = field( + component_prefix: List[tuple[str, str | None]] = field( default_factory=lambda: [ ("chassis", None), ("legs", None), @@ -1822,7 +1947,7 @@ class ArticulationCfg(ObjectBaseCfg): drive_pros: JointDrivePropertiesCfg = JointDrivePropertiesCfg(drive_type="none") """Properties to define the drive mechanism of a joint.""" - body_scale: Union[tuple, list] = (1.0, 1.0, 1.0) + body_scale: tuple | list = (1.0, 1.0, 1.0) """Scale of the articulation in the simulation world frame.""" attrs: RigidBodyAttributesCfg = RigidBodyAttributesCfg() @@ -1847,13 +1972,28 @@ class ArticulationCfg(ObjectBaseCfg): disable_self_collision: bool = True """Whether to enable or disable self-collisions.""" - init_qpos: Union[torch.Tensor, np.ndarray, Sequence[float]] = None - """Initial joint positions of the articulation. - + init_qpos: torch.Tensor | np.ndarray | Sequence[float] = None + """Initial joint positions of the articulation. + If None, the joint positions will be set to zero. If provided, it should be a array of shape (num_joints,). """ + qpos_limits: ( + torch.Tensor | np.ndarray | Sequence[float] | Dict[str, List[float]] | None + ) = None + """Override joint position limits of the articulation. + + If None, the joint position limits from the asset file (URDF/USD) are used. + If provided as a tensor/array of shape (num_joints, 2), it is applied to all + joints in the order of ``joint_names``. + If provided as a dictionary, keys are joint names or regular expressions and + values are ``[min, max]`` limits. + + This field replaces the asset limits for the articulation and can be used to + either tighten or expand the allowed range. + """ + sleep_threshold: float = 0.005 """Energy below which the articulation may go to sleep. Range: [0, max_float32]""" @@ -1882,7 +2022,7 @@ class ArticulationCfg(ObjectBaseCfg): @classmethod def from_dict( - cls, init_dict: Dict[str, Union[str, float, tuple, dict]] + cls, init_dict: Dict[str, str | float | tuple | dict] ) -> ArticulationCfg: """Initialize the configuration from a dictionary.""" cfg = cls() @@ -1949,12 +2089,12 @@ class RobotCfg(ArticulationCfg): """ # TODO: how to support one solver for multiple parts? - solver_cfg: Union[SolverCfg, Dict[str, SolverCfg], None] = None + solver_cfg: SolverCfg | Dict[str, SolverCfg] | None = None """Solver is used to compute forward and inverse kinematics for the robot. """ @classmethod - def from_dict(cls, init_dict: Dict[str, Union[str, float, tuple]]) -> RobotCfg: + def from_dict(cls, init_dict: Dict[str, str | float | tuple]) -> RobotCfg: """Initialize the configuration from a dictionary.""" if isinstance(init_dict, cls): return init_dict diff --git a/embodichain/lab/sim/diff/__init__.py b/embodichain/lab/sim/diff/__init__.py index 45483c72c..ad84e89a7 100644 --- a/embodichain/lab/sim/diff/__init__.py +++ b/embodichain/lab/sim/diff/__init__.py @@ -15,8 +15,8 @@ # ---------------------------------------------------------------------------- """Differentiable Newton stepping for EmbodiChain. -Bridges DexSim's :class:`~dexsim.engine.newton_physics.DifferentiableStepper` -into PyTorch autograd via a :class:`torch.autograd.Function`, and exposes a +Bridges DexSim's manager-owned differentiable trajectory transaction into +PyTorch autograd via a :class:`torch.autograd.Function`, and exposes a :class:`tape_context` manager for advanced users who want to compose their own Warp kernels. """ diff --git a/embodichain/lab/sim/diff/bridge.py b/embodichain/lab/sim/diff/bridge.py index ba0ee49bc..88e1c088a 100644 --- a/embodichain/lab/sim/diff/bridge.py +++ b/embodichain/lab/sim/diff/bridge.py @@ -18,6 +18,7 @@ from __future__ import annotations from contextlib import contextmanager +import math from typing import TYPE_CHECKING, Any, Callable, Iterator import torch @@ -29,6 +30,63 @@ __all__ = ["NewtonStepFunc", "differentiable_step", "tape_context"] +def _physics_dt(nm: Any, sim_state: dict[str, Any]) -> float: + """Resolve the outer Newton step duration represented by one control step.""" + physics_dt = sim_state.get("physics_dt") + if physics_dt is None: + physics_dt = float(nm.solver_dt) * int(nm.num_substeps) + try: + physics_dt = float(physics_dt) + except (TypeError, ValueError) as exc: + raise TypeError("physics_dt must be a positive finite float.") from exc + if not math.isfinite(physics_dt) or physics_dt <= 0.0: + raise ValueError("physics_dt must be a positive finite float.") + return physics_dt + + +def _resolve_step_mode(sim_state: dict[str, Any]) -> tuple[str, Callable | None]: + """Validate the explicit dynamics-versus-kinematics bridge contract.""" + step_mode = sim_state.get("step_mode", "dynamics") + if step_mode not in {"dynamics", "kinematics"}: + raise ValueError( + "step_mode must be 'dynamics' or 'kinematics', " f"got {step_mode!r}." + ) + + step_fn = sim_state.get("step_fn") + if step_mode == "dynamics" and step_fn is not None: + raise ValueError( + "step_fn is only supported when step_mode='kinematics'; " + "the dynamics route always uses Newton solver dynamics." + ) + if step_mode == "kinematics" and step_fn is None: + raise ValueError("step_mode='kinematics' requires a named step_fn.") + return step_mode, step_fn + + +def _reset_tape_then_release(tape: wp.Tape | None, trajectory: Any | None) -> None: + """End tape ownership before releasing the trajectory's model lease.""" + try: + if tape is not None: + tape.reset() + finally: + if trajectory is not None: + trajectory.release() + + +def _abort_forward( + tape: wp.Tape | None, + trajectory: Any | None, +) -> None: + """Best-effort cleanup which never masks the original forward failure.""" + try: + _reset_tape_then_release(tape, trajectory) + except BaseException: + # The active trajectory must not mask the action/solver/output failure + # which caused the abort. DexSim release is idempotent and this path is + # only entered to preserve the original exception. + pass + + @contextmanager def tape_context(manager: "SimulationManager") -> Iterator[wp.Tape]: """Open a Warp tape bound to the manager's Newton state. @@ -48,60 +106,89 @@ def tape_context(manager: "SimulationManager") -> Iterator[wp.Tape]: def differentiable_step( manager: "SimulationManager", *, - apply_control_fn: Callable[[wp.Tape], None], + apply_control_fn: Callable[[wp.Tape, Any], None], substeps: int, dt: float | None = None, -) -> dict: - """Run one EmbodiChain-level physics step inside a Warp tape. +) -> dict[str, Any]: + """Run a low-level manager-owned Newton trajectory inside a Warp tape. + + ``substeps`` remains a legacy solver-step count. It must therefore divide + evenly into whole Newton physics steps; the public trajectory transaction + owns every detached state, contact, control, and generation lease. + + The returned tape and trajectory remain active for the caller to use in a + custom backward pass. After ``tape.backward()`` (or when abandoning the + result), callers must invoke ``tape.reset()`` and then + ``trajectory.release()`` in a ``finally`` block. The helper releases both + automatically only when forward construction itself fails. Args: manager: The owning :class:`SimulationManager` (must be Newton). - apply_control_fn: Callable that writes the joint/body control - targets inside the tape. Invoked once at the start of the - step. Receives the open tape; must launch Warp kernels (or - call dexsim setters that are tape-aware) to populate - ``manager.physics.newton_manager._control``. - substeps: Number of solver substeps to run (typically - ``sim_cfg.sim_steps_per_control``). - dt: Solver dt; defaults to the manager's configured dt. + apply_control_fn: Callable that writes the trajectory-local joint/body + control targets inside the tape. It receives ``(tape, control)`` + and must launch Warp kernels targeting ``control``, never the + manager's shared control buffer. + substeps: Number of solver substeps to run. + dt: Solver dt; defaults to the manager's configured solver dt. Returns: - A dict carrying the tape and the state buffers for the caller to - save in autograd context. + A dict carrying the tape, trajectory, and detached final state for the + caller to retain through backward before resetting/releasing it. """ if not manager.is_newton_backend: raise RuntimeError("differentiable_step requires the Newton backend.") nm = manager.physics.newton_manager - stepper = manager.create_differentiable_stepper() - state_in = nm._state_0 - state_out = nm._model.state() - contacts = stepper.create_contacts() - dt_val = nm.solver_dt if dt is None else float(dt) + if isinstance(substeps, bool) or int(substeps) != substeps or substeps <= 0: + raise ValueError("substeps must be a positive integer.") + substeps = int(substeps) + num_substeps = int(nm.num_substeps) + if num_substeps <= 0: + raise ValueError("Newton num_substeps must be positive.") + if substeps % num_substeps != 0: + raise ValueError( + "substeps must be divisible by Newton num_substeps so the " + "trajectory represents whole physics steps." + ) + dt_val = float(nm.solver_dt if dt is None else dt) + if not math.isfinite(dt_val) or dt_val <= 0.0: + raise ValueError("dt must be a positive finite solver time step.") - tape = wp.Tape() - with tape: - apply_control_fn(tape) - for _ in range(substeps): - stepper.step(state_in, state_out, contacts=contacts, dt=dt_val) - state_in, state_out = state_out, state_in + trajectory = None + tape = None + try: + trajectory = nm.create_differentiable_trajectory( + physics_steps=substeps // num_substeps, + physics_dt=dt_val * num_substeps, + ) + tape = wp.Tape() + with tape: + apply_control_fn(tape, trajectory.control) + final_state = trajectory.step() + nm.commit_differentiable_trajectory(trajectory) + except BaseException: + _abort_forward(tape, trajectory) + raise - # The final state lives in state_in after the swap. return { "tape": tape, - "final_state": state_in, - "stepper": stepper, + "trajectory": trajectory, + "final_state": final_state, + "states": trajectory.states, + "contacts": trajectory.contacts, + "control": trajectory.control, } class NewtonStepFunc(torch.autograd.Function): """torch.autograd.Function bridging Warp tape autodiff to PyTorch. - Forward: launches the action-to-control Warp kernel, runs the - caller-provided ``step_fn`` (differentiable solver loop or FK bypass), - and reads observation / reward as torch tensors via ``wp.to_torch`` - (zero-copy where possible). The obs/reward kernels launched by - ``obs_reward_fn`` run INSIDE the open Warp tape so that their outputs - carry gradient back to ``action_wp``. + Forward: validates an explicit step mode before creating a tape. The + default ``dynamics`` route allocates a manager-owned detached trajectory, + launches the action-to-local-control Warp kernel, records its solver + horizon, and commits it only after tape exit. The explicitly selected + ``kinematics`` route retains its named FK ``step_fn`` escape hatch. + Observation/reward kernels run inside the tape so their outputs carry + gradient back to ``action_wp``. Backward: copies upstream PyTorch grads into the corresponding Warp ``.grad`` buffers, calls ``tape.backward()``, and returns @@ -109,13 +196,17 @@ class NewtonStepFunc(torch.autograd.Function): Callers must supply a ``sim_state`` dict with the following keys: manager: SimulationManager (Newton, requires_grad=True) - substeps: int (used by the default solver-based step_fn) - action_to_control_kernel: callable(action_wp, *kernel_args) + substeps: int control-level physics updates (used by the default + solver-based step route) + step_mode: ``"dynamics"`` (default) or explicit ``"kinematics"`` + action_to_control_kernel: dynamics callable + ``(action_wp, trajectory_control, *kernel_args)``; kinematics + retains ``(action_wp, tape, *kernel_args)`` kernel_args: tuple consumed by action_to_control_kernel obs_reward_fn: callable(final_state) -> dict with torch outputs - step_fn: optional callable() -> final Newton state; when omitted - the bridge runs the differentiable stepper for ``substeps`` - iterations (the original solver-based path) + physics_dt: optional outer Newton step duration (defaults to + ``solver_dt * num_substeps``) + step_fn: required only when ``step_mode == "kinematics"`` The ``obs_reward_fn`` must return a dict containing: _order: tuple of output names (returned in this order) @@ -124,14 +215,35 @@ class NewtonStepFunc(torch.autograd.Function): : torch tensor for each name in ``_order`` """ + @classmethod + def apply(cls, action_torch: torch.Tensor, sim_state: dict[str, Any]) -> Any: + """Capture the caller's grad mode before PyTorch enters ``forward``. + + ``torch.autograd.Function.forward`` always executes with grad mode + disabled, and ``ctx.needs_input_grad`` alone remains true when a + requires-grad action is passed through an outer ``torch.no_grad()`` + block. Passing the ambient mode as a non-differentiable argument lets + the bridge synchronously reset/release no-grad trajectories instead of + retaining an unreachable manager lease. + """ + return super().apply(action_torch, sim_state, torch.is_grad_enabled()) + @staticmethod - def forward(ctx, action_torch: torch.Tensor, sim_state: dict): + def forward( + ctx: Any, + action_torch: torch.Tensor, + sim_state: dict[str, Any], + outer_grad_enabled: bool, + ) -> tuple[torch.Tensor, ...]: manager = sim_state["manager"] substeps = int(sim_state["substeps"]) kernel = sim_state["action_to_control_kernel"] kernel_args = sim_state["kernel_args"] obs_reward_fn = sim_state["obs_reward_fn"] - step_fn = sim_state.get("step_fn") + step_mode, step_fn = _resolve_step_mode(sim_state) + tape_binder = ( + sim_state.get("_bind_dynamics_tape") if step_mode == "dynamics" else None + ) # Save the original action shape so backward can reshape the gradient. ctx.saved_action_shape = action_torch.shape @@ -139,55 +251,113 @@ def forward(ctx, action_torch: torch.Tensor, sim_state: dict): nm = manager.physics.newton_manager action_flat = action_torch.detach().clone().reshape(-1).contiguous() - action_wp = wp.from_torch(action_flat, dtype=wp.float32, requires_grad=True) + needs_action_grad = bool(outer_grad_enabled and ctx.needs_input_grad[0]) + action_wp = wp.from_torch( + action_flat, + dtype=wp.float32, + requires_grad=needs_action_grad, + ) - tape = wp.Tape() - with tape: - kernel(action_wp, *kernel_args) # writes inputs for stepping - if step_fn is not None: - final_state = step_fn() - else: - stepper = manager.create_differentiable_stepper() - state_in = nm._state_0 - state_out = nm._model.state() - contacts = stepper.create_contacts() - dt_val = nm.solver_dt - for _ in range(substeps): - stepper.step(state_in, state_out, contacts=contacts, dt=dt_val) - state_in, state_out = state_out, state_in - final_state = state_in - # Compute obs/reward INSIDE the tape so the reward/obs kernels - # participate in the Warp autodiff graph. The torch tensors - # returned by obs_reward_fn are built via wp.to_torch of - # tape-tracked Warp arrays, so they carry gradient back to - # action_wp when tape.backward() is called. - outputs = obs_reward_fn(final_state) + trajectory = None + tape = None + try: + if step_mode == "dynamics": + if substeps <= 0: + raise ValueError("substeps must be a positive integer.") + trajectory = nm.create_differentiable_trajectory( + physics_steps=substeps, + physics_dt=_physics_dt(nm, sim_state), + ) + + tape = wp.Tape() + try: + with tape: + if tape_binder is not None: + tape_binder(tape) + if step_mode == "dynamics": + kernel(action_wp, trajectory.control, *kernel_args) + final_state = trajectory.step() + else: + # The explicit FK route keeps the historical callback + # shape and receives the open tape, but never detached + # solver control. + kernel(action_wp, tape, *kernel_args) + final_state = step_fn() + + # Validate and materialize outputs inside the tape. A malformed + # output dictionary is a forward failure and must not publish a + # detached dynamics trajectory. + outputs = obs_reward_fn(final_state) + outputs_order = tuple(outputs["_order"]) + output_values = tuple(outputs[name] for name in outputs_order) + outputs_grad_track = outputs.get("_grad_track", {}) + finally: + if tape_binder is not None: + tape_binder(None) + + if trajectory is not None: + nm.commit_differentiable_trajectory(trajectory) + except BaseException: + _abort_forward(tape, trajectory) + raise + + if not needs_action_grad: + _reset_tape_then_release(tape, trajectory) + return output_values ctx.tape = tape + ctx.trajectory = trajectory ctx.action_wp = action_wp - ctx.outputs_order = outputs["_order"] - ctx.outputs_grad_track = outputs.get("_grad_track", {}) - return tuple(outputs[k] for k in outputs["_order"]) + ctx.outputs_order = outputs_order + ctx.outputs_grad_track = outputs_grad_track + ctx._bridge_released = False + return output_values @staticmethod - def backward(ctx, *grad_outputs): - # Copy each upstream grad back into the corresponding Warp .grad. - for name, grad_t in zip(ctx.outputs_order, grad_outputs): - wp_arr = ctx.outputs_grad_track.get(name) - if grad_t is None or wp_arr is None: - continue - # Warp allocates .grad lazily for arrays with requires_grad=True - # that participate in the tape; allocate defensively in case - # the array was created but never written inside the tape. - if wp_arr.grad is None: - wp_arr.grad = wp.zeros_like(wp_arr) - wp.copy( - wp_arr.grad, - wp.from_torch(grad_t.detach().clone().contiguous(), dtype=wp.float32), + def backward( + ctx: Any, + *grad_outputs: torch.Tensor | None, + ) -> tuple[torch.Tensor | None, None, None]: + if getattr(ctx, "_bridge_released", False): + raise RuntimeError( + "NewtonStepFunc backward was already consumed; create a new " + "differentiable trajectory for another backward pass." ) - ctx.tape.backward() - action_grad = wp.to_torch(ctx.action_wp.grad).clone() - ctx.tape.zero() - # Reshape to the original action layout; second input (sim_state) - # has no gradient. - return action_grad.reshape(ctx.saved_action_shape), None + + action_grad = None + try: + # Copy each upstream grad back into the corresponding Warp .grad. + for name, grad_t in zip(ctx.outputs_order, grad_outputs): + wp_arr = ctx.outputs_grad_track.get(name) + if grad_t is None or wp_arr is None: + continue + # Warp allocates .grad lazily for arrays with requires_grad=True + # that participate in the tape; allocate defensively in case + # the array was created but never written inside the tape. + if wp_arr.grad is None: + wp_arr.grad = wp.zeros_like(wp_arr) + wp.copy( + wp_arr.grad, + wp.from_torch( + grad_t.detach().clone().contiguous(), + dtype=wp.float32, + ), + ) + ctx.tape.backward() + action_wp_grad = getattr(ctx.action_wp, "grad", None) + if action_wp_grad is not None: + # Capture the action gradient before reset invalidates tape + # storage, then terminate tape ownership before releasing the + # trajectory's active manager token. + action_grad = wp.to_torch(action_wp_grad).clone() + finally: + try: + _reset_tape_then_release(ctx.tape, ctx.trajectory) + finally: + ctx._bridge_released = True + + if action_grad is None: + return None, None, None + # Reshape to the original action layout; metadata inputs have no + # gradient. + return action_grad.reshape(ctx.saved_action_shape), None, None diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index 06d4e89a5..ad1245db4 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -36,7 +36,10 @@ RigidBodyAttributesOverrideCfg, ) from dexsim.types import PhysicalAttr -from embodichain.utils.string import resolve_matching_names +from embodichain.utils.string import ( + resolve_matching_names, + resolve_matching_names_values, +) from embodichain.lab.sim.common import BatchEntity from embodichain.lab.sim.objects.backends import ( DefaultArticulationView, @@ -150,6 +153,21 @@ def __init__( self._qf = torch.zeros( (self.num_instances, max_dof), dtype=torch.float32, device=self.device ) + self._qpos_limits = torch.as_tensor( + np.array([entity.get_joint_position_limits() for entity in self.entities]), + dtype=torch.float32, + device=self.device, + ) + self._qvel_limits = torch.as_tensor( + np.array([entity.get_joint_velocity_limit() for entity in self.entities]), + dtype=torch.float32, + device=self.device, + ) + self._qf_limits = torch.as_tensor( + np.array([entity.get_joint_effort_limit() for entity in self.entities]), + dtype=torch.float32, + device=self.device, + ) @property def is_newton_backend(self) -> bool: @@ -322,49 +340,32 @@ def joint_armature(self) -> torch.Tensor: device=self.device, ) - @cached_property + @property def qpos_limits(self) -> torch.Tensor: """Get the joint position limits of the articulation. Returns: torch.Tensor: The joint position limits of the articulation with shape (N, dof, 2). """ - return torch.as_tensor( - np.array([entity.get_joint_limits() for entity in self.entities]), - dtype=torch.float32, - device=self.device, - ) + return self._qpos_limits - @cached_property + @property def qvel_limits(self) -> torch.Tensor: """Get the joint velocity limits of the articulation. Returns: torch.Tensor: The joint velocity limits of the articulation with shape (N, dof). """ - # TODO: get joint velocity limits always returns zero? - return torch.as_tensor( - np.array( - [entity.get_drive()[3] for entity in self.entities], - ), - dtype=torch.float32, - device=self.device, - ) + return self._qvel_limits - @cached_property + @property def qf_limits(self) -> torch.Tensor: """Get the joint effort limits of the articulation. Returns: torch.Tensor: The joint effort limits of the articulation with shape (N, dof). """ - return torch.as_tensor( - np.array( - [entity.get_drive()[2] for entity in self.entities], - ), - dtype=torch.float32, - device=self.device, - ) + return self._qf_limits @cached_property def link_vert_face(self) -> Dict[str, Tuple[torch.Tensor, torch.Tensor]]: @@ -505,6 +506,31 @@ def __init__( self.default_joint_max_velocity[0].cpu().numpy().tolist() ) + # Apply configured qpos limits if provided. This replaces the asset + # limits as the baseline and allows expanding the allowed range. + if self.cfg.qpos_limits is not None: + if isinstance(self.cfg.qpos_limits, dict): + indices, _, values = resolve_matching_names_values( + self.cfg.qpos_limits, self.joint_names + ) + local_joint_ids = torch.as_tensor( + indices, dtype=torch.long, device=self.device + ) + values_tensor = torch.as_tensor( + values, dtype=torch.float32, device=self.device + ).unsqueeze(0) + values_tensor = values_tensor.expand(self.num_instances, -1, -1) + self.set_qpos_limits(values_tensor, joint_ids=local_joint_ids) + else: + qpos_limits = torch.as_tensor( + self.cfg.qpos_limits, dtype=torch.float32, device=self.device + ) + if qpos_limits.dim() == 2: + qpos_limits = qpos_limits.unsqueeze(0).expand( + self.num_instances, -1, -1 + ) + self.set_qpos_limits(qpos_limits) + self.pk_chain = None if self.cfg.build_pk_chain: self.pk_chain = create_pk_chain( @@ -717,6 +743,28 @@ def _set_default_collision_filter(self) -> None: collision_filter_data[i, 1] = 1 self.set_collision_filter(collision_filter_data) + def _resolve_env_ids( + self, env_ids: Sequence[int] | torch.Tensor | None + ) -> torch.Tensor: + """Resolve environment ids to a device tensor.""" + if env_ids is None: + return torch.arange( + self.num_instances, dtype=torch.long, device=self.device + ) + if isinstance(env_ids, torch.Tensor): + return env_ids.to(device=self.device, dtype=torch.long) + return torch.as_tensor(env_ids, dtype=torch.long, device=self.device) + + def _resolve_joint_ids( + self, joint_ids: Sequence[int] | torch.Tensor | None + ) -> torch.Tensor: + """Resolve joint ids to a device tensor.""" + if joint_ids is None: + return torch.arange(self.dof, dtype=torch.long, device=self.device) + if isinstance(joint_ids, torch.Tensor): + return joint_ids.to(device=self.device, dtype=torch.long) + return torch.as_tensor(joint_ids, dtype=torch.long, device=self.device) + def set_collision_filter( self, filter_data: torch.Tensor, env_ids: Sequence[int] | None = None ) -> None: @@ -862,6 +910,177 @@ def get_qpos(self, target: bool = False) -> torch.Tensor: """ return self.body_data.qpos if not target else self.body_data.target_qpos + def get_qpos_limits( + self, + joint_ids: Sequence[int] | torch.Tensor | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> torch.Tensor: + """Get joint position limits for selected environments and joints. + + Args: + joint_ids: Joint indices to query. If None, all joints are queried. + env_ids: Environment indices to query. If None, all environments are + queried. + + Returns: + torch.Tensor: Joint position limits with shape (num_envs, num_joints, 2). + """ + local_env_ids = self._resolve_env_ids(env_ids) + local_joint_ids = self._resolve_joint_ids(joint_ids) + return self.body_data.qpos_limits[local_env_ids][:, local_joint_ids, :] + + def _coerce_pair_limit_batch( + self, + values: torch.Tensor | np.ndarray, + local_env_ids: torch.Tensor, + local_joint_ids: torch.Tensor, + ) -> torch.Tensor: + """Normalize batched pair-valued limits to ``(num_envs, num_joints, 2)``.""" + values = torch.as_tensor(values, dtype=torch.float32, device=self.device) + if values.dim() == 2 and len(local_env_ids) == 1: + values = values.unsqueeze(0) + expected_shape = (len(local_env_ids), len(local_joint_ids), 2) + if tuple(values.shape) != expected_shape: + logger.log_error( + f"Expected qpos limit shape {expected_shape}, got {tuple(values.shape)}." + ) + return values + + def _coerce_scalar_limit_batch( + self, + values: torch.Tensor | np.ndarray, + local_env_ids: torch.Tensor, + local_joint_ids: torch.Tensor, + limit_name: str, + ) -> torch.Tensor: + """Normalize batched scalar limits to ``(num_envs, num_joints)``.""" + values = torch.as_tensor(values, dtype=torch.float32, device=self.device) + if values.dim() == 1 and len(local_env_ids) == 1: + values = values.unsqueeze(0) + expected_shape = (len(local_env_ids), len(local_joint_ids)) + if tuple(values.shape) != expected_shape: + logger.log_error( + f"Expected {limit_name} shape {expected_shape}, got {tuple(values.shape)}." + ) + return values + + def set_qpos_limits( + self, + qpos_limits: torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Set joint position limits for selected environments and joints. + + Args: + qpos_limits: Joint position limits with shape (num_envs, num_joints, 2). + When a single environment is selected, a (num_joints, 2) tensor is also accepted. + joint_ids: Joint indices to update. If None, all joints are updated. + env_ids: Environment indices to update. If None, all environments are updated. + """ + local_env_ids = self._resolve_env_ids(env_ids) + local_joint_ids = self._resolve_joint_ids(joint_ids) + qpos_limits = self._coerce_pair_limit_batch( + qpos_limits, local_env_ids, local_joint_ids + ) + joint_ids_np = ( + local_joint_ids.detach().cpu().numpy().astype(np.int32, copy=False) + ) + + failed_envs = [] + for i, env_idx in enumerate(local_env_ids.detach().cpu().tolist()): + result = self._entities[env_idx].set_joint_position_limits( + qpos_limits[i].detach().cpu().numpy(), + joint_ids_np, + ) + if result == -1: + failed_envs.append(env_idx) + continue + self.body_data.qpos_limits[env_idx, local_joint_ids, :] = qpos_limits[i] + + if failed_envs: + logger.log_error( + f"set_joint_position_limits failed for envs {failed_envs} and joint_ids {joint_ids_np.tolist()}." + ) + + def set_qvel_limits( + self, + qvel_limits: torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Set joint velocity limits for selected environments and joints. + + Args: + qvel_limits: Joint velocity limits with shape (num_envs, num_joints). + When a single environment is selected, a (num_joints,) tensor is also accepted. + joint_ids: Joint indices to update. If None, all joints are updated. + env_ids: Environment indices to update. If None, all environments are updated. + """ + local_env_ids = self._resolve_env_ids(env_ids) + local_joint_ids = self._resolve_joint_ids(joint_ids) + qvel_limits = self._coerce_scalar_limit_batch( + qvel_limits, local_env_ids, local_joint_ids, "qvel limit" + ) + joint_ids_np = ( + local_joint_ids.detach().cpu().numpy().astype(np.int32, copy=False) + ) + + failed_envs = [] + for i, env_idx in enumerate(local_env_ids.detach().cpu().tolist()): + result = self._entities[env_idx].set_joint_velocity_limit( + qvel_limits[i].detach().cpu().numpy(), + joint_ids_np, + ) + if result == -1: + failed_envs.append(env_idx) + continue + self.body_data.qvel_limits[env_idx, local_joint_ids] = qvel_limits[i] + + if failed_envs: + logger.log_error( + f"set_joint_velocity_limit failed for envs {failed_envs} and joint_ids {joint_ids_np.tolist()}." + ) + + def set_qf_limits( + self, + qf_limits: torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Set joint effort limits for selected environments and joints. + + Args: + qf_limits: Joint effort limits with shape (num_envs, num_joints). + When a single environment is selected, a (num_joints,) tensor is also accepted. + joint_ids: Joint indices to update. If None, all joints are updated. + env_ids: Environment indices to update. If None, all environments are updated. + """ + local_env_ids = self._resolve_env_ids(env_ids) + local_joint_ids = self._resolve_joint_ids(joint_ids) + qf_limits = self._coerce_scalar_limit_batch( + qf_limits, local_env_ids, local_joint_ids, "qf limit" + ) + joint_ids_np = ( + local_joint_ids.detach().cpu().numpy().astype(np.int32, copy=False) + ) + + failed_envs = [] + for i, env_idx in enumerate(local_env_ids.detach().cpu().tolist()): + result = self._entities[env_idx].set_joint_effort_limit( + qf_limits[i].detach().cpu().numpy(), + joint_ids_np, + ) + if result == -1: + failed_envs.append(env_idx) + continue + self.body_data.qf_limits[env_idx, local_joint_ids] = qf_limits[i] + + if failed_envs: + logger.log_error( + f"set_joint_effort_limit failed for envs {failed_envs} and joint_ids {joint_ids_np.tolist()}." + ) + def set_qpos( self, qpos: torch.Tensor, @@ -890,18 +1109,8 @@ def set_qpos( else: qpos = qpos.to(device=self.device, dtype=torch.float32) - if joint_ids is None: - local_joint_ids = torch.arange( - self.dof, device=self.device, dtype=torch.int32 - ) - elif not isinstance(joint_ids, torch.Tensor): - local_joint_ids = torch.as_tensor( - joint_ids, dtype=torch.int32, device=self.device - ) - else: - local_joint_ids = joint_ids.to(device=self.device, dtype=torch.int32) - - local_env_ids = self._all_indices if env_ids is None else env_ids + local_joint_ids = self._resolve_joint_ids(joint_ids) + local_env_ids = self._resolve_env_ids(env_ids) # Make sure qpos is 2D tensor if qpos.dim() == 1: @@ -913,10 +1122,10 @@ def set_qpos( f"env_ids: {local_env_ids}, qpos.shape: {qpos.shape}" ) - limits = self.body_data.qpos_limits[0].T - lower_limits = limits[0][local_joint_ids] - upper_limits = limits[1][local_joint_ids] - qpos = qpos.clamp(lower_limits, upper_limits) + selected_limits = self.body_data.qpos_limits[local_env_ids][ + :, local_joint_ids, : + ] + qpos = qpos.clamp(selected_limits[..., 0], selected_limits[..., 1]) self._data.articulation_view.apply_qpos( qpos, local_env_ids, @@ -935,6 +1144,25 @@ def get_qvel(self, target: bool = False) -> torch.Tensor: """ return self.body_data.qvel if not target else self.body_data.target_qvel + def get_qvel_limits( + self, + joint_ids: Sequence[int] | torch.Tensor | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> torch.Tensor: + """Get joint velocity limits for selected environments and joints. + + Args: + joint_ids: Joint indices to query. If None, all joints are queried. + env_ids: Environment indices to query. If None, all environments are + queried. + + Returns: + torch.Tensor: Joint velocity limits with shape (num_envs, num_joints). + """ + local_env_ids = self._resolve_env_ids(env_ids) + local_joint_ids = self._resolve_joint_ids(joint_ids) + return self.body_data.qvel_limits[local_env_ids][:, local_joint_ids] + def set_qvel( self, qvel: torch.Tensor, @@ -1027,6 +1255,25 @@ def set_qf( self._data.articulation_view.apply_qf(qf, local_env_ids, local_joint_ids) + def get_qf_limits( + self, + joint_ids: Sequence[int] | torch.Tensor | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> torch.Tensor: + """Get joint effort limits for selected environments and joints. + + Args: + joint_ids: Joint indices to query. If None, all joints are queried. + env_ids: Environment indices to query. If None, all environments are + queried. + + Returns: + torch.Tensor: Joint effort limits with shape (num_envs, num_joints). + """ + local_env_ids = self._resolve_env_ids(env_ids) + local_joint_ids = self._resolve_joint_ids(joint_ids) + return self.body_data.qf_limits[local_env_ids][:, local_joint_ids] + def set_mass( self, mass: torch.Tensor, @@ -1233,6 +1480,8 @@ def set_joint_drive( """ local_env_ids = self._all_indices if env_ids is None else env_ids local_joint_ids = np.arange(self.dof) if joint_ids is None else joint_ids + cache_env_ids = self._resolve_env_ids(env_ids) + cache_joint_ids = self._resolve_joint_ids(joint_ids) for i, env_idx in enumerate(local_env_ids): drive_args = { @@ -1253,6 +1502,19 @@ def set_joint_drive( drive_args["armature"] = armature[i].cpu().numpy() self._entities[env_idx].set_drive(**drive_args) + if max_velocity is not None: + max_velocity = torch.as_tensor( + max_velocity, dtype=torch.float32, device=self.device + ) + self._data._qvel_limits[cache_env_ids[:, None], cache_joint_ids] = ( + max_velocity + ) + if max_effort is not None: + max_effort = torch.as_tensor( + max_effort, dtype=torch.float32, device=self.device + ) + self._data._qf_limits[cache_env_ids[:, None], cache_joint_ids] = max_effort + def get_joint_drive( self, joint_ids: Sequence[int] | None = None, @@ -1854,3 +2116,6 @@ def destroy(self) -> None: arenas[i].remove_skeleton(entity) else: arenas[i].remove_articulation(entity) + + +__all__ = ["ArticulationData", "Articulation"] diff --git a/embodichain/lab/sim/objects/light.py b/embodichain/lab/sim/objects/light.py index 8cea8d8ae..065267333 100644 --- a/embodichain/lab/sim/objects/light.py +++ b/embodichain/lab/sim/objects/light.py @@ -14,14 +14,19 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import torch import numpy as np -from typing import List, Sequence +from typing import TYPE_CHECKING, List, Sequence from dexsim.render import Light as _Light from embodichain.lab.sim.cfg import LightCfg from embodichain.lab.sim.common import BatchEntity from embodichain.utils import logger +if TYPE_CHECKING: + from dexsim.models import MeshObject + class Light(BatchEntity): """Light represents a batch of lights in the simulation. @@ -87,6 +92,281 @@ def set_falloff( """ self._apply_scalar(falloffs, env_ids, "set_falloff") + def set_direction( + self, directions: torch.Tensor, env_ids: Sequence[int] | None = None + ) -> None: + """Set direction for directional-type lights. + + Only applies to ``sun``, ``direction``, ``spot``, ``rect``, and ``mesh`` + light types. Logs a warning and no-ops for other types. + + Args: + directions (torch.Tensor): Tensor of shape (3,) or (M, 3), representing + (x, y, z) direction vectors. + env_ids (Sequence[int] | None): Indices of instances to set. If None: + - For shape (3,), applies to all instances. + - For shape (M, 3), M must equal num_instances, applies per-instance. + """ + if self.cfg.light_type not in ("sun", "direction", "spot", "rect", "mesh"): + logger.log_warning( + f"set_direction not applicable to light type " + f"'{self.cfg.light_type}', ignoring." + ) + return + self._apply_vector3(directions, env_ids, "set_direction") + + def set_spot_angle( + self, + inner_angles: torch.Tensor, + outer_angles: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set inner and outer cone angles for spot lights. + + Only applies to ``spot`` light type. Logs a warning and no-ops for other types. + + Args: + inner_angles (torch.Tensor): Tensor of shape () (0-dim), (1,), or (M,) + representing inner cone angle in degrees. + outer_angles (torch.Tensor): Tensor of shape () (0-dim), (1,), or (M,) + representing outer cone angle in degrees. + env_ids (Sequence[int] | None): Indices of instances to set. + """ + if self.cfg.light_type != "spot": + logger.log_warning( + f"set_spot_angle not applicable to light type " + f"'{self.cfg.light_type}', ignoring." + ) + return + + if not torch.is_tensor(inner_angles) or not torch.is_tensor(outer_angles): + logger.log_error("set_spot_angle requires torch.Tensor arguments") + return + + inner_cpu = inner_angles.detach().cpu() + outer_cpu = outer_angles.detach().cpu() + + if env_ids is None: + all_ids = list(range(self.num_instances)) + else: + all_ids = list(env_ids) + + # Both scalar (0-dim): broadcast to all_ids + if inner_cpu.ndim == 0 and outer_cpu.ndim == 0: + iv = float(inner_cpu.item()) + ov = float(outer_cpu.item()) + for i in all_ids: + self._entities[i].set_spot_angle(iv, ov) + return + + # Both 1D with matching length + if inner_cpu.ndim == 1 and outer_cpu.ndim == 1: + ilen, olen = inner_cpu.shape[0], outer_cpu.shape[0] + inner_arr = inner_cpu.numpy() + outer_arr = outer_cpu.numpy() + + if ilen == olen == self.num_instances and env_ids is None: + for i in range(self.num_instances): + self._entities[i].set_spot_angle( + float(inner_arr[i]), float(outer_arr[i]) + ) + return + + if env_ids is not None and ilen == olen == len(all_ids): + for idx, i in enumerate(all_ids): + self._entities[i].set_spot_angle( + float(inner_arr[idx]), float(outer_arr[idx]) + ) + return + + # length-1 broadcast + if ilen == olen == 1: + iv = float(inner_arr[0]) + ov = float(outer_arr[0]) + for i in all_ids: + self._entities[i].set_spot_angle(iv, ov) + return + + logger.log_error( + f"set_spot_angle: invalid tensor shapes " + f"inner={tuple(inner_cpu.shape)}, outer={tuple(outer_cpu.shape)}" + ) + + def set_rect_wh( + self, + widths: torch.Tensor, + heights: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set width and height for rectangular area lights. + + Only applies to ``rect`` light type. Logs a warning and no-ops for other types. + + Args: + widths (torch.Tensor): Tensor of shape () (0-dim), (1,), or (M,) + representing width of the rectangular light. + heights (torch.Tensor): Tensor of shape () (0-dim), (1,), or (M,) + representing height of the rectangular light. + env_ids (Sequence[int] | None): Indices of instances to set. + """ + if self.cfg.light_type != "rect": + logger.log_warning( + f"set_rect_wh not applicable to light type " + f"'{self.cfg.light_type}', ignoring." + ) + return + + if not torch.is_tensor(widths) or not torch.is_tensor(heights): + logger.log_error("set_rect_wh requires torch.Tensor arguments") + return + + w_cpu = widths.detach().cpu() + h_cpu = heights.detach().cpu() + + if env_ids is None: + all_ids = list(range(self.num_instances)) + else: + all_ids = list(env_ids) + + # Both scalar (0-dim): broadcast to all_ids + if w_cpu.ndim == 0 and h_cpu.ndim == 0: + wv = float(w_cpu.item()) + hv = float(h_cpu.item()) + for i in all_ids: + self._entities[i].set_rect_wh(wv, hv) + return + + # Both 1D with matching length + if w_cpu.ndim == 1 and h_cpu.ndim == 1: + wlen, hlen = w_cpu.shape[0], h_cpu.shape[0] + w_arr = w_cpu.numpy() + h_arr = h_cpu.numpy() + + if wlen == hlen == self.num_instances and env_ids is None: + for i in range(self.num_instances): + self._entities[i].set_rect_wh(float(w_arr[i]), float(h_arr[i])) + return + + if env_ids is not None and wlen == hlen == len(all_ids): + for idx, i in enumerate(all_ids): + self._entities[i].set_rect_wh(float(w_arr[idx]), float(h_arr[idx])) + return + + # length-1 broadcast + if wlen == hlen == 1: + wv = float(w_arr[0]) + hv = float(h_arr[0]) + for i in all_ids: + self._entities[i].set_rect_wh(wv, hv) + return + + logger.log_error( + f"set_rect_wh: invalid tensor shapes " + f"width={tuple(w_cpu.shape)}, height={tuple(h_cpu.shape)}" + ) + + def set_mesh( + self, + mesh: "MeshObject", + env_ids: Sequence[int] | None = None, + ) -> None: + """Set the mesh for mesh-type lights. + + Only applies to ``mesh`` light type. Logs a warning and no-ops for other types. + This is NOT tensor-batched — the same MeshObject is assigned to all targeted + instances. + + Args: + mesh (MeshObject): The mesh object to assign to the light. + env_ids (Sequence[int] | None): Indices of instances to set. If None, + applies to all instances. + """ + if self.cfg.light_type != "mesh": + logger.log_warning( + f"set_mesh not applicable to light type " + f"'{self.cfg.light_type}', ignoring." + ) + return + + if env_ids is None: + target_ids = list(range(self.num_instances)) + else: + target_ids = list(env_ids) + + for i in target_ids: + try: + self._entities[i].set_mesh(mesh) + except Exception as e: + logger.log_error(f"set_mesh: error for instance {i}: {e}") + + def enable_shadow( + self, + flags: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Enable or disable shadow casting. + + Applies to all light types. + + Args: + flags (torch.Tensor): Boolean tensor of shape () (0-dim), (1,), or (M,). + Non-zero values enable shadows; zero disables. + env_ids (Sequence[int] | None): Indices of instances to set. + """ + if not torch.is_tensor(flags): + logger.log_error( + f"enable_shadow requires a torch.Tensor, got {type(flags)}" + ) + return + + cpu = flags.detach().cpu() + if env_ids is None: + all_ids = list(range(self.num_instances)) + else: + all_ids = list(env_ids) + + # Scalar: broadcast + if cpu.ndim == 0: + val = bool(cpu.item() != 0) + for i in all_ids: + self._entities[i].set_shadow(val) + return + + # 1D tensor + if cpu.ndim == 1: + length = cpu.shape[0] + arr = cpu.numpy() + if length == self.num_instances and env_ids is None: + for i in range(self.num_instances): + self._entities[i].set_shadow(bool(arr[i] != 0)) + return + if env_ids is not None and length == len(all_ids): + for idx, i in enumerate(all_ids): + self._entities[i].set_shadow(bool(arr[idx] != 0)) + return + if length == 1: + val = bool(arr[0] != 0) + for i in all_ids: + self._entities[i].set_shadow(val) + return + + logger.log_error( + f"enable_shadow: tensor shape {tuple(cpu.shape)} is invalid for broadcasting" + ) + + @property + def is_global(self) -> bool: + """Whether this light is a global scene light (single instance). + + Global lights (``"sun"``, ``"direction"``) are infinite-distance + light sources that illuminate the entire scene. They are never + batched per environment. + + Returns: + bool: True if the light is a global scene light. + """ + return self.cfg.light_type in ("sun", "direction") + def set_local_pose( self, pose: torch.Tensor, @@ -104,6 +384,13 @@ def set_local_pose( - For matrix input (4,4) broadcast to all, or (M,4,4) with M == num_instances. to_matrix (bool): Interpret `pose` as full 4x4 matrix if True, else as vector(s). """ + if self.is_global: + logger.log_warning( + f"set_local_pose not applicable to '{self.cfg.light_type}' light type " + f"(infinite distance, direction only). Use set_direction() instead." + ) + return + if not torch.is_tensor(pose): logger.log_error( f"set_local_pose requires a torch.Tensor, got {type(pose)}" @@ -292,8 +579,63 @@ def _apply_scalar( ) def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Reset the light to its initial configuration state. + + Applies only the properties relevant to ``self.cfg.light_type``. + + .. attention:: + When this light has only one instance (global lights such as + ``"sun"`` and ``"direction"``, or any light in a single-environment + scene), ``env_ids`` is normalized to ``None`` because there is only + one instance to update. Passing per-environment indices is silently + ignored. + + Args: + env_ids (Sequence[int] | None): The environment IDs to reset. + If None, resets all environments. + """ self.cfg: LightCfg + light_type = self.cfg.light_type + + # Global lights have a single instance — ignore per-env indices. + if self.num_instances == 1: + env_ids = None + + # Universal properties self.set_color(torch.as_tensor(self.cfg.color), env_ids=env_ids) self.set_intensity(torch.as_tensor(self.cfg.intensity), env_ids=env_ids) - self.set_falloff(torch.as_tensor(self.cfg.radius), env_ids=env_ids) - self.set_local_pose(torch.as_tensor(self.cfg.init_pos), env_ids=env_ids) + self.enable_shadow( + torch.as_tensor(float(self.cfg.enable_shadow)), env_ids=env_ids + ) + + # Position (all types except global infinite-distance lights) + if not self.is_global: + self.set_local_pose(torch.as_tensor(self.cfg.init_pos), env_ids=env_ids) + + # Point light: falloff + if light_type == "point": + self.set_falloff(torch.as_tensor(self.cfg.radius), env_ids=env_ids) + + # Directional types: direction vector + if light_type in ("sun", "direction", "spot", "rect", "mesh"): + self.set_direction(torch.as_tensor(self.cfg.direction), env_ids=env_ids) + + # Spot light: cone angles + if light_type == "spot": + self.set_spot_angle( + torch.as_tensor(self.cfg.spot_angle_inner), + torch.as_tensor(self.cfg.spot_angle_outer), + env_ids=env_ids, + ) + + # Rect light: dimensions + if light_type == "rect": + self.set_rect_wh( + torch.as_tensor(self.cfg.rect_width), + torch.as_tensor(self.cfg.rect_height), + env_ids=env_ids, + ) + + # Mesh light: mesh_path is stored in cfg but actual mesh assignment + # is done via set_mesh() which requires a MeshObject. + # Sun-specific angular_radius/halo are reserved for future backend support. diff --git a/embodichain/lab/sim/objects/rigid_object.py b/embodichain/lab/sim/objects/rigid_object.py index c1f3357a1..cb26e7a6a 100644 --- a/embodichain/lab/sim/objects/rigid_object.py +++ b/embodichain/lab/sim/objects/rigid_object.py @@ -20,7 +20,7 @@ import dexsim import numpy as np -from dataclasses import dataclass +from dataclasses import dataclass, MISSING from typing import List, Sequence, Union from functools import cached_property @@ -34,6 +34,7 @@ is_newton_scene, ) from embodichain.lab.sim.objects.backends.base import RigidBodyViewBase +from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim import ( VisualMaterial, VisualMaterialInst, @@ -274,9 +275,15 @@ def __init__( def __str__(self) -> str: parent_str = super().__str__() + max_hull = self.cfg.max_convex_hull_num + if max_hull is MISSING: + if isinstance(self.cfg.shape, MeshCfg): + max_hull = self.cfg.shape.max_convex_hull_num + else: + max_hull = 1 return ( parent_str - + f" | body type: {self.body_type} | max_convex_hull_num: {self.cfg.max_convex_hull_num}" + + f" | body type: {self.body_type} | max_convex_hull_num: {max_hull}" ) @cached_property diff --git a/embodichain/lab/sim/objects/robot.py b/embodichain/lab/sim/objects/robot.py index dcf13b64e..b6ab3f07b 100644 --- a/embodichain/lab/sim/objects/robot.py +++ b/embodichain/lab/sim/objects/robot.py @@ -14,10 +14,12 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import torch import numpy as np -from typing import List, Dict, Tuple, Union, Sequence +from typing import List, Dict, Tuple, Sequence from dataclasses import dataclass, field from tensordict import TensorDict @@ -77,13 +79,15 @@ def __init__( self._control_groups: Dict[str, ControlGroup] = {} + # Articulation initialization may apply configured qpos limits through + # Robot.set_qpos_limits(), which synchronizes this registry. + self._solvers: Dict[str, BaseSolver] = {} + if self.cfg.control_parts: self._init_control_parts(self.cfg.control_parts) super().__init__(cfg, entities, device) - self._solvers = {} - if self.cfg.solver_cfg: self.init_solver(self.cfg.solver_cfg) @@ -128,7 +132,7 @@ def get_joint_ids( else [i for i in self._joint_ids[name] if i not in self.mimic_ids] ) - def get_link_names(self, name: str | None = None) -> Union[List[str], None]: + def get_link_names(self, name: str | None = None) -> List[str] | None: """Get the link names of the robot for a specific control part. If no control part is specified, return all link names. @@ -148,8 +152,33 @@ def get_link_names(self, name: str | None = None) -> Union[List[str], None]: ) return self._control_groups[name].link_names + def _resolve_limit_joint_ids( + self, + name: str | None, + joint_ids: Sequence[int] | None, + ) -> Sequence[int] | None: + """Resolve joint selection for robot limit APIs. + + When `name` is specified, the control-part joint ids take precedence over + explicit `joint_ids` to preserve the existing Robot setter/getter behavior. + """ + if name is None: + return joint_ids + + if not self.control_parts or name not in self.control_parts: + logger.log_error( + f"The control part '{name}' does not exist in the robot's control parts." + ) + if joint_ids is not None: + logger.log_warning("`joint_ids` is ignored when `name` is specified.") + return self.get_joint_ids(name=name) + def get_qpos_limits( - self, name: str | None = None, env_ids: Sequence[int] | None = None + self, + name: str | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + *, + joint_ids: Sequence[int] | torch.Tensor | None = None, ) -> torch.Tensor: """Get the joint position limits (qpos) of the robot for a specific control part. @@ -157,26 +186,25 @@ def get_qpos_limits( Args: name (str | None): The name of the control part to get the qpos limits for. - env_ids (Sequence[int] | None): The environment ids to get the qpos limits for. If None, all environments are used. + env_ids (Sequence[int] | torch.Tensor | None): The environment ids to get the qpos limits for. If None, all environments are used. + joint_ids (Sequence[int] | torch.Tensor | None): Joint indices to get the qpos limits for. + Must be passed as a keyword argument. Returns: torch.Tensor: Joint position limits with shape (N, dof, 2), where N is the number of environments. """ - local_env_ids = self._all_indices if env_ids is None else env_ids - - qpos_limits = self.body_data.qpos_limits - if name is None: - return qpos_limits[local_env_ids, :] - else: - if not self.control_parts or name not in self.control_parts: - logger.log_error( - f"The control part '{name}' does not exist in the robot's control parts." - ) - part_joint_ids = self.get_joint_ids(name=name) - return qpos_limits[local_env_ids][:, part_joint_ids, :] + resolved_joint_ids = self._resolve_limit_joint_ids(name, joint_ids) + return super().get_qpos_limits( + joint_ids=resolved_joint_ids, + env_ids=env_ids, + ) def get_qvel_limits( - self, name: str | None = None, env_ids: Sequence[int] | None = None + self, + name: str | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + *, + joint_ids: Sequence[int] | torch.Tensor | None = None, ) -> torch.Tensor: """Get the joint velocity limits (qvel) of the robot for a specific control part. @@ -184,26 +212,25 @@ def get_qvel_limits( Args: name (str | None): The name of the control part to get the qvel limits for. - env_ids (Sequence[int] | None): The environment ids to get the qvel limits for. If None, all environments are used. + env_ids (Sequence[int] | torch.Tensor | None): The environment ids to get the qvel limits for. If None, all environments are used. + joint_ids (Sequence[int] | torch.Tensor | None): Joint indices to get the qvel limits for. + Must be passed as a keyword argument. Returns: torch.Tensor: Joint velocity limits with shape (N, dof), where N is the number of environments. """ - local_env_ids = self._all_indices if env_ids is None else env_ids - - qvel_limits = self.body_data.qvel_limits - if name is None: - return qvel_limits[local_env_ids, :] - else: - if not self.control_parts or name not in self.control_parts: - logger.log_error( - f"The control part '{name}' does not exist in the robot's control parts." - ) - part_joint_ids = self.get_joint_ids(name=name) - return qvel_limits[local_env_ids][:, part_joint_ids] + resolved_joint_ids = self._resolve_limit_joint_ids(name, joint_ids) + return super().get_qvel_limits( + joint_ids=resolved_joint_ids, + env_ids=env_ids, + ) def get_qf_limits( - self, name: str | None = None, env_ids: Sequence[int] | None = None + self, + name: str | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + *, + joint_ids: Sequence[int] | torch.Tensor | None = None, ) -> torch.Tensor: """Get the joint effort limits (qf) of the robot for a specific control part. @@ -211,23 +238,91 @@ def get_qf_limits( Args: name (str | None): The name of the control part to get the qf limits for. - env_ids (Sequence[int] | None): The environment ids to get the qf limits for. If None, all environments are used. + env_ids (Sequence[int] | torch.Tensor | None): The environment ids to get the qf limits for. If None, all environments are used. + joint_ids (Sequence[int] | torch.Tensor | None): Joint indices to get the qf limits for. + Must be passed as a keyword argument. Returns: torch.Tensor: Joint effort limits with shape (N, dof), where N is the number of environments. """ - local_env_ids = self._all_indices if env_ids is None else env_ids + resolved_joint_ids = self._resolve_limit_joint_ids(name, joint_ids) + return super().get_qf_limits( + joint_ids=resolved_joint_ids, + env_ids=env_ids, + ) - qf_limits = self.body_data.qf_limits - if name is None: - return qf_limits[local_env_ids, :] - else: - if not self.control_parts or name not in self.control_parts: - logger.log_error( - f"The control part '{name}' does not exist in the robot's control parts." - ) - part_joint_ids = self.get_joint_ids(name=name) - return qf_limits[local_env_ids][:, part_joint_ids] + def set_qpos_limits( + self, + qpos_limits: torch.Tensor, + name: str | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + *, + joint_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Set the joint position limits (qpos) of the robot. + + Args: + qpos_limits: Joint position limits with shape (N, num_joints, 2). + name (str | None): The name of the control part to set the qpos limits for. + env_ids (Sequence[int] | torch.Tensor | None): The environment ids to set the qpos limits for. + joint_ids (Sequence[int] | torch.Tensor | None): Joint indices to set the qpos limits for. + Must be passed as a keyword argument; ignored when ``name`` is provided. + """ + resolved_joint_ids = self._resolve_limit_joint_ids(name, joint_ids) + super().set_qpos_limits( + qpos_limits=qpos_limits, + joint_ids=resolved_joint_ids, + env_ids=env_ids, + ) + self._sync_solver_limits(name=name) + + def set_qvel_limits( + self, + qvel_limits: torch.Tensor, + name: str | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + *, + joint_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Set the joint velocity limits (qvel) of the robot. + + Args: + qvel_limits: Joint velocity limits with shape (N, num_joints). + name (str | None): The name of the control part to set the qvel limits for. + env_ids (Sequence[int] | torch.Tensor | None): The environment ids to set the qvel limits for. + joint_ids (Sequence[int] | torch.Tensor | None): Joint indices to set the qvel limits for. + Must be passed as a keyword argument; ignored when ``name`` is provided. + """ + resolved_joint_ids = self._resolve_limit_joint_ids(name, joint_ids) + super().set_qvel_limits( + qvel_limits=qvel_limits, + joint_ids=resolved_joint_ids, + env_ids=env_ids, + ) + + def set_qf_limits( + self, + qf_limits: torch.Tensor, + name: str | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + *, + joint_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Set the joint effort limits (qf) of the robot. + + Args: + qf_limits: Joint effort limits with shape (N, num_joints). + name (str | None): The name of the control part to set the qf limits for. + env_ids (Sequence[int] | torch.Tensor | None): The environment ids to set the qf limits for. + joint_ids (Sequence[int] | torch.Tensor | None): Joint indices to set the qf limits for. + Must be passed as a keyword argument; ignored when ``name`` is provided. + """ + resolved_joint_ids = self._resolve_limit_joint_ids(name, joint_ids) + super().set_qf_limits( + qf_limits=qf_limits, + joint_ids=resolved_joint_ids, + env_ids=env_ids, + ) def get_proprioception(self) -> TensorDict[str, torch.Tensor]: """Gets robot proprioception information, primarily for agent state representation in robot learning scenarios. @@ -904,11 +999,31 @@ def _set_default_joint_drive(self) -> None: drive_type=drive_type, ) - def init_solver(self, cfg: Union[SolverCfg, Dict[str, SolverCfg]]) -> None: + def _sync_solver_limits(self, name: str | None = None) -> None: + """Synchronize solver joint limits with the robot's effective qpos limits.""" + if not self._solvers: + return + + if not self.control_parts: + solver = self._solvers.get("default") + if solver is not None: + solver.update_with_robot_limit(self._data.qpos_limits[0]) + return + + part_names = [name] if name is not None else list(self.control_parts.keys()) + for part_name in part_names: + solver = self._solvers.get(part_name) + if solver is None: + continue + joint_ids = self.get_joint_ids(name=part_name) + joint_limits = self._data.qpos_limits[0][joint_ids] + solver.update_with_robot_limit(joint_limits) + + def init_solver(self, cfg: SolverCfg | Dict[str, SolverCfg]) -> None: """Initialize the kinematic solver for the robot. Args: - cfg (Union[SolverCfg, Dict[str, SolverCfg]]): The configuration for the kinematic solver. + cfg (SolverCfg | Dict[str, SolverCfg]): The configuration for the kinematic solver. """ self.cfg: RobotCfg @@ -921,6 +1036,7 @@ def init_solver(self, cfg: Union[SolverCfg, Dict[str, SolverCfg]]) -> None: if cfg.urdf_path is None: cfg.urdf_path = self.cfg.fpath self._solvers["default"] = cfg.init_solver(device=self.device) + self._sync_solver_limits() elif isinstance(cfg, Dict): if isinstance(self.cfg.control_parts, Dict) is False: logger.log_error( @@ -940,10 +1056,10 @@ def init_solver(self, cfg: Union[SolverCfg, Dict[str, SolverCfg]]) -> None: or solver_cfg.joint_names is None ): solver_cfg.joint_names = self.cfg.control_parts[part_name] - self._solvers[name] = solver_cfg.init_solver(device=self.device) - joint_ids = self.get_joint_ids(name=part_name) - joint_limits = self._data.qpos_limits[0][joint_ids] - self._solvers[name].update_with_robot_limit(joint_limits) + self._solvers[part_name] = solver_cfg.init_solver( + device=self.device + ) + self._sync_solver_limits() def get_solver(self, name: str | None = None) -> BaseSolver | None: """Get the kinematic solver for a specific control part. @@ -1112,3 +1228,6 @@ def set_physical_visible( def destroy(self) -> None: return super().destroy() + + +__all__ = ["ControlGroup", "Robot"] diff --git a/embodichain/lab/sim/planners/base_planner.py b/embodichain/lab/sim/planners/base_planner.py index 5eff14cf5..c71866f34 100644 --- a/embodichain/lab/sim/planners/base_planner.py +++ b/embodichain/lab/sim/planners/base_planner.py @@ -42,6 +42,45 @@ class PlanOptions: pass +def _infer_batch_size(target_states: list[PlanState]) -> int | None: + """Return the leading batch dim B of the first tensor found in target_states, or None if none.""" + for s in target_states: + for t in (s.qpos, s.xpos, s.qvel, s.qacc): + if isinstance(t, torch.Tensor) and t.dim() >= 1: + return int(t.shape[0]) + return None + + +def _check_batch_consistency( + target_states: list[PlanState], + expected_b: int | None, + robot_num_instances: int | None, +) -> int: + """Validate that all PlanState tensors share the same leading B and match the robot.""" + bs = set() + for s in target_states: + b = _infer_batch_size([s]) + if b is not None: + bs.add(b) + if len(bs) > 1: + logger.log_error( + f"All PlanState entries must share the same batch dim B, got {sorted(bs)}", + ValueError, + ) + b = bs.pop() if bs else 1 + if expected_b is not None and b != expected_b: + logger.log_error( + f"Batch dim B={b} does not match robot.num_instances={expected_b}", + ValueError, + ) + if robot_num_instances is not None and b not in (1, robot_num_instances): + logger.log_error( + f"Batch dim B={b} must be 1 or robot.num_instances={robot_num_instances}", + ValueError, + ) + return b + + def validate_plan_options(_func=None, *, options_cls: type = PlanOptions): """Decorator (factory) that validates the ``options`` argument is a ``PlanOptions`` instance. @@ -77,6 +116,12 @@ def wrapper(self, *args, **kwargs): f"(or a subclass), but got {type(options).__name__}.", TypeError, ) + target_states = kwargs.get("target_states", args[0] if args else None) + if target_states is not None and hasattr(self, "robot"): + robot_num = getattr(self.robot, "num_instances", None) + _check_batch_consistency( + target_states, expected_b=robot_num, robot_num_instances=robot_num + ) return func(self, *args, **kwargs) return wrapper @@ -123,17 +168,21 @@ def plan( planning algorithm. Args: - target_states: List of dictionaries containing target states + target_states: list of :class:`PlanState` waypoints. Tensor fields + carry a leading batch dim ``B`` (e.g. ``qpos`` is ``(B, DOF)``). Returns: - PlanResult: An object containing: - - success: bool, whether planning succeeded - - positions: torch.Tensor (N, DOF), joint positions along trajectory - - velocities: torch.Tensor (N, DOF), joint velocities along trajectory - - accelerations: torch.Tensor (N, DOF), joint accelerations along trajectory - - times: torch.Tensor (N,), time stamps for each point - - duration: float, total trajectory duration - - error_msg: Optional error message if planning failed + PlanResult: An env-batched object containing: + - success: torch.Tensor ``(B,)`` bool, per-env success + - positions: torch.Tensor ``(B, N, DOF)``, joint positions + - velocities: torch.Tensor ``(B, N, DOF)`` or ``None``, joint + velocities. Populated by planners that compute dynamics; may be + ``None`` for planners that do not. + - accelerations: torch.Tensor ``(B, N, DOF)`` or ``None``, joint + accelerations. Populated by planners that compute dynamics; may + be ``None`` for planners that do not. + - dt: torch.Tensor ``(B, N)``, per-point time deltas + - duration: torch.Tensor ``(B,)``, total trajectory duration per env """ logger.log_error("Subclasses must implement plan() method", NotImplementedError) diff --git a/embodichain/lab/sim/planners/motion_generator.py b/embodichain/lab/sim/planners/motion_generator.py index ba09efc2e..d83004957 100644 --- a/embodichain/lab/sim/planners/motion_generator.py +++ b/embodichain/lab/sim/planners/motion_generator.py @@ -29,11 +29,16 @@ ToppraPlannerCfg, NeuralPlanner, NeuralPlannerCfg, + NeuralPlanOptions, ) from embodichain.lab.sim.utility.action_utils import interpolate_with_nums from embodichain.utils import logger, configclass from .utils import MovePart, MoveType, PlanState, PlanResult -from .utils import calculate_point_allocations, interpolate_xpos +from .utils import ( + calculate_point_allocations, + interpolate_xpos, + interpolate_xpos_batched, +) __all__ = ["MotionGenerator", "MotionGenCfg", "MotionGenOptions"] @@ -53,7 +58,7 @@ class MotionGenCfg: class MotionGenOptions: start_qpos: torch.Tensor | None = None - """Optional starting joint configuration for the trajectory. If provided, the planner will ensure that the trajectory starts from this configuration. If not provided, the planner will use the current joint configuration of the robot as the starting point.""" + """Optional starting joint configuration for the trajectory, shape (B, DOF). If provided, the planner will ensure that the trajectory starts from this configuration. If not provided, the planner will use the current joint configuration of the robot as the starting point.""" control_part: str | None = None """Name of the robot part to control, e.g. 'left_arm'. Must correspond to a valid control part defined in the robot's configuration.""" @@ -151,48 +156,55 @@ def generate( Returns: PlanResult containing the planned trajectory details. """ + if options.is_interpolate and isinstance(self.planner, NeuralPlanner): + logger.log_warning( + "is_interpolate=True is not supported with NeuralPlanner; " + "disabling interpolation." + ) + options.is_interpolate = False + if options.is_interpolate: - # interpolate trajectory to generate more waypoints for smoother motion and better constraint handling - if target_states[0].move_type == MoveType.EEF_MOVE: - xpos_list = [] - for state in target_states: - if state.move_type != MoveType.EEF_MOVE: + move_type = target_states[0].move_type + if move_type == MoveType.EEF_MOVE: + for s in target_states: + if s.move_type != move_type: logger.log_error( - f"All states must be the same. First state is {target_states[0].move_type}, but got {state.move_type}" + f"All states must share move_type; got {s.move_type}", + ValueError, ) - xpos_list.append(state.xpos) - qpos_list = None - elif target_states[0].move_type == MoveType.JOINT_MOVE: - qpos_list = [] - for state in target_states: - if state.move_type != MoveType.JOINT_MOVE: + xpos_list = torch.stack([s.xpos for s in target_states]).transpose( + 0, 1 + ) # (B, N, 4, 4) + qpos_list = None + elif move_type == MoveType.JOINT_MOVE: + for s in target_states: + if s.move_type != move_type: logger.log_error( - f"All states must be the same. First state is {target_states[0].move_type}, but got {state.move_type}" + f"All states must share move_type; got {s.move_type}", + ValueError, ) - qpos_list.append(state.qpos) - xpos_list = None + qpos_list = torch.stack([s.qpos for s in target_states]).transpose( + 0, 1 + ) # (B, N, DOF) + xpos_list = None else: logger.log_error( - f"Unsupported move type for pre-interpolation: {target_states[0].move_type}" + f"Unsupported move type for pre-interpolation: {move_type}" ) - if qpos_list is not None: - qpos_list = torch.stack(qpos_list) - if xpos_list is not None: - xpos_list = torch.stack(xpos_list) - if options.start_qpos is not None: + start = options.start_qpos + if start.dim() == 1: + start = start.unsqueeze(0) if qpos_list is not None: - qpos_list = torch.cat( - [options.start_qpos.unsqueeze(0), qpos_list], dim=0 - ) + qpos_list = torch.cat([start.unsqueeze(1), qpos_list], dim=1) if xpos_list is not None: start_xpos = self.robot.compute_fk( - qpos=options.start_qpos.unsqueeze(0), - name=options.control_part, - to_matrix=True, + qpos=start, name=options.control_part, to_matrix=True ) - xpos_list = torch.cat([start_xpos, xpos_list], dim=0) + if start_xpos.dim() == 3: + start_xpos = start_xpos.unsqueeze(1) + xpos_list = torch.cat([start_xpos, xpos_list], dim=1) qpos_interpolated, xpos_interpolated = self.interpolate_trajectory( control_part=options.control_part, @@ -200,23 +212,17 @@ def generate( qpos_list=qpos_list, options=options, ) - if not options.plan_opts: - # Directly return the interpolated trajectory if no further planning is needed return PlanResult( success=True, positions=qpos_interpolated, xpos_list=xpos_interpolated, ) - target_plan_states = [] - for qpos in qpos_interpolated: - target_plan_states.append( - PlanState( - move_type=MoveType.JOINT_MOVE, - qpos=qpos, - ) - ) + target_plan_states = [ + PlanState(move_type=MoveType.JOINT_MOVE, qpos=qpos_interpolated[:, j]) + for j in range(qpos_interpolated.shape[1]) + ] else: target_plan_states = target_states @@ -225,12 +231,20 @@ def generate( options.plan_opts = self.planner.default_plan_options() else: options.plan_opts = PlanOptions() - # Planner-specific options (e.g. NeuralPlanOptions) must be set on - # plan_opts explicitly, same as ToppraPlanOptions. - result = self.planner.plan( + + # Propagate MotionGenOptions fields into NeuralPlanOptions so that callers + # can set control_part/start_qpos at the MotionGenerator level. + if isinstance(self.planner, NeuralPlanner) and isinstance( + options.plan_opts, NeuralPlanOptions + ): + if options.plan_opts.control_part is None: + options.plan_opts.control_part = options.control_part + if options.plan_opts.start_qpos is None: + options.plan_opts.start_qpos = options.start_qpos + + return self.planner.plan( target_states=target_plan_states, options=options.plan_opts ) - return result def estimate_trajectory_sample_count( self, @@ -467,164 +481,167 @@ def interpolate_trajectory( options: MotionGenOptions = MotionGenOptions(), ) -> tuple[torch.Tensor, torch.Tensor | None]: r"""Interpolate trajectory based on provided waypoints. - This method performs interpolation on the provided waypoints to generate a smoother trajectory. - It supports both Cartesian (end-effector) and joint space interpolation based on the control part and options specified. + + This method performs interpolation on the provided waypoints to generate a + smoother trajectory. It supports both Cartesian (end-effector) and joint + space interpolation based on the control part and options specified. Args: - control_part: Name of the robot part to control, e.g. 'left_arm'. Must correspond to a valid control part defined in the robot's configuration. - xpos_list: List of end-effector poses (torch.Tensor of shape [N, 4, 4]) to interpolate through. Required if control_part is an end-effector control part. - qpos_list: List of joint positions (torch.Tensor of shape [N, DOF]) to interpolate through. Required if control_part is a joint control part. - options: MotionGenOptions containing interpolation settings such as step size and whether to use linear interpolation. + control_part: Name of the robot part to control, e.g. 'left_arm'. Must + correspond to a valid control part defined in the robot's configuration. + xpos_list: End-effector poses, shape ``(B, N, 4, 4)`` or ``(N, 4, 4)``. + Required if control_part is an end-effector control part. + qpos_list: Joint positions, shape ``(B, N, DOF)`` or ``(N, DOF)``. + Required if control_part is a joint control part. + options: MotionGenOptions containing interpolation settings such as step + size and whether to use linear interpolation. Returns: Tuple containing: - - interpolate_qpos_list: Tensor of interpolated joint positions along the trajectory, shape [M, DOF] - - feasible_pose_targets: Tensor of corresponding end-effector poses for the interpolated joint positions, shape [M, 4, 4]. This is useful for verifying the interpolation results and can be None if not applicable. + - interpolate_qpos_list: Interpolated joint positions, shape + ``(B, M, DOF)``. + - feasible_pose_targets: Corresponding end-effector poses, shape + ``(B, M, 4, 4)``, or ``None`` if not applicable. """ + # Normalize single-env inputs to batched form. + if qpos_list is not None and qpos_list.dim() == 2: + qpos_list = qpos_list.unsqueeze(0) + if xpos_list is not None and xpos_list.dim() == 3: + xpos_list = xpos_list.unsqueeze(0) + if qpos_list is not None and xpos_list is None and options.is_linear: - qpos_batch = qpos_list.unsqueeze(0) # [n_env=1, n_batch=N, dof] - xpos_batch = self.robot.compute_batch_fk( - qpos=qpos_batch, + # qpos_list is (B, N, DOF); compute_batch_fk handles batched qpos directly. + xpos_list = self.robot.compute_batch_fk( + qpos=qpos_list, name=control_part, to_matrix=True, - ) - xpos_list = xpos_batch.squeeze_(0) + ) # (B, N, 4, 4) if xpos_list is None and qpos_list is None: logger.log_error("Either xpos_list or qpos_list must be provided") - # if options.start_qpos is not None: - # start_xpos = self.robot.compute_fk( - # qpos=options.start_qpos.unsqueeze(0), name=control_part, to_matrix=True - # ) - # qpos_list = ( - # torch.cat([options.start_qpos.unsqueeze(0), qpos_list], dim=0) - # if qpos_list is not None - # else None - # ) - # if xpos_list is not None: - # xpos_list = torch.cat([start_xpos, xpos_list], dim=0) - - # Input validation - if (xpos_list is not None and len(xpos_list) < 2) or ( - qpos_list is not None and len(qpos_list) < 2 + # Input validation: the waypoint count is the second-to-last or last batch dim. + if (xpos_list is not None and xpos_list.shape[-3] < 2) or ( + qpos_list is not None and qpos_list.shape[-2] < 2 ): logger.log_error( "xpos_list and qpos_list must contain at least 2 way points" ) qpos_seed = options.start_qpos + if qpos_seed is not None and qpos_seed.dim() == 1: + qpos_seed = qpos_seed.unsqueeze(0) if qpos_seed is None and qpos_list is not None: - # first waypoint as seed - qpos_seed = qpos_list[0] + # First waypoint per env as seed. + qpos_seed = qpos_list[:, 0] # (B, DOF) if qpos_seed is None: - # fallback to current robot state as seed - qpos_seed = self.robot.get_qpos(name=control_part)[0] + # Fallback to current robot state as seed. + qpos_seed = self.robot.get_qpos(name=control_part) # (B, DOF) # Generate trajectory - interpolate_qpos_list = [] if options.is_linear or qpos_list is None: - # Calculate point allocations for interpolation - interpolated_point_allocations = calculate_point_allocations( - xpos_list, - step_size=options.interpolate_position_step, - angle_step=options.interpolate_angle_step, - device=self.device, - ) - - # Linear cartesian interpolation + # ``calculate_point_allocations`` only handles single-env (N, 4, 4), + # so compute allocations per env and use the per-segment maximum so + # all envs can share the same interpolated pose count. + per_env_allocations = [ + calculate_point_allocations( + xpos_list[b], + step_size=options.interpolate_position_step, + angle_step=options.interpolate_angle_step, + device=self.device, + ) + for b in range(xpos_list.shape[0]) + ] + n_segments = xpos_list.shape[1] - 1 + interpolated_point_allocations = [ + max(alloc[i] for alloc in per_env_allocations) + for i in range(n_segments) + ] + + # Linear cartesian interpolation, batched across B envs. total_interpolated_poses = [] - - # TODO: We may improve the computation efficiency using warp for parallel interpolation of all segments if necessary. - for i in range(len(xpos_list) - 1): - interpolated_poses = interpolate_xpos( - ( - xpos_list[i].detach().cpu().numpy() - if isinstance(xpos_list, torch.Tensor) - else xpos_list[i] - ), - ( - xpos_list[i + 1].detach().cpu().numpy() - if isinstance(xpos_list, torch.Tensor) - else xpos_list[i + 1] - ), + for i in range(n_segments): + seg = interpolate_xpos_batched( + xpos_list[:, i], + xpos_list[:, i + 1], interpolated_point_allocations[i], - ) - total_interpolated_poses.extend(interpolated_poses) - - total_interpolated_poses = torch.as_tensor( - np.asarray(total_interpolated_poses), - dtype=torch.float32, - device=self.device, - ) - - # Use batch IK for performance - # compute_batch_ik expects (n_envs, n_batch, 7) or (n_envs, n_batch, 4, 4) - # Here we assume n_envs = 1 or we want to apply this to all envs if available. - # Since MotionGenerator usually works with self.robot.device, we use its batching capabilities. - qpos_seed_repeat = ( - qpos_seed.unsqueeze(0) - .repeat(total_interpolated_poses.shape[0], 1) - .unsqueeze(0) - ) + ) # (B, seg, 4, 4) + total_interpolated_poses.append(seg) + total_interpolated_poses = torch.cat( + total_interpolated_poses, dim=1 + ) # (B, M, 4, 4) + + qpos_seed_b = qpos_seed + if qpos_seed_b.dim() == 1: + qpos_seed_b = qpos_seed_b.unsqueeze(0).repeat(xpos_list.shape[0], 1) + joint_seed = qpos_seed_b.unsqueeze(1).repeat( + 1, total_interpolated_poses.shape[1], 1 + ) # (B, M, D) success_batch, qpos_batch = self.robot.compute_batch_ik( - pose=total_interpolated_poses.unsqueeze(0), - joint_seed=qpos_seed_repeat, # Or use qpos_seed if properly shaped + pose=total_interpolated_poses, + joint_seed=joint_seed, name=control_part, - ) - - # success_batch: (n_envs, n_batch), qpos_batch: (n_envs, n_batch, dof) - success_mask = success_batch[0] # Take first env - qpos_results = qpos_batch[0] - has_nan_mask = torch.isnan(qpos_results).any(dim=-1) - - valid_mask = success_mask & (~has_nan_mask) - valid_indices = torch.where(valid_mask)[0] + ) # (B, M), (B, M, D) - interpolate_qpos_list = qpos_results[valid_indices] - feasible_pose_targets = total_interpolated_poses[valid_indices] + has_nan = torch.isnan(qpos_batch).any(dim=-1) + valid = success_batch.bool() & (~has_nan) # (B, M) # Vectorized FK feasibility check to keep only physically consistent IK outputs. - if len(interpolate_qpos_list) > 0: + if valid.any(): fk_batch = self.robot.compute_batch_fk( - qpos=interpolate_qpos_list.unsqueeze(0), + qpos=qpos_batch, name=control_part, to_matrix=True, - ).squeeze_(0) + ) # (B, M, 4, 4) pos_err = torch.norm( - fk_batch[:, :3, 3] - feasible_pose_targets[:, :3, 3], dim=-1 + fk_batch[:, :, :3, 3] - total_interpolated_poses[:, :, :3, 3], + dim=-1, ) rot_err = torch.norm( - fk_batch[:, :3, :3] - feasible_pose_targets[:, :3, :3], + fk_batch[:, :, :3, :3] - total_interpolated_poses[:, :, :3, :3], dim=(-2, -1), ) - valid_mask = (pos_err < 0.02) & (rot_err < 0.2) - interpolate_qpos_list = interpolate_qpos_list[valid_mask] - feasible_pose_targets = feasible_pose_targets[valid_mask] + fk_valid = (pos_err < 0.02) & (rot_err < 0.2) + valid = valid & fk_valid + + # Per-env filter: keep only valid rows; pad short envs by repeating last valid. + B, M, D = qpos_batch.shape + max_valid = int(valid.sum(dim=1).max().item()) + max_valid = max(max_valid, 1) + interp_q = torch.zeros( + B, max_valid, D, device=self.device, dtype=torch.float32 + ) + feasible = torch.zeros( + B, max_valid, 4, 4, device=self.device, dtype=torch.float32 + ) + for b in range(B): + v = qpos_batch[b][valid[b]] + f = total_interpolated_poses[b][valid[b]] + if v.shape[0] == 0: + v = qpos_batch[b : b + 1, 0] + f = total_interpolated_poses[b : b + 1, 0] + interp_q[b, : v.shape[0]] = v + interp_q[b, v.shape[0] :] = v[-1] + feasible[b, : f.shape[0]] = f + feasible[b, f.shape[0] :] = f[-1] + interpolate_qpos_list = interp_q + feasible_pose_targets = feasible else: - # Perform joint space interpolation directly if not linear or if end-effector poses are not provided - qpos_interpolated = qpos_list.unsqueeze_(0).permute(1, 0, 2) # [N, 1, DOF] - + # Joint-space interpolation. qpos_list is (B, N, DOF). if isinstance(options.interpolate_nums, int): - interp_nums = [options.interpolate_nums] * (len(qpos_list) - 1) + interp_nums = [options.interpolate_nums] * (qpos_list.shape[1] - 1) else: - if len(options.interpolate_nums) != len(qpos_list) - 1: + if len(options.interpolate_nums) != qpos_list.shape[1] - 1: logger.log_error( - "Length of interpolate_nums list must be equal to number of segments (len(qpos_list) - 1)" + "Length of interpolate_nums list must equal number of segments", + ValueError, ) interp_nums = options.interpolate_nums - interpolate_qpos_list = ( - interpolate_with_nums( - qpos_interpolated, - interp_nums=interp_nums, - device=self.device, - ) - .permute(1, 0, 2) - .squeeze_(0) - ) # [M, DOF] - + interpolate_qpos_list = interpolate_with_nums( + qpos_list, interp_nums=interp_nums, device=self.device + ) # (B, M, DOF) feasible_pose_targets = None return interpolate_qpos_list, feasible_pose_targets diff --git a/embodichain/lab/sim/planners/neural_planner.py b/embodichain/lab/sim/planners/neural_planner.py index 62adb0e64..90e917f9e 100644 --- a/embodichain/lab/sim/planners/neural_planner.py +++ b/embodichain/lab/sim/planners/neural_planner.py @@ -26,6 +26,7 @@ BasePlanner, BasePlannerCfg, PlanOptions, + _infer_batch_size, validate_plan_options, ) from embodichain.lab.sim.planners.utils import MoveType, PlanResult, PlanState @@ -39,6 +40,30 @@ ] +def _safe_torch_load(path: Path, map_location: torch.device) -> dict: + """Load a PyTorch checkpoint with safe deserialization when possible. + + Attempts ``weights_only=True`` first. If that fails (e.g. on older PyTorch + versions or checkpoints with unsupported pickle objects), falls back to + ``weights_only=False`` and logs a warning. + + Args: + path: Path to the checkpoint file. + map_location: Device to map tensors to. + + Returns: + The loaded checkpoint dictionary. + """ + try: + return torch.load(path, map_location=map_location, weights_only=True) + except (TypeError, RuntimeError, AttributeError) as exc: + logger.log_warning( + f"Failed to load checkpoint with weights_only=True from {path}: {exc}. " + "Falling back to weights_only=False." + ) + return torch.load(path, map_location=map_location, weights_only=False) + + def _layer_init(layer: nn.Linear, std: float = np.sqrt(2), bias_const: float = 0.0): torch.nn.init.orthogonal_(layer.weight, std) torch.nn.init.constant_(layer.bias, bias_const) @@ -203,15 +228,26 @@ class NeuralPlanOptions(PlanOptions): class NeuralPlanner(BasePlanner): + r"""Neural motion planner based on an APG waypoint transformer policy. + + The planner loads a checkpoint containing a waypoint-conditioned actor and + rolls it out in closed loop to drive the arm toward a sequence of + end-effector waypoints. Velocities and accelerations in the returned + :class:`PlanResult` are estimated via finite differences over the generated + position trajectory. + + Args: + cfg: Configuration for the neural planner. + + Raises: + ValueError: If ``checkpoint_path`` is missing or invalid. + FileNotFoundError: If the checkpoint file does not exist. + KeyError: If the checkpoint is missing required keys. + """ + def __init__(self, cfg: NeuralPlannerCfg): super().__init__(cfg) - if self.robot.num_instances > 1: - logger.log_error( - "NeuralPlanner currently supports one robot instance", - NotImplementedError, - ) - self.cfg: NeuralPlannerCfg = cfg if cfg.checkpoint_path is MISSING or not str(cfg.checkpoint_path): logger.log_error("checkpoint_path is required", ValueError) @@ -226,7 +262,7 @@ def _load_checkpoint(self, checkpoint_path: Path) -> None: f"Checkpoint not found: {checkpoint_path}", FileNotFoundError ) - ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False) + ckpt = _safe_torch_load(checkpoint_path, map_location=self.device) if "agent" not in ckpt: raise KeyError( f"Checkpoint at '{checkpoint_path}' is missing 'agent'. " @@ -313,8 +349,35 @@ def plan( target_states: list[PlanState], options: NeuralPlanOptions = NeuralPlanOptions(), ) -> PlanResult: + r"""Execute neural trajectory planning. + + Runs the waypoint transformer policy in closed loop for each environment + until all waypoints are reached or ``max_steps`` is exhausted. + + Args: + target_states: List of :class:`PlanState` waypoints. Each entry must + use :attr:`MoveType.EEF_MOVE` and carry an ``xpos`` tensor of + shape ``(B, 4, 4)``. + options: :class:`NeuralPlanOptions` with ``control_part``, + ``start_qpos``, and ``max_steps`` overrides. + + Returns: + :class:`PlanResult` containing the planned trajectory. All tensor + fields are env-batched with leading dim ``B``: ``success`` ``(B,)``, + ``positions``/``velocities``/``accelerations`` ``(B, N, DOF)``, + ``xpos_list`` ``(B, N, 4, 4)``, ``dt`` ``(B, N)``, and + ``duration`` ``(B,)``. Velocities and accelerations are computed + via finite differences and are therefore approximate. + + Raises: + ValueError: If ``control_part`` is not provided, if a target state + is not ``EEF_MOVE``, or if ``start_qpos`` has too few joints. + """ if not target_states: - return PlanResult(success=False, positions=None) + return PlanResult( + success=torch.zeros(0, dtype=torch.bool, device=self.device), + positions=None, + ) control_part = options.control_part or self.cfg.control_part if control_part is None: @@ -327,14 +390,16 @@ def plan( target_states ) qpos = self._initial_qpos(control_part, options.start_qpos) + b = qpos.shape[0] limits = self.robot.get_qpos_limits(name=control_part)[0].to(self.device) lower = limits[: self._action_dim, 0] upper = limits[: self._action_dim, 1] - last_action = torch.zeros(1, self._action_dim, device=self.device) - active_idx = 0 - positions = [qpos.squeeze(0).clone()] - xpos_list = [self._fk_matrix(qpos, control_part).squeeze(0)] + last_action = torch.zeros(b, self._action_dim, device=self.device) + active_idx = torch.zeros(b, dtype=torch.long, device=self.device) + positions = [qpos.clone()] + xpos_list = [self._fk_matrix(qpos, control_part)] + converged = torch.zeros(b, dtype=torch.bool, device=self.device) max_steps = int(options.max_steps or self._max_steps) with torch.no_grad(): @@ -350,38 +415,60 @@ def plan( last_action, ) action = self._actor(self._normalizer.normalize(obs)).clamp(-1.0, 1.0) + # Hold converged envs: zero their action so qpos does not drift. + # `converged` reflects state up to the end of the previous step, so + # once an env converged at the end of step N its action is masked + # from step N+1 onward. + not_converged = ~converged + action = torch.where( + not_converged.unsqueeze(-1), action, torch.zeros_like(action) + ) qpos[:, : self._action_dim] += action * float(self.cfg.action_scale) qpos[:, : self._action_dim] = torch.clamp( qpos[:, : self._action_dim], lower, upper ) - last_action = action - - positions.append(qpos.squeeze(0).clone()) - xpos_list.append(self._fk_matrix(qpos, control_part).squeeze(0)) + last_action = torch.where( + not_converged.unsqueeze(-1), action, last_action + ) + positions.append(qpos.clone()) + xpos_list.append(self._fk_matrix(qpos, control_part)) ee_pose = self._fk_pose_xyzw(qpos, control_part) - if self._is_active_reached( + reached = self._is_active_reached( ee_pose, waypoints_pos, waypoints_quat, active_idx, episode_k - ): - active_idx += 1 - if active_idx >= episode_k: - break + ) + active_idx = torch.where(reached, active_idx + 1, active_idx) + converged = converged | (active_idx >= episode_k) + if converged.all(): + break positions_t = torch.stack(positions) xpos_t = torch.stack(xpos_list) - success = active_idx >= episode_k dt = torch.full( (positions_t.shape[0],), float(self.cfg.dt), dtype=torch.float32, device=self.device, ) + dt = dt.unsqueeze(0).expand(b, -1) + positions_t = positions_t.permute(1, 0, 2) + xpos_t = xpos_t.permute(1, 0, 2, 3) + velocities_t, accelerations_t = self._compute_vel_acc_via_finite_diff( + positions_t, dt + ) + success = active_idx >= episode_k return PlanResult( success=success, positions=positions_t, + velocities=velocities_t, + accelerations=accelerations_t, xpos_list=xpos_t, dt=dt, - duration=float(max(positions_t.shape[0] - 1, 0) * self.cfg.dt), + duration=torch.full( + (b,), + float(max(positions_t.shape[1] - 1, 0) * self.cfg.dt), + device=self.device, + ), ) def _parse_waypoints( @@ -393,11 +480,10 @@ def _parse_waypoints( f"at most {self._num_waypoints}.", ValueError, ) - - waypoint_pos = torch.zeros(1, self._num_waypoints, 3, device=self.device) - waypoint_quat = torch.zeros(1, self._num_waypoints, 4, device=self.device) - valid_mask = torch.zeros(1, self._num_waypoints, device=self.device) - + b = _infer_batch_size(target_states) or 1 + waypoint_pos = torch.zeros(b, self._num_waypoints, 3, device=self.device) + waypoint_quat = torch.zeros(b, self._num_waypoints, 4, device=self.device) + valid_mask = torch.zeros(b, self._num_waypoints, device=self.device) for idx, target in enumerate(target_states): if target.move_type != MoveType.EEF_MOVE or target.xpos is None: logger.log_error( @@ -405,19 +491,20 @@ def _parse_waypoints( ValueError, ) xpos = torch.as_tensor(target.xpos, dtype=torch.float32, device=self.device) - waypoint_pos[0, idx] = xpos[:3, 3] - waypoint_quat[0, idx] = convert_quat( - quat_from_matrix(xpos[:3, :3].unsqueeze(0)), to="xyzw" - )[0] - valid_mask[0, idx] = 1.0 - + if xpos.dim() == 2: + xpos = xpos.unsqueeze(0) + waypoint_pos[:, idx] = xpos[:, :3, 3] + waypoint_quat[:, idx] = convert_quat( + quat_from_matrix(xpos[:, :3, :3]), to="xyzw" + ) + valid_mask[:, idx] = 1.0 return waypoint_pos, waypoint_quat, valid_mask, len(target_states) def _initial_qpos( self, control_part: str, start_qpos: torch.Tensor | None ) -> torch.Tensor: if start_qpos is None: - qpos = self.robot.get_qpos(name=control_part)[0] + qpos = self.robot.get_qpos(name=control_part) else: qpos = torch.as_tensor(start_qpos, dtype=torch.float32, device=self.device) if qpos.dim() == 1: @@ -446,24 +533,26 @@ def _build_obs( waypoint_pos: torch.Tensor, waypoint_quat: torch.Tensor, valid_mask: torch.Tensor, - active_idx: int, + active_idx: torch.Tensor, last_action: torch.Tensor, ) -> torch.Tensor: - active_idx_clamped = min(active_idx, self._num_waypoints - 1) - active_onehot = torch.zeros(1, self._num_waypoints, device=self.device) - active_onehot[0, active_idx_clamped] = 1.0 + b = joint_pos.shape[0] + active_idx_clamped = torch.clamp(active_idx, max=self._num_waypoints - 1) + active_onehot = torch.zeros(b, self._num_waypoints, device=self.device) + active_onehot.scatter_(1, active_idx_clamped.unsqueeze(1), 1.0) obs_parts = [ joint_pos, ee_pose, - waypoint_pos.reshape(1, self._num_waypoints * 3), - waypoint_quat.reshape(1, self._num_waypoints * 4), + waypoint_pos.reshape(b, self._num_waypoints * 3), + waypoint_quat.reshape(b, self._num_waypoints * 4), active_onehot, valid_mask, last_action, ] if self._use_relative_obs: - active_pos = waypoint_pos[:, active_idx_clamped] - active_quat = waypoint_quat[:, active_idx_clamped] + idx = torch.arange(b, device=self.device) + active_pos = waypoint_pos[idx, active_idx_clamped] + active_quat = waypoint_quat[idx, active_idx_clamped] obs_parts.append( torch.cat([active_pos - ee_pose[:, :3], active_quat], dim=-1) ) @@ -479,22 +568,86 @@ def _is_active_reached( ee_pose: torch.Tensor, waypoint_pos: torch.Tensor, waypoint_quat: torch.Tensor, - active_idx: int, + active_idx: torch.Tensor, episode_k: int, - ) -> bool: - active_pos = waypoint_pos[:, active_idx] - active_quat_xyzw = waypoint_quat[:, active_idx] + ) -> torch.Tensor: + b = ee_pose.shape[0] + idx = torch.arange(b, device=self.device) + active_idx_clamped = torch.clamp(active_idx, max=self._num_waypoints - 1) + active_pos = waypoint_pos[idx, active_idx_clamped] + active_quat_xyzw = waypoint_quat[idx, active_idx_clamped] pos_dist = (ee_pose[:, :3] - active_pos).norm(dim=-1) ee_quat_wxyz = convert_quat(ee_pose[:, 3:7], to="wxyz") active_quat_wxyz = convert_quat(active_quat_xyzw, to="wxyz") rot_dist = quat_error_magnitude(ee_quat_wxyz, active_quat_wxyz) - orientation_required = ( - self._intermediate_orientation or active_idx >= episode_k - 1 + orientation_required = self._intermediate_orientation | ( + active_idx >= episode_k - 1 ) - rot_ok = ( - (rot_dist < self._rot_eps) - if orientation_required - else torch.ones_like(rot_dist, dtype=torch.bool) + rot_ok = torch.where( + orientation_required, + rot_dist < self._rot_eps, + torch.ones_like(rot_dist, dtype=torch.bool), ) reached = (pos_dist < self._pos_eps) & rot_ok - return bool(reached.item()) + return reached + + @staticmethod + def _compute_vel_acc_via_finite_diff( + positions: torch.Tensor, dt: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + r"""Estimate velocities and accelerations via finite differences. + + Uses a second-order central difference for interior points and one-sided + differences at the boundaries. The estimates are approximate because the + neural policy does not produce velocities or accelerations directly. + + Args: + positions: Joint positions of shape ``(B, N, DOF)``. + dt: Per-point time deltas of shape ``(B, N)``. ``dt[:, t]`` is the + interval used to reach point ``t`` from point ``t - 1``. + + Returns: + Tuple of ``(velocities, accelerations)``, each of shape + ``(B, N, DOF)``. + """ + b, n, dof = positions.shape + if n == 1: + zeros = torch.zeros_like(positions) + return zeros, zeros + + # Forward difference for the first point: (p[1] - p[0]) / dt[1] + v_first = (positions[:, 1] - positions[:, 0]) / dt[:, 1].unsqueeze(-1) + # Backward difference for the last point: (p[N-1] - p[N-2]) / dt[N-1] + v_last = (positions[:, -1] - positions[:, -2]) / dt[:, -1].unsqueeze(-1) + + if n == 2: + velocities = torch.stack([v_first, v_last], dim=1) + return velocities, torch.zeros_like(velocities) + + # Central difference for interior points: + # (p[i+1] - p[i-1]) / (dt[i] + dt[i+1]) + p_next = positions[:, 2:] + p_prev = positions[:, :-2] + dt_sum = (dt[:, 1:-1] + dt[:, 2:]).unsqueeze(-1) + v_interior = (p_next - p_prev) / dt_sum.clamp_min(1e-12) + velocities = torch.cat( + [v_first.unsqueeze(1), v_interior, v_last.unsqueeze(1)], dim=1 + ) + + # Acceleration via second-order finite differences. + # Boundary points use a one-sided stencil; interior points use + # (p[i+1] - 2*p[i] + p[i-1]) / dt[i]^2 + a_first = (positions[:, 2] - 2.0 * positions[:, 1] + positions[:, 0]) / ( + dt[:, 1].unsqueeze(-1) ** 2 + ) + a_last = (positions[:, -1] - 2.0 * positions[:, -2] + positions[:, -3]) / ( + dt[:, -1].unsqueeze(-1) ** 2 + ) + a_interior = ( + positions[:, 2:] - 2.0 * positions[:, 1:-1] + positions[:, :-2] + ) / (dt[:, 1:-1].unsqueeze(-1) ** 2) + accelerations = torch.cat( + [a_first.unsqueeze(1), a_interior, a_last.unsqueeze(1)], dim=1 + ) + + return velocities, accelerations diff --git a/embodichain/lab/sim/planners/toppra_planner.py b/embodichain/lab/sim/planners/toppra_planner.py index 218d17ede..633a9522e 100644 --- a/embodichain/lab/sim/planners/toppra_planner.py +++ b/embodichain/lab/sim/planners/toppra_planner.py @@ -14,8 +14,11 @@ # limitations under the License. # ---------------------------------------------------------------------------- +import os + import torch import numpy as np +from concurrent.futures.process import BrokenProcessPool from embodichain.utils import logger, configclass from embodichain.lab.sim.planners.utils import TrajectorySampleMethod @@ -24,6 +27,7 @@ BasePlanner, BasePlannerCfg, PlanOptions, + _infer_batch_size, ) from .utils import PlanState, PlanResult @@ -38,6 +42,194 @@ ta.setup_logging(level="WARN") +def _build_constraint_arrays(value, acc, dofs: int) -> tuple[np.ndarray, np.ndarray]: + """Expand scalar limits to (dofs, 2) arrays; pass through arrays as-is.""" + if isinstance(value, (float, int)): + vlims = np.array([[-value, value] for _ in range(dofs)]) + else: + vlims = np.array(value) + if isinstance(acc, (float, int)): + alims = np.array([[-acc, acc] for _ in range(dofs)]) + else: + alims = np.array(acc) + return vlims, alims + + +def _toppra_solve_one_env( + waypoints: np.ndarray, + vel_constraint, + acc_constraint, + sample_method: "TrajectorySampleMethod", + sample_interval: float | int, +) -> dict: + """Solve a single-env TOPPRA trajectory. Pure numpy/scipy — picklable, no torch/robot. + + Args: + waypoints: ``(N, DOF)`` numpy array of joint waypoints. + vel_constraint / acc_constraint: scalar or per-DoF array limits. + sample_method: TIME or QUANTITY. + sample_interval: seconds (TIME) or sample count (QUANTITY). + + Returns: + dict with ``positions`` ``(N_b, DOF)``, ``velocities``, ``accelerations``, + ``dt`` ``(N_b,)``, ``success`` bool, ``n`` int, ``duration`` float. + """ + dofs = waypoints.shape[1] + vlims, alims = _build_constraint_arrays(vel_constraint, acc_constraint, dofs) + + if sample_method == TrajectorySampleMethod.TIME and sample_interval <= 0: + return _empty_failure(dofs) + if sample_method == TrajectorySampleMethod.QUANTITY and sample_interval < 2: + return _empty_failure(dofs) + + # Remove consecutive duplicate waypoints. Long plateaus of identical points + # (e.g. when start_qpos equals the first target and joint-space interpolation + # is enabled) can make TOPPRA's controllable-set computation numerically + # ill-conditioned and fail with "Instance is not controllable". + dup_tol = 1e-6 + keep = [0] + for i in range(1, len(waypoints)): + if np.max(np.abs(waypoints[i] - waypoints[keep[-1]])) >= dup_tol: + keep.append(i) + if keep[-1] != len(waypoints) - 1: + keep.append(len(waypoints) - 1) + waypoints = waypoints[keep] + + # Trivial same-waypoint shortcut + if len(waypoints) == 2 and np.sum(np.abs(waypoints[1] - waypoints[0])) < 1e-3: + pos = np.stack([waypoints[0], waypoints[1]]) + return { + "positions": pos, + "velocities": np.zeros_like(pos), + "accelerations": np.zeros_like(pos), + "dt": np.array([0.0, 0.0], dtype=np.float32), + "success": True, + "n": 2, + "duration": 0.0, + } + + ss = np.linspace(0.0, 1.0, len(waypoints)) + try: + path = ta.SplineInterpolator(ss, waypoints) + pc_vel = constraint.JointVelocityConstraint(vlims) + pc_acc = constraint.JointAccelerationConstraint(alims) + instance = ta.algorithm.TOPPRA( + [pc_vel, pc_acc], + path, + parametrizer="ParametrizeConstAccel", + gridpt_min_nb_points=max(100, 10 * len(waypoints)), + ) + jnt_traj = instance.compute_trajectory() + except Exception: + return _empty_failure(dofs) + + if jnt_traj is None: + return _empty_failure(dofs) + + duration = float(jnt_traj.duration) + if duration <= 0: + return _empty_failure(dofs) + + if sample_method == TrajectorySampleMethod.TIME: + n_points = max(2, int(np.ceil(duration / sample_interval)) + 1) + ts = np.linspace(0.0, duration, n_points) + else: + ts = np.linspace(0.0, duration, num=int(sample_interval)) + + positions = np.array([jnt_traj.eval(t) for t in ts]) + velocities = np.array([jnt_traj.evald(t) for t in ts]) + accelerations = np.array([jnt_traj.evaldd(t) for t in ts]) + dt = np.diff(ts, prepend=0.0).astype(np.float32) + return { + "positions": positions, + "velocities": velocities, + "accelerations": accelerations, + "dt": dt, + "success": True, + "n": len(ts), + "duration": duration, + } + + +def _empty_failure(dofs: int) -> dict: + z = np.zeros((2, dofs), dtype=np.float32) + return { + "positions": z, + "velocities": np.zeros_like(z), + "accelerations": np.zeros_like(z), + "dt": np.array([0.0, 0.0], dtype=np.float32), + "success": False, + "n": 2, + "duration": 0.0, + } + + +def _set_parent_death_signal() -> None: + r"""Best-effort: ask the kernel to SIGKILL this worker when its parent dies. + + This is the **only** cleanup mechanism that survives ``os._exit(0)`` in the + parent process, which is what :meth:`SimulationManager.destroy` does by + default (gated by ``EMBODICHAIN_SIM_EXIT_PROCESS``). ``os._exit`` skips + every Python-level finalizer — ``atexit`` handlers, ``__del__`` methods, + and ``concurrent.futures``' internal ``_python_exit`` that would otherwise + join/terminate daemon workers — so the worker processes would be orphaned + and reparented to init, leaving residual python processes behind. + + A kernel parent-death signal fires regardless of *how* the parent exits + (``os._exit``, a crash, or a normal return), so workers are reaped even + when no Python cleanup runs. No-op on non-Linux or if ``prctl`` is + unavailable; in that case workers still rely on ``__del__`` and the + daemon-process reaping that runs on a normal interpreter shutdown. + """ + try: + import ctypes + import signal + + libc = ctypes.CDLL(None) + PR_SET_PDEATHSIG = 1 + libc.prctl.argtypes = [ + ctypes.c_int, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ] + libc.prctl.restype = ctypes.c_int + libc.prctl(PR_SET_PDEATHSIG, signal.SIGKILL, 0, 0, 0) + except Exception: + pass + + +def _worker_init() -> None: + r"""Initializer run in each worker when the pool starts. + + Two responsibilities: + + 1. **Parent-death signal (Linux).** Install ``prctl(PR_SET_PDEATHSIG, + SIGKILL)`` so the kernel kills this worker the instant its parent + dies. This is what guarantees no residual worker processes when the + parent exits via ``os._exit(0)`` (see :func:`_set_parent_death_signal`). + Because of this, callers never need to explicitly shut the planner + down — workers self-terminate with the parent. + + 2. **atexit clearing (fork only).** When ``fork`` is used, the parent + has usually already initialized CUDA/GPU by the time the pool is + created, and ``fork`` copies the parent's atexit registry into the + child. Without clearing it the worker would try to run the parent's + cleanup routines on exit, which can deadlock or corrupt state. With + ``spawn`` this is not necessary, but the initializer still runs so + that ``fork`` remains usable for advanced callers who create the pool + before CUDA is initialized. + """ + import atexit + import multiprocessing as mp + + if mp.get_start_method() == "fork": + atexit._clear() + + _set_parent_death_signal() + + __all__ = ["ToppraPlanner", "ToppraPlannerCfg", "ToppraPlanOptions"] @@ -45,6 +237,24 @@ class ToppraPlannerCfg(BasePlannerCfg): planner_type: str = "toppra" + max_workers: int | None = None + """Worker process count for the batched fan-out. None => min(cpu_count()//2, B).""" + mp_context: str | None = None + """Multiprocessing start method for the batched fan-out. + + ``None`` (default) auto-selects based on the simulation device: + ``'fork'`` on CPU and ``'spawn'`` on GPU. ``'fork'`` is faster — workers + inherit the parent's already-loaded modules, so pool startup is + near-instant — and is safe here because the TOPPRA worker + (:func:`_toppra_solve_one_env`) is pure numpy/scipy and never touches the + parent's Vulkan/Warp/CUDA context or render threads; ``_worker_init`` + clears the inherited atexit registry and installs ``prctl(PR_SET_PDEATHSIG)`` + so workers are reaped when the parent dies (incl. the ``os._exit`` path). + ``'spawn'`` is the safer choice when the parent has initialized CUDA + physics (``device='cuda'``) — fork-after-CUDA-init is the officially + unsupported case — or if fork deadlocks are observed, at the cost of + re-importing modules per worker. + """ @configclass @@ -54,21 +264,21 @@ class ToppraPlanOptions(PlanOptions): "velocity": 0.2, "acceleration": 0.5, } - """Constraints for the planner, including velocity and acceleration limits. + """Constraints for the planner, including velocity and acceleration limits. Should be a dictionary with keys 'velocity' and 'acceleration', each containing a value or a list of limits for each joint. """ sample_method: TrajectorySampleMethod = TrajectorySampleMethod.QUANTITY - """Method for sampling the trajectory. - + """Method for sampling the trajectory. + Options are 'time' for uniform time intervals or 'quantity' for a fixed number of samples. """ sample_interval: float | int = 0.01 - """Interval for sampling the trajectory. - - If sample_method is 'time', this is the time interval in seconds. + """Interval for sampling the trajectory. + + If sample_method is 'time', this is the time interval in seconds. If sample_method is 'quantity', this is the total number of samples. """ @@ -85,11 +295,106 @@ def __init__(self, cfg: ToppraPlannerCfg): """ super().__init__(cfg) - if self.robot.num_instances > 1: - logger.log_error( - "ToppraPlanner does not support multiple robot instances", - NotImplementedError, - ) + self._pool = None + # Resolve the multiprocessing start method once, now that self.device + # (from BasePlanner) is known. None => auto: fork on CPU, spawn on GPU. + self._mp_context = self._resolve_mp_context(cfg.mp_context, self.device) + # No atexit / __del__-based shutdown is registered here: workers install + # prctl(PR_SET_PDEATHSIG) in _worker_init, so the kernel reaps them the + # moment the parent process dies — including the os._exit(0) path taken + # by SimulationManager.destroy(), which skips every Python finalizer. + # __del__ below only handles in-process GC of an abandoned planner. + + @staticmethod + def _resolve_mp_context(mp_context: str | None, device: torch.device) -> str: + """Return the multiprocessing start method to use for the worker pool. + + An explicit ``mp_context`` is honored as-is. ``None`` auto-selects: + ``'fork'`` on CPU (fast — workers inherit loaded modules; safe because + the worker is pure numpy/scipy), ``'spawn'`` everywhere else (safer + under an initialized CUDA context). + + Args: + mp_context: The cfg value; ``None`` means auto-select. + device: The physics device the planner's robot runs on. + + Returns: + One of ``'fork'`` / ``'spawn'``. + """ + if mp_context is not None: + return mp_context + return "fork" if device.type == "cpu" else "spawn" + + def _get_pool(self, batch_size: int): + if self._pool is not None: + return self._pool + import multiprocessing as mp + + max_workers = self.cfg.max_workers + if max_workers is None: + max_workers = max(1, min((os.cpu_count() or 2) // 2, batch_size)) + ctx = mp.get_context(self._mp_context) + from concurrent.futures import ProcessPoolExecutor + + self._pool = ProcessPoolExecutor( + max_workers=max_workers, + mp_context=ctx, + initializer=_worker_init, + ) + return self._pool + + def _shutdown_pool(self) -> None: + r"""Shut down the TOPPRA worker process pool (internal). + + We do **not** use ``shutdown(wait=True)``: with ``fork`` workers can + inherit the parent's CUDA/GPU context, causing them to deadlock inside + the driver at exit. ``wait=True`` would then hang the main process, + and if the user kills it the workers are left behind as residual + python processes. + + Instead we cancel pending work and then forcibly terminate/join/kill + every worker process so that this method returns only after all + workers are actually gone. + + This is **not** part of the public API. It exists for two internal + callers: the ``BrokenProcessPool`` recovery path in :meth:`plan`, and + ``__del__`` (so an abandoned planner in a long-running process does + not leak workers). Normal process exit does not rely on it — workers + install ``prctl(PR_SET_PDEATHSIG)`` in :func:`_worker_init` and are + reaped by the kernel when the parent dies, even under ``os._exit``. + """ + if self._pool is None: + return + + # Capture the worker process objects before shutdown clears them. + worker_processes = list(getattr(self._pool, "_processes", {}).values()) + + # Stop accepting new work and cancel any futures that have not started. + self._pool.shutdown(wait=False, cancel_futures=True) + + # Forcibly reap every worker. This is required for fork-based pools + # when the parent has initialized CUDA: the children inherit the + # context and may not exit cleanly on their own. + for proc in worker_processes: + if not proc.is_alive(): + continue + proc.terminate() + proc.join(timeout=5.0) + if proc.is_alive(): + proc.kill() + proc.join(timeout=1.0) + + self._pool = None + + def __del__(self): + # Only matters for in-process GC of an abandoned planner (and as a + # non-Linux fallback). Process-exit cleanup is handled by the kernel + # via PR_SET_PDEATHSIG installed in each worker, which survives the + # os._exit(0) path that SimulationManager.destroy() takes. + try: + self._shutdown_pool() + except Exception: + pass @validate_plan_options(options_cls=ToppraPlanOptions) def plan( @@ -100,167 +405,125 @@ def plan( r"""Execute trajectory planning. Args: - target_states: List of dictionaries containing target states - cfg: ToppraPlanOptions + target_states: list of :class:`PlanState` waypoints. Tensor fields + carry a leading batch dim ``B``: ``qpos`` is ``(B, DOF)``. + options: :class:`ToppraPlanOptions` with constraints and sampling. Returns: - PlanResult containing the planned trajectory details. + PlanResult containing the planned trajectory details. All tensor + fields are env-batched with leading dim ``B``: ``success`` ``(B,)``, + ``positions``/``velocities``/``accelerations`` ``(B, N, DOF)``, + ``dt`` ``(B, N)``, ``duration`` ``(B,)``. """ - - for i, target in enumerate(target_states): - if target.qpos is None: - logger.log_error(f"Target state at index {i} missing qpos") - - dofs = len(target_states[0].qpos) - - # set constraints - if isinstance(options.constraints["velocity"], float): - vlims = np.array( - [ - [ - -options.constraints["velocity"], - options.constraints["velocity"], - ] - for _ in range(dofs) - ] - ) - else: - vlims = np.array(options.constraints["velocity"]) - - if isinstance(options.constraints["acceleration"], float): - alims = np.array( - [ - [ - -options.constraints["acceleration"], - options.constraints["acceleration"], - ] - for _ in range(dofs) - ] - ) + for i, t in enumerate(target_states): + if t.qpos is None: + logger.log_error(f"Target state at index {i} missing qpos", ValueError) + + b = _infer_batch_size(target_states) or 1 + dofs = target_states[0].qpos.shape[-1] + + # Build (B, N, DOF) numpy waypoints + waypoints = np.stack( + [s.qpos.detach().cpu().numpy() for s in target_states], axis=1 + ) # (B, N, DOF) + + vc = options.constraints["velocity"] + ac = options.constraints["acceleration"] + args_per_env = [ + (waypoints[i], vc, ac, options.sample_method, options.sample_interval) + for i in range(b) + ] + + # Single-env planning never needs a process pool. + if b == 1: + results = [_toppra_solve_one_env(*a) for a in args_per_env] else: - alims = np.array(options.constraints["acceleration"]) - - sample_method = options.sample_method - sample_interval = options.sample_interval - if sample_method == TrajectorySampleMethod.TIME and sample_interval <= 0: - logger.log_error("Time interval must be positive", ValueError) - if sample_method == TrajectorySampleMethod.TIME and isinstance( - sample_interval, int - ): - logger.log_error( - "Time interval must be a float when sample_method is TIME", TypeError + # Inline fallback for max_workers==1 or a single-core machine. + max_workers = self.cfg.max_workers + use_inline = (max_workers == 1) or ( + max_workers is None and ((os.cpu_count() or 2) // 2) <= 1 ) - if sample_method == TrajectorySampleMethod.QUANTITY and sample_interval < 2: - logger.log_error("At least 2 sample points required", ValueError) - if sample_method == TrajectorySampleMethod.QUANTITY and isinstance( - sample_interval, float - ): - logger.log_error( - "Number of samples must be an integer when sample_method is QUANTITY", - TypeError, - ) - - # Check waypoints - for i, target in enumerate(target_states): - if target.qpos is None: - logger.log_error(f"Target state at index {i} missing qpos") - if len(target.qpos) != dofs: - logger.log_error(f"Target waypoints do not align at index {i}") - - if ( - len(target_states) == 2 - and torch.sum(torch.abs(target_states[1].qpos - target_states[0].qpos)) - < 1e-3 - ): - logger.log_warning("Only two same waypoints, returning trivial trajectory.") - return PlanResult( - success=True, - positions=torch.as_tensor( - np.stack([target_states[0].qpos, target_states[1].qpos]), - dtype=torch.float32, - device=self.device, - ), - velocities=torch.zeros( - (2, dofs), dtype=torch.float32, device=self.device - ), - accelerations=torch.zeros( - (2, dofs), dtype=torch.float32, device=self.device - ), - dt=torch.tensor([0.0, 0.0], dtype=torch.float32, device=self.device), - duration=0.0, - ) - - # Build waypoints - waypoints = np.array( - [target.qpos.to("cpu").numpy() for target in target_states] - ) - # Create spline interpolation - # NOTE: Suitable for dense waypoints - ss = np.linspace(0, 1, len(waypoints)) - - # NOTE: Suitable for sparse waypoints; for dense waypoints, CubicSpline may fail strict monotonicity requirement - # len_total = 0 - # len_from_start = [0] - # for i in range(len(waypoints)-1): - # len_total += np.sum(np.abs(waypoints[i+1] - waypoints[i])) - # len_from_start.append(len_total) - # ss = np.array([cur/len_total for cur in len_from_start]) - - path = ta.SplineInterpolator(ss, waypoints) - - # Set constraints - pc_vel = constraint.JointVelocityConstraint(vlims) - pc_acc = constraint.JointAccelerationConstraint(alims) - - # Create TOPPRA instance - instance = ta.algorithm.TOPPRA( - [pc_vel, pc_acc], - path, - parametrizer="ParametrizeConstAccel", - gridpt_min_nb_points=max(100, 10 * len(waypoints)), - ) - # NOTES: Important to set a large number of grid points for better performance in dense waypoint scenarios. - - # Compute parameterized trajectory - jnt_traj = instance.compute_trajectory() - if jnt_traj is None: - # raise RuntimeError("Unable to find feasible trajectory") - logger.log_warning("Unable to find feasible trajectory") - return PlanResult(success=False) - - duration = jnt_traj.duration - # Sample trajectory points - if duration <= 0: - logger.log_error(f"Duration must be positive, got {duration}", ValueError) - if sample_method == TrajectorySampleMethod.TIME: - n_points = max(2, int(np.ceil(duration / sample_interval)) + 1) - ts = np.linspace(0, duration, n_points) - else: - ts = np.linspace(0, duration, num=int(sample_interval)) - - positions = [] - velocities = [] - accelerations = [] - - for t in ts: - positions.append(jnt_traj.eval(t)) - velocities.append(jnt_traj.evald(t)) - accelerations.append(jnt_traj.evaldd(t)) + if use_inline: + results = [_toppra_solve_one_env(*a) for a in args_per_env] + else: + pool = self._get_pool(b) + results = [None] * b + try: + futures = [ + pool.submit(_toppra_solve_one_env, *a) for a in args_per_env + ] + broken = False + for i, fut in enumerate(futures): + try: + results[i] = fut.result() + except BrokenProcessPool: + logger.log_warning( + "TOPPRA process pool broke; returning failure." + ) + self._shutdown_pool() + broken = True + break + except Exception: + results[i] = _empty_failure(dofs) + if broken: + for i in range(b): + if results[i] is None: + results[i] = _empty_failure(dofs) + except BrokenProcessPool: + # pool was already broken at submit time + logger.log_warning("TOPPRA process pool broke; returning failure.") + self._shutdown_pool() + for i in range(b): + if results[i] is None: + results[i] = _empty_failure(dofs) + + return self._assemble_batched_result(results, dofs) + + def _assemble_batched_result(self, results: list[dict], dofs: int) -> PlanResult: + """Stack per-env TOPPRA results into a batched :class:`PlanResult`. + + Each entry of ``results`` is the dict returned by + :func:`_toppra_solve_one_env`. Env trajectories may have different + lengths (``n``); this method pads shorter trajectories out to the + longest by repeating their final waypoint (held pose) with zero + velocity and acceleration, so every output tensor shares the same + ``(B, N, DOF)`` / ``(B, N)`` shape. - dt = torch.as_tensor(ts, dtype=torch.float32, device=self.device) - dt = torch.diff(dt, prepend=torch.tensor([0.0], device=self.device)) + Args: + results: list of per-env result dicts (length ``B``). + dofs: per-env degrees of freedom. + Returns: + PlanResult with env-batched tensors (``success`` ``(B,)``, + ``positions``/``velocities``/``accelerations`` ``(B, N, DOF)``, + ``dt`` ``(B, N)``, ``duration`` ``(B,)``). + """ + b = len(results) + max_n = max(r["n"] for r in results) + positions = np.zeros((b, max_n, dofs), dtype=np.float32) + velocities = np.zeros((b, max_n, dofs), dtype=np.float32) + accelerations = np.zeros((b, max_n, dofs), dtype=np.float32) + dt = np.zeros((b, max_n), dtype=np.float32) + duration = np.zeros((b,), dtype=np.float32) + success = np.zeros((b,), dtype=bool) + for i, r in enumerate(results): + n = r["n"] + positions[i, :n] = r["positions"] + velocities[i, :n] = r["velocities"] + accelerations[i, :n] = r["accelerations"] + dt[i, :n] = r["dt"] + duration[i] = r["duration"] + success[i] = r["success"] + # tail-pad: repeat final waypoint for held-pose rows + if n < max_n: + positions[i, n:] = r["positions"][-1] + velocities[i, n:] = 0.0 + accelerations[i, n:] = 0.0 return PlanResult( - success=True, - positions=torch.as_tensor( - np.array(positions), dtype=torch.float32, device=self.device - ), - velocities=torch.as_tensor( - np.array(velocities), dtype=torch.float32, device=self.device - ), - accelerations=torch.as_tensor( - np.array(accelerations), dtype=torch.float32, device=self.device - ), - dt=dt, - duration=duration, + success=torch.as_tensor(success, device=self.device), + positions=torch.as_tensor(positions, device=self.device), + velocities=torch.as_tensor(velocities, device=self.device), + accelerations=torch.as_tensor(accelerations, device=self.device), + dt=torch.as_tensor(dt, device=self.device), + duration=torch.as_tensor(duration, device=self.device), ) diff --git a/embodichain/lab/sim/planners/utils.py b/embodichain/lab/sim/planners/utils.py index cfeee4437..020142248 100644 --- a/embodichain/lab/sim/planners/utils.py +++ b/embodichain/lab/sim/planners/utils.py @@ -31,6 +31,7 @@ "PlanResult", "calculate_point_allocations", "interpolate_xpos", + "interpolate_xpos_batched", ] @@ -114,33 +115,44 @@ class MoveType(Enum): @dataclass class PlanResult: - r"""Data class representing the result of a motion plan.""" + r"""Data class representing the result of a motion plan (env-batched).""" success: bool | torch.Tensor = False - """Whether planning succeeded.""" + """Per-env success, shape ``(B,)`` bool tensor (or scalar bool).""" xpos_list: torch.Tensor | None = None - """End-effector poses along trajectory with shape `(N, 4, 4)`.""" + """End-effector poses, shape ``(B, N, 4, 4)``.""" positions: torch.Tensor | None = None - """Joint positions along trajectory with shape `(N, DOF)`.""" + """Joint positions, shape ``(B, N, DOF)``.""" velocities: torch.Tensor | None = None - """Joint velocities along trajectory with shape `(N, DOF)`.""" + """Joint velocities, shape ``(B, N, DOF)``.""" accelerations: torch.Tensor | None = None - """Joint accelerations along trajectory with shape `(N, DOF)`.""" + """Joint accelerations, shape ``(B, N, DOF)``.""" dt: torch.Tensor | None = None - """Time duration between each point with shape `(N,)`.""" + """Per-env time deltas, shape ``(B, N)``.""" duration: float | torch.Tensor = 0.0 - """Total trajectory duration in seconds.""" + """Per-env total duration, shape ``(B,)``.""" + + def is_all_success(self) -> bool: + """Return True only when every env succeeded.""" + if isinstance(self.success, torch.Tensor): + return bool(torch.all(self.success).item()) + return bool(self.success) @dataclass class PlanState: - r"""Data class representing the state for a motion plan.""" + r"""Data class representing the state for a motion plan (env-batched). + + Tensor fields carry a leading batch dim ``B``: ``qpos:(B, DOF)``, + ``xpos:(B, 4, 4)``. Enum/scalar fields are shared across ``B`` (vectorized + envs share the same task skeleton). + """ move_type: MoveType = MoveType.JOINT_MOVE """Type of movement used by the plan.""" @@ -149,25 +161,72 @@ class PlanState: """Robot part that should move.""" xpos: torch.Tensor | None = None - """Target TCP pose (4x4 matrix) for `MoveType.EEF_MOVE`.""" + """Target TCP pose (Bx4x4) for ``MoveType.EEF_MOVE``.""" qpos: torch.Tensor | None = None - """Target joint angles for `MoveType.JOINT_MOVE` with shape `(DOF,)`.""" + """Target joint angles for ``MoveType.JOINT_MOVE`` with shape ``(B, DOF)``.""" qvel: torch.Tensor | None = None - """Target joint velocities for `MoveType.JOINT_MOVE` with shape `(DOF,)`.""" + """Target joint velocities for ``MoveType.JOINT_MOVE`` with shape ``(B, DOF)``.""" qacc: torch.Tensor | None = None - """Target joint accelerations for `MoveType.JOINT_MOVE` with shape `(DOF,)`.""" + """Target joint accelerations for ``MoveType.JOINT_MOVE`` with shape ``(B, DOF)``.""" is_open: bool = True - """For `MoveType.TOOL`, indicates whether to open (`True`) or close (`False`) the tool.""" + """For ``MoveType.TOOL``, indicates whether to open (``True``) or close (``False``) the tool.""" is_world_coordinate: bool = True - """`True` if the target pose is in world coordinates, `False` if relative to the current pose.""" + """``True`` if the target pose is in world coordinates, ``False`` if relative to the current pose.""" pause_seconds: float = 0.0 - """Duration of a pause when `move_type` is `MoveType.PAUSE`.""" + """Duration of a pause when ``move_type`` is ``MoveType.PAUSE``.""" + + @classmethod + def from_qpos( + cls, + qpos: torch.Tensor, + *, + move_type: MoveType = MoveType.JOINT_MOVE, + move_part: MovePart = MovePart.LEFT, + **kwargs, + ) -> "PlanState": + """Create a PlanState from batched joint positions ``(B, DOF)``.""" + return cls(move_type=move_type, move_part=move_part, qpos=qpos, **kwargs) + + @classmethod + def from_xpos( + cls, + xpos: torch.Tensor, + *, + move_type: MoveType = MoveType.EEF_MOVE, + move_part: MovePart = MovePart.LEFT, + **kwargs, + ) -> "PlanState": + """Create a PlanState from batched end-effector poses ``(B, 4, 4)``.""" + return cls(move_type=move_type, move_part=move_part, xpos=xpos, **kwargs) + + @classmethod + def single( + cls, + *, + qpos: torch.Tensor | None = None, + xpos: torch.Tensor | None = None, + move_type: MoveType = MoveType.JOINT_MOVE, + move_part: MovePart = MovePart.LEFT, + **kwargs, + ) -> "PlanState": + """B=1 convenience constructor: unsqueezes a single-env qpos/xpos. + + Already-batched tensors (2D qpos / 3D xpos) pass through unchanged + (idempotent). + """ + if qpos is not None and qpos.dim() == 1: + qpos = qpos.unsqueeze(0) + if xpos is not None and xpos.dim() == 2: + xpos = xpos.unsqueeze(0) + return cls( + move_type=move_type, move_part=move_part, qpos=qpos, xpos=xpos, **kwargs + ) def interpolate_xpos( @@ -192,6 +251,38 @@ def interpolate_xpos( return interp_poses +def interpolate_xpos_batched( + start_xpos: torch.Tensor, end_xpos: torch.Tensor, num_samples: int +) -> torch.Tensor: + """Batched pose interpolation. + + Args: + start_xpos: Start poses, shape ``(B, 4, 4)``. + end_xpos: End poses, shape ``(B, 4, 4)``. + num_samples: Number of samples to generate (clamped to at least 2). + + Returns: + Interpolated poses, shape ``(B, num_samples, 4, 4)``. + """ + num_samples = max(2, int(num_samples)) + B = start_xpos.shape[0] + out = torch.eye(4, dtype=start_xpos.dtype, device=start_xpos.device).repeat( + B, num_samples, 1, 1 + ) + # ``scipy.spatial.transform.Slerp`` interpolates a single path, so we loop + # over envs. B is typically small (number of parallel envs), so this is fine. + for b in range(B): + poses = interpolate_xpos( + start_xpos[b].detach().cpu().numpy(), + end_xpos[b].detach().cpu().numpy(), + num_samples, + ) + out[b] = torch.as_tensor( + poses, dtype=start_xpos.dtype, device=start_xpos.device + ) + return out + + def calculate_point_allocations( xpos_list: torch.Tensor | np.ndarray, step_size: float = 0.002, diff --git a/embodichain/lab/sim/robots/__init__.py b/embodichain/lab/sim/robots/__init__.py index be86d163d..085e23db6 100644 --- a/embodichain/lab/sim/robots/__init__.py +++ b/embodichain/lab/sim/robots/__init__.py @@ -16,12 +16,14 @@ from .dexforce_w1 import * from .cobotmagic import CobotMagicCfg +from .franka_panda import FrankaPandaCfg from .ur_robot import URRobotCfg from .dual_arm import DualArmRobotCfg, build_dual_arm_cfg __all__ = [ "DexforceW1Cfg", "CobotMagicCfg", + "FrankaPandaCfg", "URRobotCfg", "DualArmRobotCfg", "build_dual_arm_cfg", diff --git a/embodichain/lab/sim/robots/dexforce_w1/cfg.py b/embodichain/lab/sim/robots/dexforce_w1/cfg.py index b2515f2c0..d87e0dde0 100644 --- a/embodichain/lab/sim/robots/dexforce_w1/cfg.py +++ b/embodichain/lab/sim/robots/dexforce_w1/cfg.py @@ -239,11 +239,16 @@ def _build_default_physics_cfgs( DEFAULT_EEF_GRIPPER_JOINT_NAMES = "(LEFT|RIGHT)_FINGER[1-2]" ARM_JOINTS = "(RIGHT|LEFT)_J[0-9]" BODY_JOINTS = "(ANKLE|KNEE|BUTTOCK|WAIST)" + HEAD_JOINTS = "(NECK1|NECK2)" joint_params = { - "stiffness": {ARM_JOINTS: 1e4, BODY_JOINTS: 1e7}, - "damping": {ARM_JOINTS: 1e3, BODY_JOINTS: 1e4}, - "max_effort": {ARM_JOINTS: 1e5, BODY_JOINTS: 1e10}, + "stiffness": { + ARM_JOINTS: 1e4, + BODY_JOINTS: 1e7, + HEAD_JOINTS: 1e4, + }, + "damping": {ARM_JOINTS: 1e3, BODY_JOINTS: 1e4, HEAD_JOINTS: 1e3}, + "max_effort": {ARM_JOINTS: 1e5, BODY_JOINTS: 1e10, HEAD_JOINTS: 1e5}, } drive_pros = JointDrivePropertiesCfg(**joint_params) diff --git a/embodichain/lab/sim/robots/dexforce_w1/utils.py b/embodichain/lab/sim/robots/dexforce_w1/utils.py index e1e842854..9402b5657 100644 --- a/embodichain/lab/sim/robots/dexforce_w1/utils.py +++ b/embodichain/lab/sim/robots/dexforce_w1/utils.py @@ -516,7 +516,7 @@ def build_dexforce_w1_solver_cfg( ], component_versions: dict[DexforceW1Type, DexforceW1Version] | None = None, urdf_cfg: URDFCfg | None = None, -) -> Dict[DexforceW1Type, SolverCfg]: +) -> Dict[str, SolverCfg]: """ Build DexforceW1 solver configuration dict. @@ -527,7 +527,8 @@ def build_dexforce_w1_solver_cfg( urdf_cfg: Optional, URDFCfg object from build_dexforce_w1_assembly_urdf_cfg. Returns: - Dict[DexforceW1Type, SolverCfg] + Dict[str, SolverCfg]: solver config keyed by control part name + (e.g. ``"left_arm"``, ``"full_body"``). """ def get_version(t, default=DexforceW1Version.V021): @@ -550,7 +551,10 @@ def get_version(t, default=DexforceW1Version.V021): ) arm_version = get_version(arm_type) arm_cfg = arm_manager.get_config(arm_kind, arm_side, arm_version) - solver_cfg[arm_type] = SolverCfg.from_dict( + # Use control_parts-aligned key (e.g. "left_arm") so init_solver + # can match this entry to the corresponding control part. + solver_key = f"{arm_side.value}_arm" + solver_cfg[solver_key] = SolverCfg.from_dict( { "class_type": "PytorchSolver", "urdf_path": arm_cfg["urdf_path"], @@ -560,14 +564,14 @@ def get_version(t, default=DexforceW1Version.V021): } ) - # Use urdf_cfg.fname if provided, otherwise fallback to default path + # Use urdf_cfg.fpath if provided, otherwise fallback to default path full_body_urdf_path = ( - urdf_cfg.fname + urdf_cfg.fpath or get_data_path("DexforceW1FullBodyV021/full_body.urdf") if urdf_cfg is not None else get_data_path("DexforceW1FullBodyV021/full_body.urdf") ) - solver_cfg[DexforceW1Type.FULL_BODY] = SolverCfg.from_dict( + solver_cfg[DexforceW1Type.FULL_BODY.value] = SolverCfg.from_dict( { "class_type": "PytorchSolver", "urdf_path": full_body_urdf_path, diff --git a/embodichain/lab/sim/robots/dual_arm.py b/embodichain/lab/sim/robots/dual_arm.py index 970c1eac1..eaf614b00 100644 --- a/embodichain/lab/sim/robots/dual_arm.py +++ b/embodichain/lab/sim/robots/dual_arm.py @@ -55,6 +55,7 @@ ) from embodichain.lab.sim.solvers import SolverCfg from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg +from embodichain.lab.sim.robots.franka_panda import FrankaPandaCfg from embodichain.lab.sim.robots.ur_robot import URRobotCfg from embodichain.utils import configclass @@ -74,8 +75,9 @@ #: Registry mapping a ``base_robot`` key to ``(single-arm cfg class, init dict)``. #: Adding a new single-arm robot here is all that is needed to make it -#: dual-arm-able: e.g. ``"franka": (FrankaRobotCfg, {})``. +#: dual-arm-able: e.g. ``"franka": (FrankaPandaCfg, {"robot_type": "panda"})``. _BASE_ROBOT_REGISTRY: Dict[str, tuple] = { + "franka": (FrankaPandaCfg, {"robot_type": "panda"}), "ur3": (URRobotCfg, {"robot_type": "ur3"}), "ur3e": (URRobotCfg, {"robot_type": "ur3e"}), "ur5": (URRobotCfg, {"robot_type": "ur5"}), diff --git a/embodichain/lab/sim/robots/franka_panda.py b/embodichain/lab/sim/robots/franka_panda.py new file mode 100644 index 000000000..d66d1e6fd --- /dev/null +++ b/embodichain/lab/sim/robots/franka_panda.py @@ -0,0 +1,213 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Dict + +import numpy as np +import torch + +from embodichain.data import get_data_path +from embodichain.lab.sim.cfg import ( + JointDrivePropertiesCfg, + RigidBodyAttributesCfg, + RobotCfg, + URDFCfg, +) +from embodichain.lab.sim.solvers import PytorchSolverCfg +from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg +from embodichain.utils import configclass + +if TYPE_CHECKING: + import pytorch_kinematics as pk + +__all__ = ["FrankaPandaCfg"] + +# ``robot_type`` -> URDF directory / file name. +_FRANKA_URDF_DIR: Dict[str, str] = { + "panda": "Panda", +} + +# Default init_qpos: arm in a neutral ready pose, fingers open. +# Derived from the Isaac Lab Franka Emika Panda reference configuration. +_FRANKA_DEFAULT_INIT_QPOS = [ + 0.0, # Joint1 + -0.569, # Joint2 + 0.0, # Joint3 + -2.810, # Joint4 + 0.0, # Joint5 + 3.037, # Joint6 + 0.741, # Joint7 + 0.04, # finger_joint1 + 0.04, # finger_joint2 (mimic of finger_joint1) +] + + +@configclass +class FrankaPandaCfg(RobotCfg): + """Configuration for the Franka Emika Panda robot with Panda hand. + + The PandaWithHand URDF includes both the 7-DOF arm and the parallel-jaw + gripper in a single file. The solver defaults to + :class:`~embodichain.lab.sim.solvers.PytorchSolverCfg`. + + Example: + + cfg = FrankaPandaCfg.from_dict({"robot_type": "panda"}) + robot = sim.add_robot(cfg=cfg) + """ + + robot_type: str = "panda" + + @classmethod + def from_dict(cls, init_dict): + """Initialize ``FrankaPandaCfg`` from a dictionary. + + Args: + init_dict: Dictionary of configuration parameters. ``robot_type`` + selects the Franka variant (currently ``"panda"``). All other + keys are merged on top of the defaults via + :func:`merge_robot_cfg`. + + Returns: + A ``FrankaPandaCfg`` instance. + """ + cfg = cls() + cfg._build_defaults(init_dict) + return merge_robot_cfg(cfg, init_dict) + + def _build_defaults(self, init_dict: dict | None = None) -> None: + """Populate default urdf/control/solver/physics for the Franka variant. + + Args: + init_dict: The raw override dict passed to ``from_dict``. + ``robot_type`` is read from here (falling back to the class + default). + """ + init_dict = init_dict or {} + robot_type = init_dict.get("robot_type", self.robot_type) + if robot_type not in _FRANKA_URDF_DIR: + raise ValueError( + f"Unknown Franka robot_type: {robot_type!r}. " + f"Expected one of {sorted(_FRANKA_URDF_DIR)}." + ) + + self.robot_type = robot_type + self.uid = "FrankaPanda" + + urdf_dir = _FRANKA_URDF_DIR[robot_type] + urdf_path = get_data_path(f"Franka/{urdf_dir}/{urdf_dir}WithHand.urdf") + + self.urdf_cfg = URDFCfg( + components=[ + { + "component_type": "arm", + "urdf_path": urdf_path, + "transform": np.eye(4), + } + ] + ) + + self.control_parts = { + "arm": [f"fr3_joint{i}" for i in range(1, 8)], + "hand": ["fr3_finger_joint1", "fr3_finger_joint2"], + } + + self.solver_cfg = { + "arm": PytorchSolverCfg( + end_link_name="fr3_hand_tcp", + root_link_name="base", + tcp=[ + [1, 0, 0.0, 0.0], + [0, 1, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ], + num_samples=30, + ), + } + + self.drive_pros = JointDrivePropertiesCfg( + stiffness={ + "fr3_joint[1-7]": 1e4, + "fr3_finger_joint[1-2]": 1e3, + }, + damping={ + "fr3_joint[1-7]": 1e3, + "fr3_finger_joint[1-2]": 1e2, + }, + max_effort={ + "fr3_joint[1-7]": 1e5, + "fr3_finger_joint[1-2]": 1e4, + }, + ) + + self.init_qpos = list(_FRANKA_DEFAULT_INIT_QPOS) + + @property + def _pk_urdf_path(self) -> str: + """URDF used for the FK/IK serial chain.""" + urdf_dir = _FRANKA_URDF_DIR[self.robot_type] + return get_data_path(f"Franka/{urdf_dir}/{urdf_dir}WithHand.urdf") + + def build_pk_serial_chain( + self, device: torch.device = torch.device("cpu"), **kwargs + ) -> Dict[str, "pk.SerialChain"]: + """Build the pytorch-kinematics serial chain for the arm. + + Args: + device: The device to which the chain will be moved. Defaults to CPU. + **kwargs: Additional arguments for building the serial chain. + + Returns: + A ``{"arm": pk.SerialChain}`` mapping. + """ + from embodichain.lab.sim.utility.solver_utils import create_pk_serial_chain + + chain = create_pk_serial_chain( + urdf_path=self._pk_urdf_path, + device=device, + end_link_name="fr3_hand_tcp", + root_link_name="base", + ) + return {"arm": chain} + + +if __name__ == "__main__": + np.set_printoptions(precision=5, suppress=True) + + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg + from embodichain.lab.sim.cfg import RenderCfg + + config = SimulationManagerCfg( + headless=False, + device="cpu", + num_envs=1, + render_cfg=RenderCfg(renderer="hybrid"), + ) + sim = SimulationManager(config) + + cfg = FrankaPandaCfg.from_dict({"robot_type": "panda"}) + robot = sim.add_robot(cfg=cfg) + sim.open_window() + + if sim.is_use_gpu_physics: + sim.init_gpu_physics() + + from IPython import embed + + embed() # noqa: F401 diff --git a/embodichain/lab/sim/shapes.py b/embodichain/lab/sim/shapes.py old mode 100644 new mode 100755 index a804c42d2..08e44f587 --- a/embodichain/lab/sim/shapes.py +++ b/embodichain/lab/sim/shapes.py @@ -80,6 +80,8 @@ def from_dict(cls, init_dict: Dict[str, Any]) -> ShapeCfg: if hasattr(cfg, key): attr = getattr(cfg, key) if key == "visual_material" and isinstance(value, dict): + from embodichain.lab.sim.material import VisualMaterialCfg + setattr( cfg, key, @@ -117,6 +119,30 @@ class MeshCfg(ShapeCfg): project_direction: List[float] = [1.0, 1.0, 1.0] """Direction to project the UV coordinates. Defaults to [1.0, 1.0, 1.0].""" + max_convex_hull_num: int = 1 + """The maximum number of convex hulls that will be created for the mesh. + + If set to larger than 1, the mesh will be decomposed into multiple convex hulls + using the approximate convex decomposition method specified by :attr:`acd_method`. + Reference: https://github.com/SarahWeiii/CoACD + """ + + acd_method: str = "coacd" + """The method used for approximate convex decomposition (ACD) of the mesh. + + Currently, ``"coacd"`` and ``"vhacd"`` are supported. Only used when + :attr:`max_convex_hull_num` is set to larger than 1. + """ + + sdf_resolution: int = 0 + """Resolution for the signed distance field (SDF) of the mesh. + + The spacing of the uniformly sampled SDF is equal to the largest AABB extent + of the mesh, divided by the resolution. If ``sdf_resolution`` is set to larger + than 0, an SDF will be generated for collision detection. SDF increases the + accuracy of collision, but also takes more time to initialize and simulate. + """ + @configclass class CubeCfg(ShapeCfg): diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 8263cb7c3..952d116ad 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -80,6 +80,7 @@ validate_physics_cfg, MarkerCfg, WindowRecordCfg, + WindowCameraPoseCfg, LightCfg, RigidObjectCfg, SoftObjectCfg, @@ -124,6 +125,7 @@ def __init__( device: str | torch.device | None = None, physics_cfg: DefaultPhysicsCfg | NewtonPhysicsCfg | None = None, window_record: WindowRecordCfg | None = None, + window_camera_pose: WindowCameraPoseCfg | None = None, ) -> None: self.width = width self.height = height @@ -138,6 +140,9 @@ def __init__( self.window_record = ( WindowRecordCfg() if window_record is None else window_record ) + self.window_camera_pose = ( + WindowCameraPoseCfg() if window_camera_pose is None else window_camera_pose + ) if physics_dt is not None: self.physics_cfg.physics_dt = physics_dt @@ -196,6 +201,9 @@ def __init__( window_record: WindowRecordCfg = field(default_factory=WindowRecordCfg) """Viewer window recording settings (hotkey, paths, FPS, memory budget).""" + window_camera_pose: WindowCameraPoseCfg = field(default_factory=WindowCameraPoseCfg) + """Interactive viewer camera-pose printing settings.""" + def __post_init__(self): validate_physics_cfg(self.physics_cfg) @@ -328,6 +336,13 @@ def __init__( ) self._window_record_input_control: ObjectManipulator | None = None self._window_record_save_threads: list[threading.Thread] = [] + wcp = sim_config.window_camera_pose + self._window_camera_pose_hotkey_cfg: dict[str, object] | None = ( + {"convert_to_look_at": wcp.convert_to_look_at} + if wcp.enable_hotkey + else None + ) + self._window_camera_pose_input_control: ObjectManipulator | None = None self._world.set_delta_time(sim_config.physics_cfg.physics_dt) self._world.show_coordinate_axis(False) @@ -371,6 +386,7 @@ def __init__( self._create_default_plane() self.set_default_background() + self.set_default_global_lighting() # Set physics to manual update mode by default. self.set_manual_update(True) @@ -767,6 +783,11 @@ def open_window(self) -> None: and self._window_record_input_control is None ): self.enable_window_record_hotkey(**self._window_record_hotkey_cfg) + if ( + self._window_camera_pose_hotkey_cfg is not None + and self._window_camera_pose_input_control is None + ): + self.enable_window_camera_pose_hotkey(**self._window_camera_pose_hotkey_cfg) self.is_window_opened = True def close_window(self) -> None: @@ -776,6 +797,7 @@ def close_window(self) -> None: self._world.close_window() self._window = None self._window_record_input_control = None + self._window_camera_pose_input_control = None self.is_window_opened = False def _build_multiple_arenas(self, num: int, space: float | None = None) -> None: @@ -844,6 +866,28 @@ def _create_default_plane(self): self._default_plane.add_rigidbody(ActorType.STATIC, RigidBodyShape.PLANE, attr) self._invalidate_newton_physics() + def set_default_global_lighting(self) -> None: + """Set default global lighting for the scene. + + Configures both the environment emission (ambient) light and a + directional light to provide default scene illumination. The + directional light is a global scene light (infinite distance) + pointing downward along the -Z axis. + """ + # Environment emission light + self.set_emission_light([1.0, 1.0, 1.0], 120.0) + + # Directional light as global scene light + dir_light_cfg = LightCfg( + uid="default_global_light", + light_type="sun", + intensity=8.0, + direction=(0.0, 0.0, -1.0), + color=(1.0, 0.95, 0.85), + enable_shadow=True, + ) + self.add_light(dir_light_cfg) + def set_default_background(self) -> None: """Set default background.""" @@ -857,12 +901,10 @@ def set_default_background(self) -> None: uid=mat_name, base_color_texture=color_texture, roughness_texture=roughness_texture, - roughness=0.7, + roughness=1.0, ) ) - self.set_emission_light([1.0, 1.0, 1.0], 120.0) - self._default_plane.set_material(mat.get_instance("plane_mat").mat) self._visual_materials[mat_name] = mat @@ -940,14 +982,42 @@ def get_asset( logger.log_warning(f"Asset {uid} not found.") return None + # Light type string → dexsim LightType enum mapping + _LIGHT_TYPE_MAP: dict[str, LightType] = { + "point": LightType.POINT, + "sun": LightType.SUN, + "direction": LightType.DIRECTION, + "spot": LightType.SPOT, + "rect": LightType.RECT, + "mesh": LightType.MESH, + } + + # Light types that are created as a single global scene light (not per-environment). + _GLOBAL_LIGHT_TYPES: tuple[str, ...] = ("sun", "direction") + def add_light(self, cfg: LightCfg) -> Light: """Create a light in the scene. + Supports six light types: ``"point"``, ``"sun"``, ``"direction"``, + ``"spot"``, ``"rect"``, and ``"mesh"``. See :class:`LightCfg` for + type-specific configuration fields. + + .. attention:: + ``"sun"`` and ``"direction"`` lights are global scene lights + (infinite-distance directional light sources). They are created + as a single instance on the root environment, not batched per + environment. All other types are created as per-environment + batched lights. + Args: - cfg (LightCfg): Configuration for the light, including type, color, intensity, and radius. + cfg (LightCfg): Configuration for the light, including type, color, + intensity, and type-specific properties. Returns: Light: The created light instance. + + Raises: + RuntimeError: If ``cfg.light_type`` is not one of the supported types. """ if cfg.uid is None: uid = "light" @@ -958,22 +1028,42 @@ def add_light(self, cfg: LightCfg) -> Light: if uid in self._lights: logger.log_error(f"Light {uid} already exists.") - light_type = cfg.light_type - if light_type == "point": - light_type = LightType.POINT - else: + light_type_str = cfg.light_type + light_type = self._LIGHT_TYPE_MAP.get(light_type_str) + if light_type is None: + supported = ", ".join(self._LIGHT_TYPE_MAP.keys()) logger.log_error( - f"Unsupported light type: {light_type}. Supported types: point." + f"Unsupported light type: '{light_type_str}'. " + f"Supported types: {supported}." ) - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - light_list = [] - for i, env in enumerate(env_list): - light_name = f"{uid}_{i}" - light = env.create_light(light_name, light_type) - light_list.append(light) + # Validation warnings for type-specific constraints + if light_type_str == "mesh" and not cfg.mesh_path: + logger.log_warning( + f"Mesh light '{uid}' has no mesh_path set. " + f"Use set_mesh() to assign a MeshObject." + ) + if light_type_str == "rect" and (cfg.rect_width <= 0 or cfg.rect_height <= 0): + logger.log_warning( + f"Rect light '{uid}' has zero or negative dimensions " + f"(width={cfg.rect_width}, height={cfg.rect_height})." + ) - batch_lights = Light(cfg=cfg, entities=light_list) + if cfg.light_type in self._GLOBAL_LIGHT_TYPES: + # Global scene light: create a single instance on the root + # environment. Infinite-distance lights (sun, direction) are + # physically scene-global and should not be duplicated per arena. + light = self._env.create_light(uid, light_type) + batch_lights = Light(cfg=cfg, entities=[light]) + else: + # Per-environment batched light: one instance per arena. + env_list = [self._env] if len(self._arenas) == 0 else self._arenas + light_list = [] + for i, env in enumerate(env_list): + light_name = f"{uid}_{i}" + light = env.create_light(light_name, light_type) + light_list.append(light) + batch_lights = Light(cfg=cfg, entities=light_list) self._lights[uid] = batch_lights @@ -2491,6 +2581,145 @@ def on_key_down(self, key): ) return True + @staticmethod + def _window_camera_pose_to_look_at( + pose: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Convert a DexSim window model matrix to look-at vectors. + + DexSim stores the viewer camera model matrix with columns + ``[right, up, -forward]``. The local camera up axis changes while the + viewer orbits, but ``Windows.set_look_at`` uses a world-up reference. + Always use DexSim's default Z-up vector so a captured snippet retains + the standard viewer controls. + + Args: + pose: A 4x4 homogeneous viewer camera pose matrix. + + Returns: + The ``(eye, look_at, up)`` vectors accepted by + ``Windows.set_look_at``. + + Raises: + ValueError: If ``pose`` is not a 4x4 homogeneous matrix. + """ + matrix = np.asarray(pose, dtype=np.float64) + if matrix.shape != (4, 4): + raise ValueError( + f"Window camera pose must have shape (4, 4), got {matrix.shape}." + ) + eye = matrix[:3, 3] + look_at = eye - matrix[:3, 2] + up = np.array([0.0, 0.0, 1.0], dtype=np.float64) + return eye, look_at, up + + @staticmethod + def _format_window_camera_pose( + pose: np.ndarray, convert_to_look_at: bool = True + ) -> str: + """Format a DexSim window pose as an executable Python snippet. + + Args: + pose: A 4x4 homogeneous viewer camera pose matrix. + convert_to_look_at: Print a ``set_look_at`` call when true; + otherwise print the raw pose matrix. + + Returns: + An executable Python snippet containing the camera pose. + + Raises: + ValueError: If ``pose`` is not a 4x4 homogeneous matrix. + """ + matrix = np.asarray(pose, dtype=np.float64) + if matrix.shape != (4, 4): + raise ValueError( + f"Window camera pose must have shape (4, 4), got {matrix.shape}." + ) + + def _format_float(value: float) -> str: + if abs(value) < 1e-12: + return "0.0" + formatted = format(value, ".8g") + if "e" not in formatted and "." not in formatted: + formatted += ".0" + return formatted + + def _vector_literal(vector: np.ndarray) -> str: + values = ", ".join(_format_float(float(value)) for value in vector) + return f"np.array([{values}], dtype=np.float32)" + + if convert_to_look_at: + eye, look_at, up = SimulationManager._window_camera_pose_to_look_at(matrix) + return ( + "window.set_look_at(" + f"eye={_vector_literal(eye)}, " + f"look_at={_vector_literal(look_at)}, " + f"up={_vector_literal(up)})" + ) + + rows = ",\n ".join( + "[" + ", ".join(_format_float(float(value)) for value in row) + "]" + for row in matrix + ) + return f"window_pose = np.array([\n {rows}\n], dtype=np.float32)" + + def print_window_camera_pose(self, convert_to_look_at: bool = True) -> str | None: + """Print the current viewer camera pose as reusable Python code. + + Args: + convert_to_look_at: Print ``window.set_look_at(...)`` by default. + Set false to print the raw 4x4 pose matrix instead. + + Returns: + The printed snippet, or ``None`` when no viewer window is open. + """ + if self._window is None: + logger.log_warning("No simulation window available to print its pose.") + return None + + pose = np.asarray(self._window.get_pose_matrix(), dtype=np.float32) + snippet = self._format_window_camera_pose(pose, convert_to_look_at) + print(snippet) + return snippet + + def enable_window_camera_pose_hotkey(self, convert_to_look_at: bool = True) -> bool: + """Register ``p`` to print the current viewer camera pose. + + Args: + convert_to_look_at: Print a ``window.set_look_at(...)`` call when + true, which is the default. Set false to print the raw matrix. + + Returns: + Whether the control is registered on an available window. + """ + self._window_camera_pose_hotkey_cfg = {"convert_to_look_at": convert_to_look_at} + if self._window is None: + logger.log_warning( + "No simulation window available yet. The camera pose print " + "hotkey will be registered after `open_window()`." + ) + return False + if self._window_camera_pose_input_control is not None: + return True + + from dexsim.types import InputKey + + sim = self + hotkey_cfg = dict(self._window_camera_pose_hotkey_cfg) + + class WindowCameraPoseEvent(ObjectManipulator): + def on_key_down(self, key): + if key == InputKey.SCANCODE_P.value: + sim.print_window_camera_pose(**hotkey_cfg) + + self._window_camera_pose_input_control = WindowCameraPoseEvent() + self._window.add_input_control(self._window_camera_pose_input_control) + logger.log_info( + "Camera pose print hotkey registered. Press 'p' to print the " + "current viewer pose." + ) + return True + def create_visual_material(self, cfg: VisualMaterialCfg) -> VisualMaterial: """Create a visual material with given configuration. diff --git a/embodichain/lab/sim/solvers/base_solver.py b/embodichain/lab/sim/solvers/base_solver.py index 271ac14f0..47d9e8dcb 100644 --- a/embodichain/lab/sim/solvers/base_solver.py +++ b/embodichain/lab/sim/solvers/base_solver.py @@ -285,14 +285,19 @@ def _init_qpos_limits(self): ) def update_with_robot_limit(self, robot_qpos_limits: torch.Tensor): - """Update with robot joint limits. - Make sure the solver's joint limits are within the robot's joint limits. + """Intersect solver joint limits with the robot's effective qpos limits. + + Robot-side articulation limits are the hard physical bound. Solver-specific + limits from ``SolverCfg.user_qpos_limits`` may be even tighter for planning. + The final solver limits must satisfy both constraints. Args: - robot_qpos_limits (torch.Tensor): [DOF, 2] tensor of joint limits from the robot data + robot_qpos_limits (torch.Tensor): [DOF, 2] tensor of joint limits from + the robot data. """ robot_lower_limits = robot_qpos_limits[:, 0] robot_upper_limits = robot_qpos_limits[:, 1] + if self.lower_qpos_limits is not None: if torch.any(self.lower_qpos_limits < robot_lower_limits): logger.log_warning( @@ -303,6 +308,7 @@ def update_with_robot_limit(self, robot_qpos_limits: torch.Tensor): ) else: self.lower_qpos_limits = robot_lower_limits + if self.upper_qpos_limits is not None: if torch.any(self.upper_qpos_limits > robot_upper_limits): logger.log_warning( diff --git a/embodichain/lab/sim/solvers/ur_solver.py b/embodichain/lab/sim/solvers/ur_solver.py index f3a7bd001..4657a1aed 100644 --- a/embodichain/lab/sim/solvers/ur_solver.py +++ b/embodichain/lab/sim/solvers/ur_solver.py @@ -206,7 +206,9 @@ def get_ik( return all_solutions_validity, all_solutions # Select ik qpos based on the closest distance to the seed qpos qpos_seed_expanded = qpos_seed.unsqueeze(1).expand(-1, N_SOL, -1) - distances = torch.norm(all_solutions - qpos_seed_expanded, dim=-1) + distances = torch.norm( + self.ik_nearest_weight * (all_solutions - qpos_seed_expanded), dim=-1 + ) # fill invalid solutions with inf distance distances[~all_solutions_validity] = float("inf") closest_indices = torch.argmin(distances, dim=1) diff --git a/embodichain/lab/sim/utility/atom_action_utils.py b/embodichain/lab/sim/utility/atom_action_utils.py index e4c715178..4525d5dd6 100644 --- a/embodichain/lab/sim/utility/atom_action_utils.py +++ b/embodichain/lab/sim/utility/atom_action_utils.py @@ -201,7 +201,7 @@ def plan_trajectory( ) plan_state = [ - PlanState(qpos=torch.as_tensor(qpos), move_type=MoveType.JOINT_MOVE) + PlanState.single(qpos=torch.as_tensor(qpos), move_type=MoveType.JOINT_MOVE) for qpos in qpos_list ] @@ -215,8 +215,10 @@ def plan_trajectory( ), ) - select_qpos_traj.extend(ret.positions.numpy()) - ee_state_list_select.extend([select_arm_current_gripper_state] * len(ret.positions)) + select_qpos_traj.extend(ret.positions[0].numpy()) + ee_state_list_select.extend( + [select_arm_current_gripper_state] * ret.positions.shape[1] + ) def plan_gripper_trajectory( diff --git a/embodichain/lab/sim/utility/cfg_utils.py b/embodichain/lab/sim/utility/cfg_utils.py index ffeb1831d..12246e848 100644 --- a/embodichain/lab/sim/utility/cfg_utils.py +++ b/embodichain/lab/sim/utility/cfg_utils.py @@ -96,22 +96,49 @@ def merge_robot_cfg(base_cfg: RobotCfg, override_cfg_dict: dict[str, any]) -> Ro for key, value in override_cfg_dict.items(): if key == "solver_cfg": - # merge provided solver_cfg values into default solver config + # Per-part merge of provided solver_cfg into default solver config. + # Two modes: + # 1. Part dict includes "class_type" → new/replacement solver + # (already deserialized into a SolverCfg subclass by + # RobotCfg.from_dict above). + # 2. Part dict lacks "class_type" → attribute overrides for an + # existing solver part (e.g. {"tcp": ..., "stiffness": ...}). provided_solver_cfg = override_cfg_dict.get("solver_cfg") - if provided_solver_cfg: + if provided_solver_cfg and isinstance(provided_solver_cfg, dict): + if base_cfg.solver_cfg is None: + base_cfg.solver_cfg = {} for part, item in provided_solver_cfg.items(): - if "class_type" in provided_solver_cfg[part]: - base_cfg.solver_cfg[part] = robot_cfg.solver_cfg[part] - else: - try: - merged = merge_solver_cfg( - base_cfg.solver_cfg, provided_solver_cfg - ) - base_cfg.solver_cfg = merged - except Exception: - logger.log_error( - f"Failed to merge solver_cfg, using provided config outright." + if isinstance(item, dict) and "class_type" in item: + # New or replacement solver part — use the deserialized + # SolverCfg object produced by RobotCfg.from_dict. + parsed = ( + robot_cfg.solver_cfg.get(part) + if isinstance(robot_cfg.solver_cfg, dict) + else None + ) + if parsed is not None: + base_cfg.solver_cfg[part] = parsed + else: + logger.log_warning( + f"Failed to deserialize solver_cfg['{part}'] " + f"with class_type={item.get('class_type')!r}. " + f"Skipping." ) + elif part in base_cfg.solver_cfg: + # Existing part — merge individual attribute overrides + # in-place so other parts are preserved. + target = base_cfg.solver_cfg[part] + if isinstance(item, dict): + for attr_name, attr_val in item.items(): + if hasattr(target, attr_name): + setattr(target, attr_name, attr_val) + else: + logger.log_warning( + f"Cannot add solver part {part!r} without " + f"'class_type'. Provide 'class_type' to create a " + f"new solver entry, or ensure the part name " + f"matches an existing solver." + ) elif key == "drive_pros": # merge joint drive properties user_drive_pros_dict = override_cfg_dict.get("drive_pros") diff --git a/embodichain/lab/sim/utility/sim_utils.py b/embodichain/lab/sim/utility/sim_utils.py index 9f0257d97..fc3e0ce24 100644 --- a/embodichain/lab/sim/utility/sim_utils.py +++ b/embodichain/lab/sim/utility/sim_utils.py @@ -20,6 +20,7 @@ import dexsim import open3d as o3d +from dataclasses import MISSING from typing import List, Union from dexsim.types import ( @@ -218,6 +219,37 @@ def get_dexsim_arena_num() -> int: return len(arenas) +def _resolve_mesh_collision_params( + cfg: RigidObjectCfg, +) -> tuple[int, str, int]: + """Resolve legacy and shape-level mesh collision parameters.""" + + def is_missing(value) -> bool: + # deepcopy() can produce a distinct instance of dataclasses.MISSING. + return value is MISSING or isinstance(value, type(MISSING)) + + max_convex_hull_num = next( + value + for value in ( + cfg.max_convex_hull_num, + cfg.shape.max_convex_hull_num, + 1, + ) + if not is_missing(value) + ) + acd_method = next( + value + for value in (cfg.acd_method, cfg.shape.acd_method, "coacd") + if not is_missing(value) + ) + sdf_resolution = next( + value + for value in (cfg.sdf_resolution, cfg.shape.sdf_resolution, 0) + if not is_missing(value) + ) + return max_convex_hull_num, acd_method, sdf_resolution + + def get_dexsim_drive_type(drive_type: str) -> DriveType: """Get the dexsim drive type from a string. @@ -622,10 +654,12 @@ def _load_rigid_mesh_prototype( """Load and configure one mesh rigid-object prototype in the source arena.""" option = _mesh_load_option_from_cfg(cfg) fpath = cfg.shape.fpath - max_convex_hull_num = cfg.max_convex_hull_num + max_convex_hull_num, acd_method, sdf_resolution = _resolve_mesh_collision_params( + cfg + ) if max_convex_hull_num > 1: - obj = env.load_actor_with_coacd( + obj = env.load_actor_with_acd( fpath, duplicate=True, attach_scene=True, @@ -633,8 +667,9 @@ def _load_rigid_mesh_prototype( cache_path=cache_dir, actor_type=body_type, max_convex_hull_num=max_convex_hull_num, + method=acd_method, ) - elif cfg.sdf_resolution > 0: + elif sdf_resolution > 0: if not is_newton_backend and cfg.body_scale not in [ (1.0, 1.0, 1.0), [1.0, 1.0, 1.0], @@ -645,7 +680,7 @@ def _load_rigid_mesh_prototype( "collision." ) obj = env.load_actor(fpath, duplicate=True, attach_scene=True, option=option) - sdf_cfg = SDFConfig(resolution=cfg.sdf_resolution) + sdf_cfg = SDFConfig(resolution=sdf_resolution) obj.add_physical_body( body_type, RigidBodyShape.SDF, diff --git a/embodichain/utils/configclass.py b/embodichain/utils/configclass.py index dad8b4efd..a2d0a5542 100644 --- a/embodichain/utils/configclass.py +++ b/embodichain/utils/configclass.py @@ -172,7 +172,12 @@ def custom_post_init(obj): # duplicate data members that are mutable if not callable(value) and not isinstance(ann, property): try: - setattr(obj, key, deepcopy(value)) + # MISSING is a sentinel singleton — preserve its identity so that + # ``is MISSING`` / ``is not MISSING`` checks work on configclass + # instances. deepcopy() would create a new _MISSING_TYPE object, + # breaking identity comparison. + if value is not MISSING: + setattr(obj, key, deepcopy(value)) except AttributeError as e: from IPython import embed @@ -617,6 +622,12 @@ def _wrap(): else: return f.default_factory else: + # MISSING is a sentinel singleton — preserve its identity so that + # ``is MISSING`` / ``is not MISSING`` checks work on configclass + # instances. deepcopy() would create a new _MISSING_TYPE object, + # breaking identity comparison. + if f is MISSING: + return f return deepcopy(f) return _wrap diff --git a/embodichain/utils/device_utils.py b/embodichain/utils/device_utils.py index 7a84ea4ac..bfe85f2a0 100644 --- a/embodichain/utils/device_utils.py +++ b/embodichain/utils/device_utils.py @@ -33,7 +33,7 @@ def standardize_device_string(device: Union[str, torch.device]) -> str: else: device_str = str(device) - if device_str.startswith("cuda"): + if device_str == "cuda": device_str = "cuda:0" return device_str diff --git a/embodichain/utils/math.py b/embodichain/utils/math.py index fbbe75f6f..d5f664791 100644 --- a/embodichain/utils/math.py +++ b/embodichain/utils/math.py @@ -540,6 +540,61 @@ def _axis_angle_rotation( return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3)) +def axis_angle_to_rotation_matrix(axis_angle: torch.Tensor) -> torch.Tensor: + """Convert axis-angle representation to rotation matrix. + + Args: + axis_angle: Axis-angle representation, radian * axis. Shape is (N, 3). + + Returns: + Rotation matrices. Shape is (N, 3, 3). + """ + if axis_angle.dim() == 0 or axis_angle.shape[-1] != 3: + raise ValueError( + f"Invalid axis-angle shape {axis_angle.shape}, expected (..., 3)." + ) + + # Rodrigues formula with stable Taylor fallback near zero angle. + theta2 = torch.sum(axis_angle * axis_angle, dim=-1, keepdim=True) + theta = torch.sqrt(theta2) + eps = 1.0e-8 + + sin_over_theta = torch.where( + theta2 > eps, + torch.sin(theta) / theta, + 1.0 - theta2 / 6.0 + theta2 * theta2 / 120.0, + ) + one_minus_cos_over_theta2 = torch.where( + theta2 > eps, + (1.0 - torch.cos(theta)) / theta2, + 0.5 - theta2 / 24.0 + theta2 * theta2 / 720.0, + ) + + wx, wy, wz = torch.unbind(axis_angle, dim=-1) + skew = torch.zeros( + axis_angle.shape[:-1] + (3, 3), + dtype=axis_angle.dtype, + device=axis_angle.device, + ) + skew[..., 0, 1] = -wz + skew[..., 0, 2] = wy + skew[..., 1, 0] = wz + skew[..., 1, 2] = -wx + skew[..., 2, 0] = -wy + skew[..., 2, 1] = wx + + eye = torch.eye(3, dtype=axis_angle.dtype, device=axis_angle.device).expand( + axis_angle.shape[:-1] + (3, 3) + ) + skew2 = torch.matmul(skew, skew) + + return ( + eye + + sin_over_theta.unsqueeze(-1) * skew + + one_minus_cos_over_theta2.unsqueeze(-1) * skew2 + ) + + def matrix_from_euler( euler_angles: torch.Tensor, convention: str = "XYZ" ) -> torch.Tensor: diff --git a/examples/sim/demo/pick_up_cloth.py b/examples/sim/demo/pick_up_cloth.py index 6be3e3c18..5f5315b61 100644 --- a/examples/sim/demo/pick_up_cloth.py +++ b/examples/sim/demo/pick_up_cloth.py @@ -31,21 +31,18 @@ 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.lab.sim.solvers import PytorchSolverCfg from embodichain.data import get_data_path from embodichain.utils import logger from embodichain.lab.sim.cfg import ( RenderCfg, physics_cfg_for_backend, - JointDrivePropertiesCfg, - RobotCfg, RigidObjectCfg, RigidBodyAttributesCfg, LightCfg, ClothObjectCfg, ClothPhysicalAttributesCfg, - URDFCfg, ) +from embodichain.lab.sim.robots import URRobotCfg import os from embodichain.lab.sim.shapes import MeshCfg, CubeCfg import tempfile @@ -62,42 +59,47 @@ def create_robot(sim: SimulationManager, position=[0.0, 0.0, 0.0]): Returns: Robot: The configured robot instance added to the simulation. """ - # Retrieve URDF paths for the robot arm and hand - ur10_urdf_path = get_data_path("UniversalRobots/UR10/UR10.urdf") gripper_urdf_path = get_data_path("DH_PGC_140_50_M/DH_PGC_140_50_M.urdf") - # Configure the robot with its components and control properties - cfg = RobotCfg( - uid="UR10", - urdf_cfg=URDFCfg( - components=[ - {"component_type": "arm", "urdf_path": ur10_urdf_path}, - {"component_type": "hand", "urdf_path": gripper_urdf_path}, - ] - ), - drive_pros=JointDrivePropertiesCfg( - stiffness={"JOINT[0-9]": 1e4, "FINGER[1-2]": 1e2}, - damping={"JOINT[0-9]": 1e3, "FINGER[1-2]": 1e1}, - max_effort={"JOINT[0-9]": 1e5, "FINGER[1-2]": 1e3}, - drive_type="force", - ), - control_parts={ - "arm": ["JOINT[0-9]"], - "hand": ["FINGER[1-2]"], - }, - solver_cfg={ - "arm": PytorchSolverCfg( - end_link_name="ee_link", - root_link_name="base_link", - tcp=[ - [0.0, 1.0, 0.0, 0.0], - [-1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.12], - [0.0, 0.0, 0.0, 1.0], - ], - ) - }, - 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, + 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 sim.add_robot(cfg=cfg) diff --git a/examples/sim/demo/press_softbody.py b/examples/sim/demo/press_softbody.py index 7b6d23ac6..9d5af90f0 100644 --- a/examples/sim/demo/press_softbody.py +++ b/examples/sim/demo/press_softbody.py @@ -30,21 +30,19 @@ 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.lab.sim.solvers import PytorchSolverCfg from embodichain.data import get_data_path from embodichain.utils import logger from embodichain.lab.sim.cfg import ( RenderCfg, physics_cfg_for_backend, - RobotCfg, LightCfg, SoftObjectCfg, SoftbodyVoxelAttributesCfg, SoftbodyPhysicalAttributesCfg, - URDFCfg, ) 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.robots import URRobotCfg def parse_arguments(): @@ -94,33 +92,20 @@ def create_robot(sim: SimulationManager): Returns: Robot: The configured robot instance added to the simulation. """ - # Retrieve URDF paths for the robot arm and hand - ur10_urdf_path = get_data_path("UniversalRobots/UR10/UR10.urdf") - - # Configure the robot with its components and control properties - cfg = RobotCfg( - uid="UR10", - urdf_cfg=URDFCfg( - components=[{"component_type": "arm", "urdf_path": ur10_urdf_path}] - ), - control_parts={ - "arm": ["Joint[0-9]"], - }, - solver_cfg={ - "arm": PytorchSolverCfg( - end_link_name="ee_link", - root_link_name="base_link", - tcp=np.eye(4), - ) - }, - init_qpos=[ - 0.0, - -np.pi / 2, - -np.pi / 2, - np.pi / 2, - -np.pi / 2, - 0.0, - ], + 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) diff --git a/examples/sim/demo/scoop_ice.py b/examples/sim/demo/scoop_ice.py index 2f03afe8c..ff9493f16 100644 --- a/examples/sim/demo/scoop_ice.py +++ b/examples/sim/demo/scoop_ice.py @@ -31,22 +31,20 @@ from embodichain.lab.sim.cfg import ( RenderCfg, physics_cfg_for_backend, - JointDrivePropertiesCfg, - RobotCfg, - URDFCfg, RigidObjectCfg, RigidBodyAttributesCfg, ArticulationCfg, RigidObjectGroupCfg, + JointDrivePropertiesCfg, LightCfg, ) from embodichain.lab.sim.material import VisualMaterialCfg from embodichain.lab.sim.utility.action_utils import interpolate_with_distance from embodichain.lab.sim.shapes import MeshCfg, CubeCfg -from embodichain.lab.sim.solvers import PytorchSolverCfg from embodichain.data import get_data_path 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): @@ -114,8 +112,6 @@ def create_robot(sim): Returns: Robot: The configured robot instance added to the simulation. """ - # Retrieve URDF paths for the robot arm and hand - ur10_urdf_path = get_data_path("UniversalRobots/UR10/UR10.urdf") hand_urdf_path = get_data_path( "BrainCoHandRevo1/BrainCoLeftHand/BrainCoLeftHand.urdf" ) @@ -124,61 +120,55 @@ def create_robot(sim): hand_attach_xpos = np.eye(4) hand_attach_xpos[:3, :3] = R.from_rotvec([90, 0, 0], degrees=True).as_matrix() - # Configure the robot with its components and control properties - cfg = RobotCfg( - uid="ur10_with_brainco", - urdf_cfg=URDFCfg( - components=[ - {"component_type": "arm", "urdf_path": ur10_urdf_path}, - { - "component_type": "hand", - "urdf_path": hand_urdf_path, - "transform": hand_attach_xpos, - }, - ] - ), - control_parts={ - "arm": ["JOINT[0-9]"], - "hand": [ - "LEFT_HAND_THUMB1", - "LEFT_HAND_THUMB2", - "LEFT_HAND_INDEX", - "LEFT_HAND_MIDDLE", - "LEFT_HAND_RING", - "LEFT_HAND_PINKY", + 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, ], - }, - drive_pros=JointDrivePropertiesCfg( - stiffness={"JOINT[0-9]": 1e4, "LEFT_[A-Z|_]+[0-9]?": 1e2}, - damping={"JOINT[0-9]": 1e3, "LEFT_[A-Z|_]+[0-9]?": 1e1}, - max_effort={"JOINT[0-9]": 1e5, "LEFT_[A-Z|_]+[0-9]?": 1e3}, - drive_type="force", - ), - solver_cfg={ - "arm": PytorchSolverCfg( - end_link_name="ee_link", - root_link_name="base_link", - 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) diff --git a/examples/sim/planners/motion_generator.py b/examples/sim/planners/motion_generator.py deleted file mode 100644 index 3f118e24e..000000000 --- a/examples/sim/planners/motion_generator.py +++ /dev/null @@ -1,200 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -import time -import torch -import numpy as np -from copy import deepcopy -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim.robots import CobotMagicCfg -from embodichain.lab.sim.planners.utils import TrajectorySampleMethod - - -def move_robot_along_trajectory( - robot: Robot, arm_name: str, qpos_list: list[torch.Tensor], delay: float = 0.1 -): - """ - Set the robot joint positions sequentially along the given joint trajectory. - Args: - robot: Robot instance. - arm_name: Name of the robot arm. - qpos_list: List of joint positions (torch.Tensor). - delay: Time delay between each step (seconds). - """ - for q in qpos_list: - robot.set_qpos(qpos=q.unsqueeze(0), joint_ids=robot.get_joint_ids(arm_name)) - time.sleep(delay) - - -def create_demo_trajectory( - robot: Robot, arm_name: str -) -> tuple[list[torch.Tensor], list[np.ndarray]]: - """ - Generate a three-point trajectory (start, middle, end) for demonstration. - Args: - robot: Robot instance. - arm_name: Name of the robot arm. - Returns: - qpos_list: List of joint positions (torch.Tensor). - xpos_list: List of end-effector poses (numpy arrays). - """ - qpos_fk = torch.tensor( - [[0.0, np.pi / 4, -np.pi / 4, 0.0, np.pi / 4, 0.0]], dtype=torch.float32 - ) - xpos_begin = robot.compute_fk(name=arm_name, qpos=qpos_fk, to_matrix=True) - xpos_mid = deepcopy(xpos_begin) - xpos_mid[0, 2, 3] -= 0.1 # Move down by 0.1m in Z direction - xpos_final = deepcopy(xpos_mid) - xpos_final[0, 0, 3] += 0.2 # Move forward by 0.2m in X direction - - qpos_begin = robot.compute_ik(pose=xpos_begin, name=arm_name)[1][0] - qpos_mid = robot.compute_ik(pose=xpos_mid, name=arm_name)[1][0] - qpos_final = robot.compute_ik(pose=xpos_final, name=arm_name)[1][0] - return [qpos_begin, qpos_mid, qpos_final], [ - xpos_begin[0], - xpos_mid[0], - xpos_final[0], - ] - - -def main(interactive=False): - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) - - # Initialize simulation - sim = SimulationManager(SimulationManagerCfg(headless=False, device="cpu")) - sim.set_manual_update(False) - - # Robot configuration - 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]], - }, - }, - } - robot: Robot = sim.add_robot(cfg=CobotMagicCfg.from_dict(cfg_dict)) - arm_name = "left_arm" - - # Generate trajectory points - qpos_list, xpos_list = create_demo_trajectory(robot, arm_name) - - from embodichain.lab.sim.planners import ( - MotionGenerator, - MotionGenCfg, - MotionGenOptions, - ToppraPlannerCfg, - ToppraPlanOptions, - PlanState, - MoveType, - MovePart, - ) - - # Initialize motion generator - motion_cfg = MotionGenCfg( - planner_cfg=ToppraPlannerCfg( - robot_uid=robot.uid, - ) - ) - motion_generator = MotionGenerator(cfg=motion_cfg) - - current_qpos = robot.get_qpos(name=arm_name)[0] - options = MotionGenOptions( - control_part=arm_name, - start_qpos=current_qpos, - is_interpolate=True, - is_linear=False, - plan_opts=ToppraPlanOptions( - constraints={ - "velocity": 0.2, - "acceleration": 0.5, - }, - sample_method=TrajectorySampleMethod.QUANTITY, - sample_interval=20, - ), - ) - # Joint space trajectory - qpos_list = torch.vstack(qpos_list) - target_states = [] - for qpos in qpos_list: - target_states.append( - PlanState( - move_type=MoveType.JOINT_MOVE, - qpos=qpos, - ) - ) - plan_result = motion_generator.generate( - target_states=target_states, options=options - ) - move_robot_along_trajectory(robot, arm_name, plan_result.positions) - - # Cartesian space trajectory - options.is_linear = True - - xpos_list = torch.concatenate([xpos.unsqueeze(0) for xpos in xpos_list]) - target_states = [] - for xpos in xpos_list: - target_states.append( - PlanState( - move_type=MoveType.EEF_MOVE, - xpos=xpos, - ) - ) - plan_result = motion_generator.generate( - target_states=target_states, options=options - ) - sim.reset() - move_robot_along_trajectory(robot, arm_name, plan_result.positions) - - if interactive: - # Enter IPython interactive shell if needed - from IPython import embed - - embed() - - -if __name__ == "__main__": - main() diff --git a/examples/sim/planners/neural_planner.py b/examples/sim/planners/neural_planner.py index e0a98e96d..eb202d3b1 100644 --- a/examples/sim/planners/neural_planner.py +++ b/examples/sim/planners/neural_planner.py @@ -17,18 +17,16 @@ from __future__ import annotations import argparse -import math -import os import time import numpy as np import torch -from embodichain.data import get_data_path from embodichain.data.assets.planner_assets import download_neural_planner_checkpoint from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import MarkerCfg, RobotCfg +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 ( MotionGenCfg, MotionGenOptions, @@ -89,44 +87,8 @@ def _resolve_device(device: str) -> str: return device -def _franka_tcp() -> list[list[float]]: - c = math.cos(-math.pi / 4) - s = math.sin(-math.pi / 4) - return [ - [c, -s, 0.0, 0.0], - [s, c, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.1034], - [0.0, 0.0, 0.0, 1.0], - ] - - def create_franka(sim: SimulationManager) -> Robot: - urdf = get_data_path("Franka/Panda/PandaWithHand.urdf") - assert os.path.isfile(urdf) - - cfg_dict = { - "fpath": urdf, - "control_parts": { - "main_arm": [ - "Joint1", - "Joint2", - "Joint3", - "Joint4", - "Joint5", - "Joint6", - "Joint7", - ], - }, - "solver_cfg": { - "main_arm": { - "class_type": "PytorchSolver", - "end_link_name": "ee_link", - "root_link_name": "base_link", - "tcp": _franka_tcp(), - }, - }, - } - return sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + return sim.add_robot(cfg=FrankaPandaCfg.from_dict({"robot_type": "panda"})) def make_waypoints(start_pose: torch.Tensor, num_waypoints: int) -> torch.Tensor: @@ -179,6 +141,9 @@ def play_trajectory( delay: float = 0.0, ) -> None: joint_ids = robot.get_joint_ids(arm_name) + # ``positions`` may be env-batched (B, N, DOF); this example is single-env. + if positions.dim() == 3: + positions = positions[0] for qpos in positions: robot.set_qpos(qpos=qpos.unsqueeze(0), joint_ids=joint_ids) sim.update(step=step_repeat) @@ -201,7 +166,7 @@ def main() -> None: ) robot = create_franka(sim) - arm_name = "main_arm" + arm_name = "arm" device = robot.device if not args.headless: @@ -233,7 +198,8 @@ def main() -> None: ) ) target_states = [ - PlanState(move_type=MoveType.EEF_MOVE, xpos=waypoint) for waypoint in waypoints + PlanState.single(move_type=MoveType.EEF_MOVE, xpos=waypoint) + for waypoint in waypoints ] result = motion_generator.generate( target_states=target_states, @@ -248,7 +214,7 @@ def main() -> None: print(f"NeuralPlanner success: {result.success}") print(f"positions shape: {tuple(result.positions.shape)}") print(f"xpos_list shape: {tuple(result.xpos_list.shape)}") - print(f"duration: {result.duration:.3f}s") + print(f"duration: {result.duration.item():.3f}s") play_trajectory( sim, diff --git a/examples/sim/solvers/neural_ik_solver.py b/examples/sim/solvers/neural_ik_solver.py index ee6109551..2bd3580cf 100644 --- a/examples/sim/solvers/neural_ik_solver.py +++ b/examples/sim/solvers/neural_ik_solver.py @@ -15,18 +15,18 @@ # ---------------------------------------------------------------------------- import argparse import math -import os import time import numpy as np import torch from IPython import embed -from embodichain.data import get_data_path from embodichain.data.assets.solver_assets import download_neural_ik_checkpoint -from embodichain.lab.sim.cfg import MarkerCfg, RobotCfg -from embodichain.lab.sim.objects import Robot 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: @@ -97,55 +97,32 @@ def main(): ) sim = SimulationManager(config) - urdf = get_data_path("Franka/Panda/PandaWithHand.urdf") - assert os.path.isfile(urdf) - checkpoint_path = download_neural_ik_checkpoint() - c = math.cos(-math.pi / 4) - s = math.sin(-math.pi / 4) - tcp = [ - [c, -s, 0.0, 0.0], - [s, c, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.1034], - [0.0, 0.0, 0.0, 1.0], - ] - - cfg_dict = { - "fpath": urdf, - "control_parts": { - "main_arm": [ - "Joint1", - "Joint2", - "Joint3", - "Joint4", - "Joint5", - "Joint6", - "Joint7", - ], - }, - "solver_cfg": { - "main_arm": { - "class_type": "NeuralIKSolver", - "end_link_name": "ee_link", - "root_link_name": "base_link", - "tcp": tcp, - "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=RobotCfg.from_dict(cfg_dict)) + 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 = "main_arm" + arm_name = "arm" device = robot.device seed_qpos = torch.tensor( diff --git a/pyproject.toml b/pyproject.toml index f5506b0d5..b36950503 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dynamic = ["version"] # Core install dependencies (kept from requirements.txt). Some VCS links are # specified using PEP 508 direct references where present. dependencies = [ - "dexsim_engine==0.4.2", + "dexsim_engine==0.4.3", "setuptools>=78.1.1", "gymnasium>=0.29.1", "langchain", @@ -72,3 +72,11 @@ where = ["."] exclude = ["docs"] [tool.black] + +[tool.pytest.ini_options] +markers = [ + "slow: marks tests as slow (deselect with '-m \"not slow\"')", + "requires_sim: marks tests that require a real simulation backend", + "gpu: marks tests that execute CUDA/GPU code (run serially)", + "renderer: marks tests that exercise a rendering backend", +] diff --git a/scripts/benchmark/__main__.py b/scripts/benchmark/__main__.py index a885cdca1..3bb06936c 100644 --- a/scripts/benchmark/__main__.py +++ b/scripts/benchmark/__main__.py @@ -22,6 +22,7 @@ python -m scripts.benchmark rl --rebuild-report-only python -m scripts.benchmark robotics-kinematic-solver -s pytorch python -m scripts.benchmark planners-neural-planner --num-waypoints 1 3 5 + python -m scripts.benchmark atomic-action --smoke """ from __future__ import annotations @@ -66,6 +67,13 @@ def _run_neural_planner_cli(args: argparse.Namespace) -> None: ) +def _run_atomic_action_cli(_: argparse.Namespace) -> None: + """Run atomic action benchmark CLI entrypoint.""" + from scripts.benchmark.atomic_action.run_benchmark import main as atomic_main + + atomic_main() + + def main() -> None: """Dispatch to the appropriate benchmark sub-command CLI.""" parser = argparse.ArgumentParser( @@ -167,6 +175,16 @@ def main() -> None: ) neural_planner_parser.set_defaults(func=_run_neural_planner_cli) + # -- atomic-action ------------------------------------------------------- + atomic_action_parser = subparsers.add_parser( + "atomic-action", + help="Benchmark atomic actions over object presets and positions.", + ) + from scripts.benchmark.atomic_action.run_benchmark import add_benchmark_args + + add_benchmark_args(atomic_action_parser) + atomic_action_parser.set_defaults(func=_run_atomic_action_cli) + # -- Parse --------------------------------------------------------------- # If no sub-command is given, print help and exit. if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help"): @@ -193,3 +211,6 @@ def main() -> None: if __name__ == "__main__": main() + + +__all__ = ["main"] diff --git a/scripts/benchmark/atomic_action/__init__.py b/scripts/benchmark/atomic_action/__init__.py new file mode 100644 index 000000000..4709da3d9 --- /dev/null +++ b/scripts/benchmark/atomic_action/__init__.py @@ -0,0 +1,21 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Benchmarks for atomic actions.""" + +from __future__ import annotations + +__all__ = [] diff --git a/scripts/benchmark/atomic_action/common.py b/scripts/benchmark/atomic_action/common.py new file mode 100644 index 000000000..7220ab625 --- /dev/null +++ b/scripts/benchmark/atomic_action/common.py @@ -0,0 +1,1194 @@ +# ---------------------------------------------------------------------------- +# 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 helpers for atomic-action benchmark scripts.""" + +from __future__ import annotations + +import argparse +import math +import os +import re +import resource +import sys +import time +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Callable + +try: + import psutil +except ModuleNotFoundError: + psutil = None + +CPU_MEMORY_BACKEND = "psutil" if psutil is not None else "resource" +DEFAULT_VIDEO_DIR = Path("outputs/benchmark_videos") +DEFAULT_VIDEO_FPS = 20 +DEFAULT_VIDEO_MAX_MEMORY_MB = 2048 +DEFAULT_VIDEO_WIDTH = 640 +DEFAULT_VIDEO_HEIGHT = 480 +DEFAULT_VIDEO_HOLD_STEPS = 120 +DEFAULT_VIDEO_CASE_LIMIT = 0 +PHYSICAL_PICK_MIN_LIFT_M = 0.04 +PHYSICAL_PLACE_XY_TOLERANCE_M = 0.10 +PHYSICAL_MOVE_HELD_OBJECT_XYZ_TOLERANCE_M = 0.12 +PHYSICAL_VALIDATION_HOLD_STEPS = 80 +SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z = 0.35 +SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT = 0.5 +BENCHMARK_PROFILES = ("smoke", "coverage", "full") +DEFAULT_BENCHMARK_PROFILE = "coverage" +DEFAULT_VIDEO_LOOK_AT = ( + (-1.25, -1.15, 0.95), + (-0.25, -0.02, 0.25), + (0.0, 0.0, 1.0), +) + + +@dataclass(frozen=True) +class PositionCase: + """Initial object position case with a quadrant label.""" + + name: str + quadrant: str + xy: tuple[float, float] + + +@dataclass(frozen=True) +class MeshObjectPreset: + """Real mesh object preset used by object-conditioned benchmarks.""" + + object_type: str + material_name: str + label: str + init_rot: tuple[float, float, float] + body_scale: tuple[float, float, float] + mass: float + initial_z: float + mesh_path: str = "" + shape_type: str = "mesh" + cube_size: tuple[float, float, float] | None = None + use_usd_properties: bool = False + dynamic_friction: float = 0.97 + static_friction: float = 0.99 + restitution: float = 0.0 + contact_offset: float = 0.002 + rest_offset: float = 0.0 + linear_damping: float = 0.7 + angular_damping: float = 0.7 + max_depenetration_velocity: float = 10.0 + min_position_iters: int = 4 + min_velocity_iters: int = 1 + max_linear_velocity: float = 100.0 + max_angular_velocity: float = 100.0 + max_convex_hull_num: int = 16 + enable_ccd: bool = False + + +POSITION_CASES: dict[str, PositionCase] = { + "q1_near": PositionCase(name="q1_near", quadrant="q1", xy=(0.02, 0.18)), + "q1_far": PositionCase(name="q1_far", quadrant="q1", xy=(0.12, 0.36)), + "q2_near": PositionCase(name="q2_near", quadrant="q2", xy=(-0.42, 0.18)), + "q2_far": PositionCase(name="q2_far", quadrant="q2", xy=(-0.62, 0.36)), + "q3_near": PositionCase(name="q3_near", quadrant="q3", xy=(-0.42, -0.18)), + "q3_far": PositionCase(name="q3_far", quadrant="q3", xy=(-0.62, -0.36)), + "q4_near": PositionCase(name="q4_near", quadrant="q4", xy=(0.02, -0.18)), + "q4_far": PositionCase(name="q4_far", quadrant="q4", xy=(0.12, -0.36)), +} +FULL_POSITION_CASE_NAMES = tuple(POSITION_CASES.keys()) +COVERAGE_POSITION_CASE_NAMES = FULL_POSITION_CASE_NAMES +SMOKE_POSITION_CASE_NAMES = ("q3_near",) + +MESH_OBJECT_PRESETS: dict[str, MeshObjectPreset] = { + "sugar_box": MeshObjectPreset( + object_type="sugar_box", + material_name="cardboard", + 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, + initial_z=0.05, + use_usd_properties=False, + ), + "coffee_cup": MeshObjectPreset( + object_type="coffee_cup", + material_name="ceramic", + label="coffee_cup", + mesh_path="CoffeeCup/cup.ply", + init_rot=(0.0, 0.0, -90.0), + body_scale=(4.0, 4.0, 4.0), + mass=0.01, + initial_z=0.01, + use_usd_properties=False, + ), + "cube": MeshObjectPreset( + object_type="cube", + material_name="plastic", + label="cube", + shape_type="cube", + cube_size=(0.05, 0.05, 0.05), + init_rot=(0.0, 0.0, 0.0), + body_scale=(1.0, 1.0, 1.0), + mass=0.05, + initial_z=0.05, + use_usd_properties=False, + dynamic_friction=0.5, + static_friction=0.5, + contact_offset=0.003, + rest_offset=0.001, + max_depenetration_velocity=10.0, + min_position_iters=32, + min_velocity_iters=8, + max_convex_hull_num=1, + ), + "paper_cup": MeshObjectPreset( + object_type="paper_cup", + material_name="paper", + label="paper_cup", + mesh_path="PaperCup/paper_cup.ply", + init_rot=(0.0, 0.0, 0.0), + body_scale=(0.75, 0.75, 1.0), + mass=0.01, + initial_z=0.05, + use_usd_properties=False, + dynamic_friction=1.0, + static_friction=1.0, + contact_offset=0.003, + rest_offset=0.001, + linear_damping=2.0, + angular_damping=2.0, + max_depenetration_velocity=2.0, + min_position_iters=32, + min_velocity_iters=8, + max_linear_velocity=5.0, + max_angular_velocity=10.0, + max_convex_hull_num=8, + ), + "scanned_bottle": MeshObjectPreset( + object_type="scanned_bottle", + material_name="plastic", + label="scanned_bottle", + mesh_path="ScannedBottle/yibao_processed.ply", + init_rot=(0.0, 0.0, 0.0), + body_scale=(1.0, 1.0, 1.0), + mass=0.05, + initial_z=0.05, + use_usd_properties=False, + ), +} +COVERAGE_MESH_OBJECT_TYPES = ("sugar_box", "cube", "paper_cup") +FULL_MESH_OBJECT_TYPES = COVERAGE_MESH_OBJECT_TYPES +SMOKE_MESH_OBJECT_TYPES = ("sugar_box",) +PICKUP_APPROACH_CASES = ("top", "side") +SMOKE_PICKUP_APPROACH_CASES = ("top",) + + +def ensure_repo_root() -> None: + """Add the repository root to sys.path for module execution.""" + repo_root = Path(__file__).resolve().parents[3] + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + + +def ensure_torch(): + """Import torch or raise a clear benchmark runtime error.""" + try: + import torch + except ModuleNotFoundError as exc: + raise RuntimeError( + "Atomic action benchmark requires the EmbodiChain simulation runtime " + f"and PyTorch. Missing module: {exc.name}." + ) from exc + return torch + + +def add_common_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add common atomic-action benchmark CLI arguments.""" + add_profile_benchmark_args(parser) + parser.add_argument( + "--repeat", + type=int, + default=1, + help="Number of repeats for every benchmark case.", + ) + parser.add_argument( + "--smoke", + action="store_true", + help="Alias for --profile smoke.", + ) + parser.add_argument( + "--device", + type=str, + default="cpu", + help="Simulation device, e.g. 'cpu' or 'cuda'.", + ) + parser.add_argument( + "--renderer", + type=str, + choices=("auto", "hybrid", "fast-rt", "rt"), + default="auto", + help="Renderer backend used by SimulationManager.", + ) + add_video_benchmark_args(parser) + + +def add_video_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add optional benchmark video recording CLI arguments.""" + parser.add_argument( + "--record_video", + action="store_true", + help="Record trajectory replay videos for selected successful cases.", + ) + parser.add_argument( + "--record_failed_video", + action="store_true", + help=( + "With --record_video, also record failed cases when a partial " + "trajectory or static debug scene is available." + ), + ) + parser.add_argument( + "--video_case_limit", + type=int, + default=DEFAULT_VIDEO_CASE_LIMIT, + help="Maximum cases to record. Use 0 to record all selected cases.", + ) + parser.add_argument( + "--video_dir", + type=Path, + default=DEFAULT_VIDEO_DIR, + help="Directory for benchmark replay videos.", + ) + parser.add_argument( + "--video_fps", + type=int, + default=DEFAULT_VIDEO_FPS, + help="Recorded video frames per second.", + ) + parser.add_argument( + "--video_max_memory", + type=int, + default=DEFAULT_VIDEO_MAX_MEMORY_MB, + help="Maximum recorder frame-buffer memory in MB.", + ) + parser.add_argument( + "--video_width", + type=int, + default=DEFAULT_VIDEO_WIDTH, + help="Recorded video width in pixels.", + ) + parser.add_argument( + "--video_height", + type=int, + default=DEFAULT_VIDEO_HEIGHT, + help="Recorded video height in pixels.", + ) + parser.add_argument( + "--video_hold_steps", + type=int, + default=DEFAULT_VIDEO_HOLD_STEPS, + help="Extra simulation steps to hold the final replay pose.", + ) + + +def add_grasp_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add common grasp-affordance setup arguments.""" + parser.add_argument( + "--n_sample", + type=int, + default=10000, + help="Number of samples for antipodal grasp generation.", + ) + parser.add_argument( + "--force_reannotate", + action="store_true", + help="Force grasp region re-annotation instead of using cached data.", + ) + + +def add_profile_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add unified benchmark profile CLI arguments.""" + parser.add_argument( + "--profile", + choices=BENCHMARK_PROFILES, + default=DEFAULT_BENCHMARK_PROFILE, + help=( + "Benchmark profile: smoke is one fast case, coverage is the default " + "core object/position matrix, and full sweeps all configured cases." + ), + ) + + +def add_object_position_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add object and initial-position selection arguments.""" + parser.add_argument( + "--object_types", + nargs="+", + choices=(*MESH_OBJECT_PRESETS.keys(), "all"), + default=None, + help=( + "Real mesh object presets to benchmark. Defaults are selected by " + "--profile; use 'all' to include every default full preset." + ), + ) + parser.add_argument( + "--position_cases", + nargs="+", + choices=(*POSITION_CASES.keys(), "all"), + default=None, + help=( + "Initial object position cases to benchmark. Defaults are selected by " + "--profile; use 'all' for all near/far cases." + ), + ) + + +def add_pickup_approach_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add common PickUp approach selection arguments.""" + parser.add_argument( + "--approach_cases", + nargs="+", + choices=(*PICKUP_APPROACH_CASES, "all"), + default=None, + help=( + "PickUp approach cases to benchmark. Defaults are selected by " + "--profile; side uses the current object's initial XY direction." + ), + ) + + +def resolve_profile(args: argparse.Namespace) -> str: + """Resolve the effective benchmark profile from CLI arguments.""" + profile = getattr(args, "profile", DEFAULT_BENCHMARK_PROFILE) + if getattr(args, "smoke", False): + profile = "smoke" + if profile not in BENCHMARK_PROFILES: + raise ValueError( + f"Unsupported benchmark profile {profile!r}. " + f"Expected one of {BENCHMARK_PROFILES}." + ) + return profile + + +def default_position_case_names_for_profile(profile: str) -> tuple[str, ...]: + """Return default position case names for a profile.""" + if profile == "smoke": + return SMOKE_POSITION_CASE_NAMES + if profile == "coverage": + return COVERAGE_POSITION_CASE_NAMES + if profile == "full": + return FULL_POSITION_CASE_NAMES + raise ValueError(f"Unsupported benchmark profile: {profile}") + + +def select_position_cases( + case_names: Sequence[str] | None, + profile: str, +) -> list[PositionCase]: + """Resolve requested position cases, falling back to profile defaults.""" + names = ( + default_position_case_names_for_profile(profile) + if not case_names + else tuple(case_names) + ) + if "all" in names: + names = FULL_POSITION_CASE_NAMES + return [POSITION_CASES[name] for name in names] + + +def default_mesh_object_types_for_profile(profile: str) -> tuple[str, ...]: + """Return default real mesh object preset names for a profile.""" + if profile == "smoke": + return SMOKE_MESH_OBJECT_TYPES + if profile == "coverage": + return COVERAGE_MESH_OBJECT_TYPES + if profile == "full": + return FULL_MESH_OBJECT_TYPES + raise ValueError(f"Unsupported benchmark profile: {profile}") + + +def select_mesh_object_presets( + object_types: Sequence[str] | None, + profile: str, +) -> list[MeshObjectPreset]: + """Resolve requested mesh object presets, falling back to profile defaults.""" + names = ( + default_mesh_object_types_for_profile(profile) + if not object_types + else tuple(object_types) + ) + if "all" in names: + names = FULL_MESH_OBJECT_TYPES + return [MESH_OBJECT_PRESETS[name] for name in names] + + +def default_pickup_approach_cases_for_profile(profile: str) -> tuple[str, ...]: + """Return default PickUp approach case names for a profile.""" + if profile == "smoke": + return SMOKE_PICKUP_APPROACH_CASES + if profile in ("coverage", "full"): + return PICKUP_APPROACH_CASES + raise ValueError(f"Unsupported benchmark profile: {profile}") + + +def select_pickup_approaches( + approach_cases: Sequence[str] | None, + profile: str, +) -> list[str]: + """Resolve PickUp approach cases, falling back to profile defaults.""" + names = ( + default_pickup_approach_cases_for_profile(profile) + if not approach_cases + else tuple(approach_cases) + ) + if "all" in names: + names = PICKUP_APPROACH_CASES + return list(names) + + +def pickup_approach_direction_tuple( + approach: str, + position_case: PositionCase, +) -> tuple[float, float, float]: + """Resolve a PickUp approach name into a normalized world-frame tuple.""" + if approach == "top": + direction = (0.0, 0.0, -1.0) + elif approach == "side": + direction = (-position_case.xy[0], -position_case.xy[1], 0.0) + else: + raise ValueError(f"Unsupported PickUp approach case: {approach}") + + norm = math.sqrt(sum(value * value for value in direction)) + if norm < 1e-6: + raise ValueError(f"PickUp approach direction is zero for {position_case.name}.") + return tuple(value / norm for value in direction) + + +def resolve_pickup_approach_direction( + approach: str, + position_case: PositionCase, + device, +): + """Resolve a PickUp approach name into a normalized world-frame vector.""" + torch = ensure_torch() + return torch.tensor( + pickup_approach_direction_tuple(approach, position_case), + dtype=torch.float32, + device=device, + ) + + +def format_vector3(vector) -> str: + """Format a 3D vector-like value for benchmark reports.""" + if hasattr(vector, "detach"): + vector = vector.detach().to("cpu").tolist() + return f"({float(vector[0]):.3f},{float(vector[1]):.3f},{float(vector[2]):.3f})" + + +def _is_horizontal_approach_direction(approach_direction) -> bool: + """Return true when the requested approach is a side/horizontal approach.""" + if not hasattr(approach_direction, "detach"): + return abs(float(approach_direction[2])) < 1e-4 + direction = approach_direction.detach() + if direction.ndim > 1: + direction = direction.reshape(-1, direction.shape[-1])[0] + return abs(float(direction[2].to("cpu"))) < 1e-4 + + +def create_benchmark_object( + sim, + preset: MeshObjectPreset, + position_case: PositionCase, + uid_suffix: str, +): + """Create one benchmark object at a selected initial position.""" + from embodichain.data import get_data_path + from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg + from embodichain.lab.sim.shapes import CubeCfg, MeshCfg + + if preset.shape_type == "mesh": + shape = MeshCfg(fpath=get_data_path(preset.mesh_path)) + elif preset.shape_type == "cube": + if preset.cube_size is None: + raise ValueError(f"Cube preset {preset.object_type!r} misses cube_size.") + shape = CubeCfg(size=list(preset.cube_size)) + else: + raise ValueError( + f"Unsupported benchmark object shape_type {preset.shape_type!r}." + ) + + cfg = RigidObjectCfg( + uid=f"benchmark_{preset.label}_{position_case.name}_{uid_suffix}", + shape=shape, + attrs=RigidBodyAttributesCfg( + mass=preset.mass, + dynamic_friction=preset.dynamic_friction, + static_friction=preset.static_friction, + restitution=preset.restitution, + contact_offset=preset.contact_offset, + rest_offset=preset.rest_offset, + linear_damping=preset.linear_damping, + angular_damping=preset.angular_damping, + max_depenetration_velocity=preset.max_depenetration_velocity, + min_position_iters=preset.min_position_iters, + min_velocity_iters=preset.min_velocity_iters, + max_linear_velocity=preset.max_linear_velocity, + max_angular_velocity=preset.max_angular_velocity, + enable_ccd=preset.enable_ccd, + ), + max_convex_hull_num=preset.max_convex_hull_num, + init_pos=[position_case.xy[0], position_case.xy[1], preset.initial_z], + init_rot=preset.init_rot, + body_scale=preset.body_scale, + use_usd_properties=preset.use_usd_properties, + ) + obj = sim.add_rigid_object(cfg=cfg) + sim.update(step=10) + return obj + + +create_mesh_benchmark_object = create_benchmark_object + + +def _make_benchmark_antipodal_affordance_class(): + """Create an AntipodalAffordance subclass after project imports are available.""" + from embodichain.lab.sim.atomic_actions import AntipodalAffordance + + class BenchmarkAntipodalAffordance(AntipodalAffordance): + """Benchmark affordance that biases side grasps to horizontal closing.""" + + def get_valid_grasp_poses( + self, + obj_poses, + approach_direction, + ): + results = super().get_valid_grasp_poses( + obj_poses=obj_poses, + approach_direction=approach_direction, + ) + if not _is_horizontal_approach_direction(approach_direction): + return results + + adjusted_results = [] + for grasp_poses, costs in results: + if grasp_poses.ndim < 3 or grasp_poses.shape[0] == 0: + adjusted_results.append((grasp_poses, costs)) + continue + + opening_axis_abs_z = grasp_poses[:, :3, 0].abs()[:, 2] + keep = opening_axis_abs_z <= SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z + if bool(keep.any()): + adjusted_results.append((grasp_poses[keep], costs[keep])) + continue + + adjusted_costs = costs + ( + opening_axis_abs_z * SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT + ) + adjusted_results.append((grasp_poses, adjusted_costs)) + return adjusted_results + + return BenchmarkAntipodalAffordance + + +def create_antipodal_object_semantics( + obj, + preset: MeshObjectPreset, + args: argparse.Namespace, + build_gripper_collision_cfg: Callable[[], object], + build_grasp_generator_cfg: Callable[[argparse.Namespace], object], +): + """Create object semantics with an antipodal grasp affordance.""" + from embodichain.lab.sim.atomic_actions import ObjectSemantics + + mesh_vertices = obj.get_vertices(env_ids=[0], scale=True)[0] + mesh_triangles = obj.get_triangles(env_ids=[0])[0] + affordance_cls = _make_benchmark_antipodal_affordance_class() + return ObjectSemantics( + label=preset.label, + geometry={ + "mesh_vertices": mesh_vertices, + "mesh_triangles": mesh_triangles, + }, + affordance=affordance_cls( + mesh_vertices=mesh_vertices, + mesh_triangles=mesh_triangles, + gripper_collision_cfg=build_gripper_collision_cfg(), + generator_cfg=build_grasp_generator_cfg(args), + force_reannotate=args.force_reannotate, + ), + entity=obj, + ) + + +def describe_object_preset(preset: MeshObjectPreset) -> str: + """Describe an object preset for benchmark report notes.""" + if preset.shape_type == "cube": + return ( + f"{preset.object_type}/{preset.material_name}/" + f"CubeCfg(size={preset.cube_size})" + ) + return f"{preset.object_type}/{preset.material_name}/{preset.mesh_path}" + + +def sync_cuda() -> None: + """Synchronize CUDA stream when available.""" + torch = ensure_torch() + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def reset_peak_gpu_memory() -> None: + """Reset PyTorch peak GPU memory stats when CUDA is available.""" + torch = ensure_torch() + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + + +def peak_gpu_memory_mb() -> float: + """Return peak GPU memory allocated by PyTorch in MB.""" + torch = ensure_torch() + if not torch.cuda.is_available(): + return 0.0 + return torch.cuda.max_memory_allocated() / 1024**2 + + +def memory_snapshot() -> dict[str, float]: + """Return current process memory usage snapshot in MB.""" + torch = ensure_torch() + if psutil is not None: + process = psutil.Process(os.getpid()) + cpu_mb = process.memory_info().rss / 1024**2 + else: + cpu_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0 + gpu_mb = ( + torch.cuda.memory_allocated() / 1024**2 if torch.cuda.is_available() else 0.0 + ) + return {"cpu_mb": cpu_mb, "gpu_mb": gpu_mb} + + +def timed_call( + callable_fn: Callable[[], object], +) -> tuple[float, dict[str, float], float, object]: + """Time a callable and return elapsed seconds, memory deltas, peak GPU, result.""" + reset_peak_gpu_memory() + before = memory_snapshot() + sync_cuda() + + start = time.perf_counter() + result = callable_fn() + sync_cuda() + elapsed = time.perf_counter() - start + + after = memory_snapshot() + deltas = { + "cpu_mb": after["cpu_mb"] - before["cpu_mb"], + "gpu_mb": after["gpu_mb"] - before["gpu_mb"], + } + return elapsed, deltas, peak_gpu_memory_mb(), result + + +def reset_robot(robot, initial_qpos) -> None: + """Reset current and target robot qpos to the benchmark initial posture.""" + for target in (False, True): + robot.set_qpos(initial_qpos, target=target) + robot.clear_dynamics() + + +def reset_rigid_object(obj, initial_pose) -> None: + """Reset a rigid object pose and clear residual dynamics.""" + obj.set_local_pose(initial_pose) + obj.clear_dynamics() + + +def reset_rigid_object_xy( + obj, + base_pose, + xy: tuple[float, float], + sim=None, + settle_steps: int = 0, +): + """Reset a rigid object to a new XY position while preserving base orientation.""" + pose = base_pose.clone() + pose[:, 0, 3] = xy[0] + pose[:, 1, 3] = xy[1] + reset_rigid_object(obj, pose) + if sim is not None and settle_steps > 0: + sim.update(step=settle_steps) + return pose + + +def park_rigid_object( + obj, + base_pose, + index: int = 0, + sim=None, +) -> None: + """Move an inactive benchmark object outside the robot workspace.""" + pose = base_pose.clone() + pose[:, 0, 3] = 8.0 + float(index) + pose[:, 1, 3] = 8.0 + pose[:, 2, 3] = 1.0 + reset_rigid_object(obj, pose) + if sim is not None: + sim.update(step=1) + + +def object_position_tuple(obj) -> tuple[float, float, float]: + """Return the current object-frame origin position as a CPU tuple.""" + pose = obj.get_local_pose(to_matrix=True) + xyz = pose[0, :3, 3] + if hasattr(xyz, "detach"): + xyz = xyz.detach().to("cpu").tolist() + return (float(xyz[0]), float(xyz[1]), float(xyz[2])) + + +def xy_distance_m( + position: Sequence[float], + target: Sequence[float], +) -> float: + """Return XY Euclidean distance in meters.""" + return math.sqrt( + (float(position[0]) - float(target[0])) ** 2 + + (float(position[1]) - float(target[1])) ** 2 + ) + + +def xyz_distance_m( + position: Sequence[float], + target: Sequence[float], +) -> float: + """Return XYZ Euclidean distance in meters.""" + return math.sqrt( + (float(position[0]) - float(target[0])) ** 2 + + (float(position[1]) - float(target[1])) ** 2 + + (float(position[2]) - float(target[2])) ** 2 + ) + + +def replay_trajectory_for_physical_validation( + sim, + robot, + obj, + traj, + on_step: Callable[[int], None] | None = None, + hold_steps: int = PHYSICAL_VALIDATION_HOLD_STEPS, +) -> tuple[float, float, float] | None: + """Replay a trajectory in physics and return the final object position.""" + if traj is None or getattr(traj, "ndim", 0) < 3 or traj.shape[1] == 0: + return None + + for waypoint_index in range(traj.shape[1]): + robot.set_qpos(traj[:, waypoint_index, :]) + sim.update(step=4) + if on_step is not None: + on_step(waypoint_index) + + final_qpos = traj[:, -1, :] + for _ in range(hold_steps): + robot.set_qpos(final_qpos) + sim.update(step=2) + return object_position_tuple(obj) + + +def should_record_case( + args: argparse.Namespace, + recorded_count: int, + success: bool, +) -> bool: + """Return whether a benchmark case should emit a replay video.""" + if not getattr(args, "record_video", False): + return False + if not success and not getattr(args, "record_failed_video", False): + return False + + case_limit = getattr(args, "video_case_limit", DEFAULT_VIDEO_CASE_LIMIT) + if case_limit < 0: + raise ValueError("--video_case_limit must be non-negative.") + return case_limit == 0 or recorded_count < case_limit + + +def build_video_output_path( + args: argparse.Namespace, + benchmark_name: str, + case_id: str, +) -> Path: + """Build a deterministic, timestamped output path for one replay video.""" + output_dir = Path(getattr(args, "video_dir", DEFAULT_VIDEO_DIR)) + output_dir.mkdir(parents=True, exist_ok=True) + safe_case_id = re.sub(r"[^A-Za-z0-9_.-]+", "_", case_id).strip("_") + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + return output_dir / f"{benchmark_name}_{safe_case_id}_{timestamp}.mp4" + + +def summarize_video_recording( + args: argparse.Namespace, + results: Sequence[dict[str, object]], + video_paths: Sequence[str], +) -> list[str]: + """Build report notes that make video coverage explicit.""" + evaluated_count = len(results) + success_count = sum(1 for result in results if bool(result.get("success"))) + failure_count = evaluated_count - success_count + notes = [ + ( + "Case/video summary: " + f"evaluated={evaluated_count}, success={success_count}, " + f"failure={failure_count}, videos={len(video_paths)}" + ) + ] + if getattr(args, "record_video", False): + if getattr(args, "record_failed_video", False): + notes.append( + "Video policy: records report-success replays and failed-case " + "debug videos when trajectory/static scene capture is available." + ) + else: + notes.append( + "Video policy: records report-success replays only; failed " + "cases are reported in the tables but do not emit videos." + ) + else: + notes.append("Video policy: disabled.") + notes.append( + "Replay videos: " + (", ".join(video_paths) if video_paths else "disabled") + ) + return notes + + +def _replay_trajectory_with_recording( + sim, + robot, + traj, + args: argparse.Namespace, + video_path: Path, + on_step: Callable[[int], None] | None = None, + look_at: tuple[ + Sequence[float], + Sequence[float], + Sequence[float], + ] = DEFAULT_VIDEO_LOOK_AT, +) -> Path | None: + """Replay a planned trajectory and record it with the simulation recorder.""" + if traj is None or getattr(traj, "ndim", 0) < 3 or traj.shape[1] == 0: + return None + + video_fps = getattr(args, "video_fps", DEFAULT_VIDEO_FPS) + video_max_memory = getattr(args, "video_max_memory", DEFAULT_VIDEO_MAX_MEMORY_MB) + video_width = getattr(args, "video_width", DEFAULT_VIDEO_WIDTH) + video_height = getattr(args, "video_height", DEFAULT_VIDEO_HEIGHT) + video_hold_steps = getattr(args, "video_hold_steps", DEFAULT_VIDEO_HOLD_STEPS) + + original_width = sim.sim_config.width + original_height = sim.sim_config.height + recording_started = False + try: + sim.sim_config.width = video_width + sim.sim_config.height = video_height + recording_started = sim.start_window_record( + save_path=str(video_path), + fps=video_fps, + max_memory=video_max_memory, + look_at=look_at, + use_sim_time=True, + ) + finally: + sim.sim_config.width = original_width + sim.sim_config.height = original_height + + if not recording_started: + return None + + stop_success = False + try: + for waypoint_index in range(traj.shape[1]): + robot.set_qpos(traj[:, waypoint_index, :]) + sim.update(step=4) + if on_step is not None: + on_step(waypoint_index) + + final_qpos = traj[:, -1, :] + for _ in range(video_hold_steps): + robot.set_qpos(final_qpos) + sim.update(step=2) + finally: + if sim.is_window_recording(): + stop_success = sim.stop_window_record() + sim.wait_window_record_saves() + + return video_path if stop_success else None + + +def _record_static_scene_video( + sim, + args: argparse.Namespace, + video_path: Path, + look_at: tuple[ + Sequence[float], + Sequence[float], + Sequence[float], + ] = DEFAULT_VIDEO_LOOK_AT, +) -> Path | None: + """Record the current scene without replaying a planned trajectory.""" + video_fps = getattr(args, "video_fps", DEFAULT_VIDEO_FPS) + video_max_memory = getattr(args, "video_max_memory", DEFAULT_VIDEO_MAX_MEMORY_MB) + video_width = getattr(args, "video_width", DEFAULT_VIDEO_WIDTH) + video_height = getattr(args, "video_height", DEFAULT_VIDEO_HEIGHT) + video_hold_steps = getattr(args, "video_hold_steps", DEFAULT_VIDEO_HOLD_STEPS) + + original_width = sim.sim_config.width + original_height = sim.sim_config.height + recording_started = False + try: + sim.sim_config.width = video_width + sim.sim_config.height = video_height + recording_started = sim.start_window_record( + save_path=str(video_path), + fps=video_fps, + max_memory=video_max_memory, + look_at=look_at, + use_sim_time=True, + ) + finally: + sim.sim_config.width = original_width + sim.sim_config.height = original_height + + if not recording_started: + return None + + stop_success = False + try: + for _ in range(video_hold_steps): + sim.update(step=2) + finally: + if sim.is_window_recording(): + stop_success = sim.stop_window_record() + sim.wait_window_record_saves() + + return video_path if stop_success else None + + +def replay_trajectory_with_recording( + sim, + robot, + traj, + args: argparse.Namespace, + video_path: Path, + on_step: Callable[[int], None] | None = None, + look_at: tuple[ + Sequence[float], + Sequence[float], + Sequence[float], + ] = DEFAULT_VIDEO_LOOK_AT, +) -> Path | None: + """Best-effort replay recording that never changes benchmark success.""" + try: + return _replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=traj, + args=args, + video_path=video_path, + on_step=on_step, + look_at=look_at, + ) + except Exception as exc: + try: + if sim.is_window_recording(): + sim.stop_window_record() + sim.wait_window_record_saves() + except Exception: + pass + print( + "Warning: failed to record benchmark replay video " + f"{video_path}: {type(exc).__name__}: {exc}" + ) + return None + + +def record_static_scene_video( + sim, + args: argparse.Namespace, + video_path: Path, + look_at: tuple[ + Sequence[float], + Sequence[float], + Sequence[float], + ] = DEFAULT_VIDEO_LOOK_AT, +) -> Path | None: + """Best-effort static scene recording that never changes benchmark success.""" + try: + return _record_static_scene_video( + sim=sim, + args=args, + video_path=video_path, + look_at=look_at, + ) + except Exception as exc: + try: + if sim.is_window_recording(): + sim.stop_window_record() + sim.wait_window_record_saves() + except Exception: + pass + print( + "Warning: failed to record benchmark static debug video " + f"{video_path}: {type(exc).__name__}: {exc}" + ) + return None + + +def format_float(value: float | None, precision: int = 6) -> str: + """Format finite floats for tables and use N/A for missing values.""" + if value is None or not math.isfinite(value): + return "N/A" + return f"{value:.{precision}f}" + + +def format_markdown_table(rows: list[dict[str, object]]) -> list[str]: + """Format rows into a markdown table.""" + if not rows: + return ["No data."] + + headers = list(rows[0].keys()) + lines = [ + "| " + " | ".join(headers) + " |", + "| " + " | ".join(["---"] * len(headers)) + " |", + ] + for row in rows: + lines.append("| " + " | ".join(str(row[h]) for h in headers) + " |") + return lines + + +def write_markdown_report( + benchmark_name: str, + perf_rows: list[dict[str, object]], + metric_rows: list[dict[str, object]], + leaderboard_rows: list[dict[str, object]], + notes: list[str] | None = None, +) -> Path: + """Write benchmark results into one markdown report file.""" + output_dir = Path("outputs/benchmarks") + output_dir.mkdir(parents=True, exist_ok=True) + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + report_path = output_dir / f"{benchmark_name}_{timestamp}.md" + + lines: list[str] = [ + f"# {benchmark_name} Benchmark Report", + "", + f"Generated at: {datetime.now().isoformat(timespec='seconds')}", + "", + "## Time & Memory", + "", + ] + lines.extend(format_markdown_table(perf_rows)) + lines.extend(["", "## Success & Other Metrics", ""]) + lines.extend(format_markdown_table(metric_rows)) + lines.extend(["", "## Leaderboard", ""]) + lines.extend(format_markdown_table(leaderboard_rows)) + + if notes: + lines.extend(["", "## Notes", ""]) + lines.extend([f"- {note}" for note in notes]) + + report_path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return report_path + + +def build_single_action_leaderboard( + action_name: str, + metric_rows: list[dict[str, object]], +) -> list[dict[str, object]]: + """Aggregate rows for one action by benchmark case.""" + if not metric_rows: + return [] + + success_sum = sum(float(row["success_rate"]) for row in metric_rows) + count = len(metric_rows) + return [ + { + "rank": 1, + "algorithm": action_name, + "overall_success_rate": f"{success_sum / max(count, 1):.2%}", + "evaluated_cases": count, + } + ] + + +__all__ = [ + "BENCHMARK_PROFILES", + "CPU_MEMORY_BACKEND", + "COVERAGE_MESH_OBJECT_TYPES", + "COVERAGE_POSITION_CASE_NAMES", + "DEFAULT_BENCHMARK_PROFILE", + "add_common_benchmark_args", + "add_grasp_benchmark_args", + "add_object_position_benchmark_args", + "add_profile_benchmark_args", + "add_video_benchmark_args", + "build_video_output_path", + "build_single_action_leaderboard", + "create_antipodal_object_semantics", + "create_benchmark_object", + "create_mesh_benchmark_object", + "DEFAULT_VIDEO_DIR", + "default_pickup_approach_cases_for_profile", + "default_mesh_object_types_for_profile", + "default_position_case_names_for_profile", + "describe_object_preset", + "ensure_repo_root", + "ensure_torch", + "format_float", + "format_vector3", + "FULL_MESH_OBJECT_TYPES", + "FULL_POSITION_CASE_NAMES", + "MESH_OBJECT_PRESETS", + "MeshObjectPreset", + "PICKUP_APPROACH_CASES", + "PHYSICAL_MOVE_HELD_OBJECT_XYZ_TOLERANCE_M", + "PHYSICAL_PICK_MIN_LIFT_M", + "PHYSICAL_PLACE_XY_TOLERANCE_M", + "PHYSICAL_VALIDATION_HOLD_STEPS", + "POSITION_CASES", + "PositionCase", + "object_position_tuple", + "park_rigid_object", + "pickup_approach_direction_tuple", + "record_static_scene_video", + "replay_trajectory_for_physical_validation", + "replay_trajectory_with_recording", + "reset_rigid_object", + "reset_rigid_object_xy", + "reset_robot", + "resolve_pickup_approach_direction", + "resolve_profile", + "select_mesh_object_presets", + "select_pickup_approaches", + "select_position_cases", + "should_record_case", + "SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z", + "SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT", + "SMOKE_PICKUP_APPROACH_CASES", + "SMOKE_MESH_OBJECT_TYPES", + "SMOKE_POSITION_CASE_NAMES", + "timed_call", + "summarize_video_recording", + "write_markdown_report", + "xy_distance_m", + "xyz_distance_m", +] diff --git a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py new file mode 100644 index 000000000..6267e5418 --- /dev/null +++ b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py @@ -0,0 +1,318 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Benchmark MoveEndEffector atomic action across target poses. + +Measures planning latency, memory usage, trajectory success, and final TCP +translation error for several reachable pose targets. +Run: python -m scripts.benchmark.atomic_action.move_end_effector_benchmark +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + CPU_MEMORY_BACKEND, + add_common_benchmark_args, + build_single_action_leaderboard, + build_video_output_path, + ensure_repo_root, + ensure_torch, + format_float, + replay_trajectory_with_recording, + reset_robot, + resolve_profile, + should_record_case, + timed_call, + write_markdown_report, +) + + +@dataclass(frozen=True) +class PoseCase: + """End-effector target pose benchmark case.""" + + name: str + xyz: tuple[float, float, float] + + +POSE_CASES = { + "front_left": PoseCase("front_left", (0.30, -0.20, 0.36)), + "front_right": PoseCase("front_right", (0.30, 0.20, 0.36)), + "near_center": PoseCase("near_center", (0.18, 0.00, 0.42)), + "far_center": PoseCase("far_center", (0.42, 0.00, 0.34)), +} +DEFAULT_POSE_CASES = tuple(POSE_CASES.keys()) +MOVE_SAMPLE_INTERVAL = 80 +SUCCESS_TOLERANCE_M = 0.01 + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add MoveEndEffector benchmark CLI arguments.""" + parser.add_argument( + "--pose_cases", + nargs="+", + choices=(*POSE_CASES.keys(), "all"), + default=list(DEFAULT_POSE_CASES), + help="Target pose cases to benchmark. Use 'all' for every case.", + ) + add_common_benchmark_args(parser) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Benchmark MoveEndEffector over reachable target poses." + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _select_pose_cases(case_names: list[str]) -> list[PoseCase]: + """Resolve selected pose case names.""" + if "all" in case_names: + return list(POSE_CASES.values()) + return [POSE_CASES[name] for name in case_names] + + +def _make_pose(device, xyz: tuple[float, float, float]): + """Create a top-down target pose.""" + torch = ensure_torch() + pose = torch.eye(4, dtype=torch.float32, device=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=device, + ) + pose[:3, 3] = torch.tensor(xyz, dtype=torch.float32, device=device) + return pose + + +def _run_case( + sim, + robot, + atomic_engine, + initial_qpos, + pose_case: PoseCase, + repeat: int, + args: argparse.Namespace, + recorded_count: int, +): + """Run one MoveEndEffector case.""" + torch = ensure_torch() + from embodichain.lab.sim.atomic_actions import EndEffectorPoseTarget + + reset_robot(robot, initial_qpos) + target_pose = _make_pose(sim.device, pose_case.xyz) + + elapsed, mem_delta, peak_gpu, result = timed_call( + lambda: atomic_engine.run( + steps=[("move_end_effector", EndEffectorPoseTarget(xpos=target_pose))] + ) + ) + is_success, traj, _ = result + video_path = None + if should_record_case(args, recorded_count, bool(is_success)): + reset_robot(robot, initial_qpos) + video_path = replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=traj, + args=args, + video_path=build_video_output_path( + args, "atomic_action_move_end_effector", f"{pose_case.name}_r{repeat}" + ), + ) + reset_robot(robot, initial_qpos) + + final_error_m = None + if is_success and traj.shape[1] > 0: + arm_joint_ids = robot.get_joint_ids(name="arm") + final_arm_qpos = traj[:, -1, arm_joint_ids] + final_pose = robot.compute_fk(qpos=final_arm_qpos, name="arm", to_matrix=True)[ + 0 + ] + final_error_m = float(torch.linalg.norm(final_pose[:3, 3] - target_pose[:3, 3])) + target_reached = bool( + is_success + and final_error_m is not None + and final_error_m <= SUCCESS_TOLERANCE_M + ) + return { + "case_id": f"{pose_case.name}:r{repeat}", + "target_case": pose_case.name, + "repeat": repeat, + "planning_success": bool(is_success), + "target_reached": target_reached, + "cost_time_ms": elapsed * 1000.0, + "cpu_delta_mb": mem_delta["cpu_mb"], + "gpu_delta_mb": mem_delta["gpu_mb"], + "peak_gpu_mb": peak_gpu, + "final_error_m": final_error_m, + "trajectory_waypoints": int(traj.shape[1]) if traj.ndim >= 2 else 0, + "failure_reason": "" if target_reached else "target_not_reached", + "video_path": str(video_path) if video_path is not None else "", + } + + +def _build_rows(results: list[dict[str, object]]): + """Build report rows for MoveEndEffector.""" + perf_rows = [] + metric_rows = [] + for result in results: + perf_rows.append( + { + "sample_size": 1, + "impl": "move_end_effector", + "case_id": result["case_id"], + "target_case": result["target_case"], + "repeat": result["repeat"], + "cost_time_ms": format_float(result["cost_time_ms"]), + "cpu_delta_mb": format_float(result["cpu_delta_mb"]), + "gpu_delta_mb": format_float(result["gpu_delta_mb"]), + "peak_gpu_mb": format_float(result["peak_gpu_mb"]), + } + ) + metric_rows.append( + { + "sample_size": 1, + "impl": "move_end_effector", + "case_id": result["case_id"], + "target_case": result["target_case"], + "success_rate": f"{float(result['target_reached']):.6f}", + "planning_success_rate": f"{float(result['planning_success']):.6f}", + "translation_err_m": format_float(result["final_error_m"]), + "trajectory_waypoints": result["trajectory_waypoints"], + "failure_reason": result["failure_reason"] or "N/A", + } + ) + return perf_rows, metric_rows + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: + """Run MoveEndEffector benchmark and write a markdown report.""" + args = _parse_args() if args is None else args + if args.repeat < 1: + raise ValueError("--repeat must be at least 1.") + profile = resolve_profile(args) + + ensure_repo_root() + ensure_torch() + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + MoveEndEffector, + MoveEndEffectorCfg, + ) + from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.planners import ToppraPlannerCfg + from scripts.tutorials.atomic_action.move_end_effector import ( + create_robot, + initialize_simulation, + ) + + pose_cases = _select_pose_cases(args.pose_cases) + repeat = 1 if profile == "smoke" else args.repeat + if profile == "smoke": + pose_cases = [POSE_CASES["front_left"]] + + print("=" * 60) + print("MoveEndEffector Atomic Action Benchmark") + print("=" * 60) + print( + "Coverage: " + f"profile={profile}, {len(pose_cases)} pose case(s) x " + f"{repeat} repeat(s)" + ) + + sim = initialize_simulation(args) + robot = create_robot(sim) + initial_qpos = robot.get_qpos().clone() + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + MoveEndEffector( + motion_gen, + cfg=MoveEndEffectorCfg( + control_part="arm", sample_interval=MOVE_SAMPLE_INTERVAL + ), + ) + ) + + results: list[dict[str, object]] = [] + video_paths: list[str] = [] + print("\n=== MoveEndEffector Target Sweep ===") + for pose_case in pose_cases: + for repeat_index in range(repeat): + result = _run_case( + sim, + robot, + atomic_engine, + initial_qpos, + pose_case, + repeat_index, + args, + len(video_paths), + ) + results.append(result) + if result["video_path"]: + video_paths.append(str(result["video_path"])) + print( + f" {result['case_id']:<24} " + f"time={result['cost_time_ms']:>10.2f} ms | " + f"success={result['target_reached']} " + f"err={format_float(result['final_error_m'], precision=4)}" + ) + + perf_rows, metric_rows = _build_rows(results) + leaderboard_rows = build_single_action_leaderboard("move_end_effector", metric_rows) + report_path = write_markdown_report( + benchmark_name="atomic_action_move_end_effector", + perf_rows=perf_rows, + metric_rows=metric_rows, + leaderboard_rows=leaderboard_rows, + notes=[ + f"Profile: {profile}", + f"CPU memory backend: {CPU_MEMORY_BACKEND}", + f"Success tolerance: {SUCCESS_TOLERANCE_M} m translation error.", + "Replay videos: " + (", ".join(video_paths) if video_paths else "disabled"), + ], + ) + print(f"Markdown report saved: {report_path}") + return report_path + + +def main() -> None: + """Run the CLI entry point.""" + try: + run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/move_held_object_benchmark.py b/scripts/benchmark/atomic_action/move_held_object_benchmark.py new file mode 100644 index 000000000..cabfb1582 --- /dev/null +++ b/scripts/benchmark/atomic_action/move_held_object_benchmark.py @@ -0,0 +1,766 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Benchmark MoveHeldObject atomic action after a PickUp precondition. + +Measures MoveHeldObject-only planning latency and memory usage once a held +object state has been produced by PickUp. +Run: python -m scripts.benchmark.atomic_action.move_held_object_benchmark +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + CPU_MEMORY_BACKEND, + add_common_benchmark_args, + add_grasp_benchmark_args, + add_object_position_benchmark_args, + add_pickup_approach_benchmark_args, + build_single_action_leaderboard, + build_video_output_path, + create_antipodal_object_semantics, + create_benchmark_object, + describe_object_preset, + ensure_repo_root, + ensure_torch, + format_float, + format_vector3, + MeshObjectPreset, + object_position_tuple, + park_rigid_object, + PHYSICAL_MOVE_HELD_OBJECT_XYZ_TOLERANCE_M, + PHYSICAL_PICK_MIN_LIFT_M, + pickup_approach_direction_tuple, + PositionCase, + record_static_scene_video, + replay_trajectory_for_physical_validation, + replay_trajectory_with_recording, + reset_rigid_object, + reset_rigid_object_xy, + reset_robot, + resolve_pickup_approach_direction, + resolve_profile, + select_mesh_object_presets, + select_pickup_approaches, + select_position_cases, + should_record_case, + SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z, + SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT, + summarize_video_recording, + timed_call, + write_markdown_report, + xy_distance_m, + xyz_distance_m, +) + + +@dataclass(frozen=True) +class HeldObjectCase: + """Held-object target pose benchmark case.""" + + name: str + xyz: tuple[float, float, float] + + +HELD_OBJECT_CASES = { + "front_left": HeldObjectCase("front_left", (-0.30, -0.30, 0.50)), + "front_right": HeldObjectCase("front_right", (-0.30, 0.30, 0.50)), + "center_high": HeldObjectCase("center_high", (-0.42, 0.00, 0.58)), +} +PICK_SAMPLE_INTERVAL = 120 +MOVE_SAMPLE_INTERVAL = 60 +MOVE_HELD_OBJECT_SAMPLE_INTERVAL = 120 +HAND_INTERP_STEPS = 12 + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add MoveHeldObject benchmark CLI arguments.""" + add_pickup_approach_benchmark_args(parser) + parser.add_argument( + "--held_object_cases", + nargs="+", + choices=(*HELD_OBJECT_CASES.keys(), "all"), + default=None, + help=( + "Held-object target cases to benchmark. Defaults are selected by " + "--profile; use 'all' for every case." + ), + ) + add_object_position_benchmark_args(parser) + add_grasp_benchmark_args(parser) + add_common_benchmark_args(parser) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Benchmark MoveHeldObject after a PickUp precondition." + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _selected_cases( + case_names: list[str] | None, + profile: str, +) -> list[HeldObjectCase]: + """Resolve selected held-object case names.""" + if not case_names: + return [HELD_OBJECT_CASES["front_left"]] + if "all" in case_names: + return list(HELD_OBJECT_CASES.values()) + return [HELD_OBJECT_CASES[name] for name in case_names] + + +def _make_object_target_pose(device, xyz: tuple[float, float, float]): + """Create a held-object target pose.""" + torch = ensure_torch() + 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(xyz, dtype=torch.float32, device=device) + return pose + + +def _make_pickup_args( + args: argparse.Namespace, + object_preset: MeshObjectPreset, + profile: str, +) -> argparse.Namespace: + """Build a tutorial-compatible argparse namespace.""" + return argparse.Namespace( + object=object_preset.label, + n_sample=1000 if profile == "smoke" else args.n_sample, + force_reannotate=args.force_reannotate, + device=args.device, + renderer=args.renderer, + headless=True, + ) + + +def _prepare_held_state( + sim, + robot, + obj, + motion_gen, + args, + object_preset: MeshObjectPreset, + position_case: PositionCase, + pickup_approach: str, + profile: str, +): + """Run PickUp precondition outside the timed MoveHeldObject block.""" + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + EndEffectorPoseTarget, + GraspTarget, + MoveEndEffector, + MoveEndEffectorCfg, + PickUp, + PickUpCfg, + ) + from scripts.tutorials.atomic_action.move_held_object import ( + build_grasp_generator_cfg, + build_gripper_collision_cfg, + get_hand_open_close_qpos, + make_pre_pick_eef_pose, + ) + + hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + MoveEndEffector( + motion_gen, + cfg=MoveEndEffectorCfg( + control_part="arm", + sample_interval=MOVE_SAMPLE_INTERVAL, + ), + ) + ) + atomic_engine.register( + PickUp( + motion_gen, + cfg=PickUpCfg( + control_part="arm", + hand_control_part="hand", + hand_open_qpos=hand_open, + hand_close_qpos=hand_close, + approach_direction=resolve_pickup_approach_direction( + pickup_approach, position_case, sim.device + ), + pre_grasp_distance=0.15, + lift_height=0.16, + sample_interval=PICK_SAMPLE_INTERVAL, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ) + ) + semantics = create_antipodal_object_semantics( + obj=obj, + preset=object_preset, + args=_make_pickup_args(args, object_preset, profile), + build_gripper_collision_cfg=build_gripper_collision_cfg, + build_grasp_generator_cfg=build_grasp_generator_cfg, + ) + obj_pose = obj.get_local_pose(to_matrix=True) + move_position = obj_pose[0, :3, 3].clone() + move_position[2] = 0.36 + move_target = make_pre_pick_eef_pose(robot, move_position) + is_success, traj, state = atomic_engine.run( + steps=[ + ("move_end_effector", EndEffectorPoseTarget(xpos=move_target)), + ("pick_up", GraspTarget(semantics=semantics)), + ] + ) + if not is_success or state.held_object is None: + raise RuntimeError( + "Failed to prepare held-object state for MoveHeldObject benchmark." + ) + robot.set_qpos(state.last_qpos) + return state, hand_close, traj + + +def _run_case( + sim, + robot, + motion_gen, + initial_qpos, + args, + obj, + base_obj_pose, + object_preset: MeshObjectPreset, + position_case: PositionCase, + pickup_approach: str, + case: HeldObjectCase, + repeat: int, + recorded_count: int, + profile: str, +): + """Run one MoveHeldObject benchmark case.""" + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + HeldObjectPoseTarget, + MoveHeldObject, + MoveHeldObjectCfg, + ) + from scripts.tutorials.atomic_action.move_held_object import ( + compute_pick_close_end_step, + ) + + case_id = ( + f"{object_preset.object_type}:{position_case.name}:" + f"{pickup_approach}:{case.name}:r{repeat}" + ) + try: + reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + reset_rigid_object(obj, initial_obj_pose) + initial_obj_position = object_position_tuple(obj) + state, hand_close, precondition_traj = _prepare_held_state( + sim, + robot, + obj, + motion_gen, + args, + object_preset, + position_case, + pickup_approach, + profile, + ) + approach_direction_text = format_vector3( + pickup_approach_direction_tuple(pickup_approach, position_case) + ) + precondition_waypoints = int(precondition_traj.shape[1]) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + MoveHeldObject( + motion_gen, + cfg=MoveHeldObjectCfg( + control_part="arm", + hand_control_part="hand", + hand_close_qpos=hand_close, + sample_interval=MOVE_HELD_OBJECT_SAMPLE_INTERVAL, + ), + ) + ) + target_pose = _make_object_target_pose(sim.device, case.xyz) + elapsed, mem_delta, peak_gpu, result = timed_call( + lambda: atomic_engine.run( + steps=[ + ( + "move_held_object", + HeldObjectPoseTarget(object_target_pose=target_pose), + ) + ], + state=state, + ) + ) + is_success, traj, final_state = result + torch = ensure_torch() + precondition_obj_position = None + final_obj_position = None + object_lift_delta_m = None + held_object_xy_error_m = None + held_object_z_error_m = None + held_object_xyz_error_m = None + + if ( + getattr(precondition_traj, "ndim", 0) >= 3 + and precondition_traj.shape[1] > 0 + ): + if bool(is_success) and getattr(traj, "ndim", 0) >= 3 and traj.shape[1] > 0: + validation_traj = torch.cat((precondition_traj, traj), dim=1) + else: + validation_traj = precondition_traj + + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + post_grasp_clear_step = compute_pick_close_end_step() + should_clear_object_dynamics = True + + def _on_validation_step(waypoint_index: int) -> None: + nonlocal should_clear_object_dynamics, precondition_obj_position + if ( + should_clear_object_dynamics + and waypoint_index + 1 >= post_grasp_clear_step + ): + obj.clear_dynamics() + should_clear_object_dynamics = False + if ( + precondition_obj_position is None + and waypoint_index + 1 >= precondition_waypoints + ): + precondition_obj_position = object_position_tuple(obj) + + final_obj_position = replay_trajectory_for_physical_validation( + sim=sim, + robot=robot, + obj=obj, + traj=validation_traj, + on_step=_on_validation_step, + ) + if precondition_obj_position is not None: + object_lift_delta_m = ( + precondition_obj_position[2] - initial_obj_position[2] + ) + if bool(is_success) and final_obj_position is not None: + held_object_xy_error_m = xy_distance_m(final_obj_position, case.xyz) + held_object_z_error_m = abs(final_obj_position[2] - case.xyz[2]) + held_object_xyz_error_m = xyz_distance_m(final_obj_position, case.xyz) + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + still_held = bool(is_success and final_state.held_object is not None) + physical_pick_success = bool( + object_lift_delta_m is not None + and object_lift_delta_m >= PHYSICAL_PICK_MIN_LIFT_M + ) + physical_move_success = bool( + still_held + and physical_pick_success + and held_object_xyz_error_m is not None + and held_object_xyz_error_m <= PHYSICAL_MOVE_HELD_OBJECT_XYZ_TOLERANCE_M + ) + success = physical_move_success + if success: + failure_reason = "" + elif not is_success: + failure_reason = "planning_failed" + elif not still_held: + failure_reason = "held_object_lost" + elif not physical_pick_success: + failure_reason = "physical_pick_failed" + elif held_object_xyz_error_m is None: + failure_reason = "physical_replay_missing" + else: + failure_reason = "held_object_target_miss" + + video_path = None + if should_record_case(args, recorded_count, success): + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + video_case_suffix = ( + f"{object_preset.object_type}_{position_case.name}_" + f"{pickup_approach}_{case.name}_r{repeat}" + ) + if getattr(traj, "ndim", 0) >= 3 and traj.shape[1] > 0: + replay_traj = torch.cat((precondition_traj, traj), dim=1) + else: + replay_traj = precondition_traj + + if getattr(replay_traj, "ndim", 0) >= 3 and replay_traj.shape[1] > 0: + post_grasp_clear_step = compute_pick_close_end_step() + should_clear_object_dynamics = True + + def _on_step(waypoint_index: int) -> None: + nonlocal should_clear_object_dynamics + if ( + should_clear_object_dynamics + and waypoint_index + 1 >= post_grasp_clear_step + ): + obj.clear_dynamics() + should_clear_object_dynamics = False + + video_path = replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=replay_traj, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_move_held_object", + video_case_suffix, + ), + on_step=_on_step, + ) + else: + video_path = record_static_scene_video( + sim=sim, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_move_held_object", + f"{video_case_suffix}_failed_static", + ), + ) + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + return { + "case_id": case_id, + "object_type": object_preset.object_type, + "material": object_preset.material_name, + "quadrant": position_case.quadrant, + "position_case": position_case.name, + "init_xy": position_case.xy, + "pickup_approach": pickup_approach, + "approach_direction": approach_direction_text, + "held_object_case": case.name, + "repeat": repeat, + "planning_success": bool(is_success), + "still_held": still_held, + "physical_pick_success": physical_pick_success, + "physical_move_success": physical_move_success, + "success": success, + "cost_time_ms": elapsed * 1000.0, + "cpu_delta_mb": mem_delta["cpu_mb"], + "gpu_delta_mb": mem_delta["gpu_mb"], + "peak_gpu_mb": peak_gpu, + "object_lift_delta_m": object_lift_delta_m, + "object_final_x_m": ( + final_obj_position[0] if final_obj_position is not None else None + ), + "object_final_y_m": ( + final_obj_position[1] if final_obj_position is not None else None + ), + "object_final_z_m": ( + final_obj_position[2] if final_obj_position is not None else None + ), + "held_object_xy_error_m": held_object_xy_error_m, + "held_object_z_error_m": held_object_z_error_m, + "held_object_xyz_error_m": held_object_xyz_error_m, + "precondition_waypoints": precondition_waypoints, + "trajectory_waypoints": int(traj.shape[1]) if traj.ndim >= 2 else 0, + "failure_reason": failure_reason, + "video_path": str(video_path) if video_path is not None else "", + } + except Exception as exc: + video_path = None + if should_record_case(args, recorded_count, False): + try: + reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + video_path = record_static_scene_video( + sim=sim, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_move_held_object", + ( + f"{object_preset.object_type}_{position_case.name}_" + f"{pickup_approach}_{case.name}_r{repeat}" + "_exception_static" + ), + ), + ) + reset_rigid_object(obj, initial_obj_pose) + except Exception: + video_path = None + return { + "case_id": case_id, + "object_type": object_preset.object_type, + "material": object_preset.material_name, + "quadrant": position_case.quadrant, + "position_case": position_case.name, + "init_xy": position_case.xy, + "pickup_approach": pickup_approach, + "approach_direction": "N/A", + "held_object_case": case.name, + "repeat": repeat, + "planning_success": False, + "still_held": False, + "physical_pick_success": False, + "physical_move_success": False, + "success": False, + "cost_time_ms": 0.0, + "cpu_delta_mb": 0.0, + "gpu_delta_mb": 0.0, + "peak_gpu_mb": 0.0, + "object_lift_delta_m": None, + "object_final_x_m": None, + "object_final_y_m": None, + "object_final_z_m": None, + "held_object_xy_error_m": None, + "held_object_z_error_m": None, + "held_object_xyz_error_m": None, + "precondition_waypoints": 0, + "trajectory_waypoints": 0, + "failure_reason": f"exception:{type(exc).__name__}:{exc}", + "video_path": str(video_path) if video_path is not None else "", + } + + +def _build_rows(results: list[dict[str, object]]): + """Build report rows for MoveHeldObject.""" + perf_rows = [] + metric_rows = [] + for result in results: + perf_rows.append( + { + "sample_size": 1, + "impl": "move_held_object", + "case_id": result["case_id"], + "object_type": result["object_type"], + "material": result["material"], + "quadrant": result["quadrant"], + "position_case": result["position_case"], + "init_xy": f"({result['init_xy'][0]:.3f},{result['init_xy'][1]:.3f})", + "pickup_approach": result["pickup_approach"], + "approach_direction": result["approach_direction"], + "held_object_case": result["held_object_case"], + "repeat": result["repeat"], + "cost_time_ms": format_float(result["cost_time_ms"]), + "cpu_delta_mb": format_float(result["cpu_delta_mb"]), + "gpu_delta_mb": format_float(result["gpu_delta_mb"]), + "peak_gpu_mb": format_float(result["peak_gpu_mb"]), + } + ) + metric_rows.append( + { + "sample_size": 1, + "impl": "move_held_object", + "case_id": result["case_id"], + "object_type": result["object_type"], + "material": result["material"], + "quadrant": result["quadrant"], + "position_case": result["position_case"], + "pickup_approach": result["pickup_approach"], + "approach_direction": result["approach_direction"], + "held_object_case": result["held_object_case"], + "success_rate": f"{float(result['success']):.6f}", + "planning_success_rate": f"{float(result['planning_success']):.6f}", + "held_object_rate": f"{float(result['still_held']):.6f}", + "physical_pick_success_rate": ( + f"{float(result['physical_pick_success']):.6f}" + ), + "physical_move_success_rate": ( + f"{float(result['physical_move_success']):.6f}" + ), + "object_lift_delta_m": format_float(result["object_lift_delta_m"]), + "object_final_x_m": format_float(result["object_final_x_m"]), + "object_final_y_m": format_float(result["object_final_y_m"]), + "object_final_z_m": format_float(result["object_final_z_m"]), + "held_object_xy_error_m": format_float( + result["held_object_xy_error_m"] + ), + "held_object_z_error_m": format_float(result["held_object_z_error_m"]), + "held_object_xyz_error_m": format_float( + result["held_object_xyz_error_m"] + ), + "precondition_waypoints": result["precondition_waypoints"], + "trajectory_waypoints": result["trajectory_waypoints"], + "failure_reason": result["failure_reason"] or "N/A", + } + ) + return perf_rows, metric_rows + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: + """Run MoveHeldObject benchmark and write a markdown report.""" + args = _parse_args() if args is None else args + if args.repeat < 1: + raise ValueError("--repeat must be at least 1.") + profile = resolve_profile(args) + + ensure_repo_root() + ensure_torch() + from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.planners import ToppraPlannerCfg + from scripts.tutorials.atomic_action.move_held_object import ( + create_robot, + initialize_simulation, + ) + + object_presets = select_mesh_object_presets(args.object_types, profile) + position_cases = select_position_cases(args.position_cases, profile) + approaches = select_pickup_approaches(args.approach_cases, profile) + cases = _selected_cases(args.held_object_cases, profile) + repeat = 1 if profile == "smoke" else args.repeat + + print("=" * 60) + print("MoveHeldObject Atomic Action Benchmark") + print("=" * 60) + print( + "Coverage: " + f"profile={profile}, {len(object_presets)} object(s) x " + f"{len(position_cases)} position(s) x {len(approaches)} pickup approach(es) " + f"x {len(cases)} held target(s) " + f"x {repeat} repeat(s)" + ) + + sim = initialize_simulation(args) + robot = create_robot(sim) + initial_qpos = robot.get_qpos().clone() + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + object_pool = {} + for object_index, object_preset in enumerate(object_presets): + obj = create_benchmark_object( + sim=sim, + preset=object_preset, + position_case=position_cases[0], + uid_suffix="pool", + ) + base_pose = obj.get_local_pose(to_matrix=True).clone() + park_rigid_object(obj, base_pose, index=object_index, sim=sim) + object_pool[object_preset.object_type] = (obj, base_pose) + + results: list[dict[str, object]] = [] + video_paths: list[str] = [] + print("\n=== MoveHeldObject Object/Position/Target Sweep ===") + for object_preset in object_presets: + obj, base_pose = object_pool[object_preset.object_type] + for parked_index, parked_preset in enumerate(object_presets): + if parked_preset.object_type == object_preset.object_type: + continue + parked_obj, parked_base_pose = object_pool[parked_preset.object_type] + park_rigid_object(parked_obj, parked_base_pose, index=parked_index, sim=sim) + for position_case in position_cases: + for pickup_approach in approaches: + for case in cases: + for repeat_index in range(repeat): + result = _run_case( + sim, + robot, + motion_gen, + initial_qpos, + args, + obj, + base_pose, + object_preset, + position_case, + pickup_approach, + case, + repeat_index, + len(video_paths), + profile, + ) + results.append(result) + if result["video_path"]: + video_paths.append(str(result["video_path"])) + print( + f" {result['case_id']:<54} " + f"time={result['cost_time_ms']:>10.2f} ms | " + f"success={result['success']}" + ) + + perf_rows, metric_rows = _build_rows(results) + leaderboard_rows = build_single_action_leaderboard("move_held_object", metric_rows) + report_path = write_markdown_report( + benchmark_name="atomic_action_move_held_object", + perf_rows=perf_rows, + metric_rows=metric_rows, + leaderboard_rows=leaderboard_rows, + notes=[ + "Timed block includes MoveHeldObject only; PickUp precondition is " + "prepared outside timing.", + f"Profile: {profile}", + "Object presets: " + + ", ".join(describe_object_preset(preset) for preset in object_presets), + "Position cases: " + + ", ".join( + f"{case.name}/{case.quadrant}/xy={case.xy}" for case in position_cases + ), + "PickUp approach cases: " + ", ".join(approaches), + "Held-object target cases: " + ", ".join(case.name for case in cases), + f"CPU memory backend: {CPU_MEMORY_BACKEND}", + f"n_sample: {1000 if profile == 'smoke' else args.n_sample}", + ( + "Side grasp opening-axis rule: prefer candidates with " + f"abs(opening_axis_z) <= {SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z:.2f}; " + "fallback cost += " + f"abs(opening_axis_z) * {SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT:.2f}." + ), + ( + "Physical success rule: PickUp precondition lift delta >= " + f"{PHYSICAL_PICK_MIN_LIFT_M:.3f} m, held-object state is kept, " + "and final object XYZ error <= " + f"{PHYSICAL_MOVE_HELD_OBJECT_XYZ_TOLERANCE_M:.3f} m." + ), + *summarize_video_recording(args, results, video_paths), + ], + ) + print(f"Markdown report saved: {report_path}") + return report_path + + +def main() -> None: + """Run the CLI entry point.""" + try: + run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/move_joints_benchmark.py b/scripts/benchmark/atomic_action/move_joints_benchmark.py new file mode 100644 index 000000000..25a1f3a05 --- /dev/null +++ b/scripts/benchmark/atomic_action/move_joints_benchmark.py @@ -0,0 +1,331 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Benchmark MoveJoints atomic action across joint-space target sequences. + +Measures planning latency, memory usage, trajectory success, and final arm joint +error for named and explicit joint-position targets. +Run: python -m scripts.benchmark.atomic_action.move_joints_benchmark +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + CPU_MEMORY_BACKEND, + add_common_benchmark_args, + build_single_action_leaderboard, + build_video_output_path, + ensure_repo_root, + ensure_torch, + format_float, + replay_trajectory_with_recording, + reset_robot, + resolve_profile, + should_record_case, + timed_call, + write_markdown_report, +) + + +@dataclass(frozen=True) +class JointSequenceCase: + """MoveJoints target sequence benchmark case.""" + + name: str + sequence: tuple[str, ...] + + +READY_QPOS = (0.35, -1.20, 1.30, -1.65, -1.57, 0.20) +HOME_QPOS = (0.0, -1.57, 1.57, -1.57, -1.57, 0.0) +EXTENDED_QPOS = (0.55, -1.05, 1.10, -1.45, -1.35, 0.35) +JOINT_TARGETS = { + "ready": READY_QPOS, + "home": HOME_QPOS, + "extended": EXTENDED_QPOS, +} +SEQUENCE_CASES = { + "ready_home": JointSequenceCase("ready_home", ("ready", "home")), + "extended_home": JointSequenceCase("extended_home", ("extended", "home")), + "ready_extended_home": JointSequenceCase( + "ready_extended_home", ("ready", "extended", "home") + ), +} +DEFAULT_SEQUENCE_CASES = tuple(SEQUENCE_CASES.keys()) +MOVE_JOINTS_SAMPLE_INTERVAL = 80 +SUCCESS_TOLERANCE_RAD = 1e-4 + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add MoveJoints benchmark CLI arguments.""" + parser.add_argument( + "--sequence_cases", + nargs="+", + choices=(*SEQUENCE_CASES.keys(), "all"), + default=list(DEFAULT_SEQUENCE_CASES), + help="Joint sequence cases to benchmark. Use 'all' for every case.", + ) + add_common_benchmark_args(parser) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Benchmark MoveJoints over named and explicit qpos targets." + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _select_cases(case_names: list[str]) -> list[JointSequenceCase]: + """Resolve selected sequence case names.""" + if "all" in case_names: + return list(SEQUENCE_CASES.values()) + return [SEQUENCE_CASES[name] for name in case_names] + + +def _qpos(values, device): + """Create a torch qpos tensor.""" + torch = ensure_torch() + return torch.tensor(values, dtype=torch.float32, device=device) + + +def _targets_for_sequence(sequence_case: JointSequenceCase, device): + """Build typed MoveJoints targets for a sequence case.""" + from embodichain.lab.sim.atomic_actions import ( + JointPositionTarget, + NamedJointPositionTarget, + ) + + targets = [] + for index, name in enumerate(sequence_case.sequence): + if index == 0 and name == "ready": + targets.append(("move_joints", NamedJointPositionTarget(name="ready"))) + else: + targets.append( + ( + "move_joints", + JointPositionTarget(qpos=_qpos(JOINT_TARGETS[name], device)), + ) + ) + return targets + + +def _run_case( + sim, + robot, + atomic_engine, + initial_qpos, + case: JointSequenceCase, + repeat: int, + args: argparse.Namespace, + recorded_count: int, +): + """Run one MoveJoints case.""" + torch = ensure_torch() + reset_robot(robot, initial_qpos) + steps = _targets_for_sequence(case, sim.device) + elapsed, mem_delta, peak_gpu, result = timed_call(lambda: atomic_engine.run(steps)) + is_success, traj, _ = result + video_path = None + if should_record_case(args, recorded_count, bool(is_success)): + reset_robot(robot, initial_qpos) + video_path = replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=traj, + args=args, + video_path=build_video_output_path( + args, "atomic_action_move_joints", f"{case.name}_r{repeat}" + ), + ) + reset_robot(robot, initial_qpos) + + final_error_rad = None + if is_success and traj.shape[1] > 0: + arm_joint_ids = robot.get_joint_ids(name="arm") + final_qpos = traj[:, -1, arm_joint_ids][0] + expected = _qpos(JOINT_TARGETS[case.sequence[-1]], sim.device) + final_error_rad = float(torch.linalg.norm(final_qpos - expected)) + target_reached = bool( + is_success + and final_error_rad is not None + and final_error_rad <= SUCCESS_TOLERANCE_RAD + ) + return { + "case_id": f"{case.name}:r{repeat}", + "sequence_case": case.name, + "repeat": repeat, + "planning_success": bool(is_success), + "target_reached": target_reached, + "cost_time_ms": elapsed * 1000.0, + "cpu_delta_mb": mem_delta["cpu_mb"], + "gpu_delta_mb": mem_delta["gpu_mb"], + "peak_gpu_mb": peak_gpu, + "final_error_rad": final_error_rad, + "trajectory_waypoints": int(traj.shape[1]) if traj.ndim >= 2 else 0, + "failure_reason": "" if target_reached else "target_not_reached", + "video_path": str(video_path) if video_path is not None else "", + } + + +def _build_rows(results: list[dict[str, object]]): + """Build report rows for MoveJoints.""" + perf_rows = [] + metric_rows = [] + for result in results: + perf_rows.append( + { + "sample_size": 1, + "impl": "move_joints", + "case_id": result["case_id"], + "sequence_case": result["sequence_case"], + "repeat": result["repeat"], + "cost_time_ms": format_float(result["cost_time_ms"]), + "cpu_delta_mb": format_float(result["cpu_delta_mb"]), + "gpu_delta_mb": format_float(result["gpu_delta_mb"]), + "peak_gpu_mb": format_float(result["peak_gpu_mb"]), + } + ) + metric_rows.append( + { + "sample_size": 1, + "impl": "move_joints", + "case_id": result["case_id"], + "sequence_case": result["sequence_case"], + "success_rate": f"{float(result['target_reached']):.6f}", + "planning_success_rate": f"{float(result['planning_success']):.6f}", + "joint_err_rad": format_float(result["final_error_rad"]), + "trajectory_waypoints": result["trajectory_waypoints"], + "failure_reason": result["failure_reason"] or "N/A", + } + ) + return perf_rows, metric_rows + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: + """Run MoveJoints benchmark and write a markdown report.""" + args = _parse_args() if args is None else args + if args.repeat < 1: + raise ValueError("--repeat must be at least 1.") + profile = resolve_profile(args) + + ensure_repo_root() + ensure_torch() + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + MoveJoints, + MoveJointsCfg, + ) + from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.planners import ToppraPlannerCfg + from scripts.tutorials.atomic_action.move_joints import ( + create_robot, + initialize_simulation, + ) + + cases = _select_cases(args.sequence_cases) + repeat = 1 if profile == "smoke" else args.repeat + if profile == "smoke": + cases = [SEQUENCE_CASES["ready_home"]] + + print("=" * 60) + print("MoveJoints Atomic Action Benchmark") + print("=" * 60) + print( + "Coverage: " + f"profile={profile}, {len(cases)} sequence case(s) x " + f"{repeat} repeat(s)" + ) + + sim = initialize_simulation(args) + robot = create_robot(sim) + initial_qpos = robot.get_qpos().clone() + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + ready_qpos = _qpos(READY_QPOS, sim.device) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + MoveJoints( + motion_gen, + cfg=MoveJointsCfg( + control_part="arm", + sample_interval=MOVE_JOINTS_SAMPLE_INTERVAL, + named_joint_positions={"ready": ready_qpos}, + ), + ) + ) + + results: list[dict[str, object]] = [] + video_paths: list[str] = [] + print("\n=== MoveJoints Sequence Sweep ===") + for case in cases: + for repeat_index in range(repeat): + result = _run_case( + sim, + robot, + atomic_engine, + initial_qpos, + case, + repeat_index, + args, + len(video_paths), + ) + results.append(result) + if result["video_path"]: + video_paths.append(str(result["video_path"])) + print( + f" {result['case_id']:<24} " + f"time={result['cost_time_ms']:>10.2f} ms | " + f"success={result['target_reached']} " + f"err={format_float(result['final_error_rad'], precision=6)}" + ) + + perf_rows, metric_rows = _build_rows(results) + leaderboard_rows = build_single_action_leaderboard("move_joints", metric_rows) + report_path = write_markdown_report( + benchmark_name="atomic_action_move_joints", + perf_rows=perf_rows, + metric_rows=metric_rows, + leaderboard_rows=leaderboard_rows, + notes=[ + f"Profile: {profile}", + f"CPU memory backend: {CPU_MEMORY_BACKEND}", + f"Success tolerance: {SUCCESS_TOLERANCE_RAD} rad joint error.", + "Replay videos: " + (", ".join(video_paths) if video_paths else "disabled"), + ], + ) + print(f"Markdown report saved: {report_path}") + return report_path + + +def main() -> None: + """Run the CLI entry point.""" + try: + run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/pickup_benchmark.py b/scripts/benchmark/atomic_action/pickup_benchmark.py new file mode 100644 index 000000000..37395ca83 --- /dev/null +++ b/scripts/benchmark/atomic_action/pickup_benchmark.py @@ -0,0 +1,572 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Benchmark PickUp atomic action across grasp approach directions. + +Measures planning latency, memory usage, grasp planning success, held-object +state creation, and trajectory length for the PickUp action. +Run: python -m scripts.benchmark.atomic_action.pickup_benchmark +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + CPU_MEMORY_BACKEND, + add_common_benchmark_args, + add_grasp_benchmark_args, + add_object_position_benchmark_args, + add_pickup_approach_benchmark_args, + build_single_action_leaderboard, + build_video_output_path, + create_antipodal_object_semantics, + create_benchmark_object, + describe_object_preset, + ensure_repo_root, + ensure_torch, + format_float, + format_vector3, + MeshObjectPreset, + object_position_tuple, + park_rigid_object, + PHYSICAL_PICK_MIN_LIFT_M, + pickup_approach_direction_tuple, + PositionCase, + record_static_scene_video, + replay_trajectory_for_physical_validation, + replay_trajectory_with_recording, + reset_rigid_object, + reset_rigid_object_xy, + reset_robot, + resolve_pickup_approach_direction, + resolve_profile, + select_mesh_object_presets, + select_pickup_approaches, + select_position_cases, + should_record_case, + SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z, + SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT, + summarize_video_recording, + timed_call, + write_markdown_report, +) + +PICK_SAMPLE_INTERVAL = 120 +HAND_INTERP_STEPS = 12 + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add PickUp benchmark CLI arguments.""" + add_pickup_approach_benchmark_args(parser) + add_object_position_benchmark_args(parser) + add_grasp_benchmark_args(parser) + add_common_benchmark_args(parser) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Benchmark PickUp over grasp approach directions." + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _make_pickup_args( + args: argparse.Namespace, + approach: str, + object_preset: MeshObjectPreset, + profile: str, +) -> argparse.Namespace: + """Build a tutorial-compatible argparse namespace.""" + return argparse.Namespace( + object=object_preset.label, + n_sample=1000 if profile == "smoke" else args.n_sample, + force_reannotate=args.force_reannotate, + approach="custom" if approach == "side" else approach, + custom_approach_direction=None, + device=args.device, + renderer=args.renderer, + headless=True, + ) + + +def _run_case( + sim, + robot, + motion_gen, + initial_qpos, + args, + obj, + base_obj_pose, + object_preset: MeshObjectPreset, + position_case: PositionCase, + approach: str, + repeat: int, + recorded_count: int, + profile: str, +): + """Run one PickUp benchmark case.""" + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + GraspTarget, + PickUp, + PickUpCfg, + ) + from scripts.tutorials.atomic_action.pickup import ( + build_grasp_generator_cfg, + build_gripper_collision_cfg, + get_hand_open_close_qpos, + initialize_pre_pick_robot_pose, + compute_pick_close_end_step, + ) + + case_id = f"{object_preset.object_type}:{position_case.name}:{approach}:r{repeat}" + try: + reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + initial_obj_position = object_position_tuple(obj) + hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + case_args = _make_pickup_args(args, approach, object_preset, profile) + approach_direction = resolve_pickup_approach_direction( + approach, position_case, sim.device + ) + approach_direction_text = format_vector3( + pickup_approach_direction_tuple(approach, position_case) + ) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + PickUp( + motion_gen, + 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, + ), + ) + ) + semantics = create_antipodal_object_semantics( + obj=obj, + preset=object_preset, + args=case_args, + build_gripper_collision_cfg=build_gripper_collision_cfg, + build_grasp_generator_cfg=build_grasp_generator_cfg, + ) + elapsed, mem_delta, peak_gpu, result = timed_call( + lambda: atomic_engine.run( + steps=[("pick_up", GraspTarget(semantics=semantics))] + ) + ) + is_success, traj, final_state = result + final_obj_position = None + object_lift_delta_m = None + object_xy_drift_m = None + + if bool(is_success) and getattr(traj, "ndim", 0) >= 3 and traj.shape[1] > 0: + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + post_grasp_clear_step = compute_pick_close_end_step() + should_clear_object_dynamics = True + + def _on_validation_step(waypoint_index: int) -> None: + nonlocal should_clear_object_dynamics + if ( + should_clear_object_dynamics + and waypoint_index + 1 >= post_grasp_clear_step + ): + obj.clear_dynamics() + should_clear_object_dynamics = False + + final_obj_position = replay_trajectory_for_physical_validation( + sim=sim, + robot=robot, + obj=obj, + traj=traj, + on_step=_on_validation_step, + ) + if final_obj_position is not None: + object_lift_delta_m = final_obj_position[2] - initial_obj_position[2] + object_xy_drift_m = ( + (final_obj_position[0] - initial_obj_position[0]) ** 2 + + (final_obj_position[1] - initial_obj_position[1]) ** 2 + ) ** 0.5 + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + held_created = bool(is_success and final_state.held_object is not None) + physical_pick_success = bool( + held_created + and object_lift_delta_m is not None + and object_lift_delta_m >= PHYSICAL_PICK_MIN_LIFT_M + ) + + video_path = None + if should_record_case(args, recorded_count, physical_pick_success): + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + video_case_suffix = ( + f"{object_preset.object_type}_{position_case.name}_" + f"{approach}_r{repeat}" + ) + if getattr(traj, "ndim", 0) >= 3 and traj.shape[1] > 0: + post_grasp_clear_step = compute_pick_close_end_step() + should_clear_object_dynamics = True + + def _on_step(waypoint_index: int) -> None: + nonlocal should_clear_object_dynamics + if ( + should_clear_object_dynamics + and waypoint_index + 1 >= post_grasp_clear_step + ): + obj.clear_dynamics() + should_clear_object_dynamics = False + + video_path = replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=traj, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_pick_up", + video_case_suffix, + ), + on_step=_on_step, + ) + else: + video_path = record_static_scene_video( + sim=sim, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_pick_up", + f"{video_case_suffix}_failed_static", + ), + ) + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + lift_height_m = None + if is_success and traj.shape[1] > 0: + arm_joint_ids = robot.get_joint_ids(name="arm") + start_pose = robot.compute_fk( + qpos=traj[:, 0, arm_joint_ids], name="arm", to_matrix=True + )[0] + final_pose = robot.compute_fk( + qpos=traj[:, -1, arm_joint_ids], name="arm", to_matrix=True + )[0] + lift_height_m = float(final_pose[2, 3] - start_pose[2, 3]) + success = physical_pick_success + if success: + failure_reason = "" + elif not is_success: + failure_reason = "planning_failed" + elif not held_created: + failure_reason = "held_object_missing" + elif object_lift_delta_m is None: + failure_reason = "physical_replay_missing" + else: + failure_reason = "physical_pick_lift_too_low" + return { + "case_id": case_id, + "object_type": object_preset.object_type, + "material": object_preset.material_name, + "quadrant": position_case.quadrant, + "position_case": position_case.name, + "init_xy": position_case.xy, + "approach": approach, + "approach_direction": approach_direction_text, + "repeat": repeat, + "planning_success": bool(is_success), + "held_created": held_created, + "physical_pick_success": physical_pick_success, + "success": success, + "cost_time_ms": elapsed * 1000.0, + "cpu_delta_mb": mem_delta["cpu_mb"], + "gpu_delta_mb": mem_delta["gpu_mb"], + "peak_gpu_mb": peak_gpu, + "lift_height_m": lift_height_m, + "object_initial_z_m": initial_obj_position[2], + "object_final_z_m": ( + final_obj_position[2] if final_obj_position is not None else None + ), + "object_lift_delta_m": object_lift_delta_m, + "object_xy_drift_m": object_xy_drift_m, + "trajectory_waypoints": int(traj.shape[1]) if traj.ndim >= 2 else 0, + "failure_reason": failure_reason, + "video_path": str(video_path) if video_path is not None else "", + } + except Exception as exc: + video_path = None + if should_record_case(args, recorded_count, False): + try: + reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + video_path = record_static_scene_video( + sim=sim, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_pick_up", + ( + f"{object_preset.object_type}_{position_case.name}_" + f"{approach}_r{repeat}_exception_static" + ), + ), + ) + reset_rigid_object(obj, initial_obj_pose) + except Exception: + video_path = None + return { + "case_id": case_id, + "object_type": object_preset.object_type, + "material": object_preset.material_name, + "quadrant": position_case.quadrant, + "position_case": position_case.name, + "init_xy": position_case.xy, + "approach": approach, + "approach_direction": "N/A", + "repeat": repeat, + "planning_success": False, + "held_created": False, + "physical_pick_success": False, + "success": False, + "cost_time_ms": 0.0, + "cpu_delta_mb": 0.0, + "gpu_delta_mb": 0.0, + "peak_gpu_mb": 0.0, + "lift_height_m": None, + "object_initial_z_m": None, + "object_final_z_m": None, + "object_lift_delta_m": None, + "object_xy_drift_m": None, + "trajectory_waypoints": 0, + "failure_reason": f"exception:{type(exc).__name__}:{exc}", + "video_path": str(video_path) if video_path is not None else "", + } + + +def _build_rows(results: list[dict[str, object]]): + """Build report rows for PickUp.""" + perf_rows = [] + metric_rows = [] + for result in results: + perf_rows.append( + { + "sample_size": 1, + "impl": "pick_up", + "case_id": result["case_id"], + "object_type": result["object_type"], + "material": result["material"], + "quadrant": result["quadrant"], + "position_case": result["position_case"], + "init_xy": f"({result['init_xy'][0]:.3f},{result['init_xy'][1]:.3f})", + "approach": result["approach"], + "approach_direction": result["approach_direction"], + "repeat": result["repeat"], + "cost_time_ms": format_float(result["cost_time_ms"]), + "cpu_delta_mb": format_float(result["cpu_delta_mb"]), + "gpu_delta_mb": format_float(result["gpu_delta_mb"]), + "peak_gpu_mb": format_float(result["peak_gpu_mb"]), + } + ) + metric_rows.append( + { + "sample_size": 1, + "impl": "pick_up", + "case_id": result["case_id"], + "object_type": result["object_type"], + "material": result["material"], + "quadrant": result["quadrant"], + "position_case": result["position_case"], + "approach": result["approach"], + "approach_direction": result["approach_direction"], + "success_rate": f"{float(result['success']):.6f}", + "planning_success_rate": f"{float(result['planning_success']):.6f}", + "held_object_rate": f"{float(result['held_created']):.6f}", + "physical_pick_success_rate": ( + f"{float(result['physical_pick_success']):.6f}" + ), + "lift_height_m": format_float(result["lift_height_m"]), + "object_initial_z_m": format_float(result["object_initial_z_m"]), + "object_final_z_m": format_float(result["object_final_z_m"]), + "object_lift_delta_m": format_float(result["object_lift_delta_m"]), + "object_xy_drift_m": format_float(result["object_xy_drift_m"]), + "trajectory_waypoints": result["trajectory_waypoints"], + "failure_reason": result["failure_reason"] or "N/A", + } + ) + return perf_rows, metric_rows + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: + """Run PickUp benchmark and write a markdown report.""" + args = _parse_args() if args is None else args + if args.repeat < 1: + raise ValueError("--repeat must be at least 1.") + profile = resolve_profile(args) + + ensure_repo_root() + ensure_torch() + from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.planners import ToppraPlannerCfg + from scripts.tutorials.atomic_action.pickup import ( + create_robot, + initialize_simulation, + ) + + object_presets = select_mesh_object_presets(args.object_types, profile) + position_cases = select_position_cases(args.position_cases, profile) + approaches = select_pickup_approaches(args.approach_cases, profile) + repeat = 1 if profile == "smoke" else args.repeat + + print("=" * 60) + print("PickUp Atomic Action Benchmark") + print("=" * 60) + print( + "Coverage: " + f"profile={profile}, {len(object_presets)} object(s) x " + f"{len(position_cases)} position(s) x {len(approaches)} approach(es) " + f"x {repeat} repeat(s)" + ) + + sim = initialize_simulation(args) + robot = create_robot(sim) + initial_qpos = robot.get_qpos().clone() + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + object_pool = {} + for object_index, object_preset in enumerate(object_presets): + obj = create_benchmark_object( + sim=sim, + preset=object_preset, + position_case=position_cases[0], + uid_suffix="pool", + ) + base_pose = obj.get_local_pose(to_matrix=True).clone() + park_rigid_object(obj, base_pose, index=object_index, sim=sim) + object_pool[object_preset.object_type] = (obj, base_pose) + + results: list[dict[str, object]] = [] + video_paths: list[str] = [] + print("\n=== PickUp Object/Position/Approach Sweep ===") + for object_preset in object_presets: + obj, base_pose = object_pool[object_preset.object_type] + for parked_index, parked_preset in enumerate(object_presets): + if parked_preset.object_type == object_preset.object_type: + continue + parked_obj, parked_base_pose = object_pool[parked_preset.object_type] + park_rigid_object(parked_obj, parked_base_pose, index=parked_index, sim=sim) + for position_case in position_cases: + for approach in approaches: + for repeat_index in range(repeat): + result = _run_case( + sim, + robot, + motion_gen, + initial_qpos, + args, + obj, + base_pose, + object_preset, + position_case, + approach, + repeat_index, + len(video_paths), + profile, + ) + results.append(result) + if result["video_path"]: + video_paths.append(str(result["video_path"])) + print( + f" {result['case_id']:<42} " + f"time={result['cost_time_ms']:>10.2f} ms | " + f"success={result['success']} " + f"lift={format_float(result['lift_height_m'], precision=4)}" + ) + + perf_rows, metric_rows = _build_rows(results) + leaderboard_rows = build_single_action_leaderboard("pick_up", metric_rows) + report_path = write_markdown_report( + benchmark_name="atomic_action_pick_up", + perf_rows=perf_rows, + metric_rows=metric_rows, + leaderboard_rows=leaderboard_rows, + notes=[ + f"Profile: {profile}", + "Object presets: " + + ", ".join(describe_object_preset(preset) for preset in object_presets), + "Position cases: " + + ", ".join( + f"{case.name}/{case.quadrant}/xy={case.xy}" for case in position_cases + ), + "PickUp approach cases: " + ", ".join(approaches), + f"CPU memory backend: {CPU_MEMORY_BACKEND}", + f"n_sample: {1000 if profile == 'smoke' else args.n_sample}", + ( + "Side grasp opening-axis rule: prefer candidates with " + f"abs(opening_axis_z) <= {SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z:.2f}; " + "fallback cost += " + f"abs(opening_axis_z) * {SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT:.2f}." + ), + ( + "Physical success rule: planning success, held-object state " + f"created, and object lift delta >= {PHYSICAL_PICK_MIN_LIFT_M:.3f} m." + ), + *summarize_video_recording(args, results, video_paths), + ], + ) + print(f"Markdown report saved: {report_path}") + return report_path + + +def main() -> None: + """Run the CLI entry point.""" + try: + run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/place_benchmark.py b/scripts/benchmark/atomic_action/place_benchmark.py new file mode 100644 index 000000000..6ea98d68d --- /dev/null +++ b/scripts/benchmark/atomic_action/place_benchmark.py @@ -0,0 +1,736 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Benchmark Place atomic action after a PickUp precondition. + +Measures Place-only planning latency and memory usage once a held-object state +has been produced by the PickUp action. +Run: python -m scripts.benchmark.atomic_action.place_benchmark +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + CPU_MEMORY_BACKEND, + add_common_benchmark_args, + add_grasp_benchmark_args, + add_object_position_benchmark_args, + add_pickup_approach_benchmark_args, + build_single_action_leaderboard, + build_video_output_path, + create_antipodal_object_semantics, + create_benchmark_object, + describe_object_preset, + ensure_repo_root, + ensure_torch, + format_float, + format_vector3, + MeshObjectPreset, + object_position_tuple, + park_rigid_object, + PHYSICAL_PICK_MIN_LIFT_M, + PHYSICAL_PLACE_XY_TOLERANCE_M, + pickup_approach_direction_tuple, + PositionCase, + record_static_scene_video, + replay_trajectory_for_physical_validation, + replay_trajectory_with_recording, + reset_rigid_object, + reset_rigid_object_xy, + reset_robot, + resolve_pickup_approach_direction, + resolve_profile, + select_mesh_object_presets, + select_pickup_approaches, + select_position_cases, + should_record_case, + SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z, + SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT, + summarize_video_recording, + timed_call, + write_markdown_report, + xy_distance_m, +) + + +@dataclass(frozen=True) +class PlaceCase: + """Place target benchmark case.""" + + name: str + xyz: tuple[float, float, float] + + +PLACE_CASES = { + "left_bin": PlaceCase("left_bin", (-0.20, 0.28, 0.10)), + "right_bin": PlaceCase("right_bin", (-0.20, -0.28, 0.10)), + "center": PlaceCase("center", (-0.35, 0.00, 0.12)), +} +PICK_SAMPLE_INTERVAL = 120 +PLACE_SAMPLE_INTERVAL = 120 +HAND_INTERP_STEPS = 12 +PLACE_LIFT_HEIGHT = 0.14 + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add Place benchmark CLI arguments.""" + add_pickup_approach_benchmark_args(parser) + parser.add_argument( + "--place_cases", + nargs="+", + choices=(*PLACE_CASES.keys(), "all"), + default=None, + help=( + "Place target cases to benchmark. Defaults are selected by " + "--profile; use 'all' for every case." + ), + ) + add_object_position_benchmark_args(parser) + add_grasp_benchmark_args(parser) + add_common_benchmark_args(parser) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Benchmark Place after a PickUp precondition." + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _selected_cases( + case_names: list[str] | None, + profile: str, +) -> list[PlaceCase]: + """Resolve selected place case names.""" + if not case_names: + return [PLACE_CASES["left_bin"]] + if "all" in case_names: + return list(PLACE_CASES.values()) + return [PLACE_CASES[name] for name in case_names] + + +def _make_place_pose(device, xyz: tuple[float, float, float]): + """Create a top-down place target pose.""" + torch = ensure_torch() + pose = torch.eye(4, dtype=torch.float32, device=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=device, + ) + pose[:3, 3] = torch.tensor(xyz, dtype=torch.float32, device=device) + return pose + + +def _make_pickup_args( + args: argparse.Namespace, + object_preset: MeshObjectPreset, + profile: str, +) -> argparse.Namespace: + """Build a tutorial-compatible argparse namespace.""" + return argparse.Namespace( + object=object_preset.label, + n_sample=1000 if profile == "smoke" else args.n_sample, + force_reannotate=args.force_reannotate, + device=args.device, + renderer=args.renderer, + headless=True, + ) + + +def _prepare_held_state( + sim, + robot, + obj, + motion_gen, + args, + object_preset: MeshObjectPreset, + position_case: PositionCase, + pickup_approach: str, + profile: str, +): + """Run PickUp precondition outside the timed Place block.""" + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + GraspTarget, + PickUp, + PickUpCfg, + ) + from scripts.tutorials.atomic_action.place import ( + build_grasp_generator_cfg, + build_gripper_collision_cfg, + get_hand_open_close_qpos, + initialize_pre_pick_robot_pose, + ) + + hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + PickUp( + motion_gen, + cfg=PickUpCfg( + control_part="arm", + hand_control_part="hand", + hand_open_qpos=hand_open, + hand_close_qpos=hand_close, + approach_direction=resolve_pickup_approach_direction( + pickup_approach, position_case, sim.device + ), + pre_grasp_distance=0.15, + lift_height=0.16, + sample_interval=PICK_SAMPLE_INTERVAL, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ) + ) + semantics = create_antipodal_object_semantics( + obj=obj, + preset=object_preset, + args=_make_pickup_args(args, object_preset, profile), + build_gripper_collision_cfg=build_gripper_collision_cfg, + build_grasp_generator_cfg=build_grasp_generator_cfg, + ) + is_success, traj, state = atomic_engine.run( + steps=[("pick_up", GraspTarget(semantics=semantics))] + ) + if not is_success or state.held_object is None: + raise RuntimeError("Failed to prepare held-object state for Place benchmark.") + robot.set_qpos(state.last_qpos) + return state, hand_open, hand_close, traj + + +def _run_case( + sim, + robot, + motion_gen, + initial_qpos, + args, + obj, + base_obj_pose, + object_preset: MeshObjectPreset, + position_case: PositionCase, + pickup_approach: str, + case: PlaceCase, + repeat: int, + recorded_count: int, + profile: str, +): + """Run one Place benchmark case.""" + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + EndEffectorPoseTarget, + Place, + PlaceCfg, + ) + from scripts.tutorials.atomic_action.place import ( + compute_pick_close_end_step, + initialize_pre_pick_robot_pose, + ) + + case_id = ( + f"{object_preset.object_type}:{position_case.name}:" + f"{pickup_approach}:{case.name}:r{repeat}" + ) + try: + reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + reset_rigid_object(obj, initial_obj_pose) + initial_obj_position = object_position_tuple(obj) + state, hand_open, hand_close, precondition_traj = _prepare_held_state( + sim, + robot, + obj, + motion_gen, + args, + object_preset, + position_case, + pickup_approach, + profile, + ) + approach_direction_text = format_vector3( + pickup_approach_direction_tuple(pickup_approach, position_case) + ) + precondition_waypoints = int(precondition_traj.shape[1]) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + Place( + motion_gen, + 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, + ), + ) + ) + place_pose = _make_place_pose(sim.device, case.xyz) + elapsed, mem_delta, peak_gpu, result = timed_call( + lambda: atomic_engine.run( + steps=[("place", EndEffectorPoseTarget(xpos=place_pose))], + state=state, + ) + ) + is_success, traj, final_state = result + torch = ensure_torch() + precondition_obj_position = None + final_obj_position = None + object_lift_delta_m = None + place_xy_error_m = None + place_z_error_m = None + + if ( + getattr(precondition_traj, "ndim", 0) >= 3 + and precondition_traj.shape[1] > 0 + ): + if bool(is_success) and getattr(traj, "ndim", 0) >= 3 and traj.shape[1] > 0: + validation_traj = torch.cat((precondition_traj, traj), dim=1) + else: + validation_traj = precondition_traj + + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + post_grasp_clear_step = compute_pick_close_end_step() + should_clear_object_dynamics = True + + def _on_validation_step(waypoint_index: int) -> None: + nonlocal should_clear_object_dynamics, precondition_obj_position + if ( + should_clear_object_dynamics + and waypoint_index + 1 >= post_grasp_clear_step + ): + obj.clear_dynamics() + should_clear_object_dynamics = False + if ( + precondition_obj_position is None + and waypoint_index + 1 >= precondition_waypoints + ): + precondition_obj_position = object_position_tuple(obj) + + final_obj_position = replay_trajectory_for_physical_validation( + sim=sim, + robot=robot, + obj=obj, + traj=validation_traj, + on_step=_on_validation_step, + ) + if precondition_obj_position is not None: + object_lift_delta_m = ( + precondition_obj_position[2] - initial_obj_position[2] + ) + if bool(is_success) and final_obj_position is not None: + place_xy_error_m = xy_distance_m(final_obj_position, case.xyz) + place_z_error_m = abs(final_obj_position[2] - case.xyz[2]) + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + released = bool(is_success and final_state.held_object is None) + physical_pick_success = bool( + object_lift_delta_m is not None + and object_lift_delta_m >= PHYSICAL_PICK_MIN_LIFT_M + ) + physical_place_success = bool( + released + and physical_pick_success + and place_xy_error_m is not None + and place_xy_error_m <= PHYSICAL_PLACE_XY_TOLERANCE_M + ) + success = physical_place_success + if success: + failure_reason = "" + elif not is_success: + failure_reason = "planning_failed" + elif not released: + failure_reason = "held_object_not_released" + elif not physical_pick_success: + failure_reason = "physical_pick_failed" + elif place_xy_error_m is None: + failure_reason = "physical_replay_missing" + else: + failure_reason = "place_target_miss" + + video_path = None + if should_record_case(args, recorded_count, success): + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + video_case_suffix = ( + f"{object_preset.object_type}_{position_case.name}_" + f"{pickup_approach}_{case.name}_r{repeat}" + ) + if getattr(traj, "ndim", 0) >= 3 and traj.shape[1] > 0: + replay_traj = torch.cat((precondition_traj, traj), dim=1) + else: + replay_traj = precondition_traj + + if getattr(replay_traj, "ndim", 0) >= 3 and replay_traj.shape[1] > 0: + post_grasp_clear_step = compute_pick_close_end_step() + should_clear_object_dynamics = True + + def _on_step(waypoint_index: int) -> None: + nonlocal should_clear_object_dynamics + if ( + should_clear_object_dynamics + and waypoint_index + 1 >= post_grasp_clear_step + ): + obj.clear_dynamics() + should_clear_object_dynamics = False + + video_path = replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=replay_traj, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_place", + video_case_suffix, + ), + on_step=_on_step, + ) + else: + video_path = record_static_scene_video( + sim=sim, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_place", + f"{video_case_suffix}_failed_static", + ), + ) + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + return { + "case_id": case_id, + "object_type": object_preset.object_type, + "material": object_preset.material_name, + "quadrant": position_case.quadrant, + "position_case": position_case.name, + "init_xy": position_case.xy, + "pickup_approach": pickup_approach, + "approach_direction": approach_direction_text, + "place_case": case.name, + "repeat": repeat, + "planning_success": bool(is_success), + "released": released, + "physical_pick_success": physical_pick_success, + "physical_place_success": physical_place_success, + "success": success, + "cost_time_ms": elapsed * 1000.0, + "cpu_delta_mb": mem_delta["cpu_mb"], + "gpu_delta_mb": mem_delta["gpu_mb"], + "peak_gpu_mb": peak_gpu, + "object_lift_delta_m": object_lift_delta_m, + "object_final_x_m": ( + final_obj_position[0] if final_obj_position is not None else None + ), + "object_final_y_m": ( + final_obj_position[1] if final_obj_position is not None else None + ), + "object_final_z_m": ( + final_obj_position[2] if final_obj_position is not None else None + ), + "place_xy_error_m": place_xy_error_m, + "place_z_error_m": place_z_error_m, + "precondition_waypoints": precondition_waypoints, + "trajectory_waypoints": int(traj.shape[1]) if traj.ndim >= 2 else 0, + "failure_reason": failure_reason, + "video_path": str(video_path) if video_path is not None else "", + } + except Exception as exc: + video_path = None + if should_record_case(args, recorded_count, False): + try: + reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + video_path = record_static_scene_video( + sim=sim, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_place", + ( + f"{object_preset.object_type}_{position_case.name}_" + f"{pickup_approach}_{case.name}_r{repeat}" + "_exception_static" + ), + ), + ) + reset_rigid_object(obj, initial_obj_pose) + except Exception: + video_path = None + return { + "case_id": case_id, + "object_type": object_preset.object_type, + "material": object_preset.material_name, + "quadrant": position_case.quadrant, + "position_case": position_case.name, + "init_xy": position_case.xy, + "pickup_approach": pickup_approach, + "approach_direction": "N/A", + "place_case": case.name, + "repeat": repeat, + "planning_success": False, + "released": False, + "physical_pick_success": False, + "physical_place_success": False, + "success": False, + "cost_time_ms": 0.0, + "cpu_delta_mb": 0.0, + "gpu_delta_mb": 0.0, + "peak_gpu_mb": 0.0, + "object_lift_delta_m": None, + "object_final_x_m": None, + "object_final_y_m": None, + "object_final_z_m": None, + "place_xy_error_m": None, + "place_z_error_m": None, + "precondition_waypoints": 0, + "trajectory_waypoints": 0, + "failure_reason": f"exception:{type(exc).__name__}:{exc}", + "video_path": str(video_path) if video_path is not None else "", + } + + +def _build_rows(results: list[dict[str, object]]): + """Build report rows for Place.""" + perf_rows = [] + metric_rows = [] + for result in results: + perf_rows.append( + { + "sample_size": 1, + "impl": "place", + "case_id": result["case_id"], + "object_type": result["object_type"], + "material": result["material"], + "quadrant": result["quadrant"], + "position_case": result["position_case"], + "init_xy": f"({result['init_xy'][0]:.3f},{result['init_xy'][1]:.3f})", + "pickup_approach": result["pickup_approach"], + "approach_direction": result["approach_direction"], + "place_case": result["place_case"], + "repeat": result["repeat"], + "cost_time_ms": format_float(result["cost_time_ms"]), + "cpu_delta_mb": format_float(result["cpu_delta_mb"]), + "gpu_delta_mb": format_float(result["gpu_delta_mb"]), + "peak_gpu_mb": format_float(result["peak_gpu_mb"]), + } + ) + metric_rows.append( + { + "sample_size": 1, + "impl": "place", + "case_id": result["case_id"], + "object_type": result["object_type"], + "material": result["material"], + "quadrant": result["quadrant"], + "position_case": result["position_case"], + "pickup_approach": result["pickup_approach"], + "approach_direction": result["approach_direction"], + "place_case": result["place_case"], + "success_rate": f"{float(result['success']):.6f}", + "planning_success_rate": f"{float(result['planning_success']):.6f}", + "release_success_rate": f"{float(result['released']):.6f}", + "physical_pick_success_rate": ( + f"{float(result['physical_pick_success']):.6f}" + ), + "physical_place_success_rate": ( + f"{float(result['physical_place_success']):.6f}" + ), + "object_lift_delta_m": format_float(result["object_lift_delta_m"]), + "object_final_x_m": format_float(result["object_final_x_m"]), + "object_final_y_m": format_float(result["object_final_y_m"]), + "object_final_z_m": format_float(result["object_final_z_m"]), + "place_xy_error_m": format_float(result["place_xy_error_m"]), + "place_z_error_m": format_float(result["place_z_error_m"]), + "precondition_waypoints": result["precondition_waypoints"], + "trajectory_waypoints": result["trajectory_waypoints"], + "failure_reason": result["failure_reason"] or "N/A", + } + ) + return perf_rows, metric_rows + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: + """Run Place benchmark and write a markdown report.""" + args = _parse_args() if args is None else args + if args.repeat < 1: + raise ValueError("--repeat must be at least 1.") + profile = resolve_profile(args) + + ensure_repo_root() + ensure_torch() + from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.planners import ToppraPlannerCfg + from scripts.tutorials.atomic_action.place import ( + create_robot, + initialize_simulation, + ) + + object_presets = select_mesh_object_presets(args.object_types, profile) + position_cases = select_position_cases(args.position_cases, profile) + approaches = select_pickup_approaches(args.approach_cases, profile) + cases = _selected_cases(args.place_cases, profile) + repeat = 1 if profile == "smoke" else args.repeat + + print("=" * 60) + print("Place Atomic Action Benchmark") + print("=" * 60) + print( + "Coverage: " + f"profile={profile}, {len(object_presets)} object(s) x " + f"{len(position_cases)} position(s) x {len(approaches)} pickup approach(es) " + f"x {len(cases)} place target(s) " + f"x {repeat} repeat(s)" + ) + + sim = initialize_simulation(args) + robot = create_robot(sim) + initial_qpos = robot.get_qpos().clone() + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + object_pool = {} + for object_index, object_preset in enumerate(object_presets): + obj = create_benchmark_object( + sim=sim, + preset=object_preset, + position_case=position_cases[0], + uid_suffix="pool", + ) + base_pose = obj.get_local_pose(to_matrix=True).clone() + park_rigid_object(obj, base_pose, index=object_index, sim=sim) + object_pool[object_preset.object_type] = (obj, base_pose) + + results: list[dict[str, object]] = [] + video_paths: list[str] = [] + print("\n=== Place Object/Position/Target Sweep ===") + for object_preset in object_presets: + obj, base_pose = object_pool[object_preset.object_type] + for parked_index, parked_preset in enumerate(object_presets): + if parked_preset.object_type == object_preset.object_type: + continue + parked_obj, parked_base_pose = object_pool[parked_preset.object_type] + park_rigid_object(parked_obj, parked_base_pose, index=parked_index, sim=sim) + for position_case in position_cases: + for pickup_approach in approaches: + for case in cases: + for repeat_index in range(repeat): + result = _run_case( + sim, + robot, + motion_gen, + initial_qpos, + args, + obj, + base_pose, + object_preset, + position_case, + pickup_approach, + case, + repeat_index, + len(video_paths), + profile, + ) + results.append(result) + if result["video_path"]: + video_paths.append(str(result["video_path"])) + print( + f" {result['case_id']:<48} " + f"time={result['cost_time_ms']:>10.2f} ms | " + f"success={result['success']}" + ) + + perf_rows, metric_rows = _build_rows(results) + leaderboard_rows = build_single_action_leaderboard("place", metric_rows) + report_path = write_markdown_report( + benchmark_name="atomic_action_place", + perf_rows=perf_rows, + metric_rows=metric_rows, + leaderboard_rows=leaderboard_rows, + notes=[ + "Timed block includes Place only; PickUp precondition is prepared " + "outside timing.", + f"Profile: {profile}", + "Object presets: " + + ", ".join(describe_object_preset(preset) for preset in object_presets), + "Position cases: " + + ", ".join( + f"{case.name}/{case.quadrant}/xy={case.xy}" for case in position_cases + ), + "PickUp approach cases: " + ", ".join(approaches), + "Place target cases: " + ", ".join(case.name for case in cases), + f"CPU memory backend: {CPU_MEMORY_BACKEND}", + f"n_sample: {1000 if profile == 'smoke' else args.n_sample}", + ( + "Side grasp opening-axis rule: prefer candidates with " + f"abs(opening_axis_z) <= {SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z:.2f}; " + "fallback cost += " + f"abs(opening_axis_z) * {SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT:.2f}." + ), + ( + "Physical success rule: PickUp precondition lift delta >= " + f"{PHYSICAL_PICK_MIN_LIFT_M:.3f} m, Place released the object, " + f"and final object XY error <= {PHYSICAL_PLACE_XY_TOLERANCE_M:.3f} m." + ), + *summarize_video_recording(args, results, video_paths), + ], + ) + print(f"Markdown report saved: {report_path}") + return report_path + + +def main() -> None: + """Run the CLI entry point.""" + try: + run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/press_benchmark.py b/scripts/benchmark/atomic_action/press_benchmark.py new file mode 100644 index 000000000..a2c874a73 --- /dev/null +++ b/scripts/benchmark/atomic_action/press_benchmark.py @@ -0,0 +1,985 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Benchmark Press atomic action across object presets and start positions. + +The benchmark sweeps object presets such as bottle and mug against multiple +initial XY positions that cover all four workspace quadrants. It reports +planning latency, memory usage, planning success, and whether the generated +trajectory reaches the object's top center. +Run: python -m scripts.benchmark.atomic_action.press_benchmark +""" + +from __future__ import annotations + +import argparse +import math +import os +import resource +import sys +import time +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + add_profile_benchmark_args, + add_video_benchmark_args, + build_video_output_path, + COVERAGE_POSITION_CASE_NAMES, + park_rigid_object, + replay_trajectory_with_recording, + reset_rigid_object, + reset_rigid_object_xy, + resolve_profile, + should_record_case, + SMOKE_POSITION_CASE_NAMES, +) + +try: + import psutil +except ModuleNotFoundError: + psutil = None + +CPU_MEMORY_BACKEND = "psutil" if psutil is not None else "resource" +_RUNTIME_IMPORTS_READY = False + +# Keep these constants aligned with scripts/tutorials/atomic_action/press.py. +DEFAULT_PRESS_TOLERANCE = 0.01 +MOVE_SAMPLE_INTERVAL = 60 +PRESS_SAMPLE_INTERVAL = 90 +HAND_INTERP_STEPS = 12 +TABLE_TOP_Z = -0.045 +PRESS_CLEARANCE = 0.13 +PRESS_SURFACE_OFFSET = 0.003 + + +def _ensure_runtime_imports() -> None: + """Import simulation dependencies only when the benchmark is executed.""" + global _RUNTIME_IMPORTS_READY + if _RUNTIME_IMPORTS_READY: + return + + repo_root = Path(__file__).resolve().parents[3] + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + + try: + import torch as torch_module + from embodichain.lab.sim import SimulationManager as simulation_manager_cls + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine as atomic_action_engine_cls, + EndEffectorPoseTarget as end_effector_pose_target_cls, + MoveEndEffector as move_end_effector_cls, + MoveEndEffectorCfg as move_end_effector_cfg_cls, + Press as press_cls, + PressCfg as press_cfg_cls, + ) + from embodichain.lab.sim.cfg import ( + RigidBodyAttributesCfg as rigid_body_attributes_cfg_cls, + RigidObjectCfg as rigid_object_cfg_cls, + ) + from embodichain.lab.sim.material import ( + VisualMaterialCfg as visual_material_cfg_cls, + ) + from embodichain.lab.sim.objects import RigidObject as rigid_object_cls + from embodichain.lab.sim.objects import Robot as robot_cls + from embodichain.lab.sim.planners import ( + MotionGenerator as motion_generator_cls, + MotionGenCfg as motion_gen_cfg_cls, + ToppraPlannerCfg as toppra_planner_cfg_cls, + ) + from embodichain.lab.sim.shapes import CubeCfg as cube_cfg_cls + from scripts.tutorials.atomic_action.press import ( + create_robot as create_robot_fn, + create_table as create_table_fn, + get_hand_close_qpos as get_hand_close_qpos_fn, + initialize_simulation as initialize_simulation_fn, + make_top_down_eef_pose as make_top_down_eef_pose_fn, + settle_object as settle_object_fn, + ) + except ModuleNotFoundError as exc: + raise RuntimeError( + "Atomic action benchmark requires the EmbodiChain simulation runtime " + f"and PyTorch. Missing module: {exc.name}." + ) from exc + + globals().update( + { + "torch": torch_module, + "SimulationManager": simulation_manager_cls, + "AtomicActionEngine": atomic_action_engine_cls, + "EndEffectorPoseTarget": end_effector_pose_target_cls, + "MoveEndEffector": move_end_effector_cls, + "MoveEndEffectorCfg": move_end_effector_cfg_cls, + "Press": press_cls, + "PressCfg": press_cfg_cls, + "RigidBodyAttributesCfg": rigid_body_attributes_cfg_cls, + "RigidObjectCfg": rigid_object_cfg_cls, + "VisualMaterialCfg": visual_material_cfg_cls, + "RigidObject": rigid_object_cls, + "Robot": robot_cls, + "MotionGenerator": motion_generator_cls, + "MotionGenCfg": motion_gen_cfg_cls, + "ToppraPlannerCfg": toppra_planner_cfg_cls, + "CubeCfg": cube_cfg_cls, + "create_robot": create_robot_fn, + "create_table": create_table_fn, + "get_hand_close_qpos": get_hand_close_qpos_fn, + "initialize_simulation": initialize_simulation_fn, + "make_top_down_eef_pose": make_top_down_eef_pose_fn, + "settle_object": settle_object_fn, + } + ) + _RUNTIME_IMPORTS_READY = True + + +@dataclass(frozen=True) +class ObjectPreset: + """Primitive object preset used by the atomic-action benchmark.""" + + object_type: str + material_name: str + size: tuple[float, float, float] + base_color: tuple[float, float, float, float] + roughness: float + dynamic_friction: float = 0.8 + static_friction: float = 0.9 + + +@dataclass(frozen=True) +class PositionCase: + """Initial object position case with a quadrant label.""" + + name: str + quadrant: str + xy: tuple[float, float] + + +@dataclass(frozen=True) +class PressCaseResult: + """Result for one Press benchmark case.""" + + case_id: str + object_type: str + material_name: str + quadrant: str + position_case: str + init_xy: tuple[float, float] + repeat_index: int + planning_success: bool + center_hit: bool + cost_time_ms: float + cpu_delta_mb: float + gpu_delta_mb: float + peak_gpu_mb: float + xy_error_m: float | None + hit_step: int | None + trajectory_waypoints: int + failure_reason: str + video_path: str = "" + + +OBJECT_PRESETS: dict[str, ObjectPreset] = { + "bottle": ObjectPreset( + object_type="bottle", + material_name="green_plastic", + size=(0.06, 0.06, 0.16), + base_color=(0.10, 0.45, 0.32, 1.0), + roughness=0.55, + ), + "mug": ObjectPreset( + object_type="mug", + material_name="ceramic", + size=(0.10, 0.08, 0.10), + base_color=(0.88, 0.85, 0.78, 1.0), + roughness=0.35, + ), + "wooden_block": ObjectPreset( + object_type="wooden_block", + material_name="wood", + size=(0.12, 0.12, 0.06), + base_color=(0.58, 0.32, 0.14, 1.0), + roughness=0.85, + ), +} + +POSITION_CASES: dict[str, PositionCase] = { + "q1_near": PositionCase(name="q1_near", quadrant="q1", xy=(0.02, 0.18)), + "q1_far": PositionCase(name="q1_far", quadrant="q1", xy=(0.12, 0.36)), + "q2_near": PositionCase(name="q2_near", quadrant="q2", xy=(-0.42, 0.18)), + "q2_far": PositionCase(name="q2_far", quadrant="q2", xy=(-0.62, 0.36)), + "q3_near": PositionCase(name="q3_near", quadrant="q3", xy=(-0.42, -0.18)), + "q3_far": PositionCase(name="q3_far", quadrant="q3", xy=(-0.62, -0.36)), + "q4_near": PositionCase(name="q4_near", quadrant="q4", xy=(0.02, -0.18)), + "q4_far": PositionCase(name="q4_far", quadrant="q4", xy=(0.12, -0.36)), +} + +DEFAULT_OBJECT_TYPES = ("bottle", "mug") +FULL_OBJECT_TYPES = tuple(OBJECT_PRESETS.keys()) +SMOKE_OBJECT_TYPES = ("bottle",) + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add atomic-action benchmark arguments to an argument parser.""" + add_profile_benchmark_args(parser) + parser.add_argument( + "--object_types", + nargs="+", + choices=(*OBJECT_PRESETS.keys(), "all"), + default=None, + help=( + "Object presets to benchmark. Defaults are selected by --profile; " + "use 'all' to include every preset." + ), + ) + parser.add_argument( + "--position_cases", + nargs="+", + choices=(*POSITION_CASES.keys(), "all"), + default=None, + help=( + "Initial position cases to benchmark. Defaults are selected by " + "--profile; use 'all' for all near/far cases." + ), + ) + parser.add_argument( + "--repeat", + type=int, + default=1, + help="Number of repeats for every object-position case.", + ) + parser.add_argument( + "--smoke", + action="store_true", + help="Alias for --profile smoke.", + ) + parser.add_argument( + "--device", + type=str, + default="cpu", + help="Simulation device, e.g. 'cpu' or 'cuda'.", + ) + parser.add_argument( + "--renderer", + type=str, + choices=("auto", "hybrid", "fast-rt", "rt"), + default="auto", + help="Renderer backend used by SimulationManager.", + ) + add_video_benchmark_args(parser) + parser.add_argument( + "--press_tolerance", + type=float, + default=DEFAULT_PRESS_TOLERANCE, + help="XY tolerance in meters for the press-center check.", + ) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments for the atomic-action benchmark.""" + parser = argparse.ArgumentParser( + description=( + "Benchmark Press atomic action over object presets and initial " + "workspace quadrants." + ) + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _sync_cuda() -> None: + """Synchronize CUDA stream when available.""" + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def _reset_peak_gpu_memory() -> None: + """Reset PyTorch peak GPU memory stats when CUDA is available.""" + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + + +def _peak_gpu_memory_mb() -> float: + """Return peak GPU memory allocated by PyTorch in MB.""" + if not torch.cuda.is_available(): + return 0.0 + return torch.cuda.max_memory_allocated() / 1024**2 + + +def _memory_snapshot() -> dict[str, float]: + """Return current process memory usage snapshot in MB.""" + if psutil is not None: + process = psutil.Process(os.getpid()) + cpu_mb = process.memory_info().rss / 1024**2 + else: + cpu_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0 + gpu_mb = ( + torch.cuda.memory_allocated() / 1024**2 if torch.cuda.is_available() else 0.0 + ) + return {"cpu_mb": cpu_mb, "gpu_mb": gpu_mb} + + +def _format_float(value: float | None, precision: int = 6) -> str: + """Format finite floats for tables and use N/A for missing values.""" + if value is None or not math.isfinite(value): + return "N/A" + return f"{value:.{precision}f}" + + +def _format_markdown_table(rows: list[dict[str, object]]) -> list[str]: + """Format rows into a markdown table.""" + if not rows: + return ["No data."] + + headers = list(rows[0].keys()) + lines = [ + "| " + " | ".join(headers) + " |", + "| " + " | ".join(["---"] * len(headers)) + " |", + ] + for row in rows: + lines.append("| " + " | ".join(str(row[h]) for h in headers) + " |") + return lines + + +def _write_markdown_report( + benchmark_name: str, + perf_rows: list[dict[str, object]], + metric_rows: list[dict[str, object]], + leaderboard_rows: list[dict[str, object]], + notes: list[str] | None = None, +) -> Path: + """Write benchmark results into one markdown report file.""" + output_dir = Path("outputs/benchmarks") + output_dir.mkdir(parents=True, exist_ok=True) + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + report_path = output_dir / f"{benchmark_name}_{timestamp}.md" + + lines: list[str] = [ + f"# {benchmark_name} Benchmark Report", + "", + f"Generated at: {datetime.now().isoformat(timespec='seconds')}", + "", + "## Time & Memory", + "", + ] + lines.extend(_format_markdown_table(perf_rows)) + lines.extend(["", "## Success & Other Metrics", ""]) + lines.extend(_format_markdown_table(metric_rows)) + lines.extend(["", "## Leaderboard", ""]) + lines.extend(_format_markdown_table(leaderboard_rows)) + + if notes: + lines.extend(["", "## Notes", ""]) + lines.extend([f"- {note}" for note in notes]) + + report_path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return report_path + + +def _default_object_types_for_profile(profile: str) -> tuple[str, ...]: + """Return default Press primitive object names for a profile.""" + if profile in ("smoke", "coverage", "full"): + return SMOKE_OBJECT_TYPES + raise ValueError(f"Unsupported benchmark profile: {profile}") + + +def _select_object_presets( + object_types: list[str] | None, + profile: str, +) -> list[ObjectPreset]: + """Resolve selected object preset names.""" + if not object_types: + object_types = list(_default_object_types_for_profile(profile)) + if "all" in object_types: + return list(OBJECT_PRESETS.values()) + return [OBJECT_PRESETS[name] for name in object_types] + + +def _default_position_cases_for_profile(profile: str) -> tuple[str, ...]: + """Return default Press position case names for a profile.""" + if profile == "smoke": + return SMOKE_POSITION_CASE_NAMES + if profile in ("coverage", "full"): + return COVERAGE_POSITION_CASE_NAMES + raise ValueError(f"Unsupported benchmark profile: {profile}") + + +def _select_position_cases( + position_cases: list[str] | None, + profile: str, +) -> list[PositionCase]: + """Resolve selected position case names.""" + if not position_cases: + position_cases = list(_default_position_cases_for_profile(profile)) + if "all" in position_cases: + return list(POSITION_CASES.values()) + return [POSITION_CASES[name] for name in position_cases] + + +def _create_benchmark_object( + sim: SimulationManager, + preset: ObjectPreset, + position_case: PositionCase, + repeat_index: int, +) -> RigidObject: + """Create one static benchmark object at the requested initial position.""" + init_pos = ( + position_case.xy[0], + position_case.xy[1], + TABLE_TOP_Z + 0.5 * preset.size[2], + ) + uid = f"atomic_benchmark_{preset.object_type}_{position_case.name}_{repeat_index}" + cfg = RigidObjectCfg( + uid=uid, + shape=CubeCfg( + size=list(preset.size), + visual_material=VisualMaterialCfg( + uid=f"{preset.object_type}_{preset.material_name}_mat", + base_color=list(preset.base_color), + roughness=preset.roughness, + ), + ), + body_type="static", + attrs=RigidBodyAttributesCfg( + dynamic_friction=preset.dynamic_friction, + static_friction=preset.static_friction, + ), + init_pos=init_pos, + ) + return sim.add_rigid_object(cfg=cfg) + + +def _reset_robot(robot: Robot, initial_qpos: torch.Tensor) -> None: + """Reset current and target robot qpos to the benchmark initial posture.""" + for target in (False, True): + robot.set_qpos(initial_qpos, target=target) + robot.clear_dynamics() + + +def _build_atomic_engine( + motion_gen: MotionGenerator, + robot: Robot, + device: torch.device, +) -> AtomicActionEngine: + """Build a Press benchmark engine with MoveEndEffector pre-positioning.""" + hand_close = get_hand_close_qpos(robot, device) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + MoveEndEffector( + motion_gen, + cfg=MoveEndEffectorCfg( + control_part="arm", + sample_interval=MOVE_SAMPLE_INTERVAL, + ), + ) + ) + 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, + ), + ) + ) + return atomic_engine + + +def _make_press_targets( + obj: RigidObject, + preset: ObjectPreset, +) -> tuple[torch.Tensor, torch.Tensor]: + """Create pre-press and press poses for the object's top center.""" + obj_pose = obj.get_local_pose(to_matrix=True) + object_center = obj_pose[0, :3, 3].clone() + object_top_z = object_center[2] + 0.5 * preset.size[2] + + press_position = object_center.clone() + press_position[2] = object_top_z + PRESS_SURFACE_OFFSET + move_position = press_position.clone() + move_position[2] = object_top_z + PRESS_CLEARANCE + + return make_top_down_eef_pose(move_position), make_top_down_eef_pose(press_position) + + +def _compute_press_center_check( + robot: Robot, + traj: torch.Tensor, + obj: RigidObject, + object_height: float, + tolerance: float, +) -> tuple[bool, float, int]: + """Check whether the planned Press trajectory reaches the object top center.""" + if traj.numel() == 0: + return False, float("inf"), -1 + + arm_joint_ids = robot.get_joint_ids(name="arm") + n_down = (PRESS_SAMPLE_INTERVAL - HAND_INTERP_STEPS) // 2 + press_segment_start = MOVE_SAMPLE_INTERVAL + HAND_INTERP_STEPS + press_segment_end = min(press_segment_start + n_down, traj.shape[1]) + arm_traj = traj[:, press_segment_start:press_segment_end, arm_joint_ids] + if arm_traj.shape[1] == 0: + return False, float("inf"), -1 + + fk_pose = torch.stack( + [ + robot.compute_fk( + qpos=waypoint.unsqueeze(0), + name="arm", + to_matrix=True, + )[0] + for waypoint in arm_traj[0] + ], + dim=0, + ) + + obj_pose = obj.get_local_pose(to_matrix=True) + object_center = obj_pose[0, :3, 3] + object_top_z = object_center[2] + 0.5 * object_height + target_xy = object_center[:2] + target_z = object_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 <= tolerance, center_error, press_segment_start + best_idx + + +def _timed_atomic_run( + atomic_engine: AtomicActionEngine, + move_target: torch.Tensor, + press_target: torch.Tensor, +) -> tuple[float, dict[str, float], float, bool, torch.Tensor]: + """Run a timed atomic-action sequence and return timing/memory/results.""" + _reset_peak_gpu_memory() + mem_before = _memory_snapshot() + _sync_cuda() + + start = time.perf_counter() + is_success, traj, _ = atomic_engine.run( + steps=[ + ("move_end_effector", EndEffectorPoseTarget(xpos=move_target)), + ("press", EndEffectorPoseTarget(xpos=press_target)), + ] + ) + _sync_cuda() + elapsed = time.perf_counter() - start + + mem_after = _memory_snapshot() + deltas = { + "cpu_mb": mem_after["cpu_mb"] - mem_before["cpu_mb"], + "gpu_mb": mem_after["gpu_mb"] - mem_before["gpu_mb"], + } + return elapsed, deltas, _peak_gpu_memory_mb(), is_success, traj + + +def _run_press_case( + sim: SimulationManager, + robot: Robot, + atomic_engine: AtomicActionEngine, + initial_qpos: torch.Tensor, + obj: RigidObject, + base_obj_pose: torch.Tensor, + preset: ObjectPreset, + position_case: PositionCase, + repeat_index: int, + press_tolerance: float, + args: argparse.Namespace, + recorded_count: int, +) -> PressCaseResult: + """Run one object-position Press benchmark case.""" + case_id = f"{preset.object_type}:{position_case.name}:r{repeat_index}" + try: + _reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + move_target, press_target = _make_press_targets(obj, preset) + + elapsed, mem_delta, peak_gpu, planning_success, traj = _timed_atomic_run( + atomic_engine=atomic_engine, + move_target=move_target, + press_target=press_target, + ) + video_path = None + if should_record_case(args, recorded_count, bool(planning_success)): + _reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + video_path = replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=traj, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_press", + (f"{preset.object_type}_{position_case.name}" f"_r{repeat_index}"), + ), + ) + _reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + center_hit = False + xy_error_m: float | None = None + hit_step: int | None = None + failure_reason = "" + if planning_success: + center_hit, xy_error_m, raw_hit_step = _compute_press_center_check( + robot=robot, + traj=traj, + obj=obj, + object_height=preset.size[2], + tolerance=press_tolerance, + ) + hit_step = raw_hit_step if raw_hit_step >= 0 else None + if not center_hit: + failure_reason = "center_miss" + else: + failure_reason = "planning_failed" + + return PressCaseResult( + case_id=case_id, + object_type=preset.object_type, + material_name=preset.material_name, + quadrant=position_case.quadrant, + position_case=position_case.name, + init_xy=position_case.xy, + repeat_index=repeat_index, + planning_success=planning_success, + center_hit=center_hit, + cost_time_ms=elapsed * 1000.0, + cpu_delta_mb=mem_delta["cpu_mb"], + gpu_delta_mb=mem_delta["gpu_mb"], + peak_gpu_mb=peak_gpu, + xy_error_m=xy_error_m, + hit_step=hit_step, + trajectory_waypoints=int(traj.shape[1]) if traj.ndim >= 2 else 0, + failure_reason=failure_reason, + video_path=str(video_path) if video_path is not None else "", + ) + except Exception as exc: + return PressCaseResult( + case_id=case_id, + object_type=preset.object_type, + material_name=preset.material_name, + quadrant=position_case.quadrant, + position_case=position_case.name, + init_xy=position_case.xy, + repeat_index=repeat_index, + planning_success=False, + center_hit=False, + cost_time_ms=0.0, + cpu_delta_mb=0.0, + gpu_delta_mb=0.0, + peak_gpu_mb=0.0, + xy_error_m=None, + hit_step=None, + trajectory_waypoints=0, + failure_reason=f"exception:{type(exc).__name__}:{exc}", + ) + + +def _build_perf_rows(results: list[PressCaseResult]) -> list[dict[str, object]]: + """Build Time & Memory table rows.""" + rows: list[dict[str, object]] = [] + for result in results: + rows.append( + { + "sample_size": 1, + "impl": "press", + "case_id": result.case_id, + "object_type": result.object_type, + "material": result.material_name, + "quadrant": result.quadrant, + "position_case": result.position_case, + "init_xy": f"({result.init_xy[0]:.3f},{result.init_xy[1]:.3f})", + "repeat": result.repeat_index, + "cost_time_ms": _format_float(result.cost_time_ms), + "cpu_delta_mb": _format_float(result.cpu_delta_mb), + "gpu_delta_mb": _format_float(result.gpu_delta_mb), + "peak_gpu_mb": _format_float(result.peak_gpu_mb), + } + ) + return rows + + +def _build_metric_rows(results: list[PressCaseResult]) -> list[dict[str, object]]: + """Build Success & Other Metrics table rows.""" + rows: list[dict[str, object]] = [] + for result in results: + overall_success = result.planning_success and result.center_hit + rows.append( + { + "sample_size": 1, + "impl": "press", + "case_id": result.case_id, + "object_type": result.object_type, + "material": result.material_name, + "quadrant": result.quadrant, + "position_case": result.position_case, + "success_rate": f"{float(overall_success):.6f}", + "planning_success_rate": f"{float(result.planning_success):.6f}", + "center_hit_rate": f"{float(result.center_hit):.6f}", + "xy_error_m": _format_float(result.xy_error_m), + "hit_step": result.hit_step if result.hit_step is not None else "N/A", + "trajectory_waypoints": result.trajectory_waypoints, + "failure_reason": result.failure_reason or "N/A", + } + ) + return rows + + +def _build_leaderboard_rows(results: list[PressCaseResult]) -> list[dict[str, object]]: + """Aggregate and rank object-conditioned Press variants by success rate.""" + aggregate: dict[str, dict[str, float | set[str]]] = {} + for result in results: + algorithm = f"press:{result.object_type}" + if algorithm not in aggregate: + aggregate[algorithm] = { + "overall_success_sum": 0.0, + "planning_success_sum": 0.0, + "xy_error_sum": 0.0, + "xy_error_count": 0.0, + "cost_time_sum": 0.0, + "case_count": 0.0, + "quadrants": set(), + } + + stats = aggregate[algorithm] + stats["overall_success_sum"] = float(stats["overall_success_sum"]) + float( + result.planning_success and result.center_hit + ) + stats["planning_success_sum"] = float(stats["planning_success_sum"]) + float( + result.planning_success + ) + if result.xy_error_m is not None and math.isfinite(result.xy_error_m): + stats["xy_error_sum"] = float(stats["xy_error_sum"]) + result.xy_error_m + stats["xy_error_count"] = float(stats["xy_error_count"]) + 1.0 + stats["cost_time_sum"] = float(stats["cost_time_sum"]) + result.cost_time_ms + stats["case_count"] = float(stats["case_count"]) + 1.0 + quadrants = stats["quadrants"] + if isinstance(quadrants, set): + quadrants.add(result.quadrant) + + ranked = sorted( + aggregate.items(), + key=lambda item: ( + float(item[1]["overall_success_sum"]) + / max(float(item[1]["case_count"]), 1.0), + -float(item[1]["cost_time_sum"]) / max(float(item[1]["case_count"]), 1.0), + ), + reverse=True, + ) + + rows: list[dict[str, object]] = [] + for rank, (algorithm, stats) in enumerate(ranked, start=1): + case_count = max(float(stats["case_count"]), 1.0) + xy_error_count = float(stats["xy_error_count"]) + avg_xy_error = ( + float(stats["xy_error_sum"]) / xy_error_count + if xy_error_count > 0.0 + else None + ) + quadrants = stats["quadrants"] + quadrant_coverage = ( + ",".join(sorted(quadrants)) if isinstance(quadrants, set) else "" + ) + rows.append( + { + "rank": rank, + "algorithm": algorithm, + "overall_success_rate": ( + f"{float(stats['overall_success_sum']) / case_count:.2%}" + ), + "planning_success_rate": ( + f"{float(stats['planning_success_sum']) / case_count:.2%}" + ), + "avg_xy_error_m": _format_float(avg_xy_error), + "avg_cost_time_ms": _format_float( + float(stats["cost_time_sum"]) / case_count + ), + "evaluated_cases": int(case_count), + "quadrant_coverage": quadrant_coverage, + } + ) + return rows + + +def _print_case_result(result: PressCaseResult) -> None: + """Print one aligned case result line.""" + overall_success = result.planning_success and result.center_hit + print( + f" {result.case_id:<28} " + f"time={result.cost_time_ms:>10.2f} ms | " + f"CPU delta={result.cpu_delta_mb:+.1f} MB " + f"GPU delta={result.gpu_delta_mb:+.1f} MB " + f"peak GPU={result.peak_gpu_mb:.1f} MB | " + f"success={overall_success} " + f"xy_error={_format_float(result.xy_error_m, precision=4)}" + ) + if result.failure_reason: + print(f" reason={result.failure_reason}") + + +def _build_notes( + object_presets: list[ObjectPreset], + position_cases: list[PositionCase], + repeat: int, + video_paths: list[str], + profile: str, +) -> list[str]: + """Build report notes with benchmark coverage metadata.""" + quadrant_counts: dict[str, int] = {} + for position_case in position_cases: + quadrant_counts[position_case.quadrant] = ( + quadrant_counts.get(position_case.quadrant, 0) + 1 + ) + return [ + f"Profile: {profile}", + "Object presets: " + + ", ".join( + f"{preset.object_type}/{preset.material_name}/size={preset.size}" + for preset in object_presets + ), + "Position cases per quadrant: " + + ", ".join( + f"{quadrant}={count}" for quadrant, count in sorted(quadrant_counts.items()) + ), + f"CPU memory backend: {CPU_MEMORY_BACKEND}", + f"Repeat per object-position case: {repeat}", + "Replay videos: " + (", ".join(video_paths) if video_paths else "disabled"), + "success_rate is 1 only when planning succeeds and the Press trajectory " + "reaches the object top center.", + ] + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: + """Run all atomic-action benchmarks and write the markdown report.""" + args = _parse_args() if args is None else args + if args.repeat < 1: + raise ValueError("--repeat must be at least 1.") + profile = resolve_profile(args) + _ensure_runtime_imports() + + object_presets = _select_object_presets(args.object_types, profile) + position_cases = _select_position_cases(args.position_cases, profile) + repeat = 1 if profile == "smoke" else args.repeat + + print("=" * 60) + print("Atomic Action Press Benchmark") + print("=" * 60) + print( + "Coverage: " + f"profile={profile}, {len(object_presets)} object presets x " + f"{len(position_cases)} position cases x {repeat} repeat(s)" + ) + + sim = initialize_simulation(args) + robot = create_robot(sim) + create_table(sim) + initial_qpos = robot.get_qpos().clone() + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + atomic_engine = _build_atomic_engine(motion_gen, robot, sim.device) + object_pool = {} + for object_index, preset in enumerate(object_presets): + obj = _create_benchmark_object(sim, preset, position_cases[0], object_index) + settle_object(sim, obj, step=2) + base_pose = obj.get_local_pose(to_matrix=True).clone() + park_rigid_object(obj, base_pose, index=object_index, sim=sim) + object_pool[preset.object_type] = (obj, base_pose) + + results: list[PressCaseResult] = [] + video_paths: list[str] = [] + print("\n=== Press Object/Position Sweep ===") + for preset in object_presets: + obj, base_pose = object_pool[preset.object_type] + for parked_index, parked_preset in enumerate(object_presets): + if parked_preset.object_type == preset.object_type: + continue + parked_obj, parked_base_pose = object_pool[parked_preset.object_type] + park_rigid_object(parked_obj, parked_base_pose, index=parked_index, sim=sim) + for position_case in position_cases: + for repeat_index in range(repeat): + result = _run_press_case( + sim=sim, + robot=robot, + atomic_engine=atomic_engine, + initial_qpos=initial_qpos, + obj=obj, + base_obj_pose=base_pose, + preset=preset, + position_case=position_case, + repeat_index=repeat_index, + press_tolerance=args.press_tolerance, + args=args, + recorded_count=len(video_paths), + ) + results.append(result) + if result.video_path: + video_paths.append(result.video_path) + _print_case_result(result) + + perf_rows = _build_perf_rows(results) + metric_rows = _build_metric_rows(results) + leaderboard_rows = _build_leaderboard_rows(results) + report_path = _write_markdown_report( + benchmark_name="atomic_action_press", + perf_rows=perf_rows, + metric_rows=metric_rows, + leaderboard_rows=leaderboard_rows, + notes=_build_notes( + object_presets, + position_cases, + repeat, + video_paths, + profile, + ), + ) + + print("\n" + "=" * 60) + print("Benchmarks complete.") + print(f"Markdown report saved: {report_path}") + print("=" * 60) + return report_path + + +def main() -> None: + """Run the CLI entry point.""" + try: + run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/run_benchmark.py b/scripts/benchmark/atomic_action/run_benchmark.py new file mode 100644 index 000000000..26c9749b3 --- /dev/null +++ b/scripts/benchmark/atomic_action/run_benchmark.py @@ -0,0 +1,340 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Dispatch benchmarks for all atomic actions. + +Run a single action benchmark or all action benchmarks in sequence. +Run: python -m scripts.benchmark.atomic_action.run_benchmark --action press +""" + +from __future__ import annotations + +import argparse +import shlex +import subprocess +import sys +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + MESH_OBJECT_PRESETS, + PICKUP_APPROACH_CASES, + POSITION_CASES, + add_profile_benchmark_args, + add_video_benchmark_args, + resolve_profile, +) + +ACTION_MODULES = { + "move_end_effector": "scripts.benchmark.atomic_action.move_end_effector_benchmark", + "move_joints": "scripts.benchmark.atomic_action.move_joints_benchmark", + "pick_up": "scripts.benchmark.atomic_action.pickup_benchmark", + "move_held_object": "scripts.benchmark.atomic_action.move_held_object_benchmark", + "place": "scripts.benchmark.atomic_action.place_benchmark", + "press": "scripts.benchmark.atomic_action.press_benchmark", +} +DEFAULT_ACTIONS = tuple(ACTION_MODULES.keys()) +MESH_OBJECT_ACTIONS = {"pick_up", "move_held_object", "place"} +PRESS_OBJECT_TYPES = {"bottle", "mug", "wooden_block", "all"} +MESH_OBJECT_TYPES = {*MESH_OBJECT_PRESETS.keys(), "all"} + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add atomic-action aggregate benchmark CLI arguments.""" + parser.add_argument( + "--action", + nargs="+", + choices=(*ACTION_MODULES.keys(), "all"), + default=["press"], + help="Atomic action benchmark(s) to run. Use 'all' for every action.", + ) + parser.add_argument( + "--smoke", + action="store_true", + help="Alias for --profile smoke.", + ) + add_profile_benchmark_args(parser) + parser.add_argument( + "--object_types", + nargs="+", + default=None, + help=( + "Optional object presets forwarded to selected object-conditioned " + "benchmarks. Values must be valid for each selected action." + ), + ) + parser.add_argument( + "--position_cases", + nargs="+", + choices=(*POSITION_CASES.keys(), "all"), + default=None, + help=( + "Optional initial position cases forwarded to selected " + "object-conditioned benchmarks." + ), + ) + parser.add_argument( + "--approach_cases", + nargs="+", + choices=(*PICKUP_APPROACH_CASES, "all"), + default=None, + help=( + "Optional PickUp approach cases forwarded to pick/place/held-object " + "benchmarks." + ), + ) + parser.add_argument( + "--device", + type=str, + default="cpu", + help="Simulation device forwarded to each selected benchmark.", + ) + parser.add_argument( + "--renderer", + type=str, + choices=("auto", "hybrid", "fast-rt", "rt"), + default="auto", + help="Renderer backend forwarded to each selected benchmark.", + ) + add_video_benchmark_args(parser) + parser.add_argument( + "--in_process", + action="store_true", + help=( + "Run selected benchmarks in the current Python process. This is useful " + "for debugging but can leave DexSim resources alive across actions." + ), + ) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Run one or more atomic-action benchmarks." + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _selected_actions(actions: list[str]) -> list[str]: + """Resolve selected action names.""" + if "all" in actions: + return list(DEFAULT_ACTIONS) + return actions + + +def _validate_object_types_for_actions( + selected_actions: list[str], + object_types: list[str] | None, +) -> None: + """Validate optional object filters against selected action namespaces.""" + if not object_types: + return + + requested = set(object_types) + validators: dict[str, set[str]] = {} + for action_name in selected_actions: + if action_name in MESH_OBJECT_ACTIONS: + validators[action_name] = MESH_OBJECT_TYPES + elif action_name == "press": + validators[action_name] = PRESS_OBJECT_TYPES + + invalid_parts = [] + for action_name, valid_types in validators.items(): + invalid = sorted(requested - valid_types) + if invalid: + valid = ", ".join(sorted(valid_types)) + invalid_parts.append( + f"{action_name}: invalid {invalid}; valid values are {valid}" + ) + if invalid_parts: + raise RuntimeError( + "--object_types must be valid for every selected object-conditioned " + "benchmark. " + " | ".join(invalid_parts) + ) + + +def _make_child_args(args: argparse.Namespace) -> argparse.Namespace: + """Build minimal child benchmark arguments for aggregate dispatch.""" + profile = resolve_profile(args) + return argparse.Namespace( + smoke=profile == "smoke", + profile=profile, + repeat=1, + device=args.device, + renderer=args.renderer, + object_types=args.object_types, + position_cases=args.position_cases, + press_tolerance=0.01, + pose_cases=["all"], + sequence_cases=["all"], + approach_cases=args.approach_cases, + place_cases=None, + held_object_cases=None, + n_sample=1000 if profile == "smoke" else 10000, + force_reannotate=False, + record_video=args.record_video, + record_failed_video=args.record_failed_video, + video_case_limit=args.video_case_limit, + video_dir=args.video_dir, + video_fps=args.video_fps, + video_max_memory=args.video_max_memory, + video_width=args.video_width, + video_height=args.video_height, + video_hold_steps=args.video_hold_steps, + ) + + +def _make_child_cli_args(args: argparse.Namespace, action_name: str) -> list[str]: + """Build CLI arguments forwarded to one isolated benchmark subprocess.""" + profile = resolve_profile(args) + child_args = [ + "--profile", + profile, + "--repeat", + "1", + "--device", + args.device, + "--renderer", + args.renderer, + "--video_case_limit", + str(args.video_case_limit), + "--video_dir", + str(args.video_dir), + "--video_fps", + str(args.video_fps), + "--video_max_memory", + str(args.video_max_memory), + "--video_width", + str(args.video_width), + "--video_height", + str(args.video_height), + "--video_hold_steps", + str(args.video_hold_steps), + ] + if action_name in {"pick_up", "move_held_object", "place", "press"}: + if args.object_types: + child_args.append("--object_types") + child_args.extend(args.object_types) + if args.position_cases: + child_args.append("--position_cases") + child_args.extend(args.position_cases) + if action_name in {"pick_up", "move_held_object", "place"}: + if args.approach_cases: + child_args.append("--approach_cases") + child_args.extend(args.approach_cases) + if args.record_video: + child_args.append("--record_video") + if args.record_failed_video: + child_args.append("--record_failed_video") + return child_args + + +def _run_action_subprocess(action_name: str, args: argparse.Namespace) -> Path: + """Run one action benchmark in an isolated Python subprocess.""" + module_name = ACTION_MODULES[action_name] + command = [ + sys.executable, + "-m", + module_name, + *_make_child_cli_args(args, action_name), + ] + print("\n" + "=" * 60, flush=True) + print( + f"Running {action_name} benchmark in an isolated subprocess", + flush=True, + ) + print( + "Command: " + " ".join(shlex.quote(part) for part in command), + flush=True, + ) + print("=" * 60, flush=True) + + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + encoding="utf-8", + errors="replace", + bufsize=1, + ) + report_path: Path | None = None + if process.stdout is not None: + for line in process.stdout: + print(line, end="", flush=True) + if line.startswith("Markdown report saved:"): + report_path = Path(line.split(":", maxsplit=1)[1].strip()) + + return_code = process.wait() + if return_code != 0: + raise RuntimeError( + f"{action_name} benchmark subprocess failed with exit code " + f"{return_code}." + ) + if report_path is None: + raise RuntimeError( + f"{action_name} benchmark subprocess finished without reporting " + "a markdown report path." + ) + return report_path + + +def _run_in_process_benchmarks( + args: argparse.Namespace, + selected_actions: list[str], +) -> list[Path]: + """Run selected benchmarks in the current Python process.""" + reports: list[Path] = [] + for action_name in selected_actions: + module_name = ACTION_MODULES[action_name] + module = __import__(module_name, fromlist=["run_all_benchmarks"]) + child_args = _make_child_args(args) + report_path = module.run_all_benchmarks(child_args) + reports.append(report_path) + return reports + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> list[Path]: + """Run selected atomic-action benchmarks.""" + args = _parse_args() if args is None else args + selected_actions = _selected_actions(args.action) + _validate_object_types_for_actions(selected_actions, args.object_types) + + if args.in_process: + return _run_in_process_benchmarks(args, selected_actions) + + return [ + _run_action_subprocess(action_name, args) for action_name in selected_actions + ] + + +def main() -> None: + """Run the CLI entry point.""" + try: + reports = run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + if reports: + print("Generated reports:") + for report in reports: + print(f" {report}") + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/planners/neural_planner/run_benchmark.py b/scripts/benchmark/planners/neural_planner/run_benchmark.py index 12f5cd13f..a707b6b48 100644 --- a/scripts/benchmark/planners/neural_planner/run_benchmark.py +++ b/scripts/benchmark/planners/neural_planner/run_benchmark.py @@ -426,7 +426,8 @@ def _make_waypoints(start_pose: torch.Tensor, num_waypoints: int) -> torch.Tenso def _make_target_states(waypoints: torch.Tensor) -> list[PlanState]: return [ - PlanState(move_type=MoveType.EEF_MOVE, xpos=waypoint) for waypoint in waypoints + PlanState.single(move_type=MoveType.EEF_MOVE, xpos=waypoint) + for waypoint in waypoints ] @@ -559,7 +560,11 @@ def plan_ik_interpolate( joint_seed=qpos_seed, ) if not _all_success(success): - return PlanResult(success=False, positions=None) + return PlanResult( + success=torch.tensor([False]), + positions=None, + duration=torch.tensor([0.0]), + ) qpos_seed = qpos joint_targets.append(qpos.squeeze(0)) @@ -568,21 +573,21 @@ def plan_ik_interpolate( trajectory=trajectory, interp_num=sample_interval, device=robot.device, - ).squeeze(0) + ) return PlanResult( - success=True, + success=torch.tensor([True]), positions=positions, - duration=0.0, + duration=torch.tensor([0.0]), ) def _trajectory_fk_poses(result: PlanResult, robot: Robot) -> list[torch.Tensor]: """Return TCP poses sampled along the planned trajectory.""" if result.xpos_list is not None and result.xpos_list.shape[0] > 0: - return [pose for pose in result.xpos_list] - if result.positions is None or result.positions.shape[0] == 0: + return [pose for pose in result.xpos_list[0]] + if result.positions is None or result.positions.shape[1] == 0: return [] - qpos = result.positions + qpos = result.positions[0] if qpos.dim() == 1: qpos = qpos.unsqueeze(0) fk = robot.compute_batch_fk( @@ -638,10 +643,10 @@ def _final_eef_pose( robot: Robot, ) -> torch.Tensor | None: if result.xpos_list is not None and result.xpos_list.shape[0] > 0: - return result.xpos_list[-1] - if result.positions is None or result.positions.shape[0] == 0: + return result.xpos_list[0][-1] + if result.positions is None or result.positions.shape[1] == 0: return None - qpos = result.positions[-1] + qpos = result.positions[0][-1] if qpos.dim() == 1: qpos = qpos.unsqueeze(0) return robot.compute_fk(qpos=qpos, name=ARM_NAME, to_matrix=True)[0] @@ -652,15 +657,11 @@ def _compute_result_metrics( waypoints: torch.Tensor, robot: Robot, ) -> dict[str, object]: - success = bool( - result.success.item() if torch.is_tensor(result.success) else result.success - ) + success = bool(result.success.all().item()) rollout_steps = ( - int(result.positions.shape[0]) if result.positions is not None else 0 - ) - duration_s = float( - result.duration.item() if torch.is_tensor(result.duration) else result.duration + int(result.positions.shape[1]) if result.positions is not None else 0 ) + duration_s = float(result.duration[0].item()) trajectory_poses = _trajectory_fk_poses(result, robot) waypoint_errors = compute_waypoint_errors(trajectory_poses, waypoints) diff --git a/scripts/tutorials/atomic_action/coordinated_pickment.py b/scripts/tutorials/atomic_action/coordinated_pickment.py new file mode 100644 index 000000000..cb00e9fe3 --- /dev/null +++ b/scripts/tutorials/atomic_action/coordinated_pickment.py @@ -0,0 +1,795 @@ +# ---------------------------------------------------------------------------- +# 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 dual-arm coordinated pickment with selectable object meshes. + +The two UR5 arms pinch opposite sides of one object, lift it together, and move +the object to an object-centric target pose while both grippers stay closed. +""" + +from __future__ import annotations + +import argparse +import math +import os +import sys +import time +from dataclasses import dataclass +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 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, + AtomicActionEngine, + CoordinatedPickment, + CoordinatedPickmentCfg, + CoordinatedPickmentTarget, + ObjectSemantics, +) +from embodichain.lab.sim.cfg import ( + JointDrivePropertiesCfg, + RigidBodyAttributesCfg, + RigidObjectCfg, + RobotCfg, + URDFCfg, +) +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.utils import logger +from embodichain.utils.math import matrix_from_euler +from scripts.tutorials.atomic_action.tutorial_utils import ( + broadcast_pose_batch, + clone_local_pose_from_first_env, + create_toppra_motion_generator, + create_tutorial_simulation, + draw_axis_marker, + prepare_tutorial_scene, + replay_trajectory, +) + +ARM_URDF_PATH = "UniversalRobots/UR5/UR5.urdf" +GRIPPER_URDF_PATH = "DH_PGI_140_80/DH_PGI_140_80.urdf" +PICKMENT_ASSET_ROOT = "CoordinatedPlacementAndPickment" +GRIPPER_TCP_Z = 0.121 +ROBOT_INIT_POS = (1.95, 0.0, 0.1) +ROBOT_INIT_ROT = (0.0, 0.0, -90.0) +LEFT_ARM_HOME = (0.0, 0.0, -1.57, -1.57, 1.57, 1.57) +RIGHT_ARM_HOME = (-1.57, -1.57, -1.57, -1.57, 0.0, 0.0) +SUPPORT_SURFACE_Z = 0.65 +SUPPORT_SURFACE_SIZE = (0.60, 0.60, 0.02) +SUPPORT_SURFACE_CENTER = ( + 0.0, + 0.0, + SUPPORT_SURFACE_Z - 0.5 * SUPPORT_SURFACE_SIZE[2], +) +PICKMENT_RECORD_LOOK_AT = ( + (-0.25, 0.02, 2.5), + (0.0, 0.02, 0.75), + (0.0, 0.0, 1.0), +) + + +@dataclass(frozen=True) +class PickmentObjectPreset: + """Configuration for an object used by the coordinated pickment demo.""" + + label: str + mesh_path: str + init_xy: tuple[float, float] + init_rot: tuple[float, float, float] + surface_clearance: float + body_scale: tuple[float, float, float] + grasp_end_margin_ratio: float + grasp_z_clearance: float + target_translation: tuple[float, float, float] + target_world_yaw_deg: float + hand_close_qpos: float + grasp_z_ratio: float | None = None + + +OBJECT_PRESETS = { + "pencil": PickmentObjectPreset( + label="pencil", + mesh_path=f"{PICKMENT_ASSET_ROOT}/pencil.glb", + init_xy=(-0.02, 0.02), + # Rotate the imported pencil from its default upright orientation to a supported pose. + init_rot=(90.0, 0.0, 0.0), + surface_clearance=0.008, + body_scale=(2.0, 2.0, 2.0), + grasp_end_margin_ratio=0.12, + grasp_z_clearance=0.015, + target_translation=(-0.22, -0.04, 0.16), + target_world_yaw_deg=0.0, + hand_close_qpos=0.026, + ), + "pot": PickmentObjectPreset( + label="pot", + mesh_path=f"{PICKMENT_ASSET_ROOT}/pot.glb", + init_xy=(-0.02, 0.02), + init_rot=(-90.0, 90.0, 0.0), + surface_clearance=0.008, + body_scale=(2.0, 2.0, 2.0), + grasp_end_margin_ratio=0.08, + grasp_z_clearance=0.01, + target_translation=(-0.12, -0.03, 0.12), + target_world_yaw_deg=0.0, + hand_close_qpos=0.026, + grasp_z_ratio=0.55, + ), +} +PICKMENT_SAMPLE_INTERVAL = 96 +PICKMENT_OBJECT_MOTION_KEYFRAMES = 6 +PICKMENT_PRE_GRASP_DISTANCE = 0.11 +PICKMENT_LIFT_HEIGHT = 0.10 +PICKMENT_HAND_INTERP_STEPS = 10 +PICKMENT_HOLD_STEPS = 4 +TRAJECTORY_SIM_STEPS = 4 + + +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) + parser.set_defaults(device="cpu", renderer="hybrid") + parser.add_argument( + "--diagnose_plan", + action="store_true", + help="Plan and print diagnostics without playing the trajectory.", + ) + parser.add_argument( + "--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", + action="store_true", + help="Execute planned trajectories without opening the viewer window.", + ) + parser.add_argument( + "--object", + choices=sorted(OBJECT_PRESETS), + 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() + + +def get_cached_data_path(data_path: str) -> str: + """Resolve an asset path from the local cache before importing data helpers.""" + if os.path.isabs(data_path): + return data_path + + data_root = Path( + os.environ.get( + "EMBODICHAIN_DATA_ROOT", + str(Path.home() / ".cache" / "embodichain_data"), + ) + ) + candidates = ( + data_root / data_path, + data_root / "extract" / data_path, + ) + for candidate in candidates: + if candidate.exists(): + return str(candidate) + + from embodichain.data import get_data_path + + return get_data_path(data_path) + + +def rotation_z(yaw: float) -> np.ndarray: + """Build a 3x3 yaw rotation matrix.""" + cos_yaw = math.cos(yaw) + sin_yaw = math.sin(yaw) + return np.array( + [ + [cos_yaw, -sin_yaw, 0.0], + [sin_yaw, cos_yaw, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + + +def make_transform(xyz: tuple[float, float, float], yaw: float) -> np.ndarray: + """Build a homogeneous transform from translation and yaw.""" + transform = np.eye(4, dtype=np.float32) + transform[:3, :3] = rotation_z(yaw) + transform[:3, 3] = np.asarray(xyz, dtype=np.float32) + return transform + + +def create_dual_ur5_robot(sim: SimulationManager) -> Robot: + """Create a dual-UR5 robot with one PGI gripper on each arm.""" + arm_urdf_path = get_cached_data_path(ARM_URDF_PATH) + gripper_urdf_path = get_cached_data_path(GRIPPER_URDF_PATH) + tcp = [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, GRIPPER_TCP_Z], + [0.0, 0.0, 0.0, 1.0], + ] + cfg = RobotCfg( + uid="DualUR5CoordinatedPickment", + urdf_cfg=URDFCfg( + components=[ + { + "component_type": "left_arm", + "urdf_path": arm_urdf_path, + "transform": make_transform((-0.3, -1.45, 0.4), np.pi / 2), + }, + { + "component_type": "right_arm", + "urdf_path": arm_urdf_path, + "transform": make_transform((0.3, -1.45, 0.4), np.pi / 2), + }, + {"component_type": "left_hand", "urdf_path": gripper_urdf_path}, + {"component_type": "right_hand", "urdf_path": gripper_urdf_path}, + ], + fname="dual_ur5_coordinated_pickment", + name_case={"joint": "upper", "link": "lower"}, + ), + drive_pros=JointDrivePropertiesCfg( + stiffness={ + "LEFT_JOINT[0-9]": 1e4, + "RIGHT_JOINT[0-9]": 1e4, + "LEFT_GRIPPER_FINGER[1-2]_JOINT_1": 1e3, + "RIGHT_GRIPPER_FINGER[1-2]_JOINT_1": 1e3, + }, + damping={ + "LEFT_JOINT[0-9]": 1e3, + "RIGHT_JOINT[0-9]": 1e3, + "LEFT_GRIPPER_FINGER[1-2]_JOINT_1": 1e2, + "RIGHT_GRIPPER_FINGER[1-2]_JOINT_1": 1e2, + }, + max_effort={ + "LEFT_JOINT[0-9]": 1e5, + "RIGHT_JOINT[0-9]": 1e5, + "LEFT_GRIPPER_FINGER[1-2]_JOINT_1": 1e4, + "RIGHT_GRIPPER_FINGER[1-2]_JOINT_1": 1e4, + }, + drive_type="force", + ), + control_parts={ + "left_arm": ["LEFT_JOINT[0-9]"], + "right_arm": ["RIGHT_JOINT[0-9]"], + "dual_arm": ["LEFT_JOINT[0-9]", "RIGHT_JOINT[0-9]"], + "left_hand": ["LEFT_GRIPPER_FINGER1_JOINT_1"], + "right_hand": ["RIGHT_GRIPPER_FINGER1_JOINT_1"], + }, + solver_cfg={ + "left_arm": PytorchSolverCfg( + end_link_name="left_ee_link", + root_link_name="left_base_link", + tcp=tcp, + num_samples=30, + ), + "right_arm": PytorchSolverCfg( + end_link_name="right_ee_link", + root_link_name="right_base_link", + tcp=tcp, + num_samples=30, + ), + }, + init_pos=list(ROBOT_INIT_POS), + init_rot=list(ROBOT_INIT_ROT), + init_qpos=list(LEFT_ARM_HOME) + list(RIGHT_ARM_HOME) + [0.0, 0.0, 0.0, 0.0], + ) + return sim.add_robot(cfg=cfg) + + +def create_support_surface(sim: SimulationManager) -> RigidObject: + """Create a compact support slab under the staged object.""" + return sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="support_surface", + shape=CubeCfg(size=list(SUPPORT_SURFACE_SIZE)), + attrs=RigidBodyAttributesCfg( + mass=10.0, + dynamic_friction=0.9, + static_friction=0.95, + restitution=0.01, + ), + body_type="static", + init_pos=list(SUPPORT_SURFACE_CENTER), + ) + ) + + +def create_pickment_object( + sim: SimulationManager, + preset: PickmentObjectPreset, +) -> RigidObject: + """Create the selected object mesh on the support surface.""" + obj = sim.add_rigid_object( + cfg=RigidObjectCfg( + uid=preset.label, + shape=MeshCfg( + fpath=get_cached_data_path(preset.mesh_path), compute_uv=False + ), + attrs=RigidBodyAttributesCfg( + mass=0.01, + dynamic_friction=0.97, + static_friction=0.99, + angular_damping=1.0, + linear_damping=0.5, + contact_offset=0.001, + rest_offset=0.0, + restitution=0.01, + min_position_iters=32, + min_velocity_iters=8, + max_depenetration_velocity=2.0, + ), + max_convex_hull_num=16, + init_pos=[preset.init_xy[0], preset.init_xy[1], SUPPORT_SURFACE_Z], + init_rot=list(preset.init_rot), + body_scale=preset.body_scale, + ) + ) + obj.cfg.init_pos = compute_supported_init_pos(obj, preset) + obj.reset() + return obj + + +def settle_object(sim: SimulationManager, obj: RigidObject, step: int = 5) -> None: + """Settle an object before planning.""" + if sim.device.type == "cuda": + sim.init_gpu_physics() + obj.reset() + if step > 0: + sim.update(step=step) + obj.clear_dynamics() + + +def create_object_semantics(obj: RigidObject, label: str) -> ObjectSemantics: + """Create minimal object semantics for manually specified grasps.""" + return ObjectSemantics( + label=label, + geometry={}, + affordance=Affordance(object_label=label), + entity=obj, + ) + + +def get_hand_open_close_qpos( + robot: Robot, + hand_control_part: str, + device: torch.device, + close_qpos: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Get open and close qpos for a PGI gripper control part.""" + limits = robot.get_qpos_limits(name=hand_control_part)[0].to( + device=device, dtype=torch.float32 + ) + hand_open = limits[:, 0] + hand_close = torch.clamp( + torch.full_like(limits[:, 1], close_qpos), + min=limits[:, 0], + max=limits[:, 1], + ) + return hand_open, hand_close + + +def get_local_vertices(obj: RigidObject) -> torch.Tensor: + """Get scaled local mesh vertices.""" + return obj.get_vertices(env_ids=[0], scale=True)[0] + + +def compute_local_bounds(vertices: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Compute local mesh AABB from scaled vertices.""" + return vertices.min(dim=0).values, vertices.max(dim=0).values + + +def compute_supported_init_pos( + obj: RigidObject, + preset: PickmentObjectPreset, +) -> tuple[float, float, float]: + """Place an object so its rotated mesh bottom sits on the support surface.""" + vertices = get_local_vertices(obj) + rot = torch.as_tensor(preset.init_rot, dtype=torch.float32, device=vertices.device) + rot = rot.unsqueeze(0) * torch.pi / 180.0 + upright_rot = matrix_from_euler(rot, "XYZ")[0] + rotated_vertices = vertices @ upright_rot.T + bottom_z = rotated_vertices[:, 2].min().item() + z = SUPPORT_SURFACE_Z + preset.surface_clearance - bottom_z + return (preset.init_xy[0], preset.init_xy[1], z) + + +def invert_pose(pose: torch.Tensor) -> torch.Tensor: + """Invert batched homogeneous transforms.""" + inv_pose = pose.clone() + rot_t = pose[:, :3, :3].transpose(1, 2) + inv_pose[:, :3, :3] = rot_t + inv_pose[:, :3, 3] = -torch.bmm(rot_t, pose[:, :3, 3:4]).squeeze(-1) + return inv_pose + + +def transform_points(pose: torch.Tensor, points: torch.Tensor) -> torch.Tensor: + """Transform local points by a homogeneous pose.""" + return points @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + + +def compute_world_bounds( + object_pose: torch.Tensor, + local_vertices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute world AABB from transformed local mesh vertices.""" + world_vertices = transform_points(object_pose, local_vertices) + return world_vertices.min(dim=0).values, world_vertices.max(dim=0).values + + +def normalize_vector(vector: torch.Tensor, fallback: torch.Tensor) -> torch.Tensor: + """Normalize a vector with a deterministic fallback for degenerate cases.""" + norm = torch.linalg.norm(vector) + if norm < 1e-6: + return fallback.to(device=vector.device, dtype=vector.dtype) + return vector / norm + + +def rotate_pose_about_world_z(pose: torch.Tensor, yaw_deg: float) -> torch.Tensor: + """Rotate pose orientation about world Z while preserving translation.""" + yaw = math.radians(yaw_deg) + rot = torch.eye(3, dtype=pose.dtype, device=pose.device) + rot[0, 0] = math.cos(yaw) + rot[0, 1] = -math.sin(yaw) + rot[1, 0] = math.sin(yaw) + rot[1, 1] = math.cos(yaw) + rotated_pose = pose.clone() + rotated_pose[:3, :3] = rot @ pose[:3, :3] + return rotated_pose + + +def build_object_grasp_poses( + object_pose: torch.Tensor, + local_vertices: torch.Tensor, + preset: PickmentObjectPreset, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build left/right TCP poses that pinch opposite sides of the object.""" + local_min, local_max = compute_local_bounds(local_vertices) + extents = local_max - local_min + long_axis_idx = int(torch.argmax(extents).item()) + axis_local = torch.zeros(3, dtype=torch.float32, device=device) + axis_local[long_axis_idx] = 1.0 + long_axis = normalize_vector( + object_pose[:3, :3] @ axis_local, + torch.tensor([1.0, 0.0, 0.0], dtype=torch.float32, device=device), + ) + + local_center = 0.5 * (local_min + local_max) + margin = extents[long_axis_idx] * preset.grasp_end_margin_ratio + left_local = local_center.clone() + right_local = local_center.clone() + left_local[long_axis_idx] = local_min[long_axis_idx] + margin + right_local[long_axis_idx] = local_max[long_axis_idx] - margin + + world_min, world_max = compute_world_bounds(object_pose, local_vertices) + left_position = object_pose[:3, 3] + object_pose[:3, :3] @ left_local.to(device) + right_position = object_pose[:3, 3] + object_pose[:3, :3] @ right_local.to(device) + if preset.grasp_z_ratio is None: + grasp_z = world_max[2] + preset.grasp_z_clearance + else: + grasp_z = ( + world_min[2] + + (world_max[2] - world_min[2]) * preset.grasp_z_ratio + + preset.grasp_z_clearance + ) + left_position[2] = grasp_z + right_position[2] = grasp_z + + z_axis = torch.tensor([0.0, 0.0, -1.0], dtype=torch.float32, device=device) + x_axis = normalize_vector( + torch.cross(long_axis, z_axis, dim=0), + torch.tensor([1.0, 0.0, 0.0], dtype=torch.float32, device=device), + ) + y_axis = normalize_vector(torch.cross(z_axis, x_axis, dim=0), long_axis) + + left_pose = torch.eye(4, dtype=torch.float32, device=device) + left_pose[:3, 0] = x_axis + left_pose[:3, 1] = y_axis + left_pose[:3, 2] = z_axis + left_pose[:3, 3] = left_position + + right_pose = torch.eye(4, dtype=torch.float32, device=device) + right_pose[:3, 0] = -x_axis + right_pose[:3, 1] = -y_axis + right_pose[:3, 2] = z_axis + right_pose[:3, 3] = right_position + return left_pose, right_pose + + +def build_object_target_pose( + object_pose: torch.Tensor, + object_vertices: torch.Tensor, + preset: PickmentObjectPreset, + device: torch.device, +) -> torch.Tensor: + """Build the target pose for the whole object.""" + pose = rotate_pose_about_world_z( + object_pose.clone().to(device=device, dtype=torch.float32), + preset.target_world_yaw_deg, + ) + pose[:3, 3] += torch.tensor( + preset.target_translation, dtype=torch.float32, device=device + ) + bottom_z = compute_world_bounds(pose, object_vertices)[0][2] + pose[2, 3] += SUPPORT_SURFACE_Z + preset.surface_clearance + 0.10 - bottom_z + return pose + + +def format_tensor(tensor: torch.Tensor) -> str: + """Format tensor values for compact logging.""" + rounded = (tensor.detach().cpu() * 10000.0).round() / 10000.0 + return str(rounded.tolist()) + + +def log_action_plan( + robot: Robot, + action_name: str, + traj: torch.Tensor, + joint_ids: list[int], + segments: dict[str, int] | None = None, +) -> None: + """Log common action plan details.""" + joint_names = [robot.joint_names[joint_id] for joint_id in joint_ids] + logger.log_info(f"{action_name} joint ids: {joint_ids}") + logger.log_info(f"{action_name} joint names: {joint_names}") + logger.log_info(f"{action_name} trajectory shape: {tuple(traj.shape)}") + if segments is not None: + logger.log_info(f"{action_name} trajectory segments: {segments}") + + +def log_scene_targets( + object_label: str, + object_pose: torch.Tensor, + target_pose: torch.Tensor, + left_grasp_pose: torch.Tensor, + right_grasp_pose: torch.Tensor, +) -> None: + """Log compact object and grasp target positions.""" + logger.log_info( + "pickment scene: " + f"object={object_label}, " + f"object_origin={format_tensor(object_pose[:3, 3])}, " + f"target_origin={format_tensor(target_pose[:3, 3])}, " + f"left_grasp={format_tensor(left_grasp_pose[:3, 3])}, " + f"right_grasp={format_tensor(right_grasp_pose[:3, 3])}" + ) + + +def draw_pickment_target_axes( + sim: SimulationManager, + object_target_pose: torch.Tensor, + left_grasp_pose: torch.Tensor, + right_grasp_pose: torch.Tensor, + num_envs: int, +) -> None: + """Draw semantic axes for the target object pose and two grasp TCP poses.""" + draw_axis_marker( + sim, + "coordinated_pickment_object_target_axis", + broadcast_pose_batch(object_target_pose, num_envs=num_envs), + axis_len=0.12, + axis_size=0.005, + ) + draw_axis_marker( + sim, + "coordinated_pickment_left_grasp_axis", + broadcast_pose_batch(left_grasp_pose, num_envs=num_envs), + axis_len=0.07, + axis_size=0.0035, + ) + draw_axis_marker( + sim, + "coordinated_pickment_right_grasp_axis", + broadcast_pose_batch(right_grasp_pose, num_envs=num_envs), + axis_len=0.07, + axis_size=0.0035, + ) + + +def log_execution_state( + robot: Robot, + obj: RigidObject, + step_idx: int, + total_steps: int, +) -> None: + """Log hand and object state during execution.""" + object_pose = obj.get_local_pose(to_matrix=True) + left_hand = robot.get_qpos(name="left_hand") + right_hand = robot.get_qpos(name="right_hand") + logger.log_info( + f"step={step_idx}/{total_steps - 1}, " + f"left_hand={format_tensor(left_hand[0])}, " + f"right_hand={format_tensor(right_hand[0])}, " + f"{obj.uid}_pos={format_tensor(object_pose[0, :3, 3])}" + ) + + +def run_coordinated_pickment_demo( + args: argparse.Namespace, + sim: SimulationManager, + robot: Robot, +) -> None: + """Plan and optionally execute coordinated object pickment.""" + preset = OBJECT_PRESETS[args.object] + create_support_surface(sim) + obj = create_pickment_object(sim, preset) + settle_object(sim, obj, step=0) + object_pose_batch = clone_local_pose_from_first_env(obj) + obj.clear_dynamics() + object_pose = object_pose_batch[0].to(device=sim.device, dtype=torch.float32) + n_envs = object_pose_batch.shape[0] + object_vertices = get_local_vertices(obj) + object_semantics = create_object_semantics(obj, preset.label) + motion_gen = create_toppra_motion_generator(robot) + + left_open, left_close = get_hand_open_close_qpos( + robot, "left_hand", sim.device, preset.hand_close_qpos + ) + right_open, right_close = get_hand_open_close_qpos( + robot, "right_hand", sim.device, preset.hand_close_qpos + ) + pickment_action = CoordinatedPickment( + motion_generator=motion_gen, + cfg=CoordinatedPickmentCfg( + control_part="dual_arm", + left_arm_control_part="left_arm", + right_arm_control_part="right_arm", + left_hand_control_part="left_hand", + right_hand_control_part="right_hand", + left_hand_open_qpos=left_open, + left_hand_close_qpos=left_close, + right_hand_open_qpos=right_open, + right_hand_close_qpos=right_close, + pre_grasp_distance=PICKMENT_PRE_GRASP_DISTANCE, + lift_height=PICKMENT_LIFT_HEIGHT, + sample_interval=PICKMENT_SAMPLE_INTERVAL, + hand_interp_steps=PICKMENT_HAND_INTERP_STEPS, + hold_steps=PICKMENT_HOLD_STEPS, + object_motion_keyframes=PICKMENT_OBJECT_MOTION_KEYFRAMES, + ), + ) + engine = AtomicActionEngine(motion_generator=motion_gen) + engine.register(pickment_action) + + left_grasp_pose, right_grasp_pose = build_object_grasp_poses( + object_pose, + object_vertices, + preset, + sim.device, + ) + target_pose = build_object_target_pose( + object_pose, + object_vertices, + preset, + sim.device, + ) + log_scene_targets( + preset.label, + object_pose, + target_pose, + left_grasp_pose, + right_grasp_pose, + ) + if not args.no_vis_eef_axis: + draw_pickment_target_axes( + sim, + target_pose, + left_grasp_pose, + right_grasp_pose, + num_envs=n_envs, + ) + + left_object_to_eef = torch.bmm( + broadcast_pose_batch(invert_pose(object_pose.unsqueeze(0)), num_envs=n_envs), + broadcast_pose_batch(left_grasp_pose, num_envs=n_envs), + ) + right_object_to_eef = torch.bmm( + broadcast_pose_batch(invert_pose(object_pose.unsqueeze(0)), num_envs=n_envs), + broadcast_pose_batch(right_grasp_pose, num_envs=n_envs), + ) + pickment_target = CoordinatedPickmentTarget( + object_target_pose=broadcast_pose_batch(target_pose, num_envs=n_envs), + object_semantics=object_semantics, + left_object_to_eef=left_object_to_eef, + right_object_to_eef=right_object_to_eef, + object_initial_pose=broadcast_pose_batch(object_pose, num_envs=n_envs), + ) + + wait_for_user = prepare_tutorial_scene( + sim, args, "Inspect the scene, then press Enter to plan pickment..." + ) + + start_time = time.time() + success, traj, _ = engine.run([("coordinated_pickment", pickment_target)]) + logger.log_info( + f"Plan coordinated pickment cost time: {time.time() - start_time:.2f} seconds" + ) + if not success.all(): + logger.log_warning("Failed to plan coordinated pickment trajectory.") + return + joint_ids = list(range(robot.dof)) + log_action_plan( + robot, + "coordinated_pickment", + traj, + joint_ids, + pickment_action.get_segment_lengths(), + ) + + if args.diagnose_plan: + return + + if wait_for_user: + input("Press Enter to execute coordinated pickment...") + + def log_execution(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 + ): + log_execution_state(robot, obj, step_idx, total_steps) + + replay_trajectory( + sim, + robot, + traj, + args, + video_prefix=f"coordinated_pickment_{args.object}_auto_play", + hold_steps=0, + trajectory_sim_steps=TRAJECTORY_SIM_STEPS, + on_trajectory_step=log_execution, + look_at=PICKMENT_RECORD_LOOK_AT, + ) + if wait_for_user: + input("Press Enter to exit the simulation...") + + +def main() -> None: + """Run the coordinated pickment demo.""" + args = parse_arguments() + sim = create_tutorial_simulation( + args, + 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) + + +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/atomic_action/coordinated_placement.py b/scripts/tutorials/atomic_action/coordinated_placement.py new file mode 100644 index 000000000..f6b89073a --- /dev/null +++ b/scripts/tutorials/atomic_action/coordinated_placement.py @@ -0,0 +1,1073 @@ +# ---------------------------------------------------------------------------- +# 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 dual-arm coordinated placement with bread and pan meshes. + +The left UR5 picks up bread. The right UR5 picks up a pan and moves it to the +lower alignment pose. The left UR5 places the bread above the pan and releases +it while the right hand keeps holding the pan. +""" + +from __future__ import annotations + +import argparse +import math +import os +import sys +import time +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 numpy as np +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, + AtomicActionEngine, + CoordinatedPlacement, + CoordinatedPlacementCfg, + CoordinatedPlacementTarget, + GraspTarget, + HeldObjectState, + ObjectSemantics, + PickUp, + PickUpCfg, + WorldState, +) +from embodichain.lab.sim.cfg import ( + JointDrivePropertiesCfg, + RigidBodyAttributesCfg, + RigidObjectCfg, + RobotCfg, + URDFCfg, +) +from embodichain.lab.sim.objects import RigidObject, Robot +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.solvers import URSolverCfg +from embodichain.utils import logger +from scripts.tutorials.atomic_action.tutorial_utils import ( + broadcast_pose_batch, + clone_local_pose_from_first_env, + create_toppra_motion_generator, + create_tutorial_simulation, + draw_axis_marker, + make_ur5_solver_cfg, + prepare_tutorial_scene, + replay_trajectory, +) + +DEFAULT_MESH_FRAME_CORRECTION_EULER_DEG = (-90.0, 0.0, 0.0) +# DexSim imports this pan GLB with gym raw +Z mapped to local -Y. The +90deg X +# correction therefore makes the pan opening point upward in world Z. +PAN_MESH_FRAME_CORRECTION_EULER_DEG = (90.0, 0.0, 0.0) +PAN_WORLD_YAW_CORRECTION_DEG = 270.0 + + +def transform_baseline_pose( + init_pos: tuple[float, float, float], + init_rot: tuple[float, float, float], + *, + z_offset: float = 0.0, + mesh_frame_correction_euler_deg: tuple[float, float, float] = ( + DEFAULT_MESH_FRAME_CORRECTION_EULER_DEG + ), + world_yaw_correction_deg: float = 0.0, +) -> tuple[tuple[float, float, float], tuple[float, float, float]]: + """Apply mesh-frame correction while preserving baseline world placement.""" + pos = np.asarray(init_pos, dtype=np.float64) + pos[2] += z_offset + rot = ( + SciRotation.from_euler("Z", world_yaw_correction_deg, degrees=True) + * SciRotation.from_euler("XYZ", init_rot, degrees=True) + * SciRotation.from_euler("XYZ", mesh_frame_correction_euler_deg, degrees=True) + ).as_euler("XYZ", degrees=True) + return tuple(float(value) for value in pos), tuple(float(value) for value in rot) + + +ARM_URDF_PATH = "UniversalRobots/UR5/UR5.urdf" +GRIPPER_URDF_PATH = "DH_PGI_140_80/DH_PGI_140_80.urdf" +PLACEMENT_ASSET_ROOT = "CoordinatedPlacementAndPickment" +TABLE_MESH_PATH = f"{PLACEMENT_ASSET_ROOT}/table.glb" +BREAD_MESH_PATH = f"{PLACEMENT_ASSET_ROOT}/bread.glb" +PAN_MESH_PATH = f"{PLACEMENT_ASSET_ROOT}/pan.glb" +BREAD_LABEL = "bread" +PAN_LABEL = "pan" +GRIPPER_TCP_Z = 0.121 +PICK_SAMPLE_INTERVAL = 100 +COORDINATED_SAMPLE_INTERVAL = 120 +ROBOT_INIT_POS = (1.85, 0.0, 0.1) +ROBOT_INIT_ROT = (0.0, 0.0, -90.0) +LEFT_ARM_HOME = (0.0, 0.0, -1.57, -1.57, 1.57, 1.57) +RIGHT_ARM_HOME = (-1.57, -1.57, -1.57, -1.57, 0.0, 0.0) +TABLE_TOP_Z = 0.65 +BASELINE_TABLE_TOP_Z = 0.3621708124799265 +SCENE_Z_OFFSET = TABLE_TOP_Z - BASELINE_TABLE_TOP_Z +BASELINE_TABLE_INIT_POS = ( + 0.00014585733079742588, + 0.00023304896730074557, + -0.019599792839044783, +) +BASELINE_TABLE_INIT_ROT = ( + 0.0001074673904926984, + 0.00865572768366991, + -90.6562109309317, +) +BASELINE_BREAD_INIT_POS = ( + 0.007266042530919159, + 0.17218712515099063, + 0.38805152145807564, +) +BASELINE_BREAD_INIT_ROT = ( + 179.93952112929065, + -0.12776179446053365, + 85.59207565132371, +) +BASELINE_PAN_INIT_POS = ( + 0.0009683294205463406, + -0.14189524793277888, + 0.38900474548025743, +) +BASELINE_PAN_INIT_ROT = ( + -179.23950670370294, + -0.4795764805552328, + 98.19364391929443, +) +TABLE_INIT_POS, TABLE_INIT_ROT = transform_baseline_pose( + BASELINE_TABLE_INIT_POS, + BASELINE_TABLE_INIT_ROT, + z_offset=SCENE_Z_OFFSET, +) +BREAD_INIT_POS, BREAD_INIT_ROT = transform_baseline_pose( + BASELINE_BREAD_INIT_POS, + BASELINE_BREAD_INIT_ROT, + z_offset=SCENE_Z_OFFSET, +) +PAN_INIT_POS, PAN_INIT_ROT = transform_baseline_pose( + BASELINE_PAN_INIT_POS, + BASELINE_PAN_INIT_ROT, + z_offset=SCENE_Z_OFFSET, + mesh_frame_correction_euler_deg=PAN_MESH_FRAME_CORRECTION_EULER_DEG, + world_yaw_correction_deg=PAN_WORLD_YAW_CORRECTION_DEG, +) +PAN_INIT_POS = (PAN_INIT_POS[0], PAN_INIT_POS[1], TABLE_TOP_Z + 0.001) +PAN_TARGET_CENTER_XY = (-0.06, 0.0) +PAN_TARGET_Z_LIFT = 0.06 +BREAD_PLACE_TARGET_OFFSET_XY = (-0.06, -0.16) +BREAD_ON_PAN_CLEARANCE = 0.006 +BREAD_GRASP_Z_CLEARANCE = 0.018 +PAN_GRASP_Z_CLEARANCE = 0.0 +PAN_HANDLE_LOCAL_Z_MIN = 0.04 +PAN_BASIN_LOCAL_Z_MAX = 0.04 +PAN_HANDLE_ROOT_OFFSET = 0.035 +PAN_HANDLE_CLOSE_QPOS = 0.045 +PAN_PICK_SAMPLE_INTERVAL = 130 +PAN_PICK_HAND_INTERP_STEPS = 32 +BREAD_TARGET_WORLD_YAW_DEG = 0.0 +BREAD_TARGET_HEIGHT_OFFSET = 0.1 +SUPPORT_TARGET_HEIGHT_OFFSET = 0.0 +PICK_APPROACH_DISTANCE = 0.12 +PLACE_LIFT_HEIGHT = 0.10 +TRAJECTORY_SIM_STEPS = 8 + + +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) + parser.set_defaults(device="cuda", renderer="hybrid") + parser.add_argument( + "--diagnose_plan", + action="store_true", + help="Plan and print diagnostics without playing the trajectory.", + ) + parser.add_argument( + "--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", + action="store_true", + help="Execute planned trajectories without opening the viewer window.", + ) + return parser.parse_args() + + +def get_cached_data_path(data_path: str) -> str: + """Resolve an asset path from the local cache before importing data helpers.""" + if os.path.isabs(data_path): + return data_path + + data_root = Path( + os.environ.get( + "EMBODICHAIN_DATA_ROOT", + str(Path.home() / ".cache" / "embodichain_data"), + ) + ) + candidates = ( + data_root / data_path, + data_root / "extract" / data_path, + ) + for candidate in candidates: + if candidate.exists(): + return str(candidate) + + from embodichain.data import get_data_path + + return get_data_path(data_path) + + +def rotation_z(yaw: float) -> np.ndarray: + """Build a 3x3 yaw rotation matrix.""" + cos_yaw = math.cos(yaw) + sin_yaw = math.sin(yaw) + return np.array( + [ + [cos_yaw, -sin_yaw, 0.0], + [sin_yaw, cos_yaw, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + + +def make_transform(xyz: tuple[float, float, float], yaw: float) -> np.ndarray: + """Build a homogeneous transform from translation and yaw.""" + transform = np.eye(4, dtype=np.float32) + transform[:3, :3] = rotation_z(yaw) + transform[:3, 3] = np.asarray(xyz, dtype=np.float32) + return transform + + +def make_prefixed_ur5_solver_cfg(prefix: str) -> URSolverCfg: + """Create a UR5 solver cfg for a prefixed arm in the assembled robot.""" + cfg = make_ur5_solver_cfg(GRIPPER_TCP_Z) + cfg.root_link_name = f"{prefix}_base_link" + cfg.end_link_name = f"{prefix}_ee_link" + return cfg + + +def create_dual_ur5_robot(sim: SimulationManager) -> Robot: + """Create a dual-UR5 robot with one PGI gripper on each arm.""" + arm_urdf_path = get_cached_data_path(ARM_URDF_PATH) + gripper_urdf_path = get_cached_data_path(GRIPPER_URDF_PATH) + cfg = RobotCfg( + uid="DualUR5CoordinatedPlacement", + urdf_cfg=URDFCfg( + components=[ + { + "component_type": "left_arm", + "urdf_path": arm_urdf_path, + "transform": make_transform((-0.3, -1.45, 0.4), np.pi / 2), + }, + { + "component_type": "right_arm", + "urdf_path": arm_urdf_path, + "transform": make_transform((0.3, -1.45, 0.4), np.pi / 2), + }, + {"component_type": "left_hand", "urdf_path": gripper_urdf_path}, + {"component_type": "right_hand", "urdf_path": gripper_urdf_path}, + ], + fname="dual_ur5_coordinated_placement", + ), + drive_pros=JointDrivePropertiesCfg( + stiffness={ + "left_joint[0-9]": 1e4, + "right_joint[0-9]": 1e4, + "left_gripper_finger[1-2]_joint_1": 1e3, + "right_gripper_finger[1-2]_joint_1": 1e3, + }, + damping={ + "left_joint[0-9]": 1e3, + "right_joint[0-9]": 1e3, + "left_gripper_finger[1-2]_joint_1": 1e2, + "right_gripper_finger[1-2]_joint_1": 1e2, + }, + max_effort={ + "left_joint[0-9]": 1e5, + "right_joint[0-9]": 1e5, + "left_gripper_finger[1-2]_joint_1": 1e4, + "right_gripper_finger[1-2]_joint_1": 1e4, + }, + drive_type="force", + ), + control_parts={ + "left_arm": ["left_joint[0-9]"], + "right_arm": ["right_joint[0-9]"], + "dual_arm": ["left_joint[0-9]", "right_joint[0-9]"], + "left_hand": ["left_gripper_finger1_joint_1"], + "right_hand": ["right_gripper_finger1_joint_1"], + }, + solver_cfg={ + "left_arm": make_prefixed_ur5_solver_cfg("left"), + "right_arm": make_prefixed_ur5_solver_cfg("right"), + }, + init_pos=list(ROBOT_INIT_POS), + init_rot=list(ROBOT_INIT_ROT), + init_qpos=list(LEFT_ARM_HOME) + list(RIGHT_ARM_HOME) + [0.0, 0.0, 0.0, 0.0], + ) + return sim.add_robot(cfg=cfg) + + +def create_table(sim: SimulationManager) -> RigidObject: + """Create the table mesh from the bread-pan gym project.""" + return sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="table", + shape=MeshCfg(fpath=get_cached_data_path(TABLE_MESH_PATH)), + attrs=RigidBodyAttributesCfg( + mass=10.0, + dynamic_friction=0.9, + static_friction=0.95, + restitution=0.01, + ), + body_type="kinematic", + init_pos=list(TABLE_INIT_POS), + init_rot=list(TABLE_INIT_ROT), + ) + ) + + +def create_bread(sim: SimulationManager) -> RigidObject: + """Create the bread mesh to be placed by the left arm.""" + return sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="bread", + shape=MeshCfg( + fpath=get_cached_data_path(BREAD_MESH_PATH), compute_uv=False + ), + attrs=RigidBodyAttributesCfg( + mass=0.01, + contact_offset=0.003, + rest_offset=0.001, + restitution=0.01, + min_position_iters=32, + min_velocity_iters=8, + max_depenetration_velocity=10.0, + ), + body_scale=(1.75, 1.75, 1.75), + max_convex_hull_num=8, + init_pos=list(BREAD_INIT_POS), + init_rot=list(BREAD_INIT_ROT), + ) + ) + + +def create_pan(sim: SimulationManager) -> RigidObject: + """Create the pan mesh held below the bread by the right arm.""" + return sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="pan", + shape=MeshCfg(fpath=get_cached_data_path(PAN_MESH_PATH), compute_uv=False), + attrs=RigidBodyAttributesCfg( + mass=0.01, + dynamic_friction=0.97, + static_friction=0.99, + angular_damping=2.0, + linear_damping=1.0, + contact_offset=0.001, + rest_offset=0.0, + restitution=0.01, + min_position_iters=32, + min_velocity_iters=8, + max_depenetration_velocity=2.0, + ), + body_scale=(1.75, 1.75, 1.75), + max_convex_hull_num=16, + init_pos=list(PAN_INIT_POS), + init_rot=list(PAN_INIT_ROT), + ) + ) + + +def settle_object(sim: SimulationManager, obj: RigidObject, step: int = 5) -> None: + """Settle an object before planning.""" + if sim.device.type == "cuda": + sim.init_gpu_physics() + obj.reset() + if step > 0: + sim.update(step=step) + obj.clear_dynamics() + + +def create_object_semantics(obj: RigidObject, label: str) -> ObjectSemantics: + """Create minimal object semantics for manually specified grasps.""" + return ObjectSemantics( + label=label, + geometry={}, + affordance=Affordance(object_label=label), + entity=obj, + ) + + +def get_hand_open_close_qpos( + robot: Robot, hand_control_part: str, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + """Get open and close qpos for a PGI gripper control part.""" + limits = robot.get_qpos_limits(name=hand_control_part)[0].to( + device=device, dtype=torch.float32 + ) + hand_open = limits[:, 0] + hand_close = torch.minimum(limits[:, 1], torch.full_like(limits[:, 1], 0.030)) + return hand_open, hand_close + + +def get_pan_handle_open_close_qpos( + robot: Robot, hand_control_part: str, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + """Get right hand qpos tuned for holding the thin pan handle.""" + limits = robot.get_qpos_limits(name=hand_control_part)[0].to( + device=device, dtype=torch.float32 + ) + hand_open = limits[:, 0] + hand_close = torch.clamp( + torch.full_like(limits[:, 1], PAN_HANDLE_CLOSE_QPOS), + min=limits[:, 0], + max=limits[:, 1], + ) + return hand_open, hand_close + + +def invert_pose(pose: torch.Tensor) -> torch.Tensor: + """Invert batched homogeneous transforms.""" + inv_pose = pose.clone() + rot_t = pose[:, :3, :3].transpose(1, 2) + inv_pose[:, :3, :3] = rot_t + inv_pose[:, :3, 3] = -torch.bmm(rot_t, pose[:, :3, 3:4]).squeeze(-1) + return inv_pose + + +def get_local_vertices(obj: RigidObject) -> torch.Tensor: + """Get scaled local mesh vertices.""" + return obj.get_vertices(env_ids=[0], scale=True)[0] + + +def compute_local_bounds(vertices: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Compute local mesh AABB from scaled vertices.""" + return vertices.min(dim=0).values, vertices.max(dim=0).values + + +def transform_points(pose: torch.Tensor, points: torch.Tensor) -> torch.Tensor: + """Transform local points by a homogeneous pose.""" + return points @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + + +def compute_world_bounds( + object_pose: torch.Tensor, + local_vertices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute world AABB from transformed local mesh vertices.""" + world_vertices = transform_points(object_pose, local_vertices) + return world_vertices.min(dim=0).values, world_vertices.max(dim=0).values + + +def rotate_pose_about_world_z(pose: torch.Tensor, yaw_deg: float) -> torch.Tensor: + """Rotate pose orientation about world Z while preserving translation.""" + yaw = math.radians(yaw_deg) + rot = torch.eye(3, dtype=pose.dtype, device=pose.device) + rot[0, 0] = math.cos(yaw) + rot[0, 1] = -math.sin(yaw) + rot[1, 0] = math.sin(yaw) + rot[1, 1] = math.cos(yaw) + rotated_pose = pose.clone() + rotated_pose[:3, :3] = rot @ pose[:3, :3] + return rotated_pose + + +def get_pan_basin_vertices(pan_vertices: torch.Tensor) -> torch.Tensor: + """Select pan vertices that belong to the basin instead of the long handle.""" + basin_vertices = pan_vertices[pan_vertices[:, 2] <= PAN_BASIN_LOCAL_Z_MAX] + if basin_vertices.numel() == 0: + logger.log_warning("Pan basin vertex mask is empty; falling back to full mesh.") + return pan_vertices + return basin_vertices + + +def build_top_down_tcp_pose( + position: torch.Tensor, device: torch.device +) -> torch.Tensor: + """Build a simple top-down TCP pose for manually grasping flat objects.""" + pose = torch.eye(4, dtype=torch.float32, device=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=device, + ) + pose[:3, 3] = position + return pose + + +def build_flat_object_grasp_pose( + object_pose: torch.Tensor, + local_vertices: torch.Tensor, + local_min: torch.Tensor, + local_max: torch.Tensor, + device: torch.device, + *, + world_xy_offset: tuple[float, float] = (0.0, 0.0), + z_clearance: float = 0.02, +) -> torch.Tensor: + """Build a hand-tuned top-down grasp TCP pose over a flat object.""" + local_center = 0.5 * (local_min + local_max) + local_center = local_center.to(device=device, dtype=torch.float32) + grasp_position = object_pose[:3, 3] + object_pose[:3, :3] @ local_center + _, world_max = compute_world_bounds(object_pose, local_vertices) + grasp_position[0] += world_xy_offset[0] + grasp_position[1] += world_xy_offset[1] + grasp_position[2] = world_max[2] + z_clearance + return build_top_down_tcp_pose(grasp_position, device) + + +def normalize_vector(vector: torch.Tensor, fallback: torch.Tensor) -> torch.Tensor: + """Normalize a vector with a deterministic fallback for degenerate cases.""" + norm = torch.linalg.norm(vector) + if norm < 1e-6: + return fallback.to(device=vector.device, dtype=vector.dtype) + return vector / norm + + +def build_pan_handle_grasp_pose( + pan_pose: torch.Tensor, + pan_vertices: torch.Tensor, + device: torch.device, + *, + z_clearance: float = 0.006, +) -> torch.Tensor: + """Build a top-down TCP pose that pinches the pan handle.""" + handle_vertices = pan_vertices[pan_vertices[:, 2] > PAN_HANDLE_LOCAL_Z_MIN] + if handle_vertices.numel() == 0: + logger.log_warning( + "Pan handle vertex mask is empty; falling back to full mesh." + ) + handle_vertices = pan_vertices + + handle_world = transform_points(pan_pose, handle_vertices) + handle_min = handle_world.min(dim=0).values + handle_max = handle_world.max(dim=0).values + grasp_position = 0.5 * (handle_min + handle_max) + + pan_world = transform_points(pan_pose, pan_vertices) + pan_center_xy = pan_world[:, :2].mean(dim=0) + handle_dir_xy = grasp_position[:2] - pan_center_xy + handle_axis = normalize_vector( + torch.tensor( + [handle_dir_xy[0], handle_dir_xy[1], 0.0], + dtype=torch.float32, + device=device, + ), + torch.tensor([1.0, -0.2, 0.0], dtype=torch.float32, device=device), + ) + grasp_position[:2] -= handle_axis[:2] * PAN_HANDLE_ROOT_OFFSET + grasp_position[2] = handle_max[2] + z_clearance + + y_axis = handle_axis + z_axis = torch.tensor([0.0, 0.0, -1.0], dtype=torch.float32, device=device) + x_axis = normalize_vector( + torch.cross(y_axis, z_axis, dim=0), + torch.tensor([1.0, 0.0, 0.0], dtype=torch.float32, device=device), + ) + y_axis = normalize_vector( + torch.cross(z_axis, x_axis, dim=0), + y_axis, + ) + + pose = torch.eye(4, dtype=torch.float32, device=device) + pose[:3, 0] = x_axis + pose[:3, 1] = y_axis + pose[:3, 2] = z_axis + pose[:3, 3] = grasp_position + return pose + + +def build_support_object_target_pose( + pan_pose: torch.Tensor, + device: torch.device, +) -> torch.Tensor: + """Build target pose for the support pan.""" + pose = pan_pose.clone().to(device=device, dtype=torch.float32) + pose[0, 3] = PAN_TARGET_CENTER_XY[0] + pose[1, 3] = PAN_TARGET_CENTER_XY[1] + pose[2, 3] = TABLE_TOP_Z + 0.001 + PAN_TARGET_Z_LIFT + return pose + + +def build_placing_object_target_pose( + bread_pose: torch.Tensor, + bread_vertices: torch.Tensor, + pan_vertices: torch.Tensor, + support_target_pose: torch.Tensor, + device: torch.device, +) -> torch.Tensor: + """Build target bread pose aligned above the pan.""" + pose = rotate_pose_about_world_z( + bread_pose.clone().to(device=device, dtype=torch.float32), + BREAD_TARGET_WORLD_YAW_DEG, + ) + pan_basin_world = transform_points( + support_target_pose, get_pan_basin_vertices(pan_vertices) + ) + basin_center_xy = 0.5 * ( + pan_basin_world[:, :2].min(dim=0).values + + pan_basin_world[:, :2].max(dim=0).values + ) + pose[0, 3] = basin_center_xy[0] + BREAD_PLACE_TARGET_OFFSET_XY[0] + pose[1, 3] = basin_center_xy[1] + BREAD_PLACE_TARGET_OFFSET_XY[1] + pan_top_z = pan_basin_world[:, 2].max() + bread_bottom_z = compute_world_bounds(pose, bread_vertices)[0][2] + pose[2, 3] += pan_top_z + BREAD_ON_PAN_CLEARANCE - bread_bottom_z + return pose + + +def format_tensor(tensor: torch.Tensor) -> str: + """Format tensor values for compact logging.""" + rounded = (tensor.detach().cpu() * 10000.0).round() / 10000.0 + return str(rounded.tolist()) + + +def compute_actual_held_state( + robot: Robot, + semantics: ObjectSemantics, + object_pose: torch.Tensor, + arm_control_part: str, + device: torch.device, +) -> HeldObjectState: + """Build held-object state from current object pose and current TCP FK.""" + arm_qpos = robot.get_qpos(name=arm_control_part).to(device=device) + tcp_pose = robot.compute_fk( + arm_qpos, + name=arm_control_part, + to_matrix=True, + ).to(device=device, dtype=torch.float32) + object_pose = broadcast_pose_batch( + object_pose.to(device=device, dtype=torch.float32), + num_envs=tcp_pose.shape[0], + ) + object_to_eef = torch.bmm(invert_pose(object_pose), tcp_pose) + return HeldObjectState( + semantics=semantics, + object_to_eef=object_to_eef, + grasp_xpos=tcp_pose, + ) + + +def log_scene_targets( + bread_pose: torch.Tensor, + pan_pose: torch.Tensor, + support_target_pose: torch.Tensor | None = None, + placing_target_pose: torch.Tensor | None = None, +) -> None: + """Log compact object and target positions for diagnosis.""" + logger.log_info( + "scene objects: " + f"bread_origin={format_tensor(bread_pose[:3, 3])}, " + f"pan_origin={format_tensor(pan_pose[:3, 3])}" + ) + if support_target_pose is not None and placing_target_pose is not None: + logger.log_info( + "coordinated targets: " + f"support_pan_origin={format_tensor(support_target_pose[:3, 3])}, " + f"placing_bread_origin={format_tensor(placing_target_pose[:3, 3])}" + ) + + +def draw_coordinated_axes( + sim: SimulationManager, + support_target_pose: torch.Tensor, + placing_target_pose: torch.Tensor, + num_envs: int, +) -> None: + """Draw coordinate-frame markers for coordinated placement targets.""" + draw_axis_marker( + sim, + "support_pan_target_axis", + broadcast_pose_batch(support_target_pose, num_envs=num_envs), + axis_len=0.08, + axis_size=0.004, + ) + draw_axis_marker( + sim, + "placing_bread_target_axis", + broadcast_pose_batch(placing_target_pose, num_envs=num_envs), + axis_len=0.08, + axis_size=0.004, + ) + + +def log_action_plan( + robot: Robot, + action_name: str, + traj: torch.Tensor, + joint_ids: list[int], + segments: dict[str, int] | None = None, +) -> None: + """Log common action plan details.""" + joint_names = [robot.joint_names[joint_id] for joint_id in joint_ids] + logger.log_info(f"{action_name} joint ids: {joint_ids}") + logger.log_info(f"{action_name} joint names: {joint_names}") + logger.log_info(f"{action_name} trajectory shape: {tuple(traj.shape)}") + if segments is not None: + logger.log_info(f"{action_name} trajectory segments: {segments}") + + +def log_execution_state( + robot: Robot, + bread: RigidObject, + pan: RigidObject, + step_idx: int, + total_steps: int, +) -> None: + """Log hand and object state during execution.""" + bread_pose = bread.get_local_pose(to_matrix=True) + pan_pose = pan.get_local_pose(to_matrix=True) + left_hand = robot.get_qpos(name="left_hand") + right_hand = robot.get_qpos(name="right_hand") + logger.log_info( + f"step={step_idx}/{total_steps - 1}, " + f"left_hand={format_tensor(left_hand[0])}, " + f"right_hand={format_tensor(right_hand[0])}, " + f"bread_pos={format_tensor(bread_pose[0, :3, 3])}, " + f"pan_pos={format_tensor(pan_pose[0, :3, 3])}" + ) + + +def run_coordinated_placement_demo( + args: argparse.Namespace, sim: SimulationManager, robot: Robot +) -> None: + """Plan and optionally execute pick-up and coordinated placement.""" + create_table(sim) + bread = create_bread(sim) + pan = create_pan(sim) + settle_object(sim, bread, step=0) + settle_object(sim, pan, step=0) + bread_pose_batch = clone_local_pose_from_first_env(bread) + pan_pose_batch = clone_local_pose_from_first_env(pan) + bread.clear_dynamics() + pan.clear_dynamics() + bread_pose = bread_pose_batch[0].to(device=sim.device, dtype=torch.float32) + pan_pose = pan_pose_batch[0].to(device=sim.device, dtype=torch.float32) + n_envs = bread_pose_batch.shape[0] + bread_vertices = get_local_vertices(bread) + pan_vertices = get_local_vertices(pan) + bread_local_min, bread_local_max = compute_local_bounds(bread_vertices) + log_scene_targets(bread_pose, pan_pose) + bread_semantics = create_object_semantics(bread, BREAD_LABEL) + pan_semantics = create_object_semantics(pan, PAN_LABEL) + motion_gen = create_toppra_motion_generator(robot) + + right_open, right_close = get_pan_handle_open_close_qpos( + robot, "right_hand", sim.device + ) + left_open, left_close = get_hand_open_close_qpos(robot, "left_hand", sim.device) + left_pick_action = PickUp( + motion_generator=motion_gen, + cfg=PickUpCfg( + control_part="left_arm", + hand_control_part="left_hand", + hand_open_qpos=left_open, + hand_close_qpos=left_close, + pre_grasp_distance=PICK_APPROACH_DISTANCE, + lift_height=0.12, + sample_interval=PICK_SAMPLE_INTERVAL, + hand_interp_steps=10, + ), + ) + right_pick_action = PickUp( + motion_generator=motion_gen, + cfg=PickUpCfg( + control_part="right_arm", + hand_control_part="right_hand", + hand_open_qpos=right_open, + hand_close_qpos=right_close, + pre_grasp_distance=PICK_APPROACH_DISTANCE, + lift_height=0.10, + sample_interval=PAN_PICK_SAMPLE_INTERVAL, + hand_interp_steps=PAN_PICK_HAND_INTERP_STEPS, + ), + ) + coordinated_action = CoordinatedPlacement( + motion_generator=motion_gen, + cfg=CoordinatedPlacementCfg( + control_part="dual_arm", + placing_arm_control_part="left_arm", + support_arm_control_part="right_arm", + placing_hand_control_part="left_hand", + support_hand_control_part="right_hand", + placing_hand_open_qpos=left_open, + placing_hand_close_qpos=left_close, + support_hand_close_qpos=right_close, + release=True, + placing_height_offset=BREAD_TARGET_HEIGHT_OFFSET, + support_height_offset=SUPPORT_TARGET_HEIGHT_OFFSET, + lift_height=PLACE_LIFT_HEIGHT, + sample_interval=COORDINATED_SAMPLE_INTERVAL, + hand_interp_steps=10, + hold_steps=6, + retreat_steps=18, + ), + ) + engine = AtomicActionEngine(motion_generator=motion_gen) + engine.register(coordinated_action) + full_joint_ids = list(range(robot.dof)) + state = WorldState(last_qpos=robot.get_qpos().clone()) + + wait_for_user = prepare_tutorial_scene( + sim, args, "Inspect the scene, then press Enter to plan left pick-up..." + ) + + bread_grasp_pose = build_flat_object_grasp_pose( + bread_pose, + bread_vertices, + bread_local_min, + bread_local_max, + sim.device, + z_clearance=BREAD_GRASP_Z_CLEARANCE, + ) + start_time = time.time() + left_pick_result = left_pick_action.execute( + GraspTarget( + semantics=bread_semantics, + grasp_xpos=broadcast_pose_batch(bread_grasp_pose, num_envs=n_envs), + ), + state, + ) + logger.log_info( + f"Plan left bread pick-up cost time: {time.time() - start_time:.2f} seconds" + ) + if not left_pick_result.success.all(): + logger.log_warning("Failed to plan left bread pick-up trajectory.") + return + left_pick_traj = left_pick_result.trajectory + state = left_pick_result.next_state + bread_held_state = state.held_object + if bread_held_state is None: + raise RuntimeError("PickUp did not produce a held state for the bread.") + log_action_plan(robot, "left_pick_up", left_pick_traj, full_joint_ids) + + pan_grasp_pose = build_pan_handle_grasp_pose( + pan_pose, + pan_vertices, + sim.device, + z_clearance=PAN_GRASP_Z_CLEARANCE, + ) + start_time = time.time() + right_pick_result = right_pick_action.execute( + GraspTarget( + semantics=pan_semantics, + grasp_xpos=broadcast_pose_batch(pan_grasp_pose, num_envs=n_envs), + ), + state, + ) + logger.log_info( + f"Plan right pan pick-up cost time: {time.time() - start_time:.2f} seconds" + ) + if not right_pick_result.success.all(): + logger.log_warning("Failed to plan right pan pick-up trajectory.") + return + right_pick_traj = right_pick_result.trajectory + state = right_pick_result.next_state + pan_held_state = state.held_object + if pan_held_state is None: + raise RuntimeError("PickUp did not produce a held state for the pan.") + log_action_plan(robot, "right_pick_up", right_pick_traj, full_joint_ids) + + if args.diagnose_plan: + robot.set_qpos(state.last_qpos, joint_ids=full_joint_ids) + else: + if wait_for_user: + input("Press Enter to execute both pick-up trajectories...") + + def log_pick_execution(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 + ): + log_execution_state(robot, bread, pan, step_idx, total_steps) + + replay_trajectory( + sim, + robot, + left_pick_traj, + args, + video_prefix="", + hold_steps=0, + trajectory_sim_steps=TRAJECTORY_SIM_STEPS, + joint_ids=full_joint_ids, + on_trajectory_step=log_pick_execution, + record=False, + ) + bread.clear_dynamics() + replay_trajectory( + sim, + robot, + right_pick_traj, + args, + video_prefix="", + hold_steps=0, + trajectory_sim_steps=TRAJECTORY_SIM_STEPS, + joint_ids=full_joint_ids, + on_trajectory_step=log_pick_execution, + record=False, + ) + pan.clear_dynamics() + bread_pose_batch = clone_local_pose_from_first_env(bread).to( + device=sim.device, dtype=torch.float32 + ) + pan_pose_batch = clone_local_pose_from_first_env(pan).to( + device=sim.device, dtype=torch.float32 + ) + bread.clear_dynamics() + pan.clear_dynamics() + bread_pose = bread_pose_batch[0] + pan_pose = pan_pose_batch[0] + bread_held_state = compute_actual_held_state( + robot, + bread_semantics, + bread_pose_batch, + "left_arm", + sim.device, + ) + pan_held_state = compute_actual_held_state( + robot, + pan_semantics, + pan_pose_batch, + "right_arm", + sim.device, + ) + + support_target_pose = build_support_object_target_pose(pan_pose, sim.device) + placing_target_pose = build_placing_object_target_pose( + bread_pose, + bread_vertices, + pan_vertices, + support_target_pose, + sim.device, + ) + log_scene_targets( + bread_pose, + pan_pose, + support_target_pose, + placing_target_pose, + ) + if not args.auto_play and not args.no_vis_eef_axis: + draw_coordinated_axes( + sim, + support_target_pose, + placing_target_pose, + ) + coordinated_target = CoordinatedPlacementTarget( + placing_object_target_pose=broadcast_pose_batch( + placing_target_pose, num_envs=n_envs + ), + support_object_target_pose=broadcast_pose_batch( + support_target_pose, num_envs=n_envs + ), + placing_held_object=bread_held_state, + support_held_object=pan_held_state, + placing_height_offset=BREAD_TARGET_HEIGHT_OFFSET, + support_height_offset=SUPPORT_TARGET_HEIGHT_OFFSET, + release=True, + ) + start_time = time.time() + coordinated_success, coordinated_traj, state = engine.run( + [("coordinated_placement", coordinated_target)], state + ) + logger.log_info( + "Plan coordinated placement cost time: " + f"{time.time() - start_time:.2f} seconds" + ) + if not coordinated_success.all(): + logger.log_warning("Failed to plan coordinated placement trajectory.") + return + log_action_plan( + robot, + "coordinated_placement", + coordinated_traj, + full_joint_ids, + coordinated_action._compute_segment_lengths(coordinated_action.cfg.release), + ) + + if args.diagnose_plan: + return + + if args.auto_play and not args.no_vis_eef_axis: + draw_coordinated_axes( + sim, + support_target_pose, + placing_target_pose, + num_envs=n_envs, + ) + if wait_for_user: + input("Press Enter to execute coordinated placement...") + + def log_execution(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 + ): + log_execution_state(robot, bread, pan, step_idx, total_steps) + + replay_trajectory( + sim, + robot, + coordinated_traj, + args, + video_prefix="coordinated_placement_auto_play", + hold_steps=80, + trajectory_sim_steps=TRAJECTORY_SIM_STEPS, + joint_ids=full_joint_ids, + on_trajectory_step=log_execution, + look_at=( + (-0.25, 0.0, 2.5), + (-0.05, 0.0, 0.72), + (0.0, 0.0, 1.0), + ), + ) + if wait_for_user: + input("Press Enter to exit the simulation...") + + +def main() -> None: + """Run the coordinated placement demo.""" + args = parse_arguments() + sim = create_tutorial_simulation( + args, + 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) + + +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/atomic_action/move_end_effector.py b/scripts/tutorials/atomic_action/move_end_effector.py index 84c21d097..9ef606ad3 100644 --- a/scripts/tutorials/atomic_action/move_end_effector.py +++ b/scripts/tutorials/atomic_action/move_end_effector.py @@ -20,7 +20,6 @@ import argparse import sys -import time from pathlib import Path _REPO_ROOT = Path(__file__).resolve().parents[3] @@ -30,23 +29,23 @@ 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.atomic_actions import ( AtomicActionEngine, EndEffectorPoseTarget, MoveEndEffector, MoveEndEffectorCfg, ) -from embodichain.lab.sim.cfg import LightCfg, RenderCfg -from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - create_ur5_gripper_robot_cfg, + add_ur5_gripper_robot, + broadcast_pose_batch, + broadcast_waypoint_pose_batch, + create_toppra_motion_generator, + create_tutorial_simulation, draw_axis_marker, - get_tutorial_window_size, - start_auto_play_recording, - stop_auto_play_recording, + make_top_down_eef_pose, + prepare_tutorial_scene, + replay_trajectory, ) MOVE_SAMPLE_INTERVAL = 80 @@ -54,183 +53,74 @@ 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", - help="Run the viewer demo without waiting for keyboard input.", - ) - parser.add_argument( - "--no_vis_eef_axis", - action="store_true", - help="Do not draw the current end-effector/TCP coordinate frame before planning.", - ) + parser.add_argument("--auto_play", action="store_true") + parser.add_argument("--no_vis_eef_axis", action="store_true") return parser.parse_args() -def initialize_simulation(args: argparse.Namespace) -> SimulationManager: - width, height = get_tutorial_window_size(args) - cfg = SimulationManagerCfg( - width=width, - height=height, - headless=True, - device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), - physics_dt=1.0 / 100.0, - arena_space=2.5, - ) - sim = SimulationManager(cfg) - 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 create_robot(sim: SimulationManager, position=(0.0, 0.0, 0.0)) -> Robot: - cfg = create_ur5_gripper_robot_cfg(init_pos=position) - return sim.add_robot(cfg=cfg) - - -def make_top_down_eef_pose(device: torch.device) -> torch.Tensor: - pose = torch.eye(4, dtype=torch.float32, device=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=device, - ) - pose[:3, 3] = torch.tensor([0.30, -0.20, 0.36], dtype=torch.float32, device=device) - return pose - - -def make_side_eef_pose(device: torch.device) -> torch.Tensor: - """A second waypoint offset from the top-down pose for the multi-waypoint demo.""" - pose = torch.eye(4, dtype=torch.float32, device=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=device, - ) - pose[:3, 3] = torch.tensor([0.45, 0.10, 0.30], dtype=torch.float32, device=device) - return pose - - -def format_tensor(tensor: torch.Tensor) -> str: - rounded = (tensor.detach().cpu() * 10000.0).round() / 10000.0 - return str(rounded.tolist()) - - def main() -> None: - """Move the robot end effector through a multi-waypoint pose trajectory.""" + """Move the robot end effector through two pose waypoints.""" args = parse_arguments() - - # ------------------------------------------------------------------ # - # Step 1: Set up simulation and robot # - # ------------------------------------------------------------------ # - sim = initialize_simulation(args) - robot = create_robot(sim) - - # ------------------------------------------------------------------ # - # Step 2: Create a MotionGenerator for the robot # - # ------------------------------------------------------------------ # - motion_gen = MotionGenerator( - cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + 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), + ) ) - # ------------------------------------------------------------------ # - # Step 3: Configure the MoveEndEffector atomic action # - # ------------------------------------------------------------------ # - 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)), + ] ) - - # ------------------------------------------------------------------ # - # Step 4: Build the AtomicActionEngine # - # ------------------------------------------------------------------ # - atomic_engine = AtomicActionEngine(motion_generator=motion_gen) - atomic_engine.register(MoveEndEffector(motion_gen, cfg=move_cfg)) - - # ------------------------------------------------------------------ # - # Step 5: Define and visualize the end-effector target # - # ------------------------------------------------------------------ # - target_pose = make_top_down_eef_pose(sim.device) - side_pose = make_side_eef_pose(sim.device) - if not args.headless: - sim.open_window() - if not args.no_vis_eef_axis: - draw_axis_marker(sim, "move_end_effector_target_axis", target_pose) - draw_axis_marker(sim, "move_end_effector_side_axis", side_pose) - if not args.auto_play: - input("Inspect the robot, then press Enter to plan MoveEndEffector...") - - # ------------------------------------------------------------------ # - # Step 6: Plan the declared (name, typed_target) sequence # - # ------------------------------------------------------------------ # - # Pass a multi-waypoint trajectory (n_envs, n_waypoint, 4, 4): the - # end-effector visits `target_pose` then `side_pose` in a single plan. n_envs = robot.get_qpos().shape[0] - multi_waypoint_xpos = ( - torch.stack([target_pose, 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(target_pose[:3, 3])} -> " - f"xpos1={format_tensor(side_pose[:3, 3])}" + if not args.no_vis_eef_axis: + for name, pose in zip(("target", "side"), poses, strict=True): + draw_axis_marker( + sim, + f"move_end_effector_{name}_axis", + broadcast_pose_batch(pose, num_envs=n_envs), + ) + wait_for_user = prepare_tutorial_scene( + sim, args, "Inspect the robot, then press Enter to plan MoveEndEffector..." ) - is_success, traj, _ = atomic_engine.run( - steps=[ + + success, trajectory, _ = engine.run( + [ ( "move_end_effector", - EndEffectorPoseTarget(xpos=multi_waypoint_xpos), + EndEffectorPoseTarget(broadcast_waypoint_pose_batch(poses, n_envs)), ) ] ) - if not is_success: + if not success.all(): logger.log_warning("Failed to plan MoveEndEffector demo trajectory.") return - if not args.auto_play: + if wait_for_user: input("Press Enter to replay the MoveEndEffector demo...") - - # ------------------------------------------------------------------ # - # Step 7: Replay the planned trajectory # - # ------------------------------------------------------------------ # - recording_started = start_auto_play_recording( - sim, args, video_prefix="move_end_effector_auto_play" + replay_trajectory( + sim, + robot, + trajectory, + args, + video_prefix="move_end_effector_auto_play", + hold_steps=POST_TRAJECTORY_STEPS, ) - try: - for i in range(traj.shape[1]): - robot.set_qpos(traj[:, i, :]) - sim.update(step=4) - time.sleep(1e-2) - - final_qpos = traj[:, -1, :] - for _ in range(POST_TRAJECTORY_STEPS): - robot.set_qpos(final_qpos) - sim.update(step=2) - time.sleep(1e-2) - finally: - stop_auto_play_recording(sim, recording_started) - - if not args.auto_play: + if wait_for_user: input("Press Enter to exit the simulation...") diff --git a/scripts/tutorials/atomic_action/move_held_object.py b/scripts/tutorials/atomic_action/move_held_object.py index 0afc378aa..d0acc5afd 100644 --- a/scripts/tutorials/atomic_action/move_held_object.py +++ b/scripts/tutorials/atomic_action/move_held_object.py @@ -20,7 +20,6 @@ import argparse import sys -import time from pathlib import Path _REPO_ROOT = Path(__file__).resolve().parents[3] @@ -31,9 +30,7 @@ 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 import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.atomic_actions import ( - AntipodalAffordance, AtomicActionEngine, EndEffectorPoseTarget, GraspTarget, @@ -42,50 +39,29 @@ MoveEndEffectorCfg, MoveHeldObject, MoveHeldObjectCfg, - ObjectSemantics, PickUp, PickUpCfg, ) -from embodichain.lab.sim.cfg import ( - LightCfg, - RenderCfg, - RigidBodyAttributesCfg, - RigidObjectCfg, -) -from embodichain.lab.sim.objects import RigidObject, Robot -from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.objects import RigidObject from embodichain.lab.sim.shapes import MeshCfg -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 ( - create_ur5_gripper_robot_cfg, + add_ur5_gripper_robot, + broadcast_pose_batch, + clone_local_pose_from_first_env, + create_antipodal_semantics, + create_toppra_motion_generator, + create_tutorial_simulation, draw_axis_marker, - get_tutorial_window_size, - start_auto_play_recording, - stop_auto_play_recording, + get_hand_open_close_qpos, + make_eef_pose_at, + prepare_tutorial_scene, + replay_trajectory, ) -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_MESH_PATH = "PaperCup/paper_cup.ply" 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 @@ -94,300 +70,143 @@ def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the MoveHeldObject tutorial.""" parser = argparse.ArgumentParser( - description="Demonstrate MoveHeldObject holding a sugar box in the gripper." + 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, - help="Number of samples for antipodal grasp generation.", - ) - parser.add_argument( - "--force_reannotate", - action="store_true", - help="Force grasp region re-annotation instead of using cached data.", - ) - 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 the current end-effector/TCP coordinate frame before planning.", - ) + 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 initialize_simulation(args: argparse.Namespace) -> SimulationManager: - width, height = get_tutorial_window_size(args) - cfg = SimulationManagerCfg( - width=width, - height=height, - headless=True, - device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), - physics_dt=1.0 / 100.0, - arena_space=2.5, - ) - sim = SimulationManager(cfg) - 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), +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", + shape=MeshCfg(fpath=get_data_path(OBJECT_MESH_PATH)), + attrs=RigidBodyAttributesCfg( + mass=0.01, + 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), ) ) - return sim - - -def create_robot(sim: SimulationManager, position=(0.0, 0.0, 0.0)) -> Robot: - cfg = create_ur5_gripper_robot_cfg(init_pos=position) - return sim.add_robot(cfg=cfg) - - -def create_pick_object(sim: SimulationManager) -> RigidObject: - cfg = RigidObjectCfg( - uid=OBJECT_LABEL, - shape=MeshCfg(fpath=get_data_path(OBJECT_MESH_PATH)), - attrs=RigidBodyAttributesCfg( - mass=OBJECT_MASS, - dynamic_friction=0.97, - static_friction=0.99, - ), - max_convex_hull_num=16, - 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, - ) - obj = sim.add_rigid_object(cfg=cfg) - - sim.update( - step=10 - ) # Settle the object to ensure it is resting on the ground before planning + sim.update(step=10) + clone_local_pose_from_first_env(obj) + obj.clear_dynamics() return obj -def build_grasp_generator_cfg(args: argparse.Namespace) -> GraspGeneratorCfg: - return GraspGeneratorCfg( - viser_port=11801, - antipodal_sampler_cfg=AntipodalSamplerCfg( - n_sample=args.n_sample, - max_length=GRIPPER_MAX_OPEN_WIDTH, - min_length=0.003, - ), - is_partial_annotate=False, - is_filter_ground_collision=False, - ) - - -def build_gripper_collision_cfg() -> 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( - obj: RigidObject, args: argparse.Namespace -) -> ObjectSemantics: - return ObjectSemantics( - label=OBJECT_LABEL, - geometry={ - "mesh_vertices": obj.get_vertices(env_ids=[0], scale=True)[0], - "mesh_triangles": obj.get_triangles(env_ids=[0])[0], - }, - affordance=AntipodalAffordance( - mesh_vertices=obj.get_vertices(env_ids=[0], scale=True)[0], - mesh_triangles=obj.get_triangles(env_ids=[0])[0], - gripper_collision_cfg=build_gripper_collision_cfg(), - generator_cfg=build_grasp_generator_cfg(args), - force_reannotate=args.force_reannotate, - ), - entity=obj, - ) - - -def get_hand_open_close_qpos( - robot: Robot, device: torch.device -) -> tuple[torch.Tensor, torch.Tensor]: - hand_limits = robot.get_qpos_limits(name="hand")[0].to( - device=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 make_pre_pick_eef_pose(robot: Robot, position: torch.Tensor) -> torch.Tensor: - pose = robot.compute_fk( - qpos=robot.get_qpos(name="arm"), - name="arm", - to_matrix=True, - )[0].clone() - pose[:3, 3] = position - return pose - - 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], - ], + [[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], dtype=torch.float32, device=device) + pose[:3, 3] = torch.tensor([-0.3, -0.3, 0.5], device=device) return pose -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: - """Pick up an object and move the held object to an object-frame target.""" + """Plan MoveEndEffector -> PickUp -> MoveHeldObject.""" args = parse_arguments() - - # ------------------------------------------------------------------ # - # Step 1: Set up simulation, robot, and object # - # ------------------------------------------------------------------ # - sim = initialize_simulation(args) - robot = create_robot(sim) + 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) - if not args.headless: - sim.open_window() - - # ------------------------------------------------------------------ # - # Step 2: Create a MotionGenerator for the robot # - # ------------------------------------------------------------------ # - motion_gen = MotionGenerator( - cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) - ) - - # ------------------------------------------------------------------ # - # Step 3: Configure the three atomic actions # - # ------------------------------------------------------------------ # - hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) - move_cfg = MoveEndEffectorCfg( - control_part="arm", - sample_interval=MOVE_SAMPLE_INTERVAL, + engine = AtomicActionEngine(motion_generator=motion_gen) + engine.register( + MoveEndEffector( + motion_gen, MoveEndEffectorCfg(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=sim.device - ), - pre_grasp_distance=0.15, - lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, + 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, + ), + ) ) - move_held_object_cfg = MoveHeldObjectCfg( - control_part="arm", - hand_control_part="hand", - hand_close_qpos=hand_close, - sample_interval=MOVE_HELD_OBJECT_SAMPLE_INTERVAL, + engine.register( + MoveHeldObject( + motion_gen, + MoveHeldObjectCfg( + hand_close_qpos=hand_close, + sample_interval=MOVE_HELD_OBJECT_SAMPLE_INTERVAL, + ), + ) ) - # ------------------------------------------------------------------ # - # Step 4: Build the AtomicActionEngine # - # ------------------------------------------------------------------ # - atomic_engine = AtomicActionEngine(motion_generator=motion_gen) - atomic_engine.register(MoveEndEffector(motion_gen, cfg=move_cfg)) - atomic_engine.register(PickUp(motion_gen, cfg=pickup_cfg)) - atomic_engine.register(MoveHeldObject(motion_gen, cfg=move_held_object_cfg)) - - # ------------------------------------------------------------------ # - # Step 5: Describe the object and define the motion targets # - # ------------------------------------------------------------------ # - semantics = create_object_semantics(obj, args) - obj_pose = obj.get_local_pose(to_matrix=True) - move_position = obj_pose[0, :3, 3].clone() - move_position[2] = 0.36 - move_target = make_pre_pick_eef_pose(robot, move_position) - object_target_pose = make_object_target_pose(sim.device) - move_held_object_target = HeldObjectPoseTarget( - object_target_pose=object_target_pose + 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_pose) - if not args.auto_play: - input("Inspect the sugar box, then press Enter to plan...") + 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..." + ) - # ------------------------------------------------------------------ # - # Step 6: Plan the declared (name, typed_target) sequence # - # ------------------------------------------------------------------ # - logger.log_info("Planning move_end_effector -> pick_up -> move_held_object") - start_time = time.time() - is_success, traj, _ = atomic_engine.run( - steps=[ - ("move_end_effector", EndEffectorPoseTarget(xpos=move_target)), - ("pick_up", GraspTarget(semantics=semantics)), - ("move_held_object", move_held_object_target), + success, trajectory, _ = engine.run( + [ + ("move_end_effector", EndEffectorPoseTarget(move_target)), + ("pick_up", GraspTarget(semantics)), + ("move_held_object", HeldObjectPoseTarget(object_target)), ] ) - cost_time = time.time() - start_time - logger.log_info(f"Plan trajectory cost time: {cost_time:.2f} seconds") - if not is_success: - logger.log_warning("Failed to plan move_held_object demo trajectory.") + if not success.all(): + logger.log_warning("Failed to plan MoveHeldObject demo trajectory.") return - if not args.auto_play: - input("Press Enter to replay the move_held_object demo...") - - # ------------------------------------------------------------------ # - # Step 7: Replay the planned trajectory # - # ------------------------------------------------------------------ # - recording_started = start_auto_play_recording( - sim, args, video_prefix="move_held_object_auto_play" - ) - try: - post_grasp_clear_step = compute_pick_close_end_step() - should_clear_object_dynamics = True - for i in range(traj.shape[1]): - robot.set_qpos(traj[:, i, :]) - sim.update(step=4) - if should_clear_object_dynamics and i + 1 >= post_grasp_clear_step: - 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("MoveHeldObject keeps the sugar box suspended in the gripper.") - - final_qpos = traj[:, -1, :] - for i in range(POST_TRAJECTORY_STEPS): - robot.set_qpos(final_qpos) - sim.update(step=2) - time.sleep(1e-2) - finally: - stop_auto_play_recording(sim, recording_started) - - if not args.auto_play: + 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 + ) + 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, + ) + if wait_for_user: input("Press Enter to exit the simulation...") diff --git a/scripts/tutorials/atomic_action/move_joints.py b/scripts/tutorials/atomic_action/move_joints.py index 58c40fafa..e4dd06fd4 100644 --- a/scripts/tutorials/atomic_action/move_joints.py +++ b/scripts/tutorials/atomic_action/move_joints.py @@ -20,7 +20,6 @@ import argparse import sys -import time from pathlib import Path _REPO_ROOT = Path(__file__).resolve().parents[3] @@ -30,7 +29,6 @@ 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.atomic_actions import ( AtomicActionEngine, JointPositionTarget, @@ -38,16 +36,14 @@ MoveJointsCfg, NamedJointPositionTarget, ) -from embodichain.lab.sim.cfg import LightCfg, RenderCfg -from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - create_ur5_gripper_robot_cfg, + add_ur5_gripper_robot, + create_toppra_motion_generator, + create_tutorial_simulation, draw_axis_marker, - get_tutorial_window_size, - start_auto_play_recording, - stop_auto_play_recording, + prepare_tutorial_scene, + replay_trajectory, ) MOVE_JOINTS_SAMPLE_INTERVAL = 80 @@ -55,157 +51,76 @@ 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", - help="Run the viewer demo without waiting for keyboard input.", - ) - parser.add_argument( - "--no_vis_eef_axis", - action="store_true", - help="Do not draw the current end-effector/TCP coordinate frame before planning.", - ) + parser.add_argument("--auto_play", action="store_true") + parser.add_argument("--no_vis_eef_axis", action="store_true") return parser.parse_args() -def initialize_simulation(args: argparse.Namespace) -> SimulationManager: - width, height = get_tutorial_window_size(args) - cfg = SimulationManagerCfg( - width=width, - height=height, - headless=True, - device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), - physics_dt=1.0 / 100.0, - arena_space=2.5, - ) - sim = SimulationManager(cfg) - 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 create_robot(sim: SimulationManager, position=(0.0, 0.0, 0.0)) -> Robot: - cfg = create_ur5_gripper_robot_cfg(init_pos=position) - return sim.add_robot(cfg=cfg) - - -def make_arm_qpos(values: list[float], device: torch.device) -> torch.Tensor: - return torch.tensor(values, dtype=torch.float32, device=device) - - -def draw_start_eef_axis(sim: SimulationManager, robot: Robot) -> None: - eef_pose = robot.compute_fk( - qpos=robot.get_qpos(name="arm"), - name="arm", - to_matrix=True, - ) - draw_axis_marker(sim, "move_joints_start_eef_axis", eef_pose) - - def main() -> None: - """Move the robot arm through named and explicit joint targets.""" + """Move the robot arm through a named target and two explicit waypoints.""" args = parse_arguments() - - # ------------------------------------------------------------------ # - # Step 1: Set up simulation and robot # - # ------------------------------------------------------------------ # - sim = initialize_simulation(args) - robot = create_robot(sim) - - # ------------------------------------------------------------------ # - # Step 2: Create a MotionGenerator for the robot # - # ------------------------------------------------------------------ # - motion_gen = MotionGenerator( - cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + 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], + ) ) - - # ------------------------------------------------------------------ # - # Step 3: Configure the MoveJoints atomic action # - # ------------------------------------------------------------------ # - ready_qpos = make_arm_qpos([0.35, -1.20, 1.30, -1.65, -1.57, 0.20], sim.device) - mid_qpos = make_arm_qpos([0.15, -1.40, 1.45, -1.60, -1.57, 0.10], sim.device) - home_qpos = make_arm_qpos([0.0, -1.57, 1.57, -1.57, -1.57, 0.0], sim.device) - 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}, + ), + ) ) - # ------------------------------------------------------------------ # - # Step 4: Build the AtomicActionEngine # - # ------------------------------------------------------------------ # - atomic_engine = AtomicActionEngine(motion_generator=motion_gen) - atomic_engine.register(MoveJoints(motion_gen, cfg=move_joints_cfg)) - - # ------------------------------------------------------------------ # - # Step 5: Open the viewer and show the starting end-effector frame # - # ------------------------------------------------------------------ # - if not args.headless: - sim.open_window() if not args.no_vis_eef_axis: - draw_start_eef_axis(sim, robot) - if not args.auto_play: - input("Inspect the robot, then press Enter to plan MoveJoints...") - - # ------------------------------------------------------------------ # - # Step 6: Plan the declared (name, typed_target) sequence # - # ------------------------------------------------------------------ # - # The second MoveJoints step passes a multi-waypoint trajectory - # (n_envs, n_waypoint, control_dof): the arm visits `mid_qpos` then - # `home_qpos` in a single planned trajectory. - n_envs = robot.get_qpos().shape[0] - multi_waypoint_qpos = ( - torch.stack([mid_qpos, home_qpos], dim=0).unsqueeze(0).repeat(n_envs, 1, 1) + draw_axis_marker( + sim, + "move_joints_start_eef_axis", + robot.compute_fk(robot.get_qpos(name="arm"), name="arm", to_matrix=True), + ) + 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)" + + waypoints = ( + torch.stack([mid, home]).unsqueeze(0).repeat(robot.get_qpos().shape[0], 1, 1) ) - is_success, traj, _ = atomic_engine.run( - steps=[ - ("move_joints", NamedJointPositionTarget(name="ready")), - ("move_joints", JointPositionTarget(qpos=multi_waypoint_qpos)), + success, trajectory, _ = engine.run( + [ + ("move_joints", NamedJointPositionTarget("ready")), + ("move_joints", JointPositionTarget(waypoints)), ] ) - if not is_success: + if not success.all(): logger.log_warning("Failed to plan MoveJoints demo trajectory.") return - if not args.auto_play: + if wait_for_user: input("Press Enter to replay the MoveJoints demo...") - - # ------------------------------------------------------------------ # - # Step 7: Replay the planned trajectory # - # ------------------------------------------------------------------ # - recording_started = start_auto_play_recording( - sim, args, video_prefix="move_joints_auto_play" + replay_trajectory( + sim, + robot, + trajectory, + args, + video_prefix="move_joints_auto_play", + hold_steps=POST_TRAJECTORY_STEPS, ) - try: - for i in range(traj.shape[1]): - robot.set_qpos(traj[:, i, :]) - sim.update(step=4) - time.sleep(1e-2) - - final_qpos = traj[:, -1, :] - for _ in range(POST_TRAJECTORY_STEPS): - robot.set_qpos(final_qpos) - sim.update(step=2) - time.sleep(1e-2) - finally: - stop_auto_play_recording(sim, recording_started) - - if not args.auto_play: + if wait_for_user: input("Press Enter to exit the simulation...") diff --git a/scripts/tutorials/atomic_action/pickup.py b/scripts/tutorials/atomic_action/pickup.py index eb0300b24..2aa7b3082 100644 --- a/scripts/tutorials/atomic_action/pickup.py +++ b/scripts/tutorials/atomic_action/pickup.py @@ -14,13 +14,12 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Demonstrate PickUp on an upright object with configurable approach.""" +"""Demonstrate PickUp on a cube with a configurable approach direction.""" from __future__ import annotations import argparse import sys -import time from pathlib import Path _REPO_ROOT = Path(__file__).resolve().parents[3] @@ -29,65 +28,35 @@ 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 import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.atomic_actions import ( - AntipodalAffordance, AtomicActionEngine, GraspTarget, - ObjectSemantics, PickUp, PickUpCfg, ) -from embodichain.lab.sim.cfg import ( - LightCfg, - RenderCfg, - RigidBodyAttributesCfg, - RigidObjectCfg, -) -from embodichain.lab.sim.objects import RigidObject, Robot -from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg -from embodichain.lab.sim.shapes import MeshCfg -from embodichain.toolkits.graspkit.pg_grasp.antipodal_generator import ( - AntipodalSamplerCfg, - GraspGeneratorCfg, -) -from embodichain.toolkits.graspkit.pg_grasp.gripper_collision_checker import ( - GripperCollisionCfg, -) +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.objects import RigidObject +from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - create_ur5_gripper_robot_cfg, + add_ur5_gripper_robot, + clone_local_pose_from_first_env, + create_antipodal_semantics, + create_toppra_motion_generator, + create_tutorial_simulation, draw_axis_marker, - get_tutorial_window_size, - start_auto_play_recording, - stop_auto_play_recording, + get_hand_open_close_qpos, + initialize_pre_pick_robot_pose, + prepare_tutorial_scene, + replay_trajectory, ) -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_SIZE = (0.05, 0.05, 0.05) 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), @@ -96,343 +65,125 @@ def parse_arguments() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Demonstrate PickUp on an upright object." - ) + """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( - "--object", - choices=sorted(OBJECT_PRESETS.keys()), - default="sugar_box", - help="Object preset to pick.", - ) - parser.add_argument( - "--n_sample", - type=int, - default=10000, - help="Number of samples for antipodal grasp generation.", - ) - parser.add_argument( - "--force_reannotate", - action="store_true", - help="Force grasp region re-annotation instead of using cached data.", - ) - parser.add_argument( - "--auto_play", - action="store_true", - help="Run the viewer demo without waiting for keyboard input.", - ) - parser.add_argument( - "--approach", - choices=["top", "side", "side_y", "custom"], - default="top", - help="Pick approach direction preset.", - ) - parser.add_argument( - "--custom_approach_direction", - type=float, - nargs=3, - default=None, - metavar=("X", "Y", "Z"), - help="World-frame approach direction used when --approach custom.", - ) - parser.add_argument( - "--no_vis_eef_axis", - action="store_true", - help="Do not draw the current end-effector/TCP coordinate frame before planning.", + "--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() -def initialize_simulation(args: argparse.Namespace) -> SimulationManager: - width, height = get_tutorial_window_size(args) - cfg = SimulationManagerCfg( - width=width, - height=height, - headless=True, - device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), - physics_dt=1.0 / 100.0, - arena_space=2.5, - ) - sim = SimulationManager(cfg) - 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), +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)), + attrs=RigidBodyAttributesCfg( + mass=0.05, + dynamic_friction=0.97, + static_friction=0.99, + ), + max_convex_hull_num=16, + init_pos=[*OBJECT_XY, OBJECT_SIZE[2]], ) ) - return sim - - -def create_robot(sim: SimulationManager, position=(0.0, 0.0, 0.0)) -> Robot: - cfg = create_ur5_gripper_robot_cfg(init_pos=position) - return sim.add_robot(cfg=cfg) - - -def create_pick_object(sim: SimulationManager, 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=preset["mass"], - dynamic_friction=0.97, - static_friction=0.99, - ), - max_convex_hull_num=16, - 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"], - ) - obj = sim.add_rigid_object(cfg=cfg) - - # Settle the object to ensure it is resting on the ground before planning sim.update(step=10) + clone_local_pose_from_first_env(obj) + obj.clear_dynamics() return obj -def build_grasp_generator_cfg(args: argparse.Namespace) -> GraspGeneratorCfg: - return GraspGeneratorCfg( - viser_port=11801, - antipodal_sampler_cfg=AntipodalSamplerCfg( - n_sample=args.n_sample, - max_length=GRIPPER_MAX_OPEN_WIDTH, - min_length=0.003, - ), - is_partial_annotate=False, - is_filter_ground_collision=False, - ) - - -def build_gripper_collision_cfg() -> 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( - obj: RigidObject, args: argparse.Namespace -) -> ObjectSemantics: - label = OBJECT_PRESETS[args.object]["label"] - return ObjectSemantics( - label=label, - geometry={ - "mesh_vertices": obj.get_vertices(env_ids=[0], scale=True)[0], - "mesh_triangles": obj.get_triangles(env_ids=[0])[0], - }, - affordance=AntipodalAffordance( - mesh_vertices=obj.get_vertices(env_ids=[0], scale=True)[0], - mesh_triangles=obj.get_triangles(env_ids=[0])[0], - gripper_collision_cfg=build_gripper_collision_cfg(), - generator_cfg=build_grasp_generator_cfg(args), - force_reannotate=args.force_reannotate, - ), - entity=obj, - ) - - -def get_hand_open_close_qpos( - robot: Robot, device: torch.device -) -> tuple[torch.Tensor, torch.Tensor]: - hand_limits = robot.get_qpos_limits(name="hand")[0].to( - device=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( args: argparse.Namespace, device: torch.device ) -> torch.Tensor: - if args.approach == "custom": - if args.custom_approach_direction is None: - raise ValueError( - "--custom_approach_direction is required when --approach custom." - ) - direction = args.custom_approach_direction - else: - direction = APPROACH_DIRECTIONS[args.approach] - - approach_direction = torch.tensor(direction, dtype=torch.float32, device=device) - norm = torch.linalg.norm(approach_direction) - if norm < 1e-6: + """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." + ) + 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 approach_direction / norm - - -def make_pre_pick_eef_pose(robot: Robot, position: torch.Tensor) -> torch.Tensor: - pose = robot.compute_fk( - qpos=robot.get_qpos(name="arm"), - name="arm", - to_matrix=True, - ).clone() - pose[:, :3, 3] = position - return pose - - -def initialize_pre_pick_robot_pose( - robot: Robot, - obj: RigidObject, - hand_open: torch.Tensor, -) -> None: - obj_pose = obj.get_local_pose(to_matrix=True) - move_position = obj_pose[:, :3, 3].clone() - move_position[:, 2] = 0.36 - pre_pick_pose = make_pre_pick_eef_pose(robot, move_position) - ik_success, arm_qpos = robot.compute_ik( - pose=pre_pick_pose, - joint_seed=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 = robot.get_qpos().shape[0] - hand_qpos = hand_open.unsqueeze(0).repeat(n_envs, 1) - for target in (False, True): - robot.set_qpos(arm_qpos, name="arm", target=target) - robot.set_qpos(hand_qpos, name="hand", target=target) - robot.clear_dynamics() - - -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 format_tensor(tensor: torch.Tensor) -> str: - rounded = (tensor.detach().cpu() * 10000.0).round() / 10000.0 - return str(rounded.tolist()) - - -def draw_pick_object_axis(sim: SimulationManager, obj: RigidObject) -> None: - draw_axis_marker( - sim, - "pickup_object_axis", - obj.get_local_pose(to_matrix=True), - ) + return torch.nn.functional.normalize(approach, dim=0) def main() -> None: - """Pick up an object using an antipodal grasp affordance.""" + """Plan and replay a sampled antipodal PickUp trajectory.""" args = parse_arguments() - - # ------------------------------------------------------------------ # - # Step 1: Set up simulation, robot, and object # - # ------------------------------------------------------------------ # - sim = initialize_simulation(args) - robot = create_robot(sim) - obj = create_pick_object(sim, args.object) - - # ------------------------------------------------------------------ # - # Step 2: Create a MotionGenerator for the robot # - # ------------------------------------------------------------------ # - motion_gen = MotionGenerator( - cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) - ) - - # ------------------------------------------------------------------ # - # Step 3: Configure the PickUp atomic action # - # ------------------------------------------------------------------ # - hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) - approach_direction = resolve_approach_direction(args, sim.device) + 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) - 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, + 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, + ), + ) ) - - # ------------------------------------------------------------------ # - # Step 4: Build the AtomicActionEngine # - # ------------------------------------------------------------------ # - atomic_engine = AtomicActionEngine(motion_generator=motion_gen) - atomic_engine.register(PickUp(motion_gen, cfg=pickup_cfg)) - - # ------------------------------------------------------------------ # - # Step 5: Describe the object with ObjectSemantics # - # ------------------------------------------------------------------ # - semantics = create_object_semantics(obj, args) - - if not args.headless: - sim.open_window() - if not args.no_vis_eef_axis: - draw_pick_object_axis(sim, obj) - if not args.auto_play: - input(f"Inspect the upright {args.object}, then press Enter to plan...") - - # ------------------------------------------------------------------ # - # Step 6: Plan the declared (name, typed_target) sequence # - # ------------------------------------------------------------------ # - logger.log_info( - f"Planning pick_up for {args.object} with " - f"approach_direction={format_tensor(approach_direction)}" + semantics = create_antipodal_semantics( + obj, + label="cube", + n_sample=args.n_sample, + force_reannotate=args.force_reannotate, ) - start_time = time.time() - is_success, traj, _ = atomic_engine.run( - steps=[("pick_up", GraspTarget(semantics=semantics))] + 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..." ) - cost_time = time.time() - start_time - logger.log_info(f"Plan trajectory cost time: {cost_time:.2f} seconds") - if not is_success: - logger.log_warning("Failed to plan pickup demo trajectory.") - return - if not args.auto_play: - input("Press Enter to replay the pickup demo...") + success, trajectory, _ = engine.run([("pick_up", GraspTarget(semantics))]) + if not success.all(): + logger.log_warning("Failed to plan PickUp demo trajectory.") + return - # ------------------------------------------------------------------ # - # Step 7: Replay the planned trajectory # - # ------------------------------------------------------------------ # - recording_started = start_auto_play_recording( - sim, args, video_prefix=f"pickup_{args.object}_auto_play" + 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 ) - try: - post_grasp_clear_step = compute_pick_close_end_step() - should_clear_object_dynamics = True - for i in range(traj.shape[1]): - robot.set_qpos(traj[:, i, :]) - sim.update(step=4) - if should_clear_object_dynamics and i + 1 >= post_grasp_clear_step: - 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 {args.object} suspended in the gripper." - ) + dynamics_cleared = False - final_qpos = traj[:, -1, :] - for i in range(POST_TRAJECTORY_STEPS): - robot.set_qpos(final_qpos) - sim.update(step=2) - time.sleep(1e-2) - finally: - stop_auto_play_recording(sim, recording_started) + 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 - if not args.auto_play: + replay_trajectory( + sim, + robot, + trajectory, + args, + video_prefix="pickup_cube_auto_play", + hold_steps=POST_TRAJECTORY_STEPS, + on_trajectory_step=clear_object_dynamics, + ) + if wait_for_user: input("Press Enter to exit the simulation...") diff --git a/scripts/tutorials/atomic_action/place.py b/scripts/tutorials/atomic_action/place.py index a7b52b1b5..6a8752d5a 100644 --- a/scripts/tutorials/atomic_action/place.py +++ b/scripts/tutorials/atomic_action/place.py @@ -20,7 +20,6 @@ import argparse import sys -import time from pathlib import Path _REPO_ROOT = Path(__file__).resolve().parents[3] @@ -29,60 +28,37 @@ 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 import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.atomic_actions import ( - AntipodalAffordance, AtomicActionEngine, EndEffectorPoseTarget, GraspTarget, - ObjectSemantics, PickUp, PickUpCfg, Place, PlaceCfg, ) -from embodichain.lab.sim.cfg import ( - LightCfg, - RenderCfg, - RigidBodyAttributesCfg, - RigidObjectCfg, -) -from embodichain.lab.sim.objects import RigidObject, Robot -from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg -from embodichain.lab.sim.shapes import MeshCfg -from embodichain.toolkits.graspkit.pg_grasp.antipodal_generator import ( - AntipodalSamplerCfg, - GraspGeneratorCfg, -) -from embodichain.toolkits.graspkit.pg_grasp.gripper_collision_checker import ( - GripperCollisionCfg, -) +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.objects import RigidObject +from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - create_ur5_gripper_robot_cfg, + 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, draw_axis_marker, - get_tutorial_window_size, - start_auto_play_recording, - stop_auto_play_recording, + get_hand_open_close_qpos, + initialize_pre_pick_robot_pose, + prepare_tutorial_scene, + replay_trajectory, ) -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_SIZE = (0.05, 0.05, 0.05) 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 @@ -91,333 +67,153 @@ def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the Place tutorial.""" parser = argparse.ArgumentParser( - description="Demonstrate Place by picking an object and releasing it at a target pose." + 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, - help="Number of samples for antipodal grasp generation.", - ) - parser.add_argument( - "--force_reannotate", - action="store_true", - help="Force grasp region re-annotation instead of using cached data.", - ) - 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 the current end-effector/TCP coordinate frame before planning.", - ) + 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 initialize_simulation(args: argparse.Namespace) -> SimulationManager: - width, height = get_tutorial_window_size(args) - cfg = SimulationManagerCfg( - width=width, - height=height, - headless=True, - device=args.device, - render_cfg=RenderCfg(renderer=args.renderer), - physics_dt=1.0 / 100.0, - arena_space=2.5, - ) - sim = SimulationManager(cfg) - 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), +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)), + attrs=RigidBodyAttributesCfg( + mass=0.05, + 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]], ) ) - return sim - - -def create_robot(sim: SimulationManager, position=(0.0, 0.0, 0.0)) -> Robot: - cfg = create_ur5_gripper_robot_cfg(init_pos=position) - return sim.add_robot(cfg=cfg) - - -def create_pick_object(sim: SimulationManager) -> RigidObject: - cfg = RigidObjectCfg( - uid=OBJECT_LABEL, - shape=MeshCfg(fpath=get_data_path(OBJECT_MESH_PATH)), - attrs=RigidBodyAttributesCfg( - mass=OBJECT_MASS, - dynamic_friction=0.97, - static_friction=0.99, - enable_ccd=True, - ), - max_convex_hull_num=16, - 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, - ) - obj = sim.add_rigid_object(cfg=cfg) - - # Set the object to a stable pose on the ground by simulating a few steps. sim.update(step=10) + clone_local_pose_from_first_env(obj) + obj.clear_dynamics() return obj -def get_hand_open_close_qpos( - robot: Robot, device: torch.device -) -> tuple[torch.Tensor, torch.Tensor]: - hand_limits = robot.get_qpos_limits(name="hand")[0].to( - device=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 build_grasp_generator_cfg(args: argparse.Namespace) -> GraspGeneratorCfg: - return GraspGeneratorCfg( - viser_port=11801, - antipodal_sampler_cfg=AntipodalSamplerCfg( - n_sample=args.n_sample, - max_length=GRIPPER_MAX_OPEN_WIDTH, - min_length=0.003, - ), - is_partial_annotate=False, - is_filter_ground_collision=False, - ) - - -def build_gripper_collision_cfg() -> 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( - obj: RigidObject, args: argparse.Namespace -) -> ObjectSemantics: - return ObjectSemantics( - label=OBJECT_LABEL, - geometry={ - "mesh_vertices": obj.get_vertices(env_ids=[0], scale=True)[0], - "mesh_triangles": obj.get_triangles(env_ids=[0])[0], - }, - affordance=AntipodalAffordance( - mesh_vertices=obj.get_vertices(env_ids=[0], scale=True)[0], - mesh_triangles=obj.get_triangles(env_ids=[0])[0], - gripper_collision_cfg=build_gripper_collision_cfg(), - generator_cfg=build_grasp_generator_cfg(args), - force_reannotate=args.force_reannotate, - ), - entity=obj, - ) - - -def make_pre_pick_eef_pose(robot: Robot, position: torch.Tensor) -> torch.Tensor: - pose = robot.compute_fk( - qpos=robot.get_qpos(name="arm"), - name="arm", - to_matrix=True, - ).clone() - pose[:, :3, 3] = position - return pose - - -def initialize_pre_pick_robot_pose( - robot: Robot, - obj: RigidObject, - hand_open: torch.Tensor, -) -> None: - obj_pose = obj.get_local_pose(to_matrix=True) - move_position = obj_pose[:, :3, 3].clone() - move_position[:, 2] = 0.36 - pre_pick_pose = make_pre_pick_eef_pose(robot, move_position) - ik_success, arm_qpos = robot.compute_ik( - pose=pre_pick_pose, - joint_seed=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 = robot.get_qpos().shape[0] - hand_qpos = hand_open.unsqueeze(0).repeat(n_envs, 1) - for target in (False, True): - robot.set_qpos(arm_qpos, name="arm", target=target) - robot.set_qpos(hand_qpos, name="hand", target=target) - robot.clear_dynamics() - - def make_place_eef_poses(device: torch.device) -> 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. - """ + """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.0000, -0.9992], + [-0.0401, 0.0, -0.9992], ], dtype=torch.float32, device=device, ) - hover_pose = torch.eye(4, dtype=torch.float32, device=device) - hover_pose[:3, :3] = rotation - hover_pose[:3, 3] = torch.tensor( - [-0.40, 0.48, 0.20], dtype=torch.float32, device=device - ) - place_pose = torch.eye(4, dtype=torch.float32, device=device) - place_pose[:3, :3] = rotation - place_pose[:3, 3] = torch.tensor( - [-0.40, 0.48, 0.10], dtype=torch.float32, device=device - ) - return torch.stack([hover_pose, place_pose], dim=0) - - -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 + 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) def main() -> None: - """Pick up an object and place it at a target end-effector pose.""" + """Plan and replay PickUp followed by a multi-waypoint Place.""" args = parse_arguments() - - # ------------------------------------------------------------------ # - # Step 1: Set up simulation, robot, and object # - # ------------------------------------------------------------------ # - sim = initialize_simulation(args) - robot = create_robot(sim) + sim = create_tutorial_simulation(args) + robot = add_ur5_gripper_robot(sim) obj = create_pick_object(sim) - - if not args.headless: - sim.open_window() - - # ------------------------------------------------------------------ # - # Step 2: Create a MotionGenerator for the robot # - # ------------------------------------------------------------------ # - motion_gen = MotionGenerator( - cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) - ) - - # ------------------------------------------------------------------ # - # Step 3: Configure the PickUp and Place atomic actions # - # ------------------------------------------------------------------ # - hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) + 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) - 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=sim.device - ), - pre_grasp_distance=0.15, - lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, + + 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, + ), + ) ) - 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, + 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, + ), + ) ) - # ------------------------------------------------------------------ # - # Step 4: Build the AtomicActionEngine # - # ------------------------------------------------------------------ # - atomic_engine = AtomicActionEngine(motion_generator=motion_gen) - atomic_engine.register(PickUp(motion_gen, cfg=pickup_cfg)) - atomic_engine.register(Place(motion_gen, cfg=place_cfg)) - - # ------------------------------------------------------------------ # - # Step 5: Describe the object and define the place target # - # ------------------------------------------------------------------ # - semantics = create_object_semantics(obj, args) - place_eef_poses = make_place_eef_poses(sim.device) - + 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", place_eef_poses[-1]) - if not args.auto_play: - input("Inspect the object, then press Enter to plan PickUp -> Place...") + 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..." + ) - # ------------------------------------------------------------------ # - # Step 6: Plan the declared (name, typed_target) sequence # - # ------------------------------------------------------------------ # - # Pass a multi-waypoint trajectory (n_envs, n_waypoint, 4, 4): Place - # approaches from above the first waypoint, descends through each - # waypoint in order, opens the gripper at the last, and retracts. - n_envs = robot.get_qpos().shape[0] - multi_waypoint_xpos = 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") - is_success, traj, _ = atomic_engine.run( - steps=[ - ("pick_up", GraspTarget(semantics=semantics)), - ("place", place_target), + success, trajectory, _ = engine.run( + [ + ("pick_up", GraspTarget(semantics)), + ( + "place", + EndEffectorPoseTarget( + broadcast_waypoint_pose_batch( + place_poses, robot.get_qpos().shape[0] + ) + ), + ), ] ) - if not is_success: + if not success.all(): logger.log_warning("Failed to plan Place demo trajectory.") return - if not args.auto_play: + if wait_for_user: input("Press Enter to replay the Place demo...") - - # ------------------------------------------------------------------ # - # Step 7: Replay the planned trajectory # - # ------------------------------------------------------------------ # - recording_started = start_auto_play_recording( - sim, args, video_prefix="place_auto_play" - ) - try: - post_grasp_clear_step = compute_pick_close_end_step() - should_clear_object_dynamics = True - for i in range(traj.shape[1]): - robot.set_qpos(traj[:, i, :]) - sim.update(step=4) - if should_clear_object_dynamics and i + 1 >= post_grasp_clear_step: - 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 i in range(POST_TRAJECTORY_STEPS): - robot.set_qpos(final_qpos) - sim.update(step=2) - time.sleep(1e-2) - finally: - stop_auto_play_recording(sim, recording_started) - - if not args.auto_play: + clear_after_step = ( + round((PICK_SAMPLE_INTERVAL - HAND_INTERP_STEPS) * 0.6) + HAND_INTERP_STEPS + ) + 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, + ) + if wait_for_user: input("Press Enter to exit the simulation...") diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py new file mode 100644 index 000000000..10deaad3a --- /dev/null +++ b/scripts/tutorials/atomic_action/press.py @@ -0,0 +1,241 @@ +# ---------------------------------------------------------------------------- +# 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 Press on the center of a regular wooden block.""" + +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, + Press, + PressCfg, +) +from embodichain.lab.sim.cfg import ( + RigidBodyAttributesCfg, + RigidObjectCfg, +) +from embodichain.lab.sim.material import VisualMaterialCfg +from embodichain.lab.sim.objects import RigidObject +from embodichain.lab.sim.shapes import CubeCfg +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, +) + +MOVE_SAMPLE_INTERVAL = 60 +PRESS_SAMPLE_INTERVAL = 90 +HAND_INTERP_STEPS = 12 +POST_TRAJECTORY_STEPS = 180 +BLOCK_SIZE = (0.12, 0.12, 0.06) +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() + + +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( + uid="wooden_block", + shape=CubeCfg( + size=list(BLOCK_SIZE), + visual_material=VisualMaterialCfg( + uid="wooden_block_mat", + base_color=[0.58, 0.32, 0.14, 1.0], + roughness=0.85, + ), + ), + body_type="static", + attrs=RigidBodyAttributesCfg(dynamic_friction=0.8, static_friction=0.9), + init_pos=center, + ) + ) + + +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 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) + ) + ) + engine.register( + Press( + motion_gen, + PressCfg( + hand_close_qpos=hand_close, + sample_interval=PRESS_SAMPLE_INTERVAL, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ) + ) + + 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..." + ) + + success, trajectory, _ = engine.run( + [ + ("move_end_effector", EndEffectorPoseTarget(move_target)), + ("press", EndEffectorPoseTarget(press_target)), + ] + ) + 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) + ) + 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)}" + ) + 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, + ) + if wait_for_user: + input("Press Enter to exit the simulation...") + + +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/atomic_action/tutorial_utils.py b/scripts/tutorials/atomic_action/tutorial_utils.py index 8b8ea6ba7..ca80e4166 100644 --- a/scripts/tutorials/atomic_action/tutorial_utils.py +++ b/scripts/tutorials/atomic_action/tutorial_utils.py @@ -19,14 +19,35 @@ from __future__ import annotations import argparse -from collections.abc import Sequence +import time +from collections.abc import Callable, Sequence import torch from embodichain.data import get_data_path -from embodichain.lab.sim import SimulationManager -from embodichain.lab.sim.cfg import MarkerCfg, RobotCfg +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.atomic_actions import ( + AntipodalAffordance, + ObjectSemantics, +) +from embodichain.lab.sim.cfg import ( + LightCfg, + MarkerCfg, + RenderCfg, + RobotCfg, + physics_cfg_for_backend, +) +from embodichain.lab.sim.objects import RigidObject, Robot +from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator, ToppraPlannerCfg from embodichain.lab.sim.robots import URRobotCfg +from embodichain.lab.sim.solvers import URSolverCfg +from embodichain.toolkits.graspkit.pg_grasp.antipodal_generator import ( + AntipodalSamplerCfg, + GraspGeneratorCfg, +) +from embodichain.toolkits.graspkit.pg_grasp.gripper_collision_checker import ( + GripperCollisionCfg, +) RECORD_WIDTH = 640 RECORD_HEIGHT = 480 @@ -45,6 +66,354 @@ GRIPPER_URDF_PATH = "DH_PGI_140_80/DH_PGI_140_80.urdf" GRIPPER_HAND_JOINT_PATTERN = "gripper_finger1_joint_1" GRIPPER_TCP_Z = 0.15 +GRIPPER_MAX_OPEN_WIDTH = 0.080 +GRIPPER_FINGER_LENGTH = 0.088 +GRIPPER_ROOT_Z_WIDTH = 0.096 +GRIPPER_Y_THICKNESS = 0.040 +DEFAULT_GRIPPER_CLOSE_QPOS = 0.024 +DEFAULT_TUTORIAL_LIGHT_POS = (1.0, 0.0, 3.0) +TOP_DOWN_EEF_ROTATION = ( + (-0.0539, -0.9985, -0.0022), + (-0.9977, 0.0540, -0.0401), + (0.0401, 0.0000, -0.9992), +) + + +def make_ur5_solver_cfg(tcp_z: float) -> URSolverCfg: + """Create the UR5 arm solver cfg used by atomic-action tutorials.""" + cfg = URSolverCfg( + ur_type="ur5", + end_link_name="ee_link", + root_link_name="base_link", + tcp=[ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, tcp_z], + [0.0, 0.0, 0.0, 1.0], + ], + ) + cfg.urdf_path = None + return cfg + + +def create_tutorial_simulation( + args: argparse.Namespace, + *, + arena_space: float = 2.5, + light_pos: Sequence[float] = DEFAULT_TUTORIAL_LIGHT_POS, +) -> SimulationManager: + """Create the shared simulation setup used by atomic-action tutorials. + + Args: + args: Parsed launcher arguments containing environment count, device, + and renderer selections. + arena_space: Spacing between parallel simulation arenas in meters. + light_pos: Position of the scene's key light. + + Returns: + A simulation manager with the tutorial key light configured. + """ + width, height = get_tutorial_window_size(args) + sim = SimulationManager( + SimulationManagerCfg( + width=width, + height=height, + headless=True, + num_envs=args.num_envs, + device=args.device, + physics_cfg=physics_cfg_for_backend(args.physics), + render_cfg=RenderCfg(renderer=args.renderer), + physics_dt=1.0 / 100.0, + arena_space=arena_space, + ) + ) + sim.add_light( + cfg=LightCfg( + uid="main_light", + color=(0.6, 0.6, 0.6), + intensity=30.0, + init_pos=list(light_pos), + ) + ) + return sim + + +def add_ur5_gripper_robot( + sim: SimulationManager, + init_pos: Sequence[float] = (0.0, 0.0, 0.0), +) -> Robot: + """Add the standard UR5 plus PGI gripper tutorial robot. + + Args: + sim: Simulation manager that owns the robot. + init_pos: Root position of the robot in its arena. + + Returns: + The added robot instance. + """ + return sim.add_robot(cfg=create_ur5_gripper_robot_cfg(init_pos=init_pos)) + + +def create_toppra_motion_generator(robot: Robot) -> MotionGenerator: + """Create the standard TOPPRA motion generator for a tutorial robot. + + Args: + robot: Robot whose trajectories will be planned. + + Returns: + The configured motion generator. + """ + return MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + + +def get_hand_open_close_qpos( + robot: Robot, + *, + hand_control_part: str = "hand", + close_qpos: float = DEFAULT_GRIPPER_CLOSE_QPOS, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return the open limit and a safe closed position for a gripper. + + Args: + robot: Robot containing the gripper control part. + hand_control_part: Name of the gripper control part. + close_qpos: Desired per-joint closed position, clamped to joint limits. + + Returns: + Open and closed joint-position tensors on the robot device. + """ + hand_limits = robot.get_qpos_limits(name=hand_control_part)[0].to( + device=robot.device, dtype=torch.float32 + ) + return hand_limits[:, 0], torch.minimum( + hand_limits[:, 1], torch.full_like(hand_limits[:, 1], close_qpos) + ) + + +def create_antipodal_semantics( + obj: RigidObject, + *, + label: str, + n_sample: int, + force_reannotate: bool, +) -> ObjectSemantics: + """Describe a rigid object using the standard PGI antipodal-grasp setup. + + Args: + obj: Rigid object that will be grasped. + label: Human-readable object category. + n_sample: Number of grasp-pair samples to generate. + force_reannotate: Whether to ignore cached grasp annotations. + + Returns: + Object semantics with mesh data stored only on its affordance. + """ + vertices = obj.get_vertices(env_ids=[0], scale=True)[0] + triangles = obj.get_triangles(env_ids=[0])[0] + return ObjectSemantics( + label=label, + geometry={}, + affordance=AntipodalAffordance( + mesh_vertices=vertices, + mesh_triangles=triangles, + gripper_collision_cfg=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, + ), + generator_cfg=GraspGeneratorCfg( + viser_port=11801, + antipodal_sampler_cfg=AntipodalSamplerCfg( + n_sample=n_sample, + max_length=GRIPPER_MAX_OPEN_WIDTH, + min_length=0.003, + ), + is_partial_annotate=False, + is_filter_ground_collision=False, + ), + force_reannotate=force_reannotate, + ), + entity=obj, + ) + + +def make_top_down_eef_pose(position: torch.Tensor) -> torch.Tensor: + """Build the standard downward-facing TCP pose at ``position``. + + Args: + position: Three-dimensional position for the tool center point. + + Returns: + A homogeneous end-effector pose. + """ + pose = torch.eye(4, dtype=torch.float32, device=position.device) + pose[:3, :3] = torch.tensor( + TOP_DOWN_EEF_ROTATION, + dtype=torch.float32, + device=position.device, + ) + pose[:3, 3] = position + return pose + + +def format_tensor(tensor: torch.Tensor) -> str: + """Format a tensor compactly for tutorial log messages. + + Args: + tensor: Tensor to render for a human-readable log message. + + Returns: + A rounded Python-list representation. + """ + rounded = (tensor.detach().cpu() * 10000.0).round() / 10000.0 + return str(rounded.tolist()) + + +def make_eef_pose_at(robot: Robot, position: torch.Tensor) -> torch.Tensor: + """Reuse the current TCP orientation at a requested position. + + Args: + robot: Robot whose current arm pose supplies the orientation. + position: Position tensor with shape ``(3,)`` or ``(n_envs, 3)``. + + Returns: + A single or batched homogeneous pose matching ``position``. + """ + pose = robot.compute_fk( + qpos=robot.get_qpos(name="arm"), name="arm", to_matrix=True + ).clone() + if position.dim() == 1: + pose = pose[0] + pose[..., :3, 3] = position + return pose + + +def initialize_pre_pick_robot_pose( + robot: Robot, + obj: RigidObject, + hand_open: torch.Tensor, + *, + height: float = 0.36, +) -> None: + """Set a UR5 at a deterministic open-gripper pose above an object. + + Args: + robot: Robot to initialize. + obj: Object below the pre-pick TCP pose. + hand_open: Open gripper joint positions. + height: World-Z coordinate for the pre-pick TCP pose. + + Raises: + RuntimeError: If the pre-pick pose is unreachable. + """ + position = obj.get_local_pose(to_matrix=True)[:, :3, 3].clone() + position[:, 2] = height + ik_success, arm_qpos = robot.compute_ik( + pose=make_eef_pose_at(robot, position), + joint_seed=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.") + + hand_qpos = hand_open.unsqueeze(0).repeat(robot.get_qpos().shape[0], 1) + for target in (False, True): + robot.set_qpos(arm_qpos, name="arm", target=target) + robot.set_qpos(hand_qpos, name="hand", target=target) + robot.clear_dynamics() + + +def prepare_tutorial_scene( + sim: SimulationManager, + args: argparse.Namespace, + prompt: str, +) -> bool: + """Open an interactive scene and optionally wait before planning. + + Args: + sim: Simulation manager used by the tutorial. + args: Parsed tutorial arguments. + prompt: Prompt displayed for interactive runs. + + Returns: + Whether the caller should retain later interactive pauses. + """ + wait_for_user = should_wait_for_tutorial_input(args) + if should_open_tutorial_window(args): + sim.open_window() + if wait_for_user: + input(prompt) + return wait_for_user + + +def replay_trajectory( + sim: SimulationManager, + robot: Robot, + trajectory: torch.Tensor, + args: argparse.Namespace, + *, + video_prefix: str, + hold_steps: int, + trajectory_sim_steps: int = 4, + hold_sim_steps: int = 2, + joint_ids: list[int] | None = None, + on_trajectory_step: Callable[[int, int], None] | None = None, + look_at: tuple[Sequence[float], Sequence[float], Sequence[float]] | None = None, + record: bool = True, +) -> None: + """Replay a trajectory, optionally record it, and hold its final pose. + + Args: + sim: Simulation manager to step and record. + robot: Robot receiving full-DOF trajectory positions. + trajectory: Full robot trajectory with shape ``(n_envs, n_steps, dof)``. + args: Parsed tutorial arguments controlling auto-play recording. + video_prefix: Output video filename prefix. + hold_steps: Number of final-pose simulation updates after the trajectory. + trajectory_sim_steps: Physics steps for each trajectory waypoint. + hold_sim_steps: Physics steps for each final-pose update. + joint_ids: Optional joint IDs when controlling a robot subset. + on_trajectory_step: Optional callback run after each trajectory update. + look_at: Optional recorder camera pose for auto-play videos. + record: Whether to start auto-play recording for this replay. + """ + recording_started = ( + start_auto_play_recording( + sim, + args, + video_prefix=video_prefix, + look_at=DEFAULT_AUTO_PLAY_LOOK_AT if look_at is None else look_at, + ) + if record + else False + ) + try: + total_steps = trajectory.shape[1] + for step_idx in range(total_steps): + if joint_ids is None: + robot.set_qpos(trajectory[:, step_idx, :]) + else: + robot.set_qpos(trajectory[:, step_idx, :], joint_ids=joint_ids) + sim.update(step=trajectory_sim_steps) + if on_trajectory_step is not None: + on_trajectory_step(step_idx, total_steps) + time.sleep(1e-2) + + final_qpos = trajectory[:, -1, :] + for _ in range(hold_steps): + if joint_ids is None: + robot.set_qpos(final_qpos) + else: + robot.set_qpos(final_qpos, joint_ids=joint_ids) + sim.update(step=hold_sim_steps) + time.sleep(1e-2) + finally: + stop_auto_play_recording(sim, recording_started) def get_tutorial_window_size(args: argparse.Namespace) -> tuple[int, int]: @@ -52,6 +421,25 @@ def get_tutorial_window_size(args: argparse.Namespace) -> tuple[int, int]: return VIEWER_WIDTH, VIEWER_HEIGHT +def should_open_tutorial_window(args: argparse.Namespace) -> bool: + """Return whether an interactive viewer window should be opened.""" + return not ( + getattr(args, "headless", False) + or getattr(args, "diagnose_plan", False) + or getattr(args, "headless_play", False) + ) + + +def should_wait_for_tutorial_input(args: argparse.Namespace) -> bool: + """Return whether the tutorial should pause for terminal input.""" + return not ( + getattr(args, "auto_play", False) + or getattr(args, "headless", False) + or getattr(args, "diagnose_plan", False) + or getattr(args, "headless_play", False) + ) + + def start_auto_play_recording( sim: SimulationManager, args: argparse.Namespace, @@ -104,19 +492,79 @@ def draw_axis_marker( xpos: torch.Tensor, axis_len: float = DEFAULT_AXIS_LEN, axis_size: float = DEFAULT_AXIS_SIZE, + arena_index: int = -1, ) -> None: """Draw a named coordinate-frame marker for a semantic tutorial target.""" + arena_offsets = sim.arena_offsets + n_envs = arena_offsets.shape[0] + + # Normalize a single (4, 4) pose to (1, 4, 4) so it is not mistaken for a + # per-environment batch when num_envs happens to equal 4. + if xpos.dim() == 2: + xpos = xpos.unsqueeze(0) + + if n_envs == xpos.shape[0]: + # add arena offsets to xpos + draw_xpos = xpos.clone() + draw_xpos[:, :3, 3] += arena_offsets + else: + draw_xpos = xpos sim.draw_marker( cfg=MarkerCfg( name=name, marker_type="axis", - axis_xpos=xpos, + axis_xpos=draw_xpos, axis_size=axis_size, axis_len=axis_len, + arena_index=arena_index, ) ) +def broadcast_pose_batch(pose: torch.Tensor, num_envs: int) -> torch.Tensor: + """Expand a single pose to ``(num_envs, 4, 4)`` when needed.""" + if pose.dim() == 2 and pose.shape == (4, 4): + return pose.unsqueeze(0).repeat(num_envs, 1, 1) + if pose.dim() == 3 and pose.shape[1:] == (4, 4): + if pose.shape[0] == num_envs: + return pose + if pose.shape[0] == 1: + return pose.repeat(num_envs, 1, 1) + raise ValueError( + "Expected a pose with shape " + f"(4, 4), (1, 4, 4), or ({num_envs}, 4, 4) for num_envs={num_envs}, " + f"got {tuple(pose.shape)}." + ) + + +def broadcast_waypoint_pose_batch(pose: torch.Tensor, num_envs: int) -> torch.Tensor: + """Expand waypoint poses to ``(num_envs, n_waypoint, 4, 4)`` when needed.""" + if pose.dim() == 3 and pose.shape[1:] == (4, 4): + return pose.unsqueeze(0).repeat(num_envs, 1, 1, 1) + if pose.dim() == 4 and pose.shape[2:] == (4, 4): + if pose.shape[0] == num_envs: + return pose + if pose.shape[0] == 1: + return pose.repeat(num_envs, 1, 1, 1) + raise ValueError( + "Expected waypoint poses with shape " + f"(n_waypoint, 4, 4), (1, n_waypoint, 4, 4), or ({num_envs}, n_waypoint, 4, 4) " + f"for num_envs={num_envs}, got {tuple(pose.shape)}." + ) + + +def clone_local_pose_from_first_env(entity) -> torch.Tensor: + """Copy the first environment's local pose onto every environment.""" + pose = entity.get_local_pose(to_matrix=True) + if pose.dim() != 3 or pose.shape[1:] != (4, 4): + raise ValueError( + f"Expected entity local pose with shape (num_envs, 4, 4), got {tuple(pose.shape)}." + ) + shared_pose = pose[:1].clone().repeat(pose.shape[0], 1, 1) + entity.set_local_pose(shared_pose) + return shared_pose + + def create_ur5_gripper_robot_cfg( init_pos: Sequence[float] = (0.0, 0.0, 0.0), ) -> RobotCfg: @@ -189,11 +637,31 @@ def create_ur5_gripper_robot_cfg( "DEFAULT_AUTO_PLAY_LOOK_AT", "DEFAULT_AXIS_LEN", "DEFAULT_AXIS_SIZE", + "DEFAULT_GRIPPER_CLOSE_QPOS", + "DEFAULT_TUTORIAL_LIGHT_POS", "GRIPPER_HAND_JOINT_PATTERN", "GRIPPER_TCP_Z", "GRIPPER_URDF_PATH", + "TOP_DOWN_EEF_ROTATION", + "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", + "format_tensor", + "get_hand_open_close_qpos", + "initialize_pre_pick_robot_pose", + "make_ur5_solver_cfg", + "make_eef_pose_at", + "make_top_down_eef_pose", "get_tutorial_window_size", + "prepare_tutorial_scene", + "replay_trajectory", + "should_open_tutorial_window", + "should_wait_for_tutorial_input", "start_auto_play_recording", "stop_auto_play_recording", "draw_axis_marker", diff --git a/scripts/tutorials/grasp/grasp_generator.py b/scripts/tutorials/grasp/grasp_generator.py index 603644f93..ba6c488bc 100644 --- a/scripts/tutorials/grasp/grasp_generator.py +++ b/scripts/tutorials/grasp/grasp_generator.py @@ -159,6 +159,7 @@ def create_mug(sim: SimulationManager): static_friction=0.99, ), max_convex_hull_num=16, + acd_method="vhacd", init_pos=[0.55, 0.0, 0.01], init_rot=[0.0, 0.0, -90], body_scale=(4, 4, 4), diff --git a/scripts/tutorials/sim/motion_generator.py b/scripts/tutorials/sim/motion_generator.py index 8947d0e52..3d230e34d 100644 --- a/scripts/tutorials/sim/motion_generator.py +++ b/scripts/tutorials/sim/motion_generator.py @@ -14,149 +14,303 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + +import argparse import time -import torch +from collections.abc import Sequence + import numpy as np -from copy import deepcopy +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, physics_cfg_for_backend from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim.robots import CobotMagicCfg +from embodichain.lab.sim.planners import ( + MotionGenCfg, + MotionGenOptions, + MotionGenerator, + PlanState, + ToppraPlanOptions, + ToppraPlannerCfg, +) from embodichain.lab.sim.planners.utils import TrajectorySampleMethod -from embodichain.lab.sim.planners.motion_generator import MotionGenerator +from embodichain.lab.sim.robots import CobotMagicCfg + +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: + """Parse command line arguments for the motion-generator tutorial.""" + 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.", + ) + return parser.parse_args() + + +def compute_record_look_at( + num_envs: int, + arena_space: float, +) -> tuple[ + tuple[float, float, float], tuple[float, float, float], tuple[float, float, float] +]: + """Return a fixed camera pose that frames the full replicated arena grid.""" + if num_envs <= 0: + raise ValueError(f"num_envs must be positive, got {num_envs}.") + + scene_grid_length = int(np.ceil(np.sqrt(num_envs))) + scene_grid_rows = int(np.ceil(num_envs / scene_grid_length)) + span_x = float(max(scene_grid_length - 1, 0) * arena_space) + span_y = float(max(scene_grid_rows - 1, 0) * arena_space) + scene_extent = 0.5 * max(span_x, span_y) + + target = ( + 0.5 * span_x, + 0.5 * span_y, + DEFAULT_RECORD_TARGET_Z, + ) + eye = ( + 2.6 + target[0] + scene_extent, + -2.2 - scene_extent, + 1.6 + 0.4 * scene_extent, + ) + return eye, target, (0.0, 0.0, 1.0) def move_robot_along_trajectory( - robot: Robot, arm_name: str, qpos_list: list[torch.Tensor], delay: float = 0.1 -): - """ - Set the robot joint positions sequentially along the given joint trajectory. + sim: SimulationManager, + robot: Robot, + arm_name: str, + qpos_trajectory: torch.Tensor | Sequence[torch.Tensor], +) -> None: + """Play back a planned joint trajectory for one or more environments. + + This function assumes the simulation is in manual-update mode and calls + :meth:`SimulationManager.update` after each waypoint so physics advances. + Args: + sim: Simulation manager instance. robot: Robot instance. - arm_name: Name of the robot arm. - qpos_list: List of joint positions (torch.Tensor). - delay: Time delay between each step (seconds). + arm_name: Name of the robot arm to control. + qpos_trajectory: Joint positions shaped ``(B, N, DOF)``, ``(N, DOF)``, + or a sequence of waypoint tensors. + delay: Time delay between each step in seconds. """ - for q in qpos_list: - robot.set_qpos(qpos=q.unsqueeze(0), joint_ids=robot.get_joint_ids(arm_name)) - time.sleep(delay) + if isinstance(qpos_trajectory, Sequence): + qpos_steps = list(qpos_trajectory) + if not qpos_steps: + return + if qpos_steps[0].dim() == 1: + qpos_trajectory = torch.stack(qpos_steps, dim=0).unsqueeze(0) + else: + qpos_trajectory = torch.stack(qpos_steps, dim=1) + if qpos_trajectory.dim() == 2: + qpos_trajectory = qpos_trajectory.unsqueeze(0) + if qpos_trajectory.dim() != 3: + raise ValueError( + "qpos_trajectory must have shape (B, N, DOF) or (N, DOF), " + f"got {tuple(qpos_trajectory.shape)}." + ) + + joint_ids = robot.get_joint_ids(arm_name) + for qpos_step in qpos_trajectory.transpose(0, 1): + robot.set_qpos(qpos=qpos_step, joint_ids=joint_ids) + sim.update(step=4) def create_demo_trajectory( - robot: Robot, arm_name: str -) -> tuple[list[torch.Tensor], list[np.ndarray]]: - """ - Generate a three-point trajectory (start, middle, end) for demonstration. - Args: - robot: Robot instance. - arm_name: Name of the robot arm. - Returns: - qpos_list: List of joint positions (torch.Tensor). - xpos_list: List of end-effector poses (numpy arrays). - """ + robot: Robot, + arm_name: str, + num_envs: int = 1, +) -> tuple[list[torch.Tensor], list[torch.Tensor]]: + """Generate a three-point batched demo trajectory for the requested env count.""" + if num_envs <= 0: + raise ValueError(f"num_envs must be positive, got {num_envs}.") + + tensor_kwargs = {"dtype": torch.float32} + robot_device = getattr(robot, "device", None) + if isinstance(robot_device, (str, torch.device)): + tensor_kwargs["device"] = robot_device + qpos_fk = torch.tensor( - [[0.0, np.pi / 4, -np.pi / 4, 0.0, np.pi / 4, 0.0]], dtype=torch.float32 - ) + [[0.0, np.pi / 4, -np.pi / 4, 0.0, np.pi / 4, 0.0]], + **tensor_kwargs, + ).repeat(num_envs, 1) xpos_begin = robot.compute_fk(name=arm_name, qpos=qpos_fk, to_matrix=True) - xpos_mid = deepcopy(xpos_begin) - xpos_mid[0, 2, 3] -= 0.1 # Move down by 0.1m in Z direction - xpos_final = deepcopy(xpos_mid) - xpos_final[0, 0, 3] += 0.2 # Move forward by 0.2m in X direction - - qpos_begin = robot.compute_ik(pose=xpos_begin, name=arm_name)[1][0] - qpos_mid = robot.compute_ik(pose=xpos_mid, name=arm_name)[1][0] - qpos_final = robot.compute_ik(pose=xpos_final, name=arm_name)[1][0] - return [qpos_begin, qpos_mid, qpos_final], [ - xpos_begin[0], - xpos_mid[0], - xpos_final[0], - ] - - -def main(): + xpos_mid = xpos_begin.clone() + xpos_mid[:, 2, 3] -= 0.1 + xpos_final = xpos_mid.clone() + xpos_final[:, 0, 3] += 0.2 + + qpos_begin = robot.compute_ik(pose=xpos_begin, name=arm_name)[1] + qpos_mid = robot.compute_ik(pose=xpos_mid, name=arm_name)[1] + qpos_final = robot.compute_ik(pose=xpos_final, name=arm_name)[1] + 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) - # Initialize simulation - sim = SimulationManager(SimulationManagerCfg(headless=True, device="cpu")) - sim.set_manual_update(False) + sim = SimulationManager( + SimulationManagerCfg( + width=RECORD_WIDTH, + height=RECORD_HEIGHT, + headless=True, + physics_dt=1.0 / 100.0, + device=args.device, + physics_cfg=physics_cfg_for_backend(args.physics), + render_cfg=RenderCfg(renderer=args.renderer), + num_envs=args.num_envs, + arena_space=args.arena_space, + ) + ) - # Robot configuration - cfg_dict = {"uid": "CobotMagic"} - robot: Robot = sim.add_robot(cfg=CobotMagicCfg.from_dict(cfg_dict)) + robot: Robot = sim.add_robot(cfg=CobotMagicCfg.from_dict({"uid": "CobotMagic"})) arm_name = "left_arm" - sim.open_window() + if sim.is_use_gpu_physics: + sim.init_gpu_physics() - # # Generate trajectory points - qpos_list, xpos_list = create_demo_trajectory(robot, arm_name) + if not args.headless: + sim.open_window() - from embodichain.lab.sim.planners import ( - MotionGenerator, - MotionGenCfg, - MotionGenOptions, - ToppraPlannerCfg, - ToppraPlanOptions, - PlanState, - MoveType, - MovePart, + print( + f"[INFO]: Running motion generator tutorial with {sim.num_envs} environment(s)" ) - # Initialize motion generator - motion_cfg = MotionGenCfg( - planner_cfg=ToppraPlannerCfg( - robot_uid=robot.uid, + recording_started = start_headless_recording(sim, args) + try: + qpos_list, xpos_list = create_demo_trajectory( + robot=robot, + arm_name=arm_name, + num_envs=sim.num_envs, ) - ) - motion_generator = MotionGenerator(cfg=motion_cfg) - - # Joint space trajectory - qpos_list = torch.vstack(qpos_list) - 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, - ), - ) - target_states = [] - for qpos in qpos_list: - target_states.append( - PlanState( - move_type=MoveType.JOINT_MOVE, - move_part=MovePart.LEFT, - qpos=qpos, + motion_generator = MotionGenerator( + cfg=MotionGenCfg( + planner_cfg=ToppraPlannerCfg( + robot_uid=robot.uid, + ) ) ) - plan_result = motion_generator.generate( - target_states=target_states, options=options - ) - move_robot_along_trajectory(robot, arm_name, plan_result.positions) - - # Cartesian space trajectory - options.is_linear = True - - target_states = [] - for xpos in xpos_list: - target_states.append( - PlanState( - move_type=MoveType.EEF_MOVE, - move_part=MovePart.LEFT, - xpos=xpos, + + 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, + ) + + 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, ) - plan_result = motion_generator.generate( - target_states=target_states, options=options - ) - sim.reset() - move_robot_along_trajectory(robot, arm_name, plan_result.positions) + finally: + if sim.is_window_recording(): + sim.stop_window_record() + sim.wait_window_record_saves() + sim.destroy() if __name__ == "__main__": diff --git a/tests/conftest.py b/tests/conftest.py index d076ff132..77959c1af 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,6 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- +import inspect import os import pytest @@ -27,6 +28,12 @@ def pytest_addoption(parser): default="hybrid", help="Specify the renderer backend: hybrid, or fast-rt", ) + parser.addoption( + "--run-gpu", + action="store_true", + default=False, + help="Run tests marked gpu; these are skipped by default to limit VRAM use.", + ) def pytest_configure(config): @@ -42,34 +49,78 @@ def pytest_configure(config): cfg.DEFAULT_RENDERER = renderer - # PREVENT IMPLICIT INITIALIZATION BY EXPLICITLY INITIALIZING DEXSIM HERE - import dexsim - import dexsim.types - - # Map string to dexsim configuration types - renderer_map = { - "hybrid": dexsim.types.Renderer.HYBRID, - "fast-rt": dexsim.types.Renderer.FASTRT, - } - backend_map = { - "hybrid": dexsim.types.Backend.VULKAN, - "fast-rt": dexsim.types.Backend.VULKAN, - } - - if dexsim.get_world_num() == 0: - sim_config = dexsim.WorldConfig() - sim_config.renderer = renderer_map.get( - renderer, dexsim.types.Renderer.HYBRID + # DexSim initialization is intentionally deferred to the first real-simulation + # test. Most of the suite consists of pure-Python tests and should not acquire + # a CUDA/Vulkan context merely because pytest has started. + + +def _requires_real_sim(item): + """Return whether a test module creates a real SimulationManager.""" + if item.get_closest_marker("requires_sim") is not None: + return True + module = getattr(item, "module", None) + if module is not None and "SimulationManager" in vars(module): + return True + + # Some planner regression tests intentionally import the simulation manager + # inside the test body to keep their module import lightweight. + try: + return "SimulationManager" in inspect.getsource(item.obj) + except (OSError, TypeError): + return False + + +def _initialize_sim_engine(renderer): + """Initialize DexSim once, immediately before the first real-sim test.""" + import dexsim + import dexsim.types + + if dexsim.get_world_num() != 0: + return + + renderer_map = { + "hybrid": dexsim.types.Renderer.HYBRID, + "fast-rt": dexsim.types.Renderer.FASTRT, + } + sim_config = dexsim.WorldConfig() + sim_config.renderer = renderer_map[renderer] + sim_config.backend = dexsim.types.Backend.VULKAN + sim_config.open_windows = False + dexsim.init_sim_engine(sim_config) + + +def pytest_collection_modifyitems(config, items): + """Classify real-simulation tests for fast and resource-aware test selection.""" + for item in items: + nodeid = item.nodeid.lower() + requires_sim = _requires_real_sim(item) + if requires_sim: + item.add_marker(pytest.mark.requires_sim) + if "cuda" in nodeid or "gpu" in nodeid: + item.add_marker(pytest.mark.gpu) + if requires_sim and ( + "/sensors/" in nodeid or "hybrid" in nodeid or "fastrt" in nodeid + ): + item.add_marker(pytest.mark.renderer) + + if item.get_closest_marker("gpu") is not None and not config.getoption( + "--run-gpu" + ): + item.add_marker( + pytest.mark.skip( + reason="GPU tests require --run-gpu and should run in a serial job." + ) ) - sim_config.backend = backend_map.get(renderer, dexsim.types.Backend.VULKAN) - sim_config.open_windows = False - # This triggers initialization with the correct properties immediately. - dexsim.init_sim_engine(sim_config) @pytest.fixture(autouse=True, scope="function") -def wait_scene_destruction_after_test(): +def wait_scene_destruction_after_test(request): """Ensure C++ engine scenes are fully destructed globally after each test exits.""" + if not _requires_real_sim(request.node): + yield + return + + _initialize_sim_engine(request.config.getoption("--renderer")) yield # [Improvement - delayed destruction]: top-level dequeue and traceback cleanup. diff --git a/tests/gym/action_bank/test_configurable_action.py b/tests/gym/action_bank/test_configurable_action.py index a7fb00fcd..48dfc9203 100644 --- a/tests/gym/action_bank/test_configurable_action.py +++ b/tests/gym/action_bank/test_configurable_action.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + from embodichain.lab.gym.envs.action_bank.configurable_action import ( ActionBank, tag_node, diff --git a/tests/gym/envs/managers/test_event_functors.py b/tests/gym/envs/managers/test_event_functors.py index 981e44ab0..4c6cb8f06 100644 --- a/tests/gym/envs/managers/test_event_functors.py +++ b/tests/gym/envs/managers/test_event_functors.py @@ -117,6 +117,64 @@ def set_mass(self, link_name: str, mass: float) -> int: return 0 +class MockLight: + """Mock light for event functor tests. + + Supports both global lights (num_instances=1 for sun/direction) + and per-environment batched lights (num_instances=num_envs). + """ + + def __init__( + self, + uid: str = "test_light", + num_envs: int = 4, + light_type: str = "point", + init_pos: tuple = (0.0, 0.0, 2.0), + direction: tuple = (0.0, 0.0, -1.0), + intensity: float = 30.0, + color: tuple = (1.0, 1.0, 1.0), + ): + self.uid = uid + + # Config-like attributes + class CfgProxy: + pass + + self.cfg = CfgProxy() + self.cfg.light_type = light_type + self.cfg.init_pos = init_pos + self.cfg.direction = direction + self.cfg.intensity = intensity + self.cfg.color = color + + self.device = torch.device("cpu") + self.is_global = light_type in ("sun", "direction") + self.num_instances = 1 if self.is_global else num_envs + + # Track calls for assertions + self._last_set_color = None + self._last_set_intensity = None + self._last_set_local_pose = None + self._last_set_direction = None + self._last_env_ids = None + + def set_color(self, color, env_ids=None): + self._last_set_color = color + self._last_env_ids = env_ids + + def set_intensity(self, intensity, env_ids=None): + self._last_set_intensity = intensity + self._last_env_ids = env_ids + + def set_local_pose(self, pose, env_ids=None, to_matrix=False): + self._last_set_local_pose = pose + self._last_env_ids = env_ids + + def set_direction(self, direction, env_ids=None): + self._last_set_direction = direction + self._last_env_ids = env_ids + + class MockArticulation: """Mock articulation for event functor tests.""" @@ -238,6 +296,7 @@ def __init__(self, num_envs: int = 4): self._articulations = {} self._robots = {} self._rigid_object_groups = {} + self._lights = {} def get_rigid_object(self, uid: str): return self._rigid_objects.get(uid) @@ -290,6 +349,13 @@ def set_emission_light(self, color=None, intensity=None) -> None: self._last_emission_color = color self._last_emission_intensity = intensity + def get_light(self, uid: str): + return self._lights.get(uid) + + def add_light(self, light: MockLight): + self._lights[light.uid] = light + return light + class MockEnv: """Mock environment for event functor tests.""" @@ -333,7 +399,9 @@ def __init__(self, num_envs: int = 4, num_joints: int = 6): ) from embodichain.lab.gym.envs.managers.randomization.visual import ( randomize_indirect_lighting, + randomize_light, ) +from embodichain.lab.gym.envs.managers.cfg import SceneEntityCfg from embodichain.lab.gym.envs.managers import FunctorCfg @@ -1010,3 +1078,260 @@ def test_emissive_values_vary_across_calls(self): # Over 20 draws from [0, 1000] all values being identical is astronomically unlikely assert len(intensities) > 1 + + +# -------------------------------------------------------------------------- +# randomize_light tests +# -------------------------------------------------------------------------- + + +class TestRandomizeLight: + """Tests for the randomize_light function.""" + + # ------------------------------------------------------------------ + # Per-environment batched light tests (point light) + # ------------------------------------------------------------------ + + def test_randomize_color(self): + """color_range randomizes the light color for all envs.""" + env = MockEnv(num_envs=4) + light = MockLight("test_pt", num_envs=4, light_type="point") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_pt") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + color_range=[[0.2, 0.2, 0.2], [0.8, 0.8, 0.8]], + ) + + assert light._last_set_color is not None + assert light._last_set_color.shape == (4, 3) + # All values should be within [0.2, 0.8] + assert (light._last_set_color >= 0.2).all() + assert (light._last_set_color <= 0.8).all() + + def test_randomize_intensity(self): + """intensity_range randomizes the light intensity for all envs. + + The intensity_range is additive: random value is added to cfg.intensity. + """ + env = MockEnv(num_envs=4) + light = MockLight("test_pt", num_envs=4, light_type="point") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_pt") + env_ids = torch.tensor([0, 1, 2, 3]) + + # Use a small additive range so we can bound the result + randomize_light( + env, + env_ids, + entity_cfg, + intensity_range=(10.0, 20.0), + ) + + assert light._last_set_intensity is not None + assert light._last_set_intensity.shape == (4,) + # init_intensity=30 + rand(10,20) => [40, 50] + assert (light._last_set_intensity >= 40.0).all() + assert (light._last_set_intensity <= 50.0).all() + + def test_randomize_position(self): + """position_range randomizes the light position for batched lights.""" + env = MockEnv(num_envs=4) + light = MockLight("test_pt", num_envs=4, light_type="point") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_pt") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + position_range=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], + ) + + assert light._last_set_local_pose is not None + assert light._last_set_local_pose.shape == (4, 3) + + def test_randomize_direction(self): + """direction_range randomizes the light direction for spot lights.""" + env = MockEnv(num_envs=4) + light = MockLight("test_spot", num_envs=4, light_type="spot") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_spot") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + direction_range=[[-0.2, -0.2, -0.2], [0.2, 0.2, 0.2]], + ) + + assert light._last_set_direction is not None + assert light._last_set_direction.shape == (4, 3) + + def test_direction_range_ignored_for_point_light(self): + """direction_range is ignored for point lights (no crash).""" + env = MockEnv(num_envs=4) + light = MockLight("test_pt", num_envs=4, light_type="point") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_pt") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + direction_range=[[-0.1, -0.1, -0.1], [0.1, 0.1, 0.1]], + ) + + # direction should NOT have been set for point light + assert light._last_set_direction is None + + def test_combined_randomization(self): + """All four ranges can be combined in one call.""" + env = MockEnv(num_envs=4) + light = MockLight("test_spot", num_envs=4, light_type="spot") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_spot") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + position_range=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], + color_range=[[0.2, 0.2, 0.2], [0.8, 0.8, 0.8]], + intensity_range=(50.0, 100.0), + direction_range=[[-0.1, -0.1, -0.1], [0.1, 0.1, 0.1]], + ) + + assert light._last_set_local_pose is not None + assert light._last_set_color is not None + assert light._last_set_intensity is not None + assert light._last_set_direction is not None + + def test_light_not_found_returns_early(self): + """No crash when light UID is not registered.""" + env = MockEnv(num_envs=4) + entity_cfg = SceneEntityCfg(uid="nonexistent") + env_ids = torch.tensor([0, 1, 2, 3]) + + # Should not raise + try: + randomize_light( + env, + env_ids, + entity_cfg, + color_range=[[0.2, 0.2, 0.2], [0.8, 0.8, 0.8]], + ) + except Exception as e: + pytest.fail(f"randomize_light should not crash on missing light: {e}") + + # ------------------------------------------------------------------ + # Global scene light tests (sun, direction) + # ------------------------------------------------------------------ + + @pytest.mark.parametrize("light_type", ["sun", "direction"]) + def test_global_light_position_range_ignored(self, light_type): + """position_range is ignored for global lights (sun/direction).""" + env = MockEnv(num_envs=4) + light = MockLight("test_global", num_envs=1, light_type=light_type) + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_global") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + position_range=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], + ) + + # set_local_pose should NOT be called for global lights + assert light._last_set_local_pose is None + + @pytest.mark.parametrize("light_type", ["sun", "direction"]) + def test_global_light_color_applied_single_instance(self, light_type): + """color_range works for global lights — single instance broadcast.""" + env = MockEnv(num_envs=4) + light = MockLight("test_global", num_envs=1, light_type=light_type) + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_global") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + color_range=[[0.3, 0.3, 0.3], [0.7, 0.7, 0.7]], + ) + + assert light._last_set_color is not None + # Global light: single instance, color shape should be (1, 3) + assert light._last_set_color.shape == (1, 3) + # env_ids=None for single-instance lights + assert light._last_env_ids is None + + @pytest.mark.parametrize("light_type", ["sun", "direction"]) + def test_global_light_intensity_applied_single_instance(self, light_type): + """intensity_range works for global lights — single instance.""" + env = MockEnv(num_envs=4) + light = MockLight("test_global", num_envs=1, light_type=light_type) + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_global") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + intensity_range=(50.0, 100.0), + ) + + assert light._last_set_intensity is not None + assert light._last_set_intensity.shape == (1,) + + @pytest.mark.parametrize("light_type", ["sun", "direction"]) + def test_global_light_direction_range_applied(self, light_type): + """direction_range works for global lights (sun/direction).""" + env = MockEnv(num_envs=4) + light = MockLight("test_global", num_envs=1, light_type=light_type) + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_global") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + direction_range=[[-0.2, -0.2, -0.2], [0.2, 0.2, 0.2]], + ) + + assert light._last_set_direction is not None + assert light._last_set_direction.shape == (1, 3) + assert light._last_env_ids is None + + # ------------------------------------------------------------------ + # Edge cases + # ------------------------------------------------------------------ + + def test_no_ranges_does_nothing(self): + """Calling with no ranges set does not crash.""" + env = MockEnv(num_envs=4) + light = MockLight("test_pt", num_envs=4, light_type="point") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_pt") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light(env, env_ids, entity_cfg) + + # Nothing should have been called + assert light._last_set_color is None + assert light._last_set_intensity is None + assert light._last_set_local_pose is None + assert light._last_set_direction is None diff --git a/tests/gym/envs/managers/test_randomize_anchor_height.py b/tests/gym/envs/managers/test_randomize_anchor_height.py new file mode 100644 index 000000000..1c6acf17e --- /dev/null +++ b/tests/gym/envs/managers/test_randomize_anchor_height.py @@ -0,0 +1,492 @@ +# ---------------------------------------------------------------------------- +# 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 pytest +import torch +from unittest.mock import MagicMock + +from embodichain.lab.gym.envs.managers import EventCfg +from embodichain.lab.gym.envs.managers.randomization.spatial import ( + randomize_anchor_height, +) +from embodichain.lab.sim.objects import RigidObjectGroup + +# --------------------------------------------------------------------------- +# Shared mock classes +# --------------------------------------------------------------------------- + + +class _MockObject: + """Base mock for RigidObject / Articulation with (N, 7) pose storage.""" + + def __init__(self, uid: str, num_envs: int = 4): + self.uid = uid + self.num_envs = num_envs + self.device = torch.device("cpu") + self.cfg = MagicMock() + self.cfg.init_pos = [0.0, 0.0, 0.0] + self._pose = torch.zeros(num_envs, 7) + self._pose[:, 3] = 1.0 # identity quaternion + self._cleared = False + self._cleared_env_ids = None + + def get_local_pose(self, to_matrix: bool = False): + if to_matrix: + mat = torch.eye(4).unsqueeze(0).repeat(self.num_envs, 1, 1) + mat[:, :3, 3] = self._pose[:, :3] + return mat + return self._pose.clone() + + def set_local_pose(self, pose, env_ids=None): + if env_ids is None: + env_ids = torch.arange(self.num_envs) + if pose.ndim == 3: + self._pose[env_ids, :3] = pose[:, :3, 3] + else: + self._pose[env_ids] = pose + + def clear_dynamics(self, env_ids=None): + self._cleared = True + self._cleared_env_ids = env_ids + + +class _MockArticulation(_MockObject): + """Articulation mock — identical to _MockObject in behavior.""" + + pass + + +class _MockGroup(RigidObjectGroup): + """Mock for RigidObjectGroup with (N, M, 4, 4) pose storage. + + Inherits from RigidObjectGroup so isinstance checks pass, but skips the + real parent ``__init__`` to avoid heavy simulation dependencies. + """ + + def __init__(self, uid: str, num_envs: int = 4, num_objects: int = 2): + # Skip RigidObjectGroup.__init__ — only set what the functor needs. + self.uid = uid + self.num_envs = num_envs + self.device = torch.device("cpu") + self.cfg = MagicMock() + self.cfg.init_pos = [0.0, 0.0, 0.0] + # num_objects is a property backed by self._data.num_objects + self._data = MagicMock() + self._data.num_objects = num_objects + self._pose = ( + torch.eye(4).unsqueeze(0).unsqueeze(0).repeat(num_envs, num_objects, 1, 1) + ) + self._cleared = False + self._cleared_env_ids = None + + def get_local_pose(self, to_matrix: bool = False): + return self._pose.clone() + + def set_local_pose(self, pose, env_ids=None): + if env_ids is None: + env_ids = torch.arange(self.num_envs) + self._pose[env_ids] = pose + + def clear_dynamics(self, env_ids=None): + self._cleared = True + self._cleared_env_ids = env_ids + + +class _MockSim: + def __init__(self, num_envs: int = 4): + self.num_envs = num_envs + self.device = torch.device("cpu") + self._objects: dict[str, _MockObject] = {} + self._articulations: dict[str, _MockArticulation] = {} + self._groups: dict[str, _MockGroup] = {} + + def get_rigid_object(self, uid: str): + return self._objects.get(uid) + + def get_rigid_object_uid_list(self): + return list(self._objects.keys()) + + def get_articulation(self, uid: str): + return self._articulations.get(uid) + + def get_articulation_uid_list(self): + return list(self._articulations.keys()) + + def get_rigid_object_group(self, uid: str): + return self._groups.get(uid) + + def get_rigid_object_group_uid_list(self): + return list(self._groups.keys()) + + def add_rigid_object(self, obj): + self._objects[obj.uid] = obj + + def add_articulation(self, obj): + self._articulations[obj.uid] = obj + + def add_rigid_object_group(self, obj): + self._groups[obj.uid] = obj + + def update(self, step: int): + pass + + +class _MockEnv: + def __init__(self, num_envs: int = 4): + self.num_envs = num_envs + self.device = torch.device("cpu") + self.sim = _MockSim(num_envs) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def env(): + """Return a fresh MockEnv with 4 environments.""" + return _MockEnv() + + +@pytest.fixture +def make_functor(env): + """Create a functor instance wired to the given env.""" + + def _make(**kwargs): + return randomize_anchor_height( + EventCfg(func=randomize_anchor_height, mode="reset", **kwargs), env + ) + + return _make + + +@pytest.fixture +def env_ids(): + """Default env_ids tensor (all 4 environments).""" + return torch.arange(4) + + +@pytest.fixture +def env_with_table(env): + """Env with a table rigid object at init_pos Z=1.0.""" + table = _MockObject("table") + table.cfg.init_pos = [0.0, 0.0, 1.0] + env.sim.add_rigid_object(table) + return env + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + + +def test_missing_anchor_uid_raises(env, make_functor, env_ids): + functor = make_functor() + with pytest.raises(ValueError): + functor( + env, + env_ids, + anchor_uid="missing_table", + height_delta_range=([-0.05], [0.05]), + ) + + +def test_missing_sampling_fields_raises(env_with_table, make_functor, env_ids): + functor = make_functor() + with pytest.raises(ValueError): + functor(env_with_table, env_ids, anchor_uid="table") + + +def test_empty_candidates_raises(env_with_table, make_functor, env_ids): + functor = make_functor() + with pytest.raises(ValueError): + functor( + env_with_table, + env_ids, + anchor_uid="table", + height_delta_candidates=[], + ) + + +def test_invalid_include_groups_raises(env_with_table, make_functor, env_ids): + functor = make_functor() + with pytest.raises(ValueError, match="Invalid include_groups"): + functor( + env_with_table, + env_ids, + anchor_uid="table", + height_delta_range=([-0.05], [0.05]), + include_groups=["invalid_group", "rigid_object"], + ) + + +# --------------------------------------------------------------------------- +# Sampling +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("num_envs", [100]) +def test_range_sampling_within_bounds(env, make_functor, num_envs): + env.num_envs = num_envs + env.sim = _MockSim(num_envs) + table = _MockObject("table", num_envs) + table.cfg.init_pos = [0.0, 0.0, 1.0] + env.sim.add_rigid_object(table) + + functor = make_functor() + functor( + env, + torch.arange(num_envs), + anchor_uid="table", + height_delta_range=([-0.05], [0.05]), + store_key="table_delta", + ) + + delta = getattr(env, "table_delta") + assert delta.shape == (num_envs,) + assert (delta >= -0.05).all() + assert (delta <= 0.05).all() + + +@pytest.mark.parametrize("num_envs", [50]) +def test_discrete_sampling_only_candidates(env, make_functor, num_envs): + env.num_envs = num_envs + env.sim = _MockSim(num_envs) + table = _MockObject("table", num_envs) + env.sim.add_rigid_object(table) + + functor = make_functor() + functor( + env, + torch.arange(num_envs), + anchor_uid="table", + height_delta_candidates=[-0.05, 0.0, 0.05], + store_key="table_delta", + ) + + candidates = torch.tensor([-0.05, 0.0, 0.05]) + for val in getattr(env, "table_delta"): + assert torch.any(torch.isclose(val, candidates)), f"{val} not in {candidates}" + + +# --------------------------------------------------------------------------- +# Shift correctness +# --------------------------------------------------------------------------- + + +def test_anchor_and_objects_shifted_by_same_delta(env, make_functor, env_ids): + table = _MockObject("table") + table.cfg.init_pos = [0.0, 0.0, 1.0] + env.sim.add_rigid_object(table) + + cube = _MockObject("cube") + cube._pose[:, 2] = 1.1 + env.sim.add_rigid_object(cube) + + functor = make_functor() + functor(env, env_ids, anchor_uid="table", height_delta_range=([0.05], [0.05])) + + # Anchor: absolute -> init_pos[2] + delta = 1.0 + 0.05 = 1.05 + torch.testing.assert_close(table._pose[:, 2], torch.ones(4) * 1.05) + # Affected: relative -> current_z + delta = 1.1 + 0.05 = 1.15 + torch.testing.assert_close(cube._pose[:, 2], torch.ones(4) * 1.15) + + +def test_xy_and_rotation_unchanged(env, make_functor, env_ids): + table = _MockObject("table") + table.cfg.init_pos = [0.0, 0.0, 1.0] + env.sim.add_rigid_object(table) + + cube = _MockObject("cube") + cube._pose[:, 0] = 0.5 + cube._pose[:, 1] = -0.3 + env.sim.add_rigid_object(cube) + + original_xy = cube._pose[:, :2].clone() + original_rot = cube._pose[:, 3:7].clone() + + functor = make_functor() + functor(env, env_ids, anchor_uid="table", height_delta_range=([0.1], [0.1])) + + torch.testing.assert_close(cube._pose[:, :2], original_xy) + torch.testing.assert_close(cube._pose[:, 3:7], original_rot) + + +def test_exclude_uids_are_not_moved(env, make_functor, env_ids): + table = _MockObject("table") + table.cfg.init_pos = [0.0, 0.0, 1.0] + env.sim.add_rigid_object(table) + + cube = _MockObject("cube") + cube._pose[:, 2] = 1.1 + env.sim.add_rigid_object(cube) + + floor = _MockObject("floor") + floor._pose[:, 2] = 0.0 + env.sim.add_rigid_object(floor) + + functor = make_functor() + functor( + env, + env_ids, + anchor_uid="table", + height_delta_range=([0.1], [0.1]), + exclude_uids=["floor"], + ) + + torch.testing.assert_close(floor._pose[:, 2], torch.zeros(4)) + torch.testing.assert_close(cube._pose[:, 2], torch.ones(4) * 1.2) + + +def test_articulation_shifted(env, make_functor, env_ids): + table = _MockObject("table") + table.cfg.init_pos = [0.0, 0.0, 1.0] + env.sim.add_rigid_object(table) + + cabinet = _MockArticulation("cabinet") + cabinet._pose[:, 2] = 1.2 + env.sim.add_articulation(cabinet) + + functor = make_functor() + functor(env, env_ids, anchor_uid="table", height_delta_range=([0.1], [0.1])) + + torch.testing.assert_close(cabinet._pose[:, 2], torch.ones(4) * 1.3) + + +@pytest.mark.parametrize("num_envs", [100]) +def test_asymmetric_delta_range(env, make_functor, env_ids, num_envs): + env.num_envs = num_envs + env.sim = _MockSim(num_envs) + table = _MockObject("table", num_envs) + table.cfg.init_pos = [0.0, 0.0, 1.0] + env.sim.add_rigid_object(table) + + functor = make_functor() + functor( + env, + torch.arange(num_envs), + anchor_uid="table", + height_delta_range=([-0.1], [0.05]), + store_key="table_delta", + ) + + delta = getattr(env, "table_delta") + assert delta.shape == (num_envs,) + assert (delta >= -0.1).all() + assert (delta <= 0.05).all() + + +# --------------------------------------------------------------------------- +# Partial env_ids +# --------------------------------------------------------------------------- + + +def test_partial_env_ids(env, make_functor): + table = _MockObject("table") + table.cfg.init_pos = [0.0, 0.0, 1.0] + env.sim.add_rigid_object(table) + + cube = _MockObject("cube") + cube._pose[:, 2] = 1.1 + env.sim.add_rigid_object(cube) + + functor = make_functor() + partial_ids = torch.tensor([0, 2]) + functor( + env, + partial_ids, + anchor_uid="table", + height_delta_range=([0.1], [0.1]), + ) + + # Envs 0,2 shifted; envs 1,3 unchanged + torch.testing.assert_close(table._pose[0, 2], torch.tensor(1.1)) + torch.testing.assert_close(table._pose[2, 2], torch.tensor(1.1)) + torch.testing.assert_close(cube._pose[0, 2], torch.tensor(1.2)) + torch.testing.assert_close(cube._pose[2, 2], torch.tensor(1.2)) + torch.testing.assert_close(table._pose[1, 2], torch.tensor(0.0)) + torch.testing.assert_close(table._pose[3, 2], torch.tensor(0.0)) + torch.testing.assert_close(cube._pose[1, 2], torch.tensor(1.1)) + torch.testing.assert_close(cube._pose[3, 2], torch.tensor(1.1)) + + # clear_dynamics called with targeted env_ids + assert table._cleared and table._cleared_env_ids is not None + torch.testing.assert_close(table._cleared_env_ids, partial_ids) + assert cube._cleared and cube._cleared_env_ids is not None + torch.testing.assert_close(cube._cleared_env_ids, partial_ids) + + +# --------------------------------------------------------------------------- +# RigidObjectGroup anchor +# --------------------------------------------------------------------------- + + +def test_rigid_object_group_anchor_absolute_warning(env, make_functor, env_ids, caplog): + """When anchor is a RigidObjectGroup and absolute=True, a warning is emitted + and the relative shift is applied.""" + group = _MockGroup("group", num_objects=2) + group._pose[:, :, 2, 3] = 1.0 + env.sim.add_rigid_object_group(group) + + cube = _MockObject("cube") + cube._pose[:, 2] = 1.1 + env.sim.add_rigid_object(cube) + + functor = make_functor() + functor( + env, + env_ids, + anchor_uid="group", + height_delta_range=([0.1], [0.1]), + ) + + assert "absolute=True is not supported for RigidObjectGroup" in caplog.text + # Group: relative shift -> 1.0 + 0.1 = 1.1 + torch.testing.assert_close(group._pose[:, :, 2, 3], torch.ones(4, 2) * 1.1) + # Cube: relative shift -> 1.1 + 0.1 = 1.2 + torch.testing.assert_close(cube._pose[:, 2], torch.ones(4) * 1.2) + + +# --------------------------------------------------------------------------- +# Integration smoke test +# --------------------------------------------------------------------------- + + +@pytest.mark.skip(reason="Requires full simulation stack; run manually.") +def test_anchor_height_event_wires_into_embodied_env_cfg(): + """Smoke test: functor can be wired into an EmbodiedEnvCfg.""" + from embodichain.lab.gym.envs import EmbodiedEnvCfg + from embodichain.lab.sim.cfg import RigidObjectCfg + + cfg = EmbodiedEnvCfg() + cfg.events.anchor_height = EventCfg( + func=randomize_anchor_height, + mode="reset", + params={ + "anchor_uid": "table", + "height_delta_range": ([-0.05], [0.05]), + }, + ) + cfg.background.append( + RigidObjectCfg(uid="table", init_pos=[0.0, 0.0, 0.8], body_type="static") + ) + cfg.rigid_object.append( + RigidObjectCfg(uid="cube", init_pos=[0.1, 0.0, 0.9], body_type="dynamic") + ) + assert hasattr(cfg.events, "anchor_height") diff --git a/tests/gym/envs/test_base_env.py b/tests/gym/envs/test_base_env.py index bd353cd65..1d9db34f6 100644 --- a/tests/gym/envs/test_base_env.py +++ b/tests/gym/envs/test_base_env.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os import torch import pytest diff --git a/tests/gym/envs/test_differentiable_env.py b/tests/gym/envs/test_differentiable_env.py index 131045970..93b363cc4 100644 --- a/tests/gym/envs/test_differentiable_env.py +++ b/tests/gym/envs/test_differentiable_env.py @@ -17,16 +17,431 @@ from __future__ import annotations +from types import SimpleNamespace +from typing import Any + +import numpy as np import pytest import torch +import warp as wp from embodichain.lab.gym.envs.differentiable_env import ( DifferentiableEmbodiedEnv, ) from embodichain.lab.gym.envs.embodied_env import EmbodiedEnvCfg from embodichain.lab.sim.cfg import DefaultPhysicsCfg, NewtonPhysicsCfg +from embodichain.lab.sim.diff import NewtonStepFunc, differentiable_step +import embodichain.lab.sim.diff.bridge as diff_bridge from embodichain.lab.sim.sim_manager import SimulationManagerCfg +_CONTROL_SUBSTEPS = 3 + + +@wp.kernel +def _write_bridge_joint_force_kernel( + action: wp.array(dtype=wp.float32), + joint_f: wp.array(dtype=wp.float32), +) -> None: + """Write one tape-tracked action value into Newton joint force.""" + joint_f[0] = action[0] + + +@wp.kernel +def _bridge_terminal_loss_kernel( + body_q: wp.array(dtype=wp.transform), + body_id: int, + target: wp.vec3, + loss: wp.array(dtype=wp.float32), +) -> None: + """Measure a terminal body-position loss inside the Warp tape.""" + delta = wp.transform_get_translation(body_q[body_id]) - target + loss[0] = wp.dot(delta, delta) + + +class _FakeModel: + """Keep the pre-contract bridge path runnable for clean RED failures.""" + + def __init__(self) -> None: + self.states: list[_FakeState] = [] + + def state(self) -> "_FakeState": + state = _FakeState(f"trajectory-{len(self.states)}") + self.states.append(state) + return state + + +class _FakeState: + """State buffer with explicit detached-copy observability.""" + + def __init__(self, name: str, value: int = 0) -> None: + self.name = name + self.value = value + self.assign_sources: list[_FakeState] = [] + + def assign(self, other: "_FakeState") -> None: + """Copy state and retain every publication source for assertions.""" + self.value = other.value + self.assign_sources.append(other) + + +class _FakeStepper: + """Fallback used only while proving the old private route is rejected.""" + + def __init__(self, *, raise_on_step: bool = False) -> None: + self.calls: list[tuple[object, object, object, float]] = [] + self._raise_on_step = raise_on_step + + def create_contacts(self) -> object: + return object() + + def step( + self, + state_in: object, + state_out: object, + *, + contacts: object, + dt: float, + ) -> None: + self.calls.append((state_in, state_out, contacts, dt)) + if self._raise_on_step: + raise RuntimeError("injected trajectory-step failure") + state_out.value = state_in.value + 1 + + +class _RecordingTape: + """Expose construction, exit, and recording ownership of the fake tape.""" + + def __init__(self, warp: "_RecordingWarp") -> None: + self._warp = warp + + def __enter__(self) -> "_RecordingTape": + assert not self._warp.tape_active + self._warp.tape_active = True + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: Any, + ) -> bool: + del exc_type, exc_value, traceback + assert self._warp.tape_active + self._warp.tape_active = False + self._warp.events.append("tape.exit") + return False + + def reset(self) -> None: + """Model the tape cleanup required before releasing a trajectory.""" + assert not self._warp.tape_active + self._warp.events.append("tape.reset") + + def backward(self, *_args: Any, **_kwargs: Any) -> None: + """Provide the minimal action gradient required by bridge tests.""" + self._warp.events.append("tape.backward") + if self._warp.raise_on_tape_backward: + raise RuntimeError("injected tape backward failure") + if self._warp.last_action is not None: + self._warp.last_action.grad = torch.ones_like(self._warp.last_action) + + def zero(self) -> None: + """Keep the current bridge executable until it migrates to reset().""" + self._warp.events.append("tape.zero") + + +class _RecordingWarp: + """Tiny Warp fake that makes tape ownership observable to a manager.""" + + float32 = object() + + def __init__(self) -> None: + self.tape_active = False + self.events: list[str] = [] + self.last_action: torch.Tensor | None = None + self.raise_on_tape_backward = False + + def Tape(self) -> _RecordingTape: + """Record construction before returning a tape context manager.""" + self.events.append("tape.construct") + return _RecordingTape(self) + + def from_torch( + self, tensor: torch.Tensor, *, requires_grad: bool = False, **_: Any + ) -> torch.Tensor: + """Preserve the test tensor as the fake Warp action array.""" + action = tensor.detach().clone().requires_grad_(requires_grad) + self.last_action = action + return action + + def to_torch(self, tensor: torch.Tensor) -> torch.Tensor: + """Expose a fake Warp array to the PyTorch bridge.""" + if self.last_action is not None and tensor is self.last_action.grad: + self.events.append("action-gradient.capture") + return tensor + + +class _ManagerOwnedTrajectory: + """Fake public trajectory whose stepping must occur inside the tape.""" + + def __init__( + self, + manager: "_TrajectoryNewtonManager", + *, + physics_steps: int, + physics_dt: float, + ) -> None: + self._manager = manager + self.control = object() + self.physics_steps = physics_steps + self.physics_dt = physics_dt + self.total_solver_steps = physics_steps * manager.num_substeps + self.states = [ + _FakeState(f"trajectory-state-{index}") + for index in range(self.total_solver_steps + 1) + ] + self.states[0].assign(manager._state_0) + self.contacts = [object() for _ in range(self.total_solver_steps)] + self.step_calls = 0 + self._released = False + + @property + def final_state(self) -> _FakeState: + """Return the terminal state owned by this one taped trajectory.""" + return self.states[-1] + + def step(self) -> _FakeState: + """Advance the owned trajectory and expose tape placement.""" + self._manager.events.append("trajectory.step") + assert self._manager.warp.tape_active + self.step_calls += 1 + if self._manager.raise_on_trajectory_step: + raise RuntimeError("injected trajectory-step failure") + for state_in, state_out in zip(self.states, self.states[1:]): + state_out.value = state_in.value + 1 + return self.final_state + + def release(self) -> None: + """Release this trajectory's model lease after its tape is reset.""" + if self._released: + return + self._manager._release_differentiable_trajectory(self) + self._released = True + + +class _TrajectoryNewtonManager: + """Fake Newton manager for the manager-owned trajectory bridge contract.""" + + def __init__(self, warp: _RecordingWarp, *, num_substeps: int = 1) -> None: + self.warp = warp + self.events = warp.events + self._state_0 = _FakeState("live-state-0") + self._state_1 = _FakeState("live-state-1") + # Keep the old private path runnable so each regression fails on the + # missing public trajectory contract rather than a fake-only error. + self._model = _FakeModel() + self._control = object() + self.num_substeps = num_substeps + self.solver_dt = 0.01 + self._dt = self.solver_dt * self.num_substeps + self.physics_dt = self._dt + self.trajectory_requests: list[dict[str, Any]] = [] + self.trajectories: list[_ManagerOwnedTrajectory] = [] + self.commits: list[_ManagerOwnedTrajectory] = [] + self.commit_assignment_counts: list[tuple[int, int]] = [] + self._active_trajectory: _ManagerOwnedTrajectory | None = None + self.raise_on_trajectory_step = False + + def create_differentiable_trajectory( + self, *, physics_steps: int, physics_dt: float + ) -> _ManagerOwnedTrajectory: + """Create the public trajectory before the bridge opens its tape.""" + if physics_steps < 1: + raise ValueError("physics_steps must be positive") + if self._active_trajectory is not None: + raise RuntimeError( + "A differentiable trajectory is still active; release it after " + "backward before creating another trajectory." + ) + self.events.append("create") + trajectory = _ManagerOwnedTrajectory( + self, + physics_steps=physics_steps, + physics_dt=physics_dt, + ) + self.trajectories.append(trajectory) + self._active_trajectory = trajectory + self.trajectory_requests.append( + { + "physics_steps": physics_steps, + "physics_dt": physics_dt, + "tape_active": self.warp.tape_active, + } + ) + return trajectory + + def commit_differentiable_trajectory( + self, trajectory: _ManagerOwnedTrajectory + ) -> None: + """Record a detached post-tape publication through the manager API.""" + assert not self.warp.tape_active + assert trajectory in self.trajectories + before = (len(self._state_0.assign_sources), len(self._state_1.assign_sources)) + self._state_0.assign(trajectory.final_state) + self._state_1.assign(trajectory.final_state) + self.commits.append(trajectory) + self.commit_assignment_counts.append( + ( + len(self._state_0.assign_sources) - before[0], + len(self._state_1.assign_sources) - before[1], + ) + ) + self.events.append("commit") + + def _release_differentiable_trajectory( + self, trajectory: _ManagerOwnedTrajectory + ) -> None: + """Release the one active trajectory once tape ownership has ended.""" + assert self._active_trajectory is trajectory + self._active_trajectory = None + self.events.append("trajectory.release") + + +class _TrajectorySimulationManager: + """Bridge-facing manager exposing public and legacy test doubles.""" + + def __init__(self, warp: _RecordingWarp, *, num_substeps: int = 1) -> None: + self.is_newton_backend = True + self.physics = SimpleNamespace( + newton_manager=_TrajectoryNewtonManager(warp, num_substeps=num_substeps) + ) + self.steppers: list[_FakeStepper] = [] + + def create_differentiable_stepper(self) -> _FakeStepper: + """Keep the pre-contract bridge executable for a clean RED failure.""" + stepper = _FakeStepper( + raise_on_step=self.physics.newton_manager.raise_on_trajectory_step + ) + self.steppers.append(stepper) + return stepper + + +class _RealBridgeManager: + """Expose only the public Newton-trajectory surface to the bridge.""" + + def __init__(self, newton_manager: Any) -> None: + self.is_newton_backend = True + self.physics = SimpleNamespace(newton_manager=newton_manager) + + def create_differentiable_stepper(self) -> None: + """Fail if the bridge retains the removed SimulationManager route.""" + raise AssertionError( + "NewtonStepFunc must use NewtonManager.create_differentiable_trajectory(), " + "not SimulationManager.create_differentiable_stepper()." + ) + + +def _route_env( + manager: Any, + *, + mode: str | None = None, + control_substeps: int = _CONTROL_SUBSTEPS, +) -> tuple[DifferentiableEmbodiedEnv, list[object]]: + """Build an uninitialized environment with only the route dependencies.""" + env = object.__new__(DifferentiableEmbodiedEnv) + env.sim = manager + env.cfg = SimpleNamespace(sim_steps_per_control=control_substeps) + if mode is not None: + env.differentiable_step_mode = mode + final_states: list[object] = [] + + def _apply_dynamics_action( + _action_wp: torch.Tensor, _control: Any, tape: Any + ) -> None: + del tape + + def _apply_kinematic_action(_action_wp: torch.Tensor, tape: Any) -> None: + del tape + + def _read_outputs(final_state: object) -> dict[str, Any]: + final_states.append(final_state) + return { + "obs": torch.zeros(1, 1), + "reward": torch.zeros(1), + "terminated": torch.zeros(1, dtype=torch.bool), + "truncated": torch.zeros(1, dtype=torch.bool), + "_order": ("obs", "reward", "terminated", "truncated"), + "_grad_track": {}, + } + + env._apply_dynamics_action_kernel = _apply_dynamics_action + env._apply_action_kernel = _apply_kinematic_action + env._read_outputs = _read_outputs + return env, final_states + + +def _manager_owned_trajectory_sim_state( + manager: _TrajectorySimulationManager, + *, + action_to_control_kernel: Any, + step_mode: str | None = None, + step_fn: Any | None = None, +) -> dict[str, Any]: + """Build the narrow bridge input used by manager-owned trajectory tests.""" + nm = manager.physics.newton_manager + + def _read_outputs(final_state: _FakeState) -> dict[str, Any]: + del final_state + assert nm.warp.tape_active + nm.events.append("outputs") + return { + "obs": torch.zeros(1, 1), + "reward": torch.zeros(1), + "terminated": torch.zeros(1, dtype=torch.bool), + "truncated": torch.zeros(1, dtype=torch.bool), + "_order": ("obs", "reward", "terminated", "truncated"), + "_grad_track": {}, + } + + sim_state: dict[str, Any] = { + "manager": manager, + "substeps": _CONTROL_SUBSTEPS, + "physics_dt": nm.physics_dt, + "action_to_control_kernel": action_to_control_kernel, + "kernel_args": ("kernel-argument",), + "obs_reward_fn": _read_outputs, + } + if step_mode is not None: + sim_state["step_mode"] = step_mode + if step_fn is not None: + sim_state["step_fn"] = step_fn + return sim_state + + +def _assert_tape_reset_then_trajectory_release(events: list[str]) -> None: + """Require one terminal tape reset followed immediately by release.""" + assert events.count("tape.reset") == 1 + assert events.count("trajectory.release") == 1 + reset_index = events.index("tape.reset") + release_index = events.index("trajectory.release") + assert events.index("tape.exit") < reset_index < release_index + assert events[-2:] == ["tape.reset", "trajectory.release"] + + +def _assert_backward_captures_gradient_then_releases(events: list[str]) -> None: + """Require gradient capture before terminal tape and trajectory cleanup.""" + tracked_events = { + "tape.backward", + "action-gradient.capture", + "tape.reset", + "trajectory.release", + } + assert [event for event in events if event in tracked_events] == [ + "tape.backward", + "action-gradient.capture", + "tape.reset", + "trajectory.release", + ] + def _diff_env_cfg( requires_grad: bool = True, backend: str = "newton" @@ -47,6 +462,1016 @@ def _diff_env_cfg( return EmbodiedEnvCfg(sim_cfg=sim_cfg) +def test_default_dynamics_manager_trajectory_lifecycle_is_fully_ordered( + monkeypatch, +) -> None: + """Allocate, tape, action, step, output, and commit stay in one order.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + + def _apply_action(_action: torch.Tensor, *_args: Any) -> None: + assert warp.tape_active + nm.events.append("action") + + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=_apply_action, + step_mode="dynamics", + ), + ) + + assert nm.trajectory_requests == [ + { + "physics_steps": _CONTROL_SUBSTEPS, + "physics_dt": nm.physics_dt, + "tape_active": False, + } + ] + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert trajectory.step_calls == 1 + assert nm.events == [ + "create", + "tape.construct", + "action", + "trajectory.step", + "outputs", + "tape.exit", + "commit", + ] + assert nm.commits == [trajectory] + assert nm.commit_assignment_counts == [(1, 1)] + assert [len(state.assign_sources) for state in (nm._state_0, nm._state_1)] == [ + 1, + 1, + ] + + +def test_default_dynamics_action_hook_receives_trajectory_local_control( + monkeypatch, +) -> None: + """The taped action write never targets the manager's shared control.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + received: list[tuple[tuple[Any, ...], bool]] = [] + + def _apply_action(_action: torch.Tensor, *args: Any) -> None: + received.append((args, warp.tape_active)) + + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=_apply_action, + step_mode="dynamics", + ), + ) + + nm = manager.physics.newton_manager + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert received == [((trajectory.control, "kernel-argument"), True)] + assert received[0][0][0] is not nm._control + + +def test_dynamics_legacy_action_type_error_is_not_retried_after_creation( + monkeypatch, +) -> None: + """Dynamics propagates a legacy callback error instead of falling back.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + legacy_calls: list[tuple[torch.Tensor, Any]] = [] + + def _legacy_action(action_wp: torch.Tensor, tape: Any) -> None: + legacy_calls.append((action_wp, tape)) + raise TypeError("original legacy action TypeError") + + sim_state = _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=_legacy_action, + step_mode="dynamics", + ) + # With no extra kernel arguments, the legacy two-argument callback is + # entered once with local control in its obsolete ``tape`` position. + # Retrying after its body raises TypeError would invoke it a second time. + sim_state["kernel_args"] = () + + with pytest.raises(TypeError) as exc_info: + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + assert str(exc_info.value) == "original legacy action TypeError" + assert len(legacy_calls) == 1 + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert legacy_calls[0][1] is trajectory.control + assert trajectory.step_calls == 0 + assert nm.commits == [] + assert nm.commit_assignment_counts == [] + assert [state.assign_sources for state in (nm._state_0, nm._state_1)] == [ + [], + [], + ] + assert nm._active_trajectory is None + assert trajectory._released + _assert_tape_reset_then_trajectory_release(nm.events) + + +def test_default_dynamics_commits_manager_trajectory_once_after_tape_closes( + monkeypatch, +) -> None: + """A public commit is the sole detached publication of live state.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ), + ) + + nm = manager.physics.newton_manager + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert nm.commits == [trajectory] + assert nm.events[-1] == "commit" + assert nm.commit_assignment_counts == [(1, 1)] + assert [state.assign_sources for state in (nm._state_0, nm._state_1)] == [ + [trajectory.final_state], + [trajectory.final_state], + ] + + +@pytest.mark.parametrize("failure_site", ("action", "trajectory_step")) +def test_failed_manager_trajectory_forward_resets_and_releases_without_commit( + monkeypatch, failure_site: str +) -> None: + """A failed taped forward releases its manager lease without publishing it.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + + if failure_site == "action": + + def _apply_action(_action: torch.Tensor, *_args: Any) -> None: + assert warp.tape_active + nm.events.append("action.error") + raise RuntimeError("injected action failure") + + error_match = "injected action failure" + else: + nm.raise_on_trajectory_step = True + + def _apply_action(_action: torch.Tensor, *_args: Any) -> None: + assert warp.tape_active + nm.events.append("action") + + error_match = "injected trajectory-step failure" + + with pytest.raises(RuntimeError, match=error_match): + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=_apply_action, + step_mode="dynamics", + ), + ) + + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert nm.commits == [] + assert nm.commit_assignment_counts == [] + assert [state.assign_sources for state in (nm._state_0, nm._state_1)] == [ + [], + [], + ] + assert nm._active_trajectory is None + assert trajectory._released + _assert_tape_reset_then_trajectory_release(nm.events) + + +def test_backward_resets_tape_then_releases_manager_trajectory(monkeypatch) -> None: + """A grad-tracked trajectory resets its tape before releasing after backward.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + action = torch.zeros(1, requires_grad=True) + + outputs = NewtonStepFunc.apply( + action, + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ), + ) + outputs[0].sum().backward() + + assert action.grad is not None + _assert_backward_captures_gradient_then_releases(nm.events) + _assert_tape_reset_then_trajectory_release(nm.events) + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert nm.commits == [trajectory] + assert nm._active_trajectory is None + assert trajectory._released + + +def test_backward_exception_resets_tape_then_releases_manager_trajectory( + monkeypatch, +) -> None: + """A tape-backward failure cannot leave a manager trajectory leased.""" + warp = _RecordingWarp() + warp.raise_on_tape_backward = True + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + action = torch.zeros(1, requires_grad=True) + + outputs = NewtonStepFunc.apply( + action, + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ), + ) + + with pytest.raises(RuntimeError, match="injected tape backward failure"): + outputs[0].sum().backward() + + assert nm.events.count("tape.backward") == 1 + assert "action-gradient.capture" not in nm.events + _assert_tape_reset_then_trajectory_release(nm.events) + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert nm.commits == [trajectory] + assert nm._active_trajectory is None + assert trajectory._released + + +def test_obs_reward_failure_releases_manager_trajectory_before_fresh_forward( + monkeypatch, +) -> None: + """An output-read error rolls back its lease so the next trajectory starts.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + failing_state = _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ) + + def _raise_from_outputs(_final_state: _FakeState) -> dict[str, Any]: + assert warp.tape_active + nm.events.append("outputs.error") + raise RuntimeError("injected output-read failure") + + failing_state["obs_reward_fn"] = _raise_from_outputs + with pytest.raises(RuntimeError, match="injected output-read failure"): + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), failing_state) + + failure_events = list(nm.events) + assert "trajectory.step" in failure_events + assert "outputs.error" in failure_events + assert failure_events.index("trajectory.step") < failure_events.index( + "outputs.error" + ) + _assert_tape_reset_then_trajectory_release(failure_events) + assert len(nm.trajectories) == 1 + failed_trajectory = nm.trajectories[0] + assert failed_trajectory.step_calls == 1 + assert nm.commits == [] + assert nm.commit_assignment_counts == [] + assert [state.assign_sources for state in (nm._state_0, nm._state_1)] == [ + [], + [], + ] + assert nm._active_trajectory is None + assert failed_trajectory._released + + with torch.no_grad(): + outputs = NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ), + ) + + assert len(outputs) == 4 + assert len(nm.trajectories) == 2 + fresh_trajectory = nm.trajectories[1] + assert fresh_trajectory is not failed_trajectory + assert nm.commits == [fresh_trajectory] + assert nm._active_trajectory is None + assert fresh_trajectory._released + + +def test_no_grad_forward_resets_tape_then_releases_manager_trajectory( + monkeypatch, +) -> None: + """A non-grad forward cannot retain a trajectory lease for backward.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + + with torch.no_grad(): + outputs = NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ), + ) + + assert len(outputs) == 4 + assert not outputs[0].requires_grad + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert nm.commits == [trajectory] + assert nm._active_trajectory is None + assert trajectory._released + assert "tape.backward" not in nm.events + _assert_tape_reset_then_trajectory_release(nm.events) + + +def test_legacy_dynamics_step_fn_is_rejected_before_opening_tape(monkeypatch) -> None: + """An untrusted callback cannot silently bypass default solver dynamics.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + legacy_calls: list[None] = [] + + def _legacy_step() -> _FakeState: + legacy_calls.append(None) + return _FakeState("legacy-dynamics-final") + + with pytest.raises(ValueError, match=r"step_fn.*kinematics"): + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + step_fn=_legacy_step, + ), + ) + + assert legacy_calls == [] + assert warp.events == [] + + +def test_missing_step_mode_with_step_fn_is_rejected_before_opening_tape( + monkeypatch, +) -> None: + """Historical implicit-FK dictionaries cannot bypass solver dynamics.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + + with pytest.raises(ValueError, match=r"step_mode.*kinematics"): + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_fn=lambda: _FakeState("implicit-legacy-final"), + ), + ) + + assert warp.events == [] + + +def test_bridge_rejects_invalid_step_mode_before_opening_tape(monkeypatch) -> None: + """Direct bridge callers cannot open a tape for an unsupported mode.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + + with pytest.raises(ValueError, match=r"step_mode.*dynamics.*kinematics"): + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="unsupported", + ), + ) + + assert warp.events == [] + assert manager.physics.newton_manager.trajectory_requests == [] + assert manager.steppers == [] + + +def test_explicit_kinematics_step_fn_remains_a_supported_bridge_route( + monkeypatch, +) -> None: + """The deliberate kinematics escape hatch does not request a trajectory.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + kinematic_calls: list[None] = [] + final_state = _FakeState("kinematic-final") + + def _kinematic_step() -> _FakeState: + kinematic_calls.append(None) + return final_state + + outputs = NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="kinematics", + step_fn=_kinematic_step, + ), + ) + + assert len(outputs) == 4 + assert kinematic_calls == [None] + assert manager.physics.newton_manager.trajectory_requests == [] + + +def test_environment_sim_state_marks_default_and_explicit_kinematics_routes() -> None: + """The bridge can distinguish an explicit FK request from legacy bypasses.""" + dynamics_env, _ = _route_env(SimpleNamespace()) + kinematics_env, _ = _route_env(SimpleNamespace(), mode="kinematics") + kinematics_env._make_kinematic_step_fn = lambda: (lambda: _FakeState("fk")) + + dynamics_state = dynamics_env._build_sim_state_dict(torch.zeros(1)) + kinematics_state = kinematics_env._build_sim_state_dict(torch.zeros(1)) + + assert dynamics_state["step_mode"] == "dynamics" + assert kinematics_state["step_mode"] == "kinematics" + + +def test_environment_dynamics_hook_receives_local_control_with_migration_api() -> None: + """The default environment wrapper calls only the v1 dynamics hook.""" + env, _ = _route_env(SimpleNamespace()) + dynamics_calls: list[tuple[object, object, object]] = [] + legacy_calls: list[tuple[object, object]] = [] + action = object() + control = object() + + def _dynamics_action( + action_wp: object, trajectory_control: object, tape: object + ) -> None: + dynamics_calls.append((action_wp, trajectory_control, tape)) + + def _legacy_action(action_wp: object, tape: object) -> None: + legacy_calls.append((action_wp, tape)) + + env._apply_dynamics_action_kernel = _dynamics_action + env._apply_action_kernel = _legacy_action + sim_state = env._build_sim_state_dict(torch.zeros(1)) + sim_state["action_to_control_kernel"](action, control, "kernel-argument") + + assert dynamics_calls == [(action, control, None)] + assert legacy_calls == [] + + +def test_environment_dynamics_hook_observes_only_its_active_tape( + monkeypatch, +) -> None: + """The bridge binds the tape through a per-step wrapper closure.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + env, _ = _route_env(manager) + monkeypatch.setattr(diff_bridge, "wp", warp) + observed_tapes: list[object | None] = [] + + def _dynamics_action( + _action_wp: object, + _trajectory_control: object, + tape: object | None, + ) -> None: + observed_tapes.append(tape) + + env._apply_dynamics_action_kernel = _dynamics_action + sim_state = env._build_sim_state_dict(torch.zeros(1)) + + with torch.no_grad(): + NewtonStepFunc.apply(torch.zeros(1), sim_state) + + assert len(observed_tapes) == 1 + assert isinstance(observed_tapes[0], _RecordingTape) + + sim_state["action_to_control_kernel"](object(), object()) + assert observed_tapes[-1] is None + + +def test_environment_rejects_legacy_dynamics_action_hook_with_migration_error( + monkeypatch, +) -> None: + """Default dynamics cannot silently keep the pre-local-control hook.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + env = object.__new__(DifferentiableEmbodiedEnv) + env.sim = manager + env.cfg = SimpleNamespace(sim_steps_per_control=_CONTROL_SUBSTEPS) + env._apply_dynamics_action_kernel = None + env._apply_action_kernel = lambda _action, tape: None + env._read_outputs = lambda _state: { + "obs": torch.zeros(1, 1), + "reward": torch.zeros(1), + "terminated": torch.zeros(1, dtype=torch.bool), + "truncated": torch.zeros(1, dtype=torch.bool), + "_order": ("obs", "reward", "terminated", "truncated"), + "_grad_track": {}, + } + + with pytest.raises( + NotImplementedError, match=r"legacy.*_apply_dynamics_action_kernel" + ): + sim_state = env._build_sim_state_dict(torch.zeros(1, requires_grad=True)) + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + assert warp.events == [] + + +def test_environment_kinematics_hook_keeps_its_strict_legacy_signature( + monkeypatch, +) -> None: + """FK-only bridge execution receives action and tape, never local control.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + env, _ = _route_env(manager, mode="kinematics") + monkeypatch.setattr(diff_bridge, "wp", warp) + calls: list[tuple[torch.Tensor, object]] = [] + final_state = _FakeState("kinematic-final") + + def _kinematic_action(action_wp: torch.Tensor, tape: object) -> None: + calls.append((action_wp, tape)) + + env._apply_action_kernel = _kinematic_action + env._make_kinematic_step_fn = lambda: (lambda: final_state) + action = torch.zeros(1, requires_grad=True) + sim_state = env._build_sim_state_dict(action) + outputs = NewtonStepFunc.apply(action, sim_state) + + assert len(outputs) == 4 + assert len(calls) == 1 + assert torch.equal(calls[0][0], action) + assert isinstance(calls[0][1], _RecordingTape) + assert manager.physics.newton_manager.trajectory_requests == [] + + +def test_grad_terminal_step_defers_reset_until_after_backward(monkeypatch) -> None: + """A terminal grad step must return before touching fenced live state.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + env, _ = _route_env(manager) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + reset_calls: list[torch.Tensor] = [] + + def _terminal_outputs(_final_state: object) -> dict[str, Any]: + return { + "obs": torch.full((1, 1), 7.0), + "reward": torch.full((1,), 3.0), + "terminated": torch.ones(1, dtype=torch.bool), + "truncated": torch.zeros(1, dtype=torch.bool), + "_order": ("obs", "reward", "terminated", "truncated"), + "_grad_track": {}, + } + + def _reset(*, options: dict[str, Any]): + if nm._active_trajectory is not None: + raise RuntimeError("reset crossed an active Newton trajectory fence") + reset_ids = torch.as_tensor(options["reset_ids"]).clone() + reset_calls.append(reset_ids) + return torch.full((1, 1), -1.0), {} + + env._read_outputs = _terminal_outputs + env.reset = _reset + action = torch.zeros(1, requires_grad=True) + + obs, reward, terminated, truncated, info = env.step(action) + + assert torch.equal(obs.detach(), torch.full((1, 1), 7.0)) + assert terminated.tolist() == [True] + assert truncated.tolist() == [False] + assert reset_calls == [] + assert info["requires_reset_after_backward"] is True + assert torch.equal(info["deferred_reset_ids"], torch.tensor([0])) + assert nm._active_trajectory is not None + + reward.sum().backward() + + assert action.grad is not None + assert nm._active_trajectory is None + env.reset(options={"reset_ids": info["deferred_reset_ids"]}) + assert len(reset_calls) == 1 + assert torch.equal(reset_calls[0], torch.tensor([0])) + + +def test_no_grad_terminal_step_keeps_synchronous_auto_reset(monkeypatch) -> None: + """A terminal no-grad step may reset after its tape is released.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + env, _ = _route_env(manager) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + reset_calls: list[torch.Tensor] = [] + + env._read_outputs = lambda _state: { + "obs": torch.full((1, 1), 7.0), + "reward": torch.full((1,), 3.0), + "terminated": torch.ones(1, dtype=torch.bool), + "truncated": torch.zeros(1, dtype=torch.bool), + "_order": ("obs", "reward", "terminated", "truncated"), + "_grad_track": {}, + } + + def _reset(*, options: dict[str, Any]): + assert nm._active_trajectory is None + reset_ids = torch.as_tensor(options["reset_ids"]).clone() + reset_calls.append(reset_ids) + return torch.full((1, 1), -1.0), {} + + env.reset = _reset + with torch.no_grad(): + obs, reward, terminated, truncated, info = env.step( + torch.zeros(1, requires_grad=True) + ) + + assert torch.equal(obs, torch.full((1, 1), -1.0)) + assert not reward.requires_grad + assert terminated.tolist() == [True] + assert truncated.tolist() == [False] + assert len(reset_calls) == 1 + assert torch.equal(reset_calls[0], torch.tensor([0])) + assert "deferred_reset_ids" not in info + assert "requires_reset_after_backward" not in info + + +def test_default_dynamics_route_uses_manager_trajectory_without_bypass(monkeypatch): + """Default state construction delegates one control step to Newton.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + env, final_states = _route_env(manager) + monkeypatch.setattr(diff_bridge, "wp", warp) + + sim_state = env._build_sim_state_dict(torch.zeros(1, requires_grad=True)) + outputs = NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + nm = manager.physics.newton_manager + assert sim_state["step_mode"] == "dynamics" + assert "step_fn" not in sim_state + assert len(outputs) == 4 + assert len(nm.trajectories) == 1 + assert nm.trajectories[0].total_solver_steps == _CONTROL_SUBSTEPS + assert final_states[0].value == _CONTROL_SUBSTEPS + + +def test_dynamics_bridge_keeps_an_odd_continuous_horizon_in_one_trajectory( + monkeypatch, +): + """A continuous odd horizon is one lease-owning manager trajectory.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + env, final_states = _route_env(manager, control_substeps=5) + monkeypatch.setattr(diff_bridge, "wp", warp) + + sim_state = env._build_sim_state_dict(torch.zeros(1, requires_grad=True)) + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + nm = manager.physics.newton_manager + assert [state.value for state in final_states] == [5] + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert trajectory.total_solver_steps == 5 + assert nm._state_0.value == 5 + assert nm._state_1.value == 5 + assert nm._state_0.assign_sources == [trajectory.final_state] + assert nm._state_1.assign_sources == [trajectory.final_state] + assert len({id(state) for state in trajectory.states}) == 6 + assert all(state not in {nm._state_0, nm._state_1} for state in trajectory.states) + assert len({id(contact) for contact in trajectory.contacts}) == 5 + + +def test_dynamics_bridge_rejects_a_second_outstanding_manager_trajectory( + monkeypatch, +) -> None: + """A second grad forward requires release of the first trajectory lease.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + sim_state = _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ) + + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + with pytest.raises(RuntimeError, match=r"trajectory.*active.*release"): + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + +def test_dynamics_bridge_multiplies_control_and_newton_substeps(monkeypatch): + """One control step preserves EmbodiChain and Newton time semantics.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp, num_substeps=3) + env, _ = _route_env(manager, control_substeps=2) + monkeypatch.setattr(diff_bridge, "wp", warp) + + sim_state = env._build_sim_state_dict(torch.zeros(1, requires_grad=True)) + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + nm = manager.physics.newton_manager + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert trajectory.physics_steps == 2 + assert trajectory.physics_dt == nm.physics_dt + assert trajectory.total_solver_steps == 6 + + +def test_differentiable_step_uses_manager_owned_trajectory_and_local_control( + monkeypatch, +): + """The low-level helper also delegates state publication to Newton.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + received: list[tuple[Any, ...]] = [] + + result = differentiable_step( + manager, + apply_control_fn=lambda *args: received.append(args), + substeps=_CONTROL_SUBSTEPS, + ) + nm = manager.physics.newton_manager + + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert any(trajectory.control in args for args in received) + assert result["trajectory"] is trajectory + assert nm.commits == [trajectory] + assert nm.commit_assignment_counts == [(1, 1)] + + +def test_differentiable_step_rejects_substeps_not_divisible_by_newton_substeps( + monkeypatch, +) -> None: + """A low-level solver horizon must map to whole Newton physics steps.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp, num_substeps=2) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + control_calls: list[tuple[Any, ...]] = [] + + with pytest.raises(ValueError, match=r"substeps.*divisible.*num_substeps"): + differentiable_step( + manager, + apply_control_fn=lambda *args: control_calls.append(args), + substeps=3, + ) + + assert control_calls == [] + assert nm.trajectory_requests == [] + assert manager.steppers == [] + assert warp.events == [] + + +@pytest.mark.parametrize("substeps", (0, -1)) +def test_differentiable_step_rejects_nonpositive_substeps(substeps: int) -> None: + """The public helper rejects an invalid empty solver trajectory.""" + manager = _TrajectorySimulationManager(_RecordingWarp()) + + with pytest.raises(ValueError, match=r"positive"): + differentiable_step( + manager, + apply_control_fn=lambda *_args: None, + substeps=substeps, + ) + + +def test_cpu_newton_manager_trajectory_retains_local_control_gradient_and_fd(tmp_path): + """The real bridge keeps a local control trajectory across two steps.""" + newton = pytest.importorskip("newton") + pytest.importorskip("dexsim.engine.newton_physics") + from dexsim.engine.newton_physics import ( + NewtonCfg, + NewtonCollisionPipelineCfg, + NewtonManager, + SemiImplicitSolverCfg, + ) + + assert hasattr( + NewtonManager, "create_differentiable_trajectory" + ), "NewtonManager must publish create_differentiable_trajectory() first." + + previous_kernel_cache_dir = wp.config.kernel_cache_dir + previous_verify_access = wp.config.verify_autograd_array_access + nm = None + wp.config.kernel_cache_dir = str(tmp_path / "warp_cache") + wp.config.verify_autograd_array_access = True + try: + cfg = NewtonCfg() + cfg.device = "cpu" + cfg.dt = 1.0 / 60.0 + cfg.num_substeps = 2 + cfg.requires_grad = True + cfg.use_cuda_graph = False + cfg.solver_cfg = SemiImplicitSolverCfg() + cfg.collision_pipeline_cfg = NewtonCollisionPipelineCfg( + broad_phase="explicit", + requires_grad=True, + ) + nm = NewtonManager(cfg) + shape_cfg = newton.ModelBuilder.ShapeConfig( + ke=1.0e4, + kd=1.0e1, + kf=0.0, + mu=0.0, + ) + body_id = nm._builder.add_body( + xform=wp.transform(wp.vec3(0.0, 0.0, 0.5), wp.quat_identity()), + mass=1.0, + label="embodichain_manager_trajectory_gradient_ball", + ) + nm._builder.add_shape_sphere(body=body_id, radius=0.1, cfg=shape_cfg) + nm._builder.add_ground_plane(cfg=shape_cfg) + nm.start_simulation() + assert nm._model.joint_count == 1 + + manager = _RealBridgeManager(nm) + initial_state = nm._model.state() + initial_state.assign(nm._state_0) + target = wp.vec3(0.5, 0.0, 0.5) + + def _restore_initial_state() -> None: + nm._state_0.assign(initial_state) + nm._state_1.assign(initial_state) + + def _run( + action_value: float, *, requires_grad: bool + ) -> tuple[torch.Tensor, torch.Tensor, list[Any]]: + loss_wp = wp.zeros( + 1, + dtype=wp.float32, + device=nm._state_0.body_q.device, + requires_grad=True, + ) + local_controls: list[Any] = [] + + def _apply_control(action_wp: Any, *args: Any) -> None: + assert len(args) == 1, "Bridge must pass exactly one local control." + control = args[0] + assert control.joint_f is not None + local_controls.append(control) + wp.launch( + _write_bridge_joint_force_kernel, + dim=1, + inputs=[action_wp, control.joint_f], + device=control.joint_f.device, + ) + + def _read_reward(final_state: Any) -> dict[str, Any]: + loss_wp.zero_() + wp.launch( + _bridge_terminal_loss_kernel, + dim=1, + inputs=[final_state.body_q, body_id, target, loss_wp], + device=final_state.body_q.device, + ) + return { + "reward": wp.to_torch(loss_wp), + "_order": ("reward",), + "_grad_track": {"reward": loss_wp}, + } + + action = torch.tensor( + [action_value], dtype=torch.float32, requires_grad=requires_grad + ) + sim_state = { + "manager": manager, + "step_mode": "dynamics", + "substeps": 2, + "physics_dt": cfg.dt, + "action_to_control_kernel": _apply_control, + "kernel_args": (), + "obs_reward_fn": _read_reward, + } + return NewtonStepFunc.apply(action, sim_state)[0], action, local_controls + + reward, action, local_controls = _run(1.0, requires_grad=True) + reward.backward() + + assert len(local_controls) == 1 + assert action.grad is not None + analytic_gradient = float(action.grad[0]) + assert np.isfinite(analytic_gradient) + assert not np.isclose(analytic_gradient, 0.0) + first_final_state = nm._state_0.body_q.numpy().copy() + assert np.allclose( + nm._state_0.body_q.numpy(), nm._state_1.body_q.numpy(), atol=1.0e-6 + ) + + continuation_reward, continuation_action, continuation_controls = _run( + 1.0, requires_grad=True + ) + continuation_reward.backward() + assert len(continuation_controls) == 1 + assert continuation_action.grad is not None + assert np.isfinite(continuation_action.grad).all() + assert not np.allclose(first_final_state, nm._state_0.body_q.numpy()) + + def _reward_value(action_value: float) -> float: + _restore_initial_state() + value, _action, controls = _run(action_value, requires_grad=False) + assert len(controls) == 1 + return float(value.detach()) + + epsilon = 1.0e-3 + finite_difference_gradient = ( + _reward_value(1.0 + epsilon) - _reward_value(1.0 - epsilon) + ) / (2.0 * epsilon) + assert np.isclose( + analytic_gradient, + finite_difference_gradient, + rtol=2.0e-2, + atol=1.0e-4, + ) + finally: + if nm is not None: + nm.clear() + wp.config.verify_autograd_array_access = previous_verify_access + wp.config.kernel_cache_dir = previous_kernel_cache_dir + + +def test_dynamics_environment_does_not_expose_generic_step_helper(): + """Only the low-level bridge may accept an arbitrary dynamics callback.""" + assert "_make_step_fn" not in DifferentiableEmbodiedEnv.__dict__ + + +def test_kinematics_route_uses_only_named_kinematic_hook(): + """FK stepping is selected only through the explicit kinematics mode.""" + manager = SimpleNamespace() + env, _ = _route_env(manager, mode="kinematics") + expected_state = object() + kinematic_calls: list[None] = [] + + def _kinematic_step() -> object: + kinematic_calls.append(None) + return expected_state + + def _generic_step_fn() -> object: + raise AssertionError("The generic step helper must not route kinematics.") + + env._make_kinematic_step_fn = lambda: _kinematic_step + env._make_step_fn = _generic_step_fn + + sim_state = env._build_sim_state_dict(torch.zeros(1)) + + assert sim_state["step_mode"] == "kinematics" + assert sim_state["step_fn"]() is expected_state + assert kinematic_calls == [None] + + +def test_kinematics_route_requires_named_hook(): + """Kinematics mode rejects environments that do not define its hook.""" + manager = SimpleNamespace() + env, _ = _route_env(manager, mode="kinematics") + + with pytest.raises( + NotImplementedError, match=r"kinematics.*_make_kinematic_step_fn" + ): + env._build_sim_state_dict(torch.zeros(1)) + + +def test_invalid_differentiable_step_mode_raises_clear_error(): + """Unsupported stepping modes fail before creating a bridge callback.""" + manager = SimpleNamespace() + env, _ = _route_env(manager, mode="unsupported") + + with pytest.raises( + ValueError, match=r"differentiable_step_mode.*dynamics.*kinematics" + ): + env._build_sim_state_dict(torch.zeros(1)) + + def test_construct_without_requires_grad_raises(): with pytest.raises(Exception, match=r"requires_grad"): DifferentiableEmbodiedEnv(_diff_env_cfg(requires_grad=False)) @@ -71,6 +1496,192 @@ def _import_franka_env(): return FrankaReachApgEnv +def test_franka_kinematics_build_snapshots_live_primal_before_bridge( + monkeypatch, +) -> None: + """Franka must detach taped FK inputs before the parent opens a tape.""" + from embodichain.lab.gym.envs.tasks.special import franka_reach_apg + + env = object.__new__(franka_reach_apg.FrankaReachApgEnv) + live_joint_q = object() + snapshot_joint_q = object() + fresh_fk_state = object() + events: list[str] = [] + env.sim = SimpleNamespace( + physics=SimpleNamespace( + newton_manager=SimpleNamespace( + _state_0=SimpleNamespace(joint_q=live_joint_q), + _model=SimpleNamespace( + state=lambda: (events.append("state"), fresh_fk_state)[1] + ), + ) + ) + ) + + def _clone(array: object) -> object: + assert array is live_joint_q + events.append("clone") + return snapshot_joint_q + + def _parent_build(_self: object, _action: torch.Tensor) -> dict[str, Any]: + events.append("parent") + assert env._current_joint_q_snapshot is snapshot_joint_q + assert env._fk_state is fresh_fk_state + return {"prepared": True} + + monkeypatch.setattr(franka_reach_apg.wp, "clone", _clone) + monkeypatch.setattr( + DifferentiableEmbodiedEnv, + "_build_sim_state_dict", + _parent_build, + ) + + result = env._build_sim_state_dict(torch.zeros(1, 7)) + + assert result == {"prepared": True} + assert events == ["clone", "state", "parent"] + + +def test_franka_action_kernel_reads_snapshot_instead_of_live_state(monkeypatch) -> None: + """The recorded action kernel must not capture mutable manager state.""" + from embodichain.lab.gym.envs.tasks.special import franka_reach_apg + + env = object.__new__(franka_reach_apg.FrankaReachApgEnv) + live_joint_q = object() + snapshot_joint_q = object() + target_joint_q = object() + action_wp = object() + launch_inputs: list[object] = [] + env.sim = SimpleNamespace( + num_envs=1, + physics=SimpleNamespace( + newton_manager=SimpleNamespace( + _state_0=SimpleNamespace(joint_q=live_joint_q) + ) + ), + ) + env._current_joint_q_snapshot = snapshot_joint_q + env._n_joints_per_env = 9 + env._wp_device = "cpu" + env._limit_lo_wp = object() + env._limit_hi_wp = object() + env._action_scale = 0.2 + + monkeypatch.setattr( + franka_reach_apg.wp, + "zeros", + lambda *_args, **_kwargs: target_joint_q, + ) + + def _launch(*_args: Any, inputs: list[object], **_kwargs: Any) -> None: + launch_inputs.extend(inputs) + + monkeypatch.setattr(franka_reach_apg.wp, "launch", _launch) + + env._apply_action_kernel(action_wp, tape=object()) + + assert launch_inputs[0] is action_wp + assert launch_inputs[1] is snapshot_joint_q + assert launch_inputs[1] is not live_joint_q + assert launch_inputs[2] is target_joint_q + + +def test_franka_snapshot_keeps_gradient_after_live_state_mutation_and_matches_fd( + monkeypatch, + tmp_path, +) -> None: + """Detached FK input survives live writes before backward under strict mode.""" + from embodichain.lab.gym.envs.tasks.special import franka_reach_apg + + env = object.__new__(franka_reach_apg.FrankaReachApgEnv) + device = "cpu" + live_joint_q = wp.zeros(7, dtype=wp.float32, device=device) + env.sim = SimpleNamespace( + num_envs=1, + physics=SimpleNamespace( + newton_manager=SimpleNamespace( + _state_0=SimpleNamespace(joint_q=live_joint_q), + _model=SimpleNamespace(state=lambda: object()), + ) + ), + ) + env._wp_device = device + env._n_joints_per_env = 7 + env._limit_lo_wp = wp.array( + np.full(7, -10.0, dtype=np.float32), + dtype=wp.float32, + device=device, + ) + env._limit_hi_wp = wp.array( + np.full(7, 10.0, dtype=np.float32), + dtype=wp.float32, + device=device, + ) + env._action_scale = 0.2 + monkeypatch.setattr( + DifferentiableEmbodiedEnv, + "_build_sim_state_dict", + lambda _self, _action: {}, + ) + env._build_sim_state_dict(torch.zeros(1, 7)) + + previous_verify_access = wp.config.verify_autograd_array_access + previous_kernel_cache_dir = wp.config.kernel_cache_dir + wp.config.verify_autograd_array_access = True + wp.config.kernel_cache_dir = str(tmp_path / "warp_cache") + tape = wp.Tape() + try: + action_wp = wp.array( + np.zeros(7, dtype=np.float32), + dtype=wp.float32, + device=device, + requires_grad=True, + ) + with tape: + env._apply_action_kernel(action_wp, tape=tape) + analytic_output = env._new_joint_q + + wp.copy( + live_joint_q, + wp.array( + np.full(7, 5.0, dtype=np.float32), + dtype=wp.float32, + device=device, + ), + ) + tape.backward(grads={analytic_output: wp.ones_like(analytic_output)}) + analytic_gradient = action_wp.grad.numpy().copy() + + assert np.isfinite(analytic_gradient).all() + assert np.all(np.abs(analytic_gradient) > 0.0) + + def _loss(action_value: float) -> float: + values = np.zeros(7, dtype=np.float32) + values[0] = action_value + finite_difference_action = wp.array( + values, + dtype=wp.float32, + device=device, + ) + env._apply_action_kernel(finite_difference_action, tape=object()) + return float(env._new_joint_q.numpy().sum()) + + epsilon = 1.0e-3 + finite_difference_gradient = (_loss(epsilon) - _loss(-epsilon)) / ( + 2.0 * epsilon + ) + assert np.isclose( + analytic_gradient[0], + finite_difference_gradient, + rtol=1.0e-4, + atol=1.0e-5, + ) + finally: + tape.reset() + wp.config.verify_autograd_array_access = previous_verify_access + wp.config.kernel_cache_dir = previous_kernel_cache_dir + + def test_franka_apg_smoke_backward(): """Verify reward is autograd-tracked and action.grad flows back.""" try: diff --git a/tests/gym/envs/test_embodied_env.py b/tests/gym/envs/test_embodied_env.py index 7bad9b549..e3709d768 100644 --- a/tests/gym/envs/test_embodied_env.py +++ b/tests/gym/envs/test_embodied_env.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os import torch import pytest diff --git a/tests/gym/utils/test_gym_utils.py b/tests/gym/utils/test_gym_utils.py index 803de5c54..92876b6b5 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -17,18 +17,19 @@ from __future__ import annotations +import argparse + import pytest import torch -import argparse from tensordict import TensorDict from embodichain.lab.gym.utils.gym_utils import ( + DEFAULT_MANAGER_MODULES, add_env_launcher_args_to_parser, + config_to_cfg, init_rollout_buffer_from_config, merge_args_with_gym_config, - config_to_cfg, - DEFAULT_MANAGER_MODULES, ) from embodichain.utils.utility import load_config, save_config @@ -278,92 +279,134 @@ def test_extra_observation_with_nested_name(self): assert "value" in buffer["obs"]["custom"]["group1"] assert buffer["obs"]["custom"]["group1"]["value"].shape == (4, 100, 4) - def test_sensor_and_extra_obs_together(self): - """Test that both sensors and extra observations work together.""" - config = { - "sensor": [ - { - "uid": "camera", - "width": 320, - "height": 240, - "enable_mask": True, - } - ], - "env": { - "observations": { - "extra_vec": { - "mode": "add", - "extra": {"shape": [10]}, - } - } - }, - } - buffer = init_rollout_buffer_from_config( - config=config, - max_episode_steps=100, - batch_size=4, - state_dim=7, - device="cpu", - ) - - # Check sensor is present - assert "sensor" in buffer["obs"] - assert "camera" in buffer["obs"]["sensor"] - assert buffer["obs"]["sensor"]["camera"]["color"].shape == (4, 100, 240, 320, 4) - assert buffer["obs"]["sensor"]["camera"]["mask"].shape == (4, 100, 240, 320) - - # Check extra obs is present - assert "extra_vec" in buffer["obs"] - assert buffer["obs"]["extra_vec"].shape == (4, 100, 10) - - def test_different_batch_sizes(self): - """Test that batch_size correctly affects extra observations.""" - config = { - "sensor": [], - "env": { - "observations": { - "extra_data": { - "mode": "add", - "extra": {"shape": [5]}, - } +def test_merge_args_with_gym_config_overrides_max_episodes(): + """Test that CLI max_episodes overrides the gym config value.""" + args = argparse.Namespace( + num_envs=1, + device="cpu", + headless=False, + renderer="auto", + physics="default", + gpu_id=0, + arena_space=5.0, + max_episodes=12, + ) + gym_config = {"max_episodes": 3, "id": "Dummy-v0"} + + merged_config = merge_args_with_gym_config(args, gym_config) + + assert merged_config["max_episodes"] == 12 + assert gym_config["max_episodes"] == 3 + + +def test_merge_args_with_gym_config_keeps_default_max_episodes(): + """Test that gym config max_episodes is preserved when CLI omits it.""" + args = argparse.Namespace( + num_envs=1, + device="cpu", + headless=False, + renderer="auto", + physics="default", + gpu_id=0, + arena_space=5.0, + max_episodes=None, + ) + gym_config = {"max_episodes": 3, "id": "Dummy-v0"} + + merged_config = merge_args_with_gym_config(args, gym_config) + + assert merged_config["max_episodes"] == 3 + + +def test_sensor_and_extra_obs_together(): + """Test that both sensors and extra observations work together.""" + config = { + "sensor": [ + { + "uid": "camera", + "width": 320, + "height": 240, + "enable_mask": True, + } + ], + "env": { + "observations": { + "extra_vec": { + "mode": "add", + "extra": {"shape": [10]}, } - }, - } - - buffer = init_rollout_buffer_from_config( - config=config, - max_episode_steps=50, - batch_size=8, - state_dim=7, - device="cpu", - ) - - assert buffer["obs"]["extra_data"].shape == (8, 50, 5) - - def test_different_max_episode_steps(self): - """Test that max_episode_steps correctly affects extra observations.""" - config = { - "sensor": [], - "env": { - "observations": { - "extra_data": { - "mode": "add", - "extra": {"shape": [2]}, - } + } + }, + } + + buffer = init_rollout_buffer_from_config( + config=config, + max_episode_steps=100, + batch_size=4, + state_dim=7, + device="cpu", + ) + + # Check sensor is present + assert "sensor" in buffer["obs"] + assert "camera" in buffer["obs"]["sensor"] + assert buffer["obs"]["sensor"]["camera"]["color"].shape == (4, 100, 240, 320, 4) + assert buffer["obs"]["sensor"]["camera"]["mask"].shape == (4, 100, 240, 320) + + # Check extra obs is present + assert "extra_vec" in buffer["obs"] + assert buffer["obs"]["extra_vec"].shape == (4, 100, 10) + + +def test_different_batch_sizes(): + """Test that batch_size correctly affects extra observations.""" + config = { + "sensor": [], + "env": { + "observations": { + "extra_data": { + "mode": "add", + "extra": {"shape": [5]}, } - }, - } - - buffer = init_rollout_buffer_from_config( - config=config, - max_episode_steps=200, - batch_size=4, - state_dim=7, - device="cpu", - ) - - assert buffer["obs"]["extra_data"].shape == (4, 200, 2) + } + }, + } + + buffer = init_rollout_buffer_from_config( + config=config, + max_episode_steps=50, + batch_size=8, + state_dim=7, + device="cpu", + ) + + assert buffer["obs"]["extra_data"].shape == (8, 50, 5) + + +def test_different_max_episode_steps(): + """Test that max_episode_steps correctly affects extra observations.""" + config = { + "sensor": [], + "env": { + "observations": { + "extra_data": { + "mode": "add", + "extra": {"shape": [2]}, + } + } + }, + } + + buffer = init_rollout_buffer_from_config( + config=config, + max_episode_steps=200, + batch_size=4, + state_dim=7, + device="cpu", + ) + + assert buffer["obs"]["extra_data"].shape == (4, 200, 2) class TestConfigToCfgFromYaml: diff --git a/tests/learning/test_rl.py b/tests/learning/test_rl.py index bd4044f8f..d80fc654e 100644 --- a/tests/learning/test_rl.py +++ b/tests/learning/test_rl.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os import json import pytest diff --git a/tests/learning/test_rl_distributed.py b/tests/learning/test_rl_distributed.py index 113886a46..5a5821b9c 100644 --- a/tests/learning/test_rl_distributed.py +++ b/tests/learning/test_rl_distributed.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import json import os import subprocess diff --git a/tests/sim/atomic_actions/test_action_result_success.py b/tests/sim/atomic_actions/test_action_result_success.py new file mode 100644 index 000000000..61600fae8 --- /dev/null +++ b/tests/sim/atomic_actions/test_action_result_success.py @@ -0,0 +1,50 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Tests for ActionResult.success tensor and ActionCfg.motion_source/planner_type.""" + +from __future__ import annotations + +import pytest +import torch + +from embodichain.lab.sim.atomic_actions.core import ActionCfg, ActionResult, WorldState + + +class TestActionResultSuccess: + def test_success_all_tensor(self): + r = ActionResult( + success=torch.tensor([True, False]), + trajectory=torch.zeros(2, 0, 3), + next_state=WorldState(last_qpos=torch.zeros(2, 3)), + ) + assert r.success_all is False + + def test_bool_deprecation(self): + r = ActionResult( + success=torch.tensor([True, True]), + trajectory=torch.zeros(2, 0, 3), + next_state=WorldState(last_qpos=torch.zeros(2, 3)), + ) + with pytest.warns(DeprecationWarning): + assert bool(r) is True + + +class TestActionCfgMotionSource: + def test_defaults(self): + cfg = ActionCfg() + assert cfg.motion_source == "ik_interp" + assert cfg.planner_type is None diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index c593e0369..2ef969cd4 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -27,6 +27,10 @@ ) from embodichain.lab.sim.atomic_actions.core import ( ActionResult, + AtomicAction, + CoordinatedHeldObjectState, + CoordinatedPickmentTarget, + CoordinatedPlacementTarget, GraspTarget, HeldObjectState, HeldObjectPoseTarget, @@ -37,6 +41,10 @@ WorldState, ) from embodichain.lab.sim.atomic_actions.actions import ( + CoordinatedPickment, + CoordinatedPickmentCfg, + CoordinatedPlacement, + CoordinatedPlacementCfg, MoveEndEffector, MoveEndEffectorCfg, MoveJoints, @@ -47,12 +55,16 @@ PickUpCfg, Place, PlaceCfg, + Press, + PressCfg, ) NUM_ENVS = 2 ARM_DOF = 6 HAND_DOF = 2 TOTAL_DOF = ARM_DOF + HAND_DOF +DUAL_ARM_DOF = 12 +DUAL_TOTAL_DOF = DUAL_ARM_DOF + 2 * HAND_DOF def _make_mock_robot(): @@ -111,6 +123,57 @@ def _make_mock_motion_generator(): return mg +def _make_dual_arm_mock_robot(): + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = DUAL_TOTAL_DOF + + def get_qpos(name=None): + if name == "left_arm": + return torch.zeros(NUM_ENVS, ARM_DOF) + if name == "right_arm": + return torch.zeros(NUM_ENVS, ARM_DOF) + if name == "dual_arm": + return torch.zeros(NUM_ENVS, DUAL_ARM_DOF) + if name in ("left_hand", "right_hand"): + return torch.zeros(NUM_ENVS, HAND_DOF) + return torch.zeros(NUM_ENVS, DUAL_TOTAL_DOF) + + robot.get_qpos = get_qpos + + def get_joint_ids(name=None): + if name == "left_arm": + return list(range(0, ARM_DOF)) + if name == "right_arm": + return list(range(ARM_DOF, DUAL_ARM_DOF)) + if name == "dual_arm": + return list(range(DUAL_ARM_DOF)) + if name == "left_hand": + return list(range(DUAL_ARM_DOF, DUAL_ARM_DOF + HAND_DOF)) + if name == "right_hand": + return list(range(DUAL_ARM_DOF + HAND_DOF, DUAL_TOTAL_DOF)) + return list(range(DUAL_TOTAL_DOF)) + + robot.get_joint_ids = get_joint_ids + + def compute_ik(pose=None, name=None, joint_seed=None, qpos_seed=None): + seed = joint_seed if joint_seed is not None else qpos_seed + if seed is None: + seed = torch.zeros(NUM_ENVS, ARM_DOF) + offset = 0.1 if name == "left_arm" else 0.2 + return torch.ones(seed.shape[0], dtype=torch.bool), seed + offset + + robot.compute_ik = compute_ik + return robot + + +def _make_dual_arm_mock_motion_generator(): + mg = Mock() + mg.robot = _make_dual_arm_mock_robot() + mg.device = torch.device("cpu") + return mg + + def _hand_open(): return torch.zeros(HAND_DOF, dtype=torch.float32) @@ -143,7 +206,8 @@ def test_execute_returns_full_dof_trajectory(self): state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)) result = action.execute(EndEffectorPoseTarget(xpos=torch.eye(4)), state) assert isinstance(result, ActionResult) - assert result.success is True + assert result.success.all() + assert result.success.shape == (NUM_ENVS,) assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) # MoveEndEffector preserves held_object. assert result.next_state.held_object is None @@ -180,7 +244,8 @@ def interpolate(trajectory, interp_num, device): WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), ) - assert result.success is True + assert result.success.all() + assert result.success.shape == (NUM_ENVS,) assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) # Two waypoints -> two IK calls, in order. assert len(seen_poses) == 2 @@ -232,7 +297,8 @@ def interpolate(trajectory, interp_num, device): WorldState(last_qpos=last_qpos, held_object=held), ) - assert result.success is True + assert result.success.all() + assert result.success.shape == (NUM_ENVS,) assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) assert torch.allclose(result.trajectory[:, -1, :ARM_DOF], target_qpos) assert torch.allclose(result.trajectory[:, -1, ARM_DOF:], hand_qpos) @@ -256,7 +322,8 @@ def test_execute_with_named_qpos_resolves_cfg_target(self): NamedJointPositionTarget(name="home"), WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), ) - assert result.success is True + assert result.success.all() + assert result.success.shape == (NUM_ENVS,) assert torch.allclose( result.next_state.last_qpos[:, :ARM_DOF], torch.full((NUM_ENVS, ARM_DOF), 0.2), @@ -293,7 +360,8 @@ def interpolate(trajectory, interp_num, device): WorldState(last_qpos=last_qpos), ) - assert result.success is True + assert result.success.all() + assert result.success.shape == (NUM_ENVS,) assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) # start prepended to the two waypoints -> 3 keyframes keyframes = captured["keyframes"] @@ -366,12 +434,105 @@ def test_execute_populates_held_object_state(self): ): state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)) result = action.execute(GraspTarget(semantics=sem), state) - assert result.success is True + assert result.success.all() + assert result.success.shape == (NUM_ENVS,) assert result.trajectory.shape[0] == NUM_ENVS assert result.trajectory.shape[2] == TOTAL_DOF assert isinstance(result.next_state.held_object, HeldObjectState) assert result.next_state.held_object.semantics is sem + def test_execute_accepts_an_explicit_grasp_pose(self): + action = PickUp( + self.mg, + PickUpCfg( + hand_open_qpos=_hand_open(), + hand_close_qpos=_hand_close(), + sample_interval=12, + hand_interp_steps=4, + ), + ) + entity = Mock() + entity.get_local_pose.return_value = ( + torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1) + ) + affordance = Mock() + grasp_xpos = torch.eye(4) + + with patch( + "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", + return_value=torch.zeros(NUM_ENVS, 4, ARM_DOF), + ): + result = action.execute( + GraspTarget( + semantics=ObjectSemantics( + affordance=affordance, + geometry={}, + entity=entity, + ), + grasp_xpos=grasp_xpos, + ), + WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), + ) + + assert result.success.all() + assert result.next_state.held_object is not None + assert torch.allclose( + result.next_state.held_object.grasp_xpos, + grasp_xpos.unsqueeze(0).repeat(NUM_ENVS, 1, 1), + ) + affordance.get_valid_grasp_poses.assert_not_called() + + def test_execute_chooses_symmetric_grasp_variant_closest_to_start_pose(self): + cfg = PickUpCfg( + hand_open_qpos=_hand_open(), + hand_close_qpos=_hand_close(), + sample_interval=20, + hand_interp_steps=4, + ) + action = PickUp(self.mg, cfg) + compute_batch_ik = self.mg.robot.compute_batch_ik + self.mg.robot.compute_batch_ik = Mock(side_effect=compute_batch_ik) + + rz_pi_grasp = torch.eye(4) + rz_pi_grasp[:3, :3] = torch.diag(torch.tensor([-1.0, -1.0, 1.0])) + affordance = AntipodalAffordance() + affordance.get_valid_grasp_poses = Mock( + return_value=[ + (rz_pi_grasp.unsqueeze(0), torch.tensor([0.5])) for _ in range(NUM_ENVS) + ] + ) + + entity = Mock() + entity.get_local_pose = Mock( + return_value=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1) + ) + sem = ObjectSemantics( + affordance=affordance, + geometry={}, + label="mug", + entity=entity, + ) + + with patch( + "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", + side_effect=lambda trajectory, interp_num, device: torch.zeros( + NUM_ENVS, interp_num, ARM_DOF + ), + ): + state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)) + result = action.execute(GraspTarget(semantics=sem), state) + + assert result.success.all() + assert result.success.shape == (NUM_ENVS,) + assert isinstance(result.next_state.held_object, HeldObjectState) + expected_grasp = torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1) + assert torch.allclose(result.next_state.held_object.grasp_xpos, expected_grasp) + self.mg.robot.compute_batch_ik.assert_called_once() + ik_kwargs = self.mg.robot.compute_batch_ik.call_args.kwargs + assert ik_kwargs["pose"].shape == (NUM_ENVS, 1, 4, 4) + assert torch.allclose(ik_kwargs["pose"][:, 0], expected_grasp) + assert ik_kwargs["joint_seed"].shape == (NUM_ENVS, 1, ARM_DOF) + # --------------------------------------------------------------------------- # MoveHeldObject @@ -422,7 +583,8 @@ def test_preserves_held_object(self): result = action.execute( HeldObjectPoseTarget(object_target_pose=torch.eye(4)), state ) - assert result.success is True + assert result.success.all() + assert result.success.shape == (NUM_ENVS,) assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) assert result.next_state.held_object is held @@ -463,7 +625,8 @@ def test_execute_clears_held_object(self): ), ): result = action.execute(EndEffectorPoseTarget(xpos=torch.eye(4)), state) - assert result.success is True + assert result.success.all() + assert result.success.shape == (NUM_ENVS,) assert result.trajectory.shape[2] == TOTAL_DOF assert result.next_state.held_object is None @@ -515,7 +678,8 @@ def interpolate(trajectory, interp_num, device): ): result = action.execute(EndEffectorPoseTarget(xpos=multi_xpos), state) - assert result.success is True + assert result.success.all() + assert result.success.shape == (NUM_ENVS,) assert result.trajectory.shape[2] == TOTAL_DOF assert result.next_state.held_object is None # IK order: down phase (approach, pose0, pose1) then back phase (retract). @@ -535,3 +699,342 @@ def interpolate(trajectory, interp_num, device): ) # start prepended to the 3 down-phase IK solutions -> 4 keyframes. assert captured["down_keyframes"].shape == (NUM_ENVS, 4, ARM_DOF) + + def test_execute_preserves_release_pose_without_tcp_symmetry(self): + cfg = PlaceCfg( + hand_open_qpos=_hand_open(), + hand_close_qpos=_hand_close(), + sample_interval=20, + hand_interp_steps=4, + ) + action = Place(self.mg, cfg) + state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)) + rz_pi_pose = torch.eye(4) + rz_pi_pose[:3, :3] = torch.diag(torch.tensor([-1.0, -1.0, 1.0])) + seen_poses = [] + + def compute_ik(pose=None, name=None, joint_seed=None, **kwargs): + seen_poses.append(pose.clone()) + return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed.clone() + + def repeat_last_keyframe(trajectory, interp_num, device): + return trajectory[:, -1:, :].repeat(1, interp_num, 1) + + self.mg.robot.compute_ik = Mock(side_effect=compute_ik) + + with patch( + "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", + side_effect=repeat_last_keyframe, + ): + result = action.execute(EndEffectorPoseTarget(xpos=rz_pi_pose), state) + + assert result.success.all() + assert result.success.shape == (NUM_ENVS,) + assert len(seen_poses) == 3 + assert torch.allclose( + seen_poses[1], rz_pi_pose.unsqueeze(0).repeat(NUM_ENVS, 1, 1) + ) + + def test_execute_with_tcp_symmetry_selects_closest_release_variant(self): + cfg = PlaceCfg( + hand_open_qpos=_hand_open(), + hand_close_qpos=_hand_close(), + sample_interval=20, + hand_interp_steps=4, + lift_height=0.1, + ) + action = Place(self.mg, cfg) + state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)) + rz_pi_pose = torch.eye(4) + rz_pi_pose[:3, :3] = torch.diag(torch.tensor([-1.0, -1.0, 1.0])) + seen_poses = [] + + def compute_ik(pose=None, name=None, joint_seed=None, **kwargs): + seen_poses.append(pose.clone()) + return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed.clone() + + def repeat_last_keyframe(trajectory, interp_num, device): + return trajectory[:, -1:, :].repeat(1, interp_num, 1) + + self.mg.robot.compute_ik = Mock(side_effect=compute_ik) + + with patch( + "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", + side_effect=repeat_last_keyframe, + ): + result = action.execute( + EndEffectorPoseTarget( + xpos=rz_pi_pose, + tcp_symmetry="z_roll_180", + ), + state, + ) + + assert result.success.all() + assert result.success.shape == (NUM_ENVS,) + assert len(seen_poses) == 3 + expected_release = torch.eye(4) + expected_retract = torch.eye(4) + expected_retract[2, 3] += cfg.lift_height + assert torch.allclose( + seen_poses[0], expected_retract.unsqueeze(0).repeat(NUM_ENVS, 1, 1) + ) + assert torch.allclose( + seen_poses[1], expected_release.unsqueeze(0).repeat(NUM_ENVS, 1, 1) + ) + assert torch.allclose( + seen_poses[2], expected_retract.unsqueeze(0).repeat(NUM_ENVS, 1, 1) + ) + + +# --------------------------------------------------------------------------- +# Press +# --------------------------------------------------------------------------- + + +class TestPressAction: + def setup_method(self): + self.mg = _make_mock_motion_generator() + + def test_target_type_is_pose_target(self): + assert Press.TargetType is EndEffectorPoseTarget + + def test_default_name_is_explicit(self): + assert PressCfg(hand_close_qpos=_hand_close()).name == "press" + + def test_execute_closes_hand_and_preserves_held_object(self): + cfg = PressCfg( + hand_close_qpos=_hand_close(), + sample_interval=12, + hand_interp_steps=4, + ) + action = Press(self.mg, cfg) + sem = ObjectSemantics( + affordance=AntipodalAffordance(), geometry={}, label="mug" + ) + held = HeldObjectState( + semantics=sem, + object_to_eef=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1), + grasp_xpos=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1), + ) + start_hand_qpos = torch.full((NUM_ENVS, HAND_DOF), 0.01) + last_qpos = torch.cat([torch.zeros(NUM_ENVS, ARM_DOF), start_hand_qpos], dim=1) + state = WorldState(last_qpos=last_qpos, held_object=held) + + def interpolate(trajectory, interp_num, device): + return trajectory[:, -1:, :].repeat(1, interp_num, 1) + + with patch( + "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", + side_effect=interpolate, + ): + result = action.execute(EndEffectorPoseTarget(xpos=torch.eye(4)), state) + + assert result.success.all() + assert result.success.shape == (NUM_ENVS,) + assert result.trajectory.shape == (NUM_ENVS, 12, TOTAL_DOF) + expected_hand_qpos = _hand_close().unsqueeze(0).repeat(NUM_ENVS, 1) + assert torch.allclose(result.trajectory[:, -1, ARM_DOF:], expected_hand_qpos) + assert torch.allclose( + result.next_state.last_qpos[:, :ARM_DOF], + last_qpos[:, :ARM_DOF], + ) + assert result.next_state.held_object is held + + +# --------------------------------------------------------------------------- +# CoordinatedPickment +# --------------------------------------------------------------------------- + + +class TestCoordinatedPickmentAction: + def setup_method(self): + self.mg = _make_dual_arm_mock_motion_generator() + + def test_target_type_is_coordinated_pickment_target(self): + assert CoordinatedPickment.TargetType is CoordinatedPickmentTarget + assert CoordinatedPickment.__bases__ == (AtomicAction,) + + def test_execute_returns_full_dof_trajectory_and_dual_held_state(self): + cfg = CoordinatedPickmentCfg( + left_hand_open_qpos=_hand_open(), + left_hand_close_qpos=_hand_close(), + right_hand_open_qpos=_hand_open(), + right_hand_close_qpos=_hand_close(), + sample_interval=30, + hand_interp_steps=4, + hold_steps=2, + object_motion_keyframes=3, + ) + action = CoordinatedPickment(self.mg, cfg) + sem = ObjectSemantics( + affordance=AntipodalAffordance(), geometry={}, label="pencil" + ) + state = WorldState(last_qpos=torch.zeros(NUM_ENVS, DUAL_TOTAL_DOF)) + result = action.execute( + CoordinatedPickmentTarget( + object_target_pose=torch.eye(4), + object_semantics=sem, + left_object_to_eef=torch.eye(4), + right_object_to_eef=torch.eye(4), + object_initial_pose=torch.eye(4), + ), + state, + ) + assert result.success is True + assert result.trajectory.shape == (NUM_ENVS, 30, DUAL_TOTAL_DOF) + assert torch.allclose( + result.trajectory[:, -1, action.left_hand_joint_ids], + _hand_close().unsqueeze(0).repeat(NUM_ENVS, 1), + ) + assert torch.allclose( + result.trajectory[:, -1, action.right_hand_joint_ids], + _hand_close().unsqueeze(0).repeat(NUM_ENVS, 1), + ) + assert isinstance( + result.next_state.coordinated_held_object, + CoordinatedHeldObjectState, + ) + assert result.next_state.held_object is None + + +# --------------------------------------------------------------------------- +# CoordinatedPlacement +# --------------------------------------------------------------------------- + + +class TestCoordinatedPlacementAction: + def setup_method(self): + self.mg = _make_dual_arm_mock_motion_generator() + self.cfg = CoordinatedPlacementCfg( + placing_hand_open_qpos=_hand_open(), + placing_hand_close_qpos=_hand_close(), + support_hand_close_qpos=_hand_close(), + sample_interval=30, + hand_interp_steps=4, + hold_steps=3, + retreat_steps=5, + lift_height=0.08, + ) + self.action = CoordinatedPlacement(self.mg, cfg=self.cfg) + + def _make_target(self) -> CoordinatedPlacementTarget: + placing_pose = torch.eye(4) + placing_pose[0, 3] = 0.2 + support_pose = torch.eye(4) + support_pose[0, 3] = 0.2 + support_pose[2, 3] = -0.05 + + placing_object_to_eef = torch.eye(4) + placing_object_to_eef[2, 3] = 0.12 + support_object_to_eef = torch.eye(4) + support_object_to_eef[2, 3] = 0.10 + + placing_semantics = ObjectSemantics( + affordance=AntipodalAffordance(), geometry={}, label="placing" + ) + support_semantics = ObjectSemantics( + affordance=AntipodalAffordance(), geometry={}, label="support" + ) + return CoordinatedPlacementTarget( + placing_object_target_pose=placing_pose, + support_object_target_pose=support_pose, + placing_held_object=HeldObjectState( + semantics=placing_semantics, + object_to_eef=placing_object_to_eef, + grasp_xpos=torch.eye(4), + ), + support_held_object=HeldObjectState( + semantics=support_semantics, + object_to_eef=support_object_to_eef, + grasp_xpos=torch.eye(4), + ), + ) + + def test_target_type_is_coordinated_placement_target(self): + assert CoordinatedPlacement.TargetType is CoordinatedPlacementTarget + + def test_init_sets_dual_arm_and_hand_joint_ids(self): + assert self.action.dual_arm_joint_ids == list(range(DUAL_ARM_DOF)) + assert self.action.placing_arm_joint_ids == list(range(ARM_DOF)) + assert self.action.support_arm_joint_ids == list(range(ARM_DOF, DUAL_ARM_DOF)) + assert self.action.placing_hand_joint_ids == list( + range(DUAL_ARM_DOF, DUAL_ARM_DOF + HAND_DOF) + ) + assert self.action.support_hand_joint_ids == list( + range(DUAL_ARM_DOF + HAND_DOF, DUAL_TOTAL_DOF) + ) + assert self.action.joint_ids == list(range(DUAL_TOTAL_DOF)) + + def test_resolve_target_composes_object_and_tcp_poses(self): + target = self._make_target() + placing_xpos, support_xpos, release, held_state = self.action._resolve_target( + target + ) + assert placing_xpos.shape == (NUM_ENVS, 4, 4) + assert support_xpos.shape == (NUM_ENVS, 4, 4) + assert placing_xpos[0, 2, 3].item() == pytest.approx(0.12) + assert support_xpos[0, 2, 3].item() == pytest.approx(0.05) + assert release is True + assert held_state.semantics is target.support_held_object.semantics + assert held_state.object_to_eef.shape == (NUM_ENVS, 4, 4) + assert held_state.grasp_xpos.shape == (NUM_ENVS, 4, 4) + assert torch.allclose( + held_state.object_to_eef, + target.support_held_object.object_to_eef.unsqueeze(0).repeat( + NUM_ENVS, 1, 1 + ), + ) + + def test_segment_lengths_sum_to_sample_interval(self): + segments = self.action._compute_segment_lengths(self.cfg.release) + assert sum(segments.values()) == self.cfg.sample_interval + assert segments["approach"] >= 2 + assert segments["release"] == self.cfg.hand_interp_steps + assert segments["retreat"] == self.cfg.retreat_steps + + def test_execute_returns_full_dof_and_final_hand_states(self): + target = self._make_target() + state = WorldState(last_qpos=torch.zeros(NUM_ENVS, DUAL_TOTAL_DOF)) + + def interpolate(trajectory, interp_num, device): + weights = torch.linspace( + 0.0, + 1.0, + steps=interp_num, + dtype=trajectory.dtype, + device=trajectory.device, + ) + return torch.lerp( + trajectory[:, :1], + trajectory[:, -1:], + weights.view(1, -1, 1), + ) + + with patch( + "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", + side_effect=interpolate, + ): + result = self.action.execute(target, state) + + assert result.success is True + assert result.trajectory.shape == ( + NUM_ENVS, + self.cfg.sample_interval, + DUAL_TOTAL_DOF, + ) + assert torch.allclose( + result.trajectory[:, -1, self.action.placing_hand_joint_ids], + _hand_open().unsqueeze(0).repeat(NUM_ENVS, 1), + ) + assert torch.allclose( + result.trajectory[:, -1, self.action.support_hand_joint_ids], + _hand_close().unsqueeze(0).repeat(NUM_ENVS, 1), + ) + assert result.next_state.held_object is not None + assert ( + result.next_state.held_object.semantics + is target.support_held_object.semantics + ) + assert result.next_state.held_object.object_to_eef.shape == (NUM_ENVS, 4, 4) + assert result.next_state.held_object.grasp_xpos.shape == (NUM_ENVS, 4, 4) diff --git a/tests/sim/atomic_actions/test_core.py b/tests/sim/atomic_actions/test_core.py index 234ed6b4f..5710c60cd 100644 --- a/tests/sim/atomic_actions/test_core.py +++ b/tests/sim/atomic_actions/test_core.py @@ -27,6 +27,9 @@ from embodichain.lab.sim.atomic_actions.core import ( ActionCfg, ActionResult, + CoordinatedHeldObjectState, + CoordinatedPickmentTarget, + CoordinatedPlacementTarget, GraspTarget, HeldObjectState, HeldObjectPoseTarget, @@ -42,6 +45,17 @@ class TestTypedTargets: def test_pose_target_holds_tensor(self): x = torch.eye(4) assert EndEffectorPoseTarget(xpos=x).xpos is x + assert EndEffectorPoseTarget(xpos=x).tcp_symmetry == "none" + + def test_pose_target_can_declare_tcp_symmetry(self): + target = EndEffectorPoseTarget(xpos=torch.eye(4), tcp_symmetry="z_roll_180") + assert target.tcp_symmetry == "z_roll_180" + + def test_pose_target_rejects_unknown_tcp_symmetry(self): + with pytest.raises(ValueError, match="tcp_symmetry"): + EndEffectorPoseTarget( + xpos=torch.eye(4), tcp_symmetry="yaw_90" # type: ignore[arg-type] + ) def test_pose_target_is_frozen(self): t = EndEffectorPoseTarget(xpos=torch.eye(4)) @@ -86,6 +100,33 @@ def test_held_object_target_is_frozen(self): with pytest.raises(dataclasses.FrozenInstanceError): t.object_target_pose = torch.zeros(4, 4) # type: ignore[misc] + def test_coordinated_pickment_target_holds_object_offsets(self): + sem = ObjectSemantics(affordance=Affordance(), geometry={}, label="pencil") + target = CoordinatedPickmentTarget( + object_target_pose=torch.eye(4), + object_semantics=sem, + left_object_to_eef=torch.eye(4), + right_object_to_eef=torch.eye(4), + ) + assert target.object_semantics is sem + assert target.left_object_to_eef.shape == (4, 4) + + def test_coordinated_placement_target_holds_both_held_objects(self): + sem = ObjectSemantics(affordance=Affordance(), geometry={}, label="block") + held = HeldObjectState( + semantics=sem, + object_to_eef=torch.eye(4).unsqueeze(0), + grasp_xpos=torch.eye(4).unsqueeze(0), + ) + target = CoordinatedPlacementTarget( + placing_object_target_pose=torch.eye(4), + support_object_target_pose=torch.eye(4), + placing_held_object=held, + support_held_object=held, + ) + assert target.placing_held_object is held + assert target.support_object_target_pose.shape == (4, 4) + class TestObjectSemantics: def test_does_not_mutate_affordance_geometry(self): @@ -122,6 +163,21 @@ def test_required_fields(self): assert s.grasp_xpos.shape == (1, 4, 4) +class TestCoordinatedHeldObjectState: + def test_required_fields(self): + sem = ObjectSemantics(affordance=Affordance(), geometry={}) + s = CoordinatedHeldObjectState( + semantics=sem, + left_object_to_eef=torch.eye(4).unsqueeze(0), + right_object_to_eef=torch.eye(4).unsqueeze(0), + left_grasp_xpos=torch.eye(4).unsqueeze(0), + right_grasp_xpos=torch.eye(4).unsqueeze(0), + ) + assert s.semantics is sem + assert s.left_object_to_eef.shape == (1, 4, 4) + assert s.right_grasp_xpos.shape == (1, 4, 4) + + class TestWorldState: def test_constructs_with_last_qpos_only(self): qpos = torch.zeros(2, 6) @@ -139,6 +195,18 @@ def test_carries_held_object(self): ws = WorldState(last_qpos=torch.zeros(1, 6), held_object=held) assert ws.held_object is held + def test_carries_coordinated_held_object(self): + sem = ObjectSemantics(affordance=Affordance(), geometry={}) + held = CoordinatedHeldObjectState( + semantics=sem, + left_object_to_eef=torch.eye(4).unsqueeze(0), + right_object_to_eef=torch.eye(4).unsqueeze(0), + left_grasp_xpos=torch.eye(4).unsqueeze(0), + right_grasp_xpos=torch.eye(4).unsqueeze(0), + ) + ws = WorldState(last_qpos=torch.zeros(1, 14), coordinated_held_object=held) + assert ws.coordinated_held_object is held + class TestActionResult: def test_shape_contract(self): diff --git a/tests/sim/atomic_actions/test_engine.py b/tests/sim/atomic_actions/test_engine.py index ff941365d..3baa4d67b 100644 --- a/tests/sim/atomic_actions/test_engine.py +++ b/tests/sim/atomic_actions/test_engine.py @@ -148,13 +148,13 @@ def test_run_concatenates_trajectories(self): b = _fake_action("b", EndEffectorPoseTarget) self.engine.register(a, name="a") self.engine.register(b, name="b") - ok, traj, _ = self.engine.run( + success, traj, _ = self.engine.run( [ ("a", EndEffectorPoseTarget(torch.eye(4))), ("b", EndEffectorPoseTarget(torch.eye(4))), ] ) - assert ok is True + assert success.all().item() assert traj.shape == (NUM_ENVS, 10, TOTAL_DOF) def test_run_threads_world_state(self): @@ -165,14 +165,14 @@ def test_run_threads_world_state(self): self.engine.register(move, name="move") self.engine.register(place, name="place") sem = ObjectSemantics(affordance=Affordance(), geometry={}, label="x") - ok, _, final_state = self.engine.run( + success, _, final_state = self.engine.run( [ ("pick", GraspTarget(sem)), ("move", HeldObjectPoseTarget(torch.eye(4))), ("place", EndEffectorPoseTarget(torch.eye(4))), ] ) - assert ok is True + assert success.all().item() # The move action saw a non-None held_object (set by pick). move_state_arg = move.execute.call_args_list[0].args[1] assert move_state_arg.held_object is not None @@ -186,17 +186,17 @@ def test_run_stops_on_first_failure(self): self.engine.register(a, name="a") self.engine.register(b, name="b") self.engine.register(c, name="c") - ok, traj, _ = self.engine.run( + success, traj, _ = self.engine.run( [ ("a", EndEffectorPoseTarget(torch.eye(4))), ("b", EndEffectorPoseTarget(torch.eye(4))), ("c", EndEffectorPoseTarget(torch.eye(4))), ] ) - assert ok is False + assert not success.any().item() # `c` should not have been called. c.execute.assert_not_called() - # We still get back the partial trajectory accumulated from `a`. + # We get only the trajectory from executed steps (all-dead breaks the loop). assert traj.shape == (NUM_ENVS, 5, TOTAL_DOF) def test_run_raises_on_unknown_action_name(self): diff --git a/tests/sim/atomic_actions/test_engine_per_env.py b/tests/sim/atomic_actions/test_engine_per_env.py new file mode 100644 index 000000000..b71b3a3b2 --- /dev/null +++ b/tests/sim/atomic_actions/test_engine_per_env.py @@ -0,0 +1,77 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Per-environment failure propagation tests for AtomicActionEngine.run().""" + +from __future__ import annotations + +import torch +import pytest +from unittest.mock import Mock + +from embodichain.lab.sim.atomic_actions.engine import AtomicActionEngine +from embodichain.lab.sim.atomic_actions.core import ( + ActionResult, + AtomicAction, + WorldState, + EndEffectorPoseTarget, + ActionCfg, +) + + +class _StubAction(AtomicAction): + TargetType = EndEffectorPoseTarget + + def __init__(self, mg, success_vec, traj_len=4, dof=3): + super().__init__(mg, ActionCfg()) + self._success = torch.tensor(success_vec) + self._traj_len = traj_len + self._dof = dof + + def execute(self, target, state): + n = state.last_qpos.shape[0] + traj = torch.zeros(n, self._traj_len, self._dof) + traj[:] = state.last_qpos.unsqueeze(1) + return ActionResult( + success=self._success.clone(), + trajectory=traj, + next_state=WorldState(last_qpos=traj[:, -1, :].clone()), + ) + + +class TestRunPerEnv: + def test_failed_env_holds(self): + mg = Mock() + mg.robot.get_qpos = lambda: torch.zeros(3, 3) + mg.robot.dof = 3 + mg.device = torch.device("cpu") + eng = AtomicActionEngine(mg) + # env 1 fails step 2 + eng.register(_StubAction(mg, [True, True, True]), name="a") + eng.register(_StubAction(mg, [True, False, True]), name="b") + eng.register(_StubAction(mg, [True, True, True]), name="c") + success, traj, state = eng.run( + steps=[ + ("a", EndEffectorPoseTarget(xpos=torch.eye(4))), + ("b", EndEffectorPoseTarget(xpos=torch.eye(4))), + ("c", EndEffectorPoseTarget(xpos=torch.eye(4))), + ] + ) + assert success.tolist() == [True, False, True] + assert traj.shape[1] == 12 # 3 steps * 4 waypoints + # env 1's rows after its failure should equal its pre-failure qpos (held) + # all zeros here, so just check shape and that env 0/2 advanced + assert state.last_qpos.shape == (3, 3) diff --git a/tests/sim/atomic_actions/test_motion_source_e2e.py b/tests/sim/atomic_actions/test_motion_source_e2e.py new file mode 100644 index 000000000..a7161c36b --- /dev/null +++ b/tests/sim/atomic_actions/test_motion_source_e2e.py @@ -0,0 +1,107 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Reach-equivalence e2e test for MoveEndEffector across motion sources.""" + +from __future__ import annotations + +import torch +import pytest + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.robots import CobotMagicCfg +from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg +from embodichain.lab.sim.atomic_actions import AtomicActionEngine +from embodichain.lab.sim.atomic_actions.actions import ( + MoveEndEffector, + MoveEndEffectorCfg, +) +from embodichain.lab.sim.atomic_actions.core import EndEffectorPoseTarget + + +@pytest.mark.requires_sim +@pytest.mark.slow +class TestMotionSourceReachEquivalence: + """Verify MoveEndEffector reaches a reachable pose for ik_interp and motion_gen.""" + + CONTROL_PART = "left_arm" + ROBOT_UID = "cobot_e2e" + SAMPLE_INTERVAL = 80 + POS_TOL = 0.02 + + def _setup(self, motion_source: str): + sim = SimulationManager(SimulationManagerCfg(headless=True, device="cpu")) + robot = sim.add_robot( + cfg=CobotMagicCfg.from_dict( + { + "uid": self.ROBOT_UID, + "init_pos": [0.0, 0.0, 0.7775], + "init_qpos": [0.0] * 16, + } + ) + ) + mg = MotionGenerator( + MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=self.ROBOT_UID)) + ) + engine = AtomicActionEngine(mg) + cfg = MoveEndEffectorCfg( + motion_source=motion_source, + planner_type="toppra" if motion_source == "motion_gen" else None, + control_part=self.CONTROL_PART, + sample_interval=self.SAMPLE_INTERVAL, + ) + engine.register(MoveEndEffector(mg, cfg), name="move_end_effector") + return sim, robot, engine + + def _teardown(self, sim): + sim.destroy() + import embodichain.lab.sim as om + + om.SimulationManager.flush_cleanup_queue() + + def _reachable_target(self, robot): + """Return the current EE pose shifted 5 cm upward and the arm joint ids.""" + arm_ids = robot.get_joint_ids(name=self.CONTROL_PART) + qpos = robot.get_qpos(name=self.CONTROL_PART) + fk = robot.compute_fk(qpos=qpos, name=self.CONTROL_PART, to_matrix=True) + target = fk[0].clone() + target[2, 3] += 0.05 + return target, arm_ids + + def _run_reach_test(self, motion_source: str): + sim, robot, engine = self._setup(motion_source) + try: + target, arm_ids = self._reachable_target(robot) + success, traj, _ = engine.run( + [("move_end_effector", EndEffectorPoseTarget(xpos=target))] + ) + assert success.all().item(), f"{motion_source} reported failure" + final_q = traj[0, -1, arm_ids] + fk = robot.compute_fk( + qpos=final_q[None], name=self.CONTROL_PART, to_matrix=True + )[0] + err = torch.norm(fk[:3, 3] - target[:3, 3]) + assert ( + err < self.POS_TOL + ), f"{motion_source} EE pos error {err.item():.4f} m" + finally: + self._teardown(sim) + + def test_ik_interp_reaches_target(self): + self._run_reach_test("ik_interp") + + def test_motion_gen_toppra_reaches_target(self): + self._run_reach_test("motion_gen") diff --git a/tests/sim/atomic_actions/test_primitives_helpers.py b/tests/sim/atomic_actions/test_primitives_helpers.py new file mode 100644 index 000000000..6985edda7 --- /dev/null +++ b/tests/sim/atomic_actions/test_primitives_helpers.py @@ -0,0 +1,36 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Tests for atomic action primitive helpers.""" + +from __future__ import annotations + +import pytest +import torch + +from embodichain.lab.sim.atomic_actions.primitives._helpers import ( + resolve_object_target, +) + + +def test_resolve_object_target_uses_custom_name_in_shape_error() -> None: + with pytest.raises(ValueError, match="placing_object_target_pose"): + resolve_object_target( + torch.zeros(2, 4, 4), + n_envs=3, + device=torch.device("cpu"), + name="placing_object_target_pose", + ) diff --git a/tests/sim/atomic_actions/test_trajectory_motion_source.py b/tests/sim/atomic_actions/test_trajectory_motion_source.py new file mode 100644 index 000000000..cb3fccdca --- /dev/null +++ b/tests/sim/atomic_actions/test_trajectory_motion_source.py @@ -0,0 +1,99 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Tests for TrajectoryBuilder motion_source branching.""" + +from __future__ import annotations + +import torch +import pytest +from unittest.mock import Mock + +from embodichain.lab.sim.atomic_actions.trajectory import TrajectoryBuilder +from embodichain.lab.sim.planners.utils import PlanState, MoveType +from embodichain.lab.sim.atomic_actions.core import ActionCfg + + +def _mock_mg(num_envs=2, arm_dof=6): + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = arm_dof + robot.get_qpos = lambda name=None: torch.zeros(num_envs, arm_dof) + + def compute_ik(pose=None, name=None, joint_seed=None, **kw): + return torch.ones(num_envs, dtype=torch.bool), torch.zeros(num_envs, arm_dof) + + robot.compute_ik = compute_ik + mg = Mock() + mg.robot = robot + mg.device = torch.device("cpu") + return mg + + +class TestPlanArmTrajMotionGen: + def test_motion_gen_path_delegates_to_generate(self): + mg = _mock_mg(num_envs=3, arm_dof=6) + from embodichain.lab.sim.planners.utils import PlanResult + + mg.generate.return_value = PlanResult( + success=torch.ones(3, dtype=torch.bool), + positions=torch.zeros(3, 12, 6), + ) + builder = TrajectoryBuilder(mg) + cfg = ActionCfg( + motion_source="motion_gen", planner_type="toppra", control_part="arm" + ) + start_qpos = torch.zeros(3, 6) + # per-env list[list[PlanState]] with single-env PlanStates (action contract) + target_states_list = [ + [ + PlanState(xpos=torch.eye(4), move_type=MoveType.EEF_MOVE), + PlanState(xpos=torch.eye(4), move_type=MoveType.EEF_MOVE), + ] + for _ in range(3) + ] + ok, traj = builder.plan_arm_traj( + target_states_list, + start_qpos, + 12, + control_part="arm", + arm_dof=6, + cfg=cfg, + ) + assert ok.shape == (3,) + assert ok.all().item() + assert traj.shape == (3, 12, 6) + mg.generate.assert_called_once() + + def test_ik_interp_path_unchanged(self): + mg = _mock_mg(num_envs=2, arm_dof=6) + builder = TrajectoryBuilder(mg) + cfg = ActionCfg(motion_source="ik_interp", control_part="arm") + start_qpos = torch.zeros(2, 6) + target_states_list = [ + [PlanState(xpos=torch.eye(4), move_type=MoveType.EEF_MOVE)] + for _ in range(2) + ] + ok, traj = builder.plan_arm_traj( + target_states_list, + start_qpos, + 10, + control_part="arm", + arm_dof=6, + cfg=cfg, + ) + assert ok.all().item() + assert traj.shape[0] == 2 diff --git a/tests/sim/atomic_actions/test_tutorial_utils.py b/tests/sim/atomic_actions/test_tutorial_utils.py new file mode 100644 index 000000000..693e5903c --- /dev/null +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -0,0 +1,145 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Tests for atomic-action tutorial helpers.""" + +from __future__ import annotations + +from argparse import Namespace +from unittest.mock import MagicMock + +import pytest +import torch + +from scripts.tutorials.atomic_action.tutorial_utils import ( + broadcast_pose_batch, + broadcast_waypoint_pose_batch, + clone_local_pose_from_first_env, + create_antipodal_semantics, + should_wait_for_tutorial_input, +) + + +def test_should_wait_for_tutorial_input_is_disabled_for_headless_modes() -> None: + assert ( + should_wait_for_tutorial_input( + Namespace( + auto_play=False, + headless=True, + diagnose_plan=False, + headless_play=False, + ) + ) + is False + ) + assert ( + should_wait_for_tutorial_input( + Namespace( + auto_play=False, + headless=False, + diagnose_plan=True, + headless_play=False, + ) + ) + is False + ) + assert ( + should_wait_for_tutorial_input( + Namespace( + auto_play=False, + headless=False, + diagnose_plan=False, + headless_play=True, + ) + ) + is False + ) + + +def test_broadcast_pose_batch_repeats_single_pose_for_each_env() -> None: + pose = torch.eye(4, dtype=torch.float32) + + batched = broadcast_pose_batch(pose, num_envs=3) + + assert batched.shape == (3, 4, 4) + assert torch.allclose(batched[0], pose) + assert torch.allclose(batched[1], pose) + assert torch.allclose(batched[2], pose) + + +def test_broadcast_waypoint_pose_batch_repeats_waypoints_for_each_env() -> None: + waypoints = torch.stack( + [torch.eye(4, dtype=torch.float32), 2.0 * torch.eye(4, dtype=torch.float32)], + dim=0, + ) + + batched = broadcast_waypoint_pose_batch(waypoints, num_envs=2) + + assert batched.shape == (2, 2, 4, 4) + assert torch.allclose(batched[0], waypoints) + assert torch.allclose(batched[1], waypoints) + + +def test_clone_local_pose_from_first_env_sets_shared_pose() -> None: + first_pose = torch.eye(4, dtype=torch.float32) + first_pose[0, 3] = 0.2 + poses = torch.stack( + [ + first_pose, + 2.0 * torch.eye(4, dtype=torch.float32), + 3.0 * torch.eye(4, dtype=torch.float32), + ], + dim=0, + ) + entity = MagicMock() + entity.get_local_pose.return_value = poses + + shared = clone_local_pose_from_first_env(entity) + + expected = first_pose.unsqueeze(0).repeat(3, 1, 1) + assert torch.allclose(shared, expected) + entity.set_local_pose.assert_called_once() + assert torch.allclose(entity.set_local_pose.call_args.args[0], expected) + + +def test_create_antipodal_semantics_keeps_mesh_data_on_affordance() -> None: + vertices = torch.tensor([[0.0, 0.0, 0.0], [0.1, 0.0, 0.0]]) + triangles = torch.tensor([[0, 1, 1]]) + obj = MagicMock() + obj.get_vertices.return_value = vertices.unsqueeze(0) + obj.get_triangles.return_value = triangles.unsqueeze(0) + + semantics = create_antipodal_semantics( + obj, + label="cube", + n_sample=64, + force_reannotate=True, + ) + + assert semantics.entity is obj + assert semantics.label == "cube" + assert semantics.geometry == {} + assert torch.equal(semantics.affordance.mesh_vertices, vertices) + assert torch.equal(semantics.affordance.mesh_triangles, triangles) + assert semantics.affordance.force_reannotate is True + assert semantics.affordance.generator_cfg.antipodal_sampler_cfg.n_sample == 64 + + +def test_broadcast_pose_batch_rejects_wrong_env_count() -> None: + poses = torch.eye(4, dtype=torch.float32).unsqueeze(0).repeat(2, 1, 1) + + with pytest.raises(ValueError, match="num_envs"): + broadcast_pose_batch(poses, num_envs=3) diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index 80eb1ef3e..f09033aac 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os import torch import pytest @@ -50,6 +52,20 @@ def _link_static_friction(art: Articulation, link_name: str, env_idx: int = 0) - return art._entities[env_idx].get_physical_attr(link_name).static_friction +class _EntityMethodOverride: + """Delegate every entity method except one overridden setter.""" + + def __init__(self, entity, method_name: str, override): + self._entity = entity + self._method_name = method_name + self._override = override + + def __getattr__(self, name: str): + if name == self._method_name: + return self._override + return getattr(self._entity, name) + + class TestRigidBodyAttributesOverride: """Pure-Python tests for per-link physics config merging.""" @@ -317,6 +333,483 @@ def test_get_joint_drive_with_joint_ids(self): armature, expected_armature, atol=1e-5 ), "FAIL: armature does not match expected filtered values" + def test_joint_limit_getters_support_env_and_joint_filters(self): + """Test joint limit getters support joint_ids and env_ids filtering.""" + all_qpos_limits = self.art.body_data.qpos_limits + ( + _stiffness, + _damping, + all_qf_limits, + all_qvel_limits, + _friction, + _armature, + ) = self.art.get_joint_drive() + + joint_ids = [0, self.art.dof - 1] if self.art.dof >= 2 else [0] + env_ids = [0, 2, 4] if NUM_ARENAS >= 5 else [0] + + qpos_limits = self.art.get_qpos_limits(joint_ids=joint_ids, env_ids=env_ids) + qvel_limits = self.art.get_qvel_limits(joint_ids=joint_ids, env_ids=env_ids) + qf_limits = self.art.get_qf_limits(joint_ids=joint_ids, env_ids=env_ids) + + expected_qpos_limits = all_qpos_limits[env_ids][:, joint_ids, :] + expected_qvel_limits = all_qvel_limits[env_ids][:, joint_ids] + expected_qf_limits = all_qf_limits[env_ids][:, joint_ids] + + expected_qpos_shape = (len(env_ids), len(joint_ids), 2) + expected_joint_shape = (len(env_ids), len(joint_ids)) + + assert torch.allclose( + self.art.body_data.qvel_limits, all_qvel_limits, atol=1e-5 + ), "FAIL: qvel_limits backing tensor does not match post-init joint drive state" + assert torch.allclose( + self.art.body_data.qf_limits, all_qf_limits, atol=1e-5 + ), "FAIL: qf_limits backing tensor does not match post-init joint drive state" + + assert ( + qpos_limits.shape == expected_qpos_shape + ), f"FAIL: Expected qpos_limits shape {expected_qpos_shape}, got {qpos_limits.shape}" + assert ( + qvel_limits.shape == expected_joint_shape + ), f"FAIL: Expected qvel_limits shape {expected_joint_shape}, got {qvel_limits.shape}" + assert ( + qf_limits.shape == expected_joint_shape + ), f"FAIL: Expected qf_limits shape {expected_joint_shape}, got {qf_limits.shape}" + + assert torch.allclose( + qpos_limits, expected_qpos_limits, atol=1e-5 + ), "FAIL: qpos_limits does not match expected filtered values" + assert torch.allclose( + qvel_limits, expected_qvel_limits, atol=1e-5 + ), "FAIL: qvel_limits does not match expected filtered values" + assert torch.allclose( + qf_limits, expected_qf_limits, atol=1e-5 + ), "FAIL: qf_limits does not match expected filtered values" + + def test_joint_limit_cache_tracks_set_joint_drive_updates(self): + """Test qvel/qf limit caches stay aligned with set_joint_drive writes.""" + ( + _stiffness_before, + _damping_before, + all_qf_limits_before, + all_qvel_limits_before, + _friction_before, + _armature_before, + ) = self.art.get_joint_drive() + + joint_ids = [0, self.art.dof - 1] if self.art.dof >= 2 else [0] + env_ids = [0, 2, 4] if NUM_ARENAS >= 5 else [0] + env_ids_tensor = torch.as_tensor( + env_ids, dtype=torch.long, device=self.sim.device + ) + joint_ids_tensor = torch.as_tensor( + joint_ids, dtype=torch.long, device=self.sim.device + ) + + new_qvel_limits = torch.full( + (len(env_ids), len(joint_ids)), + 321.0, + dtype=torch.float32, + device=self.sim.device, + ) + new_qf_limits = torch.full( + (len(env_ids), len(joint_ids)), + 654.0, + dtype=torch.float32, + device=self.sim.device, + ) + + self.art.set_joint_drive( + max_effort=new_qf_limits, + max_velocity=new_qvel_limits, + joint_ids=joint_ids, + env_ids=env_ids, + ) + + ( + _stiffness_after, + _damping_after, + all_qf_limits_after, + all_qvel_limits_after, + _friction_after, + _armature_after, + ) = self.art.get_joint_drive() + qvel_limits = self.art.get_qvel_limits(joint_ids=joint_ids, env_ids=env_ids) + qf_limits = self.art.get_qf_limits(joint_ids=joint_ids, env_ids=env_ids) + + expected_qvel_limits = all_qvel_limits_before.clone() + expected_qvel_limits[env_ids_tensor[:, None], joint_ids_tensor] = ( + new_qvel_limits + ) + expected_qf_limits = all_qf_limits_before.clone() + expected_qf_limits[env_ids_tensor[:, None], joint_ids_tensor] = new_qf_limits + + assert torch.allclose( + self.art.body_data.qvel_limits, expected_qvel_limits, atol=1e-5 + ), "FAIL: qvel_limits backing tensor did not track set_joint_drive max_velocity" + assert torch.allclose( + self.art.body_data.qf_limits, expected_qf_limits, atol=1e-5 + ), "FAIL: qf_limits backing tensor did not track set_joint_drive max_effort" + assert torch.allclose( + all_qvel_limits_after, expected_qvel_limits, atol=1e-5 + ), "FAIL: live qvel limits did not match expected post-write state" + assert torch.allclose( + all_qf_limits_after, expected_qf_limits, atol=1e-5 + ), "FAIL: live qf limits did not match expected post-write state" + assert torch.allclose( + qvel_limits, new_qvel_limits, atol=1e-5 + ), "FAIL: filtered qvel_limits did not return the updated max_velocity values" + assert torch.allclose( + qf_limits, new_qf_limits, atol=1e-5 + ), "FAIL: filtered qf_limits did not return the updated max_effort values" + + def test_joint_limit_setters_update_selected_envs_and_body_data(self): + """Test joint limit setters update selected envs and cached body data.""" + joint_ids = [0, self.art.dof - 1] if self.art.dof >= 2 else [0] + env_ids = [0, 2] if NUM_ARENAS >= 3 else [0] + + original_qpos_limits = self.art.body_data.qpos_limits.clone() + original_qvel_limits = self.art.body_data.qvel_limits.clone() + original_qf_limits = self.art.body_data.qf_limits.clone() + + qpos_limits = self.art.get_qpos_limits( + joint_ids=joint_ids, env_ids=env_ids + ).clone() + scale = torch.arange( + 1, + len(env_ids) * len(joint_ids) + 1, + dtype=torch.float32, + device=self.sim.device, + ).reshape(len(env_ids), len(joint_ids)) + tighten = torch.minimum( + 0.001 * scale, + 0.25 * (qpos_limits[:, :, 1] - qpos_limits[:, :, 0]), + ) + qpos_limits[:, :, 0] += tighten + qpos_limits[:, :, 1] -= tighten + qvel_limits = 0.5 + 0.05 * scale + qf_limits = 1.0 + 0.25 * scale + + self.art.set_qpos_limits(qpos_limits, joint_ids=joint_ids, env_ids=env_ids) + self.art.set_qvel_limits(qvel_limits, joint_ids=joint_ids, env_ids=env_ids) + self.art.set_qf_limits(qf_limits, joint_ids=joint_ids, env_ids=env_ids) + + updated_qpos_limits = self.art.get_qpos_limits( + joint_ids=joint_ids, env_ids=env_ids + ) + updated_qvel_limits = self.art.get_qvel_limits( + joint_ids=joint_ids, env_ids=env_ids + ) + updated_qf_limits = self.art.get_qf_limits(joint_ids=joint_ids, env_ids=env_ids) + + assert torch.allclose( + updated_qpos_limits, qpos_limits, atol=1e-5 + ), "FAIL: filtered qpos_limits did not return the written values" + assert torch.allclose( + updated_qvel_limits, qvel_limits, atol=1e-5 + ), "FAIL: filtered qvel_limits did not return the written values" + assert torch.allclose( + updated_qf_limits, qf_limits, atol=1e-5 + ), "FAIL: filtered qf_limits did not return the written values" + assert torch.allclose( + self.art.body_data.qpos_limits[env_ids][:, joint_ids, :], + qpos_limits, + atol=1e-5, + ), "FAIL: body_data qpos_limits cache did not update for the selected slice" + assert torch.allclose( + self.art.body_data.qvel_limits[env_ids][:, joint_ids], + qvel_limits, + atol=1e-5, + ), "FAIL: body_data qvel_limits cache did not update for the selected slice" + assert torch.allclose( + self.art.body_data.qf_limits[env_ids][:, joint_ids], + qf_limits, + atol=1e-5, + ), "FAIL: body_data qf_limits cache did not update for the selected slice" + + non_selected_joint_ids = [ + joint_id for joint_id in range(self.art.dof) if joint_id not in joint_ids + ] + if non_selected_joint_ids: + assert torch.allclose( + self.art.body_data.qpos_limits[env_ids][:, non_selected_joint_ids, :], + original_qpos_limits[env_ids][:, non_selected_joint_ids, :], + atol=1e-5, + ), "FAIL: qpos_limits changed for non-selected joints in targeted environments" + assert torch.allclose( + self.art.body_data.qvel_limits[env_ids][:, non_selected_joint_ids], + original_qvel_limits[env_ids][:, non_selected_joint_ids], + atol=1e-5, + ), "FAIL: qvel_limits changed for non-selected joints in targeted environments" + assert torch.allclose( + self.art.body_data.qf_limits[env_ids][:, non_selected_joint_ids], + original_qf_limits[env_ids][:, non_selected_joint_ids], + atol=1e-5, + ), "FAIL: qf_limits changed for non-selected joints in targeted environments" + + untouched_env_ids = [ + env_id for env_id in range(NUM_ARENAS) if env_id not in env_ids + ] + if untouched_env_ids: + assert torch.allclose( + self.art.body_data.qpos_limits[untouched_env_ids], + original_qpos_limits[untouched_env_ids], + atol=1e-5, + ), "FAIL: qpos_limits changed for untouched environments" + assert torch.allclose( + self.art.body_data.qvel_limits[untouched_env_ids], + original_qvel_limits[untouched_env_ids], + atol=1e-5, + ), "FAIL: qvel_limits changed for untouched environments" + assert torch.allclose( + self.art.body_data.qf_limits[untouched_env_ids], + original_qf_limits[untouched_env_ids], + atol=1e-5, + ), "FAIL: qf_limits changed for untouched environments" + + def test_joint_limit_setters_accept_single_env_convenience_shapes(self): + """Test single-env convenience shapes for joint limit setters.""" + env_ids = [0] + joint_ids = [0, self.art.dof - 1] if self.art.dof >= 2 else [0] + + original_qpos_limits = self.art.body_data.qpos_limits.clone() + original_qvel_limits = self.art.body_data.qvel_limits.clone() + original_qf_limits = self.art.body_data.qf_limits.clone() + + qpos_limits = self.art.get_qpos_limits(joint_ids=joint_ids, env_ids=env_ids)[ + 0 + ].clone() + scale = torch.arange( + 1, + len(joint_ids) + 1, + dtype=torch.float32, + device=self.sim.device, + ) + tighten = torch.minimum( + 0.001 * scale, + 0.25 * (qpos_limits[:, 1] - qpos_limits[:, 0]), + ) + qpos_limits[:, 0] += tighten + qpos_limits[:, 1] -= tighten + qvel_limits = 0.6 + 0.05 * scale + qf_limits = 1.5 + 0.1 * scale + + self.art.set_qpos_limits(qpos_limits, joint_ids=joint_ids, env_ids=env_ids) + self.art.set_qvel_limits(qvel_limits, joint_ids=joint_ids, env_ids=env_ids) + self.art.set_qf_limits(qf_limits, joint_ids=joint_ids, env_ids=env_ids) + + assert torch.allclose( + self.art.get_qpos_limits(joint_ids=joint_ids, env_ids=env_ids), + qpos_limits.unsqueeze(0), + atol=1e-5, + ), "FAIL: single-env qpos convenience shape did not write expected values" + assert torch.allclose( + self.art.get_qvel_limits(joint_ids=joint_ids, env_ids=env_ids), + qvel_limits.unsqueeze(0), + atol=1e-5, + ), "FAIL: single-env qvel convenience shape did not write expected values" + assert torch.allclose( + self.art.get_qf_limits(joint_ids=joint_ids, env_ids=env_ids), + qf_limits.unsqueeze(0), + atol=1e-5, + ), "FAIL: single-env qf convenience shape did not write expected values" + + non_selected_joint_ids = [ + joint_id for joint_id in range(self.art.dof) if joint_id not in joint_ids + ] + if non_selected_joint_ids: + assert torch.allclose( + self.art.body_data.qpos_limits[env_ids][:, non_selected_joint_ids, :], + original_qpos_limits[env_ids][:, non_selected_joint_ids, :], + atol=1e-5, + ), "FAIL: single-env qpos write changed non-selected joints" + assert torch.allclose( + self.art.body_data.qvel_limits[env_ids][:, non_selected_joint_ids], + original_qvel_limits[env_ids][:, non_selected_joint_ids], + atol=1e-5, + ), "FAIL: single-env qvel write changed non-selected joints" + assert torch.allclose( + self.art.body_data.qf_limits[env_ids][:, non_selected_joint_ids], + original_qf_limits[env_ids][:, non_selected_joint_ids], + atol=1e-5, + ), "FAIL: single-env qf write changed non-selected joints" + + untouched_env_ids = [ + env_id for env_id in range(NUM_ARENAS) if env_id not in env_ids + ] + if untouched_env_ids: + assert torch.allclose( + self.art.body_data.qpos_limits[untouched_env_ids], + original_qpos_limits[untouched_env_ids], + atol=1e-5, + ), "FAIL: single-env qpos write changed untouched environments" + assert torch.allclose( + self.art.body_data.qvel_limits[untouched_env_ids], + original_qvel_limits[untouched_env_ids], + atol=1e-5, + ), "FAIL: single-env qvel write changed untouched environments" + assert torch.allclose( + self.art.body_data.qf_limits[untouched_env_ids], + original_qf_limits[untouched_env_ids], + atol=1e-5, + ), "FAIL: single-env qf write changed untouched environments" + + def test_set_qpos_limits_failure_does_not_update_cache(self): + """Test a failed DexSim qpos limit write leaves the Python cache unchanged.""" + env_ids = [0] + joint_ids = [0, self.art.dof - 1] if self.art.dof >= 2 else [0] + original_qpos_limits = self.art.body_data.qpos_limits.clone() + + qpos_limits = self.art.get_qpos_limits( + joint_ids=joint_ids, env_ids=env_ids + ).clone() + scale = torch.arange( + 1, + len(joint_ids) + 1, + dtype=torch.float32, + device=self.sim.device, + ).unsqueeze(0) + tighten = torch.minimum( + 0.001 * scale, + 0.25 * (qpos_limits[:, :, 1] - qpos_limits[:, :, 0]), + ) + qpos_limits[:, :, 0] += tighten + qpos_limits[:, :, 1] -= tighten + + original_entity = self.art._entities[0] + + def _fail_set_joint_limits(_limits, _joint_ids): + return -1 + + self.art._entities[0] = _EntityMethodOverride( + original_entity, + "set_joint_position_limits", + _fail_set_joint_limits, + ) + try: + with pytest.raises(RuntimeError, match="set_joint_position_limits failed"): + self.art.set_qpos_limits( + qpos_limits, joint_ids=joint_ids, env_ids=env_ids + ) + finally: + self.art._entities[0] = original_entity + + assert torch.allclose( + self.art.body_data.qpos_limits, + original_qpos_limits, + atol=1e-5, + ), "FAIL: qpos_limits cache changed after a failed DexSim write" + + def test_set_qpos_clamps_against_updated_qpos_limits(self): + """Test set_qpos clamps to the selected environments' exact updated limits.""" + env_ids = [0, 1] if NUM_ARENAS >= 2 else [0] + joint_ids = [0] + if len(env_ids) == 2: + qpos_limits = torch.tensor( + [[[-0.02, 0.02]], [[-0.05, 0.01]]], + dtype=torch.float32, + device=self.sim.device, + ) + requested_qpos = torch.tensor( + [[0.5], [-0.5]], + dtype=torch.float32, + device=self.sim.device, + ) + expected_qpos = torch.tensor( + [[0.02], [-0.05]], + dtype=torch.float32, + device=self.sim.device, + ) + else: + qpos_limits = torch.tensor( + [[[-0.02, 0.02]]], + dtype=torch.float32, + device=self.sim.device, + ) + requested_qpos = torch.tensor( + [[0.5]], + dtype=torch.float32, + device=self.sim.device, + ) + expected_qpos = torch.tensor( + [[0.02]], + dtype=torch.float32, + device=self.sim.device, + ) + + self.art.set_qpos_limits(qpos_limits, joint_ids=joint_ids, env_ids=env_ids) + + self.art.set_qpos( + requested_qpos, + joint_ids=joint_ids, + env_ids=env_ids, + target=False, + ) + + clamped_qpos = self.art.get_qpos()[env_ids][:, joint_ids] + + assert torch.allclose( + clamped_qpos, expected_qpos, atol=1e-5 + ), f"FAIL: qpos did not clamp to the per-env updated limits: {clamped_qpos.tolist()}" + + def test_qpos_limits_from_cfg_dict_can_tighten(self): + """Test qpos_limits can be set from ArticulationCfg with a regex dictionary.""" + from embodichain.lab.sim.cfg import ArticulationCfg + + cfg = ArticulationCfg( + uid="drawer_cfg_qpos_limits", + fpath=get_data_path(ART_PATH), + drive_pros=JointDrivePropertiesCfg(drive_type="force"), + qpos_limits={".*": [-0.05, 0.05]}, + ) + art: Articulation = self.sim.add_articulation(cfg=cfg) + limits = art.get_qpos_limits() + assert torch.all( + limits[..., 0] >= -0.0501 + ), "FAIL: cfg qpos_limits lower bound not applied" + assert torch.all( + limits[..., 1] <= 0.0501 + ), "FAIL: cfg qpos_limits upper bound not applied" + + def test_qpos_limits_from_cfg_can_expand(self): + """Test qpos_limits from ArticulationCfg can expand joint limits.""" + from embodichain.lab.sim.cfg import ArticulationCfg + + joint_name = self.art.joint_names[0] + asset_limits = self.art.get_qpos_limits()[:, 0, :] + expanded_lower = asset_limits[:, 0].min().item() - 0.1 + expanded_upper = asset_limits[:, 1].max().item() + 0.1 + + cfg = ArticulationCfg( + uid="drawer_expanded_limits", + fpath=get_data_path(ART_PATH), + drive_pros=JointDrivePropertiesCfg(drive_type="force"), + qpos_limits={joint_name: [expanded_lower, expanded_upper]}, + ) + art: Articulation = self.sim.add_articulation(cfg=cfg) + limits = art.get_qpos_limits()[:, 0, :] + assert torch.allclose( + limits, + torch.tensor( + [expanded_lower, expanded_upper], + device=self.sim.device, + dtype=torch.float32, + ), + atol=1e-4, + ), f"FAIL: cfg qpos_limits not applied: {limits.tolist()}" + + # set_qpos should clamp to the expanded limits, not the asset limits. + requested_qpos = torch.full( + (NUM_ARENAS, 1), expanded_upper, device=self.sim.device + ) + art.set_qpos(requested_qpos, joint_ids=[0], target=False) + actual_qpos = art.get_qpos()[:, 0] + assert torch.allclose( + actual_qpos, + torch.full_like(actual_qpos, expanded_upper), + atol=1e-4, + ), f"FAIL: set_qpos did not use expanded limits: {actual_qpos.tolist()}" + def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() diff --git a/tests/sim/objects/test_cloth_object.py b/tests/sim/objects/test_cloth_object.py index 7b3aa3130..3f0f6f81b 100644 --- a/tests/sim/objects/test_cloth_object.py +++ b/tests/sim/objects/test_cloth_object.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os from dexsim.utility.path import get_resources_data_path from embodichain.lab.sim import SimulationManager, SimulationManagerCfg diff --git a/tests/sim/objects/test_dual_arm.py b/tests/sim/objects/test_dual_arm.py index 93a5bc06e..1daef1df2 100644 --- a/tests/sim/objects/test_dual_arm.py +++ b/tests/sim/objects/test_dual_arm.py @@ -182,9 +182,17 @@ def test_dual_arm_from_dict_explicit_base_robot(): assert cfg.solver_cfg["left_arm"].ur_type == "ur5" +def test_dual_arm_from_dict_franka_base_robot(): + cfg = DualArmRobotCfg.from_dict( + {"base_robot": "franka", "mount": {"preset": "side_by_side", "separation": 0.6}} + ) + assert set(["left_arm", "right_arm", "dual_arm"]).issubset(cfg.control_parts.keys()) + assert set(["left_hand", "right_hand"]).issubset(cfg.control_parts.keys()) + + def test_dual_arm_unknown_base_robot_raises(): with pytest.raises(ValueError): - DualArmRobotCfg.from_dict({"base_robot": "franka"}) + DualArmRobotCfg.from_dict({"base_robot": "telepathic"}) def test_dual_arm_roundtrip(): diff --git a/tests/sim/objects/test_light.py b/tests/sim/objects/test_light.py index 2840567bc..e8ea7ed57 100644 --- a/tests/sim/objects/test_light.py +++ b/tests/sim/objects/test_light.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import pytest import torch from embodichain.lab.sim import SimulationManager, SimulationManagerCfg @@ -31,7 +33,7 @@ def setup_method(self): "light_type": "point", "color": [0.1, 0.1, 0.1], "radius": 10.0, - "position": [0.0, 0.0, 2.0], + "init_pos": [0.0, 0.0, 2.0], "uid": "point_light", } self.light = self.sim.add_light(cfg=LightCfg.from_dict(cfg_dict)) @@ -159,3 +161,297 @@ def teardown_method(self): import gc gc.collect() + + +class TestLightTypes: + """Tests for all six supported light types.""" + + @pytest.fixture(autouse=True) + def setup(self): + """Create a SimulationManager for each test.""" + config = SimulationManagerCfg(headless=True, device="cpu", num_envs=4) + self.sim = SimulationManager(config) + yield + self.sim.destroy() + import embodichain.lab.sim as om + + om.SimulationManager.flush_cleanup_queue() + self.__dict__.clear() + import gc + + gc.collect() + + # ------------------------------------------------------------------ + # Creation tests + # ------------------------------------------------------------------ + + @pytest.mark.parametrize( + "light_type, expected_num_instances", + [ + ("point", 4), + ("sun", 1), + ("direction", 1), + ("spot", 4), + ("rect", 4), + ("mesh", 4), + ], + ) + def test_create_each_light_type(self, light_type, expected_num_instances): + """Each of the 6 light types can be created without error. + + ``sun`` and ``direction`` are global scene lights with a single instance. + All other types are per-environment batched lights. + """ + cfg = LightCfg( + uid=f"test_{light_type}", + light_type=light_type, + init_pos=(0.0, 0.0, 3.0), + ) + light = self.sim.add_light(cfg=cfg) + assert light is not None, f"Failed to create light of type '{light_type}'" + assert light.num_instances == expected_num_instances + if light_type in ("sun", "direction"): + assert light.is_global, f"{light_type} should be a global light" + + def test_unknown_light_type_errors(self): + """Passing an invalid light_type raises RuntimeError.""" + cfg = LightCfg(uid="bad", light_type="invalid") + with pytest.raises(RuntimeError, match="Unsupported light type"): + self.sim.add_light(cfg=cfg) + + def test_mesh_light_empty_path_warns(self): + """Mesh light with empty mesh_path warns at creation but succeeds.""" + cfg = LightCfg( + uid="mesh_no_path", + light_type="mesh", + init_pos=(0.0, 0.0, 2.0), + ) + light = self.sim.add_light(cfg=cfg) + assert light is not None, "Mesh light should be created even without path" + + # ------------------------------------------------------------------ + # Setter type-validation tests + # ------------------------------------------------------------------ + + def test_set_direction_on_point_warns(self): + """Calling set_direction on a point light logs a warning and no-ops.""" + cfg = LightCfg(uid="pt", light_type="point", init_pos=(0, 0, 2)) + light = self.sim.add_light(cfg=cfg) + direction = torch.tensor([1.0, 0.0, 0.0]) + # Should not raise; should log a warning + try: + light.set_direction(direction) + except Exception as e: + pytest.fail(f"set_direction on point light should warn, not crash: {e}") + + def test_set_spot_angle_on_point_warns(self): + """Calling set_spot_angle on a non-spot light warns and no-ops.""" + cfg = LightCfg(uid="pt2", light_type="point", init_pos=(0, 0, 2)) + light = self.sim.add_light(cfg=cfg) + inner = torch.tensor(15.0) + outer = torch.tensor(30.0) + try: + light.set_spot_angle(inner, outer) + except Exception as e: + pytest.fail(f"set_spot_angle on point light should warn, not crash: {e}") + + def test_set_rect_wh_on_point_warns(self): + """Calling set_rect_wh on a non-rect light warns and no-ops.""" + cfg = LightCfg(uid="pt3", light_type="point", init_pos=(0, 0, 2)) + light = self.sim.add_light(cfg=cfg) + w = torch.tensor(2.0) + h = torch.tensor(3.0) + try: + light.set_rect_wh(w, h) + except Exception as e: + pytest.fail(f"set_rect_wh on point light should warn, not crash: {e}") + + def test_set_local_pose_on_direction_warns(self): + """Calling set_local_pose on a direction light warns and no-ops.""" + cfg = LightCfg( + uid="dir_light", + light_type="direction", + direction=(0.0, -1.0, 0.0), + ) + light = self.sim.add_light(cfg=cfg) + pose = torch.tensor([1.0, 2.0, 3.0]) + try: + light.set_local_pose(pose) + except Exception as e: + pytest.fail( + f"set_local_pose on direction light should warn, not crash: {e}" + ) + + def test_set_local_pose_on_sun_warns(self): + """Calling set_local_pose on a sun light warns and no-ops.""" + cfg = LightCfg( + uid="sun_light", + light_type="sun", + direction=(0.0, -1.0, 0.0), + ) + light = self.sim.add_light(cfg=cfg) + pose = torch.tensor([1.0, 2.0, 3.0]) + try: + light.set_local_pose(pose) + except Exception as e: + pytest.fail(f"set_local_pose on sun light should warn, not crash: {e}") + + # ------------------------------------------------------------------ + # Type-specific property tests + # ------------------------------------------------------------------ + + def test_spot_light_has_cone_angles(self): + """Spot light creation applies inner/outer cone angles from cfg.""" + cfg = LightCfg( + uid="spot_test", + light_type="spot", + init_pos=(0.0, 0.0, 3.0), + direction=(0.0, 0.0, -1.0), + spot_angle_inner=20.0, + spot_angle_outer=40.0, + ) + light = self.sim.add_light(cfg=cfg) + assert light is not None + # Should not crash on creation; spot_angle applied during reset() + + def test_rect_light_has_dimensions(self): + """Rect light creation applies width/height from cfg.""" + cfg = LightCfg( + uid="rect_test", + light_type="rect", + init_pos=(0.0, 0.0, 3.0), + direction=(0.0, 0.0, -1.0), + rect_width=2.0, + rect_height=1.5, + ) + light = self.sim.add_light(cfg=cfg) + assert light is not None + + @pytest.mark.parametrize("light_type", ["sun", "direction"]) + def test_global_light_no_position_in_reset(self, light_type): + """Global light (sun/direction) reset does not set position.""" + cfg = LightCfg( + uid=f"{light_type}_test", + light_type=light_type, + direction=(0.0, -1.0, 0.0), + ) + light = self.sim.add_light(cfg=cfg) + assert light is not None + assert light.is_global + assert light.num_instances == 1 + + @pytest.mark.parametrize("light_type", ["sun", "direction"]) + def test_reset_global_light_with_env_ids(self, light_type): + """reset_objects_state with env_ids does not crash for global lights.""" + cfg = LightCfg( + uid=f"global_{light_type}", + light_type=light_type, + direction=(0.0, -1.0, 0.0), + ) + self.sim.add_light(cfg=cfg) + try: + self.sim.reset_objects_state(env_ids=[1, 2]) + except Exception as e: + pytest.fail( + f"reset_objects_state with global light ({light_type}) failed: {e}" + ) + + # ------------------------------------------------------------------ + # Broadcasting tests + # ------------------------------------------------------------------ + + def test_set_direction_broadcasting(self): + """set_direction with (3,) tensor broadcasts to all instances.""" + cfg = LightCfg( + uid="spot_broadcast", + light_type="spot", + init_pos=(0, 0, 3), + direction=(0.0, 0.0, -1.0), + ) + light = self.sim.add_light(cfg=cfg) + new_dir = torch.tensor([1.0, 0.0, 0.0]) + try: + light.set_direction(new_dir) + except Exception as e: + pytest.fail(f"set_direction broadcast failed: {e}") + + def test_set_direction_with_env_ids(self): + """set_direction with (M, 3) tensor applies per-instance.""" + cfg = LightCfg( + uid="spot_env", + light_type="spot", + init_pos=(0, 0, 3), + direction=(0.0, 0.0, -1.0), + ) + light = self.sim.add_light(cfg=cfg) + env_ids = [0, 2] + dirs = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + try: + light.set_direction(dirs, env_ids=env_ids) + except Exception as e: + pytest.fail(f"set_direction per-instance failed: {e}") + + def test_enable_shadow(self): + """enable_shadow sets shadow flag on all instances.""" + cfg = LightCfg(uid="shadow_test", light_type="point", init_pos=(0, 0, 2)) + light = self.sim.add_light(cfg=cfg) + try: + light.enable_shadow(torch.tensor(0.0)) # disable + light.enable_shadow(torch.tensor(1.0)) # enable + except Exception as e: + pytest.fail(f"enable_shadow failed: {e}") + + # ------------------------------------------------------------------ + # from_dict compatibility + # ------------------------------------------------------------------ + + def test_from_dict_new_types(self): + """LightCfg.from_dict works with new type fields.""" + cfg = LightCfg.from_dict( + { + "light_type": "spot", + "color": [1.0, 0.9, 0.8], + "intensity": 100.0, + "init_pos": [0.0, 0.0, 3.0], + "direction": [0.0, 0.0, -1.0], + "spot_angle_inner": 25.0, + "spot_angle_outer": 50.0, + "uid": "from_dict_spot", + } + ) + assert cfg.light_type == "spot" + assert cfg.spot_angle_inner == 25.0 + assert cfg.spot_angle_outer == 50.0 + assert cfg.direction == (0.0, 0.0, -1.0) or cfg.direction == [0.0, 0.0, -1.0] + + # ------------------------------------------------------------------ + # Backward compatibility + # ------------------------------------------------------------------ + + def test_point_light_backward_compat(self): + """Existing point-light config pattern still works.""" + cfg_dict = { + "light_type": "point", + "color": [0.1, 0.1, 0.1], + "radius": 10.0, + "init_pos": [0.0, 0.0, 2.0], + "uid": "point_compat", + } + cfg = LightCfg.from_dict(cfg_dict) + assert cfg.init_pos == [0.0, 0.0, 2.0] + light = self.sim.add_light(cfg=cfg) + assert light is not None + assert light.num_instances == 4 + + # set_color should still work + base_color = torch.tensor([0.1, 0.1, 0.1]) + try: + light.set_color(base_color) + except Exception as e: + pytest.fail(f"set_color failed: {e}") + + # set_falloff should still work + try: + light.set_falloff(torch.tensor(100.0)) + except Exception as e: + pytest.fail(f"set_falloff failed: {e}") diff --git a/tests/sim/objects/test_rigid_object.py b/tests/sim/objects/test_rigid_object.py index 176618c9e..681cb9c89 100644 --- a/tests/sim/objects/test_rigid_object.py +++ b/tests/sim/objects/test_rigid_object.py @@ -15,6 +15,8 @@ # ---------------------------------------------------------------------------- from __future__ import annotations +from __future__ import annotations + import os import pytest @@ -370,8 +372,7 @@ def test_add_sdf_mesh(self): sdf = self.sim.add_rigid_object( cfg=RigidObjectCfg( uid="duck_sdf", - shape=MeshCfg(fpath=duck_path), - sdf_resolution=128, + shape=MeshCfg(fpath=duck_path, sdf_resolution=128), body_type="dynamic", ) ) diff --git a/tests/sim/objects/test_rigid_object_group.py b/tests/sim/objects/test_rigid_object_group.py index fd2d6f7f4..f8a6266e4 100644 --- a/tests/sim/objects/test_rigid_object_group.py +++ b/tests/sim/objects/test_rigid_object_group.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os import torch import pytest diff --git a/tests/sim/objects/test_robot.py b/tests/sim/objects/test_robot.py index 3cfb48e4f..f62a47fb0 100644 --- a/tests/sim/objects/test_robot.py +++ b/tests/sim/objects/test_robot.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os import torch import pytest @@ -258,6 +260,178 @@ def test_setter_and_getter_with_control_part(self): ) self.robot.set_qpos(qpos=dummy_qpos, name="left_arm") + def test_joint_limit_apis_support_control_parts_and_joint_ids(self): + left_arm_joint_ids = self.robot.get_joint_ids("left_arm") + selected_joint_ids = left_arm_joint_ids[:2] + + qpos_limits = self.robot.get_qpos_limits(joint_ids=selected_joint_ids) + qvel_limits = self.robot.get_qvel_limits(joint_ids=selected_joint_ids) + qf_limits = self.robot.get_qf_limits(joint_ids=selected_joint_ids) + + assert qpos_limits.shape == (10, len(selected_joint_ids), 2) + assert qvel_limits.shape == (10, len(selected_joint_ids)) + assert qf_limits.shape == (10, len(selected_joint_ids)) + + left_arm_qpos_limits = self.robot.get_qpos_limits( + name="left_arm", joint_ids=[0] + ) + left_arm_qvel_limits = self.robot.get_qvel_limits( + name="left_arm", joint_ids=[0] + ) + left_arm_qf_limits = self.robot.get_qf_limits(name="left_arm", joint_ids=[0]) + + assert left_arm_qpos_limits.shape == (10, len(left_arm_joint_ids), 2) + assert left_arm_qvel_limits.shape == (10, len(left_arm_joint_ids)) + assert left_arm_qf_limits.shape == (10, len(left_arm_joint_ids)) + + def test_joint_limit_setters_ignore_joint_ids_when_name_is_provided(self): + left_arm_joint_ids = self.robot.get_joint_ids("left_arm") + selected_joint_ids = left_arm_joint_ids[:2] + + qpos_limits = self.robot.get_qpos_limits(name="left_arm").clone() + qpos_limits[..., 0] = qpos_limits[..., 0] + 0.01 + qpos_limits[..., 1] = qpos_limits[..., 1] - 0.01 + self.robot.set_qpos_limits( + qpos_limits, + name="left_arm", + joint_ids=[left_arm_joint_ids[0]], + ) + assert torch.allclose( + self.robot.get_qpos_limits(name="left_arm"), + qpos_limits, + atol=1e-5, + ) + + qvel_limits = torch.full( + (10, len(left_arm_joint_ids)), + 0.5, + device=self.sim.device, + ) + qf_limits = torch.full( + (10, len(left_arm_joint_ids)), + 1.5, + device=self.sim.device, + ) + self.robot.set_qvel_limits( + qvel_limits, + name="left_arm", + joint_ids=[left_arm_joint_ids[0]], + ) + self.robot.set_qf_limits( + qf_limits, + name="left_arm", + joint_ids=[left_arm_joint_ids[0]], + ) + + assert torch.allclose( + self.robot.get_qvel_limits(name="left_arm"), + qvel_limits, + atol=1e-5, + ) + assert torch.allclose( + self.robot.get_qf_limits(name="left_arm"), + qf_limits, + atol=1e-5, + ) + + joint_qpos_limits = self.robot.get_qpos_limits( + joint_ids=selected_joint_ids + ).clone() + joint_qpos_limits[..., 0] = joint_qpos_limits[..., 0] + 0.02 + joint_qpos_limits[..., 1] = joint_qpos_limits[..., 1] - 0.02 + joint_qvel_limits = torch.full( + (10, len(selected_joint_ids)), + 0.65, + device=self.sim.device, + ) + joint_qf_limits = torch.full( + (10, len(selected_joint_ids)), + 1.65, + device=self.sim.device, + ) + self.robot.set_qpos_limits(joint_qpos_limits, joint_ids=selected_joint_ids) + self.robot.set_qvel_limits(joint_qvel_limits, joint_ids=selected_joint_ids) + self.robot.set_qf_limits(joint_qf_limits, joint_ids=selected_joint_ids) + + assert torch.allclose( + self.robot.get_qpos_limits(joint_ids=selected_joint_ids), + joint_qpos_limits, + atol=1e-5, + ) + assert torch.allclose( + self.robot.get_qvel_limits(joint_ids=selected_joint_ids), + joint_qvel_limits, + atol=1e-5, + ) + assert torch.allclose( + self.robot.get_qf_limits(joint_ids=selected_joint_ids), + joint_qf_limits, + atol=1e-5, + ) + + def test_qpos_limits_update_solver_limits(self): + """Test qpos limit updates are propagated to the control-part solver.""" + arm_name = "left_arm" + solver = self.robot.get_solver(arm_name) + assert solver is not None, "FAIL: expected left_arm solver to be initialized" + + asset_limits = self.robot.get_qpos_limits(name=arm_name).clone() + updated_limits = asset_limits.clone() + margin = 0.05 + updated_limits[..., 0] = torch.clamp( + updated_limits[..., 0] + margin, + asset_limits[..., 0], + asset_limits[..., 1], + ) + updated_limits[..., 1] = torch.clamp( + updated_limits[..., 1] - margin, + asset_limits[..., 0], + asset_limits[..., 1], + ) + + self.robot.set_qpos_limits(updated_limits, name=arm_name) + + solver_limits = solver.get_qpos_limits() + assert torch.allclose( + torch.tensor(solver_limits["lower_qpos_limits"], device=self.sim.device), + updated_limits[0, :, 0], + atol=1e-5, + ), "FAIL: solver lower_qpos_limits did not update with robot qpos limits" + assert torch.allclose( + torch.tensor(solver_limits["upper_qpos_limits"], device=self.sim.device), + updated_limits[0, :, 1], + atol=1e-5, + ), "FAIL: solver upper_qpos_limits did not update with robot qpos limits" + + def test_configured_qpos_limits_sync_to_solver_after_initialization(self): + """Test configured robot limits sync to the solver after it is created.""" + configured_limits = [-0.05, 0.05] + cfg = DexforceW1Cfg.from_dict( + { + "uid": "dexforce_w1_solver_limit_sync", + "version": "v021", + "arm_kind": "anthropomorphic", + "qpos_limits": {"LEFT_J[1-7]": configured_limits}, + } + ) + robot: Robot = self.sim.add_robot(cfg=cfg) + + solver = robot.get_solver("left_arm") + assert solver is not None, "FAIL: expected left_arm solver to be initialized" + + solver_limits = solver.get_qpos_limits() + expected_limits = robot.get_qpos_limits(name="left_arm")[0] + assert torch.allclose( + torch.tensor(solver_limits["lower_qpos_limits"], device=self.sim.device), + expected_limits[:, 0], + atol=1e-5, + ), "FAIL: solver lower_qpos_limits did not sync configured robot limits" + assert torch.allclose( + torch.tensor(solver_limits["upper_qpos_limits"], device=self.sim.device), + expected_limits[:, 1], + atol=1e-5, + ), "FAIL: solver upper_qpos_limits did not sync configured robot limits" + def test_robot_cfg_merge(self): from copy import deepcopy from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg diff --git a/tests/sim/objects/test_soft_object.py b/tests/sim/objects/test_soft_object.py index 081ab526a..ad27f0041 100644 --- a/tests/sim/objects/test_soft_object.py +++ b/tests/sim/objects/test_soft_object.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os from dexsim.utility.path import get_resources_data_path from embodichain.lab.sim import SimulationManager, SimulationManagerCfg diff --git a/tests/sim/objects/test_usd.py b/tests/sim/objects/test_usd.py index 8d8f95d18..7a79d2099 100644 --- a/tests/sim/objects/test_usd.py +++ b/tests/sim/objects/test_usd.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import pytest import torch diff --git a/tests/sim/planners/test_motion_generator.py b/tests/sim/planners/test_motion_generator.py index 08758e790..c38b939fc 100644 --- a/tests/sim/planners/test_motion_generator.py +++ b/tests/sim/planners/test_motion_generator.py @@ -13,6 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- + +from __future__ import annotations + import time import torch import pytest @@ -219,7 +222,7 @@ def test_create_trajectory_with_xpos(self, is_linear): target_states = [] for xpos in self.xpos_list: target_states.append( - PlanState( + PlanState.single( move_type=MoveType.EEF_MOVE, move_part=MovePart.LEFT, xpos=xpos, @@ -229,12 +232,17 @@ def test_create_trajectory_with_xpos(self, is_linear): plan_result = self.motion_gen.generate( target_states=target_states, options=options ) + assert plan_result.positions.shape == ( + 1, + self.sample_num, + 6, + ), f"Shape mismatch: {plan_result.positions.shape} != (1, {self.sample_num}, 6)" out_qpos_list = to_numpy(plan_result.positions) assert ( len(out_qpos_list) == self.sample_num ), f"Sample number mismatch: {len(out_qpos_list)} != {self.sample_num}" test_xpos = self.robot.compute_fk( - qpos=plan_result.positions[-1].unsqueeze(0), + qpos=plan_result.positions[0, -1].unsqueeze(0), name=self.arm_name, to_matrix=True, )[0] @@ -255,6 +263,7 @@ def test_create_trajectory_with_qpos(self, is_linear): control_part=self.arm_name, is_linear=is_linear, is_interpolate=True, + interpolate_nums=1, plan_opts=ToppraPlanOptions( sample_method=TrajectorySampleMethod.QUANTITY, sample_interval=20, @@ -268,7 +277,7 @@ def test_create_trajectory_with_qpos(self, is_linear): target_states = [] for qpos in self.qpos_list: target_states.append( - PlanState( + PlanState.single( move_type=MoveType.JOINT_MOVE, move_part=MovePart.LEFT, qpos=qpos, @@ -277,6 +286,11 @@ def test_create_trajectory_with_qpos(self, is_linear): plan_result = self.motion_gen.generate( target_states=target_states, options=options ) + assert plan_result.positions.shape == ( + 1, + self.sample_num, + 6, + ), f"Shape mismatch: {plan_result.positions.shape} != (1, {self.sample_num}, 6)" out_qpos_list = to_numpy(plan_result.positions) assert ( len(out_qpos_list) == self.sample_num diff --git a/tests/sim/planners/test_motion_generator_batched.py b/tests/sim/planners/test_motion_generator_batched.py new file mode 100644 index 000000000..4b950d216 --- /dev/null +++ b/tests/sim/planners/test_motion_generator_batched.py @@ -0,0 +1,80 @@ +# ---------------------------------------------------------------------------- +# 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 torch +import pytest +from unittest.mock import Mock, patch + +from embodichain.lab.sim.planners.motion_generator import ( + MotionGenerator, + MotionGenOptions, +) +from embodichain.lab.sim.planners.utils import PlanState, PlanResult, MoveType + + +def _mock_planner(b=3, n=15, dofs=6): + planner = Mock() + planner.robot.num_instances = b + planner.robot.device = torch.device("cpu") + planner.plan.return_value = PlanResult( + success=torch.ones(b, dtype=torch.bool), + positions=torch.zeros(b, n, dofs), + ) + planner.default_plan_options.return_value = None + return planner + + +class TestGenerateBatched: + def test_generate_passes_batched_states_to_planner(self): + planner = _mock_planner() + mg = MotionGenerator.__new__(MotionGenerator) + mg.planner = planner + mg.robot = planner.robot + mg.device = torch.device("cpu") + + B, dofs = 3, 6 + states = [ + PlanState.from_qpos(torch.zeros(B, dofs)), + PlanState.from_qpos(torch.ones(B, dofs)), + ] + r = mg.generate(states, MotionGenOptions(plan_opts=Mock())) + assert r.success.shape == (B,) + assert r.positions.shape == (B, 15, dofs) + # planner.plan received the batched states list as-is + _, kwargs = planner.plan.call_args + assert ( + kwargs["target_states"] is states or planner.plan.call_args[0][0] is states + ) + + +class TestInterpolateBatched: + def test_interpolate_joint_space_batched(self): + planner = _mock_planner(b=3, n=10, dofs=6) + mg = MotionGenerator.__new__(MotionGenerator) + mg.planner = planner + mg.robot = planner.robot + mg.device = torch.device("cpu") + B, N, D = 3, 4, 6 + qpos_list = torch.zeros(B, N, D) + qpos_interpolated, _ = mg.interpolate_trajectory( + control_part="arm", + xpos_list=None, + qpos_list=qpos_list, + options=MotionGenOptions(is_linear=False, interpolate_nums=10), + ) + assert qpos_interpolated.shape[0] == B diff --git a/tests/sim/planners/test_neural_batched.py b/tests/sim/planners/test_neural_batched.py new file mode 100644 index 000000000..78055af88 --- /dev/null +++ b/tests/sim/planners/test_neural_batched.py @@ -0,0 +1,197 @@ +# ---------------------------------------------------------------------------- +# 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 torch +import pytest + + +class TestNeuralParseWaypoints: + def test_parse_waypoints_batched(self): + from embodichain.lab.sim.planners.neural_planner import NeuralPlanner + from embodichain.lab.sim.planners.utils import PlanState, MoveType + + # Build a minimal planner shell by stubbing __init__ + planner = NeuralPlanner.__new__(NeuralPlanner) + planner.device = torch.device("cpu") + planner._num_waypoints = 4 + B = 3 + states = [ + PlanState.from_xpos( + torch.eye(4).unsqueeze(0).repeat(B, 1, 1) * 1.0, + move_type=MoveType.EEF_MOVE, + ) + for _ in range(2) + ] + pos, quat, mask, k = planner._parse_waypoints(states) + assert pos.shape == (B, 4, 3) + assert quat.shape == (B, 4, 4) + assert mask.shape == (B, 4) + assert k == 2 + + +class TestNeuralPlanBatched: + def test_plan_returns_batched_success(self, monkeypatch): + from embodichain.lab.sim.planners.neural_planner import ( + NeuralPlanner, + NeuralPlanOptions, + ) + from embodichain.lab.sim.planners.utils import PlanState, MoveType + + planner = NeuralPlanner.__new__(NeuralPlanner) + planner.device = torch.device("cpu") + planner._num_waypoints = 4 + planner._action_dim = 7 + planner._max_steps = 5 + planner._pos_eps = 1e9 # always reached + planner._rot_eps = 1e9 + planner._intermediate_orientation = True + planner._use_relative_obs = False + planner._obs_dim = 57 # 7+7+4*3+4*4+4+4+7 + planner.cfg = type( + "c", + (), + { + "action_scale": 0.0, + "dt": 0.01, + "control_part": "arm", + "num_arm_joints": 7, + }, + )() + + # stub actor: returns zeros so qpos never changes but eps is huge -> reached + planner._actor = lambda obs: torch.zeros(obs.shape[0], 7) + planner._normalizer = type("n", (), {"normalize": lambda self, o: o})() + + # stub robot FK + limits + class _Robot: + num_instances = 3 + device = torch.device("cpu") + + def get_qpos(self, name=None): + return torch.zeros(3, 7) + + def get_qpos_limits(self, name=None): + return (torch.zeros(7, 2),) + + def compute_fk(self, qpos=None, name=None, to_matrix=True): + m = torch.eye(4).unsqueeze(0).repeat(qpos.shape[0], 1, 1) + return ( + m + if to_matrix + else torch.cat([m[:, :3, 3], torch.zeros(qpos.shape[0], 4)], dim=-1) + ) + + planner.robot = _Robot() + + B = 3 + states = [ + PlanState.from_xpos( + torch.eye(4).unsqueeze(0).repeat(B, 1, 1), move_type=MoveType.EEF_MOVE + ) + for _ in range(2) + ] + r = planner.plan(states, NeuralPlanOptions(control_part="arm", max_steps=3)) + assert r.success.shape == (B,) + assert r.success.all().item() + assert r.positions.shape[0] == B + + +class TestNeuralEarlyConvergenceHold: + def test_converged_env_holds_qpos(self): + from embodichain.lab.sim.planners.neural_planner import ( + NeuralPlanner, + NeuralPlanOptions, + ) + from embodichain.lab.sim.planners.utils import PlanState, MoveType + + planner = NeuralPlanner.__new__(NeuralPlanner) + planner.device = torch.device("cpu") + planner._num_waypoints = 4 + planner._action_dim = 7 + planner._max_steps = 5 + planner._pos_eps = 1e9 + planner._rot_eps = 1e9 + planner._intermediate_orientation = True + planner._use_relative_obs = False + # joint(7)+ee(7)+wp_pos(4*3)+wp_quat(4*4)+onehot(4)+valid(4)+last_act(7)=57 + planner._obs_dim = 57 + planner.cfg = type( + "c", + (), + { + "action_scale": 0.1, + "dt": 0.01, + "control_part": "arm", + "num_arm_joints": 7, + }, + )() + + # actor: non-trivial action so qpos would drift if not masked + planner._actor = lambda obs: torch.ones(obs.shape[0], 7) * 0.5 + planner._normalizer = type("n", (), {"normalize": lambda self, o: o})() + + class _Robot: + num_instances = 3 + device = torch.device("cpu") + + def get_qpos(self, name=None): + return torch.zeros(3, 7) + + def get_qpos_limits(self, name=None): + return (torch.tensor([[-3.14, 3.14]] * 7),) + + def compute_fk(self, qpos=None, name=None, to_matrix=True): + m = torch.eye(4).unsqueeze(0).repeat(qpos.shape[0], 1, 1) + return ( + m + if to_matrix + else torch.cat([m[:, :3, 3], torch.zeros(qpos.shape[0], 4)], dim=-1) + ) + + planner.robot = _Robot() + + # Force env 0 to reach every step, envs 1 & 2 to never reach. + # episode_k = 2 (2 waypoints). Env 0: reached=True every step -> + # active_idx 0->1 (step0), 1->2 (step1) => converges after step 1. + # Envs 1,2: reached=False always -> never converge within max_steps. + def fake_reach(ee_pose, wp_pos, wp_quat, active_idx, episode_k): + r = torch.zeros(ee_pose.shape[0], dtype=torch.bool, device=ee_pose.device) + r[0] = True + return r + + planner._is_active_reached = fake_reach + + B = 3 + states = [ + PlanState.from_xpos( + torch.eye(4).unsqueeze(0).repeat(B, 1, 1), move_type=MoveType.EEF_MOVE + ) + for _ in range(2) + ] + r = planner.plan(states, NeuralPlanOptions(control_part="arm", max_steps=5)) + # Env 0 converges (success True); envs 1,2 do not (success False) + assert r.success[0].item() is True + assert r.success[1].item() is False + assert r.success[2].item() is False + # Env 0 converges at end of step 1 (active_idx 0->1 on step0, 1->2 on + # step1 => converged after step1). From step 2 onward, env 0's qpos + # must not change (held). + env0_traj = r.positions[0] # (T, 7) + assert torch.allclose(env0_traj[2:], env0_traj[2:3].expand(env0_traj[2:].shape)) + # Meanwhile env 1 (never converges) keeps drifting every step. + env1_traj = r.positions[1] + assert not torch.allclose(env1_traj[0], env1_traj[-1]) diff --git a/tests/sim/planners/test_neural_planner.py b/tests/sim/planners/test_neural_planner.py index 878c54ffd..df14d7545 100644 --- a/tests/sim/planners/test_neural_planner.py +++ b/tests/sim/planners/test_neural_planner.py @@ -132,7 +132,7 @@ def test_neural_planner_generate_with_fake_checkpoint(tmp_path, monkeypatch): ) ) - target_state = PlanState(move_type=MoveType.EEF_MOVE, xpos=torch.eye(4)) + target_state = PlanState.single(move_type=MoveType.EEF_MOVE, xpos=torch.eye(4)) result = motion_generator.generate( target_states=[target_state], options=MotionGenOptions( @@ -143,7 +143,7 @@ def test_neural_planner_generate_with_fake_checkpoint(tmp_path, monkeypatch): ), ) - assert result.success is True + assert result.success.all().item() assert result.positions is not None assert result.positions.shape[-1] == NUM_ARM_JOINTS assert torch.isfinite(result.positions).all() @@ -169,7 +169,9 @@ def test_neural_planner_uses_plan_opts_start_qpos(tmp_path, monkeypatch): ) custom_qpos = torch.ones(NUM_ARM_JOINTS) result = motion_generator.generate( - target_states=[PlanState(move_type=MoveType.EEF_MOVE, xpos=torch.eye(4))], + target_states=[ + PlanState.single(move_type=MoveType.EEF_MOVE, xpos=torch.eye(4)) + ], options=MotionGenOptions( plan_opts=NeuralPlanOptions( control_part="main_arm", @@ -178,8 +180,8 @@ def test_neural_planner_uses_plan_opts_start_qpos(tmp_path, monkeypatch): ), ) - assert result.success is True - assert torch.allclose(result.positions[0], custom_qpos) + assert result.success.all().item() + assert torch.allclose(result.positions[0, 0], custom_qpos) def test_neural_planner_rejects_short_start_qpos(tmp_path, monkeypatch): @@ -201,7 +203,9 @@ def test_neural_planner_rejects_short_start_qpos(tmp_path, monkeypatch): with pytest.raises(ValueError, match="policy expects"): motion_generator.generate( - target_states=[PlanState(move_type=MoveType.EEF_MOVE, xpos=torch.eye(4))], + target_states=[ + PlanState.single(move_type=MoveType.EEF_MOVE, xpos=torch.eye(4)) + ], options=MotionGenOptions( plan_opts=NeuralPlanOptions( control_part="main_arm", @@ -209,3 +213,148 @@ def test_neural_planner_rejects_short_start_qpos(tmp_path, monkeypatch): ), ), ) + + +def test_neural_planner_returns_velocities_and_accelerations(tmp_path, monkeypatch): + checkpoint_path = _create_fake_checkpoint(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + + motion_generator = MotionGenerator( + cfg=MotionGenCfg( + planner_cfg=NeuralPlannerCfg( + robot_uid="fake_robot", + checkpoint_path=checkpoint_path, + control_part="main_arm", + ) + ) + ) + + result = motion_generator.generate( + target_states=[ + PlanState.single(move_type=MoveType.EEF_MOVE, xpos=torch.eye(4)) + ], + options=MotionGenOptions( + plan_opts=NeuralPlanOptions( + control_part="main_arm", + start_qpos=torch.zeros(NUM_ARM_JOINTS), + ), + ), + ) + + assert result.velocities is not None + assert result.accelerations is not None + assert result.velocities.shape == result.positions.shape + assert result.accelerations.shape == result.positions.shape + assert torch.isfinite(result.velocities).all() + assert torch.isfinite(result.accelerations).all() + + +def test_neural_planner_finite_diff_helper(): + """Finite-difference estimates match a known polynomial trajectory.""" + b, n, dof = 2, 5, 7 + dt_value = 0.01 + positions = torch.zeros(b, n, dof) + for t in range(n): + positions[:, t] = (t * dt_value) ** 2 + dt = torch.full((b, n), dt_value) + + velocities, accelerations = NeuralPlanner._compute_vel_acc_via_finite_diff( + positions, dt + ) + + # For p(t) = t^2, v(t) = 2t, a(t) = 2. Interior central differences should + # match exactly; boundary one-sided differences deviate slightly. + expected_v = torch.zeros(b, n, dof) + for t in range(n): + expected_v[:, t] = 2.0 * (t * dt_value) + expected_a = torch.full((b, n, dof), 2.0) + + assert velocities.shape == (b, n, dof) + assert accelerations.shape == (b, n, dof) + assert torch.allclose(velocities[:, 1:-1], expected_v[:, 1:-1], atol=1e-5) + assert torch.allclose(accelerations[:, 1:-1], expected_a[:, 1:-1], atol=1e-3) + # Boundary points should still be finite and have the right sign/order. + assert torch.all(velocities[:, 0] >= 0.0) + assert torch.all(velocities[:, -1] >= velocities[:, 0]) + + +def test_motion_generator_neural_propagates_motion_gen_options(tmp_path, monkeypatch): + checkpoint_path = _create_fake_checkpoint(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + + motion_generator = MotionGenerator( + cfg=MotionGenCfg( + planner_cfg=NeuralPlannerCfg( + robot_uid="fake_robot", + checkpoint_path=checkpoint_path, + ) + ) + ) + custom_qpos = torch.ones(NUM_ARM_JOINTS) + result = motion_generator.generate( + target_states=[ + PlanState.single(move_type=MoveType.EEF_MOVE, xpos=torch.eye(4)) + ], + options=MotionGenOptions( + control_part="main_arm", + start_qpos=custom_qpos, + ), + ) + + assert result.success.all().item() + assert torch.allclose(result.positions[0, 0], custom_qpos) + + +def test_motion_generator_neural_auto_disables_interpolation( + tmp_path, monkeypatch, caplog +): + checkpoint_path = _create_fake_checkpoint(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + + motion_generator = MotionGenerator( + cfg=MotionGenCfg( + planner_cfg=NeuralPlannerCfg( + robot_uid="fake_robot", + checkpoint_path=checkpoint_path, + control_part="main_arm", + ) + ) + ) + + import logging + + with caplog.at_level(logging.WARNING): + result = motion_generator.generate( + target_states=[ + PlanState.single(move_type=MoveType.EEF_MOVE, xpos=torch.eye(4)) + ], + options=MotionGenOptions( + is_interpolate=True, + control_part="main_arm", + start_qpos=torch.zeros(NUM_ARM_JOINTS), + ), + ) + + assert result.success.all().item() + assert "is_interpolate=True is not supported with NeuralPlanner" in caplog.text + + +def test_safe_torch_load_roundtrip(tmp_path): + checkpoint = {"agent": torch.tensor([1.0, 2.0, 3.0])} + path = tmp_path / "checkpoint.pt" + torch.save(checkpoint, path) + + from embodichain.lab.sim.planners.neural_planner import _safe_torch_load + + loaded = _safe_torch_load(path, map_location=torch.device("cpu")) + assert "agent" in loaded + assert torch.equal(loaded["agent"], checkpoint["agent"]) diff --git a/tests/sim/planners/test_plan_state_batched.py b/tests/sim/planners/test_plan_state_batched.py new file mode 100644 index 000000000..a48e8e2a0 --- /dev/null +++ b/tests/sim/planners/test_plan_state_batched.py @@ -0,0 +1,87 @@ +# ---------------------------------------------------------------------------- +# 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 pytest +import torch + +from embodichain.lab.sim.planners.utils import PlanState, PlanResult, MoveType, MovePart + + +class TestPlanStateBatched: + def test_from_qpos_batched(self): + qpos = torch.zeros(4, 7) + ps = PlanState.from_qpos( + qpos, move_type=MoveType.JOINT_MOVE, move_part=MovePart.LEFT + ) + assert ps.qpos.shape == (4, 7) + assert ps.move_type == MoveType.JOINT_MOVE + + def test_from_xpos_batched(self): + xpos = torch.eye(4).unsqueeze(0).repeat(3, 1, 1) + ps = PlanState.from_xpos(xpos, move_type=MoveType.EEF_MOVE) + assert ps.xpos.shape == (3, 4, 4) + + def test_single_ctor_unsqueezes(self): + ps = PlanState.single(qpos=torch.zeros(7), move_type=MoveType.JOINT_MOVE) + assert ps.qpos.shape == (1, 7) + + +class TestPlanResultBatched: + def test_is_all_success_tensor(self): + r = PlanResult(success=torch.tensor([True, True, False])) + assert r.is_all_success() is False + + def test_is_all_success_scalar(self): + r = PlanResult(success=True) + assert r.is_all_success() is True + + def test_batched_shapes(self): + r = PlanResult( + success=torch.tensor([True, False]), + positions=torch.zeros(2, 10, 7), + velocities=torch.zeros(2, 10, 7), + accelerations=torch.zeros(2, 10, 7), + dt=torch.zeros(2, 10), + duration=torch.tensor([1.0, 0.0]), + ) + assert r.positions.shape == (2, 10, 7) + assert r.dt.shape == (2, 10) + assert r.duration.shape == (2,) + + +class TestValidateBatchConsistency: + def test_rejects_inconsistent_B(self): + # B=2 then B=3 -> inconsistent + states = [ + PlanState.from_qpos(torch.zeros(2, 7)), + PlanState.from_qpos(torch.zeros(3, 7)), + ] + from embodichain.lab.sim.planners import base_planner as bp + + with pytest.raises(ValueError): + bp._check_batch_consistency(states, expected_b=None, robot_num_instances=2) + + def test_rejects_expected_b_mismatch(self): + from embodichain.lab.sim.planners import base_planner as bp + + states = [ + PlanState.from_qpos(torch.zeros(2, 7)), + PlanState.from_qpos(torch.zeros(2, 7)), + ] + with pytest.raises(ValueError): + bp._check_batch_consistency(states, expected_b=4, robot_num_instances=4) diff --git a/tests/sim/planners/test_toppra_batched.py b/tests/sim/planners/test_toppra_batched.py new file mode 100644 index 000000000..06c65966e --- /dev/null +++ b/tests/sim/planners/test_toppra_batched.py @@ -0,0 +1,410 @@ +# ---------------------------------------------------------------------------- +# 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 gc +import math + +import numpy as np +import pytest +import torch + +from embodichain.lab.sim.planners.toppra_planner import _toppra_solve_one_env +from embodichain.lab.sim.planners.utils import TrajectorySampleMethod + + +# Note: TOPPRA's ParametrizeConstAccel robustly avoids returning None for smooth +# splines (verified empirically across tiny limits, huge displacements, zigzags, +# and plateaus), so the ``jnt_traj is None`` branch in ``_toppra_solve_one_env`` +# is defensive and not directly unit-tested here; infeasible inputs instead +# raise during path/constraint construction and are caught by the +# ``except Exception`` -> ``_empty_failure`` path. +class TestToppraWorker: + def test_solve_one_env_quantity(self): + # 2-waypoint, 6-DOF + wp = np.array([[0.0] * 6, [0.5] * 6]) + out = _toppra_solve_one_env( + waypoints=wp, + vel_constraint=1.0, + acc_constraint=2.0, + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=20, + ) + assert out["success"] is True + assert out["positions"].shape == (20, 6) + assert out["velocities"].shape == (20, 6) + assert out["dt"].shape == (20,) + + def test_solve_one_env_infeasible_exception(self): + # Single waypoint -> SplineInterpolator raises -> caught, returns failure + wp = np.array([[0.0] * 6]) + out = _toppra_solve_one_env( + waypoints=wp, + vel_constraint=1.0, + acc_constraint=2.0, + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=10, + ) + assert out["success"] is False + + def test_solve_one_env_time_sampling(self): + wp = np.array([[0.0] * 6, [0.5] * 6]) + out = _toppra_solve_one_env( + waypoints=wp, + vel_constraint=1.0, + acc_constraint=2.0, + sample_method=TrajectorySampleMethod.TIME, + sample_interval=0.05, + ) + assert out["success"] is True + assert out["positions"].shape[0] == out["n"] + assert out["n"] >= 2 + + def test_solve_one_env_same_waypoint_shortcut(self): + wp = np.array([[0.3] * 6, [0.3] * 6]) # identical + out = _toppra_solve_one_env( + waypoints=wp, + vel_constraint=1.0, + acc_constraint=2.0, + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=20, + ) + assert out["success"] is True + assert out["n"] == 2 + assert out["duration"] == 0.0 + + def test_solve_one_env_duplicate_plateau(self): + # Long plateaus of identical waypoints (e.g. from interpolating a + # segment where start_qpos equals the first target) must not make + # TOPPRA's controllable-set computation fail. + wp = np.array([[0.3] * 6] * 12 + [[0.5] * 6] * 12) + out = _toppra_solve_one_env( + waypoints=wp, + vel_constraint=1.0, + acc_constraint=2.0, + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=20, + ) + assert out["success"] is True + assert out["positions"].shape == (20, 6) + + +class TestToppraCfgFields: + def test_cfg_defaults(self): + from embodichain.lab.sim.planners.toppra_planner import ToppraPlannerCfg + + cfg = ToppraPlannerCfg(robot_uid="x") + assert cfg.max_workers is None + assert cfg.mp_context is None # None => auto-select by device + + def test_mp_context_auto_resolution(self): + # _resolve_mp_context is a pure function of (cfg value, device) — no sim + # needed. None auto-selects: fork on CPU, spawn on GPU; explicit values + # are honored as-is. + import torch + + from embodichain.lab.sim.planners.toppra_planner import ToppraPlanner + + resolve = ToppraPlanner._resolve_mp_context + # Auto: CPU -> fork, CUDA -> spawn (creating a torch.device does NOT + # initialize CUDA, so this is safe to assert without a GPU). + assert resolve(None, torch.device("cpu")) == "fork" + assert resolve(None, torch.device("cuda", 0)) == "spawn" + assert resolve(None, torch.device("cuda")) == "spawn" + # Explicit override wins regardless of device. + assert resolve("spawn", torch.device("cpu")) == "spawn" + assert resolve("fork", torch.device("cuda", 0)) == "fork" + + +class TestToppraPlanBatched: + def _make_planner(self): + from embodichain.lab.sim.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ) + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg + from embodichain.lab.sim.robots import CobotMagicCfg + + sim = SimulationManager( + SimulationManagerCfg(headless=True, device="cpu", num_envs=3) + ) + robot = sim.add_robot( + cfg=CobotMagicCfg.from_dict( + {"uid": "t", "init_pos": [0, 0, 0.7775], "init_qpos": [0.0] * 16} + ) + ) + planner = ToppraPlanner(ToppraPlannerCfg(robot_uid="t", max_workers=1)) + return planner, sim + + def test_plan_batched_quantity_uniform_N(self): + from embodichain.lab.sim.planners.utils import PlanState, TrajectorySampleMethod + from embodichain.lab.sim.planners.toppra_planner import ToppraPlanOptions + + planner, sim = self._make_planner() + try: + B, dofs = 3, 6 + wp = torch.zeros(B, dofs) + wp[:, 0] = torch.linspace(0.0, 0.4, B) + states = [ + PlanState.from_qpos(torch.zeros(B, dofs)), + PlanState.from_qpos(wp), + ] + opts = ToppraPlanOptions( + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=15, + constraints={"velocity": 1.0, "acceleration": 2.0}, + ) + r = planner.plan(states, opts) + assert r.success.shape == (B,) + assert r.success.all().item() + assert r.positions.shape == (B, 15, dofs) + finally: + del planner + gc.collect() + sim.destroy() + import embodichain.lab.sim as om + + om.SimulationManager.flush_cleanup_queue() + + def test_plan_batched_time_tailpads(self): + from embodichain.lab.sim.planners.utils import PlanState, TrajectorySampleMethod + from embodichain.lab.sim.planners.toppra_planner import ToppraPlanOptions + + planner, sim = self._make_planner() + try: + B, dofs = 3, 6 + wp = torch.zeros(B, dofs) + wp[:, 0] = torch.tensor([0.1, 0.4, 0.9]) # different durations + states = [ + PlanState.from_qpos(torch.zeros(B, dofs)), + PlanState.from_qpos(wp), + ] + opts = ToppraPlanOptions( + sample_method=TrajectorySampleMethod.TIME, + sample_interval=0.05, + constraints={"velocity": 1.0, "acceleration": 2.0}, + ) + r = planner.plan(states, opts) + assert r.success.shape == (B,) + assert r.positions.shape[0] == B + assert r.duration.shape == (B,) + + # The env with the SHORTEST duration got tail-padded; its trailing + # padded rows must equal its last real waypoint (held pose) and the + # padded tail must be constant. + shortest_env = int(r.duration.argmin().item()) + longest_env = int(r.duration.argmax().item()) + tail = r.positions[shortest_env] # (max_n, DOF) + max_n = tail.shape[0] + # Reconstruct n_real for the shortest env the same way the worker does. + n_real = max(2, int(math.ceil(r.duration[shortest_env].item() / 0.05)) + 1) + # The longest env should not be padded (its n_real == max_n). + assert r.positions[longest_env].shape[0] == max_n + # Shortest env must actually have been padded (otherwise the test + # isn't exercising the tail-pad branch). + assert n_real < max_n, "expected shortest env to be tail-padded" + held = tail[n_real - 1] # last real waypoint + padded = tail[n_real:] # all padded rows + assert torch.allclose(padded, held.expand(padded.shape)) + finally: + del planner + gc.collect() + sim.destroy() + import embodichain.lab.sim as om + + om.SimulationManager.flush_cleanup_queue() + + @pytest.mark.parametrize("mp_context", ["fork", "spawn"]) + def test_plan_batched_pool_path(self, mp_context): + # Exercise the real ProcessPoolExecutor branch (max_workers=2, B=3) + # under both start methods: 'fork' is the CPU default, 'spawn' is the + # GPU fallback. Both must plan correctly and reap workers on GC. + from embodichain.lab.sim.planners.utils import PlanState, TrajectorySampleMethod + from embodichain.lab.sim.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg + from embodichain.lab.sim.robots import CobotMagicCfg + + sim = SimulationManager( + SimulationManagerCfg(headless=True, device="cpu", num_envs=3) + ) + sim.add_robot( + cfg=CobotMagicCfg.from_dict( + {"uid": "p", "init_pos": [0, 0, 0.7775], "init_qpos": [0.0] * 16} + ) + ) + planner = ToppraPlanner( + ToppraPlannerCfg(robot_uid="p", max_workers=2, mp_context=mp_context) + ) + assert planner._mp_context == mp_context + try: + B, dofs = 3, 6 + wp = torch.zeros(B, dofs) + wp[:, 0] = torch.linspace(0.1, 0.5, B) + states = [ + PlanState.from_qpos(torch.zeros(B, dofs)), + PlanState.from_qpos(wp), + ] + opts = ToppraPlanOptions( + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=12, + constraints={"velocity": 1.0, "acceleration": 2.0}, + ) + r = planner.plan(states, opts) + assert r.success.shape == (B,) + assert r.success.all().item() + assert r.positions.shape == (B, 12, dofs) + finally: + del planner + gc.collect() + sim.destroy() + import embodichain.lab.sim as om + + om.SimulationManager.flush_cleanup_queue() + + @pytest.mark.parametrize("mp_context", ["fork", "spawn"]) + def test_workers_reaped_on_gc(self, mp_context): + # Regression: abandoning the planner must reap its worker processes, + # under both 'fork' (CPU default) and 'spawn' (GPU fallback). + # + # Workers install prctl(PR_SET_PDEATHSIG) in _worker_init, so the kernel + # kills them the moment the parent process dies — including the + # os._exit(0) path taken by SimulationManager.destroy(), which skips + # every Python finalizer. That covers process exit. For in-process GC + # of an abandoned planner (e.g. between tests in a long-running + # session), __del__ -> _shutdown_pool() must terminate workers + # synchronously so none leak. + from embodichain.lab.sim.planners.utils import PlanState, TrajectorySampleMethod + from embodichain.lab.sim.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg + from embodichain.lab.sim.robots import CobotMagicCfg + + sim = SimulationManager( + SimulationManagerCfg(headless=True, device="cpu", num_envs=3) + ) + sim.add_robot( + cfg=CobotMagicCfg.from_dict( + { + "uid": "close_reap", + "init_pos": [0, 0, 0.7775], + "init_qpos": [0.0] * 16, + } + ) + ) + planner = ToppraPlanner( + ToppraPlannerCfg( + robot_uid="close_reap", max_workers=2, mp_context=mp_context + ) + ) + assert planner._mp_context == mp_context + worker_processes = [] + try: + B, dofs = 3, 6 + wp = torch.zeros(B, dofs) + wp[:, 0] = torch.linspace(0.1, 0.5, B) + states = [ + PlanState.from_qpos(torch.zeros(B, dofs)), + PlanState.from_qpos(wp), + ] + opts = ToppraPlanOptions( + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=12, + constraints={"velocity": 1.0, "acceleration": 2.0}, + ) + r = planner.plan(states, opts) + assert r.success.all().item() + + # Capture the worker processes created by the executor. + worker_processes = list(planner._pool._processes.values()) + assert len(worker_processes) == 2 + finally: + # Abandon the planner; __del__ -> _shutdown_pool must reap workers. + del planner + gc.collect() + for proc in worker_processes: + assert not proc.is_alive(), ( + f"TOPPRA worker process {proc.pid} survived planner GC; " + "worker processes must be reaped when the planner is collected." + ) + sim.destroy() + import embodichain.lab.sim as om + + om.SimulationManager.flush_cleanup_queue() + + +@pytest.mark.slow +class TestToppraNumericalRegression: + def test_batched_equals_inline_single(self): + from embodichain.lab.sim.planners.toppra_planner import ( + _toppra_solve_one_env, + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim.planners.utils import PlanState, TrajectorySampleMethod + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg + from embodichain.lab.sim.robots import CobotMagicCfg + + sim = SimulationManager( + SimulationManagerCfg(headless=True, device="cpu", num_envs=4) + ) + sim.add_robot( + cfg=CobotMagicCfg.from_dict( + {"uid": "r", "init_pos": [0, 0, 0.7775], "init_qpos": [0.0] * 16} + ) + ) + planner = ToppraPlanner(ToppraPlannerCfg(robot_uid="r", max_workers=1)) + try: + B, dofs = 4, 6 + wp = torch.zeros(B, dofs) + wp[:, 0] = torch.linspace(0.1, 0.6, B) + states = [ + PlanState.from_qpos(torch.zeros(B, dofs)), + PlanState.from_qpos(wp), + ] + opts = ToppraPlanOptions( + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=20, + constraints={"velocity": 1.0, "acceleration": 2.0}, + ) + r = planner.plan(states, opts) + # Compare each env to the inline single-env solve + for b in range(B): + single = _toppra_solve_one_env( + np.stack([np.zeros(dofs), wp[b].numpy()]), + 1.0, + 2.0, + TrajectorySampleMethod.QUANTITY, + 20, + ) + assert np.allclose( + r.positions[b].cpu().numpy(), single["positions"], atol=1e-5 + ) + finally: + del planner + gc.collect() + sim.destroy() + import embodichain.lab.sim as om + + om.SimulationManager.flush_cleanup_queue() diff --git a/tests/sim/planners/test_toppra_planner.py b/tests/sim/planners/test_toppra_planner.py index 31662fc8e..517f9cb84 100644 --- a/tests/sim/planners/test_toppra_planner.py +++ b/tests/sim/planners/test_toppra_planner.py @@ -3,8 +3,19 @@ # # 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 torch import pytest import numpy as np @@ -61,7 +72,10 @@ def test_initialization(self): assert self.planner.device == torch.device("cpu") def test_plan_basic(self): - target_states = [PlanState(qpos=torch.zeros(6)), PlanState(qpos=torch.zeros(6))] + target_states = [ + PlanState.single(qpos=torch.zeros(6)), + PlanState.single(qpos=torch.zeros(6)), + ] opts = ToppraPlanOptions( sample_method=TrajectorySampleMethod.TIME, @@ -69,10 +83,11 @@ def test_plan_basic(self): constraints={"velocity": 1.0, "acceleration": 2.0}, ) result = self.planner.plan(target_states, options=opts) - assert result.success is True + assert result.success.all().item() assert result.positions is not None assert result.velocities is not None assert result.accelerations is not None + assert result.positions.shape[0] == 1 # Check constraints is_satisfied = self.planner.is_satisfied_constraint( @@ -81,7 +96,27 @@ def test_plan_basic(self): assert is_satisfied is True def test_trivial_trajectory(self): - target_states = [PlanState(qpos=torch.zeros(6)), PlanState(qpos=torch.zeros(6))] + target_states = [ + PlanState.single(qpos=torch.zeros(6)), + PlanState.single(qpos=torch.zeros(6)), + ] + + opts = ToppraPlanOptions( + sample_method=TrajectorySampleMethod.TIME, + sample_interval=0.1, + constraints={"velocity": 1.0, "acceleration": 2.0}, + ) + result = self.planner.plan(target_states, options=opts) + assert result.success.all().item() + assert result.positions.shape == (1, 2, 6) + assert result.duration.item() == 0.0 + + def test_single_env_does_not_spawn_pool(self): + # Single-env plans must stay inline and never create a ProcessPoolExecutor. + target_states = [ + PlanState.single(qpos=torch.zeros(6)), + PlanState.single(qpos=torch.zeros(6)), + ] opts = ToppraPlanOptions( sample_method=TrajectorySampleMethod.TIME, @@ -89,9 +124,8 @@ def test_trivial_trajectory(self): constraints={"velocity": 1.0, "acceleration": 2.0}, ) result = self.planner.plan(target_states, options=opts) - assert result.success is True - assert len(result.positions) == 2 - assert result.duration == 0.0 + assert result.success.all().item() + assert self.planner._pool is None if __name__ == "__main__": diff --git a/tests/sim/sensors/test_camera.py b/tests/sim/sensors/test_camera.py index 3e9af436f..03cacf576 100644 --- a/tests/sim/sensors/test_camera.py +++ b/tests/sim/sensors/test_camera.py @@ -14,6 +14,8 @@ # limitations under the License. # ----------------------------------------------------------------------------, +from __future__ import annotations + import pytest import torch import os @@ -26,29 +28,42 @@ from embodichain.lab.sim.cfg import ArticulationCfg, RenderCfg from embodichain.data import get_data_path -NUM_ENVS = 4 +FULL_NUM_ENVS = 4 +FULL_WIDTH = 640 +FULL_HEIGHT = 480 +SMOKE_NUM_ENVS = 1 +SMOKE_WIDTH = 160 +SMOKE_HEIGHT = 120 ART_PATH = "SlidingBoxDrawer/SlidingBoxDrawer.urdf" class CameraTest: - def setup_simulation(self, device, renderer="hybrid"): + def setup_simulation( + self, + sim_device, + renderer="hybrid", + num_envs=FULL_NUM_ENVS, + width=FULL_WIDTH, + height=FULL_HEIGHT, + enable_auxiliary_data=True, + ): # Setup SimulationManager config = SimulationManagerCfg( headless=True, - device=device, + device=sim_device, render_cfg=RenderCfg(renderer=renderer), - num_envs=NUM_ENVS, + num_envs=num_envs, ) self.sim = SimulationManager(config) # Create batch of cameras cfg_dict = { "sensor_type": "Camera", - "width": 640, - "height": 480, - "enable_mask": True, - "enable_depth": True, - "enable_normal": True, - "enable_position": True, + "width": width, + "height": height, + "enable_mask": enable_auxiliary_data, + "enable_depth": enable_auxiliary_data, + "enable_normal": enable_auxiliary_data, + "enable_position": enable_auxiliary_data, } cfg = SensorCfg.from_dict(cfg_dict) self.camera: Camera = self.sim.add_sensor(cfg) @@ -68,25 +83,34 @@ def test_get_data(self): assert key in data, f"Missing key in camera data: {key}" # Check if the data shape matches the expected shape - assert data["color"].shape == (NUM_ENVS, 480, 640, 4), "RGB data shape mismatch" + assert data["color"].shape == ( + FULL_NUM_ENVS, + FULL_HEIGHT, + FULL_WIDTH, + 4, + ), "RGB data shape mismatch" assert data["depth"].shape == ( - NUM_ENVS, - 480, - 640, + FULL_NUM_ENVS, + FULL_HEIGHT, + FULL_WIDTH, ), "Depth data shape mismatch" assert data["normal"].shape == ( - NUM_ENVS, - 480, - 640, + FULL_NUM_ENVS, + FULL_HEIGHT, + FULL_WIDTH, 3, ), "Normal data shape mismatch" assert data["position"].shape == ( - NUM_ENVS, - 480, - 640, + FULL_NUM_ENVS, + FULL_HEIGHT, + FULL_WIDTH, 3, ), "Position data shape mismatch" - assert data["mask"].shape == (NUM_ENVS, 480, 640), "Mask data shape mismatch" + assert data["mask"].shape == ( + FULL_NUM_ENVS, + FULL_HEIGHT, + FULL_WIDTH, + ), "Mask data shape mismatch" # Check if the data types are correct assert data["color"].dtype == torch.uint8, "Color data type mismatch" @@ -131,7 +155,7 @@ def test_set_intrinsics(self): device=self.sim.device, ) .unsqueeze(0) - .repeat(NUM_ENVS, 1) + .repeat(FULL_NUM_ENVS, 1) ) # Set new intrinsic parameters for all environments @@ -155,30 +179,36 @@ def teardown_method(self): gc.collect() -class TestCameraHybrid(CameraTest): - def setup_method(self): - - self.setup_simulation("cpu", renderer="hybrid") - - class TestCameraHybridCUDA(CameraTest): def setup_method(self): self.setup_simulation("cuda", renderer="hybrid") -class TestCameraFastRT(CameraTest): - def setup_method(self): - self.setup_simulation("cpu", renderer="fast-rt") - - -class TestCameraFastRTCUDA(CameraTest): - def setup_method(self): - - self.setup_simulation("cuda", renderer="fast-rt") +@pytest.mark.parametrize( + ("sim_device", "renderer"), + [("cpu", "hybrid"), ("cpu", "fast-rt"), ("cuda", "fast-rt")], +) +def test_camera_backend_smoke(sim_device, renderer): + """Check that each remaining backend/device pair renders a color frame.""" + test = CameraTest() + test.setup_simulation( + sim_device, + renderer, + num_envs=SMOKE_NUM_ENVS, + width=SMOKE_WIDTH, + height=SMOKE_HEIGHT, + enable_auxiliary_data=False, + ) + try: + test.camera.update() + data = test.camera.get_data() + assert data["color"].shape == (SMOKE_NUM_ENVS, SMOKE_HEIGHT, SMOKE_WIDTH, 4) + finally: + test.teardown_method() if __name__ == "__main__": - test = TestCameraFastRT() + test = TestCameraHybridCUDA() test.setup_method() test.test_attach_to_parent() diff --git a/tests/sim/sensors/test_contact.py b/tests/sim/sensors/test_contact.py index 6e0742647..e134b47ed 100644 --- a/tests/sim/sensors/test_contact.py +++ b/tests/sim/sensors/test_contact.py @@ -14,6 +14,8 @@ # limitations under the License. # ----------------------------------------------------------------------------, +from __future__ import annotations + import pytest import torch from embodichain.lab.sim import SimulationManager, SimulationManagerCfg @@ -36,13 +38,15 @@ from embodichain.lab.sim.robots import URRobotCfg NUM_ENVS = 4 +CONTACT_TEST_WIDTH = 320 +CONTACT_TEST_HEIGHT = 240 class ContactTest: def setup_simulation(self, device, renderer="hybrid"): sim_cfg = SimulationManagerCfg( - width=1920, - height=1080, + width=CONTACT_TEST_WIDTH, + height=CONTACT_TEST_HEIGHT, num_envs=2, headless=True, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) diff --git a/tests/sim/sensors/test_stereo.py b/tests/sim/sensors/test_stereo.py index c59b8cb14..71818f828 100644 --- a/tests/sim/sensors/test_stereo.py +++ b/tests/sim/sensors/test_stereo.py @@ -14,6 +14,8 @@ # limitations under the License. # ----------------------------------------------------------------------------, +from __future__ import annotations + import pytest import torch @@ -22,28 +24,39 @@ from embodichain.lab.sim.sensors import StereoCamera, SensorCfg NUM_ENVS = 4 +SMOKE_NUM_ENVS = 1 +SMOKE_WIDTH = 160 +SMOKE_HEIGHT = 120 class StereoCameraTest: - def setup_simulation(self, device, renderer="hybrid"): + def setup_simulation( + self, + sim_device, + renderer="hybrid", + num_envs=NUM_ENVS, + width=640, + height=480, + enable_auxiliary_data=True, + ): # Setup SimulationManager config = SimulationManagerCfg( headless=True, - device=device, - num_envs=NUM_ENVS, + device=sim_device, + num_envs=num_envs, render_cfg=RenderCfg(renderer=renderer), ) self.sim = SimulationManager(config) # Create batch of cameras cfg_dict = { "sensor_type": "StereoCamera", - "width": 640, - "height": 480, - "enable_mask": True, - "enable_depth": True, - "enable_normal": True, - "enable_position": True, - "enable_disparity": True, + "width": width, + "height": height, + "enable_mask": enable_auxiliary_data, + "enable_depth": enable_auxiliary_data, + "enable_normal": enable_auxiliary_data, + "enable_position": enable_auxiliary_data, + "enable_disparity": enable_auxiliary_data, "left_to_right_pos": (0.1, 0.0, 0.0), } cfg = SensorCfg.from_dict(cfg_dict) @@ -158,25 +171,30 @@ def teardown_method(self): gc.collect() -class TestStereoCameraHybrid(StereoCameraTest): - def setup_method(self): - - self.setup_simulation("cpu", renderer="hybrid") - - class TestStereoCameraHybridCUDA(StereoCameraTest): def setup_method(self): self.setup_simulation("cuda", renderer="hybrid") -class TestStereoCameraFastRT(StereoCameraTest): - def setup_method(self): - - self.setup_simulation("cpu", renderer="fast-rt") - - -class TestStereoCameraFastRTCUDA(StereoCameraTest): - def setup_method(self): - - self.setup_simulation("cuda", renderer="fast-rt") +@pytest.mark.parametrize( + ("sim_device", "renderer"), + [("cpu", "hybrid"), ("cpu", "fast-rt"), ("cuda", "fast-rt")], +) +def test_stereo_camera_backend_smoke(sim_device, renderer): + """Check that each remaining backend/device pair renders a color frame.""" + test = StereoCameraTest() + test.setup_simulation( + sim_device, + renderer, + num_envs=SMOKE_NUM_ENVS, + width=SMOKE_WIDTH, + height=SMOKE_HEIGHT, + enable_auxiliary_data=False, + ) + try: + test.camera.update() + data = test.camera.get_data() + assert data["color"].shape == (SMOKE_NUM_ENVS, SMOKE_HEIGHT, SMOKE_WIDTH, 4) + finally: + test.teardown_method() diff --git a/tests/sim/solvers/test_differential_solver.py b/tests/sim/solvers/test_differential_solver.py index 1c49c8e9d..e705c8754 100644 --- a/tests/sim/solvers/test_differential_solver.py +++ b/tests/sim/solvers/test_differential_solver.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os import torch import pytest @@ -122,6 +124,7 @@ def test_differential_solver(self, arm_name: str): def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() + SimulationManager.flush_cleanup_queue() class TestDifferentialSolver(BaseSolverTest): diff --git a/tests/sim/solvers/test_neural_ik_solver.py b/tests/sim/solvers/test_neural_ik_solver.py index eee25a399..40766c340 100644 --- a/tests/sim/solvers/test_neural_ik_solver.py +++ b/tests/sim/solvers/test_neural_ik_solver.py @@ -15,29 +15,17 @@ # ---------------------------------------------------------------------------- from __future__ import annotations -import math -import os - import numpy as np import pytest import torch -from embodichain.data import get_data_path from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RobotCfg from embodichain.lab.sim.objects import Robot +from embodichain.lab.sim.robots.franka_panda import FrankaPandaCfg +from embodichain.lab.sim.solvers import NeuralIKSolverCfg from embodichain.lab.sim.solvers.neural_ik_solver import _build_mlp from embodichain.utils.utility import reset_all_seeds -_c = math.cos(-math.pi / 4) -_s = math.sin(-math.pi / 4) -TCP = [ - [_c, -_s, 0.0, 0.0], - [_s, _c, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.1034], - [0.0, 0.0, 0.0, 1.0], -] - NUM_ARM_JOINTS = 7 OBS_DIM = 2 * NUM_ARM_JOINTS + 14 # 28 HIDDEN_DIMS = [256, 256] @@ -67,49 +55,38 @@ def _setup(self, tmp_path): config = SimulationManagerCfg(headless=True, device="cpu") self.sim = SimulationManager(config) - urdf = get_data_path("Franka/Panda/PandaWithHand.urdf") - assert os.path.isfile(urdf) - - cfg_dict = { - "fpath": urdf, - "control_parts": { - "main_arm": [ - "Joint1", - "Joint2", - "Joint3", - "Joint4", - "Joint5", - "Joint6", - "Joint7", - ], - }, - "solver_cfg": { - "main_arm": { - "class_type": "NeuralIKSolver", - "end_link_name": "ee_link", - "root_link_name": "base_link", - "tcp": TCP, - "checkpoint_path": checkpoint_path, - "num_arm_joints": NUM_ARM_JOINTS, - "max_steps": 30, - "action_scale": 0.2, - "hidden_dims": HIDDEN_DIMS, - "pos_eps": 0.1, - "rot_eps": 0.5, - }, - }, - } - - self.robot: Robot = self.sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + 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=NUM_ARM_JOINTS, + max_steps=30, + action_scale=0.2, + hidden_dims=HIDDEN_DIMS, + pos_eps=0.1, + rot_eps=0.5, + ) + + self.robot = self.sim.add_robot(cfg=cfg) self.sim.update(step=100) def teardown_method(self): if self.sim is not None: self.sim.destroy() + SimulationManager.flush_cleanup_queue() + self.sim = None + self.robot = None def _make_solver_input(self): """Create a standard qpos and its FK target for solver tests.""" - arm_name = "main_arm" + 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, @@ -123,7 +100,7 @@ def test_ik_interface(self, tmp_path): """Verify compute_ik returns correct shapes and types.""" reset_all_seeds(0) self._setup(tmp_path) - arm_name = "main_arm" + 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], diff --git a/tests/sim/solvers/test_opw_solver.py b/tests/sim/solvers/test_opw_solver.py index 8636f3549..e9336dc9c 100644 --- a/tests/sim/solvers/test_opw_solver.py +++ b/tests/sim/solvers/test_opw_solver.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import torch import pytest import numpy as np @@ -189,6 +191,7 @@ def test_ik(self, arm_name: str): def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() + SimulationManager.flush_cleanup_queue() class TestOPWSolver(BaseSolverTest): diff --git a/tests/sim/solvers/test_pink_solver.py b/tests/sim/solvers/test_pink_solver.py index 50c510b97..e1b34fc3e 100644 --- a/tests/sim/solvers/test_pink_solver.py +++ b/tests/sim/solvers/test_pink_solver.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os import torch import pytest @@ -126,6 +128,7 @@ def test_differential_solver(self): def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() + SimulationManager.flush_cleanup_queue() @pytest.mark.skip(reason="Skipping Pink tests temporarily") diff --git a/tests/sim/solvers/test_pinocchio_solver.py b/tests/sim/solvers/test_pinocchio_solver.py index bd0236d84..9daa63327 100644 --- a/tests/sim/solvers/test_pinocchio_solver.py +++ b/tests/sim/solvers/test_pinocchio_solver.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os import torch import pytest @@ -107,6 +109,7 @@ def test_ik(self, arm_name: str): def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() + SimulationManager.flush_cleanup_queue() class TestPinocchioSolver(BaseSolverTest): diff --git a/tests/sim/solvers/test_pytorch_solver.py b/tests/sim/solvers/test_pytorch_solver.py index e3657648f..af10acfbb 100644 --- a/tests/sim/solvers/test_pytorch_solver.py +++ b/tests/sim/solvers/test_pytorch_solver.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os import torch import pytest @@ -171,6 +173,7 @@ def test_ik(self, arm_name: str): def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() + SimulationManager.flush_cleanup_queue() class TestPytorchSolver(BaseSolverTest): diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index 6b9fdbc22..139d89545 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os import torch @@ -116,6 +118,75 @@ def test_ik( fk_xpos, ik_xpos, atol=1e-3, rtol=1e-3 ), f"FK and IK results do not match for {solver_key}" + def test_update_with_robot_limit_intersects_existing_solver_limits(self): + """Test robot limit sync only tightens solver limits and never widens them.""" + solver_key = next(iter(self.solver)) + solver = self.solver[solver_key] + solver_limits = solver.get_qpos_limits() + + configured_lower = torch.tensor( + solver_limits["lower_qpos_limits"], + dtype=torch.float32, + device=solver.device, + ) + configured_upper = torch.tensor( + solver_limits["upper_qpos_limits"], + dtype=torch.float32, + device=solver.device, + ) + + looser_robot_limits = torch.stack( + (configured_lower - 0.1, configured_upper + 0.1), dim=-1 + ) + solver.update_with_robot_limit(looser_robot_limits) + looser_sync_limits = solver.get_qpos_limits() + assert torch.allclose( + torch.tensor( + looser_sync_limits["lower_qpos_limits"], + dtype=torch.float32, + device=solver.device, + ), + configured_lower, + atol=1e-5, + ), "FAIL: robot sync widened solver lower_qpos_limits" + assert torch.allclose( + torch.tensor( + looser_sync_limits["upper_qpos_limits"], + dtype=torch.float32, + device=solver.device, + ), + configured_upper, + atol=1e-5, + ), "FAIL: robot sync widened solver upper_qpos_limits" + + margin = torch.minimum( + torch.full_like(configured_lower, 0.05), + 0.25 * (configured_upper - configured_lower), + ) + tighter_robot_limits = torch.stack( + (configured_lower + margin, configured_upper - margin), dim=-1 + ) + solver.update_with_robot_limit(tighter_robot_limits) + tighter_sync_limits = solver.get_qpos_limits() + assert torch.allclose( + torch.tensor( + tighter_sync_limits["lower_qpos_limits"], + dtype=torch.float32, + device=solver.device, + ), + tighter_robot_limits[:, 0], + atol=1e-5, + ), "FAIL: robot sync did not tighten solver lower_qpos_limits" + assert torch.allclose( + torch.tensor( + tighter_sync_limits["upper_qpos_limits"], + dtype=torch.float32, + device=solver.device, + ), + tighter_robot_limits[:, 1], + atol=1e-5, + ), "FAIL: robot sync did not tighten solver upper_qpos_limits" + @classmethod def teardown_class(cls): if cls.solver is not None: @@ -274,6 +345,7 @@ def test_robot_ik(self, arm_name: str): def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() + SimulationManager.flush_cleanup_queue() class TestSRSCPUSolver(BaseSolverTest): diff --git a/tests/sim/solvers/test_ur_solver.py b/tests/sim/solvers/test_ur_solver.py index 3bd6b1962..620cd527e 100644 --- a/tests/sim/solvers/test_ur_solver.py +++ b/tests/sim/solvers/test_ur_solver.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import torch import pytest import numpy as np @@ -189,6 +191,7 @@ def test_ik(self): def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() + SimulationManager.flush_cleanup_queue() class TestURSolverCUDA(BaseSolverTest): diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index aa07dfc22..e981e8ecd 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -104,6 +104,18 @@ def _make_sim_manager(window: object | None = None) -> SimulationManager: return sim +def test_window_camera_pose_to_look_at_uses_dexsim_world_up() -> None: + """Captured look-at snippets preserve DexSim's default Z-up controls.""" + pose = np.eye(4, dtype=np.float32) + pose[:3, 3] = np.array([1.0, 2.0, 3.0], dtype=np.float32) + + eye, look_at, up = SimulationManager._window_camera_pose_to_look_at(pose) + + np.testing.assert_allclose(eye, [1.0, 2.0, 3.0]) + np.testing.assert_allclose(look_at, [1.0, 2.0, 2.0]) + np.testing.assert_allclose(up, [0.0, 0.0, 1.0]) + + def test_start_window_record_rejects_invalid_parameters() -> None: sim = _make_sim_manager() diff --git a/tests/sim/test_sim_manager_cfg.py b/tests/sim/test_sim_manager_cfg.py index 9fb6c616c..5236b311e 100644 --- a/tests/sim/test_sim_manager_cfg.py +++ b/tests/sim/test_sim_manager_cfg.py @@ -19,7 +19,7 @@ import torch from embodichain.lab.sim import SimulationManagerCfg -from embodichain.lab.sim.cfg import NewtonPhysicsCfg +from embodichain.lab.sim.cfg import NewtonPhysicsCfg, WindowCameraPoseCfg def test_physics_runtime_fields_are_stored_on_physics_cfg() -> None: @@ -51,6 +51,17 @@ def test_simulation_manager_cfg_keeps_legacy_physics_accessors() -> None: assert cfg.physics_cfg.device == "cuda:0" +def test_simulation_manager_cfg_initializes_window_camera_pose() -> None: + window_camera_pose = WindowCameraPoseCfg( + enable_hotkey=False, + convert_to_look_at=False, + ) + + cfg = SimulationManagerCfg(window_camera_pose=window_camera_pose) + + assert cfg.window_camera_pose == window_camera_pose + + def test_newton_physics_cfg_uses_device() -> None: cfg = NewtonPhysicsCfg(device="cuda:1") diff --git a/tests/sim/utility/test_sim_utils.py b/tests/sim/utility/test_sim_utils.py new file mode 100644 index 000000000..5ccdbbd26 --- /dev/null +++ b/tests/sim/utility/test_sim_utils.py @@ -0,0 +1,77 @@ +# ---------------------------------------------------------------------------- +# 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 embodichain.lab.sim.cfg import RigidObjectCfg +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.utility.sim_utils import _load_rigid_mesh_prototype + + +class _FakeActor: + def add_rigidbody(self, *args, **kwargs) -> None: + pass + + +class _FakeArena: + def __init__(self) -> None: + self.load_actor_called = False + self.acd_method: str | None = None + + def load_actor(self, *args, **kwargs) -> _FakeActor: + self.load_actor_called = True + return _FakeActor() + + def load_actor_with_acd(self, *args, method: str, **kwargs) -> _FakeActor: + self.acd_method = method + return _FakeActor() + + +def test_load_rigid_mesh_uses_shape_collision_defaults() -> None: + arena = _FakeArena() + cfg = RigidObjectCfg(uid="mesh", shape=MeshCfg(fpath="mesh.obj")) + + _load_rigid_mesh_prototype( + arena, + cfg, + cache_dir=None, + body_type=None, + is_newton_backend=False, + ) + + assert arena.load_actor_called + + +def test_load_rigid_mesh_forwards_shape_acd_method() -> None: + arena = _FakeArena() + cfg = RigidObjectCfg( + uid="mesh", + shape=MeshCfg( + fpath="mesh.obj", + max_convex_hull_num=2, + acd_method="vhacd", + ), + ) + + _load_rigid_mesh_prototype( + arena, + cfg, + cache_dir=None, + body_type=None, + is_newton_backend=False, + ) + + assert arena.acd_method == "vhacd" diff --git a/tests/sim/utility/test_workspace_analyze.py b/tests/sim/utility/test_workspace_analyze.py index 1c952899e..cae280aa3 100644 --- a/tests/sim/utility/test_workspace_analyze.py +++ b/tests/sim/utility/test_workspace_analyze.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import torch import pytest import numpy as np @@ -79,6 +81,7 @@ def setup_simulation(self): def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() + SimulationManager.flush_cleanup_queue() class TestJointWorkspaceAnalyzeTest(BaseWorkspaceAnalyzeTest): diff --git a/tests/toolkits/test_batch_convex_collision.py b/tests/toolkits/test_batch_convex_collision.py index 291e15e16..67116ff5f 100644 --- a/tests/toolkits/test_batch_convex_collision.py +++ b/tests/toolkits/test_batch_convex_collision.py @@ -13,6 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- + +from __future__ import annotations + import torch from embodichain.data import get_data_path import trimesh diff --git a/tests/toolkits/test_grasp_pose_generator.py b/tests/toolkits/test_grasp_pose_generator.py index 28dc77a09..a4c4d5cd3 100644 --- a/tests/toolkits/test_grasp_pose_generator.py +++ b/tests/toolkits/test_grasp_pose_generator.py @@ -19,9 +19,12 @@ in a simulated environment using the SimulationManager and grasp planning utilities. """ +from __future__ import annotations + import argparse import numpy as np import time +import pytest import torch from embodichain.lab.sim import SimulationManager, SimulationManagerCfg @@ -196,66 +199,73 @@ def get_grasp_traj(sim: SimulationManager, robot: Robot, grasp_xpos: torch.Tenso return interp_trajectory +@pytest.mark.requires_sim +@pytest.mark.gpu +@pytest.mark.slow def test_grasp_pose_generator(): sim = initialize_simulation() - robot = create_robot(sim, position=[0.0, 0.0, 0.0]) - mug = create_mug(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, - ) + try: + robot = create_robot(sim, position=[0.0, 0.0, 0.0]) + mug = create_mug(sim) - gripper_collision_cfg = GripperCollisionCfg( - max_open_length=0.088, finger_length=0.078, point_sample_dense=0.012 - ) + # 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, + ) - vertices = mug.get_vertices(env_ids=[0], scale=True)[0] - triangles = mug.get_triangles(env_ids=[0])[0] - grasp_generator = GraspGenerator( - vertices=vertices, - triangles=triangles, - cfg=grasp_cfg, - gripper_collision_cfg=gripper_collision_cfg, - ) + gripper_collision_cfg = GripperCollisionCfg( + max_open_length=0.088, finger_length=0.078, point_sample_dense=0.012 + ) - # Annotate grasp region (populates internal antipodal point pairs) - grasp_generator.annotate() + vertices = mug.get_vertices(env_ids=[0], scale=True)[0] + triangles = mug.get_triangles(env_ids=[0])[0] + grasp_generator = GraspGenerator( + vertices=vertices, + triangles=triangles, + cfg=grasp_cfg, + gripper_collision_cfg=gripper_collision_cfg, + ) - # Compute grasp poses per environment - approach_direction = torch.tensor( - [0, 0, -1], dtype=torch.float32, device=sim.device - ) - obj_poses = mug.get_local_pose(to_matrix=True) - grasp_xpos_list = [] - - 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=False, + # Annotate grasp region (populates internal antipodal point pairs) + grasp_generator.annotate() + + # Compute grasp poses per environment + 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)) - - grasp_xpos = torch.cat(grasp_xpos_list, dim=0) - grab_traj = get_grasp_traj(sim, robot, grasp_xpos) - assert grasp_xpos.shape == torch.Size([1, 4, 4]) - assert grab_traj.shape == torch.Size([1, 200, 8]) + obj_poses = mug.get_local_pose(to_matrix=True) + grasp_xpos_list = [] + + 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=False, + ) + 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)) + + grasp_xpos = torch.cat(grasp_xpos_list, dim=0) + grab_traj = get_grasp_traj(sim, robot, grasp_xpos) + assert grasp_xpos.shape == torch.Size([1, 4, 4]) + assert grab_traj.shape == torch.Size([1, 200, 8]) + finally: + sim.destroy() + SimulationManager.flush_cleanup_queue() if __name__ == "__main__": diff --git a/tests/utils/test_device_utils.py b/tests/utils/test_device_utils.py new file mode 100644 index 000000000..94f5034e5 --- /dev/null +++ b/tests/utils/test_device_utils.py @@ -0,0 +1,48 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest +import torch + +from embodichain.utils.device_utils import standardize_device_string + + +class TestStandardizeDeviceString: + def test_cpu_string(self): + assert standardize_device_string("cpu") == "cpu" + + def test_cpu_torch_device(self): + assert standardize_device_string(torch.device("cpu")) == "cpu" + + def test_cuda_without_index(self): + assert standardize_device_string("cuda") == "cuda:0" + + def test_cuda_without_index_torch_device(self): + assert standardize_device_string(torch.device("cuda")) == "cuda:0" + + def test_cuda_zero(self): + assert standardize_device_string("cuda:0") == "cuda:0" + + def test_cuda_zero_torch_device(self): + assert standardize_device_string(torch.device("cuda:0")) == "cuda:0" + + def test_cuda_non_zero_index(self): + assert standardize_device_string("cuda:4") == "cuda:4" + + def test_cuda_non_zero_index_torch_device(self): + assert standardize_device_string(torch.device("cuda:4")) == "cuda:4"