From 2319d6ecd833b4b41d298f0977618447c1178d33 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Fri, 31 Jul 2026 04:51:50 +0000 Subject: [PATCH] refactor(atomic-actions): clarify target contracts Move action-specific target classes beside their primitives, add a shared object target base, and track held objects by control part. Update tutorials, benchmarks, tests, documentation, and agent context for the new contracts. --- agent_context/MAP.yaml | 7 + .../topics/atomic-actions/atomic-actions.md | 72 ++++-- ...hain.lab.sim.atomic_actions.primitives.rst | 15 ++ .../embodichain.lab.sim.atomic_actions.rst | 31 ++- ...dichain.lab.sim.atomic_actions.targets.rst | 21 ++ .../sim/atomic_actions/builtin_actions.md | 45 ++-- .../overview/sim/atomic_actions/index.md | 91 ++++--- docs/source/tutorial/atomic_actions.rst | 36 ++- .../lab/sim/atomic_actions/__init__.py | 41 +++- embodichain/lab/sim/atomic_actions/core.py | 225 +++++++----------- embodichain/lab/sim/atomic_actions/engine.py | 25 +- .../sim/atomic_actions/primitives/__init__.py | 48 +++- .../primitives/coordinated_pickment.py | 100 ++++++-- .../primitives/coordinated_placement.py | 108 ++++++++- .../primitives/move_end_effector.py | 30 ++- .../primitives/move_held_object.py | 37 ++- .../atomic_actions/primitives/move_joints.py | 58 ++++- .../sim/atomic_actions/primitives/pick_up.py | 39 ++- .../sim/atomic_actions/primitives/place.py | 59 ++++- .../sim/atomic_actions/primitives/press.py | 27 ++- embodichain/lab/sim/atomic_actions/targets.py | 47 ++++ .../move_held_object_benchmark.py | 4 +- .../atomic_action/pickup_benchmark.py | 4 +- .../atomic_action/place_benchmark.py | 8 +- .../atomic_action/press_benchmark.py | 4 +- .../atomic_action/coordinated_pickment.py | 9 +- .../atomic_action/coordinated_placement.py | 13 +- .../atomic_action/move_end_effector.py | 3 +- .../atomic_action/move_held_object.py | 3 +- .../tutorials/atomic_action/move_joints.py | 3 +- scripts/tutorials/atomic_action/pickup.py | 3 +- scripts/tutorials/atomic_action/place.py | 7 +- scripts/tutorials/atomic_action/press.py | 6 +- .../tutorials/atomic_action/tutorial_utils.py | 24 ++ tests/sim/atomic_actions/test_actions.py | 192 +++++++++------ tests/sim/atomic_actions/test_core.py | 170 ++++++++++--- .../test_curobo_motion_source_e2e.py | 6 +- tests/sim/atomic_actions/test_engine.py | 51 ++-- .../sim/atomic_actions/test_engine_per_env.py | 6 +- .../atomic_actions/test_motion_source_e2e.py | 6 +- 40 files changed, 1210 insertions(+), 474 deletions(-) create mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.targets.rst create mode 100644 embodichain/lab/sim/atomic_actions/targets.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 81b995e2..41c2acfb 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -414,6 +414,7 @@ topics: - motion primitive - action primitive - AtomicAction + - ActionTarget - AtomicActionEngine - TrajectoryBuilder - 原子动作 @@ -425,6 +426,11 @@ topics: - TrajectoryBuilder - ActionResult - WorldState + - held_objects + - ObjectActionTarget + - PlaceTarget + - PressTarget + - target ownership - ActionCfg - motion_source - plan_arm_traj @@ -433,6 +439,7 @@ topics: - topics/atomic-actions/atomic-actions.md source_of_truth: - embodichain/lab/sim/atomic_actions/core.py + - embodichain/lab/sim/atomic_actions/targets.py - embodichain/lab/sim/atomic_actions/engine.py - embodichain/lab/sim/atomic_actions/trajectory.py - embodichain/lab/sim/atomic_actions/primitives/ diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index 420b4fde..30cec725 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -4,10 +4,11 @@ | What | Path | |---|---| -| Base classes, typed targets, configs | `embodichain/lab/sim/atomic_actions/core.py` | +| Base classes, configs, runtime state | `embodichain/lab/sim/atomic_actions/core.py` | +| Shared target contracts | `embodichain/lab/sim/atomic_actions/targets.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/` | +| Built-in primitives and their targets | `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` | @@ -21,7 +22,7 @@ AtomicActionEngine │ │ │ └── TrajectoryBuilder ← IK/interpolation and MotionGenerator dispatch │ - └── WorldState ← last_qpos + held_object/coordinated_held_object + └── WorldState ← last_qpos + per-control-part held-object maps ``` All tensor shapes carry a leading batch dim `B = n_envs`. @@ -30,26 +31,46 @@ All tensor shapes carry a leading batch dim `B = n_envs`. ### Typed Targets -Frozen dataclasses accepted by actions via their `TargetType` class variable. +Frozen, identity-equality dataclasses accepted by actions via their `TargetType` +class variable. Built-in and third-party targets inherit the open `ActionTarget` +marker; `BuiltinTarget` is only the closed union of targets shipped by EmbodiChain. +Each action-exclusive target is defined beside its owning action and config in +`primitives/.py`. The package root re-exports all targets, so callers +should import them from `embodichain.lab.sim.atomic_actions`. A genuinely shared +target belongs in a neutral target module, not in one primitive that another +primitive must import. + +`ObjectActionTarget(semantics)` is the neutral shared base for actions operating +on a semantic object. It intentionally does not define a generic pose field: +object poses, single-arm grasp poses, and dual-arm grasp pairs are distinct +contracts. | Target | Holds | Used by | |---|---|---| -| `EndEffectorPoseTarget(xpos)` | `(4,4)`, `(B,4,4)` or `(B,n_waypoint,4,4)` EEF pose | `MoveEndEffector`, `Place`, `Press` | +| `EndEffectorPoseTarget(xpos)` | `(4,4)`, `(B,4,4)` or `(B,n_waypoint,4,4)` EEF pose | `MoveEndEffector` | +| `PlaceTarget(xpos, tcp_symmetry)` | Release EEF pose plus optional TCP z-roll symmetry | `Place` | +| `PressTarget(xpos)` | One `(4,4)` or `(B,4,4)` contact pose | `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` | +| `ObjectActionTarget(semantics)` | Shared semantic-object contract; no generic pose | Base of object-centric targets | +| `GraspTarget(semantics)` | Object semantics plus optional single-arm `grasp_xpos` | `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` | +| `CoordinatedPickTarget(semantics, ...)` | Shared object + target object pose + left/right object-to-EEF transforms | `CoordinatedPickment` | +| `CoordinatedPlacementTarget(...)` | Placing/support target poses and per-call offsets | `CoordinatedPlacement` | + +`CoordinatedPickmentTarget` remains an alias of `CoordinatedPickTarget`. ### 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. +- `held_objects: dict[str, HeldObjectState]` — independently held objects keyed by arm/control part. +- `coordinated_held_objects: dict[tuple[str, str], CoordinatedHeldObjectState]` — jointly held objects keyed by an ordered control-part pair. `HeldObjectState` stores the object's semantics plus the object-to-EEF transform and grasp pose (both `(B, 4, 4)`). +Use `get_held_object(control_part)`, `get_coordinated_held_object(first, second)`, +and `with_updates(...)`; `with_updates` copies both maps so successor actions do +not alias their containers. ### ActionResult @@ -116,24 +137,28 @@ success, traj, final_state = engine.run(steps=[("move_end_effector", target)]) |---|---|---| | `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 | +| `PickUp` | `GraspTarget` | Approach → close gripper → lift; populates `held_objects[cfg.control_part]` | +| `MoveHeldObject` | `HeldObjectPoseTarget` | Moves the object at `held_objects[cfg.control_part]` | +| `Place` | `PlaceTarget` | Lower → open gripper → retract; clears its control-part entry | +| `Press` | `PressTarget` | Close gripper → press down → return | +| `CoordinatedPickment` | `CoordinatedPickTarget` | Replaces the two individual entries with one coordinated held state | +| `CoordinatedPlacement` | `CoordinatedPlacementTarget` | Reads both individual held states from `WorldState` | ## 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`. +2. Define an action-exclusive `@dataclass(frozen=True, slots=True, eq=False)` + target beside the action. Reuse or promote a target to a neutral module only + when the contract is genuinely shared. Inherit `ObjectActionTarget` when + multiple object-centric actions share only `semantics`; keep pose roles in + the concrete target. +3. Subclass `AtomicAction[YourTarget]` directly (do not inherit from another action). Set `TargetType` for runtime checking 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`. + - Advance `last_qpos` with `state.with_updates(...)` and preserve/update the held-object maps. 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`. @@ -143,4 +168,13 @@ success, traj, final_state = engine.run(steps=[("move_end_effector", target)]) - **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. +- **Putting runtime held state into a target** — desired state belongs in the + target; objects already held by a control part belong in `WorldState`. +- **Using a target dataclass with default Tensor equality** — use `eq=False`; + generated dataclass equality is invalid for multi-element tensors. +- **Importing a target from a sibling primitive** — give the action its own + contract or promote the shared contract to a neutral module. +- **Putting a generic `xpos` on a shared object target** — use explicit names + such as `object_target_pose`, `grasp_xpos`, or left/right grasp transforms; + their frames and cardinalities are not interchangeable. - **`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. 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 index b08c6cac..68326d4d 100644 --- 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 @@ -37,6 +37,21 @@ resulting trajectories along the time axis. CoordinatedPlacementCfg CoordinatedPlacement + .. rubric:: Built-in Target Contracts + + .. autosummary:: + + EndEffectorPoseTarget + JointPositionTarget + NamedJointPositionTarget + GraspTarget + HeldObjectPoseTarget + PlaceTarget + PressTarget + CoordinatedPickTarget + CoordinatedPickmentTarget + CoordinatedPlacementTarget + .. currentmodule:: embodichain.lab.sim.atomic_actions.primitives MoveEndEffector 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 fb8360b7..9bc4fe46 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 @@ -11,14 +11,20 @@ embodichain.lab.sim.atomic_actions AntipodalAffordance InteractionPoints ObjectSemantics + ActionTarget + ObjectActionTarget EndEffectorPoseTarget + PlaceTarget + PressTarget JointPositionTarget NamedJointPositionTarget GraspTarget HeldObjectPoseTarget + CoordinatedPickTarget CoordinatedPickmentTarget CoordinatedPlacementTarget Target + BuiltinTarget HeldObjectState CoordinatedHeldObjectState WorldState @@ -48,6 +54,7 @@ embodichain.lab.sim.atomic_actions :maxdepth: 1 :hidden: + embodichain.lab.sim.atomic_actions.targets embodichain.lab.sim.atomic_actions.primitives .. currentmodule:: embodichain.lab.sim.atomic_actions @@ -58,6 +65,8 @@ 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 +shared target contracts live in +``embodichain.lab.sim.atomic_actions.targets``. ``embodichain.lab.sim.atomic_actions.actions`` remains a compatibility re-export for existing imports. @@ -80,10 +89,26 @@ Core :members: :show-inheritance: +.. autoclass:: ActionTarget + :members: + :show-inheritance: + +.. autoclass:: ObjectActionTarget + :members: + :show-inheritance: + .. autoclass:: EndEffectorPoseTarget :members: :show-inheritance: +.. autoclass:: PlaceTarget + :members: + :show-inheritance: + +.. autoclass:: PressTarget + :members: + :show-inheritance: + .. autoclass:: JointPositionTarget :members: :show-inheritance: @@ -100,16 +125,20 @@ Core :members: :show-inheritance: -.. autoclass:: CoordinatedPickmentTarget +.. autoclass:: CoordinatedPickTarget :members: :show-inheritance: +.. autodata:: CoordinatedPickmentTarget + .. autoclass:: CoordinatedPlacementTarget :members: :show-inheritance: .. autodata:: Target +.. autodata:: BuiltinTarget + .. autoclass:: HeldObjectState :members: :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.targets.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.targets.rst new file mode 100644 index 00000000..2f60da71 --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.targets.rst @@ -0,0 +1,21 @@ +embodichain.lab.sim.atomic_actions.targets +========================================= + +.. automodule:: embodichain.lab.sim.atomic_actions.targets + +Overview +-------- + +Shared target contracts for object-centric atomic actions. These contracts +contain only fields whose meaning is consistent across multiple actions. +Action-specific pose roles remain on the concrete target dataclasses in +``embodichain.lab.sim.atomic_actions.primitives``. + +.. currentmodule:: embodichain.lab.sim.atomic_actions.targets + +Object Target +------------- + +.. autoclass:: ObjectActionTarget + :members: + :show-inheritance: diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index 4587267e..bd3c6393 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -17,9 +17,9 @@ The built-in atomic actions currently support gripper-based manipulation only. D | `MoveJoints` | Single | `JointPositionTarget` or `NamedJointPositionTarget` — qpos | Interpolate control-part joints | MoveJoints | | `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 | +| `Place` | Single | `PlaceTarget` — EEF release pose | Lower → open gripper → retract | Place | +| `Press` | Single | `PressTarget` — EEF contact pose | Close gripper → press down → return | Press | +| `CoordinatedPickment` | Dual | `CoordinatedPickTarget` — 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 | --- @@ -78,7 +78,9 @@ Three-phase grasp motion: *approach → close gripper → lift*. **Target:** `GraspTarget(semantics=...)` — an `ObjectSemantics` whose `affordance` is an `AntipodalAffordance`. The grasp pose is solved from the affordance and the entity's live pose at execute time. On success, the returned `WorldState` carries a populated -`held_object` (`HeldObjectState`). +`held_objects[control_part]` (`HeldObjectState`). +`GraspTarget` inherits the shared `ObjectActionTarget(semantics)` contract and +adds only its optional single-arm `grasp_xpos` override. ![PickUp demo](../../../_static/atomic_actions/pickup.gif) @@ -87,7 +89,8 @@ pose at execute time. On success, the returned `WorldState` carries a populated ## `MoveHeldObject` Moves a held object to an object-centric target pose while preserving the grasp. It requires -the `HeldObjectState` populated by a prior `PickUp` (read from `WorldState.held_object`) +the `HeldObjectState` populated by a prior `PickUp` (read from +`WorldState.held_objects[control_part]`) and preserves it in its successor state. `HeldObjectState` and `HeldObjectPoseTarget` are intentionally kept separate from @@ -114,7 +117,8 @@ Three-phase release motion: *lower → open gripper → retract*. Mirrors `PickU `PlaceCfg` carries its own gripper fields directly (it inherits `ActionCfg`, not a shared grasp-cfg base). The `approach_direction` field is not used — the arm moves straight -down to the target pose. On success, the returned `WorldState` clears `held_object` to `None`. +down to the target pose. On success, the returned `WorldState` removes the +entry for `PlaceCfg.control_part` from `held_objects`. | Config field | Default | Description | |---|---|---| @@ -125,7 +129,7 @@ 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=..., tcp_symmetry="none")` — the EEF pose at +**Target:** `PlaceTarget(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 @@ -143,7 +147,7 @@ Three-phase contact motion: *close gripper → press down → return*. This is u 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 +`Press` does not create or clear `WorldState.held_objects`; it preserves the state threaded into it. | Config field | Default | Description | @@ -153,7 +157,7 @@ threaded into it. | `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` +**Target:** `PressTarget(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) @@ -165,8 +169,8 @@ of shape `(4, 4)` or `(n_envs, 4, 4)`. 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`. +an entry in `coordinated_held_objects[(left_arm, right_arm)]` +(`CoordinatedHeldObjectState`) and removes individual held entries for those arms. | Config field | Default | Description | |---|---|---| @@ -178,8 +182,12 @@ while keeping both grippers closed. On success, the returned `WorldState` carrie | `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 +**Target:** `CoordinatedPickTarget(...)` with a target object pose, object semantics, and left/right object-to-EEF transforms. +It inherits the same `ObjectActionTarget(semantics)` base as `GraspTarget`, but +keeps the dual-arm pose fields in its own action-specific contract. + +`CoordinatedPickmentTarget` remains a compatibility alias. **Tutorial:** `scripts/tutorials/atomic_action/coordinated_pickment.py` @@ -193,10 +201,10 @@ Dual-arm object-centric placement. The support arm moves its held object to a lo 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. +`CoordinatedPlacement` reads both held objects from +`WorldState.held_objects`, keyed by `placing_arm_control_part` and +`support_arm_control_part`. The target contains desired poses and per-call +overrides only. | Config field | Default | Description | |---|---|---| @@ -218,8 +226,9 @@ held objects. | `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. +poses and optional height/release overrides. On success, the support arm's +entry remains in `WorldState.held_objects`; the placing arm's entry is removed +when `release=True`. **Tutorial:** `scripts/tutorials/atomic_action/coordinated_placement.py` diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index f1f120a7..b80ef1bb 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -21,7 +21,7 @@ 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/coordinated_held_object) + (last_qpos + per-control-part held-object maps) ``` Each action receives a typed target and a `WorldState`, runs its planning pipeline, and @@ -32,10 +32,12 @@ trajectories into one contiguous sequence: ``` GraspTarget(semantics, grasp_xpos=None) ──► AtomicAction.execute(target, state) EndEffectorPoseTarget(xpos) │ +PlaceTarget(xpos, tcp_symmetry) │ +PressTarget(xpos) │ JointPositionTarget(qpos) ├─ IK solve when pose-based NamedJointPositionTarget(name) ├─ Motion plan / interpolation HeldObjectPoseTarget(pose) └─ Gripper interpolation when needed -CoordinatedPickmentTarget(...) │ +CoordinatedPickTarget(...) │ CoordinatedPlacementTarget(...) │ │ ActionResult @@ -58,21 +60,38 @@ the held object's semantics and object-to-end-effector transform so later action object without recomputing the grasp. It is intentionally separate from `ObjectSemantics`, which remains a reusable object description rather than per-execution robot state. -**Typed targets** describe *where* an action should go. Each one is a small frozen dataclass, -and every action declares the target type, or tuple of target types, it accepts via its -`TargetType` class variable: +**Typed targets** describe *where* an action should go. Each one is a small frozen, +identity-equality dataclass inheriting the open `ActionTarget` marker. Every action +declares the target type, or tuple of target types, it accepts via its `TargetType` +class variable. Object-centric actions can share the semantic-object contract +provided by `ObjectActionTarget`, while keeping their pose roles action-specific: | Target | Constructor | Used by | |---|---|---| -| `EndEffectorPoseTarget` | `EndEffectorPoseTarget(xpos, tcp_symmetry="none")` | `MoveEndEffector`, `Place`, `Press` | +| `EndEffectorPoseTarget` | `EndEffectorPoseTarget(xpos)` | `MoveEndEffector` | +| `PlaceTarget` | `PlaceTarget(xpos, tcp_symmetry="none")` | `Place` | +| `PressTarget` | `PressTarget(xpos)` | `Press` | | `JointPositionTarget` | `JointPositionTarget(qpos)` | `MoveJoints` | | `NamedJointPositionTarget` | `NamedJointPositionTarget(name)` | `MoveJoints` | | `GraspTarget` | `GraspTarget(semantics, grasp_xpos=None)` | `PickUp` | | `HeldObjectPoseTarget` | `HeldObjectPoseTarget(object_target_pose)` | `MoveHeldObject` | -| `CoordinatedPickmentTarget` | `CoordinatedPickmentTarget(...)` | `CoordinatedPickment` | +| `CoordinatedPickTarget` | `CoordinatedPickTarget(...)` | `CoordinatedPickment` | | `CoordinatedPlacementTarget` | `CoordinatedPlacementTarget(...)` | `CoordinatedPlacement` | -`Target` is the union of these typed target dataclasses. +`Target` is a compatibility alias of the open `ActionTarget` marker. +`BuiltinTarget` is the closed union of target types shipped by EmbodiChain. +`CoordinatedPickmentTarget` remains an alias of `CoordinatedPickTarget`. +`GraspTarget` and `CoordinatedPickTarget` inherit +`ObjectActionTarget(semantics)`, but only `GraspTarget` defines the optional +single-arm `grasp_xpos`. + +Action-exclusive target classes live beside their owning primitive under +`atomic_actions/primitives/`. They are re-exported from +`embodichain.lab.sim.atomic_actions`, which remains the stable public import +surface. Shared target contracts such as `ObjectActionTarget` live in the +neutral `atomic_actions/targets.py` module instead of introducing dependencies +between primitive implementations. The shared base deliberately has no generic +`xpos`: object poses and single-/dual-arm EEF poses have different meanings. **`Affordance`** is a data class that encodes a specific interaction capability. The built-in affordance types are: @@ -103,25 +122,32 @@ action's `TargetType` before calling `execute`: | Target | Holds | Accepted by | |---|---|---| -| `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` | +| `EndEffectorPoseTarget(xpos)` | EEF pose tensor `(4,4)`, `(n_envs,4,4)` or `(n_envs, n_waypoint, 4, 4)` | `MoveEndEffector` | +| `PlaceTarget(xpos, tcp_symmetry="none")` | Release EEF pose; placement may opt into TCP z-roll 180 equivalence | `Place` | +| `PressTarget(xpos)` | One EEF contact pose `(4,4)` or `(n_envs,4,4)` | `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, 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` | +| `CoordinatedPickTarget(...)` | Shared object semantics plus left/right grasp transforms and target object pose | `CoordinatedPickment` | +| `CoordinatedPlacementTarget(...)` | Object-centric placing/support target poses and per-call overrides | `CoordinatedPlacement` | -`WorldState` is threaded between actions and carries the robot's `last_qpos` plus optional -`held_object: HeldObjectState` and `coordinated_held_object: CoordinatedHeldObjectState`. +Both object-grasp targets expose `semantics` through `ObjectActionTarget`. +`CoordinatedPickTarget` names its object and dual-arm pose roles explicitly, so +an unsupported single-arm `grasp_xpos` cannot be silently ignored. + +`WorldState` is threaded between actions and carries the robot's `last_qpos` plus +`held_objects: dict[str, HeldObjectState]` and +`coordinated_held_objects: dict[tuple[str, str], CoordinatedHeldObjectState]`. The built-in actions update it as follows: -| Action | Effect on `held_object` | +| Action | Effect on held-object maps | |---|---| -| `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 | +| `PickUp` | Populates `held_objects[cfg.control_part]` | +| `MoveHeldObject` | Requires and preserves `held_objects[cfg.control_part]` | +| `Place` | Removes its `cfg.control_part` entry | +| `CoordinatedPickment` | Removes both individual arm entries and populates their coordinated pair | +| `CoordinatedPlacement` | Reads both individual arm entries; removes only the placing entry when releasing | | `MoveEndEffector` | Leaves it unchanged | | `MoveJoints` | Leaves it unchanged | | `Press` | Leaves it unchanged | @@ -140,6 +166,7 @@ from embodichain.lab.sim.atomic_actions import ( AntipodalAffordance, GraspTarget, EndEffectorPoseTarget, + PlaceTarget, JointPositionTarget, NamedJointPositionTarget, HeldObjectPoseTarget, @@ -213,7 +240,7 @@ is_success, traj, final_state = engine.run( steps=[ ("pick_up", GraspTarget(semantics=semantics)), ("move_held_object", HeldObjectPoseTarget(object_target_pose=carry_pose)), - ("place", EndEffectorPoseTarget(xpos=place_pose)), + ("place", PlaceTarget(xpos=place_pose)), ("move_joints", NamedJointPositionTarget(name="home")), ] ) @@ -247,14 +274,23 @@ class PushCfg(ActionCfg): ### Step 2 — Implement the action ```python +from dataclasses import dataclass import torch from typing import ClassVar from embodichain.lab.sim.atomic_actions import ( - AtomicAction, ActionResult, EndEffectorPoseTarget, Target, WorldState, TrajectoryBuilder, + ActionTarget, AtomicAction, ActionResult, WorldState, TrajectoryBuilder, ) -class Push(AtomicAction): - TargetType: ClassVar[type] = EndEffectorPoseTarget +@dataclass(frozen=True, slots=True, eq=False) +class PushTarget(ActionTarget): + xpos: torch.Tensor + + def __post_init__(self) -> None: + if self.xpos.shape[-2:] != (4, 4) or self.xpos.dim() not in (2, 3): + raise ValueError("xpos must have shape (4, 4) or (n_envs, 4, 4)") + +class Push(AtomicAction[PushTarget]): + TargetType: ClassVar[type] = PushTarget def __init__(self, motion_generator, cfg: PushCfg | None = None): super().__init__(motion_generator, cfg or PushCfg()) @@ -263,15 +299,14 @@ class Push(AtomicAction): self.robot_dof = self.robot.dof self.n_envs = self.robot.get_qpos().shape[0] - def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionResult: + def execute(self, target: PushTarget, state: WorldState) -> ActionResult: # ... your planning logic, using self.builder for IK / interpolation ... # full must be shaped (n_envs, n_waypoints, robot.dof) return ActionResult( success=is_success, trajectory=full, - next_state=WorldState( + next_state=state.with_updates( last_qpos=full[:, -1, :].clone(), - held_object=state.held_object, # push does not change what is held ), ) ``` @@ -281,11 +316,9 @@ class Push(AtomicAction): Register an instance with the engine so it can be referenced by name in `run()`: ```python -from embodichain.lab.sim.atomic_actions import EndEffectorPoseTarget - engine.register(Push(motion_gen, cfg=PushCfg(push_distance=0.08))) is_success, traj, final_state = engine.run( - steps=[("push", EndEffectorPoseTarget(xpos=target_pose))] + steps=[("push", PushTarget(xpos=target_pose))] ) ``` diff --git a/docs/source/tutorial/atomic_actions.rst b/docs/source/tutorial/atomic_actions.rst index 5644792b..f5d6f53c 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -15,8 +15,12 @@ Key Features - **Typed targets** — every action accepts a small target dataclass such as ``EndEffectorPoseTarget``, ``JointPositionTarget``, ``NamedJointPositionTarget``, - ``GraspTarget`` (wrapping an ``ObjectSemantics``), or ``HeldObjectPoseTarget``. The + ``PlaceTarget``, ``PressTarget``, ``GraspTarget`` (wrapping an + ``ObjectSemantics``), or + ``HeldObjectPoseTarget``. The engine checks each step's target against the action's declared ``TargetType`` before running. + Object-centric targets may inherit ``ObjectActionTarget`` to share the + ``semantics`` contract without sharing action-specific pose fields. - **Built-in primitives** — ``MoveEndEffector``, ``MoveJoints``, ``PickUp``, ``MoveHeldObject``, ``Place``, ``Press``, ``CoordinatedPickment``, and ``CoordinatedPlacement`` cover the most common tabletop manipulation workflows out of the box. @@ -25,7 +29,7 @@ Key Features ``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`` / ``coordinated_held_object``) + ``held_objects`` / ``coordinated_held_objects``) from one action into the next, and returns a single concatenated full-DOF trajectory ready to replay in the simulator. @@ -159,7 +163,11 @@ Executing a pick-place-end-effector sequence .. code-block:: python - from embodichain.lab.sim.atomic_actions import GraspTarget, EndEffectorPoseTarget + from embodichain.lab.sim.atomic_actions import ( + GraspTarget, + EndEffectorPoseTarget, + PlaceTarget, + ) place_xpos = ... # torch.Tensor [4, 4] — target placement pose rest_xpos = ... # torch.Tensor [4, 4] — resting pose after placing @@ -167,7 +175,7 @@ Executing a pick-place-end-effector sequence is_success, trajectory, _ = engine.run( steps=[ ("pick_up", GraspTarget(semantics=semantics)), - ("place", EndEffectorPoseTarget(xpos=place_xpos)), + ("place", PlaceTarget(xpos=place_xpos)), ("move_end_effector", EndEffectorPoseTarget(xpos=rest_xpos)), ] ) @@ -204,7 +212,7 @@ Moving a held object ~~~~~~~~~~~~~~~~~~~~ ``MoveHeldObject`` consumes the runtime ``HeldObjectState`` produced by a previous -``PickUp`` (read from the threaded ``WorldState.held_object``). The target is +``PickUp`` (read from ``WorldState.held_objects[control_part]``). The target is object-centric: the caller specifies where the held object should move, and the action converts that pose into an end-effector target via the stored object-to-EEF transform while keeping the gripper closed. @@ -238,19 +246,25 @@ Registering custom actions .. code-block:: python + from dataclasses import dataclass from typing import ClassVar + import torch from embodichain.lab.sim.atomic_actions import ( - AtomicAction, ActionResult, ActionCfg, EndEffectorPoseTarget, WorldState, TrajectoryBuilder, + ActionTarget, AtomicAction, ActionResult, WorldState, TrajectoryBuilder, ) - class Push(AtomicAction): - TargetType: ClassVar[type] = EndEffectorPoseTarget + @dataclass(frozen=True, slots=True, eq=False) + class PushTarget(ActionTarget): + xpos: torch.Tensor + + class Push(AtomicAction[PushTarget]): + TargetType: ClassVar[type] = PushTarget def __init__(self, motion_generator, cfg: PushCfg | None = None): super().__init__(motion_generator, cfg or PushCfg()) self.builder = TrajectoryBuilder(motion_generator) - def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionResult: + def execute(self, target: PushTarget, state: WorldState) -> ActionResult: # ... your planning logic, using self.builder ... return ActionResult(success=is_success, trajectory=full, next_state=...) @@ -272,4 +286,8 @@ Notes & Best Practices - For static (non-physics) playback, iterate over ``trajectory[:, i]`` and call ``robot.set_qpos`` directly; for physics-enabled playback, feed waypoints through your controller or gym wrapper instead. +- Define an action-specific target beside its action implementation. If multiple actions + genuinely share part of a target contract, extract only that minimal base into + ``atomic_actions/targets.py`` instead of importing from another primitive. Do not use a + generic ``xpos`` for object, single-arm EEF, and dual-arm EEF poses; name each pose role. - To add a new action type, see :doc:`/overview/sim/atomic_actions/index`. diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index 381bfd8c..baa2f6ad 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -31,19 +31,13 @@ InteractionPoints, ) from .core import ( + ActionTarget, ActionCfg, ActionResult, AtomicAction, CoordinatedHeldObjectState, - CoordinatedPickmentTarget, - CoordinatedPlacementTarget, - GraspTarget, HeldObjectState, - HeldObjectPoseTarget, - JointPositionTarget, - NamedJointPositionTarget, ObjectSemantics, - EndEffectorPoseTarget, Target, WorldState, ) @@ -53,42 +47,75 @@ unregister_action, get_registered_actions, ) +from .targets import ObjectActionTarget from .primitives import ( + CoordinatedPickTarget, CoordinatedPickment, CoordinatedPickmentCfg, + CoordinatedPickmentTarget, CoordinatedPlacement, CoordinatedPlacementCfg, + CoordinatedPlacementTarget, + EndEffectorPoseTarget, + GraspTarget, + HeldObjectPoseTarget, + JointPositionTarget, MoveEndEffector, MoveEndEffectorCfg, MoveHeldObject, MoveHeldObjectCfg, MoveJoints, MoveJointsCfg, + NamedJointPositionTarget, PickUp, PickUpCfg, Place, PlaceCfg, + PlaceTarget, Press, PressCfg, + PressTarget, ) from .trajectory import TrajectoryBuilder +BuiltinTarget = ( + EndEffectorPoseTarget + | JointPositionTarget + | NamedJointPositionTarget + | GraspTarget + | HeldObjectPoseTarget + | PlaceTarget + | PressTarget + | CoordinatedPickTarget + | CoordinatedPlacementTarget +) +"""Union of target types shipped by EmbodiChain. + +Use :class:`ActionTarget` rather than this closed union at extension boundaries. +""" + __all__ = [ # Core classes + "ActionTarget", "Affordance", "AntipodalAffordance", "InteractionPoints", "ObjectSemantics", + "ObjectActionTarget", "HeldObjectState", "CoordinatedHeldObjectState", "HeldObjectPoseTarget", "JointPositionTarget", "NamedJointPositionTarget", "EndEffectorPoseTarget", + "PlaceTarget", + "PressTarget", + "CoordinatedPickTarget", "CoordinatedPickmentTarget", "CoordinatedPlacementTarget", "GraspTarget", "Target", + "BuiltinTarget", "WorldState", "ActionResult", "ActionCfg", diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index 2bad1e0d..fccb39be 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, Literal, TYPE_CHECKING +from typing import Any, ClassVar, Generic, Mapping, TYPE_CHECKING, TypeVar from embodichain.lab.sim.common import BatchEntity from embodichain.utils import configclass @@ -62,143 +62,47 @@ def __post_init__(self) -> None: # ============================================================================= -# Typed targets +# Target foundation # ============================================================================= -@dataclass(frozen=True) -class EndEffectorPoseTarget: - """End-effector pose target. Used by MoveEndEffector, Place, and Press.""" +class ActionTarget: + """Open marker base for atomic-action target value objects. - xpos: torch.Tensor - """Target end-effector homogeneous transform. - - Accepts: - - - ``(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. (Consumed as multi-waypoint by MoveEndEffector and - Place.) + Third-party actions should define a target dataclass that inherits from this + class. The engine performs the action-specific runtime check using + :attr:`AtomicAction.TargetType`; the marker keeps the public engine contract + open to targets outside the built-in set. """ - 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: - """Joint-space target for a configured robot control part.""" - - qpos: torch.Tensor - """Target joint positions. - - Accepts: - - - ``(control_dof,)`` or ``(n_envs, control_dof)`` — a single waypoint. - - ``(n_envs, n_waypoint, control_dof)`` — a multi-waypoint trajectory; - waypoints are visited in order. - """ - - -@dataclass(frozen=True) -class NamedJointPositionTarget: - """Named joint-space target resolved from ``MoveJointsCfg``.""" - - name: str - """Name of a joint-position target in ``MoveJointsCfg.named_joint_positions``.""" - - -@dataclass(frozen=True) -class GraspTarget: - """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: - """Move the currently-held object to a desired object pose.""" - - object_target_pose: torch.Tensor - """(4, 4) or (n_envs, 4, 4) target pose for the held object.""" + __slots__ = () -@dataclass(frozen=True) -class CoordinatedPickmentTarget: - """Object-centric target for picking and moving one object with two hands.""" +TargetT = TypeVar("TargetT", bound=ActionTarget) - 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.""" +def _validate_pose_tensor( + value: torch.Tensor, + name: str, + *, + allow_waypoints: bool, +) -> None: + """Validate the environment-independent part of a pose tensor contract.""" + if not isinstance(value, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(value).__name__}.") + valid_dims = {2, 3, 4} if allow_waypoints else {2, 3} + if value.dim() not in valid_dims or value.shape[-2:] != (4, 4): + supported = "(4, 4), (n_envs, 4, 4)" + if allow_waypoints: + supported += ", or (n_envs, n_waypoint, 4, 4)" + raise ValueError( + f"{name} must have shape {supported}, got {tuple(value.shape)}." + ) -Target = ( - EndEffectorPoseTarget - | JointPositionTarget - | NamedJointPositionTarget - | GraspTarget - | HeldObjectPoseTarget - | CoordinatedPickmentTarget - | CoordinatedPlacementTarget -) +# ``Target`` used to be a closed union of built-in target classes. Keep the +# public name as an open compatibility alias so extension targets are accepted. +Target = ActionTarget # ============================================================================= @@ -206,7 +110,7 @@ class CoordinatedPlacementTarget: # ============================================================================= -@dataclass +@dataclass(slots=True, eq=False) class HeldObjectState: """State of an object currently held by the robot.""" @@ -220,7 +124,7 @@ class HeldObjectState: """Batched end-effector pose used to grasp the object, shape [n_envs, 4, 4].""" -@dataclass +@dataclass(slots=True, eq=False) class CoordinatedHeldObjectState: """State of a single object jointly held by two robot hands.""" @@ -240,21 +144,59 @@ class CoordinatedHeldObjectState: """Right end-effector grasp pose for the shared object, shape ``[n_envs, 4, 4]``.""" -@dataclass +@dataclass(slots=True, eq=False) class WorldState: """State the engine threads through a sequence of actions.""" last_qpos: torch.Tensor """Robot joint positions at the start of the next action, shape [n_envs, robot.dof].""" - held_object: HeldObjectState | None = None - """Object currently held by the gripper, or None.""" + held_objects: dict[str, HeldObjectState] = field(default_factory=dict) + """Objects held by individual control parts, keyed by control-part name.""" - coordinated_held_object: CoordinatedHeldObjectState | None = None - """Object currently held by two grippers, or None.""" + coordinated_held_objects: dict[tuple[str, str], CoordinatedHeldObjectState] = field( + default_factory=dict + ) + """Objects jointly held by two control parts, keyed by their ordered pair.""" + def get_held_object(self, control_part: str) -> HeldObjectState | None: + """Return the object held by ``control_part``, if any.""" + return self.held_objects.get(control_part) -@dataclass + def get_coordinated_held_object( + self, + first_control_part: str, + second_control_part: str, + ) -> CoordinatedHeldObjectState | None: + """Return the object jointly held by an ordered control-part pair.""" + return self.coordinated_held_objects.get( + (first_control_part, second_control_part) + ) + + def with_updates( + self, + *, + last_qpos: torch.Tensor | None = None, + held_objects: Mapping[str, HeldObjectState] | None = None, + coordinated_held_objects: ( + Mapping[tuple[str, str], CoordinatedHeldObjectState] | None + ) = None, + ) -> WorldState: + """Return a successor state without aliasing held-state dictionaries.""" + return WorldState( + last_qpos=self.last_qpos if last_qpos is None else last_qpos, + held_objects=dict( + self.held_objects if held_objects is None else held_objects + ), + coordinated_held_objects=dict( + self.coordinated_held_objects + if coordinated_held_objects is None + else coordinated_held_objects + ), + ) + + +@dataclass(slots=True, eq=False) class ActionResult: """Return value of every AtomicAction.execute call.""" @@ -318,7 +260,7 @@ def __post_init__(self) -> None: # ============================================================================= -class AtomicAction(ABC): +class AtomicAction(Generic[TargetT], ABC): """Abstract base for atomic actions. Subclasses declare ``TargetType`` to advertise the concrete target dataclass @@ -326,7 +268,7 @@ class AtomicAction(ABC): dropped from the contract in this redesign. """ - TargetType: ClassVar[type | tuple[type, ...]] + TargetType: ClassVar[type[ActionTarget] | tuple[type[ActionTarget], ...]] """Concrete target dataclass or dataclasses accepted by ``execute``.""" def __init__( @@ -341,7 +283,7 @@ def __init__( self.control_part = self.cfg.control_part @abstractmethod - def execute(self, target: Target, state: WorldState) -> ActionResult: + def execute(self, target: TargetT, state: WorldState) -> ActionResult: """Plan and return a full-DoF trajectory for this action. Args: @@ -354,19 +296,14 @@ def execute(self, target: Target, state: WorldState) -> ActionResult: __all__ = [ + "ActionTarget", "ActionCfg", "ActionResult", "AtomicAction", "CoordinatedHeldObjectState", - "CoordinatedPickmentTarget", - "CoordinatedPlacementTarget", - "GraspTarget", "HeldObjectState", - "HeldObjectPoseTarget", - "JointPositionTarget", - "NamedJointPositionTarget", "ObjectSemantics", - "EndEffectorPoseTarget", "Target", + "TargetT", "WorldState", ] diff --git a/embodichain/lab/sim/atomic_actions/engine.py b/embodichain/lab/sim/atomic_actions/engine.py index 59574af4..b1b0b391 100644 --- a/embodichain/lab/sim/atomic_actions/engine.py +++ b/embodichain/lab/sim/atomic_actions/engine.py @@ -22,9 +22,9 @@ from embodichain.utils import logger from .core import ( + ActionTarget, ActionResult, AtomicAction, - Target, WorldState, ) @@ -83,12 +83,26 @@ def actions(self) -> dict[str, AtomicAction]: def register(self, action: AtomicAction, *, name: str | None = None) -> None: """Register an action instance under ``name`` or its ``cfg.name``.""" + declared_target_type = getattr(action, "TargetType", None) + target_types = ( + declared_target_type + if isinstance(declared_target_type, tuple) + else (declared_target_type,) + ) + if not target_types or not all( + isinstance(target_type, type) and issubclass(target_type, ActionTarget) + for target_type in target_types + ): + logger.log_error( + "AtomicAction.TargetType must contain ActionTarget subclasses.", + TypeError, + ) key = name if name is not None else action.cfg.name self._actions[key] = action def run( self, - steps: Iterable[tuple[str, Target]], + steps: Iterable[tuple[str, ActionTarget]], state: WorldState | None = None, ) -> tuple[torch.Tensor, torch.Tensor, WorldState]: """Run a sequence of named actions, threading WorldState through. @@ -146,9 +160,10 @@ def run( 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 + state = result.next_state.with_updates( + last_qpos=torch.where( + alive[:, None], result.next_state.last_qpos, prev_last_qpos + ) ) return alive, full_traj, state diff --git a/embodichain/lab/sim/atomic_actions/primitives/__init__.py b/embodichain/lab/sim/atomic_actions/primitives/__init__.py index 0998f4d1..a857e2c2 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/__init__.py +++ b/embodichain/lab/sim/atomic_actions/primitives/__init__.py @@ -18,30 +18,62 @@ 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 +from .coordinated_pickment import ( + CoordinatedPickTarget, + CoordinatedPickment, + CoordinatedPickmentCfg, + CoordinatedPickmentTarget, +) +from .coordinated_placement import ( + CoordinatedPlacement, + CoordinatedPlacementCfg, + CoordinatedPlacementTarget, +) +from .move_end_effector import ( + EndEffectorPoseTarget, + MoveEndEffector, + MoveEndEffectorCfg, +) +from .move_held_object import ( + HeldObjectPoseTarget, + MoveHeldObject, + MoveHeldObjectCfg, +) +from .move_joints import ( + JointPositionTarget, + MoveJoints, + MoveJointsCfg, + NamedJointPositionTarget, +) +from .pick_up import GraspTarget, PickUp, PickUpCfg +from .place import Place, PlaceCfg, PlaceTarget +from .press import Press, PressCfg, PressTarget __all__ = [ + "CoordinatedPickTarget", "CoordinatedPickment", "CoordinatedPickmentCfg", + "CoordinatedPickmentTarget", "CoordinatedPlacement", "CoordinatedPlacementCfg", + "CoordinatedPlacementTarget", + "EndEffectorPoseTarget", + "GraspTarget", + "HeldObjectPoseTarget", + "JointPositionTarget", "MoveEndEffector", "MoveEndEffectorCfg", "MoveHeldObject", "MoveHeldObjectCfg", "MoveJoints", "MoveJointsCfg", + "NamedJointPositionTarget", "PickUp", "PickUpCfg", "Place", "PlaceCfg", + "PlaceTarget", "Press", "PressCfg", + "PressTarget", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py index 2b0fb757..5cce3e69 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py @@ -18,6 +18,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import ClassVar import torch @@ -30,12 +31,58 @@ ActionResult, AtomicAction, CoordinatedHeldObjectState, - CoordinatedPickmentTarget, WorldState, + _validate_pose_tensor, ) +from ..targets import ObjectActionTarget from ..trajectory import TrajectoryBuilder +@dataclass(frozen=True, slots=True, eq=False) +class CoordinatedPickTarget(ObjectActionTarget): + """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)``.""" + + 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 ``semantics.entity`` pose.""" + + def __post_init__(self) -> None: + ObjectActionTarget.__post_init__(self) + _validate_pose_tensor( + self.object_target_pose, + "object_target_pose", + allow_waypoints=False, + ) + _validate_pose_tensor( + self.left_object_to_eef, + "left_object_to_eef", + allow_waypoints=False, + ) + _validate_pose_tensor( + self.right_object_to_eef, + "right_object_to_eef", + allow_waypoints=False, + ) + if self.object_initial_pose is not None: + _validate_pose_tensor( + self.object_initial_pose, + "object_initial_pose", + allow_waypoints=False, + ) + + +# Backward-compatible spelling retained for the existing action class and users. +CoordinatedPickmentTarget = CoordinatedPickTarget + + @configclass class CoordinatedPickmentCfg(ActionCfg): name: str = "coordinated_pickment" @@ -384,10 +431,10 @@ def _interpolate_object_pose( return poses -class CoordinatedPickment(AtomicAction): +class CoordinatedPickment(AtomicAction[CoordinatedPickTarget]): """Pick and move a single object pinched by two hands.""" - TargetType: ClassVar[type] = CoordinatedPickmentTarget + TargetType: ClassVar[type] = CoordinatedPickTarget _assemble_phase = _DualArmHelpers._assemble_phase _compose_dual_arm_trajectory = _DualArmHelpers._compose_dual_arm_trajectory @@ -465,24 +512,24 @@ def _validate_hand_qpos_cfg(self) -> None: ) def _resolve_object_initial_pose( - self, target: CoordinatedPickmentTarget + self, target: CoordinatedPickTarget ) -> 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: + if target.semantics.entity is None: logger.log_error( - "CoordinatedPickmentTarget requires object_initial_pose when " - "object_semantics.entity is not provided.", + "CoordinatedPickTarget requires object_initial_pose when " + "semantics.entity is not provided.", ValueError, ) return self._resolve_pose( - target.object_semantics.entity.get_local_pose(to_matrix=True), + target.semantics.entity.get_local_pose(to_matrix=True), "object_initial_pose", ) def _resolve_target( self, - target: CoordinatedPickmentTarget, + target: CoordinatedPickTarget, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -508,7 +555,7 @@ def _resolve_target( 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, + semantics=target.semantics, left_object_to_eef=left_object_to_eef, right_object_to_eef=right_object_to_eef, left_grasp_xpos=left_grasp_xpos, @@ -711,9 +758,7 @@ def _plan_synchronized_object_motion( self._interpolate_qpos_keyframes(right_traj, keyframe_indices, n_waypoints), ) - def execute( - self, target: CoordinatedPickmentTarget, state: WorldState - ) -> ActionResult: + def execute(self, target: CoordinatedPickTarget, state: WorldState) -> ActionResult: ( object_initial_pose, object_target_pose, @@ -867,15 +912,36 @@ def execute( left_grasp_xpos=left_target_xpos, right_grasp_xpos=right_target_xpos, ) + involved_control_parts = { + self.cfg.left_arm_control_part, + self.cfg.right_arm_control_part, + } + held_objects = { + key: value + for key, value in state.held_objects.items() + if key not in involved_control_parts + } + coordinated_held_objects = dict(state.coordinated_held_objects) + coordinated_held_objects[ + ( + self.cfg.left_arm_control_part, + self.cfg.right_arm_control_part, + ) + ] = coordinated_held_object return ActionResult( success=success_mask, trajectory=full, - next_state=WorldState( + next_state=state.with_updates( last_qpos=full[:, -1, :].clone(), - held_object=None, - coordinated_held_object=coordinated_held_object, + held_objects=held_objects, + coordinated_held_objects=coordinated_held_objects, ), ) -__all__ = ["CoordinatedPickment", "CoordinatedPickmentCfg"] +__all__ = [ + "CoordinatedPickTarget", + "CoordinatedPickment", + "CoordinatedPickmentCfg", + "CoordinatedPickmentTarget", +] diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py index ba503220..7f77328a 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py @@ -18,6 +18,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import ClassVar import torch @@ -27,16 +28,49 @@ from ._helpers import resolve_object_target from ..core import ( + ActionTarget, ActionCfg, ActionResult, AtomicAction, - CoordinatedPlacementTarget, HeldObjectState, WorldState, + _validate_pose_tensor, ) from ..trajectory import TrajectoryBuilder +@dataclass(frozen=True, slots=True, eq=False) +class CoordinatedPlacementTarget(ActionTarget): + """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_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.""" + + def __post_init__(self) -> None: + _validate_pose_tensor( + self.placing_object_target_pose, + "placing_object_target_pose", + allow_waypoints=False, + ) + _validate_pose_tensor( + self.support_object_target_pose, + "support_object_target_pose", + allow_waypoints=False, + ) + + @configclass class CoordinatedPlacementCfg(ActionCfg): name: str = "coordinated_placement" @@ -91,7 +125,7 @@ class CoordinatedPlacementCfg(ActionCfg): """Number of waypoints used for the placing-arm lift retreat.""" -class CoordinatedPlacement(AtomicAction): +class CoordinatedPlacement(AtomicAction[CoordinatedPlacementTarget]): """Coordinate two held objects: support object below, placing object above.""" TargetType: ClassVar[type] = CoordinatedPlacementTarget @@ -159,9 +193,13 @@ def __init__( def execute( self, target: CoordinatedPlacementTarget, state: WorldState ) -> ActionResult: - placing_xpos, support_xpos, release, support_held_object = self._resolve_target( - target - ) + ( + placing_xpos, + support_xpos, + release, + placing_held_object, + support_held_object, + ) = self._resolve_target(target, state) placing_start_qpos, support_start_qpos = self._resolve_start_qpos(state) segments = self._compute_segment_lengths(release) @@ -258,12 +296,28 @@ def execute( ], dim=1, ) + held_objects = dict(state.held_objects) + if release: + held_objects.pop(self.cfg.placing_arm_control_part, None) + else: + held_objects[self.cfg.placing_arm_control_part] = placing_held_object + held_objects[self.cfg.support_arm_control_part] = support_held_object + involved_control_parts = { + self.cfg.placing_arm_control_part, + self.cfg.support_arm_control_part, + } + coordinated_held_objects = { + key: value + for key, value in state.coordinated_held_objects.items() + if involved_control_parts.isdisjoint(key) + } return ActionResult( success=True, trajectory=full, - next_state=WorldState( + next_state=state.with_updates( last_qpos=full[:, -1, :].clone(), - held_object=support_held_object, + held_objects=held_objects, + coordinated_held_objects=coordinated_held_objects, ), ) @@ -341,7 +395,28 @@ def _resolve_held_state( def _resolve_target( self, target: CoordinatedPlacementTarget, - ) -> tuple[torch.Tensor, torch.Tensor, bool, HeldObjectState]: + state: WorldState, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + bool, + HeldObjectState, + HeldObjectState, + ]: + placing_held_object = state.get_held_object(self.cfg.placing_arm_control_part) + if placing_held_object is None: + logger.log_error( + "CoordinatedPlacement requires an object held by placing control " + f"part {self.cfg.placing_arm_control_part!r}.", + ValueError, + ) + support_held_object = state.get_held_object(self.cfg.support_arm_control_part) + if support_held_object is None: + logger.log_error( + "CoordinatedPlacement requires an object held by support control " + f"part {self.cfg.support_arm_control_part!r}.", + ValueError, + ) placing_height_offset = ( self.cfg.placing_height_offset if target.placing_height_offset is None @@ -363,11 +438,11 @@ def _resolve_target( "support_object_target_pose", ) placing_object_to_eef = self._resolve_object_to_eef( - target.placing_held_object, + placing_held_object, "placing_held_object", ) support_object_to_eef = self._resolve_object_to_eef( - target.support_held_object, + support_held_object, "support_held_object", ) placing_xpos = torch.bmm(placing_object_pose, placing_object_to_eef) @@ -378,7 +453,12 @@ def _resolve_target( support_xpos, release, self._resolve_held_state( - target.support_held_object, + placing_held_object, + "placing_held_object", + placing_object_to_eef, + ), + self._resolve_held_state( + support_held_object, "support_held_object", support_object_to_eef, ), @@ -480,4 +560,8 @@ def _fail(self, state: WorldState) -> ActionResult: ) -__all__ = ["CoordinatedPlacement", "CoordinatedPlacementCfg"] +__all__ = [ + "CoordinatedPlacement", + "CoordinatedPlacementCfg", + "CoordinatedPlacementTarget", +] diff --git a/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py b/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py index 1fe9694f..49e7e7f4 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py @@ -18,6 +18,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import ClassVar import torch @@ -27,15 +28,34 @@ from ._helpers import arm_qpos_from_state from ..core import ( + ActionTarget, ActionCfg, ActionResult, AtomicAction, - EndEffectorPoseTarget, WorldState, + _validate_pose_tensor, ) from ..trajectory import TrajectoryBuilder +@dataclass(frozen=True, slots=True, eq=False) +class EndEffectorPoseTarget(ActionTarget): + """End-effector pose target used by :class:`MoveEndEffector`.""" + + xpos: torch.Tensor + """Target end-effector homogeneous transform. + + Accepts: + + - ``(4, 4)`` or ``(n_envs, 4, 4)`` — a single waypoint. + - ``(n_envs, n_waypoint, 4, 4)`` — a multi-waypoint trajectory whose + waypoints are visited in order. + """ + + def __post_init__(self) -> None: + _validate_pose_tensor(self.xpos, "xpos", allow_waypoints=True) + + @configclass class MoveEndEffectorCfg(ActionCfg): name: str = "move_end_effector" @@ -48,7 +68,7 @@ class MoveEndEffectorCfg(ActionCfg): """Optional planner-specific options copied for each motion-generator call.""" -class MoveEndEffector(AtomicAction): +class MoveEndEffector(AtomicAction[EndEffectorPoseTarget]): """Plan a free-space end-effector move to a target pose. The :class:`EndEffectorPoseTarget` may carry either a single waypoint @@ -94,10 +114,8 @@ def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionRes return ActionResult( success=success, trajectory=full, - next_state=WorldState( + next_state=state.with_updates( last_qpos=full[:, -1, :].clone(), - held_object=state.held_object, - coordinated_held_object=state.coordinated_held_object, ), ) @@ -139,4 +157,4 @@ def _fail(self, state: WorldState) -> ActionResult: ) -__all__ = ["MoveEndEffector", "MoveEndEffectorCfg"] +__all__ = ["EndEffectorPoseTarget", "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 index ba261703..c2668335 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py @@ -18,6 +18,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import ClassVar import torch @@ -28,15 +29,31 @@ from ._helpers import arm_qpos_from_state, resolve_object_target from ..core import ( + ActionTarget, ActionCfg, ActionResult, AtomicAction, - HeldObjectPoseTarget, WorldState, + _validate_pose_tensor, ) from ..trajectory import TrajectoryBuilder +@dataclass(frozen=True, slots=True, eq=False) +class HeldObjectPoseTarget(ActionTarget): + """Desired pose for the object held by this action's control part.""" + + object_target_pose: torch.Tensor + """Target object pose, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" + + def __post_init__(self) -> None: + _validate_pose_tensor( + self.object_target_pose, + "object_target_pose", + allow_waypoints=False, + ) + + @configclass class MoveHeldObjectCfg(ActionCfg): name: str = "move_held_object" @@ -58,7 +75,7 @@ class MoveHeldObjectCfg(ActionCfg): """Optional rotation in radians used by the legacy upright transport mode.""" -class MoveHeldObject(AtomicAction): +class MoveHeldObject(AtomicAction[HeldObjectPoseTarget]): """Move the held object to a target object pose; keep the gripper closed.""" TargetType: ClassVar[type] = HeldObjectPoseTarget @@ -83,9 +100,11 @@ def __init__( 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: + held_object = state.get_held_object(self.cfg.control_part) + if held_object is None: logger.log_error( - "MoveHeldObject requires WorldState.held_object - run PickUp first.", + "MoveHeldObject requires an object held by control part " + f"{self.cfg.control_part!r} - run PickUp first.", ValueError, ) object_target_pose = resolve_object_target( @@ -104,9 +123,9 @@ def execute(self, target: HeldObjectPoseTarget, state: WorldState) -> ActionResu self._apply_configured_upright_rotation( object_target_pose, end_arm_xpos, - state.held_object.semantics.entity.get_local_pose(to_matrix=True), + held_object.semantics.entity.get_local_pose(to_matrix=True), ) - object_to_eef = state.held_object.object_to_eef.to( + object_to_eef = held_object.object_to_eef.to( device=self.device, dtype=torch.float32 ) if object_to_eef.shape == (4, 4): @@ -141,10 +160,8 @@ def execute(self, target: HeldObjectPoseTarget, state: WorldState) -> ActionResu return ActionResult( success=success, trajectory=full, - next_state=WorldState( + next_state=state.with_updates( last_qpos=full[:, -1, :].clone(), - held_object=state.held_object, - coordinated_held_object=state.coordinated_held_object, ), ) @@ -239,4 +256,4 @@ def _fail(self, state: WorldState) -> ActionResult: ) -__all__ = ["MoveHeldObject", "MoveHeldObjectCfg"] +__all__ = ["HeldObjectPoseTarget", "MoveHeldObject", "MoveHeldObjectCfg"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py index 845d6bb5..e619c113 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py @@ -18,6 +18,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import ClassVar import torch @@ -25,16 +26,56 @@ from embodichain.utils import configclass, logger from ..core import ( + ActionTarget, ActionCfg, ActionResult, AtomicAction, - JointPositionTarget, - NamedJointPositionTarget, WorldState, ) from ..trajectory import TrajectoryBuilder +@dataclass(frozen=True, slots=True, eq=False) +class JointPositionTarget(ActionTarget): + """Joint-space target for a configured robot control part.""" + + qpos: torch.Tensor + """Target joint positions. + + Accepts: + + - ``(control_dof,)`` or ``(n_envs, control_dof)`` — a single waypoint. + - ``(n_envs, n_waypoint, control_dof)`` — a multi-waypoint trajectory; + waypoints are visited in order. + """ + + def __post_init__(self) -> None: + if not isinstance(self.qpos, torch.Tensor): + raise TypeError( + f"qpos must be a torch.Tensor, got {type(self.qpos).__name__}." + ) + if self.qpos.dim() not in (1, 2, 3) or self.qpos.shape[-1] == 0: + raise ValueError( + "qpos must have shape (control_dof,), (n_envs, control_dof), " + "or (n_envs, n_waypoint, control_dof), " + f"got {tuple(self.qpos.shape)}." + ) + + +@dataclass(frozen=True, slots=True, eq=False) +class NamedJointPositionTarget(ActionTarget): + """Named joint-space target resolved from :class:`MoveJointsCfg`.""" + + name: str + """Name of a joint-position target in ``MoveJointsCfg.named_joint_positions``.""" + + def __post_init__(self) -> None: + if not isinstance(self.name, str): + raise TypeError(f"name must be a str, got {type(self.name).__name__}.") + if not self.name.strip(): + raise ValueError("name must not be empty.") + + @configclass class MoveJointsCfg(ActionCfg): name: str = "move_joints" @@ -47,7 +88,7 @@ class MoveJointsCfg(ActionCfg): """Optional named joint targets resolved by ``NamedJointPositionTarget``.""" -class MoveJoints(AtomicAction): +class MoveJoints(AtomicAction[JointPositionTarget | NamedJointPositionTarget]): """Plan a joint-space move for the configured control part. The :class:`JointPositionTarget` may carry either a single waypoint @@ -104,10 +145,8 @@ def execute( return ActionResult( success=success, trajectory=full, - next_state=WorldState( + next_state=state.with_updates( last_qpos=full[:, -1, :].clone(), - held_object=state.held_object, - coordinated_held_object=state.coordinated_held_object, ), ) @@ -138,4 +177,9 @@ def _embed( return full -__all__ = ["MoveJoints", "MoveJointsCfg"] +__all__ = [ + "JointPositionTarget", + "MoveJoints", + "MoveJointsCfg", + "NamedJointPositionTarget", +] diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index 3f88d35f..89a0af82 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -19,6 +19,7 @@ from __future__ import annotations import math +from dataclasses import dataclass from typing import ClassVar import torch @@ -38,14 +39,33 @@ ActionCfg, ActionResult, AtomicAction, - GraspTarget, HeldObjectState, ObjectSemantics, WorldState, + _validate_pose_tensor, ) +from ..targets import ObjectActionTarget from ..trajectory import TrajectoryBuilder +@dataclass(frozen=True, slots=True, eq=False) +class GraspTarget(ObjectActionTarget): + """Pickup target with an affordance-selected or supplied grasp pose.""" + + 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. + """ + + def __post_init__(self) -> None: + ObjectActionTarget.__post_init__(self) + if self.grasp_xpos is not None: + _validate_pose_tensor(self.grasp_xpos, "grasp_xpos", allow_waypoints=False) + + @configclass class PickUpCfg(ActionCfg): name: str = "pick_up" @@ -88,7 +108,7 @@ class PickUpCfg(ActionCfg): """Optional rotation (radians) about the grasp x-axis to apply after grasp selection.""" -class PickUp(AtomicAction): +class PickUp(AtomicAction[GraspTarget]): """Approach a grasp pose, close the gripper, lift.""" TargetType: ClassVar[type] = GraspTarget @@ -241,13 +261,20 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: held = HeldObjectState( semantics=sem, object_to_eef=object_to_eef, grasp_xpos=grasp_xpos ) + held_objects = dict(state.held_objects) + held_objects[self.cfg.control_part] = held + coordinated_held_objects = { + key: value + for key, value in state.coordinated_held_objects.items() + if self.cfg.control_part not in key + } return ActionResult( success=success, trajectory=full, - next_state=WorldState( + next_state=state.with_updates( last_qpos=full[:, -1, :].clone(), - held_object=held, - coordinated_held_object=state.coordinated_held_object, + held_objects=held_objects, + coordinated_held_objects=coordinated_held_objects, ), ) @@ -483,4 +510,4 @@ def _upright_adjusted_grasp_poses( return adjusted_grasp_xpos -__all__ = ["PickUp", "PickUpCfg"] +__all__ = ["GraspTarget", "PickUp", "PickUpCfg"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/place.py b/embodichain/lab/sim/atomic_actions/primitives/place.py index 5198d873..1a391d40 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/place.py +++ b/embodichain/lab/sim/atomic_actions/primitives/place.py @@ -18,7 +18,8 @@ from __future__ import annotations -from typing import ClassVar +from dataclasses import dataclass +from typing import ClassVar, Literal import torch @@ -28,14 +29,45 @@ from ._helpers import arm_qpos_from_state from ..core import ( + ActionTarget, ActionCfg, ActionResult, AtomicAction, - EndEffectorPoseTarget, WorldState, + _validate_pose_tensor, ) from ..trajectory import TrajectoryBuilder +TcpSymmetry = Literal["none", "z_roll_180"] + + +@dataclass(frozen=True, slots=True, eq=False) +class PlaceTarget(ActionTarget): + """End-effector release-pose target used by :class:`Place`.""" + + xpos: torch.Tensor + """Target end-effector release pose. + + Accepts ``(4, 4)``, ``(n_envs, 4, 4)``, or + ``(n_envs, n_waypoint, 4, 4)``. + """ + + tcp_symmetry: TcpSymmetry = "none" + """Optional TCP-frame symmetry allowed by the placement semantics. + + ``"none"`` preserves the pose exactly. ``"z_roll_180"`` lets placement + 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: + _validate_pose_tensor(self.xpos, "xpos", allow_waypoints=True) + 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}" + ) + @configclass class PlaceCfg(ActionCfg): @@ -67,10 +99,10 @@ class PlaceCfg(ActionCfg): """Number of fixed-orientation Cartesian keyframes per translation segment.""" -class Place(AtomicAction): +class Place(AtomicAction[PlaceTarget]): """Lower the held object to a place pose, open the gripper, retract. - The :class:`EndEffectorPoseTarget` may carry either a single waypoint + The :class:`PlaceTarget` 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 @@ -79,7 +111,7 @@ class Place(AtomicAction): joint positions are inherited from ``WorldState.last_qpos``. """ - TargetType: ClassVar[type] = EndEffectorPoseTarget + TargetType: ClassVar[type] = PlaceTarget def __init__( self, @@ -105,7 +137,7 @@ def __init__( if self.cfg.cartesian_waypoint_count < 1: logger.log_error("cartesian_waypoint_count must be at least 1.", ValueError) - def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionResult: + def execute(self, target: PlaceTarget, 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) @@ -200,13 +232,20 @@ def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionRes full[:, n_down_actual + n_open :, self.arm_joint_ids] = back_arm full[:, n_down_actual + n_open :, self.hand_joint_ids] = self.hand_open_qpos + held_objects = dict(state.held_objects) + held_objects.pop(self.cfg.control_part, None) + coordinated_held_objects = { + key: value + for key, value in state.coordinated_held_objects.items() + if self.cfg.control_part not in key + } return ActionResult( success=success, trajectory=full, - next_state=WorldState( + next_state=state.with_updates( last_qpos=full[:, -1, :].clone(), - held_object=None, - coordinated_held_object=state.coordinated_held_object, + held_objects=held_objects, + coordinated_held_objects=coordinated_held_objects, ), ) @@ -296,4 +335,4 @@ def _select_tcp_symmetric_place_variant( ] -__all__ = ["Place", "PlaceCfg"] +__all__ = ["Place", "PlaceCfg", "PlaceTarget"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/press.py b/embodichain/lab/sim/atomic_actions/primitives/press.py index 82d6fd44..f1606d54 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/press.py +++ b/embodichain/lab/sim/atomic_actions/primitives/press.py @@ -18,6 +18,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import ClassVar import torch @@ -27,15 +28,27 @@ from ._helpers import arm_qpos_from_state from ..core import ( + ActionTarget, ActionCfg, ActionResult, AtomicAction, - EndEffectorPoseTarget, WorldState, + _validate_pose_tensor, ) from ..trajectory import TrajectoryBuilder +@dataclass(frozen=True, slots=True, eq=False) +class PressTarget(ActionTarget): + """Single end-effector contact pose used by :class:`Press`.""" + + xpos: torch.Tensor + """Contact pose, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" + + def __post_init__(self) -> None: + _validate_pose_tensor(self.xpos, "xpos", allow_waypoints=False) + + @configclass class PressCfg(ActionCfg): name: str = "press" @@ -54,10 +67,10 @@ class PressCfg(ActionCfg): """Joint positions for the closed hand state, shape ``[hand_dof,]``.""" -class Press(AtomicAction): +class Press(AtomicAction[PressTarget]): """Close the gripper, press down to a target pose, then return.""" - TargetType: ClassVar[type] = EndEffectorPoseTarget + TargetType: ClassVar[type] = PressTarget def __init__( self, @@ -83,7 +96,7 @@ def __init__( hand_dof=self.hand_dof, ) - def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionResult: + def execute(self, target: PressTarget, 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), @@ -149,10 +162,8 @@ def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionRes return ActionResult( success=success, trajectory=full, - next_state=WorldState( + next_state=state.with_updates( last_qpos=full[:, -1, :].clone(), - held_object=state.held_object, - coordinated_held_object=state.coordinated_held_object, ), ) @@ -186,4 +197,4 @@ def _fail(self, state: WorldState) -> ActionResult: ) -__all__ = ["Press", "PressCfg"] +__all__ = ["Press", "PressCfg", "PressTarget"] diff --git a/embodichain/lab/sim/atomic_actions/targets.py b/embodichain/lab/sim/atomic_actions/targets.py new file mode 100644 index 00000000..414a432e --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/targets.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. +# ---------------------------------------------------------------------------- + +"""Shared target contracts for object-centric atomic actions.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .core import ActionTarget, ObjectSemantics + + +@dataclass(frozen=True, slots=True, eq=False) +class ObjectActionTarget(ActionTarget): + """Base target for atomic actions operating on a semantic object. + + Concrete actions add only the pose roles and constraints they actually + consume. This shared contract deliberately does not define a generic pose: + an object pose, a single-arm grasp pose, and a dual-arm grasp pair have + different meanings and shapes. + """ + + semantics: ObjectSemantics + """Semantic description of the object on which the action operates.""" + + def __post_init__(self) -> None: + if not isinstance(self.semantics, ObjectSemantics): + raise TypeError( + "semantics must be an ObjectSemantics, " + f"got {type(self.semantics).__name__}." + ) + + +__all__ = ["ObjectActionTarget"] diff --git a/scripts/benchmark/atomic_action/move_held_object_benchmark.py b/scripts/benchmark/atomic_action/move_held_object_benchmark.py index 09364365..97e7d14e 100644 --- a/scripts/benchmark/atomic_action/move_held_object_benchmark.py +++ b/scripts/benchmark/atomic_action/move_held_object_benchmark.py @@ -236,7 +236,7 @@ def _prepare_held_state( ("pick_up", GraspTarget(semantics=semantics)), ] ) - if not is_success or state.held_object is None: + if not is_success or state.get_held_object("arm") is None: raise RuntimeError( "Failed to prepare held-object state for MoveHeldObject benchmark." ) @@ -380,7 +380,7 @@ def _on_validation_step(waypoint_index: int) -> None: 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) + still_held = bool(is_success and final_state.get_held_object("arm") is not None) physical_pick_success = bool( object_lift_delta_m is not None and object_lift_delta_m >= PHYSICAL_PICK_MIN_LIFT_M diff --git a/scripts/benchmark/atomic_action/pickup_benchmark.py b/scripts/benchmark/atomic_action/pickup_benchmark.py index 65350e88..9c8f56ff 100644 --- a/scripts/benchmark/atomic_action/pickup_benchmark.py +++ b/scripts/benchmark/atomic_action/pickup_benchmark.py @@ -222,7 +222,9 @@ def _on_validation_step(waypoint_index: int) -> None: 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) + held_created = bool( + is_success and final_state.get_held_object("arm") is not None + ) physical_pick_success = bool( held_created and object_lift_delta_m is not None diff --git a/scripts/benchmark/atomic_action/place_benchmark.py b/scripts/benchmark/atomic_action/place_benchmark.py index 3f120fa6..39d30100 100644 --- a/scripts/benchmark/atomic_action/place_benchmark.py +++ b/scripts/benchmark/atomic_action/place_benchmark.py @@ -217,7 +217,7 @@ def _prepare_held_state( is_success, traj, state = atomic_engine.run( steps=[("pick_up", GraspTarget(semantics=semantics))] ) - if not is_success or state.held_object is None: + if not is_success or state.get_held_object("arm") 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 @@ -242,9 +242,9 @@ def _run_case( """Run one Place benchmark case.""" from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, - EndEffectorPoseTarget, Place, PlaceCfg, + PlaceTarget, ) from scripts.tutorials.atomic_action.place import ( compute_pick_close_end_step, @@ -299,7 +299,7 @@ def _run_case( 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))], + steps=[("place", PlaceTarget(xpos=place_pose))], state=state, ) ) @@ -357,7 +357,7 @@ def _on_validation_step(waypoint_index: int) -> None: reset_robot(robot, initial_qpos) reset_rigid_object(obj, initial_obj_pose) - released = bool(is_success and final_state.held_object is None) + released = bool(is_success and final_state.get_held_object("arm") is None) physical_pick_success = bool( object_lift_delta_m is not None and object_lift_delta_m >= PHYSICAL_PICK_MIN_LIFT_M diff --git a/scripts/benchmark/atomic_action/press_benchmark.py b/scripts/benchmark/atomic_action/press_benchmark.py index e9dc9a0a..54942934 100644 --- a/scripts/benchmark/atomic_action/press_benchmark.py +++ b/scripts/benchmark/atomic_action/press_benchmark.py @@ -87,6 +87,7 @@ def _ensure_runtime_imports() -> None: MoveEndEffectorCfg as move_end_effector_cfg_cls, Press as press_cls, PressCfg as press_cfg_cls, + PressTarget as press_target_cls, ) from embodichain.lab.sim.cfg import ( RigidBodyAttributesCfg as rigid_body_attributes_cfg_cls, @@ -127,6 +128,7 @@ def _ensure_runtime_imports() -> None: "MoveEndEffectorCfg": move_end_effector_cfg_cls, "Press": press_cls, "PressCfg": press_cfg_cls, + "PressTarget": press_target_cls, "RigidBodyAttributesCfg": rigid_body_attributes_cfg_cls, "RigidObjectCfg": rigid_object_cfg_cls, "VisualMaterialCfg": visual_material_cfg_cls, @@ -580,7 +582,7 @@ def _timed_atomic_run( is_success, traj, _ = atomic_engine.run( steps=[ ("move_end_effector", EndEffectorPoseTarget(xpos=move_target)), - ("press", EndEffectorPoseTarget(xpos=press_target)), + ("press", PressTarget(xpos=press_target)), ] ) _sync_cuda() diff --git a/scripts/tutorials/atomic_action/coordinated_pickment.py b/scripts/tutorials/atomic_action/coordinated_pickment.py index cb00e9fe..a9f1a06a 100644 --- a/scripts/tutorials/atomic_action/coordinated_pickment.py +++ b/scripts/tutorials/atomic_action/coordinated_pickment.py @@ -42,9 +42,9 @@ from embodichain.lab.sim.atomic_actions import ( Affordance, AtomicActionEngine, + CoordinatedPickTarget, CoordinatedPickment, CoordinatedPickmentCfg, - CoordinatedPickmentTarget, ObjectSemantics, ) from embodichain.lab.sim.cfg import ( @@ -67,6 +67,7 @@ draw_axis_marker, prepare_tutorial_scene, replay_trajectory, + run_tutorial, ) ARM_URDF_PATH = "UniversalRobots/UR5/UR5.urdf" @@ -723,9 +724,9 @@ def run_coordinated_pickment_demo( 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( + pickment_target = CoordinatedPickTarget( + semantics=object_semantics, 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), @@ -792,4 +793,4 @@ def main() -> None: if __name__ == "__main__": - main() + run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/coordinated_placement.py b/scripts/tutorials/atomic_action/coordinated_placement.py index b717d6a4..27f46d47 100644 --- a/scripts/tutorials/atomic_action/coordinated_placement.py +++ b/scripts/tutorials/atomic_action/coordinated_placement.py @@ -73,6 +73,7 @@ make_ur5_solver_cfg, prepare_tutorial_scene, replay_trajectory, + run_tutorial, ) DEFAULT_MESH_FRAME_CORRECTION_EULER_DEG = (-90.0, 0.0, 0.0) @@ -871,7 +872,7 @@ def run_coordinated_placement_demo( return left_pick_traj = left_pick_result.trajectory state = left_pick_result.next_state - bread_held_state = state.held_object + bread_held_state = state.get_held_object("left_arm") 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) @@ -898,7 +899,7 @@ def run_coordinated_placement_demo( return right_pick_traj = right_pick_result.trajectory state = right_pick_result.next_state - pan_held_state = state.held_object + pan_held_state = state.get_held_object("right_arm") 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) @@ -965,6 +966,10 @@ def log_pick_execution(step_idx: int, total_steps: int) -> None: "right_arm", sim.device, ) + held_objects = dict(state.held_objects) + held_objects["left_arm"] = bread_held_state + held_objects["right_arm"] = pan_held_state + state = state.with_updates(held_objects=held_objects) support_target_pose = build_support_object_target_pose(pan_pose, sim.device) placing_target_pose = build_placing_object_target_pose( @@ -994,8 +999,6 @@ def log_pick_execution(step_idx: int, total_steps: int) -> None: 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, @@ -1071,4 +1074,4 @@ def main() -> None: if __name__ == "__main__": - main() + run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/move_end_effector.py b/scripts/tutorials/atomic_action/move_end_effector.py index 9ef606ad..dff45def 100644 --- a/scripts/tutorials/atomic_action/move_end_effector.py +++ b/scripts/tutorials/atomic_action/move_end_effector.py @@ -46,6 +46,7 @@ make_top_down_eef_pose, prepare_tutorial_scene, replay_trajectory, + run_tutorial, ) MOVE_SAMPLE_INTERVAL = 80 @@ -125,4 +126,4 @@ def main() -> None: if __name__ == "__main__": - main() + run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/move_held_object.py b/scripts/tutorials/atomic_action/move_held_object.py index d0acc5af..d30b8c7d 100644 --- a/scripts/tutorials/atomic_action/move_held_object.py +++ b/scripts/tutorials/atomic_action/move_held_object.py @@ -58,6 +58,7 @@ make_eef_pose_at, prepare_tutorial_scene, replay_trajectory, + run_tutorial, ) OBJECT_MESH_PATH = "PaperCup/paper_cup.ply" @@ -211,4 +212,4 @@ def clear_object_dynamics(step_idx: int, _: int) -> None: if __name__ == "__main__": - main() + run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/move_joints.py b/scripts/tutorials/atomic_action/move_joints.py index e4dd06fd..2e0addb1 100644 --- a/scripts/tutorials/atomic_action/move_joints.py +++ b/scripts/tutorials/atomic_action/move_joints.py @@ -44,6 +44,7 @@ draw_axis_marker, prepare_tutorial_scene, replay_trajectory, + run_tutorial, ) MOVE_JOINTS_SAMPLE_INTERVAL = 80 @@ -125,4 +126,4 @@ def main() -> None: if __name__ == "__main__": - main() + run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/pickup.py b/scripts/tutorials/atomic_action/pickup.py index 2aa7b308..cf0acef4 100644 --- a/scripts/tutorials/atomic_action/pickup.py +++ b/scripts/tutorials/atomic_action/pickup.py @@ -50,6 +50,7 @@ initialize_pre_pick_robot_pose, prepare_tutorial_scene, replay_trajectory, + run_tutorial, ) OBJECT_SIZE = (0.05, 0.05, 0.05) @@ -188,4 +189,4 @@ def clear_object_dynamics(step_idx: int, _: int) -> None: if __name__ == "__main__": - main() + run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/place.py b/scripts/tutorials/atomic_action/place.py index 6a8752d5..fb3ae0ef 100644 --- a/scripts/tutorials/atomic_action/place.py +++ b/scripts/tutorials/atomic_action/place.py @@ -31,12 +31,12 @@ from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, - EndEffectorPoseTarget, GraspTarget, PickUp, PickUpCfg, Place, PlaceCfg, + PlaceTarget, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.lab.sim.objects import RigidObject @@ -55,6 +55,7 @@ initialize_pre_pick_robot_pose, prepare_tutorial_scene, replay_trajectory, + run_tutorial, ) OBJECT_SIZE = (0.05, 0.05, 0.05) @@ -179,7 +180,7 @@ def main() -> None: ("pick_up", GraspTarget(semantics)), ( "place", - EndEffectorPoseTarget( + PlaceTarget( broadcast_waypoint_pose_batch( place_poses, robot.get_qpos().shape[0] ) @@ -218,4 +219,4 @@ def clear_object_dynamics(step_idx: int, _: int) -> None: if __name__ == "__main__": - main() + run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py index 10deaad3..39cd829d 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -36,6 +36,7 @@ MoveEndEffectorCfg, Press, PressCfg, + PressTarget, ) from embodichain.lab.sim.cfg import ( RigidBodyAttributesCfg, @@ -56,6 +57,7 @@ make_top_down_eef_pose, prepare_tutorial_scene, replay_trajectory, + run_tutorial, ) MOVE_SAMPLE_INTERVAL = 60 @@ -192,7 +194,7 @@ def main() -> None: success, trajectory, _ = engine.run( [ ("move_end_effector", EndEffectorPoseTarget(move_target)), - ("press", EndEffectorPoseTarget(press_target)), + ("press", PressTarget(press_target)), ] ) if not success.all(): @@ -238,4 +240,4 @@ def log_state(step_idx: int, total_steps: int) -> None: if __name__ == "__main__": - main() + run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/tutorial_utils.py b/scripts/tutorials/atomic_action/tutorial_utils.py index e36d58d8..3fe3db7e 100644 --- a/scripts/tutorials/atomic_action/tutorial_utils.py +++ b/scripts/tutorials/atomic_action/tutorial_utils.py @@ -133,6 +133,29 @@ def create_tutorial_simulation( return sim +def run_tutorial(main: Callable[[], None]) -> None: + """Run a tutorial entry point and release its simulation at top level. + + ``SimulationManager.destroy`` defers native cleanup because action and + planner locals may still retain wrapped C++ objects. Calling it only after + ``main`` has unwound makes the cleanup deterministic and avoids native + teardown during Python interpreter finalization. + + Args: + main: Zero-argument tutorial entry point. + """ + try: + main() + finally: + if SimulationManager.is_instantiated(): + sim = SimulationManager.get_instance() + if sim.is_window_recording(): + sim.stop_window_record() + sim.wait_window_record_saves() + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + def add_ur5_gripper_robot( sim: SimulationManager, init_pos: Sequence[float] = (0.0, 0.0, 0.0), @@ -657,6 +680,7 @@ def create_ur5_gripper_robot_cfg( "get_tutorial_window_size", "prepare_tutorial_scene", "replay_trajectory", + "run_tutorial", "should_open_tutorial_window", "should_wait_for_tutorial_input", "start_auto_play_recording", diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index 28d9f1fb..b63d5216 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -25,20 +25,24 @@ from embodichain.lab.sim.atomic_actions.affordance import ( AntipodalAffordance, ) +from embodichain.lab.sim.atomic_actions import ( + CoordinatedPickTarget, + CoordinatedPlacementTarget, + EndEffectorPoseTarget, + GraspTarget, + HeldObjectPoseTarget, + JointPositionTarget, + NamedJointPositionTarget, + PlaceTarget, + PressTarget, +) from embodichain.lab.sim.planners.utils import MoveType, PlanResult from embodichain.lab.sim.atomic_actions.core import ( ActionResult, AtomicAction, CoordinatedHeldObjectState, - CoordinatedPickmentTarget, - CoordinatedPlacementTarget, - GraspTarget, HeldObjectState, - HeldObjectPoseTarget, - JointPositionTarget, - NamedJointPositionTarget, ObjectSemantics, - EndEffectorPoseTarget, WorldState, ) from embodichain.lab.sim.atomic_actions.actions import ( @@ -238,8 +242,8 @@ def test_execute_returns_full_dof_trajectory(self): 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 + # MoveEndEffector preserves held-object mappings. + assert result.next_state.held_objects == {} def test_execute_with_multi_waypoint_visits_each_waypoint(self): action = MoveEndEffector(self.mg, MoveEndEffectorCfg(sample_interval=10)) @@ -323,7 +327,7 @@ def interpolate(trajectory, interp_num, device): ): result = action.execute( JointPositionTarget(qpos=target_qpos), - WorldState(last_qpos=last_qpos, held_object=held), + WorldState(last_qpos=last_qpos, held_objects={"arm": held}), ) assert result.success.all() @@ -331,7 +335,7 @@ def interpolate(trajectory, interp_num, device): 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) - assert result.next_state.held_object is held + assert result.next_state.get_held_object("arm") is held def test_execute_with_named_qpos_resolves_cfg_target(self): action = MoveJoints( @@ -470,8 +474,9 @@ def test_execute_populates_held_object_state(self): 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 + held_object = result.next_state.get_held_object("arm") + assert isinstance(held_object, HeldObjectState) + assert held_object.semantics is sem def test_execute_accepts_an_explicit_grasp_pose(self): action = PickUp( @@ -507,9 +512,10 @@ def test_execute_accepts_an_explicit_grasp_pose(self): ) assert result.success.all() - assert result.next_state.held_object is not None + held_object = result.next_state.get_held_object("arm") + assert held_object is not None assert torch.allclose( - result.next_state.held_object.grasp_xpos, + held_object.grasp_xpos, grasp_xpos.unsqueeze(0).repeat(NUM_ENVS, 1, 1), ) affordance.get_valid_grasp_poses.assert_not_called() @@ -556,9 +562,10 @@ def test_execute_chooses_symmetric_grasp_variant_closest_to_start_pose(self): assert result.success.all() assert result.success.shape == (NUM_ENVS,) - assert isinstance(result.next_state.held_object, HeldObjectState) + held_object = result.next_state.get_held_object("arm") + assert isinstance(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) + assert torch.allclose(held_object.grasp_xpos, expected_grasp) assert self.mg.robot.compute_batch_ik.call_count == 3 for call in self.mg.robot.compute_batch_ik.call_args_list: assert call.kwargs["pose"].shape == (NUM_ENVS, 2, 4, 4) @@ -617,7 +624,10 @@ def test_preserves_held_object(self): 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), ) - state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF), held_object=held) + state = WorldState( + last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF), + held_objects={"arm": held}, + ) with patch( "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", return_value=torch.zeros(NUM_ENVS, 10, ARM_DOF), @@ -628,7 +638,7 @@ def test_preserves_held_object(self): 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 + assert result.next_state.get_held_object("arm") is held def test_automatic_rotation_adjustment_is_isolated_per_environment(self): action = MoveHeldObject( @@ -663,7 +673,7 @@ def test_automatic_rotation_adjustment_is_isolated_per_environment(self): HeldObjectPoseTarget(object_target_pose=torch.eye(4)), WorldState( last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF), - held_object=held, + held_objects={"arm": held}, ), ) @@ -683,7 +693,7 @@ def setup_method(self): self.mg = _make_mock_motion_generator() def test_target_type_is_pose_target(self): - assert Place.TargetType is EndEffectorPoseTarget + assert Place.TargetType is PlaceTarget def test_rejects_non_positive_cartesian_waypoint_count(self): with pytest.raises(Exception, match="cartesian_waypoint_count"): @@ -712,18 +722,21 @@ def test_execute_clears_held_object(self): 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), ) - state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF), held_object=held) + state = WorldState( + last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF), + held_objects={"arm": held}, + ) 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 ), ): - result = action.execute(EndEffectorPoseTarget(xpos=torch.eye(4)), state) + result = action.execute(PlaceTarget(xpos=torch.eye(4)), state) 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 + assert result.next_state.get_held_object("arm") is None def test_execute_with_multi_waypoint_visits_each_waypoint(self): cfg = PlaceCfg( @@ -742,7 +755,10 @@ def test_execute_with_multi_waypoint_visits_each_waypoint(self): 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), ) - state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF), held_object=held) + state = WorldState( + last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF), + held_objects={"arm": held}, + ) pose0 = torch.eye(4) pose1 = torch.eye(4) @@ -771,12 +787,12 @@ def interpolate(trajectory, interp_num, device): "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", side_effect=interpolate, ): - result = action.execute(EndEffectorPoseTarget(xpos=multi_xpos), state) + result = action.execute(PlaceTarget(xpos=multi_xpos), state) 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 + assert result.next_state.get_held_object("arm") is None # IK order: down phase (approach, pose0, pose1) then back phase (retract). assert len(seen_poses) == 4 lift_height = cfg.lift_height @@ -831,7 +847,7 @@ def compute_ik(pose=None, name=None, joint_seed=None, **kwargs): ].repeat(1, interp_num, 1), ): result = action.execute( - EndEffectorPoseTarget(xpos=release_pose), + PlaceTarget(xpos=release_pose), WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), ) @@ -872,7 +888,7 @@ def compute_ik(pose=None, name=None, joint_seed=None, **kwargs): ].repeat(1, interp_num, 1), ): result = action.execute( - EndEffectorPoseTarget(xpos=target), + PlaceTarget(xpos=target), WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), ) @@ -908,7 +924,7 @@ def repeat_last_keyframe(trajectory, interp_num, device): "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", side_effect=repeat_last_keyframe, ): - result = action.execute(EndEffectorPoseTarget(xpos=rz_pi_pose), state) + result = action.execute(PlaceTarget(xpos=rz_pi_pose), state) assert result.success.all() assert result.success.shape == (NUM_ENVS,) @@ -945,7 +961,7 @@ def repeat_last_keyframe(trajectory, interp_num, device): side_effect=repeat_last_keyframe, ): result = action.execute( - EndEffectorPoseTarget( + PlaceTarget( xpos=rz_pi_pose, tcp_symmetry="z_roll_180", ), @@ -978,8 +994,8 @@ 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_target_type_is_press_target(self): + assert Press.TargetType is PressTarget def test_default_name_is_explicit(self): assert PressCfg(hand_close_qpos=_hand_close()).name == "press" @@ -1001,7 +1017,7 @@ def test_execute_closes_hand_and_preserves_held_object(self): ) 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) + state = WorldState(last_qpos=last_qpos, held_objects={"arm": held}) def interpolate(trajectory, interp_num, device): return trajectory[:, -1:, :].repeat(1, interp_num, 1) @@ -1010,7 +1026,7 @@ def interpolate(trajectory, interp_num, device): "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", side_effect=interpolate, ): - result = action.execute(EndEffectorPoseTarget(xpos=torch.eye(4)), state) + result = action.execute(PressTarget(xpos=torch.eye(4)), state) assert result.success.all() assert result.success.shape == (NUM_ENVS,) @@ -1021,7 +1037,7 @@ def interpolate(trajectory, interp_num, device): result.next_state.last_qpos[:, :ARM_DOF], last_qpos[:, :ARM_DOF], ) - assert result.next_state.held_object is held + assert result.next_state.get_held_object("arm") is held # --------------------------------------------------------------------------- @@ -1034,7 +1050,7 @@ 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.TargetType is CoordinatedPickTarget assert CoordinatedPickment.__bases__ == (AtomicAction,) def test_execute_returns_full_dof_trajectory_and_dual_held_state(self): @@ -1054,9 +1070,9 @@ def test_execute_returns_full_dof_trajectory_and_dual_held_state(self): ) state = WorldState(last_qpos=torch.zeros(NUM_ENVS, DUAL_TOTAL_DOF)) result = action.execute( - CoordinatedPickmentTarget( + CoordinatedPickTarget( + semantics=sem, 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), @@ -1074,11 +1090,11 @@ def test_execute_returns_full_dof_trajectory_and_dual_held_state(self): 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, + held_object = result.next_state.get_coordinated_held_object( + "left_arm", "right_arm" ) - assert result.next_state.held_object is None + assert isinstance(held_object, CoordinatedHeldObjectState) + assert result.next_state.held_objects == {} def test_execute_freezes_only_environment_with_partial_ik_failure(self): action = CoordinatedPickment( @@ -1119,9 +1135,9 @@ def fail_second_env_during_move( ) result = action.execute( - CoordinatedPickmentTarget( + CoordinatedPickTarget( + semantics=semantics, object_target_pose=target_pose, - object_semantics=semantics, left_object_to_eef=torch.eye(4), right_object_to_eef=torch.eye(4), object_initial_pose=torch.eye(4), @@ -1158,7 +1174,14 @@ def setup_method(self): ) self.action = CoordinatedPlacement(self.mg, cfg=self.cfg) - def _make_target(self) -> CoordinatedPlacementTarget: + def _make_target_and_state( + self, + ) -> tuple[ + CoordinatedPlacementTarget, + WorldState, + HeldObjectState, + HeldObjectState, + ]: placing_pose = torch.eye(4) placing_pose[0, 3] = 0.2 support_pose = torch.eye(4) @@ -1176,20 +1199,28 @@ def _make_target(self) -> CoordinatedPlacementTarget: support_semantics = ObjectSemantics( affordance=AntipodalAffordance(), geometry={}, label="support" ) - return CoordinatedPlacementTarget( + 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), + ) + target = 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), - ), ) + state = WorldState( + last_qpos=torch.zeros(NUM_ENVS, DUAL_TOTAL_DOF), + held_objects={ + "left_arm": placing_held_object, + "right_arm": support_held_object, + }, + ) + return target, state, placing_held_object, support_held_object def test_target_type_is_coordinated_placement_target(self): assert CoordinatedPlacement.TargetType is CoordinatedPlacementTarget @@ -1207,25 +1238,34 @@ def test_init_sets_dual_arm_and_hand_joint_ids(self): 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 - ) + target, state, placing_source, support_source = self._make_target_and_state() + ( + placing_xpos, + support_xpos, + release, + placing_held_state, + support_held_state, + ) = self.action._resolve_target(target, state) 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 placing_held_state.semantics is placing_source.semantics + assert support_held_state.semantics is support_source.semantics + assert support_held_state.object_to_eef.shape == (NUM_ENVS, 4, 4) + assert support_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 - ), + support_held_state.object_to_eef, + support_source.object_to_eef.unsqueeze(0).repeat(NUM_ENVS, 1, 1), ) + def test_resolve_target_requires_both_held_objects_in_world_state(self): + target, state, _, _ = self._make_target_and_state() + state.held_objects.pop("left_arm") + with pytest.raises(ValueError, match="left_arm"): + self.action._resolve_target(target, state) + 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 @@ -1234,8 +1274,7 @@ def test_segment_lengths_sum_to_sample_interval(self): 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)) + target, state, _, support_source = self._make_target_and_state() def interpolate(trajectory, interp_num, device): weights = torch.linspace( @@ -1271,13 +1310,12 @@ def interpolate(trajectory, interp_num, device): 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) + assert result.next_state.get_held_object("left_arm") is None + support_held_object = result.next_state.get_held_object("right_arm") + assert support_held_object is not None + assert support_held_object.semantics is support_source.semantics + assert support_held_object.object_to_eef.shape == (NUM_ENVS, 4, 4) + assert support_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 5710c60c..9fabb369 100644 --- a/tests/sim/atomic_actions/test_core.py +++ b/tests/sim/atomic_actions/test_core.py @@ -14,49 +14,140 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Tests for atomic_actions.core (typed targets, WorldState, ActionResult, ObjectSemantics).""" +"""Tests for atomic-action target contracts and shared core state.""" from __future__ import annotations import dataclasses +from typing import get_args import pytest import torch +import embodichain.lab.sim.atomic_actions.core as core_module from embodichain.lab.sim.atomic_actions.affordance import Affordance -from embodichain.lab.sim.atomic_actions.core import ( - ActionCfg, - ActionResult, - CoordinatedHeldObjectState, +from embodichain.lab.sim.atomic_actions import ( + BuiltinTarget, + CoordinatedPickTarget, CoordinatedPickmentTarget, CoordinatedPlacementTarget, + EndEffectorPoseTarget, GraspTarget, - HeldObjectState, HeldObjectPoseTarget, JointPositionTarget, NamedJointPositionTarget, + ObjectActionTarget, + PlaceTarget, + PressTarget, +) +from embodichain.lab.sim.atomic_actions.core import ( + ActionTarget, + ActionCfg, + ActionResult, + CoordinatedHeldObjectState, + HeldObjectState, ObjectSemantics, - EndEffectorPoseTarget, WorldState, ) class TestTypedTargets: + def test_core_does_not_own_concrete_target_types(self): + assert not hasattr(core_module, "GraspTarget") + + def test_builtin_target_contains_press_contract(self): + assert PressTarget in get_args(BuiltinTarget) + + def test_object_action_target_owns_shared_semantics_contract(self): + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + label="shared-object", + ) + target = ObjectActionTarget(semantics=semantics) + assert target.semantics is semantics + assert not hasattr(target, "xpos") + + def test_object_action_target_rejects_non_semantics_value(self): + with pytest.raises(TypeError, match="semantics"): + ObjectActionTarget(semantics=object()) # type: ignore[arg-type] + + def test_object_action_target_lives_in_neutral_module(self): + assert ( + ObjectActionTarget.__module__ + == "embodichain.lab.sim.atomic_actions.targets" + ) + + def test_object_action_target_is_not_a_builtin_executable_contract(self): + assert ObjectActionTarget not in get_args(BuiltinTarget) + + @pytest.mark.parametrize( + ("target_type", "owner_module"), + [ + ( + EndEffectorPoseTarget, + "embodichain.lab.sim.atomic_actions.primitives.move_end_effector", + ), + ( + JointPositionTarget, + "embodichain.lab.sim.atomic_actions.primitives.move_joints", + ), + ( + NamedJointPositionTarget, + "embodichain.lab.sim.atomic_actions.primitives.move_joints", + ), + (GraspTarget, "embodichain.lab.sim.atomic_actions.primitives.pick_up"), + ( + HeldObjectPoseTarget, + "embodichain.lab.sim.atomic_actions.primitives.move_held_object", + ), + (PlaceTarget, "embodichain.lab.sim.atomic_actions.primitives.place"), + (PressTarget, "embodichain.lab.sim.atomic_actions.primitives.press"), + ( + CoordinatedPickTarget, + "embodichain.lab.sim.atomic_actions.primitives.coordinated_pickment", + ), + ( + CoordinatedPlacementTarget, + "embodichain.lab.sim.atomic_actions.primitives.coordinated_placement", + ), + ], + ) + def test_target_is_defined_by_owning_primitive( + self, + target_type: type[ActionTarget], + owner_module: str, + ): + assert target_type.__module__ == owner_module + 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") + def test_place_target_can_declare_tcp_symmetry(self): + target = PlaceTarget(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): + def test_place_target_rejects_unknown_tcp_symmetry(self): with pytest.raises(ValueError, match="tcp_symmetry"): - EndEffectorPoseTarget( + PlaceTarget( xpos=torch.eye(4), tcp_symmetry="yaw_90" # type: ignore[arg-type] ) + def test_press_target_rejects_multiple_waypoints(self): + with pytest.raises(ValueError, match="xpos"): + PressTarget(xpos=torch.eye(4).reshape(1, 1, 4, 4)) + + def test_pose_target_rejects_invalid_shape(self): + with pytest.raises(ValueError, match="xpos"): + EndEffectorPoseTarget(xpos=torch.zeros(3, 3)) + + def test_pose_targets_use_identity_equality(self): + first = EndEffectorPoseTarget(xpos=torch.eye(4)) + second = EndEffectorPoseTarget(xpos=torch.eye(4)) + assert first == first + assert first != second + def test_pose_target_is_frozen(self): t = EndEffectorPoseTarget(xpos=torch.eye(4)) with pytest.raises(dataclasses.FrozenInstanceError): @@ -81,7 +172,9 @@ def test_named_joint_position_target_is_frozen(self): def test_grasp_target_holds_semantics(self): sem = ObjectSemantics(affordance=Affordance(), geometry={}, label="mug") - assert GraspTarget(semantics=sem).semantics is sem + target = GraspTarget(semantics=sem) + assert target.semantics is sem + assert isinstance(target, ObjectActionTarget) def test_grasp_target_is_frozen(self): sem = ObjectSemantics(affordance=Affordance(), geometry={}, label="mug") @@ -100,31 +193,26 @@ 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): + def test_coordinated_pick_target_holds_object_offsets(self): sem = ObjectSemantics(affordance=Affordance(), geometry={}, label="pencil") - target = CoordinatedPickmentTarget( + target = CoordinatedPickTarget( + semantics=sem, 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.semantics is sem + assert isinstance(target, ObjectActionTarget) assert target.left_object_to_eef.shape == (4, 4) + assert CoordinatedPickmentTarget is CoordinatedPickTarget - 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), - ) + def test_coordinated_placement_target_only_holds_desired_state(self): 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 isinstance(target, ActionTarget) + assert not hasattr(target, "placing_held_object") assert target.support_object_target_pose.shape == (4, 4) @@ -183,7 +271,8 @@ def test_constructs_with_last_qpos_only(self): qpos = torch.zeros(2, 6) ws = WorldState(last_qpos=qpos) assert ws.last_qpos is qpos - assert ws.held_object is None + assert ws.held_objects == {} + assert ws.coordinated_held_objects == {} def test_carries_held_object(self): sem = ObjectSemantics(affordance=Affordance(), geometry={}) @@ -192,8 +281,12 @@ def test_carries_held_object(self): object_to_eef=torch.eye(4).unsqueeze(0), grasp_xpos=torch.eye(4).unsqueeze(0), ) - ws = WorldState(last_qpos=torch.zeros(1, 6), held_object=held) - assert ws.held_object is held + ws = WorldState( + last_qpos=torch.zeros(1, 6), + held_objects={"left_arm": held}, + ) + assert ws.get_held_object("left_arm") is held + assert ws.get_held_object("right_arm") is None def test_carries_coordinated_held_object(self): sem = ObjectSemantics(affordance=Affordance(), geometry={}) @@ -204,8 +297,21 @@ def test_carries_coordinated_held_object(self): 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 + ws = WorldState( + last_qpos=torch.zeros(1, 14), + coordinated_held_objects={("left_arm", "right_arm"): held}, + ) + assert ws.get_coordinated_held_object("left_arm", "right_arm") is held + + def test_with_updates_does_not_alias_held_state_dictionaries(self): + ws = WorldState(last_qpos=torch.zeros(1, 6)) + successor = ws.with_updates(last_qpos=torch.ones(1, 6)) + successor.held_objects["arm"] = HeldObjectState( + semantics=ObjectSemantics(affordance=Affordance(), geometry={}), + object_to_eef=torch.eye(4).unsqueeze(0), + grasp_xpos=torch.eye(4).unsqueeze(0), + ) + assert ws.held_objects == {} class TestActionResult: diff --git a/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py index 5737ce9d..bb12e25c 100644 --- a/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py +++ b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py @@ -45,12 +45,14 @@ CuroboPlannerCfg, CuroboWorldCfg, ) -from embodichain.lab.sim.atomic_actions import AtomicActionEngine # noqa: E402 +from embodichain.lab.sim.atomic_actions import ( # noqa: E402 + AtomicActionEngine, + EndEffectorPoseTarget, +) from embodichain.lab.sim.atomic_actions.actions import ( # noqa: E402 MoveEndEffector, MoveEndEffectorCfg, ) -from embodichain.lab.sim.atomic_actions.core import EndEffectorPoseTarget # noqa: E402 ROBOT_UID = "curobo_franka" CONTROL_PART = "arm" diff --git a/tests/sim/atomic_actions/test_engine.py b/tests/sim/atomic_actions/test_engine.py index 3baa4d67..1d25fdb6 100644 --- a/tests/sim/atomic_actions/test_engine.py +++ b/tests/sim/atomic_actions/test_engine.py @@ -23,16 +23,20 @@ from unittest.mock import Mock from embodichain.lab.sim.atomic_actions.affordance import Affordance -from embodichain.lab.sim.atomic_actions.core import ( - ActionResult, - AtomicAction, +from embodichain.lab.sim.atomic_actions import ( + EndEffectorPoseTarget, GraspTarget, - HeldObjectState, HeldObjectPoseTarget, JointPositionTarget, NamedJointPositionTarget, + PlaceTarget, +) +from embodichain.lab.sim.atomic_actions.core import ( + ActionTarget, + ActionResult, + AtomicAction, + HeldObjectState, ObjectSemantics, - EndEffectorPoseTarget, WorldState, ) from embodichain.lab.sim.atomic_actions.engine import ( @@ -104,23 +108,23 @@ def execute(target, state): trajectory=torch.empty(NUM_ENVS, 0, TOTAL_DOF), next_state=state, ) - held = state.held_object + held_objects = dict(state.held_objects) if sets_held: sem = ObjectSemantics(affordance=Affordance(), geometry={}, label="x") - held = HeldObjectState( + held_objects["arm"] = 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), ) if clears_held: - held = None + held_objects.pop("arm", None) traj = torch.zeros(NUM_ENVS, 5, TOTAL_DOF) return ActionResult( success=True, trajectory=traj, - next_state=WorldState( + next_state=state.with_updates( last_qpos=traj[:, -1, :].clone(), - held_object=held, + held_objects=held_objects, ), ) @@ -160,7 +164,7 @@ def test_run_concatenates_trajectories(self): def test_run_threads_world_state(self): pick = _fake_action("pick", GraspTarget, sets_held=True) move = _fake_action("move", HeldObjectPoseTarget) - place = _fake_action("place", EndEffectorPoseTarget, clears_held=True) + place = _fake_action("place", PlaceTarget, clears_held=True) self.engine.register(pick, name="pick") self.engine.register(move, name="move") self.engine.register(place, name="place") @@ -169,15 +173,15 @@ def test_run_threads_world_state(self): [ ("pick", GraspTarget(sem)), ("move", HeldObjectPoseTarget(torch.eye(4))), - ("place", EndEffectorPoseTarget(torch.eye(4))), + ("place", PlaceTarget(torch.eye(4))), ] ) assert success.all().item() - # The move action saw a non-None held_object (set by pick). + # The move action saw the 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 + assert move_state_arg.get_held_object("arm") is not None # Final state cleared by place. - assert final_state.held_object is None + assert final_state.get_held_object("arm") is None def test_run_stops_on_first_failure(self): a = _fake_action("a", EndEffectorPoseTarget) @@ -225,7 +229,22 @@ def test_run_seeds_state_from_robot_when_none_provided(self): # First call's state argument state_arg = a.execute.call_args_list[0].args[1] assert state_arg.last_qpos.shape == (NUM_ENVS, TOTAL_DOF) - assert state_arg.held_object is None + assert state_arg.held_objects == {} assert state_arg.last_qpos.data_ptr() != seed_qpos.data_ptr() seed_qpos.fill_(1.0) assert torch.equal(state_arg.last_qpos, torch.zeros(NUM_ENVS, TOTAL_DOF)) + + def test_run_accepts_third_party_action_target(self): + class CustomTarget(ActionTarget): + pass + + action = _fake_action("custom", CustomTarget) + self.engine.register(action) + success, traj, _ = self.engine.run([("custom", CustomTarget())]) + assert success.all().item() + assert traj.shape == (NUM_ENVS, 5, TOTAL_DOF) + + def test_register_rejects_non_action_target_type(self): + action = _fake_action("invalid", str) + with pytest.raises(TypeError, match="ActionTarget"): + self.engine.register(action) diff --git a/tests/sim/atomic_actions/test_engine_per_env.py b/tests/sim/atomic_actions/test_engine_per_env.py index b71b3a3b..622d93e3 100644 --- a/tests/sim/atomic_actions/test_engine_per_env.py +++ b/tests/sim/atomic_actions/test_engine_per_env.py @@ -22,14 +22,14 @@ import pytest from unittest.mock import Mock -from embodichain.lab.sim.atomic_actions.engine import AtomicActionEngine +from embodichain.lab.sim.atomic_actions import EndEffectorPoseTarget from embodichain.lab.sim.atomic_actions.core import ( + ActionCfg, ActionResult, AtomicAction, WorldState, - EndEffectorPoseTarget, - ActionCfg, ) +from embodichain.lab.sim.atomic_actions.engine import AtomicActionEngine class _StubAction(AtomicAction): diff --git a/tests/sim/atomic_actions/test_motion_source_e2e.py b/tests/sim/atomic_actions/test_motion_source_e2e.py index 87b23d12..93f6535d 100644 --- a/tests/sim/atomic_actions/test_motion_source_e2e.py +++ b/tests/sim/atomic_actions/test_motion_source_e2e.py @@ -24,12 +24,14 @@ 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 import ( + AtomicActionEngine, + EndEffectorPoseTarget, +) from embodichain.lab.sim.atomic_actions.actions import ( MoveEndEffector, MoveEndEffectorCfg, ) -from embodichain.lab.sim.atomic_actions.core import EndEffectorPoseTarget @pytest.mark.requires_sim