diff --git a/.agents/skills/add-atomic-action/SKILL.md b/.agents/skills/add-atomic-action/SKILL.md index 20b13b301..a007497c4 100644 --- a/.agents/skills/add-atomic-action/SKILL.md +++ b/.agents/skills/add-atomic-action/SKILL.md @@ -1,341 +1,253 @@ --- name: add-atomic-action -description: Use when adding a new simulation atomic action or motion primitive to EmbodiChain's AtomicActionEngine. +description: Add a new simulation atomic action or motion primitive to EmbodiChain's typed planning and execution framework. Use when implementing a new skill, goal contract, action planner, symbolic effect, registration entry, documentation, and tests for AtomicActionEngine. --- # Add Atomic Action -Scaffold a new atomic action following EmbodiChain's `AtomicAction` pattern: a typed -target, a `WorldState` threaded across actions, and an `ActionResult` carrying a -full-DoF trajectory. +Add an action-owned goal and a side-effect-free `AtomicAction.plan()` +implementation. The engine owns all motion-planning resources; action +constructors accept only optional typed default options. Keep task-graph/MLLM +logic, simulator stepping, controller I/O, and physical-effect commits outside +the action. -## When to Use +## Read the current contracts -- User asks to add a new motion primitive (push, wipe, insert, hand-over, …) -- User says "add a new atomic action", "create a custom action", "implement a push action" -- User wants to extend `AtomicActionEngine` with a behaviour not covered by the built-ins - -## Key Files +Inspect only the files relevant to the requested skill: | Purpose | Path | -|---------|------| -| Base classes (`ActionCfg`, `AtomicAction`, `WorldState`, `ActionResult`, typed targets, `ObjectSemantics`) | `embodichain/lab/sim/atomic_actions/core.py` | -| Affordance types (`Affordance`, `AntipodalAffordance`, `InteractionPoints`) | `embodichain/lab/sim/atomic_actions/affordance.py` | -| Stateless trajectory helpers (`TrajectoryBuilder`) | `embodichain/lab/sim/atomic_actions/trajectory.py` | -| Built-in action primitives (reference implementations) | `embodichain/lab/sim/atomic_actions/primitives/` | -| Backward-compatible action re-export module | `embodichain/lab/sim/atomic_actions/actions.py` | -| Engine + global registry (`register_action`, `AtomicActionEngine.register` / `run`) | `embodichain/lab/sim/atomic_actions/engine.py` | -| Public API exports | `embodichain/lab/sim/atomic_actions/__init__.py` | -| Reference docs | `docs/source/overview/sim/atomic_actions/index.md`, `docs/source/overview/sim/atomic_actions/builtin_actions.md` | - -## The Contract (read first) - -Every atomic action is a **sibling** inheriting `AtomicAction` directly — do **not** -inherit from `MoveEndEffector` or any other action. Each action: - -1. Declares `TargetType: ClassVar[type | tuple[type, ...]]` — the concrete target dataclass, - or tuple of dataclasses, it accepts. -2. Holds `self.builder = TrajectoryBuilder(motion_generator)` for shared trajectory math. -3. Implements exactly one method: `execute(self, target, state: WorldState) -> ActionResult`. - - `target` is an instance of `self.TargetType`. - - `state.last_qpos` is the full-robot qpos `(n_envs, robot.dof)` to plan from; - `state.held_object` is the object currently grasped (or `None`). - - Returns `ActionResult(success, trajectory, next_state)` where `trajectory` is - full-DoF shaped `(n_envs, n_waypoints, robot.dof)` and `next_state` is the - successor `WorldState` (advance `last_qpos` to the trajectory's final row; - set/clear/preserve `held_object` per the action's semantics). - - `success` is a per-environment boolean tensor of shape `(n_envs,)` (or a - scalar bool). Use `ActionResult.success_all` (or `.success.all()`) when you - need a single aggregate boolean. - -There is **no** `validate` method, **no** `**kwargs`, **no** `start_qpos` parameter, -**no** `updates_held_object_state` flag, and **no** `get_held_object_state`. The -`WorldState` is the single channel for inter-action state. - -`ActionCfg` (and therefore every action cfg) carries one motion-source field used -by `TrajectoryBuilder.plan_arm_traj`: -- `motion_source: str = "ik_interp"` — `"ik_interp"` (batched IK + interpolation) - or `"motion_gen"` (delegates to the batched `MotionGenerator`). - -The `MotionGenerator` owns exactly one planner, selected by -`MotionGenCfg.planner_cfg`; action configs do not repeat or override that planner -type. - -## Steps - -### 1. Define the config - -Add a `@configclass`-decorated class that extends `ActionCfg` **directly** (the cfg -hierarchy is flat — do not inherit from another action's cfg). For a built-in -primitive, place the config beside the action class in -`embodichain/lab/sim/atomic_actions/primitives/.py`. +|---|---| +| Base action and descriptors | `embodichain/lab/sim/atomic_actions/core.py` | +| Goals and dynamic pose references | `embodichain/lab/sim/atomic_actions/goals.py` | +| Role-to-resource binding | `embodichain/lab/sim/atomic_actions/bindings.py` | +| Invocation, options, and resolved request | `embodichain/lab/sim/atomic_actions/invocation.py` | +| Control-part semantic commands | `embodichain/lab/sim/atomic_actions/control.py` | +| Invocation policies | `embodichain/lab/sim/atomic_actions/policies.py` | +| Robot/task/scene state | `embodichain/lab/sim/atomic_actions/state.py` | +| Effects and plans | `embodichain/lab/sim/atomic_actions/effects.py`, `plans.py` | +| Trajectory helpers | `embodichain/lab/sim/atomic_actions/trajectory.py` | +| Engine-owned planning resources | `embodichain/lab/sim/atomic_actions/runtime.py` | +| Reference implementations | `embodichain/lab/sim/atomic_actions/primitives/` | +| Static compiler and execution session | `engine.py`, `execution.py` | + +The public contract is: ```python -from __future__ import annotations +plan = engine.plan(invocation: ActionInvocation[Goal], context: PlanningContext) +``` -import torch +Use `engine.plan_action(action, invocation, context)` for a configured action +that is intentionally not in the stable skill registry. Never pass a motion +generator to an action constructor. + +Do not add compatibility code for `ActionTarget`, `WorldState`, `ActionResult`, +`execute()`, or `AtomicActionEngine.run()`. + +## 1. Define the goal + +Place a frozen action-owned dataclass beside the action. Add a stable +`goal_kind: ClassVar[str]`. Do not inherit a marker base merely for dispatch; +declare the accepted type on the action instead. -from embodichain.utils import configclass -from embodichain.lab.sim.atomic_actions.core import ActionCfg +```python +from dataclasses import dataclass +from typing import ClassVar + +import torch -@configclass -class PushCfg(ActionCfg): - name: str = "push" # must match the engine registration key - push_distance: float = 0.05 # metres to push forward - sample_interval: int = 30 # waypoints for the push phase - control_part: str = "arm" +@dataclass(frozen=True, slots=True, eq=False) +class PushGoal: + goal_kind: ClassVar[str] = "push" + contact_pose: torch.Tensor ``` -**Rules:** -- `name` must be unique and match the key used to register the action with the engine. -- Inherit from `ActionCfg` directly. If the action needs hand open/close fields, - declare them on this cfg (see `PickUpCfg` for the pattern) — do not invent a - shared `GraspActionCfg` parent. -- All fields must have defaults. +Keep only semantic intent in the goal. Do not include arm/hand names, planner +options, retry counts, live state, or a generic optional field bag. Use +`SceneEntityPose` for a pose that must be resolved again when the scene moves. +Use `ObjectActionGoal` only when the shared `semantics` field is genuinely +required. -### 2. Define a typed target (if needed) +## 2. Define runtime options and control commands -Reuse an existing target when it fits (`EndEffectorPoseTarget(xpos)` for an EEF-pose target, -`JointPositionTarget(qpos)` for an explicit control-part qpos target, -`NamedJointPositionTarget(name)` for a named qpos target resolved by action config, -`GraspTarget(semantics)` for a pickup, `HeldObjectPoseTarget(object_target_pose)` for -moving a grasped object). Only define a new frozen dataclass target when the action -needs inputs the existing targets don't carry. Put new targets in `core.py`. +Define a frozen `ActionOptions` subclass only when skill behavior may vary by +invocation. Examples include distances, grasp constraints, and phase split +counts. If no such behavior exists, use the base `ActionOptions`. + +Do not put `motion_source`, planner choice, sample count, control period, +velocity limits, collision policy, or recovery thresholds in skill options; +those belong to `MotionPolicy` or `RecoveryPolicy`. ```python from dataclasses import dataclass -@dataclass(frozen=True) -class PushTarget: - contact_pose: torch.Tensor # (4, 4) or (n_envs, 4, 4) EEF contact pose +from embodichain.lab.sim.atomic_actions import ActionOptions + + +@dataclass(frozen=True, slots=True, eq=False) +class PushOptions(ActionOptions): + push_distance: float = 0.05 ``` -### 3. Implement the action class +Do not put arm/hand names, hand qpos, or named robot postures in options. Bind +participants with `ActionBinding`. Register embodiment-specific commands such +as `open`, `grasp`, or `ready` on `ControlPartCommandProfile`; use +`ActionControlOverrides` only for one invocation revision. -Subclass `AtomicAction` directly, declare `TargetType`, compose a `TrajectoryBuilder`, -and implement `execute`. +## 3. Implement the planner -```python -from __future__ import annotations +Inherit `AtomicAction[PushGoal, PushOptions]` directly. Declare stable metadata and resolve +resources from semantic binding roles. -import torch +```python from typing import ClassVar -from embodichain.lab.sim.planners import PlanState, MoveType -from embodichain.lab.sim.atomic_actions.core import ( - ActionCfg, - ActionResult, +from embodichain.lab.sim.atomic_actions import ( + ResolvedActionRequest, + ActionPlan, AtomicAction, - WorldState, + PlanningContext, + StateDelta, ) -from embodichain.lab.sim.atomic_actions.trajectory import TrajectoryBuilder -from embodichain.utils import logger -class Push(AtomicAction): - """Push an object forward by a fixed distance from a contact pose.""" - - TargetType: ClassVar[type] = PushTarget # set to EndEffectorPoseTarget if you reused it - - def __init__(self, motion_generator, cfg: PushCfg | None = None): - super().__init__(motion_generator, cfg or PushCfg()) - self.builder = TrajectoryBuilder(motion_generator) - self.n_envs = self.robot.get_qpos().shape[0] - self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.arm_dof = len(self.arm_joint_ids) - self.robot_dof = self.robot.dof - - def execute(self, target: PushTarget, state: WorldState) -> ActionResult: - # 1. Resolve the batched contact pose (n_envs, 4, 4). - contact_xpos = self.builder.resolve_pose_target( - target.contact_pose, n_envs=self.n_envs - ) - - # 2. Resolve the arm start qpos from the threaded WorldState. - start_arm_qpos = self.builder.resolve_start_qpos( - state.last_qpos[:, self.arm_joint_ids], - n_envs=self.n_envs, - arm_dof=self.arm_dof, - control_part=self.cfg.control_part, - ) - - # 3. Plan the arm trajectory via the builder (uses IK + interpolation by default; - # set cfg.motion_source="motion_gen" to use the MotionGenerator instead). - target_states = [ - [PlanState(xpos=contact_xpos[i], move_type=MoveType.EEF_MOVE)] - for i in range(self.n_envs) - ] - success, arm_traj = self.builder.plan_arm_traj( +class Push(AtomicAction[PushGoal, PushOptions]): + skill_id: ClassVar[str] = "push" + GoalType: ClassVar[type] = PushGoal + OptionsType: ClassVar[type] = PushOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + + def __init__(self, default_options: PushOptions | None = None) -> None: + super().__init__(default_options) + + def plan( + self, + request: ResolvedActionRequest[PushGoal, PushOptions], + context: PlanningContext, + ) -> ActionPlan: + goal = self.require_goal(request) + options = request.skill_options + manipulator = request.binding.manipulator("primary") + control_part = manipulator.name + joint_ids = list(manipulator.joint_ids) + start_qpos = context.robot.qpos[:, joint_ids] + + # Build planner states and generate controlled-joint motion using + # request.motion_policy. Embed it into full robot DoF. + result = self.builder.generate_arm_plan( target_states, - start_arm_qpos, - self.cfg.sample_interval, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=self.cfg, + start_qpos, + request.motion_policy.sample_count, + control_part=control_part, + arm_dof=manipulator.dof, + cfg=request.motion_policy, ) - - # 4. Embed the arm slice into a full-DoF trajectory (n_envs, n_wp, robot.dof). - full = torch.empty( - (self.n_envs, arm_traj.shape[1], self.robot_dof), - dtype=torch.float32, - device=self.device, + success, trajectory = self.builder.to_full_robot_trajectory( + result, + base_qpos=context.robot.qpos, + joint_ids=joint_ids, + env_ids=context.env_ids, + control_dt=request.motion_policy.control_dt, ) - full[:, :, :] = state.last_qpos.unsqueeze(1) - full[:, :, self.arm_joint_ids] = arm_traj - - return ActionResult( + return self.build_plan( + request, + context, success=success, - trajectory=full, - next_state=WorldState( - last_qpos=full[:, -1, :].clone(), - held_object=state.held_object, # push does not grasp - ), - ) - - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=torch.zeros(self.n_envs, dtype=torch.bool, device=self.device), - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, + trajectory=trajectory, + expected_effects=StateDelta(), ) ``` -**Rules:** -- `execute()` returns an `ActionResult` — never a bare tuple. -- `trajectory` shape is always `(n_envs, n_waypoints, robot.dof)` (full robot DoF). -- Pass `cfg=self.cfg` to every `self.builder.plan_arm_traj(...)` call so the builder - reads `motion_source` and per-action planning options from the action config. -- Use `self.builder.` for all trajectory math (`resolve_pose_target`, - `resolve_joint_target`, `resolve_start_qpos`, `apply_local_offset`, `plan_arm_traj`, - `plan_joint_traj`, `split_three_phase`, `interpolate_hand_qpos`). Do not reimplement - that math inline. -- Thread `WorldState` explicitly: advance `last_qpos` to the final trajectory row; - set/clear/preserve `held_object` per what the action does to the grasp. -- Use `logger.log_error(msg, ValueError)` for contract violations (wrong target type, - missing cfg fields); use `logger.log_warning` + `_fail(state)` for soft planning - failures. -- Call `super().__init__()` — it sets `self.robot`, `self.motion_generator`, - `self.device`, `self.cfg`, `self.control_part`. - -### 4. Register the action - -Register an **instance** with the engine so `run()` can dispatch it by name. +Follow these invariants: -```python -from embodichain.lab.sim.atomic_actions import AtomicActionEngine, Push +- Let the engine supply `self.robot`, `self.motion_generator`, and the shared + `self.builder`; use `_on_bind()` only for robot/device-dependent setup. +- Call `require_goal()` before planning. +- Plan from `context.robot.qpos`, never an implicit live robot start state. +- Return full-robot `(B, N, robot.dof)` motion as a tensor or + `TimedTrajectory` with matching `env_ids`. +- Preserve backend timing/derivatives when available. +- Return `failed_plan(request, context, message=...)` for an expected soft + planning failure. +- Never mutate the context, step simulation, send commands, or claim a physical + effect occurred. +- Declare attachment/task changes with `StateDelta`; the execution runtime + applies them only after verification. +- Set `scene_dependencies` indirectly by using `SceneEntityPose` in the goal; + `build_plan()` records them for dynamic invalidation. -engine = AtomicActionEngine(motion_generator=motion_gen) -engine.register(Push(motion_gen, cfg=PushCfg())) # keyed by cfg.name "push" -``` +## 4. Register and invoke -For third-party / plugin actions that should be discoverable without the caller -constructing them, register the **class** in the global registry: +Register an instance by its class-level `skill_id`: ```python -from embodichain.lab.sim.atomic_actions import register_action -register_action("push", Push) +engine.register(Push()) ``` -### 5. Export from the public API - -Add the config, action class, and any new target to the package exports. For a -built-in primitive, first export it from -`embodichain/lab/sim/atomic_actions/primitives/__init__.py`: +Use the global registry only for discoverable third-party classes: ```python -from .push import Push, PushCfg - -__all__ = [ - ..., - "Push", - "PushCfg", -] +register_action(Push) ``` -Then export it from the public API in -`embodichain/lab/sim/atomic_actions/__init__.py`: +Construct a grounded invocation explicitly: ```python -from .primitives import Push, PushCfg -# (and from .core import PushTarget if you defined one) - -__all__ = [ - ..., - "Push", - "PushCfg", -] -``` - -Keep `embodichain/lab/sim/atomic_actions/actions.py` as a compatibility facade; -update it only if the new built-in should also be available from the legacy -`embodichain.lab.sim.atomic_actions.actions` import path. - -### 6. Update the supported actions table - -Add a row to the table in `docs/source/overview/sim/atomic_actions/builtin_actions.md`: - -```markdown -| `Push` | Single | `PushTarget` — contact pose | Approach → push forward | Add a demo asset or `N/A` | -``` - -### 7. Write a test - -Add a test in `tests/sim/atomic_actions/` (append to `test_actions.py` or create a new -file). Mock the `MotionGenerator` (see the `_make_mock_motion_generator` helper in -`test_actions.py`) and assert on behaviour: target type, full-DoF trajectory shape, -and the `WorldState` contract. - -```python -def test_push_action_cfg_defaults(): - cfg = PushCfg() - assert cfg.name == "push" - assert cfg.push_distance == 0.05 - -def test_push_action_returns_full_dof_trajectory(): - mg = _make_mock_motion_generator() - action = Push(mg, PushCfg(sample_interval=10)) - state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)) - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - return_value=torch.zeros(NUM_ENVS, 10, ARM_DOF), - ): - result = action.execute(PushTarget(contact_pose=torch.eye(4)), state) - assert isinstance(result, ActionResult) - assert result.success_all is True - assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) - # push preserves held_object - assert result.next_state.held_object is state.held_object +invocation = ActionInvocation( + skill_id="push", + goal=PushGoal(contact_pose), + binding=ActionBinding(manipulators={"primary": "left_arm"}), + motion_policy=MotionPolicy(sample_count=60), + recovery_policy=RecoveryPolicy(max_replans=2), +) +compiled = engine.compile((invocation,)) ``` -## Common Mistakes - -| Mistake | Fix | -|---------|-----| -| Inheriting from `MoveEndEffector` | Inherit `AtomicAction` directly and compose a `TrajectoryBuilder`. Actions are siblings, not a tree. | -| Returning `(bool, Tensor, joint_ids)` | Return an `ActionResult` with a full-DoF `(n_envs, n_wp, robot.dof)` trajectory. | -| Declaring `validate` / `updates_held_object_state` / `get_held_object_state` | These were removed. State flows only through `WorldState` and `ActionResult.next_state`. | -| `execute(target, start_qpos=None, **kwargs)` | Signature is `execute(self, target, state: WorldState) -> ActionResult`. No `**kwargs`, no `start_qpos`. | -| Reimplementing IK / interpolation inline | Use `self.builder.plan_arm_traj(...)`, `self.builder.plan_joint_traj(...)`, and friends. | -| Returning arm-only or arm+hand trajectory | Always embed into full `robot.dof` before returning. | -| Forgetting `cfg=self.cfg` in `plan_arm_traj` | The builder defaults to `motion_source="ik_interp"`; pass `cfg=self.cfg` to opt into `motion_gen` and per-action planning options. | -| Treating `ActionResult.success` as a scalar | It is `(n_envs,)` for batched actions; use `.success_all` or `.success.all()` for a single bool. | -| `name` not matching the engine registration key | Keep `cfg.name` identical to the key passed to `engine.register(...)` / `register_action(...)`. | -| Forgetting to export from `__init__.py` | Users import from the public API — missing exports cause `ImportError`. | -| Inheriting another action's cfg | Cfgs are flat; extend `ActionCfg` directly and declare the fields you need. | - -## Quick Reference - -| Step | Action | -|------|--------| -| 1 | Define a flat `@configclass` extending `ActionCfg` with a unique `name` | -| 2 | Define a typed target (or reuse `EndEffectorPoseTarget` / `JointPositionTarget` / `NamedJointPositionTarget` / `GraspTarget` / `HeldObjectPoseTarget`) | -| 3 | Subclass `AtomicAction` directly, set `TargetType`, compose `TrajectoryBuilder`, implement `execute(target, state) -> ActionResult` (pass `cfg=self.cfg` to `plan_arm_traj` and return per-env `success`) | -| 4 | Register: `engine.register(Push(mg, cfg=...))` (instance) or `register_action("push", Push)` (class) | -| 5 | Export config + action (+ target) from `primitives/__init__.py` and `atomic_actions/__init__.py` | -| 6 | Add a row to the supported-actions table in `builtin_actions.md` and update API reference docs | -| 7 | Write behavioural tests (target type, full-DoF shape, `WorldState` contract) | +Use `engine.start(...).tick(...)` instead when dynamic scene updates or online +error recovery are required. + +## 5. Export and document + +Export the goal, options, and action from: + +1. `embodichain/lab/sim/atomic_actions/primitives/__init__.py` +2. `embodichain/lab/sim/atomic_actions/__init__.py` + +Add the stable skill ID, goal, roles, and effect to +`docs/source/overview/sim/atomic_actions/builtin_actions.md`. Update API docs for +new public classes. Do not create a compatibility re-export module or a closed +built-in-goal union. + +## 6. Test behavior + +Add pure pytest tests under `tests/sim/atomic_actions/`. Cover: + +- descriptor `skill_id`, `GoalType`, and required roles; +- invalid goal and missing binding rejection; +- per-environment planning success/failure masks; +- full-robot trajectory shape, `env_ids`, timing, and failed-row hold behavior; +- side-effect-free context handling; +- masked `StateDelta` application for task effects; +- `SceneEntityPose` replanning when the action accepts a dynamic goal; +- effect verification when the action declares a non-empty delta. + +Run focused tests, format changed Python files with the pinned Black version, +then use the `pre-commit-check` skill before committing. + +## Common mistakes + +| Mistake | Required correction | +|---|---| +| Inherit another action | Inherit `AtomicAction` directly; compose helpers. | +| Add one generic target with many optional fields | Define a narrow action-owned goal. | +| Put hardware names in the goal | Bind semantic roles through `ActionBinding`. | +| Put arm/hand control-part names in skill options | Use `ActionBinding` as their only source. | +| Bind a joint, link, TCP frame, or arbitrary name | Every binding value must be a key in `RobotCfg.control_parts`. | +| Put hand qpos or named robot postures in skill options | Register semantic commands on the concrete control-part profile. | +| Put planner/recovery knobs in skill options | Move them to invocation policies. | +| Pass a motion generator to each action | Pass it once to `AtomicActionEngine`; construct actions from default options only. | +| Read `robot.get_qpos()` inside `plan()` | Use `context.robot.qpos`. | +| Return an arm-only tensor | Embed into full robot DoF. | +| Mutate held state after planning | Declare a `StateDelta`. | +| Treat `plan_success` as physical success | Verify effects during execution. | +| Step the simulator from the action | Emit plans; let the caller own execution. | diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 8389f34f7..c9219adeb 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -422,7 +422,7 @@ topics: - motion primitive - action primitive - AtomicAction - - ActionTarget + - ActionInvocation - AtomicActionEngine - TrajectoryBuilder - 原子动作 @@ -432,26 +432,46 @@ topics: - AtomicAction - AtomicActionEngine - TrajectoryBuilder - - ActionResult - - WorldState + - ActionPlan + - PlanningContext + - ExecutionSession + - StateDelta - held_objects - - ObjectActionTarget - - PlaceTarget - - PressTarget - - target ownership - - ActionCfg - - motion_source + - ActionBinding + - SceneEntityPose + - dynamic goal + - error recovery + - ActionOptions + - ResolvedActionRequest + - ControlPartCommandProfile + - ActionControlOverrides + - JointPositionCommand + - invocation revision + - MotionPolicy + - RecoveryPolicy - plan_arm_traj - register_action + - BUILTIN_ACTION_TYPES + - load_builtins + - engine.plan + - engine.compile + - engine.start paths: - 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/goals.py + - embodichain/lab/sim/atomic_actions/bindings.py + - embodichain/lab/sim/atomic_actions/control.py + - embodichain/lab/sim/atomic_actions/invocation.py + - embodichain/lab/sim/atomic_actions/policies.py + - embodichain/lab/sim/atomic_actions/runtime.py + - embodichain/lab/sim/atomic_actions/state.py + - embodichain/lab/sim/atomic_actions/plans.py + - embodichain/lab/sim/atomic_actions/execution.py - embodichain/lab/sim/atomic_actions/engine.py - embodichain/lab/sim/atomic_actions/trajectory.py - embodichain/lab/sim/atomic_actions/primitives/ - - embodichain/lab/sim/atomic_actions/actions.py - embodichain/lab/sim/atomic_actions/__init__.py related_topics: - motion-planning diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index 30cec7254..18e57fff9 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -1,180 +1,183 @@ -# Atomic Actions +# Atomic actions -## Entry Points +## Current contract -| What | Path | -|---|---| -| 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 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` | +Atomic actions are side-effect-free, environment-batched planners: -## Overview +```python +plan = engine.plan(invocation: ActionInvocation, context: PlanningContext) +``` + +There is no `ActionTarget`, `WorldState`, `ActionResult`, `execute()`, or +`AtomicActionEngine.run()` compatibility surface. + +`ActionInvocation` separates: + +- an action-owned typed goal (`goal_kind` is its stable discriminator); +- `ActionBinding`, which maps semantic roles to names from the engine robot's + `control_parts` mapping; +- reusable `MotionPolicy` planner/timing choices; +- bounded `RecoveryPolicy` thresholds and retry budgets; +- optional typed `skill_options` and role-scoped `control_overrides` for one + invocation revision. + +`PlanningContext` separates measured `RobotObservation`, verified symbolic +`TaskState`, versioned `SceneSnapshot`, and environment IDs. An `ActionPlan` +contains per-environment planning success, one or more `PlannedPhase` objects, +full-robot `TimedTrajectory` data, diagnostics, completion conditions, and an +uncommitted `StateDelta`. + +Each `AtomicActionEngine` exclusively owns one `ActionPlanningServices` +instance, which contains its robot, motion generator, planner backend, shared +`TrajectoryBuilder`, and control-part command profiles. Actions retain only an +owned copy of typed default options and borrow engine services. Engine +construction creates and binds a fresh instance of every type in +`BUILTIN_ACTION_TYPES`; use `load_builtins=False` only for isolated tests or a +fully custom action set. A bound action cannot be reused by another engine. + +## Engine entry points + +Choose the public engine entry point by lifecycle, not by skill type: + +| Entry point | Use | Result and state behavior | +|---|---|---| +| `engine.plan(invocation, context)` | Inspect or plan one registered action | Returns one `ActionPlan`; does not project a context for another action | +| `engine.compile(invocations, context)` | Plan an ordered sequence against a fixed scene | Returns a concatenated `CompiledTrajectory`; propagates hypothetical qpos and expected effects through `projected_context` | +| `engine.start(invocations, context)` | Execute incrementally from observations | Returns an `ExecutionSession`; `tick(latest_context)` emits commands and performs bounded recovery | + +None steps simulation directly. `compile()` never observes physical execution; +split compilation at observation boundaries when later goals depend on measured +results. Use `start()` when observation, effect verification, and replanning +must remain active during execution. -Atomic actions are env-batched motion primitives chained by `AtomicActionEngine`. Each action receives a typed target and a `WorldState`, plans a full-DoF trajectory for all environments, and returns an `ActionResult`. The engine threads `WorldState` from one action to the next and concatenates trajectories along the time axis. +`AtomicAction.plan(request, context)` is the skill implementation hook called +by the engine, not a fourth application entry point. `engine.plan_action(...)` +is only an extension/testing escape hatch for an unregistered instance. +## Static compilation + +Built-ins are already registered by their class-level stable `skill_id`; call: + +```python +compiled = engine.compile(invocations, context=None) ``` -AtomicActionEngine - ├─ AtomicAction(s) ← one primitive per class, e.g. MoveEndEffector, PickUp - │ │ - │ └── TrajectoryBuilder ← IK/interpolation and MotionGenerator dispatch - │ - └── WorldState ← last_qpos + per-control-part held-object maps + +Compilation does not step simulation. It concatenates timed trajectories and +applies successful expected effects only to `compiled.projected_context`, so a +following action can be checked against hypothetical state. Failed rows hold +their last successful qpos. + +Use invocation `skill_options` for multiple variants with the same stable +`skill_id`; do not create per-variant built-in instances. + +## Dynamic execution and recovery + +`SceneEntityPose(entity_id, relative_pose)` is resolved from the latest scene +snapshot every time the action plans. Its entity ID is recorded in +`PhaseSpec.scene_dependencies`. + +```python +session = engine.start(invocations, initial_context) +tick = session.tick(latest_context, effect_success=None) ``` -All tensor shapes carry a leading batch dim `B = n_envs`. +An `ExecutionSession` emits at most one `JointCommand` per tick and monitors: -## Core Types +- joint tracking error against the previous command; +- translation/rotation drift of referenced scene entities; +- phase timeout; +- planner and semantic-effect failure. -### Typed Targets +It replans from the latest observation within per-environment budgets. Unknown +or exhausted failures are reported as structured `ExecutionEvent` objects. A +non-empty `StateDelta` is not committed until the caller supplies an external +`effect_success` mask. -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. +Recovery replans reuse the current immutable `ResolvedActionRequest`. To change +a goal, option, policy, binding, or control command during execution, submit a +strictly newer revision explicitly: -`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. +```python +session.revise_current(revised_invocation) +``` -| Target | Holds | Used by | -|---|---|---| -| `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` | -| `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` | -| `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_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 - -Every `execute()` returns: -- `success: bool | torch.Tensor` — per-env boolean tensor of shape `(B,)` for batched actions. -- `trajectory: torch.Tensor` — full-robot trajectory `(B, n_waypoints, robot.dof)`. -- `next_state: WorldState` — state to feed into the next action. - -Helpers: -- `ActionResult.success_all` — `True` only when every env succeeded. -- `bool(action_result)` — deprecated; delegates to `success_all` and emits a `DeprecationWarning`. - -## Action Configuration - -`ActionCfg` (base for all action configs): - -| Field | Type | Default | Notes | -|---|---|---|---| -| `name` | `str` | `"default"` | Engine registration key | -| `control_part` | `str` | `"arm"` | Robot control part to move | -| `interpolation_type` | `str` | `"linear"` | Interpolation flavor | -| `velocity_limit` | `float \| None` | `None` | Used on the `motion_gen` path | -| `acceleration_limit` | `float \| None` | `None` | Used on the `motion_gen` path | -| `motion_source` | `str` | `"ik_interp"` | `"ik_interp"` (batched IK + interpolation) or `"motion_gen"` (batched `MotionGenerator`) | - -The base config is flat: every action cfg extends `ActionCfg` directly, even if it also carries hand open/close fields (see `PickUpCfg` / `PlaceCfg`). -The engine's `MotionGenerator` owns exactly one planner, so action configs do not -declare a planner type. On the `motion_gen` path, `TrajectoryBuilder` derives -planner-specific options from that owned planner. - -## TrajectoryBuilder - -Stateless helper owned by each action. Key methods: - -| Method | Purpose | -|---|---| -| `resolve_pose_target(target, n_envs)` | Broadcast EEF target to `(B,4,4)` or `(B,n,4,4)` | -| `resolve_joint_target(target, n_envs, joint_dof, control_part)` | Broadcast joint target to `(B,dof)` or `(B,n,dof)` | -| `resolve_start_qpos(start_qpos, n_envs, arm_dof, control_part)` | Broadcast start qpos to `(B, arm_dof)` | -| `plan_arm_traj(target_states_list, start_qpos, n_waypoints, control_part, arm_dof, cfg=None)` | Returns `(success:(B,), trajectory:(B,n_waypoints,arm_dof))`. Selects `ik_interp` or `motion_gen` from `cfg.motion_source`. | -| `plan_joint_traj(start_qpos, target_qpos, n_waypoints)` | Joint-space interpolation; always succeeds. | -| `split_three_phase(...)` | Split sample interval into motion / hand-interp / motion phases. | -| `interpolate_hand_qpos(...)` | Interpolate gripper qpos between two states. | - -`plan_arm_traj` input contract for actions: `target_states_list` is `list[list[PlanState]]` where the outer list is per-env and the inner list is per-waypoint. The builder internally converts to a batched `list[PlanState]` (each carrying `(B, ...)` tensors) when dispatching to `MotionGenerator`. - -## AtomicActionEngine +The replacement must keep the active `skill_id` and `invocation_id`. The +session resolves a new snapshot, resets that revision's recovery budgets, and +replans from the latest context. + +## Parameter ownership + +Goal dataclasses carry only semantic task intent. They do not carry robot part +names, planner configuration, retry policy, or runtime state. + +`MotionPolicy` owns planner selection, motion source, sample count, fallback +control period, limits, and typed planner options. `RecoveryPolicy` owns +tracking/dynamic-goal thresholds, timeouts, and budgets. Each built-in has a +frozen `*Options` value for invocation-varying phase counts, offsets, and grasp +selection behavior. An action constructor may accept `default_options`; an +invocation's `skill_options` replaces them for that call. There is no +`ActionCfg` or built-in `*Cfg` layer. + +`engine.register(action)` is reserved for custom skill implementations. A +built-in can be replaced only with explicit `replace=True`. Registration means +an implementation is installed; it does not prove that the current embodiment +has compatible control parts, profiles, bindings, or task state. Capability +adapters must filter registered descriptors before exposing skills to an Agent. +The module-level `register_action()` API is a process-wide extension-type +discovery catalog only; it neither binds actions nor changes an engine's +default built-in set. + +Every `ActionBinding` value is a `RobotCfg.control_parts` key. It is not a link, +TCP-frame, joint, or scene-object name. Planning services validate those names +and resolve immutable `ResolvedControlPart` values containing full-robot joint +indices. Built-ins use the binding as the only source for participating arm and +hand names; attachment state and `StateDelta` keys use the bound manipulator. + +Embodiment-specific joint commands do not belong to Action options. Register +them once by actual control-part name: ```python -engine = AtomicActionEngine(motion_generator) -engine.register(MoveEndEffector(motion_generator, cfg=MoveEndEffectorCfg())) -success, traj, final_state = engine.run(steps=[("move_end_effector", target)]) +engine = AtomicActionEngine( + motion_generator, + control_profiles={ + "left_hand": ControlPartCommandProfile.joint_positions( + open=left_open_qpos, + grasp=left_grasp_qpos, + ), + "left_arm": ControlPartCommandProfile.joint_positions(ready=ready_qpos), + }, +) ``` -`run(steps, state=None) -> (success, full_traj, final_state)`: -- `success` is a `(B,)` bool tensor indicating which environments completed every step. -- Failed environments hold their last successful joint position in both `full_traj` and `final_state.last_qpos` for the remainder of the sequence. -- If all envs fail, the loop stops early. -- `state` defaults to `WorldState(last_qpos=robot.get_qpos().clone())`. +Actions request semantic commands (`open`, `grasp`, or a named joint target) +from the `ResolvedControlPart`. `ActionControlOverrides` may replace commands +by semantic binding role for one invocation revision. Joint limits constrain +commands but do not define semantic open/grasp states; a robot integration or +tutorial may derive a simple profile from limits explicitly. -## Built-in Primitives +## Built-ins -| Action | Target | Notes | +| Skill ID | Goal type | Roles | |---|---|---| -| `MoveEndEffector` | `EndEffectorPoseTarget` | EEF pose move | -| `MoveJoints` | `JointPositionTarget` / `NamedJointPositionTarget` | Joint-space interpolation | -| `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. 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` 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`. - -## Common Failure Modes - -- **Forgetting `cfg=self.cfg` in `plan_arm_traj`** — without it, `motion_source` defaults to `"ik_interp"` and per-action planning options are ignored. -- **Treating `success` as scalar** — `ActionResult.success` is `(B,)` for all built-ins; use `success_all` or `success.all()` for a single bool. -- **Using `bool(action_result)` in new code** — still works but emits a `DeprecationWarning`; prefer `.success_all`. -- **Returning arm-only trajectory** — actions must embed into `(B, n_wp, robot.dof)` before returning. -- **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. +| `move_end_effector` | `EndEffectorPoseGoal` | manipulator `primary` | +| `move_joints` | `JointPositionGoal` (`target` is explicit qpos or a profile command name) | manipulator `primary` | +| `pick_up` | `GraspGoal` | manipulator/end effector `primary` | +| `move_held_object` | `HeldObjectPoseGoal` | manipulator/end effector `primary` | +| `place` | `PlaceGoal`, `AssembleGoal` | manipulator/end effector `primary` | +| `press` | `PressGoal` | manipulator/end effector `primary` | +| `coordinated_pickment` | `CoordinatedPickGoal` | `left`, `right` | +| `coordinated_placement` | `CoordinatedPlacementGoal` | `placing`, `support` | +| `hand_over` | `GraspGoal` | `source`, `destination` | + +## Extension rules + +1. Define a frozen action-owned goal dataclass with `goal_kind`. +2. Define a frozen `ActionOptions` subclass only when runtime behavior exists. +3. Declare `skill_id`, `GoalType`, `OptionsType`, and required semantic roles. +4. Validate with `require_goal(request)` and consume only the resolved binding. +5. Plan from `context.robot.qpos`; never read an implicit live start state. +6. Return full-robot positions or a `TimedTrajectory` through `build_plan()`. +7. Declare symbolic changes with `StateDelta`; do not mutate context or commit + physical effects during planning. +8. Keep scene stepping, controller I/O, and task-graph/MLLM logic outside the + atomic action. 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 68326d4df..32e46c1e5 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 @@ -1,5 +1,5 @@ embodichain.lab.sim.atomic_actions.primitives -============================================ +============================================= .. automodule:: embodichain.lab.sim.atomic_actions.primitives @@ -8,49 +8,47 @@ Overview Concrete implementations of the built-in atomic-action primitives. Each primitive is an :class:`~embodichain.lab.sim.atomic_actions.AtomicAction` that -accepts a typed target and a -:class:`~embodichain.lab.sim.atomic_actions.WorldState`, plans a full-DoF -trajectory for all parallel environments, and returns an -:class:`~embodichain.lab.sim.atomic_actions.ActionResult`. The primitives are -chained by :class:`~embodichain.lab.sim.atomic_actions.AtomicActionEngine`, -which threads ``WorldState`` from one action to the next and concatenates the -resulting trajectories along the time axis. +accepts a :class:`~embodichain.lab.sim.atomic_actions.ResolvedActionRequest` and a +:class:`~embodichain.lab.sim.atomic_actions.PlanningContext`. Planning returns a +side-effect-free :class:`~embodichain.lab.sim.atomic_actions.ActionPlan` with a +full-robot timed trajectory and uncommitted expected effects. .. rubric:: Built-in Primitive Actions .. autosummary:: - MoveEndEffectorCfg MoveEndEffector - MoveJointsCfg + MoveEndEffectorOptions MoveJoints - PickUpCfg + MoveJointsOptions PickUp - MoveHeldObjectCfg + PickUpOptions MoveHeldObject - PlaceCfg + MoveHeldObjectOptions Place - PressCfg + PlaceOptions Press - CoordinatedPickmentCfg + PressOptions CoordinatedPickment - CoordinatedPlacementCfg + CoordinatedPickmentOptions CoordinatedPlacement + CoordinatedPlacementOptions + HandOver + HandOverOptions - .. rubric:: Built-in Target Contracts + .. rubric:: Built-in Goal Contracts .. autosummary:: - EndEffectorPoseTarget - JointPositionTarget - NamedJointPositionTarget - GraspTarget - HeldObjectPoseTarget - PlaceTarget - PressTarget - CoordinatedPickTarget - CoordinatedPickmentTarget - CoordinatedPlacementTarget + EndEffectorPoseGoal + JointPositionGoal + GraspGoal + HeldObjectPoseGoal + PlaceGoal + AssembleGoal + PressGoal + CoordinatedPickGoal + CoordinatedPlacementGoal .. currentmodule:: embodichain.lab.sim.atomic_actions.primitives @@ -117,3 +115,11 @@ CoordinatedPlacement :members: :show-inheritance: :exclude-members: __init__, copy, replace, to_dict + +HandOver +-------- + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives.hand_over + :members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst index 9bc4fe46b..29d8b335b 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 @@ -3,257 +3,174 @@ embodichain.lab.sim.atomic_actions .. automodule:: embodichain.lab.sim.atomic_actions - .. rubric:: Classes + .. rubric:: Planning contracts + + .. autosummary:: + + ActionGoal + ActionBinding + ResolvedActionBinding + ResolvedControlPart + ControlCommand + JointPositionCommand + ControlPartCommandProfile + ActionControlOverrides + ActionInvocation + ResolvedActionRequest + ActionOptions + MotionPolicy + RecoveryPolicy + RobotObservation + TaskState + SceneSnapshot + SceneEntityPose + PlanningContext + StateDelta + TimedTrajectory + PhaseSpec + PlannedPhase + ActionPlan + CompiledTrajectory + + .. rubric:: Execution contracts .. autosummary:: - Affordance - AntipodalAffordance - InteractionPoints - ObjectSemantics - ActionTarget - ObjectActionTarget - EndEffectorPoseTarget - PlaceTarget - PressTarget - JointPositionTarget - NamedJointPositionTarget - GraspTarget - HeldObjectPoseTarget - CoordinatedPickTarget - CoordinatedPickmentTarget - CoordinatedPlacementTarget - Target - BuiltinTarget - HeldObjectState - CoordinatedHeldObjectState - WorldState - ActionResult - ActionCfg AtomicAction - TrajectoryBuilder - MoveEndEffectorCfg + AtomicActionEngine + ExecutionSession + ExecutionTick + JointCommand + ExecutionEvent + + .. rubric:: Built-in goals and actions + + .. autosummary:: + + EndEffectorPoseGoal + JointPositionGoal + GraspGoal + HeldObjectPoseGoal + PlaceGoal + AssembleGoal + PressGoal + CoordinatedPickGoal + CoordinatedPlacementGoal MoveEndEffector - MoveJointsCfg MoveJoints - PickUpCfg PickUp - MoveHeldObjectCfg MoveHeldObject - PlaceCfg Place - PressCfg Press - CoordinatedPickmentCfg CoordinatedPickment - CoordinatedPlacementCfg CoordinatedPlacement - AtomicActionEngine + HandOver .. toctree:: :maxdepth: 1 :hidden: - embodichain.lab.sim.atomic_actions.targets embodichain.lab.sim.atomic_actions.primitives .. currentmodule:: embodichain.lab.sim.atomic_actions -Layout ------- +Planning and state +------------------ -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. +.. autoclass:: ActionBinding + :members: -Core ----- +.. autoclass:: ResolvedActionBinding + :members: -.. autoclass:: Affordance - :members: - :show-inheritance: +.. autoclass:: ResolvedControlPart + :members: -.. autoclass:: AntipodalAffordance - :members: - :show-inheritance: +.. autoclass:: ControlCommand + :members: -.. autoclass:: InteractionPoints - :members: - :show-inheritance: +.. autoclass:: JointPositionCommand + :members: -.. autoclass:: ObjectSemantics - :members: - :show-inheritance: +.. autoclass:: ControlPartCommandProfile + :members: -.. autoclass:: ActionTarget - :members: - :show-inheritance: +.. autoclass:: ActionControlOverrides + :members: -.. autoclass:: ObjectActionTarget - :members: - :show-inheritance: +.. autoclass:: ActionInvocation + :members: -.. autoclass:: EndEffectorPoseTarget - :members: - :show-inheritance: +.. autoclass:: ResolvedActionRequest + :members: -.. autoclass:: PlaceTarget - :members: - :show-inheritance: +.. autoclass:: ActionOptions + :members: -.. autoclass:: PressTarget - :members: - :show-inheritance: +.. autoclass:: MotionPolicy + :members: + :exclude-members: __init__, copy, replace, to_dict -.. autoclass:: JointPositionTarget - :members: - :show-inheritance: +.. autoclass:: RecoveryPolicy + :members: + :exclude-members: __init__, copy, replace, to_dict -.. autoclass:: NamedJointPositionTarget - :members: - :show-inheritance: +.. autoclass:: PlanningContext + :members: -.. autoclass:: GraspTarget - :members: - :show-inheritance: +.. autoclass:: RobotObservation + :members: -.. autoclass:: HeldObjectPoseTarget - :members: - :show-inheritance: +.. autoclass:: TaskState + :members: -.. autoclass:: CoordinatedPickTarget - :members: - :show-inheritance: +.. autoclass:: SceneSnapshot + :members: -.. autodata:: CoordinatedPickmentTarget +.. autoclass:: SceneEntityPose + :members: -.. autoclass:: CoordinatedPlacementTarget - :members: - :show-inheritance: +.. autoclass:: StateDelta + :members: -.. autodata:: Target +.. autoclass:: TimedTrajectory + :members: -.. autodata:: BuiltinTarget +.. autoclass:: ActionPlan + :members: -.. autoclass:: HeldObjectState - :members: - :show-inheritance: +Engine and execution +-------------------- -.. autoclass:: CoordinatedHeldObjectState - :members: - :show-inheritance: +.. autoclass:: AtomicAction + :members: -.. autoclass:: WorldState - :members: - :show-inheritance: +.. autoclass:: AtomicActionEngine + :members: -.. autoclass:: ActionResult - :members: - :show-inheritance: +.. autoclass:: ExecutionSession + :members: -.. autoclass:: ActionCfg - :members: - :exclude-members: __init__, copy, replace, to_dict +.. autoclass:: ExecutionTick + :members: -.. autoclass:: AtomicAction - :members: - :show-inheritance: +.. autoclass:: JointCommand + :members: -Trajectory helpers ------------------- +.. autoclass:: ExecutionEvent + :members: -.. autoclass:: TrajectoryBuilder - :members: - :show-inheritance: - -Actions -------- - -.. autoclass:: MoveEndEffectorCfg - :members: - :exclude-members: __init__, copy, replace, to_dict - :show-inheritance: - -.. autoclass:: MoveEndEffector - :members: - :show-inheritance: - -.. autoclass:: MoveJointsCfg - :members: - :exclude-members: __init__, copy, replace, to_dict - :show-inheritance: - -.. autoclass:: MoveJoints - :members: - :show-inheritance: - -.. autoclass:: PickUpCfg - :members: - :exclude-members: __init__, copy, replace, to_dict - :show-inheritance: - -.. autoclass:: PickUp - :members: - :show-inheritance: - -.. autoclass:: MoveHeldObjectCfg - :members: - :exclude-members: __init__, copy, replace, to_dict - :show-inheritance: - -.. autoclass:: MoveHeldObject - :members: - :show-inheritance: - -.. autoclass:: PlaceCfg - :members: - :exclude-members: __init__, copy, replace, to_dict - :show-inheritance: - -.. autoclass:: Place - :members: - :show-inheritance: - -.. autoclass:: PressCfg - :members: - :exclude-members: __init__, copy, replace, to_dict - :show-inheritance: - -.. autoclass:: Press - :members: - :show-inheritance: - -.. autoclass:: CoordinatedPickmentCfg - :members: - :exclude-members: __init__, copy, replace, to_dict - :show-inheritance: - -.. autoclass:: CoordinatedPickment - :members: - :show-inheritance: - -.. autoclass:: CoordinatedPlacementCfg - :members: - :exclude-members: __init__, copy, replace, to_dict - :show-inheritance: - -.. autoclass:: CoordinatedPlacement - :members: - :show-inheritance: - -Engine & Registry ------------------ +Semantic objects and helpers +---------------------------- -.. autoclass:: AtomicActionEngine - :members: - :show-inheritance: +.. autoclass:: ObjectSemantics + :members: -.. autofunction:: register_action +.. autoclass:: HeldObjectState + :members: -.. autofunction:: unregister_action +.. autoclass:: CoordinatedHeldObjectState + :members: -.. autofunction:: get_registered_actions +.. autoclass:: TrajectoryBuilder + :members: 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 deleted file mode 100644 index 2f60da71c..000000000 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.targets.rst +++ /dev/null @@ -1,21 +0,0 @@ -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 bc163f630..87889a968 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -1,314 +1,560 @@ -(builtin_actions)= +(builtin-actions)= -# Built-in Actions +# Built-in atomic actions ```{currentmodule} embodichain.lab.sim.atomic_actions ``` -The following actions are available out of the box: +EmbodiChain ships nine built-in action implementations with stable skill IDs; +`AtomicActionEngine` creates and registers a fresh instance of every built-in by +default. Applications select them by stable skill ID rather than registering +routine instances themselves. +`Place` additionally accepts an `AssembleGoal`, so assembly reuses the same +release primitive instead of introducing a tenth skill ID. + +All built-ins implement +`plan(request, context) -> ActionPlan`, where `request` is the engine-resolved +snapshot of an invocation revision. Constructors accept only optional typed +default `*Options`; the owning `AtomicActionEngine` supplies the shared motion +generator, trajectory builder, and control-part command profiles when it binds +the built-in catalog. Generic motion and recovery choices belong to the +invocation, and per-call primitive behavior belongs to `skill_options`. + +Registration only installs an implementation. Whether a built-in is executable +for a particular call still depends on its binding roles, the robot's control +parts, semantic command profiles, and task-state preconditions. Action Agent +adapters must also honor `agent_visible` and filter by embodiment capability. ```{note} -The built-in atomic actions currently support gripper-based manipulation only. Dexterous-hand manipulation is not supported yet. +The current manipulation primitives consume semantic `open` and `grasp` +commands through the control-part command abstraction. The shipped command +implementation is `JointPositionCommand`. A dexterous hand may register +calibrated joint-position commands immediately; non-position hand policies or +multi-stage in-hand manipulation require additional command types and phases. ``` -| Action | Arm | Target type | Motion phases | Demo | -|---|---|---|---|---| -| `MoveEndEffector` | Single | `EndEffectorPoseTarget` — EEF pose | Move end-effector to pose | MoveEndEffector | -| `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 | `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 | -| `HandOver` | Dual | `GraspTarget` — object semantics | Move to handover pose → receive grasp → close receiving hand → release transferring hand → deliver and retreat | HandOver | +## Visual catalog ---- +The animations below are the focused simulator demos under +`scripts/tutorials/atomic_action/`. -## `MoveEndEffector` +::::{grid} 1 2 2 2 +:gutter: 2 -Moves the end-effector to a target pose in free space. +:::{grid-item-card} `MoveEndEffector` +:link: builtin-move-end-effector +:link-type: ref -| Config field | Default | Description | -|---|---|---| -| `control_part` | `"arm"` | Robot control part to move | -| `sample_interval` | `50` | Number of waypoints in the trajectory | -| `plan_opts` | `None` | Optional planner-specific options; copied before each motion-generator call | +`move_end_effector` · free-space EEF pose motion -**Target:** `EndEffectorPoseTarget(xpos=...)` where `xpos` is a `torch.Tensor` of shape `(4, 4)`, `(n_envs, 4, 4)` or `(n_envs, n_waypoint, 4, 4)` — a homogeneous EEF pose. +MoveEndEffector demo +::: -![MoveEndEffector demo](../../../_static/atomic_actions/move_end_effector.gif) +:::{grid-item-card} `MoveJoints` +:link: builtin-move-joints +:link-type: ref ---- +`move_joints` · explicit or named joint-space motion -## `MoveJoints` +MoveJoints demo +::: + +:::{grid-item-card} `PickUp` +:link: builtin-pick-up +:link-type: ref + +`pick_up` · approach, close, and lift + +PickUp demo +::: + +:::{grid-item-card} `MoveHeldObject` +:link: builtin-move-held-object +:link-type: ref + +`move_held_object` · object-centric transport + +MoveHeldObject demo +::: + +:::{grid-item-card} `Place` +:link: builtin-place +:link-type: ref + +`place` · approach, release, and retract + +Place demo +::: + +:::{grid-item-card} Assembly through `Place` +:link: builtin-assemble +:link-type: ref + +`place` + `AssembleGoal` · base-relative placement + +Assembly demo +::: + +:::{grid-item-card} `Press` +:link: builtin-press +:link-type: ref + +`press` · close, contact, and return + +Press demo +::: + +:::{grid-item-card} `CoordinatedPickment` +:link: builtin-coordinated-pickment +:link-type: ref -Moves a configured control part directly in joint space. Use this for known safe poses, -home poses, recovery motions, or any motion where a qpos target is clearer than an EEF pose. +`coordinated_pickment` · dual-arm shared-object pick -| Config field | Default | Description | +CoordinatedPickment demo +::: + +:::{grid-item-card} `CoordinatedPlacement` +:link: builtin-coordinated-placement +:link-type: ref + +`coordinated_placement` · align two held objects + +CoordinatedPlacement demo +::: + +:::{grid-item-card} `HandOver` +:link: builtin-hand-over +:link-type: ref + +`hand_over` · transfer an attachment between arms + +HandOver demo +::: + +:::: + +## Capability matrix + +| Skill ID | Accepted goal | Required binding roles | Required profile commands | Required task state | Expected task effect | +|---|---|---|---|---|---| +| `move_end_effector` | `EndEffectorPoseGoal` | manipulator `primary` | none | none | none | +| `move_joints` | `JointPositionGoal` | manipulator `primary` | named target only: command matching `target` | none | none | +| `pick_up` | `GraspGoal` | manipulator + end effector `primary` | primary: `open`, `grasp` | semantic object/entity | attach object to `primary` manipulator | +| `move_held_object` | `HeldObjectPoseGoal` | manipulator + end effector `primary` | primary: `grasp` | object held by `primary` | preserve attachment | +| `place` | `PlaceGoal`, `AssembleGoal` | manipulator + end effector `primary` | primary: `open`, `grasp` | `AssembleGoal` requires an object held by `primary`; ordinary `PlaceGoal` has no planner-enforced attachment precondition | detach object | +| `press` | `PressGoal` | manipulator + end effector `primary` | primary: `grasp` | none | none | +| `coordinated_pickment` | `CoordinatedPickGoal` | manipulator + end effector `left`, `right` | both: `open`, `grasp` | semantic object/entity | create coordinated attachment; clear individual attachments | +| `coordinated_placement` | `CoordinatedPlacementGoal` | manipulator + end effector `placing`, `support` | placing: `open`, `grasp`; support: `grasp` | one individually held object per arm | optionally detach placing object; preserve support attachment | +| `hand_over` | `GraspGoal` | manipulator + end effector `source`, `destination` | both: `open`, `grasp` | object held by source arm | transfer attachment to destination arm | + +### Binding role meanings + +Roles are action-local semantic participant slots. They are keys declared by an +action, while the corresponding `ActionBinding` values are concrete +`Robot.control_parts` keys. A role that appears in both binding maps identifies +the manipulator and actuated hand/tool serving the same functional participant; +it does not make the two maps interchangeable. + +| Role | Used by | Meaning | |---|---|---| -| `control_part` | `"arm"` | Robot control part to move | -| `sample_interval` | `50` | Number of waypoints in the interpolated trajectory | -| `named_joint_positions` | `None` | Optional `dict[str, torch.Tensor]` for named qpos targets | +| `primary` | Single-participant skills | Principal participant for this invocation; it has no inherent left/right or default-robot meaning | +| `source` | `hand_over` | Participant that initially holds and transfers the object | +| `destination` | `hand_over` | Participant that receives the object | +| `left`, `right` | `coordinated_pickment` | Participants associated with the goal's left/right grasp transforms | +| `placing` | `coordinated_placement` | Participant that aligns and optionally releases the placing object | +| `support` | `coordinated_placement` | Participant that keeps holding and positioning the support object | + +The action's `manipulator_roles` and `end_effector_roles` declarations determine +which entries are required. The engine checks that those entries exist and that +every value resolves through `Robot.control_parts`; the caller or capability +binder must select a physically compatible arm and hand/tool combination. + +`MoveJoints` is intentionally `agent_visible=False`: it is useful for home, +recovery, calibration, and scripted postures, but is not exposed to an Action +Agent by default. + +## Shared goal and configuration rules + +### Pose values and dynamic references + +Pose-valued fields use `PoseGoalValue`, which accepts either an explicit tensor +or a late-bound scene reference: + +```python +fixed = EndEffectorPoseGoal(xpos=target_pose) +tracked = EndEffectorPoseGoal( + xpos=SceneEntityPose("tray", relative_pose=tray_to_tcp) +) +``` -**Targets:** -- `JointPositionTarget(qpos=...)` where `qpos` is a `torch.Tensor` of shape `(control_dof,)`, `(n_envs, control_dof)` or `(n_envs, n_waypoint, control_dof)`. -- `NamedJointPositionTarget(name=...)` where `name` is resolved from - `MoveJointsCfg.named_joint_positions`. +Explicit pose tensors use `(4, 4)` or `(B, 4, 4)`. Waypoint-capable fields in +`EndEffectorPoseGoal` and `PlaceGoal` also accept `(B, N, 4, 4)`. +`SceneEntityPose` resolves to the latest `(B, 4, 4)` pose from each +`SceneSnapshot`, checks optional perception confidence, and registers that +entity as a recovery dependency. + +| Skill / field | `SceneEntityPose` accepted | Automatic scene-motion replan | +|---|---:|---:| +| `MoveEndEffector.xpos` | yes | yes | +| `MoveHeldObject.object_target_pose` | yes | yes | +| `Place.xpos` | yes | yes | +| `Press.xpos` | yes | yes | +| `CoordinatedPickGoal.object_target_pose` / `object_initial_pose` | yes | yes | +| `CoordinatedPlacementGoal` placing/support poses | yes | yes | +| `PickUp` / `HandOver` semantic entity lookup | not through `SceneEntityPose` | no automatic scene dependency | +| `AssembleGoal` base entity lookup | not through `SceneEntityPose` | latest pose is used when replanning, but base movement alone does not trigger it | + +### Parameter ownership + +Use this rule when configuring a built-in or adding a new one: + +- the **goal** carries only the requested outcome; +- the **binding** carries semantic-role mappings to control-part names selected + for this call; every value must be a key in the engine robot's + `control_parts` mapping; +- typed **skill options** carry phase-specific behavior that may vary by + invocation; an action may provide defaults; +- the engine's **control-part profiles** carry embodiment-specific semantic + commands such as `open`, `grasp`, and named postures; +- `MotionPolicy` carries sample count, timing, motion source, limits, + collision choice, and planner options; +- `RecoveryPolicy` carries all replan/retry thresholds and budgets. + +All built-ins resolve participating arm and hand names exclusively from +`ActionBinding`. The engine then resolves the selected control part's profile +and checks each joint-position command against its DoF. Invocation-level +`ActionControlOverrides` may replace a command by binding role for one explicit +revision. + +### Planning and effect semantics + +Every action returns a per-environment `plan_success` mask and one or more +full-robot trajectories. `plan_success=True` means motion planning succeeded; +it does not prove contact or object transfer. Actions that change attachment +state declare a `StateDelta`. Offline `compile()` projects it hypothetically; +closed-loop execution commits it only after external effect verification. + +(builtin-move-end-effector)= + +## `MoveEndEffector` -![MoveJoints demo](../../../_static/atomic_actions/move_joints.gif) +Plans a free-space motion for a bound manipulator to reach one EEF pose or an +ordered set of pose waypoints. ---- +| Contract | Value | +|---|---| +| Skill ID | `move_end_effector` | +| Goal | `EndEffectorPoseGoal(xpos=...)` | +| Binding | manipulator role `primary` | +| Motion | EEF planning from observed arm qpos; output expanded to full robot DoF | +| Completion | `EEF_GOAL_REACHED` | +| Effect | none | +| Skill options | none; reusable motion choices are in `MotionPolicy` | + +Use an explicit pose for a fixed target, `SceneEntityPose` for a tracked target, +or `(B, N, 4, 4)` for intermediate waypoints. The action does not command an end +effector/hand resource. + +**Example:** `scripts/tutorials/atomic_action/move_end_effector.py` + +(builtin-move-joints)= + +## `MoveJoints` + +Plans directly in joint space. This is appropriate for known safe postures, +homing, scripted recovery, or motions whose desired outcome is a qpos rather +than an EEF pose. + +| Contract | Value | +|---|---| +| Skill ID | `move_joints` | +| Goal | `JointPositionGoal(target=...)` | +| Binding | manipulator role `primary` | +| Motion | joint planning/interpolation from observed qpos; supports joint waypoints | +| Completion | `JOINT_GOAL_REACHED` | +| Effect | none | +| Agent visibility | hidden by default (`agent_visible=False`) | + +`target` accepts an explicit qpos tensor with shape `(control_dof,)`, +`(B, control_dof)`, or `(B, N, control_dof)`, or a non-empty string resolved +from the bound manipulator's `ControlPartCommandProfile`. Named poses remain +embodiment knowledge without becoming separate goal types: + +```python +engine = AtomicActionEngine( + motion_generator, + control_profiles={ + "left_arm": ControlPartCommandProfile.joint_positions(home=home_qpos), + }, +) + +explicit_goal = JointPositionGoal(target=home_qpos) +named_goal = JointPositionGoal(target="home") +``` + +**Example:** `scripts/tutorials/atomic_action/move_joints.py` + +(builtin-pick-up)= ## `PickUp` -Three-phase grasp motion: *approach → close gripper → lift*. +Plans **approach -> close hand -> lift** and declares the object attached to the +bound manipulator. -| Config field | Default | Description | -|---|---|---| -| `approach_direction` | `[0, 0, -1]` | Gripper approach direction in object frame | -| `pre_grasp_distance` | `0.15` | Hover distance before descending (m) | -| `lift_height` | `0.10` | Lift height after grasping (m) | -| `hand_open_qpos` | `None` | **Required.** Gripper open joint positions | -| `hand_close_qpos` | `None` | **Required.** Gripper closed joint positions | -| `hand_control_part` | `"hand"` | Robot control part for the gripper | -| `hand_interp_steps` | `5` | Waypoints for the gripper close phase | -| `sample_interval` | `80` | Total waypoints across all three phases | - -**Target:** `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_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) - ---- +| Contract | Value | +|---|---| +| Skill ID | `pick_up` | +| Goal | `GraspGoal(semantics=..., grasp_xpos=None)` | +| Binding | manipulator + end effector role `primary` | +| Precondition | `ObjectSemantics.entity` is set; an `AntipodalAffordance` is required when no explicit grasp pose is supplied | +| Effect | write `HeldObjectState` for the bound manipulator and clear overlapping coordinated attachment state | +| Verification | the attachment effect must be verified during closed-loop execution | -## `MoveHeldObject` +`grasp_xpos` may be `(4, 4)` or `(B, 4, 4)`. When omitted, the action samples +valid affordance grasps, evaluates reachability, and stores the selected +`object_to_eef` transform in the expected held-object state. Later +object-centric skills reuse that transform. -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_objects[control_part]`) -and preserves it in its successor state. +`PickUp` requires `open` and `grasp` commands on the bound end-effector profile. +Important `PickUpOptions` fields: -`HeldObjectState` and `HeldObjectPoseTarget` are intentionally kept separate from -`ObjectSemantics`: `ObjectSemantics` describes the object and affordances, while these -types describe runtime held-object state and an action-specific target pose. +| Field | Purpose | +|---|---| +| `pre_grasp_distance`, `approach_direction` | Pre-grasp offset and approach direction | +| `lift_height`, `hand_interp_steps` | Lift distance and close-phase discretization | +| `pick_object_part` | Affordance region: currently `center`, `top`, or `bottom` | +| `approach_alignment_max_angle` | Optional TCP approach-alignment filter | +| `downstream_object_target_poses` | Optional future reachability constraints used in grasp selection | +| `obj_upright_direction`, `rotate_upright` | Optional orientation-selection behavior | -| Config field | Default | Description | -|---|---|---| -| `hand_close_qpos` | `None` | **Required.** Gripper closed joint positions | -| `hand_control_part` | `"hand"` | Robot control part for the gripper | -| `sample_interval` | `50` | Number of waypoints in the trajectory | +The semantic entity pose is read when planning occurs, but it is not currently a +`SceneEntityPose` dependency. Object motion alone therefore does not trigger +automatic dynamic-goal replanning. -**Target:** `HeldObjectPoseTarget(object_target_pose=...)` where `object_target_pose` is a -`torch.Tensor` of shape `(4, 4)` or `(n_envs, 4, 4)` — the desired pose of the held object. -The action converts this to an EEF target via the stored object-to-EEF transform. +**Example:** `scripts/tutorials/atomic_action/pickup.py` -![MoveHeldObject demo](../../../_static/atomic_actions/move_held_object.gif) +(builtin-move-held-object)= + +## `MoveHeldObject` ---- +Moves an already attached object to an object-frame target while keeping the +hand closed. The caller specifies the desired **object pose**, not an EEF pose; +the action derives `target_object_pose @ object_to_eef` from verified task state. + +| Contract | Value | +|---|---| +| Skill ID | `move_held_object` | +| Goal | `HeldObjectPoseGoal(object_target_pose=...)` | +| Binding | manipulator + end effector role `primary` | +| Precondition | a `HeldObjectState` exists for the bound manipulator, normally from `PickUp` | +| Motion | single object-centric transport phase with closed-hand qpos | +| Effect | none; the existing attachment is preserved | +| Dynamic target | explicit pose or `SceneEntityPose` | + +The bound end-effector profile must provide `grasp`; optional upright-transport +settings belong to `MoveHeldObjectOptions`. The arm and hand are selected by +`ActionBinding`; generic timing and trajectory sampling remain in +`MotionPolicy`. + +**Example:** `scripts/tutorials/atomic_action/move_held_object.py` + +(builtin-place)= ## `Place` -Three-phase release motion: *lower → open gripper → retract*. Mirrors `PickUp`. +Plans **approach/descend -> open hand -> retract**. A multi-waypoint +`PlaceGoal` visits all supplied release waypoints in order and opens at the last +one. -`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` removes the -entry for `PlaceCfg.control_part` from `held_objects`. +| Contract | Value | +|---|---| +| Skill ID | `place` | +| Goal | `PlaceGoal(xpos=..., tcp_symmetry="none")` | +| Binding | manipulator + end effector role `primary` | +| State | consumes the bound manipulator's attachment when present | +| Effect | detach the object and clear overlapping coordinated attachment state | +| Verification | release must be verified during closed-loop execution | +| Dynamic target | explicit pose/waypoints or `SceneEntityPose` | -| Config field | Default | Description | -|---|---|---| -| `lift_height` | `0.10` | Retract height after opening the gripper (m) | -| `hand_open_qpos` | `None` | **Required.** Gripper open joint positions | -| `hand_close_qpos` | `None` | **Required.** Gripper closed joint positions | -| `hand_control_part` | `"hand"` | Robot control part for the gripper | -| `hand_interp_steps` | `5` | Waypoints for the gripper open phase | -| `sample_interval` | `80` | Total waypoints across all three phases | - -**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 -`tcp_symmetry="z_roll_180"` only when releasing with TCP x/y flipped is physically -equivalent; `Place` then chooses the closer TCP z-roll 180 variant from -`WorldState.last_qpos` and applies that same variant across all release waypoints. - -![Place demo](../../../_static/atomic_actions/place.gif) - -### Object assembly - -`Place` also accepts an `AssembleTarget` in place of a `PlaceTarget` to place a -held object onto a base object at a declared relative pose. There is no separate -assembly action — the same `Place` primitive consumes an `AssembleAffordance` -and derives the release pose from the base object's live pose. - -An `AssembleAffordance` anchors the assemble object (the part that is picked up -and placed) to a base object (the assembly anchor). The base object's world pose -is read at planning time from `base_object_entity`, so the target tracks a moved -base. The assemble-object target pose is `base_pose @ assemble_to_base_pose`, -which `Place` converts to an EEF release pose through the held object's -`object_to_eef` — the transform a prior `PickUp` writes into -`WorldState.held_objects[control_part]`. - -| `AssembleAffordance` field | Default | Description | -|---|---|---| -| `base_object_label` | `""` | Label of the base object the assemble object is placed onto | -| `base_object_entity` | `None` | **Required.** Simulation entity for the base object; its pose anchors the assembly | -| `assemble_object_label` | `""` | Label of the assemble object that is picked up and placed | -| `assemble_object_entity` | `None` | Optional simulation entity for the assemble object (reference/logging) | -| `assemble_to_base_pose` | `torch.eye(4)` | Pose of the assemble object relative to the base object frame, shape `(4, 4)` or `(n_envs, 4, 4)` | +Set `tcp_symmetry="z_roll_180"` only if TCP x/y can be flipped while TCP z and +translation remain physically equivalent. The action selects the closer +orientation variant from the observed starting state and uses it consistently +across all waypoints. -**Target:** `AssembleTarget(affordance=...)` wrapping an `AssembleAffordance`. -The release EEF pose is `base_pose @ assemble_to_base_pose @ object_to_eef`, -reusing the held-object transform populated by the prior `PickUp`. +The bound end-effector profile must provide `open` and `grasp`. Important +`PlaceOptions` fields: -**Tutorial:** `scripts/tutorials/atomic_action/assemble.py` +| Field | Purpose | +|---|---| +| `lift_height` | Approach and retract height | +| `hand_interp_steps` | Open-phase discretization | +| `max_approach_retract_z` | Optional world-Z ceiling for approach/retract poses | +| `cartesian_waypoint_count` | Fixed-orientation translation keyframes per segment | -![Assemble demo](../../../_static/atomic_actions/assemble.gif) +**Example:** `scripts/tutorials/atomic_action/place.py` ---- +(builtin-assemble)= -## `Press` +### Assembly through `Place` -Three-phase contact motion: *close gripper → press down → return*. This is useful -for button-like or contact-based interactions where the end-effector should reach a -target pose and then return to the pre-press arm pose. +`Place` also accepts `AssembleGoal(affordance=...)`. There is no separate +assembly skill: it derives the assemble-object target from the base object's +live pose and reuses the normal place/release phases. -`Press` does not create or clear `WorldState.held_objects`; it preserves the state -threaded into it. +```text +base_object_pose @ assemble_to_base_pose = assemble_object_target_pose +assemble_object_target_pose @ held.object_to_eef = release_eef_pose +``` -| Config field | Default | Description | -|---|---|---| -| `hand_close_qpos` | `None` | **Required.** Gripper closed joint positions | -| `hand_control_part` | `"hand"` | Robot control part for the gripper | -| `hand_interp_steps` | `5` | Waypoints for the gripper close phase | -| `sample_interval` | `80` | Total waypoints across all three phases | +The `AssembleAffordance` identifies the base and assemble objects, stores the +relative pose, and must provide `base_object_entity`. A prior verified `PickUp` +must have populated the held object's `object_to_eef` transform. Planning then +declares the same detach effect as a normal place. -**Target:** `PressTarget(xpos=...)` — the EEF pose to press, a `torch.Tensor` -of shape `(4, 4)` or `(n_envs, 4, 4)`. +The base entity's current pose is read each time `plan()` runs. Because the +goal does not yet encode that entity through `SceneEntityPose`, base movement by +itself does not invalidate an executing plan; another recovery trigger is +required before the newer pose is resolved. -![Press demo](../../../_static/atomic_actions/press.gif) +**Example:** `scripts/tutorials/atomic_action/assemble.py` ---- +(builtin-press)= + +## `Press` + +Plans **close hand -> move to contact pose -> return to the observed starting +arm qpos**. It is intended for button-like or contact interactions where the +arm should retreat along its planned path after reaching the target. + +| Contract | Value | +|---|---| +| Skill ID | `press` | +| Goal | `PressGoal(xpos=...)` | +| Binding | manipulator + end effector role `primary` | +| Motion | close, press, joint-space return | +| Effect | none; existing attachment state is unchanged | +| Dynamic target | explicit pose or `SceneEntityPose` | + +The bound end-effector profile must provide `grasp`, while +`PressOptions.hand_interp_steps` controls the close interpolation. The arm and +hand control parts come from `ActionBinding`. Contact detection is not itself a +symbolic effect in the current action; applications that require force/contact +confirmation should verify it externally. + +**Example:** `scripts/tutorials/atomic_action/press.py` + +(builtin-coordinated-pickment)= ## `CoordinatedPickment` -Dual-arm grasp motion for one shared object. Both arms move to object-relative -grasp poses, close both grippers, lift the object, and move it to an object pose -while keeping both grippers closed. On success, the returned `WorldState` carries -an entry in `coordinated_held_objects[(left_arm, right_arm)]` -(`CoordinatedHeldObjectState`) and removes individual held entries for those arms. +Coordinates two arms around one shared object: **approach both grasps -> close +both hands -> lift -> move object -> hold**. -| Config field | Default | Description | -|---|---|---| -| `control_part` | `"dual_arm"` | Combined arm control part | -| `left_arm_control_part` / `right_arm_control_part` | `"left_arm"` / `"right_arm"` | Arm control parts for each grasp | -| `left_hand_control_part` / `right_hand_control_part` | `"left_hand"` / `"right_hand"` | Hand control parts for each gripper | -| `pre_grasp_distance` | `0.10` | Distance to back away from each grasp TCP | -| `lift_height` | `0.08` | World-Z lift distance before moving to the target pose | -| `object_motion_keyframes` | `6` | Sparse object-pose IK keyframes for synchronized motion | -| `sample_interval` | `120` | Total waypoints across all phases | +| Contract | Value | +|---|---| +| Skill ID | `coordinated_pickment` | +| Goal | `CoordinatedPickGoal` | +| Binding | manipulator + end effector roles `left` and `right` | +| Goal geometry | shared-object target pose plus left/right `object_to_eef` transforms; optional initial object pose | +| Effect | clear individual left/right attachments and create `CoordinatedHeldObjectState[(left, right)]` | +| Verification | coordinated attachment must be externally verified | + +The object target and optional initial pose may use `SceneEntityPose`. When no +initial pose is supplied, `ObjectSemantics.entity` provides the object's current +pose. -**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. +Both bound end-effector profiles must provide `open` and `grasp`. Important +`CoordinatedPickmentOptions` fields group into: -`CoordinatedPickmentTarget` remains a compatibility alias. +- `pre_grasp_distance` and `lift_height`; +- `object_motion_keyframes`, `hand_interp_steps`, and `hold_steps`. -**Tutorial:** `scripts/tutorials/atomic_action/coordinated_pickment.py` +The left/right arms and hands come exclusively from the corresponding binding +roles. Coordinated dual-arm planning with `motion_source="motion_gen"` is not +supported by the cuRobo backend; use the supported IK/interpolation path for +this primitive. -![CoordinatedPickment demo](../../../_static/atomic_actions/coordinated_pickment.gif) +**Example:** `scripts/tutorials/atomic_action/coordinated_pickment.py` ---- +(builtin-coordinated-placement)= ## `CoordinatedPlacement` -Dual-arm object-centric placement. The support arm moves its held object to a lower -target pose and keeps its gripper closed. The placing arm moves its held object to -the aligned upper target pose, optionally opens the placing hand, then lifts away. +Moves a support object and a placing object together: **align both objects -> +hold -> optionally release the placing hand -> retreat the placing arm**. -`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. +| Contract | Value | +|---|---| +| Skill ID | `coordinated_placement` | +| Goal | `CoordinatedPlacementGoal` | +| Binding | manipulator + end effector roles `placing` and `support` | +| Precondition | separate `HeldObjectState` entries exist for both bound arms | +| Goal geometry | placing/support object target poses, optional height offsets, optional release override | +| Effect | preserve support attachment; remove or preserve placing attachment according to `release`; clear overlapping coordinated state | -| Config field | Default | Description | -|---|---|---| -| `control_part` | `"dual_arm"` | Robot control part containing both arms | -| `placing_arm_control_part` | `"left_arm"` | Arm that releases the placed object | -| `support_arm_control_part` | `"right_arm"` | Arm that keeps holding the support object | -| `placing_hand_control_part` | `"left_hand"` | Placing gripper control part | -| `support_hand_control_part` | `"right_hand"` | Support gripper control part | -| `placing_hand_open_qpos` | `None` | **Required.** Placing gripper open joint positions | -| `placing_hand_close_qpos` | `None` | **Required.** Placing gripper closed joint positions | -| `support_hand_close_qpos` | `None` | **Required.** Support gripper closed joint positions | -| `release` | `True` | Whether to open the placing gripper | -| `placing_height_offset` | `0.0` | World-Z offset applied to the placing object target pose | -| `support_height_offset` | `0.0` | World-Z offset applied to the support object target pose | -| `lift_height` | `0.08` | Placing-arm lift distance after release (m) | -| `hand_interp_steps` | `10` | Waypoints for placing-hand release | -| `hold_steps` | `4` | Alignment hold waypoints before release | -| `retreat_steps` | `16` | Placing-arm retreat waypoints | -| `sample_interval` | `100` | Total waypoints across all phases | - -**Target:** `CoordinatedPlacementTarget(...)` with placing/support object target -poses 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` - -![CoordinatedPlacement demo](../../../_static/atomic_actions/coordinated_placement.gif) - ---- +Both object targets may use `SceneEntityPose`, so either can participate in +dynamic-goal invalidation. Goal-level height/release values override +`CoordinatedPlacementOptions` for that invocation. -## `HandOver` +The placing profile must provide `open` and `grasp`; the support profile must +provide `grasp`. Important `CoordinatedPlacementOptions` fields group into: -Dual-arm object handover. The transferring arm (already holding the object) -moves it to a middle handover pose, the receiving arm approaches and grasps a -different part of the object, the transferring arm releases and retreats, and -the receiving arm carries the object to a final pose. +- default `release`, placing/support height offsets, and `lift_height`; +- `hand_interp_steps`, `hold_steps`, and `retreat_steps`. -`HandOver` requires a prior `PickUp`: it reads the `HeldObjectState` for the -transferring arm from `WorldState.held_objects[transfer_arm_control_part]` to -recover the object-to-EEF transform. On success, it removes that entry and -writes a new `HeldObjectState` under `receive_arm_control_part`, so the -receiving arm now holds the object. +The placing/support arms and hands come exclusively from the corresponding +binding roles. The same cuRobo restriction as coordinated pickment applies to dual-arm +`motion_source="motion_gen"` planning. -| Config field | Default | Description | -|---|---|---| -| `control_part` | `"dual_arm"` | Combined control part containing both the transferring and receiving arms | -| `transfer_arm_control_part` | `"left_arm"` | Arm that already holds the object and hands it over | -| `receive_arm_control_part` | `"right_arm"` | Arm that grasps the object and carries it away | -| `transfer_hand_control_part` | `"left_hand"` | Hand attached to the transferring arm | -| `receive_hand_control_part` | `"right_hand"` | Hand attached to the receiving arm | -| `transfer_hand_open_qpos` | `None` | **Required.** Transferring-hand qpos for the open (released) state, shape `[hand_dof,]` | -| `transfer_hand_close_qpos` | `None` | **Required.** Transferring-hand qpos for the closed (holding) state, shape `[hand_dof,]` | -| `receive_hand_open_qpos` | `None` | **Required.** Receiving-hand qpos for the open state, shape `[hand_dof,]` | -| `receive_hand_close_qpos` | `None` | **Required.** Receiving-hand qpos for the closed state, shape `[hand_dof,]` | -| `receive_pick_object_part` | `"bottom"` | Object part the receiving arm grasps during the handover | -| `middle_object_pose` | `None` | **Required.** Object pose at the handover point, shape `(4, 4)` or `(n_envs, 4, 4)` | -| `final_object_pose` | `None` | **Required.** Object pose the receiving arm delivers to, shape `(4, 4)` or `(n_envs, 4, 4)` | -| `receive_approach_direction` | `[0, 0, -1]` | World-frame approach direction used to sample the receiving grasp | -| `pre_grasp_distance` | `0.10` | Distance to offset back from the receiving grasp pose (m) | -| `lift_height` | `0.08` | World-Z lift distance for the transferring arm after release (m) | -| `sample_interval` | `120` | Total waypoints for the full handover trajectory | -| `hand_interp_steps` | `10` | Waypoints for the receiving-hand close and transferring-hand release | -| `hold_steps` | `4` | Waypoints to hold the handoff pose before releasing | -| `retreat_steps` | `24` | Waypoints for the final deliver/retreat phase | - -**Target:** `GraspTarget(semantics=...)` with an `ObjectSemantics` whose -`affordance` is an `AntipodalAffordance`. The receiving grasp is solved from -the affordance and `receive_pick_object_part` at the middle handover pose; the -transferring arm reuses the object-to-EEF transform stored by the prior -`PickUp`. - -**Tutorial:** `scripts/tutorials/atomic_action/hand_over.py` - -![HandOver demo](../../../_static/atomic_actions/hand_over.gif) +**Example:** `scripts/tutorials/atomic_action/coordinated_placement.py` + +(builtin-hand-over)= + +## `HandOver` + +Transfers an already held object from one arm to another: **move source to the +handover pose -> destination approaches and grasps -> source releases and +retreats -> destination delivers**. + +| Contract | Value | +|---|---| +| Skill ID | `hand_over` | +| Goal | `GraspGoal(semantics=...)` | +| Binding | manipulator + end effector roles `source` and `destination` | +| Precondition | source arm has a verified `HeldObjectState`; semantic object supports destination grasp selection | +| Effect | remove source attachment and create destination `HeldObjectState` | +| Verification | attachment transfer must be externally verified | + +Both source and destination end-effector profiles must provide `open` and +`grasp`. `HandOverOptions` owns the destination grasp region and approach +direction, middle/final object poses, and phase distances/counts. The +source/destination arm and hand control parts come exclusively from the +corresponding `ActionBinding` roles. + +The middle and final poses are currently option tensors rather than +`SceneEntityPose` goal fields. Consequently, handover supports tracking-error +and timeout recovery, but does not automatically invalidate a moving handover +point. An application can submit a newer invocation revision with updated +`HandOverOptions`; the action also queries the semantic object's live +orientation when replanning and preserves it at the supplied middle/final +positions. + +As with the other coordinated primitive, cuRobo does not currently support its +dual-arm `motion_source="motion_gen"` path. + +**Example:** `scripts/tutorials/atomic_action/hand_over.py` + +## Running the demos + +Every focused script is interactive by default. Add `--auto_play` to skip +keyboard prompts and combine it with `--headless --device cpu` for a headless +run that records video under `outputs/videos`: + +```bash +python scripts/tutorials/atomic_action/move_end_effector.py --headless --auto_play --device cpu +python scripts/tutorials/atomic_action/pickup.py --headless --auto_play --device cpu +python scripts/tutorials/atomic_action/hand_over.py --headless --auto_play --device cpu +``` + +See {doc}`/tutorial/atomic_actions` for engine setup, static compilation, +closed-loop execution, effect verification, and custom-action guidance. diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index b80ef1bb1..e56bf960d 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -1,369 +1,588 @@ -# Atomic Actions +(atomic-actions)= + +# Atomic actions + +```{toctree} +:hidden: + +builtin_actions +``` ```{currentmodule} embodichain.lab.sim.atomic_actions ``` -Atomic actions are the building blocks for automated robot motion generation. Each action encapsulates a complete, self-contained motion primitive — such as picking up an object or moving to a pose — that can be chained together to form complex manipulation workflows. +Atomic actions are the typed planning and execution boundary between a semantic +task request and robot joint commands. A caller describes **what** should happen +with an action-owned goal, grounds semantic roles onto robot resources, and +supplies the latest measured context. The action returns a full-robot, +time-aware plan without stepping simulation or claiming that a physical effect +has occurred. ```{note} -Atomic actions currently support gripper-based manipulation only. Dexterous-hand manipulation is not supported yet. +The current built-ins focus on arm-and-gripper manipulation. They already emit +full-robot-DoF trajectories, but dexterous-hand policies, lower-body locomotion, +and whole-body control are not implemented by this module yet. ``` -## Design Overview +## Architecture and responsibility boundary + +```text ++--------------------------------+ +--------------------------------+ +| Action Agent / semantic graph | | User-authored application | +| skill call + object references | | typed goal + binding + policy | ++---------------+----------------+ +---------------+----------------+ + | | + v | + agent adapter: schema validation, | + scene grounding, capability binding | + | | + +------------------+------------------+ + | + | ActionInvocation + | + PlanningContext + v ++-------------------------------------------------------------+ +| AtomicActionEngine | +| | +| owns exactly one ActionPlanningServices | +| +-- Robot | +| +-- MotionGenerator / planner backend | +| +-- device and shared TrajectoryBuilder | +| | +| resolves requests and calls AtomicAction.plan(...) | ++------------------------------+------------------------------+ + | + +-------------+-------------+ + | | | + v v v + engine.plan() engine.compile() engine.start()/tick() + one ActionPlan fixed projection observed closed loop + | | + v v + CompiledTrajectory JointCommand + events +``` -The module is organized into three layers: +The boundary is deliberate: -``` -AtomicActionEngine ← orchestrates a sequence of (name, typed_target) steps - │ - ├── AtomicAction(s) ← each action plans one motion primitive - │ │ - │ └── MotionGenerator ← low-level trajectory planner (IK + trajectory optimization) - │ - └── WorldState ← threaded action-to-action - (last_qpos + per-control-part held-object maps) -``` +| Concern | Owner | Contract | +|---|---|---| +| Task intent and sequencing | Action Agent, task graph, or user-authored application | Selects skills, goals, and execution order | +| Invocation construction | Agent adapter or user-authored code/config loader | Produces the same typed `ActionInvocation`; the engine has no agent-only interface | +| Perception and grounding | Agent adapter or user application | Builds scene snapshots and resource bindings, or supplies already-grounded values directly | +| Deterministic motion planning | Atomic action module | Produces an `ActionPlan` from an invocation and context | +| Motion-generation resources | `AtomicActionEngine` | Owns one robot, motion generator, planner backend, device, trajectory builder, and control-part command profiles | +| Robot/simulator stepping | Application control loop | Consumes `JointCommand`; the session never steps the simulator itself | +| Physical-effect verification | Application observer | Verifies grasp, release, handover, and other symbolic effects | + +### Caller entry points -Each action receives a typed target and a `WorldState`, runs its planning pipeline, and -returns an `ActionResult` whose trajectory covers the full robot DOF. The engine threads -the `next_state` of each action as the input state of the next, then concatenates all -trajectories into one contiguous sequence: +The engine supports two first-class caller paths. An Action Agent emits a +semantic skill call that an adapter validates, grounds, and converts into an +`ActionInvocation`. A user can instead author the typed invocation directly in +Python or load it from an application-owned configuration layer: +```python +manual_invocation = ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal(xpos=target_pose), + binding=ActionBinding(manipulators={"primary": "left_arm"}), + motion_policy=MotionPolicy(sample_count=80, control_dt=1.0 / 60.0), + recovery_policy=RecoveryPolicy(max_replans=2), +) + +# Choose one entry point according to the planning/execution requirement. +single_plan = engine.plan(manual_invocation, latest_context) +static_program = engine.compile((manual_invocation,), latest_context) +live_session = engine.start((manual_invocation,), latest_context) ``` -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 -CoordinatedPickTarget(...) │ -CoordinatedPlacementTarget(...) │ - │ - ActionResult - (success, full-DoF traj, next_state) - │ -AtomicActionEngine ◄───────────────────┘ -(run(steps, state) → (is_success, traj, final_state)) + +A manual caller may bypass the semantic-schema adapter only when its target and +robot-resource binding are already grounded. Scene-relative goals still need a +current `PlanningContext`, and object names or semantic roles still need to be +resolved by the user application (or by reusing the same grounding adapter as +the Agent path). + +Both paths converge at `ActionInvocation + PlanningContext`. They therefore use +the same goal validation, capability checks, planning backend, execution +events, bounded recovery, and physical-effect verification. Manual authoring is +an alternative orchestration entry point, not a lower-level path around the +engine contracts. + +## Choosing an engine entry point + +Application code normally chooses between these three public entry points: + +| API | Choose it when | Returns | State and observation behavior | +|---|---|---|---| +| `engine.plan(invocation, context)` | You need to inspect or plan exactly one registered action | `ActionPlan` | Reads one context; does not project its terminal qpos or expected task effect for another action | +| `engine.compile(invocations, context)` | All goals for an ordered static sequence are known before execution | `CompiledTrajectory` | Plans in order and propagates hypothetical qpos and expected effects through `projected_context`; never observes execution | +| `engine.start(invocations, context)` | Commands must be issued incrementally from fresh observations with bounded recovery | `ExecutionSession` | `tick(latest_context)` consumes measured state, emits at most one command, requests effect verification, and can replan | + +The short selection rule is: + +```text +one action to inspect or plan -> plan +one or more actions in a fixed scene -> compile +observed execution and error recovery -> start, then tick ``` -### Core Concepts +All three leave simulator stepping and controller I/O to the application. +`plan()` and `compile()` only return planning data. An `ExecutionSession` also +does not step the simulator itself; its `tick()` method returns commands for the +application to send. -**`ObjectSemantics`** describes an interaction target. It bundles: -- `affordance` — *how* to interact with the object (e.g. an `AntipodalAffordance` carrying mesh data and grasp-generation config) -- `geometry` — plain geometric metadata (e.g. a bounding box). Mesh tensors live on the affordance, not here -- `label` — object category string (also bound onto the affordance for convenience) -- `entity` — a live reference to the simulation object, so actions can read its current pose +Calling `compile()` with one invocation is valid and gives a uniform +`CompiledTrajectory` result, but it is not required for a single action. More +importantly, `compile()` cannot observe physical execution. If a later goal +depends on the measured result of an earlier action, end the compiled phase, +observe a new `PlanningContext`, and plan or compile the next phase. Use +`start()` when that observe/replan loop should be managed continuously by an +`ExecutionSession`. -**`HeldObjectState`** is runtime state produced after a successful `PickUp`. It stores -the held object's semantics and object-to-end-effector transform so later actions can move the -object without recomputing the grasp. It is intentionally separate from `ObjectSemantics`, -which remains a reusable object description rather than per-execution robot state. +## Core contracts -**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: +The public contracts separate values with different owners and lifetimes. This +keeps goals small and prevents robot-specific or planner-specific parameters +from leaking into an Action Agent schema. -| Target | Constructor | Used by | -|---|---|---| -| `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` | -| `CoordinatedPickTarget` | `CoordinatedPickTarget(...)` | `CoordinatedPickment` | -| `CoordinatedPlacementTarget` | `CoordinatedPlacementTarget(...)` | `CoordinatedPlacement` | - -`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: - -| Class | Use case | -|---|---| -| `AntipodalAffordance` | Parallel-jaw grasping via antipodal point pairs | -| `InteractionPoints` | Contact-based interactions (push, poke, touch) | - -`AntipodalAffordance` takes its inputs as direct fields — `mesh_vertices`, `mesh_triangles`, -`gripper_collision_cfg`, `generator_cfg`, and `force_reannotate` — rather than a nested config dict. - -**`AtomicAction`** is the abstract base class for all motion primitives. Subclasses declare a -`TargetType` class variable and implement a single method: -- `execute(target, state) -> ActionResult` — plan and return a full-DOF trajectory plus the - successor `WorldState` - -**`AtomicActionEngine`** holds a name-keyed registry of action instances and runs a sequence of -`(name, typed_target)` steps via `run(steps, state)`, threading `WorldState` from one action into -the next. - ---- - -## Typed Targets & State Threading - -The engine takes a sequence of `(name, typed_target)` steps. Each target is a small -frozen dataclass, and the engine checks that each step's target matches the registered -action's `TargetType` before calling `execute`: - -| Target | Holds | Accepted by | +| Contract | Contains | Does not contain | |---|---|---| -| `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` | -| `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` | - -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 maps | -|---|---| -| `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 | - -If a step fails, `run()` returns `success=False` with the partial trajectory concatenated up -to (but not including) the failed step, and the `WorldState` going into that step. - ---- - -## Typical Workflow +| `ActionGoal` | Action-specific desired outcome, such as an EEF pose or object pose | Arm names, planner instances, recovery counters | +| `ActionBinding` | Semantic-role mappings to keys from the engine robot's `control_parts`, such as `primary -> left_arm` and `primary -> left_hand` | Link/TCP names, arbitrary scene objects, motion settings, or task geometry | +| `ActionOptions` / built-in `*Options` | Frozen invocation-varying skill behavior: phase counts, offsets, grasp-selection rules | Robot resource names, hand qpos, planner backend | +| `ControlPartCommandProfile` | Embodiment-specific semantic commands such as `open`, `grasp`, and `ready`, keyed by actual control-part name | Action roles, task goals, recovery state | +| `ActionControlOverrides` | Optional role-scoped command replacements for one invocation revision | Persistent robot configuration | +| `MotionPolicy` | Motion source, sample count, timing, limits, collision option, typed planner options | Skill semantics or robot-resource names | +| `RecoveryPolicy` | Replan/retry budgets, tracking and dynamic-goal thresholds, phase timeout | Controller state or mutable counters | +| `PlanningContext` | Measured `RobotObservation`, verified `TaskState`, versioned `SceneSnapshot`, stable environment IDs | Hypothetical simulator mutation | +| `ActionPlan` | Per-environment planning result, scene-bound phases, timed trajectories, diagnostics, expected `StateDelta` | Proof that a grasp/release/contact physically succeeded | + +Goals follow the structural `ActionGoal` protocol: each action owns one or more +frozen dataclasses with a stable `goal_kind`. There is no shared `ActionTarget` +base class and no closed union that must change whenever a skill is added. + +### Semantic resource binding + +A **role** is an action-owned semantic participant slot: it describes the job a +robot resource performs in that action, not the identity of the resource. Each +`AtomicAction` declares its required slots through `manipulator_roles` and +`end_effector_roles`; the same declarations are exposed through its +`SkillDescriptor` so an Agent adapter or manual caller can construct a complete +binding before planning. + +Role names are local to both the skill and the resource category. For example, +`primary` in `manipulators` and `primary` in `end_effectors` are two separate +slots. Using the same role name expresses that the selected arm and hand/tool +serve the same functional participant in the action: ```python -from embodichain.lab.sim.atomic_actions import ( - AtomicActionEngine, - ObjectSemantics, - AntipodalAffordance, - GraspTarget, - EndEffectorPoseTarget, - PlaceTarget, - JointPositionTarget, - NamedJointPositionTarget, - HeldObjectPoseTarget, - PickUp, - PickUpCfg, - MoveJoints, - MoveJointsCfg, - MoveHeldObject, - MoveHeldObjectCfg, - Place, - PlaceCfg, - Press, - PressCfg, - MoveEndEffector, - MoveEndEffectorCfg, +binding = ActionBinding( + manipulators={"primary": "left_arm"}, + end_effectors={"primary": "left_hand"}, ) +``` -# 1. Configure each action -pickup_cfg = PickUpCfg( - control_part="arm", - hand_control_part="hand", - hand_open_qpos=torch.tensor([0.0, 0.0]), - hand_close_qpos=torch.tensor([0.025, 0.025]), -) -place_cfg = PlaceCfg( - control_part="arm", - hand_control_part="hand", - hand_open_qpos=torch.tensor([0.0, 0.0]), - hand_close_qpos=torch.tensor([0.025, 0.025]), -) -move_held_object_cfg = MoveHeldObjectCfg( - control_part="arm", - hand_control_part="hand", - hand_close_qpos=torch.tensor([0.025, 0.025]), -) -move_cfg = MoveEndEffectorCfg(control_part="arm") -press_cfg = PressCfg( - control_part="arm", - hand_control_part="hand", - hand_close_qpos=torch.tensor([0.025, 0.025]), -) -move_joints_cfg = MoveJointsCfg( - control_part="arm", - named_joint_positions={"home": torch.zeros(6)}, +In this example, `primary` is the role and `left_arm` / `left_hand` are the +bound resources. `primary` does not mean left, right, the first configured +arm, or a globally preferred arm; it simply denotes the principal participant +of a single-participant skill. Changing the values can bind the same action to +another compatible arm and tool without changing its goal or implementation. + +Every bound value is the name of a control part declared by the engine-owned +robot. Both `left_arm` and `left_hand` must therefore be keys in +`robot.control_parts` (originating from `RobotCfg.control_parts`). They are not +joint names, link names, TCP frame names, or scene-object identifiers. +`end_effectors` specifically selects the actuated tool/hand control part; the +manipulator's IK/TCP frame remains part of the robot and solver configuration. +The engine validates every name and resolves its full-robot joint indices +before calling the action planner. + +The validation boundary is intentionally narrow: the engine verifies required +roles, `control_parts` membership, resolvable joint indices, command type, and +command dimensions. The Agent adapter or application binder remains responsible +for capability compatibility, such as pairing an arm with the hand mounted on +it and choosing a semantic command supported by that tool. + +Role names should describe action responsibilities rather than robot-specific +joint, link, or model names. Single-resource skills use `primary`; handover uses +`source` and `destination`; coordinated placement uses `placing` and `support`. +The current coordinated-pick contract uses `left` and `right` because its goal +geometry also distinguishes left/right grasps. New skills should prefer +functional roles unless a spatial distinction is intrinsic to their semantics. + +All built-ins resolve participating arm and hand control parts from the binding. +They obtain hardware-specific `open` and `grasp` commands from the resolved +end-effector profile; no action or option duplicates arm names, hand names, or +hand qpos. Attachment state and expected effects are keyed by the bound +manipulator control-part name. + +### Control-part semantic commands + +Register embodiment commands once when constructing the engine. The keys are +concrete names from `robot.control_parts`; the command names remain semantic: + +```python +engine = AtomicActionEngine( + motion_generator, + control_profiles={ + "left_hand": ControlPartCommandProfile.joint_positions( + open=left_open_qpos, + grasp=left_grasp_qpos, + ), + "left_arm": ControlPartCommandProfile.joint_positions( + ready=left_ready_qpos, + ), + }, ) +``` + +`MoveJoints(JointPositionGoal("ready"))` resolves `ready` from its bound +manipulator. Manipulation primitives resolve `open` and/or `grasp` from their +bound end effectors. A one-dimensional `JointPositionCommand` broadcasts over +the planning batch; a two-dimensional value must match the selected batch. -# 2. Build the engine and register each action instance by name -engine = AtomicActionEngine(motion_generator=motion_gen) -engine.register(PickUp(motion_gen, cfg=pickup_cfg)) -engine.register(MoveHeldObject(motion_gen, cfg=move_held_object_cfg)) -engine.register(Place(motion_gen, cfg=place_cfg)) -engine.register(MoveEndEffector(motion_gen, cfg=move_cfg)) -engine.register(Press(motion_gen, cfg=press_cfg)) -engine.register(MoveJoints(motion_gen, cfg=move_joints_cfg)) - -# 3. Describe the object to pick -semantics = ObjectSemantics( - affordance=AntipodalAffordance( - mesh_vertices=obj.get_vertices(env_ids=[0], scale=True)[0], - mesh_triangles=obj.get_triangles(env_ids=[0])[0], - gripper_collision_cfg=gripper_cfg, - generator_cfg=generator_cfg, +For a one-off change, override by action role rather than by concrete robot +name: + +```python +invocation = ActionInvocation( + skill_id="pick_up", + goal=goal, + binding=binding, + control_overrides=ActionControlOverrides( + end_effectors={ + "primary": { + "grasp": JointPositionCommand(object_specific_grasp_qpos), + } + } ), - geometry={}, - label="mug", - entity=obj, + revision=1, ) +``` + +The engine merges the override after resolving `primary` and captures the +result in `ResolvedActionRequest`. Automatic recovery for revision 1 sees the +same command snapshot. Joint limits remain constraints; they do not define the +semantic meaning of `open` or `grasp`. Tutorials may explicitly derive a simple +profile from limits, while a robot integration should normally provide +calibrated commands. + +### Engine-owned planning resources + +One engine owns one motion generator. At initialization it creates a fresh +instance of every action type in `BUILTIN_ACTION_TYPES` and binds those +instances to the engine's planning services: + +```python +engine = AtomicActionEngine(motion_generator, control_profiles=profiles) + +# All nine built-ins are immediately usable by stable skill ID. +assert "move_end_effector" in engine.actions +assert "pick_up" in engine.actions +``` + +Consequences of this ownership model: -# 4. Plan the full sequence — steps are (name, typed_target) pairs -is_success, traj, final_state = engine.run( - steps=[ - ("pick_up", GraspTarget(semantics=semantics)), - ("move_held_object", HeldObjectPoseTarget(object_target_pose=carry_pose)), - ("place", PlaceTarget(xpos=place_pose)), - ("move_joints", NamedJointPositionTarget(name="home")), - ] +- built-in action constructors use their typed default options unless an + invocation supplies `skill_options`; +- every action in an engine sees the same robot, device, backend, caches, and + collision world; +- an action instance cannot be silently reused by a different engine; +- one registered instance exists per stable `skill_id` in an engine. + +Registration means that an implementation is installed, not that every robot +can execute it. Required roles, control parts, profiles, and task-state +preconditions are validated while an invocation is resolved and planned. Agent +adapters must additionally filter the catalog by `agent_visible` and +embodiment capability instead of exposing every `engine.actions` entry blindly. + +Use invocation `skill_options` whenever behavior varies per call. Two variants +with the same stable skill ID therefore share one built-in implementation: + +```python +left_invocation = ActionInvocation( + skill_id="pick_up", + goal=left_goal, + binding=left_binding, + skill_options=left_pick_options, +) +right_invocation = ActionInvocation( + skill_id="pick_up", + goal=right_goal, + binding=right_binding, + skill_options=right_pick_options, ) -# traj: (n_envs, n_waypoints, robot.dof) + +left_plan = engine.plan(left_invocation, latest_context) +right_plan = engine.plan(right_invocation, latest_context) +``` + +`register()` remains the extension point for a custom skill. Replacing a +built-in implementation requires an explicit `replace=True`; isolated tests or +fully custom engines can opt out of the catalog with `load_builtins=False`: + +```python +custom_engine = AtomicActionEngine(motion_generator, load_builtins=False) +custom_engine.register(MyAction()) + +engine.register(CustomPickUp(), replace=True) ``` ---- +The module-level `register_action()` catalog is only for process-wide extension +type discovery. It does not mutate existing engines or join their default +built-in set; instantiate a discovered extension and pass it to +`engine.register()` explicitly. -## How to Extend: Adding a Custom Action +### Implementation and advanced APIs -You can add any motion primitive by subclassing `AtomicAction`, composing a -`TrajectoryBuilder` for the shared planning math, and registering an instance with the engine. -Built-in primitives live one action per module under -`embodichain/lab/sim/atomic_actions/primitives/`, while -`embodichain.lab.sim.atomic_actions` remains the public import surface and -`embodichain.lab.sim.atomic_actions.actions` stays as a compatibility re-export. +The similarly named `AtomicAction.plan()` method is not a fourth application +entry point. It is the polymorphic method implemented by each skill and called +by the engine after resolving an invocation: -### Step 1 — Define the config +| API | Intended caller | Behavior | +|---|---|---| +| `AtomicAction.plan(request, context)` | Atomic-action implementer | Consumes an immutable `ResolvedActionRequest` and returns an `ActionPlan` | +| `engine.plan_action(action, invocation, context)` | Extension or isolated test | Temporarily binds and plans an unregistered action instance; built-in parameter variants should use invocation `skill_options` instead | +| `session.revise_current(invocation)` | Runtime orchestrator or Action Agent | Replaces the active logical call with a newer revision and replans from the latest observed context | + +Application code should start with `engine.plan()`, `engine.compile()`, or +`engine.start()` unless it specifically needs one of these extension points. + +## Planning one action + +Use `engine.plan()` when one registered action needs to be inspected, tested, +or integrated into application-owned orchestration: ```python -from embodichain.utils import configclass -from embodichain.lab.sim.atomic_actions import ActionCfg - -@configclass -class PushCfg(ActionCfg): - name: str = "push" - push_distance: float = 0.05 # metres to push forward - push_speed: int = 30 # waypoints for the push phase +invocation = ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal(xpos=target_pose), + binding=ActionBinding(manipulators={"primary": "left_arm"}), + motion_policy=MotionPolicy(sample_count=80, control_dt=1.0 / 60.0), +) + +plan = engine.plan(invocation, latest_context) +if plan.plan_success.all(): + positions = plan.trajectory.positions ``` -### Step 2 — Implement the action +The result contains that action's trajectory, diagnostics, completion +conditions, and uncommitted expected effects. `plan()` does not automatically +create a next context. If another action must be planned against this action's +hypothetical result, use `compile()` instead of manually reproducing its state +projection rules. + +## Static compilation + +`compile()` plans invocations in order. For every successful action it projects +the terminal qpos and expected task-state effect into a new context so the next +action can be checked against a hypothetical result. The observed context and +simulator remain unchanged. ```python -from dataclasses import dataclass -import torch -from typing import ClassVar from embodichain.lab.sim.atomic_actions import ( - ActionTarget, AtomicAction, ActionResult, WorldState, TrajectoryBuilder, + ActionBinding, + ActionInvocation, + AtomicActionEngine, + EndEffectorPoseGoal, + MotionPolicy, ) -@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()) - self.builder = TrajectoryBuilder(motion_generator) - self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.robot_dof = self.robot.dof - self.n_envs = self.robot.get_qpos().shape[0] - - 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=state.with_updates( - last_qpos=full[:, -1, :].clone(), - ), - ) +engine = AtomicActionEngine(motion_generator) +binding = ActionBinding(manipulators={"primary": "left_arm"}) +motion_policy = MotionPolicy(sample_count=80, control_dt=1.0 / 60.0) + +approach = ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal(xpos=approach_pose), + binding=binding, + motion_policy=motion_policy, +) +retreat = ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal(xpos=retreat_pose), + binding=binding, + motion_policy=motion_policy, +) + +initial_context = engine.initial_context() +compiled = engine.compile((approach, retreat), initial_context) +if compiled.plan_success.all(): + positions = compiled.trajectory.positions # (B, N, robot_dof) + approach_plan, retreat_plan = compiled.action_plans + final_context = compiled.projected_context ``` -### Step 3 — Register and use +When no context is supplied, the engine captures robot qpos/qvel and creates an +empty task state and scene snapshot. Supply an explicit context whenever goals +depend on perceived entities or a previous verified attachment. -Register an instance with the engine so it can be referenced by name in `run()`: +Do not compile across a boundary where execution feedback changes a later goal. +For example, `scripts/tutorials/atomic_action/coordinated_placement.py` compiles +the two pick-ups, executes them, rebuilds the held-object state from measured +poses, and only then compiles placement. + +## Dynamic goals and closed-loop recovery + +Pose-valued goals can use `SceneEntityPose` instead of freezing an object pose +at invocation creation time: ```python -engine.register(Push(motion_gen, cfg=PushCfg(push_distance=0.08))) -is_success, traj, final_state = engine.run( - steps=[("push", PushTarget(xpos=target_pose))] +from embodichain.lab.sim.atomic_actions import ( + ExecutionStatus, + RecoveryPolicy, + SceneEntityPose, +) + +moving_goal = ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal( + xpos=SceneEntityPose( + entity_id="moving_tray", + relative_pose=tray_to_tcp, + minimum_confidence=0.8, + ) + ), + binding=ActionBinding(manipulators={"primary": "left_arm"}), + recovery_policy=RecoveryPolicy( + max_replans=3, + max_phase_retries=2, + tracking_error_threshold=0.05, + goal_translation_threshold=0.02, + goal_rotation_threshold=0.087, + phase_timeout=30.0, + ), ) + +latest_context = initial_context +session = engine.start((moving_goal,), latest_context) +while session.status is ExecutionStatus.RUNNING: + tick = session.tick(latest_context) + if tick.command is not None: + send_joint_command(tick.command) + latest_context = observe_context() ``` -To publish an action class for third-party discovery (independent of any engine instance), -use the global registry: +On each tick, the session can detect: + +- joint tracking error relative to the previously emitted command; +- translation or rotation of a `SceneEntityPose` dependency beyond policy + thresholds; +- phase timeout; +- planning or terminal-goal failure for individual batch rows. + +Recovery is bounded. A session replans from the latest observation, retries an +action only within the configured budgets, freezes ineligible environment rows, +and emits structured events when recovery is exhausted. + +Recovery does not re-read a mutable Action object or invocation. The engine +resolves each call once into a `ResolvedActionRequest` containing its binding, +policies, options, control commands, invocation ID, and revision. Every local +replan for that revision reuses the same request and varies only the measured +context. + +Use an explicit newer revision when the application or Action Agent decides to +change runtime behavior: ```python -from embodichain.lab.sim.atomic_actions import register_action, unregister_action, get_registered_actions +revised = ActionInvocation( + skill_id=current.skill_id, + goal=updated_goal, + binding=current.binding, + motion_policy=updated_motion_policy, + recovery_policy=current.recovery_policy, + skill_options=updated_options, + control_overrides=updated_commands, + invocation_id=current.invocation_id, + revision=current.revision + 1, +) +session.revise_current(revised) +``` -register_action("push", Push) # registers the class under "push" -unregister_action("push") # removes it -all_actions = get_registered_actions() # dict[str, type[AtomicAction]] +`skill_id` and `invocation_id` must still identify the active logical call. +Revision replacement preserves verified task state and environment eligibility, +resets the new revision's local recovery counters, emits +`INVOCATION_REVISED`, and replans from the latest context. + +```{attention} +Automatic dynamic-goal invalidation is dependency-driven. A goal must contain a +`SceneEntityPose` for the session to track that scene entity. A primitive that +directly queries a simulation entity during planning will use its latest pose +when planning happens, but that query alone does not trigger scene-motion +replanning. ``` -> **Tip:** `execute()` always returns an `ActionResult`. Its `trajectory` is full-robot-DOF -> shaped `(n_envs, n_waypoints, robot.dof)`, and `next_state` carries the `WorldState` the -> engine will feed into the following step. Use `TrajectoryBuilder` for pose broadcasting, -> three-phase splitting, IK/FK, and hand-qpos interpolation so your action matches the -> built-ins. +## Planning success versus physical success ---- +`ActionPlan.plan_success` only means a valid trajectory was produced for an +environment row. Pick, place, handover, and coordinated skills also return an +uncommitted `StateDelta` describing the attachment state expected after +execution. -```{toctree} -:maxdepth: 1 +At the terminal waypoint, an `ExecutionSession` requests an external +per-environment verification mask before committing a non-empty effect: -builtin_actions +```python +tick = session.tick(latest_context) +if any(event.kind is ExecutionEventKind.EFFECT_VERIFICATION_REQUIRED for event in tick.events): + effect_success = verify_grasp_or_release() + tick = session.tick(latest_context, effect_success=effect_success) ``` -## Further Reading - -- {doc}`../planners/motion_generator` — the trajectory planner used by every action -- {doc}`../sim_robot` — how control parts and IK solvers are configured -- Focused primitive demos: - - `scripts/tutorials/atomic_action/move_end_effector.py` - - `scripts/tutorials/atomic_action/move_joints.py` - - `scripts/tutorials/atomic_action/pickup.py` - - `scripts/tutorials/atomic_action/move_held_object.py` - - `scripts/tutorials/atomic_action/place.py` - - `scripts/tutorials/atomic_action/press.py` - - `scripts/tutorials/atomic_action/coordinated_pickment.py` - - `scripts/tutorials/atomic_action/coordinated_placement.py` - -Run a demo in headless CPU mode with `--auto_play --headless --device cpu` to record -an MP4 under `outputs/videos`. For example: - -```bash -python scripts/tutorials/atomic_action/move_end_effector.py --headless --auto_play --device cpu +This prevents a collision-free plan or well-tracked trajectory from being +misreported as a successful grasp, release, or handover. + +## Action Agent integration + +An MLLM should not construct `ActionInvocation` by copying arbitrary JSON into +runtime objects. An adapter should expose stable `SkillDescriptor` metadata and +agent-facing goal schemas, validate the semantic call, resolve object references +and embodiment capabilities, then produce the typed invocation: + +```text +MLLM SkillCallSpec + -> schema validation + -> object / scene grounding + -> capability and role binding + -> safe skill-option selection + -> semantic command selection (never raw qpos) + -> ActionInvocation + -> AtomicActionEngine + -> ActionPlan / execution events ``` + +The adapter may expose a curated subset of `OptionsType`, but engine-only +profiles, `JointPositionCommand` payloads, planner instances, and concrete joint +groups should remain outside the MLLM schema. If the agent needs an +object-specific grasp mode, it should choose a semantic command or capability; +the grounding layer turns that choice into `ActionControlOverrides`. Invocation +IDs and monotonic revisions correlate updated agent decisions with planner +diagnostics and execution events, providing structured feedback for the next +decision without mutating an in-flight request implicitly. + +## Extending the module + +A new primitive should: + +1. define a frozen, action-owned goal dataclass with a stable `goal_kind`; +2. define a frozen `ActionOptions` subclass only for behavior that can vary per + invocation; +3. declare `skill_id`, `GoalType`, `OptionsType`, required semantic roles, and + agent visibility; +4. put reusable embodiment commands on control-part profiles and generic + motion/recovery choices in invocation policies; +5. implement side-effect-free `plan(request, context)` using the + engine-owned planning services; +6. return full-robot timed motion, per-environment planning success, + diagnostics, and uncommitted effects; +7. add registration coverage, contract tests, execution/recovery tests, a + runnable example, and documentation. + +See {doc}`builtin_actions` for the shipped skill catalog and visual demos, and +{doc}`/tutorial/atomic_actions` for complete usage patterns and runnable scripts. + +## Further reading + +- {doc}`../planners/motion_generator` — the motion generator owned by the engine +- {doc}`../sim_robot` — robot control parts and kinematic configuration +- `scripts/tutorials/atomic_action/` — focused examples for every built-in skill diff --git a/docs/source/overview/sim/planners/curobo_planner.md b/docs/source/overview/sim/planners/curobo_planner.md index 68a3ce2c8..b1f9e2926 100644 --- a/docs/source/overview/sim/planners/curobo_planner.md +++ b/docs/source/overview/sim/planners/curobo_planner.md @@ -231,9 +231,9 @@ Panda) so planning stays fast; raise it for tighter collision coverage. MotionGenerator passes start_qpos and control_part to the cuRobo backend. For Cartesian goals, leave EmbodiChain pre-interpolation disabled: cuRobo must receive the original pose. By default the returned collision-checked samples are -arc-length resampled to the action's `sample_interval` waypoint count (so -`MoveEndEffectorCfg.sample_interval` controls the trajectory length, as for the -other planners); set `CuroboPlannerCfg.preserve_plan_samples=True` to keep +arc-length resampled to the invocation's `MotionPolicy.sample_count` waypoint +count (so the same runtime policy controls trajectory length across planners); +set `CuroboPlannerCfg.preserve_plan_samples=True` to keep cuRobo's own samples (whose count is derived from `interpolation_dt` and the trajectory duration). diff --git a/docs/source/tutorial/atomic_actions.rst b/docs/source/tutorial/atomic_actions.rst index f5d6f53c9..fc3b306f8 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -1,293 +1,370 @@ -.. _tutorial_atomic_actions: - -Atomic Actions +Atomic actions ============== -EmbodiChain's **atomic action** layer provides a high-level, composable interface for common -manipulation primitives such as *move end-effector*, *move joints*, *pick up*, -*move held object*, *place*, and *press*. Each action -encapsulates the full planning pipeline — grasp-pose estimation, IK, trajectory generation, and -gripper interpolation — behind a single ``execute(target, state)`` call, making it straightforward -to chain multiple actions together into complex robot behaviours. - -Key Features ------------- - -- **Typed targets** — every action accepts a small target dataclass such as - ``EndEffectorPoseTarget``, ``JointPositionTarget``, ``NamedJointPositionTarget``, - ``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. - See :doc:`/overview/sim/atomic_actions/index` for configs and target types. -- **Extensible registry** — custom action *classes* can be registered globally with - ``register_action``; action *instances* are registered per-engine under a name. -- **Engine orchestration** — ``AtomicActionEngine.run(steps, state)`` sequences named - ``(name, typed_target)`` steps, threads a ``WorldState`` (``last_qpos`` + - ``held_objects`` / ``coordinated_held_objects``) - from one action into the next, and returns a single concatenated full-DOF trajectory - ready to replay in the simulator. - -For the full design overview, architecture diagram, and extension guide see -:doc:`/overview/sim/atomic_actions/index`. - -The Code --------- - -Focused demo scripts are available for the built-in primitives in the -``scripts/tutorials/atomic_action`` directory: - -- ``move_end_effector.py`` -- ``move_joints.py`` -- ``pickup.py`` -- ``move_held_object.py`` -- ``place.py`` -- ``press.py`` -- ``coordinated_pickment.py`` -- ``coordinated_placement.py`` - -Each script supports interactive inspection by default. Add ``--auto_play`` to skip -keyboard prompts, and combine it with ``--headless --device cpu`` to record an MP4 under -``outputs/videos``: +Atomic actions are typed, side-effect-free motion planners. The engine resolves +a grounded :class:`~embodichain.lab.sim.atomic_actions.ActionInvocation` into a +:class:`~embodichain.lab.sim.atomic_actions.ResolvedActionRequest`; an action +combines that snapshot with the latest +:class:`~embodichain.lab.sim.atomic_actions.PlanningContext` and returns an +:class:`~embodichain.lab.sim.atomic_actions.ActionPlan`. + +For the complete architecture and ownership model, see +:doc:`/overview/sim/atomic_actions/index`. For the capability matrix and visual +demonstrations of every built-in skill, see +:doc:`/overview/sim/atomic_actions/builtin_actions`. + +The contracts deliberately separate six concerns: + +* a **goal** describes what should happen; +* an **ActionBinding** maps semantic roles such as ``primary`` or ``source`` to + names declared in the engine robot's ``control_parts`` mapping; +* a **ControlPartCommandProfile** maps embodiment-specific meanings such as + ``open``, ``grasp``, or ``ready`` to typed commands; +* typed **ActionOptions** contain behavior that may vary for one skill call; +* a **MotionPolicy** and **RecoveryPolicy** describe reusable planning and + bounded-recovery choices; +* a **PlanningContext** contains measured robot state, verified task state, and + a versioned scene snapshot. + +Binding values are keys from ``RobotCfg.control_parts``. They are not joint, +link, TCP-frame, or scene-object names. The engine validates them and resolves +their full-robot joint indices before planning. The ``end_effectors`` map names +an actuated hand/tool control part rather than an IK end frame. + +A role is an action-defined semantic participant slot, not a control part. In +``{"primary": "left_arm"}``, ``primary`` means the principal participant of +that single-participant action, while ``left_arm`` is the concrete control-part +key. It has no inherent left/right or default-arm meaning. Actions publish their +required slots through ``manipulator_roles`` and ``end_effector_roles``. When a +role such as ``primary`` occurs in both maps, the entries select the arm and +hand/tool serving the same functional participant, but the caller is still +responsible for choosing a physically compatible pair. + +The engine exclusively owns the ``MotionGenerator``, shared trajectory builder, +and control-part profiles. It creates and binds all built-in actions by default; +callers select them by stable ``skill_id`` without a separate ``register()`` +step. Put invocation-varying behavior in ``ActionInvocation.skill_options``. +``register()`` remains available for custom implementations, and +``load_builtins=False`` creates an isolated or fully custom engine. + +Choosing an engine entry point +------------------------------ + +Application code normally uses one of three engine entry points: + +.. list-table:: + :header-rows: 1 + :widths: 18 26 22 34 + + * - API + - Choose it when + - Returns + - State and observation behavior + * - ``engine.plan()`` + - Planning or inspecting one action + - ``ActionPlan`` + - Reads one context and does not project a next context + * - ``engine.compile()`` + - Planning a fixed sequence whose goals are already known + - ``CompiledTrajectory`` + - Propagates hypothetical qpos and expected effects, without observing execution + * - ``engine.start()`` + - Executing from fresh observations with bounded recovery + - ``ExecutionSession`` + - ``tick()`` consumes measured context, emits commands, requests effect verification, and can replan + +As a short rule: use ``plan`` for one action, ``compile`` for a static action +sequence, and ``start`` followed by ``tick`` for observed execution and error +recovery. None of these APIs steps the simulator directly. The application +sends commands returned by an execution session and supplies new observations. + +``AtomicAction.plan(request, context)`` is different from ``engine.plan()``. +It is the implementation method overridden by an atomic-action author, not an +additional application execution entry point. Similarly, +``engine.plan_action()`` is reserved for extensions and isolated tests that +need to plan an unregistered instance. + +Runnable examples +----------------- + +Focused examples live under ``scripts/tutorials/atomic_action``: + +* ``move_end_effector.py`` +* ``move_joints.py`` +* ``pickup.py`` +* ``move_held_object.py`` +* ``place.py`` +* ``assemble.py`` +* ``press.py`` +* ``coordinated_pickment.py`` +* ``coordinated_placement.py`` +* ``hand_over.py`` + +The scripts are interactive by default. Add ``--auto_play`` to skip prompts; +combine it with ``--headless --device cpu`` for a headless run that records +video under ``outputs/videos``: .. code-block:: bash python scripts/tutorials/atomic_action/move_end_effector.py --headless --auto_play --device cpu - python scripts/tutorials/atomic_action/move_joints.py --headless --auto_play --device cpu python scripts/tutorials/atomic_action/pickup.py --headless --auto_play --device cpu - python scripts/tutorials/atomic_action/move_held_object.py --headless --auto_play --device cpu - python scripts/tutorials/atomic_action/place.py --headless --auto_play --device cpu - python scripts/tutorials/atomic_action/press.py --headless --auto_play --device cpu - python scripts/tutorials/atomic_action/coordinated_pickment.py --headless --auto_play --device cpu - python scripts/tutorials/atomic_action/coordinated_placement.py --headless --auto_play --device cpu + python scripts/tutorials/atomic_action/assemble.py --headless --auto_play --device cpu + python scripts/tutorials/atomic_action/hand_over.py --headless --auto_play --device cpu -The concrete implementations are organized one primitive per module under -``embodichain/lab/sim/atomic_actions/primitives``. Public imports from -``embodichain.lab.sim.atomic_actions`` remain the recommended API, and -``embodichain.lab.sim.atomic_actions.actions`` stays as a compatibility -re-export surface. +The ``motion_generator`` variable in the snippets below is a configured +:class:`~embodichain.lab.sim.planners.MotionGenerator`; its robot, planner, +device, cache, and collision world become the resources owned by the engine. -Typical Usage -------------- +Control-part commands +--------------------- -Setting up the engine -~~~~~~~~~~~~~~~~~~~~~ +Hand qpos and named robot postures are robot knowledge rather than action +configuration. Register them by concrete ``Robot.control_parts`` key when the +engine is built: .. code-block:: python - import torch - from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, - PickUp, PickUpCfg, - Place, PlaceCfg, - MoveEndEffector, MoveEndEffectorCfg, - MoveJoints, MoveJointsCfg, + ControlPartCommandProfile, ) - motion_gen = MotionGenerator(cfg=MotionGenCfg(...)) - - hand_open = torch.tensor([0.00, 0.00], dtype=torch.float32, device=device) - hand_close = torch.tensor([0.025, 0.025], dtype=torch.float32, device=device) - - pickup_cfg = PickUpCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - control_part="arm", - hand_control_part="hand", - approach_direction=torch.tensor([0.0, 0.0, -1.0], dtype=torch.float32, device=device), - pre_grasp_distance=0.15, - lift_height=0.15, - ) - place_cfg = PlaceCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - control_part="arm", - hand_control_part="hand", - lift_height=0.15, - ) - move_cfg = MoveEndEffectorCfg(control_part="arm") - move_joints_cfg = MoveJointsCfg( - control_part="arm", - named_joint_positions={"home": torch.zeros(6, dtype=torch.float32, device=device)}, + engine = AtomicActionEngine( + motion_generator, + control_profiles={ + "left_hand": ControlPartCommandProfile.joint_positions( + open=left_open_qpos, + grasp=left_grasp_qpos, + ), + "left_arm": ControlPartCommandProfile.joint_positions( + ready=left_ready_qpos, + ), + }, ) - # The engine takes only the motion generator; register each action instance - # by name (defaults to the action's cfg.name). - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register(PickUp(motion_gen, cfg=pickup_cfg)) - engine.register(Place(motion_gen, cfg=place_cfg)) - engine.register(MoveEndEffector(motion_gen, cfg=move_cfg)) - engine.register(MoveJoints(motion_gen, cfg=move_joints_cfg)) +``PickUp``, ``Place``, and the other manipulation skills resolve ``open`` and +``grasp`` from their bound end effector. ``MoveJoints`` resolves a string target +from its bound manipulator. Joint limits validate possible commands, but do not +define their semantic meaning; supply calibrated robot commands in production. -Defining object semantics -~~~~~~~~~~~~~~~~~~~~~~~~~ +Planning one action +------------------- + +Use :meth:`~embodichain.lab.sim.atomic_actions.AtomicActionEngine.plan` when one +registered action needs to be inspected, tested, or passed through +application-owned orchestration: .. code-block:: python from embodichain.lab.sim.atomic_actions import ( - ObjectSemantics, - AntipodalAffordance, - ) - from embodichain.toolkits.graspkit.pg_grasp import GraspGeneratorCfg, AntipodalSamplerCfg - from embodichain.toolkits.graspkit.pg_grasp.gripper_collision_checker import GripperCollisionCfg - - affordance = AntipodalAffordance( - mesh_vertices=mug.get_vertices(env_ids=[0], scale=True)[0], - mesh_triangles=mug.get_triangles(env_ids=[0])[0], - gripper_collision_cfg=GripperCollisionCfg( - max_open_length=0.088, finger_length=0.078, point_sample_dense=0.012 - ), - generator_cfg=GraspGeneratorCfg( - antipodal_sampler_cfg=AntipodalSamplerCfg( - n_sample=20000, max_length=0.088, min_length=0.003 - ) - ), - force_reannotate=False, + ActionBinding, + ActionInvocation, + EndEffectorPoseGoal, + MotionPolicy, ) - semantics = ObjectSemantics( - affordance=affordance, - geometry={}, # plain metadata; mesh data lives on the affordance - label="mug", # also bound onto affordance.object_label - entity=mug, # required so the action can query the live object pose + invocation = ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal(xpos=target_pose), + binding=ActionBinding(manipulators={"primary": "left_arm"}), + motion_policy=MotionPolicy(sample_count=80, control_dt=1.0 / 60.0), ) -Executing a pick-place-end-effector sequence -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + plan = engine.plan(invocation, latest_context) + if plan.plan_success.all(): + trajectory = plan.trajectory.positions + phase_diagnostics = tuple(phase.diagnostics for phase in plan.phases) + +The returned :class:`~embodichain.lab.sim.atomic_actions.ActionPlan` describes +only that invocation. Its expected effects are not committed, and ``plan`` does +not produce a projected context for a following action. Use ``compile`` when +the engine should propagate hypothetical state through a sequence. + +Static compilation +------------------ + +Use :meth:`~embodichain.lab.sim.atomic_actions.AtomicActionEngine.compile` when +the scene is treated as fixed and all goals in a sequence are known during +planning: .. code-block:: python from embodichain.lab.sim.atomic_actions import ( - GraspTarget, - EndEffectorPoseTarget, - PlaceTarget, + ActionBinding, + ActionInvocation, + AtomicActionEngine, + EndEffectorPoseGoal, + MotionPolicy, ) - place_xpos = ... # torch.Tensor [4, 4] — target placement pose - rest_xpos = ... # torch.Tensor [4, 4] — resting pose after placing + engine = AtomicActionEngine(motion_generator) + binding = ActionBinding(manipulators={"primary": "left_arm"}) + motion_policy = MotionPolicy(sample_count=80, control_dt=1.0 / 60.0) - is_success, trajectory, _ = engine.run( - steps=[ - ("pick_up", GraspTarget(semantics=semantics)), - ("place", PlaceTarget(xpos=place_xpos)), - ("move_end_effector", EndEffectorPoseTarget(xpos=rest_xpos)), - ] + approach = ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal(xpos=approach_pose), + binding=binding, + motion_policy=motion_policy, + ) + retreat = ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal(xpos=retreat_pose), + binding=binding, + motion_policy=motion_policy, ) - # trajectory: [n_envs, n_waypoints, robot_dof] - for i in range(trajectory.shape[1]): - robot.set_qpos(trajectory[:, i]) - sim.update(step=4) + initial_context = engine.initial_context() + compiled = engine.compile((approach, retreat), initial_context) + if compiled.plan_success.all(): + trajectory = compiled.trajectory.positions + approach_plan, retreat_plan = compiled.action_plans + final_context = compiled.projected_context + +``compile`` never steps the simulator. It applies each plan's expected +:class:`~embodichain.lab.sim.atomic_actions.StateDelta` only to the returned +``projected_context`` so a following action can be planned against hypothetical +state. Calling it with one invocation is valid, but ``plan`` is simpler when a +projected context and sequence-shaped result are unnecessary. + +Do not compile across a point where later targets depend on physical execution. +The coordinated-placement tutorial, for example, compiles both pick-ups, +executes them, rebuilds held-object state from measured poses, and then compiles +the placement phase. Use ``start`` when that observation and recovery loop +should remain active throughout execution. -Moving in joint space -~~~~~~~~~~~~~~~~~~~~~ +Dynamic goals and closed-loop execution +--------------------------------------- -``MoveJoints`` is separate from ``MoveEndEffector`` so a plan says clearly whether the -target is a pose or a qpos. It accepts either an explicit ``JointPositionTarget`` or a -``NamedJointPositionTarget`` resolved from ``MoveJointsCfg.named_joint_positions``. +Use :class:`~embodichain.lab.sim.atomic_actions.SceneEntityPose` when a goal +must be resolved from the latest scene snapshot: .. code-block:: python from embodichain.lab.sim.atomic_actions import ( - JointPositionTarget, - NamedJointPositionTarget, + EndEffectorPoseGoal, + RecoveryPolicy, + SceneEntityPose, ) - home_qpos = torch.zeros(6, dtype=torch.float32, device=device) - - is_success, trajectory, _ = engine.run( - steps=[ - ("move_joints", NamedJointPositionTarget(name="home")), - ("move_joints", JointPositionTarget(qpos=home_qpos)), - ] + invocation = ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal( + xpos=SceneEntityPose("moving_tray", relative_pose=tray_to_tcp) + ), + binding=ActionBinding(manipulators={"primary": "left_arm"}), + recovery_policy=RecoveryPolicy( + max_replans=3, + tracking_error_threshold=0.05, + goal_translation_threshold=0.02, + ), ) -Moving a held object -~~~~~~~~~~~~~~~~~~~~ + latest_context = initial_context + session = engine.start((invocation,), latest_context) + while session.status.value == "running": + tick = session.tick(latest_context) + if tick.command is not None: + send_joint_command(tick.command) + latest_context = observe_context() -``MoveHeldObject`` consumes the runtime ``HeldObjectState`` produced by a previous -``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. +The session emits one command per tick. It compares observations with the last +command, detects material motion of referenced scene entities, enforces phase +timeouts, and replans from the latest observation within the recovery budget. +It does not own the simulator or controller loop. + +Recovery replans reuse one immutable invocation-revision snapshot. If an +application intentionally changes the goal, options, policy, binding, or a +control command while the action is active, submit a strictly newer revision: .. code-block:: python - from embodichain.lab.sim.atomic_actions import ( - MoveHeldObject, MoveHeldObjectCfg, - GraspTarget, HeldObjectPoseTarget, + revised = ActionInvocation( + skill_id=invocation.skill_id, + goal=updated_goal, + binding=invocation.binding, + motion_policy=invocation.motion_policy, + recovery_policy=invocation.recovery_policy, + skill_options=updated_options, + control_overrides=updated_control_commands, + invocation_id=invocation.invocation_id, + revision=invocation.revision + 1, ) + session.revise_current(revised) - move_held_object_cfg = MoveHeldObjectCfg( - hand_close_qpos=hand_close, - control_part="arm", - hand_control_part="hand", - ) - engine.register(MoveHeldObject(motion_gen, cfg=move_held_object_cfg)) +The session replans from its latest context and emits an +``invocation_revised`` event. ``skill_id`` and ``invocation_id`` must still +identify the active logical call. - object_target_pose = torch.eye(4, dtype=torch.float32, device=device) - object_target_pose[:3, 3] = torch.tensor([0.3, -0.2, 0.25], device=device) +Only entities referenced through ``SceneEntityPose`` become automatic +scene-motion dependencies. A skill may query a simulation entity's live pose +when it plans, but that query alone does not cause an executing session to +replan when the entity moves. - is_success, trajectory, _ = engine.run( - steps=[ - ("pick_up", GraspTarget(semantics=semantics)), - ("move_held_object", HeldObjectPoseTarget(object_target_pose=object_target_pose)), - ] - ) +Task-state effects +------------------ -Registering custom actions -~~~~~~~~~~~~~~~~~~~~~~~~~~ +Pick, place, handover, and coordinated skills declare attachment changes as a +:class:`~embodichain.lab.sim.atomic_actions.StateDelta`. Planning does not commit +those changes. During closed-loop execution, a non-empty effect requires an +external per-environment verification mask: + +.. code-block:: python + + tick = session.tick(latest_context) + if any(event.kind.value == "effect_verification_required" for event in tick.events): + verified = verify_grasp_or_release() + tick = session.tick(latest_context, effect_success=verified) + +This prevents a successful trajectory plan from being mistaken for a successful +physical grasp or release. + +Adding an action +---------------- + +Define an action-owned frozen goal dataclass with a stable ``goal_kind``. Then +define typed runtime options when needed, implement ``plan(request, context)``, +and declare the stable skill metadata: .. code-block:: python from dataclasses import dataclass from typing import ClassVar - import torch - from embodichain.lab.sim.atomic_actions import ( - ActionTarget, AtomicAction, ActionResult, WorldState, TrajectoryBuilder, - ) - @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: PushTarget, state: WorldState) -> ActionResult: - # ... your planning logic, using self.builder ... - return ActionResult(success=is_success, trajectory=full, next_state=...) - - # Register an instance with an engine so it can be referenced by name in run(): - engine.register(Push(motion_gen, cfg=PushCfg())) - - # Or publish the class globally for third-party discovery: - from embodichain.lab.sim.atomic_actions import register_action - register_action("push", Push) - -Notes & Best Practices ----------------------- - -- ``PickUp`` expects an ``AntipodalAffordance`` with valid mesh data - (``mesh_vertices`` / ``mesh_triangles``) so the grasp generator can annotate the object. - Set ``force_reannotate=False`` (the default) to reuse cached annotations across episodes. -- ``ObjectSemantics.entity`` must be set when using semantic targets so the action can read - the object's current world pose at planning time. -- 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`. + @dataclass(frozen=True, slots=True) + class PushGoal: + goal_kind: ClassVar[str] = "push" + contact_pose: torch.Tensor + + @dataclass(frozen=True, slots=True) + class PushOptions(ActionOptions): + retreat_distance: float = 0.1 + + class Push(AtomicAction[PushGoal, PushOptions]): + skill_id: ClassVar[str] = "push" + GoalType: ClassVar[type] = PushGoal + OptionsType: ClassVar[type] = PushOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + + def __init__(self, default_options: PushOptions | None = None) -> None: + super().__init__(default_options) + + def plan( + self, + request: ResolvedActionRequest[PushGoal, PushOptions], + context: PlanningContext, + ) -> ActionPlan: + goal = self.require_goal(request) + options = request.skill_options + # Resolve the bound resource, plan from context.robot.qpos, and + # return a full-robot TimedTrajectory or position tensor. + return self.build_plan( + request, + context, + success=success_mask, + trajectory=full_robot_positions, + ) + +Do not step simulation, mutate ``PlanningContext``, commit ``StateDelta``, or +expose planner-specific configuration through the goal. See the in-repository +``add-atomic-action`` skill for the complete checklist. diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index f713bccc5..7564a52b4 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -14,13 +14,15 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Atomic action abstraction layer for embodied AI motion generation. +"""Typed planning contracts and built-in atomic actions. -This module provides a unified interface for the atomic motion primitives -(``move_end_effector``, ``move_joints``, ``pick_up``, ``move_held_object``, -``place``, ``press``, ``coordinated_pickment``, ``coordinated_placement``, -``hand_over``), with typed targets, a ``WorldState`` threaded across sequenced -actions, and extensible custom action registration. +The engine resolves an :class:`ActionInvocation` into a +:class:`ResolvedActionRequest`, which an action combines with a +:class:`PlanningContext` through :meth:`AtomicAction.plan`. Planning is +side-effect free: it returns an :class:`ActionPlan` with timed motion, +completion criteria, diagnostics, and uncommitted expected task-state effects. +:class:`AtomicActionEngine` can compile a static sequence; closed-loop execution +belongs to an execution session. """ from __future__ import annotations @@ -31,126 +33,166 @@ AssembleAffordance, InteractionPoints, ) -from .core import ( - ActionTarget, - ActionCfg, - ActionResult, - AtomicAction, - CoordinatedHeldObjectState, - HeldObjectState, - ObjectSemantics, - Target, - WorldState, +from .bindings import ActionBinding, ResolvedActionBinding, ResolvedControlPart +from .control import ( + ActionControlOverrides, + ControlCommand, + ControlPartCommandProfile, + GRASP_COMMAND, + JointPositionCommand, + OPEN_COMMAND, ) +from .core import AtomicAction, ObjectSemantics, SkillDescriptor +from .effects import StateDelta from .engine import ( AtomicActionEngine, + get_registered_actions, register_action, unregister_action, - get_registered_actions, ) -from .targets import ObjectActionTarget +from .execution import ( + ExecutionEvent, + ExecutionEventKind, + ExecutionSession, + ExecutionStatus, + ExecutionTick, + JointCommand, +) +from .goals import ActionGoal, ObjectActionGoal, PoseGoalValue, SceneEntityPose +from .invocation import ActionInvocation, ActionOptions, ResolvedActionRequest +from .plans import ( + ActionPlan, + CompiledTrajectory, + CompletionCondition, + CompletionConditionKind, + PhaseSpec, + PlannedPhase, + PlannerDiagnostics, + TimedTrajectory, +) +from .policies import MotionPolicy, RecoveryPolicy +from .runtime import ActionPlanningServices from .primitives import ( - AssembleTarget, - CoordinatedPickTarget, + AssembleGoal, + BUILTIN_ACTION_TYPES, + CoordinatedPickGoal, CoordinatedPickment, - CoordinatedPickmentCfg, - CoordinatedPickmentTarget, + CoordinatedPickmentOptions, CoordinatedPlacement, - CoordinatedPlacementCfg, - CoordinatedPlacementTarget, - EndEffectorPoseTarget, - GraspTarget, + CoordinatedPlacementGoal, + CoordinatedPlacementOptions, + EndEffectorPoseGoal, + GraspGoal, HandOver, - HandOverCfg, - HeldObjectPoseTarget, - JointPositionTarget, + HandOverOptions, + HeldObjectPoseGoal, + JointPositionGoal, MoveEndEffector, - MoveEndEffectorCfg, + MoveEndEffectorOptions, MoveHeldObject, - MoveHeldObjectCfg, + MoveHeldObjectOptions, MoveJoints, - MoveJointsCfg, - NamedJointPositionTarget, + MoveJointsOptions, PickUp, - PickUpCfg, + PickUpOptions, Place, - PlaceCfg, - PlaceTarget, + PlaceGoal, + PlaceOptions, Press, - PressCfg, - PressTarget, + PressGoal, + PressOptions, ) -from .trajectory import TrajectoryBuilder - -BuiltinTarget = ( - EndEffectorPoseTarget - | JointPositionTarget - | NamedJointPositionTarget - | GraspTarget - | HeldObjectPoseTarget - | PlaceTarget - | PressTarget - | CoordinatedPickTarget - | CoordinatedPlacementTarget - | AssembleTarget +from .state import ( + CoordinatedHeldObjectState, + EntityState, + HeldObjectState, + PlanningContext, + RobotObservation, + SceneSnapshot, + TaskState, ) -"""Union of target types shipped by EmbodiChain. - -Use :class:`ActionTarget` rather than this closed union at extension boundaries. -""" +from .trajectory import TrajectoryBuilder __all__ = [ - # Core classes - "ActionTarget", + "ActionBinding", + "ActionControlOverrides", + "ActionGoal", + "ActionInvocation", + "ActionOptions", + "ActionPlan", + "ActionPlanningServices", "Affordance", "AntipodalAffordance", "AssembleAffordance", - "InteractionPoints", - "ObjectSemantics", - "ObjectActionTarget", - "HeldObjectState", - "CoordinatedHeldObjectState", - "HeldObjectPoseTarget", - "JointPositionTarget", - "NamedJointPositionTarget", - "EndEffectorPoseTarget", - "PlaceTarget", - "PressTarget", - "CoordinatedPickTarget", - "CoordinatedPickmentTarget", - "CoordinatedPlacementTarget", - "AssembleTarget", - "GraspTarget", - "Target", - "BuiltinTarget", - "WorldState", - "ActionResult", - "ActionCfg", + "AssembleGoal", "AtomicAction", - # Action implementations + "AtomicActionEngine", + "BUILTIN_ACTION_TYPES", + "CompiledTrajectory", + "CompletionCondition", + "CompletionConditionKind", + "ControlCommand", + "ControlPartCommandProfile", + "CoordinatedHeldObjectState", + "CoordinatedPickGoal", "CoordinatedPickment", + "CoordinatedPickmentOptions", "CoordinatedPlacement", + "CoordinatedPlacementGoal", + "CoordinatedPlacementOptions", + "EndEffectorPoseGoal", + "EntityState", + "ExecutionEvent", + "ExecutionEventKind", + "ExecutionSession", + "ExecutionStatus", + "ExecutionTick", + "GRASP_COMMAND", + "GraspGoal", "HandOver", + "HandOverOptions", + "HeldObjectPoseGoal", + "HeldObjectState", + "InteractionPoints", + "JointPositionGoal", + "JointCommand", + "JointPositionCommand", + "MotionPolicy", "MoveEndEffector", - "MoveJoints", + "MoveEndEffectorOptions", "MoveHeldObject", + "MoveHeldObjectOptions", + "MoveJoints", + "MoveJointsOptions", + "ObjectActionGoal", + "ObjectSemantics", + "OPEN_COMMAND", + "PhaseSpec", "PickUp", + "PickUpOptions", "Place", + "PlaceGoal", + "PlaceOptions", + "PlannedPhase", + "PlannerDiagnostics", + "PlanningContext", + "PoseGoalValue", "Press", - "CoordinatedPickmentCfg", - "CoordinatedPlacementCfg", - "HandOverCfg", - "MoveEndEffectorCfg", - "MoveJointsCfg", - "MoveHeldObjectCfg", - "PickUpCfg", - "PlaceCfg", - "PressCfg", - # Engine - "AtomicActionEngine", + "PressGoal", + "PressOptions", + "RecoveryPolicy", + "ResolvedActionRequest", + "ResolvedActionBinding", + "ResolvedControlPart", + "RobotObservation", + "SceneSnapshot", + "SceneEntityPose", + "SkillDescriptor", + "StateDelta", + "TaskState", + "TimedTrajectory", + "TrajectoryBuilder", + "get_registered_actions", "register_action", "unregister_action", - "get_registered_actions", - # Trajectory helpers - "TrajectoryBuilder", ] diff --git a/embodichain/lab/sim/atomic_actions/actions.py b/embodichain/lab/sim/atomic_actions/actions.py deleted file mode 100644 index 37911831b..000000000 --- a/embodichain/lab/sim/atomic_actions/actions.py +++ /dev/null @@ -1,61 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Compatibility re-exports for built-in atomic actions. - -Concrete action implementations live in ``atomic_actions.primitives``. Importing -from ``atomic_actions.actions`` remains supported for existing callers. -""" - -from __future__ import annotations - -from .primitives import ( - CoordinatedPickment, - CoordinatedPickmentCfg, - CoordinatedPlacement, - CoordinatedPlacementCfg, - MoveEndEffector, - MoveEndEffectorCfg, - MoveHeldObject, - MoveHeldObjectCfg, - MoveJoints, - MoveJointsCfg, - PickUp, - PickUpCfg, - Place, - PlaceCfg, - Press, - PressCfg, -) - -__all__ = [ - "CoordinatedPickment", - "CoordinatedPickmentCfg", - "CoordinatedPlacement", - "CoordinatedPlacementCfg", - "MoveEndEffector", - "MoveEndEffectorCfg", - "MoveHeldObject", - "MoveHeldObjectCfg", - "MoveJoints", - "MoveJointsCfg", - "PickUp", - "PickUpCfg", - "Place", - "PlaceCfg", - "Press", - "PressCfg", -] diff --git a/embodichain/lab/sim/atomic_actions/bindings.py b/embodichain/lab/sim/atomic_actions/bindings.py new file mode 100644 index 000000000..5257c5035 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/bindings.py @@ -0,0 +1,292 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Semantic-role to robot control-part bindings for atomic actions.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Mapping + +import torch + +from .control import ControlCommand, JointPositionCommand + + +def _normalize_resource_map( + values: Mapping[str, str], + *, + field_name: str, +) -> Mapping[str, str]: + """Validate and freeze a semantic-role resource mapping.""" + if not isinstance(values, Mapping): + raise TypeError(f"{field_name} must be a mapping.") + normalized: dict[str, str] = {} + for role, resource in values.items(): + if not isinstance(role, str) or not role.strip(): + raise ValueError(f"{field_name} roles must be non-empty strings.") + if not isinstance(resource, str) or not resource.strip(): + raise ValueError(f"{field_name} resources must be non-empty strings.") + normalized[role] = resource + return MappingProxyType(normalized) + + +@dataclass(frozen=True, slots=True) +class ActionBinding: + """Bind semantic action roles to names from ``Robot.control_parts``. + + A role such as ``primary``, ``source`` or ``destination`` is an + action-defined semantic participant slot. It describes the responsibility + a resource has within that action and is not itself a robot resource. + Actions publish their required slots through ``manipulator_roles`` and + ``end_effector_roles``. Role names are scoped independently to those two + maps, so matching names associate an arm and hand/tool with the same + functional participant without making the maps interchangeable. + + ``primary`` has no inherent left/right, ordering, or default-control-part + meaning. Only the compiler or application binding layer needs to map it to + concrete robot control-part names such as ``left_arm`` and ``left_hand``. + + Every mapping value is a key from the current robot's ``control_parts`` + configuration. This value object validates the mapping shape; the + :class:`~embodichain.lab.sim.atomic_actions.AtomicActionEngine` validates + the names against its owned robot before planning. ``end_effectors`` refers + to actuated tool/hand control parts, not TCP or kinematic frame names. + """ + + manipulators: Mapping[str, str] = field(default_factory=dict) + """Manipulator control-part names keyed by semantic role.""" + + end_effectors: Mapping[str, str] = field(default_factory=dict) + """Tool or hand control-part names keyed by semantic role.""" + + def __post_init__(self) -> None: + object.__setattr__( + self, + "manipulators", + _normalize_resource_map(self.manipulators, field_name="manipulators"), + ) + object.__setattr__( + self, + "end_effectors", + _normalize_resource_map(self.end_effectors, field_name="end_effectors"), + ) + + def manipulator(self, role: str = "primary") -> str: + """Return the manipulator control-part name bound to ``role``. + + Args: + role: Semantic manipulator role. + + Returns: + Key from the current robot's ``control_parts`` mapping. + + Raises: + KeyError: If the requested role is not bound. + """ + try: + return self.manipulators[role] + except KeyError as exc: + raise KeyError(f"No manipulator is bound to role {role!r}.") from exc + + def end_effector(self, role: str = "primary") -> str: + """Return the tool/hand control-part name bound to ``role``. + + Args: + role: Semantic end-effector role. + + Returns: + Key from the current robot's ``control_parts`` mapping. + + Raises: + KeyError: If the requested role is not bound. + """ + try: + return self.end_effectors[role] + except KeyError as exc: + raise KeyError(f"No end effector is bound to role {role!r}.") from exc + + +@dataclass(frozen=True, slots=True) +class ResolvedControlPart: + """One engine-validated robot control part. + + Instances are produced by engine-owned planning services. They keep + robot-specific indices out of :class:`ActionBinding` and agent-facing + invocation schemas. + """ + + name: str + """Key from ``Robot.control_parts``.""" + + joint_ids: tuple[int, ...] + """Full-robot joint indices belonging to this control part.""" + + commands: Mapping[str, ControlCommand] = field(default_factory=dict) + """Engine-profile commands, including invocation-level overrides.""" + + def __post_init__(self) -> None: + if not isinstance(self.name, str) or not self.name.strip(): + raise ValueError("ResolvedControlPart.name must be a non-empty string.") + joint_ids = tuple(self.joint_ids) + if not joint_ids or not all( + isinstance(joint_id, int) and joint_id >= 0 for joint_id in joint_ids + ): + raise ValueError( + "ResolvedControlPart.joint_ids must contain non-negative integers." + ) + if len(set(joint_ids)) != len(joint_ids): + raise ValueError("ResolvedControlPart.joint_ids must be unique.") + object.__setattr__(self, "joint_ids", joint_ids) + if not isinstance(self.commands, Mapping): + raise TypeError("ResolvedControlPart.commands must be a mapping.") + commands: dict[str, ControlCommand] = {} + for name, command in self.commands.items(): + if not isinstance(name, str) or not name.strip(): + raise ValueError("Control command names must be non-empty strings.") + if not isinstance(command, ControlCommand): + raise TypeError( + "ResolvedControlPart.commands values must be ControlCommand " + "instances." + ) + commands[name] = command.snapshot() + object.__setattr__(self, "commands", MappingProxyType(commands)) + + @property + def dof(self) -> int: + """Return the number of joints in this control part.""" + return len(self.joint_ids) + + def with_command_overrides( + self, + overrides: Mapping[str, ControlCommand], + ) -> ResolvedControlPart: + """Return a snapshot with role-local semantic command overrides.""" + merged = dict(self.commands) + merged.update(overrides) + return ResolvedControlPart( + name=self.name, + joint_ids=self.joint_ids, + commands=merged, + ) + + def command(self, name: str) -> ControlCommand: + """Return an owned semantic command snapshot. + + Args: + name: Semantic command name, for example ``open`` or ``grasp``. + + Raises: + KeyError: If this control part does not define ``name``. + """ + try: + command = self.commands[name] + except KeyError as exc: + raise KeyError( + f"Control part {self.name!r} has no command {name!r}. " + f"Available commands: {sorted(self.commands)}." + ) from exc + return command.snapshot() + + def joint_positions( + self, + name: str, + *, + n_envs: int, + device: torch.device | str, + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + """Resolve a named joint-position command for a planning batch.""" + try: + command = self.commands[name] + except KeyError as exc: + raise KeyError( + f"Control part {self.name!r} has no command {name!r}. " + f"Available commands: {sorted(self.commands)}." + ) from exc + if not isinstance(command, JointPositionCommand): + raise TypeError( + f"Control command {name!r} on {self.name!r} is " + f"{type(command).__name__}, not JointPositionCommand." + ) + return command.resolve( + n_envs=n_envs, + control_dof=self.dof, + device=device, + dtype=dtype, + ) + + +def _normalize_resolved_map( + values: Mapping[str, ResolvedControlPart], + *, + field_name: str, +) -> Mapping[str, ResolvedControlPart]: + """Validate and freeze a resolved semantic-role mapping.""" + if not isinstance(values, Mapping): + raise TypeError(f"{field_name} must be a mapping.") + normalized: dict[str, ResolvedControlPart] = {} + for role, resource in values.items(): + if not isinstance(role, str) or not role.strip(): + raise ValueError(f"{field_name} roles must be non-empty strings.") + if not isinstance(resource, ResolvedControlPart): + raise TypeError( + f"{field_name} values must be ResolvedControlPart instances." + ) + normalized[role] = resource + return MappingProxyType(normalized) + + +@dataclass(frozen=True, slots=True) +class ResolvedActionBinding: + """Runtime control parts resolved from an :class:`ActionBinding`.""" + + manipulators: Mapping[str, ResolvedControlPart] = field(default_factory=dict) + end_effectors: Mapping[str, ResolvedControlPart] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__( + self, + "manipulators", + _normalize_resolved_map( + self.manipulators, field_name="resolved manipulators" + ), + ) + object.__setattr__( + self, + "end_effectors", + _normalize_resolved_map( + self.end_effectors, field_name="resolved end_effectors" + ), + ) + + def manipulator(self, role: str = "primary") -> ResolvedControlPart: + """Return the resolved manipulator for ``role``.""" + try: + return self.manipulators[role] + except KeyError as exc: + raise KeyError(f"No manipulator is bound to role {role!r}.") from exc + + def end_effector(self, role: str = "primary") -> ResolvedControlPart: + """Return the resolved tool/hand control part for ``role``.""" + try: + return self.end_effectors[role] + except KeyError as exc: + raise KeyError(f"No end effector is bound to role {role!r}.") from exc + + +__all__ = ["ActionBinding", "ResolvedActionBinding", "ResolvedControlPart"] diff --git a/embodichain/lab/sim/atomic_actions/control.py b/embodichain/lab/sim/atomic_actions/control.py new file mode 100644 index 000000000..80be94029 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/control.py @@ -0,0 +1,247 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Semantic command profiles for robot control parts.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Mapping + +import torch + +OPEN_COMMAND = "open" +"""Conventional semantic command for an open end effector.""" + +GRASP_COMMAND = "grasp" +"""Conventional semantic command for an object-holding end effector.""" + + +class ControlCommand(ABC): + """Immutable-by-ownership command associated with one control part. + + Command subclasses own their payload and must return another owned value + from :meth:`snapshot`. This keeps engine profiles and resolved invocation + requests isolated from caller-owned mutable tensors. + """ + + @abstractmethod + def snapshot(self) -> ControlCommand: + """Return an independently owned copy of this command.""" + + +@dataclass(frozen=True, slots=True, eq=False, init=False) +class JointPositionCommand(ControlCommand): + """A semantic command represented by one or batched joint positions. + + ``positions`` has shape ``(control_dof,)`` or + ``(n_envs, control_dof)``. A one-dimensional command is broadcast to the + planning batch when resolved. + """ + + _positions: torch.Tensor + + def __init__(self, positions: torch.Tensor) -> None: + if not isinstance(positions, torch.Tensor): + raise TypeError("positions must be a torch.Tensor.") + if positions.dim() not in (1, 2) or positions.shape[-1] == 0: + raise ValueError( + "positions must have shape (control_dof,) or " + "(n_envs, control_dof), got " + f"{tuple(positions.shape)}." + ) + if not torch.isfinite(positions).all().item(): + raise ValueError("positions must contain only finite values.") + object.__setattr__(self, "_positions", positions.detach().clone()) + + @property + def positions(self) -> torch.Tensor: + """Return an owned copy of the command payload.""" + return self._positions.clone() + + def snapshot(self) -> JointPositionCommand: + """Return an independently owned command snapshot.""" + return JointPositionCommand(self._positions) + + def resolve( + self, + *, + n_envs: int, + control_dof: int, + device: torch.device | str, + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + """Validate, move, and broadcast this command for a planning batch. + + Args: + n_envs: Number of selected environments. + control_dof: Joint count of the resolved control part. + device: Target planning device. + dtype: Optional target dtype. + + Returns: + Independently owned tensor with shape ``(n_envs, control_dof)``. + + Raises: + ValueError: If the command shape does not match the control part or + selected environment batch. + """ + if not isinstance(n_envs, int) or n_envs < 1: + raise ValueError("n_envs must be a positive integer.") + if not isinstance(control_dof, int) or control_dof < 1: + raise ValueError("control_dof must be a positive integer.") + if self._positions.shape[-1] != control_dof: + raise ValueError( + f"Joint-position command has {self._positions.shape[-1]} joints, " + f"but the resolved control part has {control_dof}." + ) + resolved = self._positions.to(device=device, dtype=dtype) + if resolved.dim() == 1: + return resolved.unsqueeze(0).expand(n_envs, -1).clone() + if resolved.shape[0] != n_envs: + raise ValueError( + f"Batched joint-position command has {resolved.shape[0]} " + f"environments, expected {n_envs}." + ) + return resolved.clone() + + +def _snapshot_commands( + commands: Mapping[str, ControlCommand], + *, + field_name: str, +) -> Mapping[str, ControlCommand]: + """Validate and freeze a semantic command mapping.""" + if not isinstance(commands, Mapping): + raise TypeError(f"{field_name} must be a mapping.") + snapshots: dict[str, ControlCommand] = {} + for name, command in commands.items(): + if not isinstance(name, str) or not name.strip(): + raise ValueError(f"{field_name} keys must be non-empty strings.") + if not isinstance(command, ControlCommand): + raise TypeError(f"{field_name} values must be ControlCommand instances.") + snapshots[name] = command.snapshot() + return MappingProxyType(snapshots) + + +@dataclass(frozen=True, slots=True) +class ControlPartCommandProfile: + """Reusable semantic commands for one named robot control part. + + Profiles are registered once on :class:`AtomicActionEngine`, keyed by + names from ``Robot.control_parts``. They describe embodiment-specific + meanings such as ``open``, ``grasp`` or ``ready`` without coupling those + values to an action implementation. + """ + + commands: Mapping[str, ControlCommand] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__( + self, + "commands", + _snapshot_commands(self.commands, field_name="commands"), + ) + + @classmethod + def joint_positions( + cls, + **commands: torch.Tensor, + ) -> ControlPartCommandProfile: + """Build a profile whose entries are joint-position commands.""" + return cls( + commands={ + name: JointPositionCommand(positions) + for name, positions in commands.items() + } + ) + + def snapshot(self) -> ControlPartCommandProfile: + """Return an independently owned profile snapshot.""" + return ControlPartCommandProfile(commands=self.commands) + + +def _snapshot_role_commands( + values: Mapping[str, Mapping[str, ControlCommand]], + *, + field_name: str, +) -> Mapping[str, Mapping[str, ControlCommand]]: + """Validate and freeze role-scoped invocation command overrides.""" + if not isinstance(values, Mapping): + raise TypeError(f"{field_name} must be a mapping.") + snapshots: dict[str, Mapping[str, ControlCommand]] = {} + for role, commands in values.items(): + if not isinstance(role, str) or not role.strip(): + raise ValueError(f"{field_name} roles must be non-empty strings.") + snapshots[role] = _snapshot_commands( + commands, + field_name=f"{field_name}[{role!r}]", + ) + return MappingProxyType(snapshots) + + +@dataclass(frozen=True, slots=True) +class ActionControlOverrides: + """Per-invocation semantic command overrides keyed by binding role. + + The outer keys are action roles such as ``primary``, ``source`` or + ``destination``. The inner keys are semantic command names. The engine + applies these values after resolving the role to a concrete control part, + and the resulting commands are captured in the invocation revision's + immutable planning snapshot. + """ + + manipulators: Mapping[str, Mapping[str, ControlCommand]] = field( + default_factory=dict + ) + end_effectors: Mapping[str, Mapping[str, ControlCommand]] = field( + default_factory=dict + ) + + def __post_init__(self) -> None: + object.__setattr__( + self, + "manipulators", + _snapshot_role_commands( + self.manipulators, + field_name="manipulators", + ), + ) + object.__setattr__( + self, + "end_effectors", + _snapshot_role_commands( + self.end_effectors, + field_name="end_effectors", + ), + ) + + @property + def is_empty(self) -> bool: + """Whether this invocation defines no command overrides.""" + return not self.manipulators and not self.end_effectors + + +__all__ = [ + "ActionControlOverrides", + "ControlCommand", + "ControlPartCommandProfile", + "GRASP_COMMAND", + "JointPositionCommand", + "OPEN_COMMAND", +] diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index fccb39be8..f9143d41e 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -14,296 +14,496 @@ # limitations under the License. # ---------------------------------------------------------------------------- +"""Core semantic objects and planning contract for atomic actions.""" + from __future__ import annotations -import torch from abc import ABC, abstractmethod +from copy import deepcopy from dataclasses import dataclass, field -from typing import Any, ClassVar, Generic, Mapping, TYPE_CHECKING, TypeVar +from typing import Any, ClassVar, Generic, TYPE_CHECKING + +import torch from embodichain.lab.sim.common import BatchEntity -from embodichain.utils import configclass from .affordance import Affordance +from .effects import StateDelta +from .goals import collect_scene_dependencies +from .invocation import ( + ActionInvocation, + ActionOptions, + GoalT, + OptionsT, + ResolvedActionRequest, +) +from .plans import ( + ActionPlan, + CompletionCondition, + CompletionConditionKind, + PhaseSpec, + PlannedPhase, + PlannerDiagnostics, + TimedTrajectory, +) if TYPE_CHECKING: + from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.planners import MotionGenerator + from .runtime import ActionPlanningServices + from .state import PlanningContext + from .trajectory import TrajectoryBuilder + + +def resolve_runtime_device(device: torch.device | str) -> torch.device: + """Resolve an indexless CUDA device to the active concrete GPU index. -# ============================================================================= -# ObjectSemantics -# ============================================================================= + Args: + device: PyTorch device or device string. + + Returns: + Concrete runtime device. + """ + resolved = torch.device(device) + if resolved.type == "cuda" and resolved.index is None: + return torch.device(f"cuda:{torch.cuda.current_device()}") + return resolved @dataclass class ObjectSemantics: - """Semantic information about an interaction target.""" + """Semantic and geometric information about an interaction object.""" affordance: Affordance - """Affordance data describing how the object can be interacted with.""" + """Affordance data describing supported interactions.""" geometry: dict[str, Any] - """Non-affordance geometric metadata (e.g., bounding_box). Mesh tensors live - on AntipodalAffordance, not here.""" + """Non-affordance geometric metadata.""" properties: dict[str, Any] = field(default_factory=dict) - """Physical properties: mass, friction, etc.""" + """Physical properties such as mass and friction.""" label: str = "none" - """Object category label (e.g., 'mug', 'apple').""" + """Semantic object category.""" entity: BatchEntity | None = None - """Optional reference to the simulation entity for this object.""" + """Optional simulation entity used by deterministic grounding.""" def __post_init__(self) -> None: - # Bind only the label onto the affordance for convenience. DO NOT - # alias the geometry dict — that was the footgun fixed by this redesign. + if not isinstance(self.affordance, Affordance): + raise TypeError("affordance must be an Affordance instance.") + if not isinstance(self.geometry, dict): + raise TypeError("geometry must be a dict.") + if not isinstance(self.properties, dict): + raise TypeError("properties must be a dict.") + if not isinstance(self.label, str) or not self.label: + raise ValueError("label must be a non-empty string.") self.affordance.object_label = self.label -# ============================================================================= -# Target foundation -# ============================================================================= - - -class ActionTarget: - """Open marker base for atomic-action target value objects. - - 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. - """ - - __slots__ = () - - -TargetT = TypeVar("TargetT", bound=ActionTarget) +@dataclass(frozen=True, slots=True) +class SkillDescriptor: + """Machine-readable metadata for one registered atomic skill.""" + skill_id: str + goal_type: type[Any] | tuple[type[Any], ...] + options_type: type[ActionOptions] + manipulator_roles: tuple[str, ...] = () + end_effector_roles: tuple[str, ...] = () + agent_visible: bool = True -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)}." + def __post_init__(self) -> None: + if not isinstance(self.skill_id, str) or not self.skill_id: + raise ValueError("SkillDescriptor.skill_id must be non-empty.") + goal_types = ( + self.goal_type if isinstance(self.goal_type, tuple) else (self.goal_type,) ) + if not goal_types or not all(isinstance(item, type) for item in goal_types): + raise TypeError("SkillDescriptor.goal_type must contain concrete types.") + if not isinstance(self.options_type, type) or not issubclass( + self.options_type, ActionOptions + ): + raise TypeError( + "SkillDescriptor.options_type must be an ActionOptions subclass." + ) + for field_name in ("manipulator_roles", "end_effector_roles"): + roles = tuple(getattr(self, field_name)) + if len(set(roles)) != len(roles) or not all( + isinstance(role, str) and role for role in roles + ): + raise ValueError(f"{field_name} must contain unique non-empty roles.") + object.__setattr__(self, field_name, roles) -# ``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 - - -# ============================================================================= -# World state threaded between actions -# ============================================================================= - - -@dataclass(slots=True, eq=False) -class HeldObjectState: - """State of an object currently held by the robot.""" - - semantics: ObjectSemantics - """Semantics of the held object.""" - - object_to_eef: torch.Tensor - """Batched transform from object frame to end-effector frame, shape [n_envs, 4, 4].""" - - grasp_xpos: torch.Tensor - """Batched end-effector pose used to grasp the object, shape [n_envs, 4, 4].""" - - -@dataclass(slots=True, eq=False) -class CoordinatedHeldObjectState: - """State of a single object jointly held by two robot hands.""" - - semantics: ObjectSemantics - """Semantic object currently held by the two grippers.""" - - left_object_to_eef: torch.Tensor - """Transform from object frame to left end-effector frame, shape ``[n_envs, 4, 4]``.""" - - right_object_to_eef: torch.Tensor - """Transform from object frame to right end-effector frame, shape ``[n_envs, 4, 4]``.""" - - left_grasp_xpos: torch.Tensor - """Left end-effector grasp pose for the shared object, shape ``[n_envs, 4, 4]``.""" - - right_grasp_xpos: torch.Tensor - """Right end-effector grasp pose for the shared object, shape ``[n_envs, 4, 4]``.""" +class AtomicAction(Generic[GoalT, OptionsT], ABC): + """Side-effect-free planner for one semantically meaningful robot skill. + Actions own only typed default runtime options. An + :class:`~embodichain.lab.sim.atomic_actions.engine.AtomicActionEngine` binds + its shared planning services before an action is invoked. + """ -@dataclass(slots=True, eq=False) -class WorldState: - """State the engine threads through a sequence of actions.""" + skill_id: ClassVar[str] + """Stable registry identifier for this skill.""" - last_qpos: torch.Tensor - """Robot joint positions at the start of the next action, shape [n_envs, robot.dof].""" + GoalType: ClassVar[type[Any] | tuple[type[Any], ...]] + """Concrete goal dataclass or dataclasses accepted by this skill.""" - held_objects: dict[str, HeldObjectState] = field(default_factory=dict) - """Objects held by individual control parts, keyed by control-part name.""" + OptionsType: ClassVar[type[ActionOptions]] = ActionOptions + """Concrete per-invocation runtime options accepted by this skill.""" - coordinated_held_objects: dict[tuple[str, str], CoordinatedHeldObjectState] = field( - default_factory=dict - ) - """Objects jointly held by two control parts, keyed by their ordered pair.""" + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + """Required semantic manipulator roles.""" - 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) + end_effector_roles: ClassVar[tuple[str, ...]] = () + """Required semantic end-effector roles.""" - 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) - ) + agent_visible: ClassVar[bool] = True + """Whether an Action Agent should expose this skill by default.""" - def with_updates( + def __init__( 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 - ), + default_options: OptionsT | None = None, + ) -> None: + selected_options = ( + self.OptionsType() if default_options is None else default_options ) + if not isinstance(selected_options, self.OptionsType): + raise TypeError( + f"{type(self).__name__} expects default_options of type " + f"{self.OptionsType.__name__}, got " + f"{type(selected_options).__name__}." + ) + self._default_options: OptionsT = deepcopy(selected_options) + self._planning_services: ActionPlanningServices | None = None + @property + def default_options(self) -> OptionsT: + """Return an owned copy of the action's default runtime options.""" + return deepcopy(self._default_options) -@dataclass(slots=True, eq=False) -class ActionResult: - """Return value of every AtomicAction.execute call.""" + @property + def is_bound(self) -> bool: + """Whether an engine has supplied this action's planning resources.""" + return self._planning_services is not None - success: bool | torch.Tensor - """Whether the action produced a valid full-DoF trajectory. - Can be a bool or a per-environment boolean tensor of shape (n_envs,).""" + @property + def planning_services(self) -> ActionPlanningServices: + """Return the engine-owned services borrowed by this action. - trajectory: torch.Tensor - """Full-robot trajectory, shape (n_envs, n_waypoints, robot.dof).""" + Raises: + RuntimeError: If the action has not been registered or planned by + an + :class:`~embodichain.lab.sim.atomic_actions.engine.AtomicActionEngine`. + """ + if self._planning_services is None: + raise RuntimeError( + f"Atomic action {self.skill_id!r} is not bound to an " + "AtomicActionEngine. Register it or call engine.plan_action()." + ) + return self._planning_services - next_state: WorldState - """World state to feed into the next action.""" + @property + def motion_generator(self) -> MotionGenerator: + """Return the engine-owned motion generator borrowed by this action.""" + return self.planning_services.motion_generator @property - def success_all(self) -> bool: - """True only if all environments succeeded.""" - if isinstance(self.success, torch.Tensor): - return bool(torch.all(self.success).item()) - return bool(self.success) - - def __bool__(self) -> bool: - import warnings as _w - - _w.warn( - "ActionResult bool() is deprecated; use .success_all", - DeprecationWarning, - stacklevel=2, - ) - return self.success_all + def robot(self) -> Robot: + """Return the robot associated with the owning engine.""" + return self.planning_services.robot + @property + def device(self) -> torch.device: + """Return the concrete runtime device associated with the engine.""" + return self.planning_services.device -# ============================================================================= -# Configuration base -# ============================================================================= + @property + def builder(self) -> TrajectoryBuilder: + """Return the engine-owned shared trajectory builder.""" + return self.planning_services.trajectory_builder + + def _bind(self, services: ActionPlanningServices) -> None: + """Bind engine-owned planning services exactly once.""" + if self._planning_services is services: + return + if self._planning_services is not None: + raise ValueError( + f"Atomic action {self.skill_id!r} is already bound to another " + "AtomicActionEngine." + ) + self._planning_services = services + try: + self._on_bind() + except Exception: + self._planning_services = None + raise + + def _on_bind(self) -> None: + """Initialize implementation state that depends on engine resources.""" + + @classmethod + def descriptor(cls) -> SkillDescriptor: + """Return stable metadata used by registries and Action Agent adapters.""" + return SkillDescriptor( + skill_id=cls.skill_id, + goal_type=cls.GoalType, + options_type=cls.OptionsType, + manipulator_roles=cls.manipulator_roles, + end_effector_roles=cls.end_effector_roles, + agent_visible=cls.agent_visible, + ) + def resolve_request( + self, + invocation: ActionInvocation[GoalT, OptionsT], + ) -> ResolvedActionRequest[GoalT, OptionsT]: + """Validate and snapshot an invocation through engine-owned resources. -@configclass -class ActionCfg: - """Configuration shared by all atomic actions.""" + Args: + invocation: Caller-owned invocation to resolve. - name: str = "default" - control_part: str = "arm" - interpolation_type: str = "linear" - velocity_limit: float | None = None - acceleration_limit: float | None = None - motion_source: str = "ik_interp" - """Trajectory source: 'ik_interp' (default, batched IK + linear interp) - or 'motion_gen' (batched MotionGenerator).""" + Returns: + Immutable request reused by planning and recovery replans. - def __post_init__(self) -> None: - valid_sources = {"ik_interp", "motion_gen"} - if self.motion_source not in valid_sources: + Raises: + ValueError: If the stable skill identifier does not match. + TypeError: If the goal or options type is incompatible. + KeyError: If a required binding role is missing. + """ + if invocation.skill_id != self.skill_id: + raise ValueError( + f"Invocation skill_id {invocation.skill_id!r} does not match " + f"{self.skill_id!r}." + ) + if not isinstance(invocation.goal, self.GoalType): + expected = ( + " | ".join(item.__name__ for item in self.GoalType) + if isinstance(self.GoalType, tuple) + else self.GoalType.__name__ + ) + raise TypeError( + f"Skill {self.skill_id!r} expects goal {expected}, got " + f"{type(invocation.goal).__name__}." + ) + for role in self.manipulator_roles: + invocation.binding.manipulator(role) + for role in self.end_effector_roles: + invocation.binding.end_effector(role) + options = ( + self._default_options + if invocation.skill_options is None + else invocation.skill_options + ) + if not isinstance(options, self.OptionsType): + raise TypeError( + f"Skill {self.skill_id!r} expects options " + f"{self.OptionsType.__name__}, got {type(options).__name__}." + ) + required_planner = invocation.motion_policy.planner + configured_planner_name = self.planning_services.planner_name + if required_planner is not None and required_planner != configured_planner_name: raise ValueError( - f"motion_source must be one of {sorted(valid_sources)}, " - f"but got {self.motion_source!r}." + f"Motion policy requires planner {required_planner!r}, but this " + f"action uses {configured_planner_name!r}." ) + return ResolvedActionRequest( + skill_id=invocation.skill_id, + goal=invocation.goal, + binding=self.planning_services.resolve_binding( + invocation.binding, + invocation.control_overrides, + ), + motion_policy=invocation.motion_policy, + recovery_policy=invocation.recovery_policy, + skill_options=options, + invocation_id=invocation.invocation_id, + revision=invocation.revision, + ) + def require_goal( + self, + request: ResolvedActionRequest[GoalT, OptionsT], + ) -> GoalT: + """Validate a resolved request and return its concrete goal.""" + if request.skill_id != self.skill_id: + raise ValueError( + f"Request skill_id {request.skill_id!r} does not match " + f"{self.skill_id!r}." + ) + if not isinstance(request.goal, self.GoalType): + raise TypeError( + f"Skill {self.skill_id!r} received incompatible goal " + f"{type(request.goal).__name__}." + ) + if not isinstance(request.skill_options, self.OptionsType): + raise TypeError( + f"Skill {self.skill_id!r} received incompatible options " + f"{type(request.skill_options).__name__}." + ) + for role in self.manipulator_roles: + request.binding.manipulator(role) + for role in self.end_effector_roles: + request.binding.end_effector(role) + return request.goal -# ============================================================================= -# AtomicAction ABC (slim) -# ============================================================================= + def build_plan( + self, + request: ResolvedActionRequest[GoalT, OptionsT], + context: PlanningContext, + *, + success: bool | torch.Tensor, + trajectory: TimedTrajectory | torch.Tensor, + expected_effects: StateDelta | None = None, + phase_name: str | None = None, + replannable: bool = True, + completion_kind: CompletionConditionKind = ( + CompletionConditionKind.TRAJECTORY_COMPLETE + ), + completion_tolerance: float | None = None, + diagnostics: PlannerDiagnostics | None = None, + ) -> ActionPlan: + """Build a validated single-phase plan for a primitive implementation. + Args: + request: Resolved invocation snapshot being planned. + context: Planning input used for the plan. + success: Per-environment planning success or scalar planner result. + trajectory: Full-robot timed trajectory or position tensor. + expected_effects: Symbolic effects to verify after execution. + phase_name: Optional phase name; defaults to the stable skill id. + replannable: Whether the execution runtime may replan this phase. + completion_kind: Completion condition category. + completion_tolerance: Optional numerical completion tolerance. + diagnostics: Optional retained planner diagnostics. -class AtomicAction(Generic[TargetT], ABC): - """Abstract base for atomic actions. + Returns: + Side-effect-free action plan. + """ + self.require_goal(request) + if isinstance(success, bool): + success_mask = torch.full( + (context.batch_size,), + success, + dtype=torch.bool, + device=self.device, + ) + elif isinstance(success, torch.Tensor): + success_mask = success.to(device=self.device) + if success_mask.dtype != torch.bool: + raise TypeError("Planning success must have dtype torch.bool.") + if success_mask.dim() == 0 or success_mask.shape == (1,): + success_mask = success_mask.reshape(1).expand(context.batch_size) + if success_mask.shape != (context.batch_size,): + raise ValueError( + "Planning success must have shape " + f"({context.batch_size},), got {tuple(success_mask.shape)}." + ) + success_mask = success_mask.clone() + else: + raise TypeError("Planning success must be bool or torch.Tensor.") + + if isinstance(trajectory, torch.Tensor): + timed = TimedTrajectory.from_positions( + trajectory, + env_ids=context.env_ids, + control_dt=request.motion_policy.control_dt, + ) + elif isinstance(trajectory, TimedTrajectory): + timed = trajectory + else: + raise TypeError("trajectory must be TimedTrajectory or torch.Tensor.") + if timed.batch_size != context.batch_size: + raise ValueError("Trajectory and planning context batch sizes must match.") + if timed.robot_dof != context.robot.robot_dof: + raise ValueError("Trajectory robot_dof must match the planning context.") + + if diagnostics is None: + diagnostics = PlannerDiagnostics( + backend=self.planning_services.planner_name + ) + phase = PlannedPhase( + spec=PhaseSpec( + name=phase_name or self.skill_id, + goal=request.goal, + replannable=replannable, + completion_condition=CompletionCondition( + kind=completion_kind, + tolerance=completion_tolerance, + ), + recovery_policy=request.recovery_policy, + scene_dependencies=collect_scene_dependencies(request.goal), + ), + trajectory=timed, + planned_scene_version=context.scene.version, + diagnostics=diagnostics, + ) + return ActionPlan( + skill_id=self.skill_id, + plan_success=success_mask, + phases=(phase,), + expected_effects=expected_effects or StateDelta(), + invocation_id=request.invocation_id, + invocation_revision=request.revision, + ) - Subclasses declare ``TargetType`` to advertise the concrete target dataclass - they accept. ``execute`` is the only required method; ``validate`` has been - dropped from the contract in this redesign. - """ + def failed_plan( + self, + request: ResolvedActionRequest[GoalT, OptionsT], + context: PlanningContext, + *, + message: str | None = None, + ) -> ActionPlan: + """Build a failed empty plan without changing task state. - TargetType: ClassVar[type[ActionTarget] | tuple[type[ActionTarget], ...]] - """Concrete target dataclass or dataclasses accepted by ``execute``.""" + Args: + request: Resolved invocation that failed to plan. + context: Planning input used for the attempt. + message: Optional diagnostic message. - def __init__( - self, - motion_generator: MotionGenerator, - cfg: ActionCfg | None = None, - ) -> None: - self.motion_generator = motion_generator - self.cfg = cfg if cfg is not None else ActionCfg() - self.robot = motion_generator.robot - self.device = self.robot.device - self.control_part = self.cfg.control_part + Returns: + Failed action plan with an empty phase trajectory. + """ + return self.build_plan( + request, + context, + success=torch.zeros( + context.batch_size, dtype=torch.bool, device=self.device + ), + trajectory=TimedTrajectory.empty( + batch_size=context.batch_size, + robot_dof=context.robot.robot_dof, + device=self.device, + env_ids=context.env_ids, + ), + replannable=True, + diagnostics=PlannerDiagnostics( + backend=self.planning_services.planner_name, + messages=(() if message is None else (message,)), + ), + ) @abstractmethod - def execute(self, target: TargetT, state: WorldState) -> ActionResult: - """Plan and return a full-DoF trajectory for this action. + def plan( + self, + request: ResolvedActionRequest[GoalT, OptionsT], + context: PlanningContext, + ) -> ActionPlan: + """Plan one invocation without stepping simulation or committing state. Args: - target: Typed target dataclass; must be an instance of ``self.TargetType``. - state: World state inherited from the previous action (or the engine seed). + request: Immutable, typed, and embodiment-resolved action request. + context: Latest observed robot, task, and scene state. Returns: - ActionResult with the planned trajectory and the successor world state. + Scene-bound action plan with expected, uncommitted effects. """ __all__ = [ - "ActionTarget", - "ActionCfg", - "ActionResult", "AtomicAction", - "CoordinatedHeldObjectState", - "HeldObjectState", "ObjectSemantics", - "Target", - "TargetT", - "WorldState", + "SkillDescriptor", + "resolve_runtime_device", ] diff --git a/embodichain/lab/sim/atomic_actions/effects.py b/embodichain/lab/sim/atomic_actions/effects.py new file mode 100644 index 000000000..f9c6507aa --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/effects.py @@ -0,0 +1,284 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Expected symbolic effects produced by side-effect-free action planning.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Mapping + +import torch + +from .state import ( + CoordinatedHeldObjectState, + HeldObjectState, + TaskState, + _normalize_coordinated_held, + _normalize_held, + _normalize_mask, +) + + +def _with_held_mask( + value: HeldObjectState, + env_mask: torch.Tensor, +) -> HeldObjectState: + """Copy a held-object relation with a replacement mask.""" + return HeldObjectState( + semantics=value.semantics, + object_to_eef=value.object_to_eef, + grasp_xpos=value.grasp_xpos, + env_mask=env_mask, + ) + + +def _with_coordinated_mask( + value: CoordinatedHeldObjectState, + env_mask: torch.Tensor, +) -> CoordinatedHeldObjectState: + """Copy a coordinated held-object relation with a replacement mask.""" + return CoordinatedHeldObjectState( + semantics=value.semantics, + left_object_to_eef=value.left_object_to_eef, + right_object_to_eef=value.right_object_to_eef, + left_grasp_xpos=value.left_grasp_xpos, + right_grasp_xpos=value.right_grasp_xpos, + env_mask=env_mask, + ) + + +def _merge_held( + previous: HeldObjectState | None, + candidate: HeldObjectState | None, + update_mask: torch.Tensor, +) -> HeldObjectState | None: + """Apply one optional held-object update per environment.""" + if previous is None and candidate is None: + return None + if previous is None: + assert candidate is not None and candidate.env_mask is not None + env_mask = candidate.env_mask & update_mask + return _with_held_mask(candidate, env_mask) if env_mask.any() else None + assert previous.env_mask is not None + if candidate is None: + env_mask = previous.env_mask & ~update_mask + return _with_held_mask(previous, env_mask) if env_mask.any() else None + assert candidate.env_mask is not None + + previous_retained = bool((previous.env_mask & ~update_mask).any().item()) + candidate_applied = bool((candidate.env_mask & update_mask).any().item()) + if ( + previous_retained + and candidate_applied + and previous.semantics is not candidate.semantics + ): + raise ValueError( + "Cannot merge different held-object semantics for one resource " + "across environments." + ) + env_mask = torch.where(update_mask, candidate.env_mask, previous.env_mask) + if not env_mask.any(): + return None + selector = update_mask[:, None, None] + return HeldObjectState( + semantics=candidate.semantics if candidate_applied else previous.semantics, + object_to_eef=torch.where( + selector, candidate.object_to_eef, previous.object_to_eef + ), + grasp_xpos=torch.where(selector, candidate.grasp_xpos, previous.grasp_xpos), + env_mask=env_mask, + ) + + +def _merge_coordinated( + previous: CoordinatedHeldObjectState | None, + candidate: CoordinatedHeldObjectState | None, + update_mask: torch.Tensor, +) -> CoordinatedHeldObjectState | None: + """Apply one optional coordinated relation update per environment.""" + if previous is None and candidate is None: + return None + if previous is None: + assert candidate is not None and candidate.env_mask is not None + env_mask = candidate.env_mask & update_mask + return _with_coordinated_mask(candidate, env_mask) if env_mask.any() else None + assert previous.env_mask is not None + if candidate is None: + env_mask = previous.env_mask & ~update_mask + return _with_coordinated_mask(previous, env_mask) if env_mask.any() else None + assert candidate.env_mask is not None + + previous_retained = bool((previous.env_mask & ~update_mask).any().item()) + candidate_applied = bool((candidate.env_mask & update_mask).any().item()) + if ( + previous_retained + and candidate_applied + and previous.semantics is not candidate.semantics + ): + raise ValueError( + "Cannot merge different coordinated held-object semantics for one " + "resource pair across environments." + ) + env_mask = torch.where(update_mask, candidate.env_mask, previous.env_mask) + if not env_mask.any(): + return None + selector = update_mask[:, None, None] + return CoordinatedHeldObjectState( + semantics=candidate.semantics if candidate_applied else previous.semantics, + left_object_to_eef=torch.where( + selector, candidate.left_object_to_eef, previous.left_object_to_eef + ), + right_object_to_eef=torch.where( + selector, candidate.right_object_to_eef, previous.right_object_to_eef + ), + left_grasp_xpos=torch.where( + selector, candidate.left_grasp_xpos, previous.left_grasp_xpos + ), + right_grasp_xpos=torch.where( + selector, candidate.right_grasp_xpos, previous.right_grasp_xpos + ), + env_mask=env_mask, + ) + + +@dataclass(frozen=True, slots=True, eq=False) +class StateDelta: + """Expected task-state changes that require post-execution verification. + + A mapping value of ``None`` removes the corresponding relation. Planning + only declares this delta; an execution runtime applies it after verifying + the semantic effect for the successful environment rows. + """ + + held_object_updates: Mapping[str, HeldObjectState | None] = field( + default_factory=dict + ) + """Per-resource attachment replacements or removals.""" + + coordinated_held_object_updates: Mapping[ + tuple[str, str], CoordinatedHeldObjectState | None + ] = field(default_factory=dict) + """Per-resource-pair coordinated attachment replacements or removals.""" + + def __post_init__(self) -> None: + held = dict(self.held_object_updates) + coordinated = dict(self.coordinated_held_object_updates) + for resource, value in held.items(): + if not isinstance(resource, str) or not resource: + raise ValueError( + "held_object_updates keys must be non-empty resource names." + ) + if value is not None and not isinstance(value, HeldObjectState): + raise TypeError( + "held_object_updates values must be HeldObjectState or None." + ) + for resources, value in coordinated.items(): + if ( + not isinstance(resources, tuple) + or len(resources) != 2 + or not all(isinstance(item, str) and item for item in resources) + ): + raise ValueError( + "coordinated_held_object_updates keys must be resource pairs." + ) + if value is not None and not isinstance(value, CoordinatedHeldObjectState): + raise TypeError( + "coordinated_held_object_updates values must be " + "CoordinatedHeldObjectState or None." + ) + object.__setattr__(self, "held_object_updates", MappingProxyType(held)) + object.__setattr__( + self, + "coordinated_held_object_updates", + MappingProxyType(coordinated), + ) + + @property + def is_empty(self) -> bool: + """Whether this delta declares no symbolic state changes.""" + return not self.held_object_updates and not self.coordinated_held_object_updates + + def apply( + self, + state: TaskState, + update_mask: torch.Tensor, + ) -> TaskState: + """Apply expected effects to selected environment rows. + + This operation is used for hypothetical state propagation while + compiling a sequence. A runtime must apply the same delta only after + effect verification. + + Args: + state: Input task state. + update_mask: Successful and verified rows, shape ``(n_envs,)``. + + Returns: + New task state with masked updates. + """ + if not isinstance(state, TaskState): + raise TypeError("state must be a TaskState.") + mask = _normalize_mask( + update_mask, + batch_size=state.batch_size, + device=state.device, + name="update_mask", + ) + held = dict(state.held_objects) + for resource, candidate in self.held_object_updates.items(): + normalized = ( + None + if candidate is None + else _normalize_held( + candidate, + batch_size=state.batch_size, + device=state.device, + ) + ) + merged = _merge_held(held.get(resource), normalized, mask) + if merged is None: + held.pop(resource, None) + else: + held[resource] = merged + + coordinated = dict(state.coordinated_held_objects) + for resources, candidate in self.coordinated_held_object_updates.items(): + normalized = ( + None + if candidate is None + else _normalize_coordinated_held( + candidate, + batch_size=state.batch_size, + device=state.device, + ) + ) + merged = _merge_coordinated(coordinated.get(resources), normalized, mask) + if merged is None: + coordinated.pop(resources, None) + else: + coordinated[resources] = merged + + return TaskState( + batch_size=state.batch_size, + device=state.device, + held_objects=held, + coordinated_held_objects=coordinated, + ) + + +__all__ = ["StateDelta"] diff --git a/embodichain/lab/sim/atomic_actions/engine.py b/embodichain/lab/sim/atomic_actions/engine.py index b1b0b3915..38a5190c0 100644 --- a/embodichain/lab/sim/atomic_actions/engine.py +++ b/embodichain/lab/sim/atomic_actions/engine.py @@ -14,159 +14,444 @@ # limitations under the License. # ---------------------------------------------------------------------------- +"""Registry and offline compiler for side-effect-free atomic actions.""" + from __future__ import annotations -import torch -from typing import Iterable, TYPE_CHECKING +from typing import Iterable, Mapping, TYPE_CHECKING -from embodichain.utils import logger +import torch -from .core import ( - ActionTarget, - ActionResult, - AtomicAction, - WorldState, -) +from .core import AtomicAction +from .control import ControlPartCommandProfile +from .invocation import ActionInvocation, ResolvedActionRequest +from .plans import ActionPlan, CompiledTrajectory, TimedTrajectory +from .runtime import ActionPlanningServices +from .state import PlanningContext, RobotObservation, SceneSnapshot, TaskState if TYPE_CHECKING: + from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.planners import MotionGenerator + from .execution import ExecutionSession -# ============================================================================= -# Global action registry (kept for third-party extensions) -# ============================================================================= +_global_extension_registry: dict[str, type[AtomicAction]] = {} -_global_action_registry: dict[str, type[AtomicAction]] = {} +def register_action(action_class: type[AtomicAction]) -> None: + """Register an extension action type for process-wide discovery. -def _target_type_name(target_type: type | tuple[type, ...]) -> str: - """Return a readable name for one accepted target type or a tuple of them.""" - if isinstance(target_type, tuple): - return " | ".join(t.__name__ for t in target_type) - return target_type.__name__ + This catalog does not bind the type to an engine or automatically load it. + Built-in types live in ``BUILTIN_ACTION_TYPES`` and are loaded separately + by each :class:`AtomicActionEngine`. + Args: + action_class: Concrete :class:`AtomicAction` subclass. -def register_action(name: str, action_class: type[AtomicAction]) -> None: - """Register a custom AtomicAction subclass globally under ``name``.""" - _global_action_registry[name] = action_class - + Raises: + TypeError: If ``action_class`` is not an AtomicAction subclass. + ValueError: If another class already owns the same skill identifier. + """ + if not isinstance(action_class, type) or not issubclass(action_class, AtomicAction): + raise TypeError("action_class must be an AtomicAction subclass.") + descriptor = action_class.descriptor() + existing = _global_extension_registry.get(descriptor.skill_id) + if existing is not None and existing is not action_class: + raise ValueError( + f"Skill id {descriptor.skill_id!r} is already registered by " + f"{existing.__name__}." + ) + _global_extension_registry[descriptor.skill_id] = action_class -def unregister_action(name: str) -> None: - """Remove a previously-registered action class. No-op if absent.""" - _global_action_registry.pop(name, None) +def unregister_action(skill_id: str) -> None: + """Remove a globally discoverable extension action type if present. -def get_registered_actions() -> dict[str, type[AtomicAction]]: - """Return a copy of the global action-class registry.""" - return _global_action_registry.copy() + Args: + skill_id: Stable registered skill identifier. + """ + _global_extension_registry.pop(skill_id, None) -# ============================================================================= -# AtomicActionEngine -# ============================================================================= +def get_registered_actions() -> dict[str, type[AtomicAction]]: + """Return a copy of the process-wide extension action-type registry.""" + return dict(_global_extension_registry) class AtomicActionEngine: - """Sequences typed atomic actions while threading WorldState through them.""" + """Own planning resources and coordinate side-effect-free atomic actions.""" + + def __init__( + self, + motion_generator: MotionGenerator, + control_profiles: Mapping[str, ControlPartCommandProfile] | None = None, + *, + load_builtins: bool = True, + ) -> None: + """Initialize one engine and bind its built-in action implementations. - def __init__(self, motion_generator: MotionGenerator) -> None: - self.motion_generator = motion_generator - self.robot = motion_generator.robot - self.device = motion_generator.device + Args: + motion_generator: Engine-owned motion-generation backend. + control_profiles: Semantic commands keyed by robot control-part name. + load_builtins: Whether to instantiate and register every built-in + action. Disable this for isolated tests or fully custom engines. + """ + self._planning_services = ActionPlanningServices( + motion_generator, + control_profiles=control_profiles, + ) self._actions: dict[str, AtomicAction] = {} + if load_builtins: + self._load_builtin_actions() + + @property + def motion_generator(self) -> MotionGenerator: + """Return the single motion generator owned by this engine.""" + return self._planning_services.motion_generator + + @property + def robot(self) -> Robot: + """Return the robot controlled by this engine.""" + return self._planning_services.robot + + @property + def device(self) -> torch.device: + """Return the concrete planning device used by this engine.""" + return self._planning_services.device + + @property + def planning_services(self) -> ActionPlanningServices: + """Engine-owned resources shared by every bound atomic action.""" + return self._planning_services + + @property + def control_profiles(self) -> Mapping[str, ControlPartCommandProfile]: + """Semantic command profiles registered for robot control parts.""" + return self._planning_services.control_profiles @property def actions(self) -> dict[str, AtomicAction]: - """Registered actions keyed by name (read-only copy).""" + """Registered action instances keyed by stable skill identifier.""" return dict(self._actions) - 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, + def register(self, action: AtomicAction, *, replace: bool = False) -> None: + """Register one action instance using its descriptor. + + Args: + action: Configured action instance. + replace: Whether to replace an implementation already registered + under the same stable skill identifier. Replacement is always + explicit so extensions cannot silently shadow built-ins. + + Raises: + TypeError: If ``action`` is not an AtomicAction. + ValueError: If it belongs to another engine or its skill identifier + conflicts with an existing action. + """ + if not isinstance(action, AtomicAction): + raise TypeError("action must be an AtomicAction instance.") + descriptor = action.descriptor() + existing = self._actions.get(descriptor.skill_id) + if existing is not None and existing is not action and not replace: + raise ValueError( + f"Skill id {descriptor.skill_id!r} is already registered in this engine." ) - key = name if name is not None else action.cfg.name - self._actions[key] = action + action._bind(self._planning_services) + self._actions[descriptor.skill_id] = action + + def _load_builtin_actions(self) -> None: + """Create and bind fresh built-in action instances for this engine.""" + # Import lazily to keep the engine/core dependency independent from the + # concrete primitive modules and to avoid package import cycles. + from .primitives import BUILTIN_ACTION_TYPES + + for action_type in BUILTIN_ACTION_TYPES: + self.register(action_type()) - def run( + def plan_action( self, - steps: Iterable[tuple[str, ActionTarget]], - state: WorldState | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, WorldState]: - """Run a sequence of named actions, threading WorldState through. + action: AtomicAction, + invocation: ActionInvocation, + context: PlanningContext, + ) -> ActionPlan: + """Plan with a configured action using this engine's resources. + + Unlike :meth:`plan`, the supplied action does not need to be in the + skill registry. This is an advanced extension and testing escape hatch; + built-in parameter variants should use ``ActionInvocation.skill_options`` + with the engine's registered implementation. Args: - steps: Iterable of ``(action_name, typed_target)`` pairs. - state: Initial world state. If None, seeded from ``robot.get_qpos()``. + action: Configured action implementation to invoke. + invocation: Grounded request matching the action's skill identifier. + context: Latest measured planning state. Returns: - ``(success, concatenated_full_dof_trajectory, final_state)``. + Validated side-effect-free action plan. + + Raises: + TypeError: If ``action`` is not an :class:`AtomicAction`. + ValueError: If the action, invocation, context, or plan is invalid. + """ + if not isinstance(action, AtomicAction): + raise TypeError("action must be an AtomicAction instance.") + self._validate_context(context) + action._bind(self._planning_services) + request = action.resolve_request(invocation) + plan = action.plan(request, context) + self._validate_plan(plan, context, request) + return plan + + def resolve( + self, + invocation: ActionInvocation, + ) -> ResolvedActionRequest: + """Resolve one registered invocation into an engine-owned snapshot. + + The returned request owns its policy and skill-option values. Closed-loop + recovery reuses the same request so a replan cannot observe later + mutations of caller-owned configuration objects. - ``success`` is a ``(B,)`` boolean tensor indicating which - environments completed every step. Failed environments hold their - last successful joint position in both ``full_traj`` and - ``final_state.last_qpos`` for the remainder of the sequence. + Args: + invocation: Grounded request for a registered skill. - An empty ``steps`` iterable is a successful no-op returning an - empty trajectory and the seed state. + Returns: + Validated and embodiment-resolved request snapshot. + + Raises: + KeyError: If the invocation references an unregistered skill. + """ + action = self._actions.get(invocation.skill_id) + if action is None: + raise KeyError( + f"No atomic action registered for skill {invocation.skill_id!r}." + ) + return action.resolve_request(invocation) + + def plan_request( + self, + request: ResolvedActionRequest, + context: PlanningContext | None = None, + ) -> ActionPlan: + """Plan an already-resolved request without rebuilding its snapshot. + + This is the planning entry point used by execution recovery. Callers + normally use :meth:`plan`, while an execution session resolves once and + calls this method for every replan. + + Args: + request: Immutable request previously returned by :meth:`resolve`. + context: Optional latest planning state; captured when omitted. + + Returns: + Validated side-effect-free action plan. """ - if state is None: - state = WorldState(last_qpos=self.robot.get_qpos().clone()) - - b = state.last_qpos.shape[0] - full_traj = torch.empty( - (b, 0, self.robot.dof), - dtype=torch.float32, - device=self.device, + if not isinstance(request, ResolvedActionRequest): + raise TypeError("request must be a ResolvedActionRequest.") + action = self._actions.get(request.skill_id) + if action is None: + raise KeyError( + f"No atomic action registered for skill {request.skill_id!r}." + ) + current = self.initial_context() if context is None else context + self._validate_context(current) + plan = action.plan(request, current) + self._validate_plan(plan, current, request) + return plan + + def plan( + self, + invocation: ActionInvocation, + context: PlanningContext | None = None, + ) -> ActionPlan: + """Plan one registered invocation through the engine-owned backend. + + Args: + invocation: Grounded request for a registered skill. + context: Optional latest planning state; captured when omitted. + + Returns: + Validated action plan. + + Raises: + KeyError: If the invocation references an unregistered skill. + """ + current = self.initial_context() if context is None else context + request = self.resolve(invocation) + return self.plan_request(request, current) + + def initial_context( + self, + *, + task: TaskState | None = None, + scene: SceneSnapshot | None = None, + timestamp: float = 0.0, + ) -> PlanningContext: + """Capture the robot state needed to start offline compilation. + + Args: + task: Optional symbolic task state; an empty state is used otherwise. + scene: Optional scene snapshot; an empty snapshot is used otherwise. + timestamp: Timestamp assigned to the captured robot observation. + + Returns: + Planning context containing owned robot tensors. + """ + qpos = self.robot.get_qpos().to(self.device).clone() + qvel_value = None + get_qvel = getattr(self.robot, "get_qvel", None) + if callable(get_qvel): + candidate = get_qvel() + if isinstance(candidate, torch.Tensor): + qvel_value = candidate.to(self.device) + qvel = torch.zeros_like(qpos) if qvel_value is None else qvel_value + batch_size = int(qpos.shape[0]) + if task is None: + task = TaskState.empty(batch_size=batch_size, device=self.device) + if scene is None: + scene = SceneSnapshot.empty() + return PlanningContext( + robot=RobotObservation(timestamp=timestamp, qpos=qpos, qvel=qvel), + task=task, + scene=scene, + env_ids=torch.arange(batch_size, dtype=torch.long, device=self.device), ) - alive = torch.ones(b, dtype=torch.bool, device=self.device) - - for name, target in steps: - if name not in self._actions: - logger.log_error(f"No action registered under name '{name}'", KeyError) - action = self._actions[name] - if not isinstance(target, action.TargetType): - logger.log_error( - f"Action '{name}' expects target of type " - f"{_target_type_name(action.TargetType)}, got {type(target).__name__}", - TypeError, - ) + + def compile( + self, + invocations: Iterable[ActionInvocation], + context: PlanningContext | None = None, + ) -> CompiledTrajectory: + """Compile a static sequence of grounded invocations. + + Planning is side-effect free. Expected effects are applied only to the + returned hypothetical ``projected_context`` so following actions can be + checked against the expected state. No simulator or observed task state + is mutated. + + Args: + invocations: Grounded action requests in execution order. + context: Optional initial planning context captured by the caller. + + Returns: + Concatenated timed trajectory, individual plans, and projected state. + + Raises: + KeyError: If an invocation references an unregistered skill. + ValueError: If context, plan, or trajectory dimensions are incompatible. + """ + if context is None: + context = self.initial_context() + self._validate_context(context) + + alive = torch.ones(context.batch_size, dtype=torch.bool, device=self.device) + plans: list[ActionPlan] = [] + trajectories: list[TimedTrajectory] = [] + projected = context + + for invocation in invocations: if not alive.any(): - # All envs dead: no further motion to plan. break - prev_last_qpos = state.last_qpos.clone() - result: ActionResult = action.execute(target, state) - step_success = ( - result.success - if isinstance(result.success, torch.Tensor) - else torch.tensor(bool(result.success), device=self.device) + previous_qpos = projected.robot.qpos + plan = self.plan(invocation, projected) + step_success = alive & plan.plan_success.to(self.device) + trajectory = plan.trajectory.hold_rows(step_success, previous_qpos) + plans.append(plan) + trajectories.append(trajectory) + + candidate_qpos = ( + trajectory.positions[:, -1] + if trajectory.waypoint_count > 0 + else previous_qpos + ) + next_qpos = torch.where( + step_success[:, None], candidate_qpos, previous_qpos + ) + next_task = plan.expected_effects.apply(projected.task, step_success) + projected = projected.project(qpos=next_qpos, task=next_task) + alive = step_success + + compiled = TimedTrajectory.concatenate(trajectories, empty_like=context) + return CompiledTrajectory( + plan_success=alive, + trajectory=compiled, + action_plans=tuple(plans), + projected_context=projected, + ) + + def start( + self, + invocations: Iterable[ActionInvocation], + context: PlanningContext | None = None, + ) -> ExecutionSession: + """Start closed-loop execution for a grounded invocation sequence. + + Args: + invocations: Grounded action requests in execution order. + context: Initial measured state and scene snapshot. The engine + captures one when omitted. + + Returns: + Stateful execution session advanced by ``session.tick(...)``. + """ + from .execution import ExecutionSession + + initial = self.initial_context() if context is None else context + return ExecutionSession(self, tuple(invocations), initial) + + def _validate_context(self, context: PlanningContext) -> None: + """Validate an externally supplied planning context.""" + if context.robot.robot_dof != self.robot.dof: + raise ValueError( + "PlanningContext robot_dof must match the engine robot, " + f"got {context.robot.robot_dof} and {self.robot.dof}." ) - step_success = step_success.to(self.device) - alive = alive & step_success - # Failed envs freeze at their last successful qpos for this step's trajectory. - traj = result.trajectory - held_rows = prev_last_qpos.unsqueeze(1).repeat(1, traj.shape[1], 1) - traj = torch.where(alive[:, None, None], traj, held_rows) - full_traj = torch.cat([full_traj, traj], dim=1) - state = result.next_state.with_updates( - last_qpos=torch.where( - alive[:, None], result.next_state.last_qpos, prev_last_qpos - ) + robot_qpos = self.robot.get_qpos() + if context.batch_size != int(robot_qpos.shape[0]): + raise ValueError( + "PlanningContext batch size must match the engine robot, " + f"got {context.batch_size} and {robot_qpos.shape[0]}." ) + if context.robot.qpos.device != self.device: + raise ValueError("PlanningContext and engine must share a device.") - return alive, full_traj, state + def _validate_plan( + self, + plan: ActionPlan, + context: PlanningContext, + request: ResolvedActionRequest, + ) -> None: + """Validate one action result before it is composed.""" + if plan.skill_id != request.skill_id: + raise ValueError( + "ActionPlan.skill_id must match its request, " + f"got {plan.skill_id!r} and {request.skill_id!r}." + ) + if plan.invocation_id != request.invocation_id: + raise ValueError( + "ActionPlan.invocation_id must preserve the invocation correlation id." + ) + if plan.invocation_revision != request.revision: + raise ValueError( + "ActionPlan.invocation_revision must preserve the request revision." + ) + trajectory = plan.trajectory + if trajectory.batch_size != context.batch_size: + raise ValueError("Action plan batch size does not match the context.") + if trajectory.robot_dof != self.robot.dof: + raise ValueError("Action plan robot_dof does not match the engine robot.") + if trajectory.positions.device != self.device: + raise ValueError("Action plan and engine must share a device.") + if not torch.equal(trajectory.env_ids, context.env_ids): + raise ValueError("Action plan and context must share ordered env_ids.") + if any( + phase.planned_scene_version != context.scene.version + for phase in plan.phases + ): + raise ValueError( + "Every action phase must record the planning scene version." + ) __all__ = [ diff --git a/embodichain/lab/sim/atomic_actions/execution.py b/embodichain/lab/sim/atomic_actions/execution.py new file mode 100644 index 000000000..748cebed8 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/execution.py @@ -0,0 +1,756 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Closed-loop execution session for dynamic atomic-action plans.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING + +import torch + +from .invocation import ActionInvocation, ResolvedActionRequest +from .plans import ActionPlan, PlannedPhase +from .state import EntityState, PlanningContext, SceneSnapshot, TaskState + +if TYPE_CHECKING: + from .engine import AtomicActionEngine + + +class ExecutionStatus(str, Enum): + """Lifecycle status of an execution session.""" + + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + + +class ExecutionEventKind(str, Enum): + """Structured event categories emitted by :meth:`ExecutionSession.tick`.""" + + ACTION_PLANNED = "action_planned" + INVOCATION_REVISED = "invocation_revised" + REPLANNED = "replanned" + TRACKING_ERROR = "tracking_error" + DYNAMIC_GOAL_CHANGED = "dynamic_goal_changed" + PHASE_TIMEOUT = "phase_timeout" + PHASE_COMPLETED = "phase_completed" + EFFECT_VERIFICATION_REQUIRED = "effect_verification_required" + ACTION_RETRY = "action_retry" + ACTION_COMPLETED = "action_completed" + RECOVERY_EXHAUSTED = "recovery_exhausted" + SESSION_COMPLETED = "session_completed" + + +@dataclass(frozen=True, slots=True, eq=False) +class ExecutionEvent: + """One timestamped execution or recovery event.""" + + kind: ExecutionEventKind + timestamp: float + skill_id: str | None + invocation_id: str | None + invocation_revision: int + invocation_index: int + env_mask: torch.Tensor + message: str = "" + + def __post_init__(self) -> None: + if self.timestamp < 0.0: + raise ValueError("ExecutionEvent.timestamp must be non-negative.") + if self.invocation_index < 0: + raise ValueError("invocation_index must be non-negative.") + if self.invocation_revision < 0: + raise ValueError("invocation_revision must be non-negative.") + if self.env_mask.dtype != torch.bool or self.env_mask.dim() != 1: + raise ValueError("ExecutionEvent.env_mask must be a 1D bool tensor.") + object.__setattr__(self, "env_mask", self.env_mask.clone()) + + +@dataclass(frozen=True, slots=True, eq=False) +class JointCommand: + """Full-robot command produced by one session tick.""" + + positions: torch.Tensor + velocities: torch.Tensor | None + active_mask: torch.Tensor + env_ids: torch.Tensor + + def __post_init__(self) -> None: + if self.positions.dim() != 2: + raise ValueError("JointCommand.positions must have shape (B, robot_dof).") + if ( + self.velocities is not None + and self.velocities.shape != self.positions.shape + ): + raise ValueError("JointCommand.velocities must match positions shape.") + if self.active_mask.dtype != torch.bool or self.active_mask.shape != ( + self.positions.shape[0], + ): + raise ValueError("JointCommand.active_mask must be bool with shape (B,).") + if self.env_ids.dtype != torch.long or self.env_ids.shape != ( + self.positions.shape[0], + ): + raise ValueError("JointCommand.env_ids must be int64 with shape (B,).") + if self.active_mask.device != self.positions.device: + raise ValueError("JointCommand tensors must share a device.") + if self.env_ids.device != self.positions.device: + raise ValueError("JointCommand tensors must share a device.") + object.__setattr__(self, "positions", self.positions.clone()) + if self.velocities is not None: + object.__setattr__(self, "velocities", self.velocities.clone()) + object.__setattr__(self, "active_mask", self.active_mask.clone()) + object.__setattr__(self, "env_ids", self.env_ids.clone()) + + +@dataclass(frozen=True, slots=True, eq=False) +class ExecutionTick: + """Result returned after one closed-loop execution update.""" + + status: ExecutionStatus + eligible_mask: torch.Tensor + command: JointCommand | None + events: tuple[ExecutionEvent, ...] + task_state: TaskState + + def __post_init__(self) -> None: + if self.eligible_mask.dtype != torch.bool or self.eligible_mask.dim() != 1: + raise ValueError("eligible_mask must be a 1D bool tensor.") + object.__setattr__(self, "eligible_mask", self.eligible_mask.clone()) + object.__setattr__(self, "events", tuple(self.events)) + + +class ExecutionSession: + """Execute grounded invocations incrementally with bounded local recovery. + + The session never steps a simulator itself. Each :meth:`tick` consumes the + latest observation and scene snapshot and emits at most one full-robot + command. Expected symbolic effects are committed only after the caller + supplies ``effect_success`` for a non-empty :class:`StateDelta`. + """ + + def __init__( + self, + engine: AtomicActionEngine, + invocations: tuple[ActionInvocation, ...], + context: PlanningContext, + ) -> None: + if not invocations: + raise ValueError("ExecutionSession requires at least one invocation.") + engine._validate_context(context) + self._engine = engine + self._requests: tuple[ResolvedActionRequest, ...] = tuple( + engine.resolve(invocation) for invocation in invocations + ) + self._task_state = context.task + self._context = context + self._invocation_index = 0 + self._phase_index = 0 + self._waypoint_index = 0 + self._plan: ActionPlan | None = None + self._planned_scene = context.scene + self._phase_started_at = context.robot.timestamp + self._last_command: torch.Tensor | None = None + self._last_command_mask = torch.zeros( + context.batch_size, dtype=torch.bool, device=context.robot.qpos.device + ) + self._eligible = torch.ones_like(self._last_command_mask) + self._pending = self._eligible.clone() + self._action_retries = torch.zeros( + context.batch_size, dtype=torch.long, device=context.robot.qpos.device + ) + self._replans = torch.zeros_like(self._action_retries) + self._effect_wait_emitted = False + self._status = ExecutionStatus.RUNNING + self._queued_events: list[ExecutionEvent] = [] + self._plan_current(context, ExecutionEventKind.ACTION_PLANNED) + + @property + def status(self) -> ExecutionStatus: + """Current session status.""" + return self._status + + @property + def eligible_mask(self) -> torch.Tensor: + """Rows still eligible to complete the full invocation sequence. + + This is deliberately not named ``success_mask``: while the session is + running, eligibility does not imply that execution or semantic effects + have succeeded. + """ + return self._eligible.clone() + + @property + def task_state(self) -> TaskState: + """Verified symbolic task state accumulated by this session.""" + return self._task_state + + def revise_current(self, invocation: ActionInvocation) -> None: + """Replace and replan the current invocation with a newer revision. + + The replacement is resolved into a new immutable request snapshot from + the latest observation. Retry and replan budgets restart for the new + revision, while verified task state, the current batch barrier, and + per-environment eligibility are preserved. Ordinary recovery replans + continue to reuse this snapshot until another explicit revision. + + Args: + invocation: Grounded replacement for the currently active skill. + Its ``revision`` must be strictly greater than the active one, + and its ``skill_id`` and ``invocation_id`` must identify the + same logical call. + + Raises: + TypeError: If ``invocation`` is not an ActionInvocation. + RuntimeError: If the session is no longer running. + ValueError: If the replacement identifies another invocation or + does not advance the revision. + """ + if not isinstance(invocation, ActionInvocation): + raise TypeError("invocation must be an ActionInvocation.") + if self._status is not ExecutionStatus.RUNNING: + raise RuntimeError("Only a running execution session can be revised.") + current = self._requests[self._invocation_index] + if invocation.skill_id != current.skill_id: + raise ValueError( + f"Revision skill_id {invocation.skill_id!r} does not match " + f"the active skill {current.skill_id!r}." + ) + if invocation.invocation_id != current.invocation_id: + raise ValueError( + "Revision invocation_id must match the active invocation_id." + ) + if invocation.revision <= current.revision: + raise ValueError( + f"Revision must advance beyond {current.revision}, got " + f"{invocation.revision}." + ) + + replacement = self._engine.resolve(invocation) + replacement_plan = self._engine.plan_request(replacement, self._context) + requests = list(self._requests) + requests[self._invocation_index] = replacement + self._requests = tuple(requests) + self._phase_index = 0 + self._waypoint_index = 0 + self._action_retries.zero_() + self._replans.zero_() + self._install_plan( + replacement_plan, + self._context, + ExecutionEventKind.INVOCATION_REVISED, + ) + + def tick( + self, + context: PlanningContext, + *, + effect_success: torch.Tensor | None = None, + ) -> ExecutionTick: + """Advance execution by one observation/command cycle. + + Args: + context: Latest measured robot and versioned scene state. Its task + state is replaced by the session's verified task state. + effect_success: Optional per-environment semantic-effect verification + for an action waiting at its terminal waypoint. + + Returns: + Status, optional command, events, and current verified task state. + """ + self._engine._validate_context(context) + if context.robot.timestamp < self._context.robot.timestamp: + raise ValueError("Execution tick timestamps must be monotonic.") + if context.scene.timestamp < self._context.scene.timestamp: + raise ValueError("Scene snapshot timestamps must be monotonic.") + if context.scene.version < self._context.scene.version: + raise ValueError("Scene snapshot versions must be monotonic.") + if not torch.equal(context.env_ids, self._context.env_ids): + raise ValueError("Execution tick env_ids must remain stable and ordered.") + self._context = PlanningContext( + robot=context.robot, + task=self._task_state, + scene=context.scene, + env_ids=context.env_ids, + ) + events = self._drain_events() + if self._status is not ExecutionStatus.RUNNING: + return self._tick_result(command=None, events=events) + + assert self._plan is not None + phase = self._current_phase() + execution_mask = self._pending & self._plan.plan_success + recovery_events = self._recover_if_needed(phase, execution_mask) + events.extend(recovery_events) + if self._status is not ExecutionStatus.RUNNING: + return self._tick_result(command=None, events=events) + if recovery_events and any( + event.kind + in { + ExecutionEventKind.REPLANNED, + ExecutionEventKind.RECOVERY_EXHAUSTED, + } + for event in recovery_events + ): + assert self._plan is not None + phase = self._current_phase() + execution_mask = self._pending & self._plan.plan_success + + trajectory = phase.trajectory + if self._waypoint_index < trajectory.waypoint_count: + command = self._command_at(phase, self._waypoint_index, execution_mask) + self._waypoint_index += 1 + return self._tick_result(command=command, events=events) + + terminal_error = self._terminal_error(phase) + not_reached = execution_mask & ( + terminal_error > phase.spec.recovery_policy.tracking_error_threshold + ) + if not_reached.any(): + events.extend( + self._attempt_replan( + not_reached, + ExecutionEventKind.TRACKING_ERROR, + "Terminal command has not been reached.", + ) + ) + if self._status is not ExecutionStatus.RUNNING: + return self._tick_result(command=None, events=events) + phase = self._current_phase() + execution_mask = self._pending & self._plan.plan_success + command = self._command_at(phase, 0, execution_mask) + self._waypoint_index = 1 + return self._tick_result(command=command, events=events) + + events.append( + self._event( + ExecutionEventKind.PHASE_COMPLETED, + execution_mask, + f"Phase {phase.spec.name!r} completed.", + ) + ) + if self._phase_index + 1 < len(self._plan.phases): + self._phase_index += 1 + self._waypoint_index = 0 + self._replans.zero_() + self._phase_started_at = self._context.robot.timestamp + self._last_command = None + return self._tick_result(command=self._hold_command(), events=events) + + command, completion_events = self._finish_action( + execution_mask, + effect_success, + ) + events.extend(completion_events) + return self._tick_result(command=command, events=events) + + def _plan_current( + self, + context: PlanningContext, + event_kind: ExecutionEventKind, + ) -> None: + """Plan the current invocation from the latest observation.""" + request = self._requests[self._invocation_index] + plan = self._engine.plan_request(request, context) + self._install_plan(plan, context, event_kind) + + def _install_plan( + self, + plan: ActionPlan, + context: PlanningContext, + event_kind: ExecutionEventKind, + ) -> None: + """Install an already validated plan as the current execution plan.""" + self._plan = plan + self._phase_index = min(self._phase_index, len(plan.phases) - 1) + self._waypoint_index = 0 + self._planned_scene = context.scene + self._phase_started_at = context.robot.timestamp + self._last_command = None + self._last_command_mask.zero_() + self._effect_wait_emitted = False + planned_mask = self._pending & plan.plan_success + self._queued_events.append( + self._event(event_kind, planned_mask, "Planned from the latest context.") + ) + + def _current_phase(self) -> PlannedPhase: + """Return the currently active planned phase.""" + assert self._plan is not None + return self._plan.phases[self._phase_index] + + def _recover_if_needed( + self, + phase: PlannedPhase, + execution_mask: torch.Tensor, + ) -> list[ExecutionEvent]: + """Detect tracking, scene, and timeout invalidation.""" + events: list[ExecutionEvent] = [] + if not execution_mask.any(): + return events + if ( + self._context.robot.timestamp - self._phase_started_at + > phase.spec.recovery_policy.phase_timeout + ): + return self._attempt_action_retry( + execution_mask, + ExecutionEventKind.PHASE_TIMEOUT, + "Phase timeout exceeded.", + ) + if self._last_command is not None: + tracking_error = torch.amax( + torch.abs(self._context.robot.qpos - self._last_command), dim=1 + ) + tracking_mask = ( + execution_mask + & self._last_command_mask + & (tracking_error > phase.spec.recovery_policy.tracking_error_threshold) + ) + if tracking_mask.any(): + return self._attempt_replan( + tracking_mask, + ExecutionEventKind.TRACKING_ERROR, + "Observed joint tracking error exceeded the policy threshold.", + ) + scene_mask = self._dynamic_scene_change_mask(phase) + if (execution_mask & scene_mask).any(): + return self._attempt_replan( + execution_mask & scene_mask, + ExecutionEventKind.DYNAMIC_GOAL_CHANGED, + "A referenced scene entity moved beyond the policy threshold.", + ) + return events + + def _attempt_replan( + self, + trigger_mask: torch.Tensor, + reason: ExecutionEventKind, + message: str, + ) -> list[ExecutionEvent]: + """Apply per-row replan budgets and regenerate the current action plan.""" + phase = self._current_phase() + events = [self._event(reason, trigger_mask, message)] + allowed = ( + trigger_mask + & phase.spec.replannable + & (self._replans < phase.spec.recovery_policy.max_replans) + ) + exhausted = trigger_mask & ~allowed + if exhausted.any(): + self._eligible &= ~exhausted + self._pending &= ~exhausted + events.append( + self._event( + ExecutionEventKind.RECOVERY_EXHAUSTED, + exhausted, + "Local replan budget exhausted.", + ) + ) + if allowed.any(): + self._replans[allowed] += 1 + self._plan_current(self._context, ExecutionEventKind.REPLANNED) + events.extend(self._drain_events()) + self._update_terminal_status() + return events + + def _attempt_action_retry( + self, + trigger_mask: torch.Tensor, + reason: ExecutionEventKind, + message: str, + ) -> list[ExecutionEvent]: + """Retry the current action or permanently fail exhausted rows.""" + policy = self._current_phase().spec.recovery_policy + events = [self._event(reason, trigger_mask, message)] + allowed = trigger_mask & (self._action_retries < policy.max_phase_retries) + exhausted = trigger_mask & ~allowed + if exhausted.any(): + self._eligible &= ~exhausted + self._pending &= ~exhausted + events.append( + self._event( + ExecutionEventKind.RECOVERY_EXHAUSTED, + exhausted, + "Action retry budget exhausted.", + ) + ) + if allowed.any(): + self._action_retries[allowed] += 1 + self._replans.zero_() + events.append( + self._event( + ExecutionEventKind.ACTION_RETRY, + allowed, + "Retrying the action from the latest observation.", + ) + ) + self._phase_index = 0 + self._plan_current(self._context, ExecutionEventKind.REPLANNED) + events.extend(self._drain_events()) + self._update_terminal_status() + return events + + def _finish_action( + self, + execution_mask: torch.Tensor, + effect_success: torch.Tensor | None, + ) -> tuple[JointCommand | None, list[ExecutionEvent]]: + """Verify effects, update symbolic state, and advance the action barrier.""" + assert self._plan is not None + events: list[ExecutionEvent] = [] + if self._plan.expected_effects.is_empty: + verified = execution_mask + elif effect_success is None: + if not self._effect_wait_emitted: + events.append( + self._event( + ExecutionEventKind.EFFECT_VERIFICATION_REQUIRED, + execution_mask, + "Expected symbolic effects require external verification.", + ) + ) + self._effect_wait_emitted = True + return self._hold_command(), events + else: + verified_input = self._normalize_mask(effect_success, "effect_success") + verified = execution_mask & verified_input + + if verified.any(): + self._task_state = self._plan.expected_effects.apply( + self._task_state, verified + ) + self._context = PlanningContext( + robot=self._context.robot, + task=self._task_state, + scene=self._context.scene, + env_ids=self._context.env_ids, + ) + self._pending &= ~verified + failed_effect = execution_mask & ~verified + planning_failed = self._pending & ~self._plan.plan_success + retry_mask = failed_effect | planning_failed + if retry_mask.any(): + events.extend( + self._attempt_action_retry( + retry_mask, + ExecutionEventKind.ACTION_RETRY, + "Planning or expected-effect verification failed.", + ) + ) + if self._status is not ExecutionStatus.RUNNING: + return None, events + return self._hold_command(), events + + if self._pending.any(): + return self._hold_command(), events + events.append( + self._event( + ExecutionEventKind.ACTION_COMPLETED, + self._eligible, + "Action completed at the batch barrier.", + ) + ) + self._invocation_index += 1 + if self._invocation_index >= len(self._requests): + self._status = ( + ExecutionStatus.COMPLETED + if self._eligible.any() + else ExecutionStatus.FAILED + ) + events.append( + self._event( + ExecutionEventKind.SESSION_COMPLETED, + self._eligible, + "Invocation sequence completed.", + ) + ) + return None, events + + self._pending = self._eligible.clone() + self._action_retries.zero_() + self._replans.zero_() + self._phase_index = 0 + self._plan_current(self._context, ExecutionEventKind.ACTION_PLANNED) + events.extend(self._drain_events()) + return self._hold_command(), events + + def _command_at( + self, + phase: PlannedPhase, + waypoint_index: int, + active_mask: torch.Tensor, + ) -> JointCommand: + """Build one command and retain it for tracking-error monitoring.""" + positions = phase.trajectory.positions[:, waypoint_index] + hold = self._context.robot.qpos + positions = torch.where(active_mask[:, None], positions, hold) + velocities = None + if phase.trajectory.velocities is not None: + values = phase.trajectory.velocities[:, waypoint_index] + velocities = torch.where( + active_mask[:, None], values, torch.zeros_like(values) + ) + self._last_command = positions.clone() + self._last_command_mask = active_mask.clone() + return JointCommand( + positions=positions, + velocities=velocities, + active_mask=active_mask, + env_ids=phase.trajectory.env_ids, + ) + + def _hold_command(self) -> JointCommand: + """Build a passive hold command from the latest observation.""" + return JointCommand( + positions=self._context.robot.qpos, + velocities=torch.zeros_like(self._context.robot.qpos), + active_mask=torch.zeros_like(self._eligible), + env_ids=self._context.env_ids, + ) + + def _terminal_error(self, phase: PlannedPhase) -> torch.Tensor: + """Return per-row max joint error to the phase terminal command.""" + if phase.trajectory.waypoint_count == 0: + return torch.full_like(self._eligible, float("inf"), dtype=torch.float32) + return torch.amax( + torch.abs(self._context.robot.qpos - phase.trajectory.positions[:, -1]), + dim=1, + ) + + def _dynamic_scene_change_mask(self, phase: PlannedPhase) -> torch.Tensor: + """Detect material motion of entities referenced by the phase goal.""" + dependencies = phase.spec.scene_dependencies + changed = torch.zeros_like(self._eligible) + if ( + not dependencies + or self._context.scene.version == self._planned_scene.version + ): + return changed + policy = phase.spec.recovery_policy + for entity_id in dependencies: + previous = self._planned_scene.entities.get(entity_id) + current = self._context.scene.entities.get(entity_id) + if previous is None or current is None: + changed |= self._eligible + continue + previous_pose = self._batched_entity_pose(previous) + current_pose = self._batched_entity_pose(current) + translation = torch.linalg.vector_norm( + current_pose[:, :3, 3] - previous_pose[:, :3, 3], dim=1 + ) + relative_rotation = torch.bmm( + previous_pose[:, :3, :3].transpose(1, 2), + current_pose[:, :3, :3], + ) + cosine = ( + (relative_rotation.diagonal(dim1=1, dim2=2).sum(dim=1) - 1.0) / 2.0 + ).clamp(-1.0, 1.0) + rotation = torch.acos(cosine) + changed |= (translation > policy.goal_translation_threshold) | ( + rotation > policy.goal_rotation_threshold + ) + return changed + + def _batched_entity_pose(self, state: EntityState) -> torch.Tensor: + """Broadcast an entity pose to the session batch.""" + pose = state.pose.to( + device=self._context.robot.qpos.device, + dtype=self._context.robot.qpos.dtype, + ) + if pose.shape == (4, 4): + return pose.unsqueeze(0).expand(self._context.batch_size, -1, -1) + if pose.shape != (self._context.batch_size, 4, 4): + raise ValueError("Scene entity pose batch does not match the session.") + return pose + + def _normalize_mask(self, value: torch.Tensor, name: str) -> torch.Tensor: + """Validate and copy a per-environment boolean mask.""" + if value.dtype != torch.bool or value.shape != (self._context.batch_size,): + raise ValueError( + f"{name} must be bool with shape ({self._context.batch_size},)." + ) + return value.to(self._context.robot.qpos.device).clone() + + def _event( + self, + kind: ExecutionEventKind, + env_mask: torch.Tensor, + message: str, + ) -> ExecutionEvent: + """Create an event correlated with the current invocation.""" + skill_id = ( + self._requests[self._invocation_index].skill_id + if self._invocation_index < len(self._requests) + else None + ) + invocation_id = ( + self._requests[self._invocation_index].invocation_id + if self._invocation_index < len(self._requests) + else None + ) + invocation_revision = ( + self._requests[self._invocation_index].revision + if self._invocation_index < len(self._requests) + else 0 + ) + return ExecutionEvent( + kind=kind, + timestamp=self._context.robot.timestamp, + skill_id=skill_id, + invocation_id=invocation_id, + invocation_revision=invocation_revision, + invocation_index=min(self._invocation_index, len(self._requests) - 1), + env_mask=env_mask, + message=message, + ) + + def _drain_events(self) -> list[ExecutionEvent]: + """Return and clear events queued during planning.""" + events = self._queued_events + self._queued_events = [] + return events + + def _update_terminal_status(self) -> None: + """Mark the session failed when no environment can continue.""" + if not self._eligible.any(): + self._status = ExecutionStatus.FAILED + + def _tick_result( + self, + *, + command: JointCommand | None, + events: list[ExecutionEvent], + ) -> ExecutionTick: + """Build an immutable tick result.""" + return ExecutionTick( + status=self._status, + eligible_mask=self._eligible, + command=command, + events=tuple(events), + task_state=self._task_state, + ) + + +__all__ = [ + "ExecutionEvent", + "ExecutionEventKind", + "ExecutionSession", + "ExecutionStatus", + "ExecutionTick", + "JointCommand", +] diff --git a/embodichain/lab/sim/atomic_actions/goals.py b/embodichain/lab/sim/atomic_actions/goals.py new file mode 100644 index 000000000..bdf38811a --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/goals.py @@ -0,0 +1,215 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Goal contracts shared by atomic actions.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, fields, is_dataclass +from typing import Any, ClassVar, Protocol, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from .core import ObjectSemantics + from .state import PlanningContext + + +class ActionGoal(Protocol): + """Structural protocol implemented by atomic-action goal value objects. + + Goals are action-owned dataclasses. They do not have to inherit from a + marker base class; the owning action declares its concrete ``GoalType``. + ``goal_kind`` supplies the stable semantic discriminator needed by skill + catalogs and agent-facing schemas. + """ + + goal_kind: ClassVar[str] + + +@dataclass(frozen=True, slots=True, eq=False) +class SceneEntityPose: + """Late-bound pose derived from a versioned scene entity. + + The semantic request remains stable while each call to + :meth:`AtomicAction.plan` resolves the latest scene pose. This is the + bridge used by an execution session to replan moving goals. + """ + + entity_id: str + """Stable scene entity identifier.""" + + relative_pose: torch.Tensor | None = None + """Optional transform applied as ``entity_pose @ relative_pose``.""" + + minimum_confidence: float = 0.0 + """Minimum accepted perception confidence.""" + + def __post_init__(self) -> None: + if not isinstance(self.entity_id, str) or not self.entity_id.strip(): + raise ValueError("entity_id must be a non-empty string.") + if self.relative_pose is not None: + validate_pose_tensor( + self.relative_pose, + "relative_pose", + allow_waypoints=False, + ) + if not 0.0 <= self.minimum_confidence <= 1.0: + raise ValueError("minimum_confidence must be in [0, 1].") + + +PoseGoalValue = torch.Tensor | SceneEntityPose +"""Explicit pose tensor or a pose resolved from the latest scene snapshot.""" + + +def validate_pose_tensor( + value: torch.Tensor, + name: str, + *, + allow_waypoints: bool, +) -> None: + """Validate the environment-independent part of a pose goal. + + Args: + value: Pose tensor to validate. + name: Field name used in validation errors. + allow_waypoints: Whether a batched waypoint dimension is accepted. + + Raises: + TypeError: If ``value`` is not a tensor. + ValueError: If the tensor shape is not a supported pose shape. + """ + 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)}." + ) + + +def validate_pose_goal( + value: PoseGoalValue, + name: str, + *, + allow_waypoints: bool, +) -> None: + """Validate an explicit or late-bound pose goal.""" + if isinstance(value, SceneEntityPose): + return + validate_pose_tensor(value, name, allow_waypoints=allow_waypoints) + + +def resolve_pose_goal( + value: PoseGoalValue, + context: PlanningContext, + *, + name: str, +) -> torch.Tensor: + """Resolve a pose goal against a planning context. + + Args: + value: Explicit tensor or scene-entity reference. + context: Latest observed planning context. + name: Field name used in validation errors. + + Returns: + Explicit pose tensor. Scene references always return shape ``(B, 4, 4)``. + """ + if isinstance(value, torch.Tensor): + return value + try: + entity = context.scene.entities[value.entity_id] + except KeyError as exc: + raise KeyError( + f"{name} references unknown scene entity {value.entity_id!r}." + ) from exc + if entity.confidence < value.minimum_confidence: + raise ValueError( + f"Scene entity {value.entity_id!r} confidence {entity.confidence} is " + f"below {value.minimum_confidence}." + ) + pose = entity.pose.to(device=context.robot.qpos.device, dtype=torch.float32) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).expand(context.batch_size, -1, -1) + elif pose.shape != (context.batch_size, 4, 4): + raise ValueError( + f"Scene entity {value.entity_id!r} pose must match planning batch size." + ) + if value.relative_pose is None: + return pose.clone() + relative = value.relative_pose.to(device=pose.device, dtype=pose.dtype) + if relative.shape == (4, 4): + relative = relative.unsqueeze(0).expand(context.batch_size, -1, -1) + elif relative.shape != (context.batch_size, 4, 4): + raise ValueError(f"{name}.relative_pose must match planning batch size.") + return torch.bmm(pose, relative) + + +def collect_scene_dependencies(value: Any) -> tuple[str, ...]: + """Collect stable scene entity identifiers referenced by a goal value.""" + found: set[str] = set() + + def visit(item: Any) -> None: + if isinstance(item, SceneEntityPose): + found.add(item.entity_id) + elif is_dataclass(item) and not isinstance(item, type): + for data_field in fields(item): + visit(getattr(item, data_field.name)) + elif isinstance(item, Mapping): + for key, nested in item.items(): + visit(key) + visit(nested) + elif isinstance(item, Sequence) and not isinstance( + item, (str, bytes, torch.Tensor) + ): + for nested in item: + visit(nested) + + visit(value) + return tuple(sorted(found)) + + +@dataclass(frozen=True, slots=True, eq=False) +class ObjectActionGoal: + """Shared semantic-object goal contract for object-centric skills.""" + + goal_kind: ClassVar[str] = "semantic_object" + + semantics: ObjectSemantics + """Semantic and geometric description of the object.""" + + def __post_init__(self) -> None: + from .core import ObjectSemantics + + if not isinstance(self.semantics, ObjectSemantics): + raise TypeError("semantics must be an ObjectSemantics instance.") + + +__all__ = [ + "ActionGoal", + "ObjectActionGoal", + "PoseGoalValue", + "SceneEntityPose", + "collect_scene_dependencies", + "resolve_pose_goal", + "validate_pose_goal", + "validate_pose_tensor", +] diff --git a/embodichain/lab/sim/atomic_actions/invocation.py b/embodichain/lab/sim/atomic_actions/invocation.py new file mode 100644 index 000000000..5e5ad06eb --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/invocation.py @@ -0,0 +1,159 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Grounded action invocations consumed by the deterministic skill layer.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass, field +from typing import Generic, TypeVar + +from .bindings import ActionBinding, ResolvedActionBinding +from .control import ActionControlOverrides +from .goals import ActionGoal +from .policies import MotionPolicy, RecoveryPolicy + +GoalT = TypeVar("GoalT", bound=ActionGoal) + + +@dataclass(frozen=True, slots=True, eq=False) +class ActionOptions: + """Marker base for immutable, skill-specific runtime options. + + Subclasses belong to action modules and contain only behavior that may vary + between invocations. Robot resources and semantic targets do not belong in + this object. + """ + + +OptionsT = TypeVar("OptionsT", bound=ActionOptions) + + +@dataclass(frozen=True, slots=True) +class ActionInvocation(Generic[GoalT, OptionsT]): + """One fully typed and embodiment-bound atomic skill request. + + This is a runtime-domain object, not the JSON protocol emitted by an MLLM. + An action compiler is responsible for converting a semantic ``SkillCallSpec`` + into this grounded representation. + """ + + skill_id: str + """Stable registered skill identifier.""" + + goal: GoalT + """Action-specific goal value object.""" + + binding: ActionBinding + """Semantic-role bindings to keys in the selected robot's control parts.""" + + motion_policy: MotionPolicy = field(default_factory=MotionPolicy) + """Reusable motion-generation settings.""" + + recovery_policy: RecoveryPolicy = field(default_factory=RecoveryPolicy) + """Bounded local execution recovery settings.""" + + skill_options: OptionsT | None = None + """Optional per-invocation behavior override for the selected skill.""" + + control_overrides: ActionControlOverrides = field( + default_factory=ActionControlOverrides + ) + """Optional semantic control commands for this invocation revision.""" + + invocation_id: str | None = None + """Optional correlation identifier propagated into execution traces.""" + + revision: int = 0 + """Monotonic revision used when replacing a runtime invocation.""" + + def __post_init__(self) -> None: + if not isinstance(self.skill_id, str) or not self.skill_id.strip(): + raise ValueError("skill_id must be a non-empty string.") + goal_kind = getattr(type(self.goal), "goal_kind", None) + if not isinstance(goal_kind, str) or not goal_kind: + raise TypeError( + "goal must implement the ActionGoal protocol with a non-empty " + "goal_kind class variable." + ) + if not isinstance(self.binding, ActionBinding): + raise TypeError("binding must be an ActionBinding.") + if not isinstance(self.motion_policy, MotionPolicy): + raise TypeError("motion_policy must be a MotionPolicy.") + if not isinstance(self.recovery_policy, RecoveryPolicy): + raise TypeError("recovery_policy must be a RecoveryPolicy.") + if self.skill_options is not None and not isinstance( + self.skill_options, ActionOptions + ): + raise TypeError("skill_options must be an ActionOptions instance.") + if not isinstance(self.control_overrides, ActionControlOverrides): + raise TypeError("control_overrides must be an ActionControlOverrides.") + if self.invocation_id is not None and ( + not isinstance(self.invocation_id, str) or not self.invocation_id.strip() + ): + raise ValueError("invocation_id must be a non-empty string when set.") + if not isinstance(self.revision, int) or self.revision < 0: + raise ValueError("revision must be a non-negative integer.") + + +@dataclass(frozen=True, slots=True) +class ResolvedActionRequest(Generic[GoalT, OptionsT]): + """Engine-owned immutable planning snapshot for one invocation revision. + + Recovery replans reuse this object verbatim and vary only the + :class:`PlanningContext`. Deep-copying policies and skill options severs + references to caller-owned runtime objects before planning starts. + """ + + skill_id: str + goal: GoalT + binding: ResolvedActionBinding + motion_policy: MotionPolicy + recovery_policy: RecoveryPolicy + skill_options: OptionsT + invocation_id: str | None = None + revision: int = 0 + + def __post_init__(self) -> None: + if not isinstance(self.skill_id, str) or not self.skill_id.strip(): + raise ValueError("skill_id must be a non-empty string.") + if not isinstance(self.binding, ResolvedActionBinding): + raise TypeError("binding must be a ResolvedActionBinding.") + if not isinstance(self.motion_policy, MotionPolicy): + raise TypeError("motion_policy must be a MotionPolicy.") + if not isinstance(self.recovery_policy, RecoveryPolicy): + raise TypeError("recovery_policy must be a RecoveryPolicy.") + if not isinstance(self.skill_options, ActionOptions): + raise TypeError("skill_options must be an ActionOptions instance.") + if self.invocation_id is not None and ( + not isinstance(self.invocation_id, str) or not self.invocation_id.strip() + ): + raise ValueError("invocation_id must be a non-empty string when set.") + if not isinstance(self.revision, int) or self.revision < 0: + raise ValueError("revision must be a non-negative integer.") + object.__setattr__(self, "motion_policy", deepcopy(self.motion_policy)) + object.__setattr__(self, "recovery_policy", deepcopy(self.recovery_policy)) + object.__setattr__(self, "skill_options", deepcopy(self.skill_options)) + + +__all__ = [ + "ActionInvocation", + "ActionOptions", + "GoalT", + "OptionsT", + "ResolvedActionRequest", +] diff --git a/embodichain/lab/sim/atomic_actions/plans.py b/embodichain/lab/sim/atomic_actions/plans.py new file mode 100644 index 000000000..f1f9dc3e6 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/plans.py @@ -0,0 +1,477 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Side-effect-free plans produced by atomic actions.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from types import MappingProxyType +from typing import Any, Mapping, Sequence + +import torch + +from .effects import StateDelta +from .goals import ActionGoal +from .policies import RecoveryPolicy +from .state import PlanningContext + + +def _validate_optional_trajectory_field( + value: torch.Tensor | None, + positions: torch.Tensor, + name: str, +) -> None: + """Validate an optional velocity or acceleration tensor.""" + if value is None: + return + if not isinstance(value, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor or None.") + if value.shape != positions.shape: + raise ValueError(f"{name} must match positions shape.") + if value.device != positions.device: + raise ValueError(f"{name} must share the positions device.") + if not torch.isfinite(value).all(): + raise ValueError(f"{name} must contain only finite values.") + + +@dataclass(frozen=True, slots=True, eq=False) +class TimedTrajectory: + """Full-robot joint trajectory with per-environment timing metadata.""" + + positions: torch.Tensor + velocities: torch.Tensor | None + accelerations: torch.Tensor | None + dt: torch.Tensor + duration: torch.Tensor + env_ids: torch.Tensor + + def __post_init__(self) -> None: + if not isinstance(self.positions, torch.Tensor): + raise TypeError("positions must be a torch.Tensor.") + if self.positions.dim() != 3: + raise ValueError("positions must have shape (batch, waypoint, robot_dof).") + if self.positions.shape[0] == 0 or self.positions.shape[2] == 0: + raise ValueError( + "positions batch and robot_dof dimensions must be non-zero." + ) + if not torch.isfinite(self.positions).all(): + raise ValueError("positions must contain only finite values.") + _validate_optional_trajectory_field( + self.velocities, self.positions, "velocities" + ) + _validate_optional_trajectory_field( + self.accelerations, self.positions, "accelerations" + ) + batch_size, waypoint_count, _ = self.positions.shape + if not isinstance(self.dt, torch.Tensor) or self.dt.shape != ( + batch_size, + waypoint_count, + ): + raise ValueError(f"dt must have shape ({batch_size}, {waypoint_count}).") + if self.dt.device != self.positions.device: + raise ValueError("dt must share the positions device.") + if not torch.isfinite(self.dt).all() or (self.dt < 0).any(): + raise ValueError("dt must contain finite non-negative values.") + if not isinstance(self.duration, torch.Tensor) or self.duration.shape != ( + batch_size, + ): + raise ValueError(f"duration must have shape ({batch_size},).") + if self.duration.device != self.positions.device: + raise ValueError("duration must share the positions device.") + if not torch.isfinite(self.duration).all() or (self.duration < 0).any(): + raise ValueError("duration must contain finite non-negative values.") + if not torch.allclose( + self.duration, + self.dt.sum(dim=1), + rtol=1e-4, + atol=1e-6, + ): + raise ValueError("duration must equal the sum of dt for each environment.") + if not isinstance(self.env_ids, torch.Tensor): + raise TypeError("env_ids must be a torch.Tensor.") + if self.env_ids.dtype != torch.long or self.env_ids.shape != (batch_size,): + raise ValueError(f"env_ids must be int64 with shape ({batch_size},).") + if self.env_ids.device != self.positions.device: + raise ValueError("env_ids must share the positions device.") + object.__setattr__(self, "env_ids", self.env_ids.clone()) + + @property + def batch_size(self) -> int: + """Number of environment rows.""" + return int(self.positions.shape[0]) + + @property + def waypoint_count(self) -> int: + """Number of trajectory samples.""" + return int(self.positions.shape[1]) + + @property + def robot_dof(self) -> int: + """Number of full-robot command columns.""" + return int(self.positions.shape[2]) + + @classmethod + def from_positions( + cls, + positions: torch.Tensor, + *, + env_ids: torch.Tensor, + control_dt: float, + velocities: torch.Tensor | None = None, + accelerations: torch.Tensor | None = None, + dt: torch.Tensor | None = None, + duration: torch.Tensor | float | None = None, + ) -> TimedTrajectory: + """Build a timed trajectory and synthesize missing timing metadata. + + Args: + positions: Full-robot positions, shape ``(B, N, D)``. + env_ids: Environment identifiers, shape ``(B,)``. + control_dt: Fallback interval used when ``dt`` is absent. + velocities: Optional joint velocities. + accelerations: Optional joint accelerations. + dt: Optional per-sample time deltas. + duration: Optional duration used to synthesize or validate ``dt``. + + Returns: + Validated timed trajectory. + """ + if control_dt <= 0.0: + raise ValueError("control_dt must be greater than zero.") + if not isinstance(positions, torch.Tensor) or positions.dim() != 3: + raise ValueError("positions must have shape (B, N, D).") + batch_size, waypoint_count, _ = positions.shape + if dt is None: + dt = torch.zeros( + (batch_size, waypoint_count), + dtype=torch.float32, + device=positions.device, + ) + if waypoint_count > 1: + if duration is None: + dt[:, 1:] = control_dt + else: + duration_tensor = torch.as_tensor( + duration, dtype=torch.float32, device=positions.device + ) + if duration_tensor.dim() == 0: + duration_tensor = duration_tensor.expand(batch_size) + if duration_tensor.shape != (batch_size,): + raise ValueError(f"duration must have shape ({batch_size},).") + dt[:, 1:] = duration_tensor[:, None] / (waypoint_count - 1) + else: + dt = dt.to(device=positions.device, dtype=torch.float32) + computed_duration = dt.sum(dim=1) + if duration is not None: + duration_tensor = torch.as_tensor( + duration, dtype=torch.float32, device=positions.device + ) + if duration_tensor.dim() == 0: + duration_tensor = duration_tensor.expand(batch_size) + if duration_tensor.shape != (batch_size,): + raise ValueError(f"duration must have shape ({batch_size},).") + if not torch.allclose( + computed_duration, duration_tensor, rtol=1e-4, atol=1e-6 + ): + raise ValueError("duration does not match the supplied dt.") + return cls( + positions=positions, + velocities=velocities, + accelerations=accelerations, + dt=dt, + duration=computed_duration, + env_ids=env_ids, + ) + + @classmethod + def empty( + cls, + *, + batch_size: int, + robot_dof: int, + device: torch.device | str, + env_ids: torch.Tensor, + ) -> TimedTrajectory: + """Create an empty trajectory with explicit batch and DoF dimensions.""" + resolved = torch.device(device) + return cls( + positions=torch.empty( + (batch_size, 0, robot_dof), dtype=torch.float32, device=resolved + ), + velocities=None, + accelerations=None, + dt=torch.empty((batch_size, 0), dtype=torch.float32, device=resolved), + duration=torch.zeros(batch_size, dtype=torch.float32, device=resolved), + env_ids=env_ids, + ) + + def hold_rows( + self, + active_mask: torch.Tensor, + hold_qpos: torch.Tensor, + ) -> TimedTrajectory: + """Replace inactive rows with a fixed hold command. + + Args: + active_mask: Rows allowed to execute this trajectory. + hold_qpos: Hold positions, shape ``(B, D)``. + + Returns: + New trajectory with inactive rows frozen and derivatives zeroed. + """ + if active_mask.dtype != torch.bool or active_mask.shape != (self.batch_size,): + raise ValueError("active_mask must be bool with shape (batch_size,).") + if hold_qpos.shape != (self.batch_size, self.robot_dof): + raise ValueError("hold_qpos must have shape (batch_size, robot_dof).") + active_mask = active_mask.to(self.positions.device) + held = ( + hold_qpos.to(self.positions.device).unsqueeze(1).expand_as(self.positions) + ) + positions = torch.where(active_mask[:, None, None], self.positions, held) + + def mask_derivative(value: torch.Tensor | None) -> torch.Tensor | None: + if value is None: + return None + return torch.where( + active_mask[:, None, None], value, torch.zeros_like(value) + ) + + return TimedTrajectory( + positions=positions, + velocities=mask_derivative(self.velocities), + accelerations=mask_derivative(self.accelerations), + dt=self.dt, + duration=self.duration, + env_ids=self.env_ids, + ) + + @classmethod + def concatenate( + cls, + trajectories: Sequence[TimedTrajectory], + *, + empty_like: PlanningContext | None = None, + ) -> TimedTrajectory: + """Concatenate trajectories along their waypoint dimension. + + Args: + trajectories: Compatible trajectories in execution order. + empty_like: Context used only when ``trajectories`` is empty. + + Returns: + Concatenated full-robot trajectory. + """ + if not trajectories: + if empty_like is None: + raise ValueError("empty_like is required for an empty concatenation.") + return cls.empty( + batch_size=empty_like.batch_size, + robot_dof=empty_like.robot.robot_dof, + device=empty_like.robot.qpos.device, + env_ids=empty_like.env_ids, + ) + first = trajectories[0] + for trajectory in trajectories[1:]: + if trajectory.batch_size != first.batch_size: + raise ValueError("All trajectories must share a batch size.") + if trajectory.robot_dof != first.robot_dof: + raise ValueError("All trajectories must share robot_dof.") + if not torch.equal(trajectory.env_ids, first.env_ids): + raise ValueError("All trajectories must share env_ids.") + if trajectory.positions.device != first.positions.device: + raise ValueError("All trajectories must share a device.") + + def concatenate_optional(name: str) -> torch.Tensor | None: + values = [getattr(item, name) for item in trajectories] + if any(value is None for value in values): + return None + return torch.cat(values, dim=1) # type: ignore[arg-type] + + dt = torch.cat([item.dt for item in trajectories], dim=1) + return cls( + positions=torch.cat([item.positions for item in trajectories], dim=1), + velocities=concatenate_optional("velocities"), + accelerations=concatenate_optional("accelerations"), + dt=dt, + duration=dt.sum(dim=1), + env_ids=first.env_ids, + ) + + +class CompletionConditionKind(str, Enum): + """Built-in phase completion-condition categories.""" + + TRAJECTORY_COMPLETE = "trajectory_complete" + JOINT_GOAL_REACHED = "joint_goal_reached" + EEF_GOAL_REACHED = "eef_goal_reached" + EFFECT_VERIFIED = "effect_verified" + + +@dataclass(frozen=True, slots=True) +class CompletionCondition: + """Declarative phase completion condition.""" + + kind: CompletionConditionKind = CompletionConditionKind.TRAJECTORY_COMPLETE + tolerance: float | None = None + + def __post_init__(self) -> None: + if self.tolerance is not None and self.tolerance <= 0.0: + raise ValueError("Completion-condition tolerance must be positive.") + + +@dataclass(frozen=True, slots=True) +class PlannerDiagnostics: + """Planner metadata retained for debugging and recovery decisions.""" + + backend: str + messages: tuple[str, ...] = () + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.backend, str) or not self.backend: + raise ValueError("PlannerDiagnostics.backend must be non-empty.") + object.__setattr__(self, "messages", tuple(self.messages)) + object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata))) + + +@dataclass(frozen=True, slots=True) +class PhaseSpec: + """Semantic and runtime contract for one sequential action phase.""" + + name: str + goal: ActionGoal + replannable: bool + completion_condition: CompletionCondition + recovery_policy: RecoveryPolicy + scene_dependencies: tuple[str, ...] = () + """Scene entities whose motion can invalidate this phase plan.""" + + def __post_init__(self) -> None: + if not isinstance(self.name, str) or not self.name: + raise ValueError("PhaseSpec.name must be non-empty.") + dependencies = tuple(self.scene_dependencies) + if len(set(dependencies)) != len(dependencies) or not all( + isinstance(entity_id, str) and entity_id for entity_id in dependencies + ): + raise ValueError( + "scene_dependencies must contain unique non-empty entity ids." + ) + object.__setattr__(self, "scene_dependencies", dependencies) + + +@dataclass(frozen=True, slots=True) +class PlannedPhase: + """One scene-bound phase trajectory and its diagnostics.""" + + spec: PhaseSpec + trajectory: TimedTrajectory + planned_scene_version: int + diagnostics: PlannerDiagnostics + + def __post_init__(self) -> None: + if self.planned_scene_version < 0: + raise ValueError("planned_scene_version must be non-negative.") + + +@dataclass(frozen=True, slots=True, eq=False) +class ActionPlan: + """Planning result for one grounded atomic action invocation.""" + + skill_id: str + plan_success: torch.Tensor + phases: tuple[PlannedPhase, ...] + expected_effects: StateDelta = field(default_factory=StateDelta) + invocation_id: str | None = None + invocation_revision: int = 0 + + def __post_init__(self) -> None: + if not isinstance(self.skill_id, str) or not self.skill_id: + raise ValueError("ActionPlan.skill_id must be non-empty.") + if ( + not isinstance(self.invocation_revision, int) + or self.invocation_revision < 0 + ): + raise ValueError("invocation_revision must be a non-negative integer.") + if not isinstance(self.plan_success, torch.Tensor): + raise TypeError("plan_success must be a torch.Tensor.") + if self.plan_success.dtype != torch.bool or self.plan_success.dim() != 1: + raise ValueError("plan_success must be a 1D bool tensor.") + if not self.phases: + raise ValueError("ActionPlan must contain at least one phase.") + phases = tuple(self.phases) + first = phases[0].trajectory + if first.batch_size != self.plan_success.shape[0]: + raise ValueError("plan_success batch must match phase trajectories.") + if first.positions.device != self.plan_success.device: + raise ValueError("plan_success and phase trajectories must share a device.") + for phase in phases[1:]: + trajectory = phase.trajectory + if trajectory.batch_size != first.batch_size: + raise ValueError("All action phases must share a batch size.") + if trajectory.robot_dof != first.robot_dof: + raise ValueError("All action phases must share robot_dof.") + if not torch.equal(trajectory.env_ids, first.env_ids): + raise ValueError("All action phases must share env_ids.") + object.__setattr__(self, "plan_success", self.plan_success.clone()) + object.__setattr__(self, "phases", phases) + + @property + def trajectory(self) -> TimedTrajectory: + """Concatenate the sequential phase trajectories.""" + return TimedTrajectory.concatenate( + tuple(phase.trajectory for phase in self.phases) + ) + + @property + def success_all(self) -> bool: + """Whether every environment row planned successfully.""" + return bool(self.plan_success.all().item()) + + +@dataclass(frozen=True, slots=True, eq=False) +class CompiledTrajectory: + """Offline compilation result for a sequence of action invocations.""" + + plan_success: torch.Tensor + trajectory: TimedTrajectory + action_plans: tuple[ActionPlan, ...] + projected_context: PlanningContext + + def __post_init__(self) -> None: + if self.plan_success.dtype != torch.bool or self.plan_success.shape != ( + self.trajectory.batch_size, + ): + raise ValueError("Compiled plan_success must be bool with shape (batch,).") + if self.plan_success.device != self.trajectory.positions.device: + raise ValueError( + "Compiled plan_success and trajectory must share a device." + ) + object.__setattr__(self, "plan_success", self.plan_success.clone()) + object.__setattr__(self, "action_plans", tuple(self.action_plans)) + + +__all__ = [ + "ActionPlan", + "CompiledTrajectory", + "CompletionCondition", + "CompletionConditionKind", + "PhaseSpec", + "PlannedPhase", + "PlannerDiagnostics", + "TimedTrajectory", +] diff --git a/embodichain/lab/sim/atomic_actions/policies.py b/embodichain/lab/sim/atomic_actions/policies.py new file mode 100644 index 000000000..8d56ab333 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/policies.py @@ -0,0 +1,126 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Motion-generation and bounded-recovery policies for atomic actions.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from embodichain.lab.sim.planners import PlanOptions + + +@dataclass(frozen=True, slots=True) +class MotionPolicy: + """Immutable motion-generation policy for one action invocation. + + The policy is a runtime value object rather than application configuration. + ``plan_opts`` is copied on construction so a caller-owned planner config + cannot change an invocation after it has been created. + """ + + planner: str | None = None + """Optional required planner backend name; ``None`` accepts the configured one.""" + + motion_source: str = "ik_interp" + """Trajectory source: ``ik_interp`` or ``motion_gen``.""" + + interpolation: str = "linear" + """Interpolation policy. Only linear interpolation is currently supported.""" + + sample_count: int = 50 + """Requested trajectory sample count when the backend does not preserve samples.""" + + control_dt: float = 1.0 / 60.0 + """Fallback command period in seconds when a planner supplies no timing.""" + + velocity_limit: float | None = None + """Optional planner velocity limit.""" + + acceleration_limit: float | None = None + """Optional planner acceleration limit.""" + + collision_check: bool = True + """Whether collision-aware backends should enable collision checking.""" + + plan_opts: PlanOptions | None = None + """Optional typed planner-specific options.""" + + def __post_init__(self) -> None: + valid_sources = {"ik_interp", "motion_gen"} + if self.motion_source not in valid_sources: + raise ValueError( + f"motion_source must be one of {sorted(valid_sources)}, " + f"got {self.motion_source!r}." + ) + if self.interpolation != "linear": + raise ValueError( + "interpolation currently supports only 'linear', " + f"got {self.interpolation!r}." + ) + if self.sample_count < 2: + raise ValueError("sample_count must be at least 2.") + if self.control_dt <= 0.0: + raise ValueError("control_dt must be greater than zero.") + if self.velocity_limit is not None and self.velocity_limit <= 0.0: + raise ValueError("velocity_limit must be greater than zero when set.") + if self.acceleration_limit is not None and self.acceleration_limit <= 0.0: + raise ValueError("acceleration_limit must be greater than zero when set.") + object.__setattr__(self, "plan_opts", deepcopy(self.plan_opts)) + + +@dataclass(frozen=True, slots=True) +class RecoveryPolicy: + """Bounded local recovery policy used by the execution runtime.""" + + max_replans: int = 3 + """Maximum current-phase replans.""" + + max_phase_retries: int = 2 + """Maximum retries of a failed phase.""" + + tracking_error_threshold: float = 0.05 + """Joint tracking-error threshold in radians.""" + + goal_translation_threshold: float = 0.02 + """Dynamic-goal translation threshold in metres.""" + + goal_rotation_threshold: float = 0.0872664626 + """Dynamic-goal rotation threshold in radians (five degrees by default).""" + + phase_timeout: float = 30.0 + """Maximum phase execution time in seconds.""" + + def __post_init__(self) -> None: + if self.max_replans < 0: + raise ValueError("max_replans must be non-negative.") + if self.max_phase_retries < 0: + raise ValueError("max_phase_retries must be non-negative.") + threshold_fields = ( + "tracking_error_threshold", + "goal_translation_threshold", + "goal_rotation_threshold", + "phase_timeout", + ) + for name in threshold_fields: + if getattr(self, name) <= 0.0: + raise ValueError(f"{name} must be greater than zero.") + + +__all__ = ["MotionPolicy", "RecoveryPolicy"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/__init__.py b/embodichain/lab/sim/atomic_actions/primitives/__init__.py index 81ada969f..85de2c985 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/__init__.py +++ b/embodichain/lab/sim/atomic_actions/primitives/__init__.py @@ -18,66 +18,73 @@ from __future__ import annotations +from ..core import AtomicAction from .coordinated_pickment import ( - CoordinatedPickTarget, + CoordinatedPickGoal, CoordinatedPickment, - CoordinatedPickmentCfg, - CoordinatedPickmentTarget, + CoordinatedPickmentOptions, ) from .coordinated_placement import ( CoordinatedPlacement, - CoordinatedPlacementCfg, - CoordinatedPlacementTarget, + CoordinatedPlacementGoal, + CoordinatedPlacementOptions, ) -from .hand_over import HandOver, HandOverCfg +from .hand_over import HandOver, HandOverOptions from .move_end_effector import ( - EndEffectorPoseTarget, + EndEffectorPoseGoal, MoveEndEffector, - MoveEndEffectorCfg, + MoveEndEffectorOptions, ) from .move_held_object import ( - HeldObjectPoseTarget, + HeldObjectPoseGoal, MoveHeldObject, - MoveHeldObjectCfg, + MoveHeldObjectOptions, ) -from .move_joints import ( - JointPositionTarget, +from .move_joints import JointPositionGoal, MoveJoints, MoveJointsOptions +from .pick_up import GraspGoal, PickUp, PickUpOptions +from .place import AssembleGoal, Place, PlaceGoal, PlaceOptions +from .press import Press, PressGoal, PressOptions + +BUILTIN_ACTION_TYPES: tuple[type[AtomicAction], ...] = ( + MoveEndEffector, MoveJoints, - MoveJointsCfg, - NamedJointPositionTarget, + PickUp, + MoveHeldObject, + Place, + Press, + CoordinatedPickment, + CoordinatedPlacement, + HandOver, ) -from .pick_up import GraspTarget, PickUp, PickUpCfg -from .place import AssembleTarget, Place, PlaceCfg, PlaceTarget -from .press import Press, PressCfg, PressTarget +"""Built-in action implementations instantiated once per action engine.""" __all__ = [ - "AssembleTarget", - "CoordinatedPickTarget", + "AssembleGoal", + "BUILTIN_ACTION_TYPES", + "CoordinatedPickGoal", "CoordinatedPickment", - "CoordinatedPickmentCfg", - "CoordinatedPickmentTarget", + "CoordinatedPickmentOptions", "CoordinatedPlacement", - "CoordinatedPlacementCfg", - "CoordinatedPlacementTarget", - "EndEffectorPoseTarget", - "GraspTarget", + "CoordinatedPlacementGoal", + "CoordinatedPlacementOptions", + "EndEffectorPoseGoal", + "GraspGoal", "HandOver", - "HandOverCfg", - "HeldObjectPoseTarget", - "JointPositionTarget", + "HandOverOptions", + "HeldObjectPoseGoal", + "JointPositionGoal", "MoveEndEffector", - "MoveEndEffectorCfg", + "MoveEndEffectorOptions", "MoveHeldObject", - "MoveHeldObjectCfg", + "MoveHeldObjectOptions", "MoveJoints", - "MoveJointsCfg", - "NamedJointPositionTarget", + "MoveJointsOptions", "PickUp", - "PickUpCfg", + "PickUpOptions", "Place", - "PlaceCfg", - "PlaceTarget", + "PlaceGoal", + "PlaceOptions", "Press", - "PressCfg", - "PressTarget", + "PressGoal", + "PressOptions", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/_helpers.py b/embodichain/lab/sim/atomic_actions/primitives/_helpers.py index 2aecb8cad..fe3f4c60f 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/_helpers.py +++ b/embodichain/lab/sim/atomic_actions/primitives/_helpers.py @@ -22,7 +22,7 @@ from embodichain.utils import logger -from ..core import WorldState +from ..state import PlanningContext def resolve_object_target( @@ -33,7 +33,7 @@ def resolve_object_target( name: str = "object_target_pose", ) -> torch.Tensor: """Broadcast an object target pose to ``(n_envs, 4, 4)`` or validate it.""" - target = target.to(device=device, dtype=torch.float32) + target = target.to(device=device, dtype=torch.float32).clone() if target.shape == (4, 4): target = target.unsqueeze(0).repeat(n_envs, 1, 1) if target.shape != (n_envs, 4, 4): @@ -44,9 +44,12 @@ def resolve_object_target( return target -def arm_qpos_from_state(state: WorldState, arm_joint_ids: list[int]) -> torch.Tensor: - """Extract the arm slice of the full-DoF ``last_qpos`` carried in state.""" - return state.last_qpos[:, arm_joint_ids] +def arm_qpos_from_state( + context: PlanningContext, + arm_joint_ids: list[int], +) -> torch.Tensor: + """Extract the arm slice of the measured planning-start joint positions.""" + return context.robot.qpos[:, arm_joint_ids] __all__ = ["arm_qpos_from_state", "resolve_object_target"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py index 5cce3e692..f8765a595 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py @@ -23,26 +23,32 @@ import torch -from embodichain.utils import configclass, logger +from embodichain.utils import logger from embodichain.utils.math import matrix_from_quat, quat_from_matrix -from ..core import ( - ActionCfg, - ActionResult, - AtomicAction, - CoordinatedHeldObjectState, - WorldState, - _validate_pose_tensor, +from ..bindings import ResolvedControlPart +from ..control import GRASP_COMMAND, OPEN_COMMAND +from ..core import AtomicAction, ObjectSemantics +from ..effects import StateDelta +from ..goals import ( + ObjectActionGoal, + PoseGoalValue, + resolve_pose_goal, + validate_pose_goal, + validate_pose_tensor, ) -from ..targets import ObjectActionTarget -from ..trajectory import TrajectoryBuilder +from ..invocation import ActionOptions, ResolvedActionRequest +from ..plans import ActionPlan +from ..state import CoordinatedHeldObjectState, PlanningContext @dataclass(frozen=True, slots=True, eq=False) -class CoordinatedPickTarget(ObjectActionTarget): +class CoordinatedPickGoal(ObjectActionGoal): """Object-centric target for picking and moving one object with two hands.""" - object_target_pose: torch.Tensor + goal_kind: ClassVar[str] = "coordinated_pick" + + object_target_pose: PoseGoalValue """Target pose for the shared object, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" left_object_to_eef: torch.Tensor @@ -51,69 +57,37 @@ class CoordinatedPickTarget(ObjectActionTarget): right_object_to_eef: torch.Tensor """Transform from object frame to right end-effector frame.""" - object_initial_pose: torch.Tensor | None = None + object_initial_pose: PoseGoalValue | None = None """Optional initial object pose. Defaults to ``semantics.entity`` pose.""" def __post_init__(self) -> None: - ObjectActionTarget.__post_init__(self) - _validate_pose_tensor( + ObjectActionGoal.__post_init__(self) + validate_pose_goal( self.object_target_pose, "object_target_pose", allow_waypoints=False, ) - _validate_pose_tensor( + validate_pose_tensor( self.left_object_to_eef, "left_object_to_eef", allow_waypoints=False, ) - _validate_pose_tensor( + 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( + validate_pose_goal( 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" - """Name of the action, used for identification and logging.""" - - control_part: str = "dual_arm" - """Combined control part containing left and right arm joints.""" - - left_arm_control_part: str = "left_arm" - """Left arm control part used to grasp one end of the object.""" - - right_arm_control_part: str = "right_arm" - """Right arm control part used to grasp the other end of the object.""" - - left_hand_control_part: str = "left_hand" - """Hand attached to the left arm.""" - - right_hand_control_part: str = "right_hand" - """Hand attached to the right arm.""" - - left_hand_open_qpos: torch.Tensor | None = None - """Left hand qpos for the open state.""" - - left_hand_close_qpos: torch.Tensor | None = None - """Left hand qpos for the closed state.""" - - right_hand_open_qpos: torch.Tensor | None = None - """Right hand qpos for the open state.""" - - right_hand_close_qpos: torch.Tensor | None = None - """Right hand qpos for the closed state.""" +@dataclass(frozen=True, slots=True, eq=False) +class CoordinatedPickmentOptions(ActionOptions): + """Per-invocation coordinated pickup behavior.""" object_motion_keyframes: int = 6 """Number of object-pose keyframes solved by IK before joint-space interpolation.""" @@ -124,87 +98,40 @@ class CoordinatedPickmentCfg(ActionCfg): lift_height: float = 0.08 """World-Z lift distance before moving to the object target pose.""" - sample_interval: int = 120 - """Number of waypoints for the full coordinated pickment trajectory.""" - hand_interp_steps: int = 10 """Number of waypoints used for the simultaneous hand close phase.""" hold_steps: int = 4 """Number of waypoints to hold the final object target pose.""" + def __post_init__(self) -> None: + if self.object_motion_keyframes < 2: + raise ValueError("object_motion_keyframes must be at least 2.") + if self.pre_grasp_distance < 0.0: + raise ValueError("pre_grasp_distance must be non-negative.") + if self.lift_height < 0.0: + raise ValueError("lift_height must be non-negative.") + for name in ("hand_interp_steps", "hold_steps"): + if getattr(self, name) < 0: + raise ValueError(f"{name} must be non-negative.") -class _DualArmHelpers: - """Shared trajectory helpers for dual-arm coordinated actions.""" - def _init_dual_arm_parts( - self, - *, - first_arm_control_part: str, - second_arm_control_part: str, - first_hand_control_part: str, - second_hand_control_part: str, - ) -> None: - self.builder = TrajectoryBuilder(self.motion_generator) - self.n_envs = self.robot.get_qpos().shape[0] - self.robot_dof = self.robot.dof - self.dual_arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.first_arm_joint_ids = self.robot.get_joint_ids(name=first_arm_control_part) - self.second_arm_joint_ids = self.robot.get_joint_ids( - name=second_arm_control_part - ) - self.first_hand_joint_ids = self.robot.get_joint_ids( - name=first_hand_control_part - ) - self.second_hand_joint_ids = self.robot.get_joint_ids( - name=second_hand_control_part - ) - self.first_arm_dof = len(self.first_arm_joint_ids) - self.second_arm_dof = len(self.second_arm_joint_ids) - self.dual_arm_dof = len(self.dual_arm_joint_ids) - self.first_hand_dof = len(self.first_hand_joint_ids) - self.second_hand_dof = len(self.second_hand_joint_ids) - self._dual_id_to_col = { - joint_id: col for col, joint_id in enumerate(self.dual_arm_joint_ids) - } - self._first_arm_cols = self._lookup_joint_columns( - self.first_arm_joint_ids, - self._dual_id_to_col, - first_arm_control_part, - ) - self._second_arm_cols = self._lookup_joint_columns( - self.second_arm_joint_ids, - self._dual_id_to_col, - second_arm_control_part, - ) +@dataclass(frozen=True, slots=True, eq=False) +class _CoordinatedPickResources: + """Invocation-bound control parts and compatible hand commands.""" - @staticmethod - def _lookup_joint_columns( - joint_ids: list[int], - joint_id_to_col: dict[int, int], - control_part: str, - ) -> list[int]: - missing = [ - joint_id for joint_id in joint_ids if joint_id not in joint_id_to_col - ] - if missing: - logger.log_error( - f"Joints {missing} from '{control_part}' are not included in " - "the configured dual-arm control part.", - ValueError, - ) - return [joint_id_to_col[joint_id] for joint_id in joint_ids] + left_arm: ResolvedControlPart + right_arm: ResolvedControlPart + left_hand: ResolvedControlPart + right_hand: ResolvedControlPart + left_hand_open_qpos: torch.Tensor + left_hand_close_qpos: torch.Tensor + right_hand_open_qpos: torch.Tensor + right_hand_close_qpos: torch.Tensor - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=False, - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, - ) + +class _DualArmHelpers: + """Shared trajectory helpers for dual-arm coordinated actions.""" def _expand_qpos(self, qpos: torch.Tensor, dof: int, name: str) -> torch.Tensor: qpos = qpos.to(device=self.device, dtype=torch.float32) @@ -233,14 +160,13 @@ def _resolve_pose(self, pose: torch.Tensor, name: str) -> torch.Tensor: def _resolve_dual_arm_start( self, - state: WorldState, + state: PlanningContext, + resources: _CoordinatedPickResources, ) -> tuple[torch.Tensor, torch.Tensor]: - dual_start = state.last_qpos[:, self.dual_arm_joint_ids].to( - device=self.device, dtype=torch.float32 - ) + start_qpos = state.last_qpos.to(device=self.device, dtype=torch.float32) return ( - dual_start[:, self._first_arm_cols], - dual_start[:, self._second_arm_cols], + start_qpos[:, list(resources.left_arm.joint_ids)], + start_qpos[:, list(resources.right_arm.joint_ids)], ) def _plan_named_arm_trajectory( @@ -283,28 +209,15 @@ def _plan_named_arm_trajectory( else self._interpolate_keyframe_qpos(trajectory, n_waypoints) ) - def _compose_dual_arm_trajectory( - self, - first_arm_traj: torch.Tensor, - second_arm_traj: torch.Tensor, - ) -> torch.Tensor: - n_waypoints = first_arm_traj.shape[1] - dual_arm_traj = torch.zeros( - (self.n_envs, n_waypoints, self.dual_arm_dof), - dtype=torch.float32, - device=self.device, - ) - dual_arm_traj[:, :, self._first_arm_cols] = first_arm_traj - dual_arm_traj[:, :, self._second_arm_cols] = second_arm_traj - return dual_arm_traj - def _assemble_phase( self, - state: WorldState, + state: PlanningContext, first_arm_traj: torch.Tensor, second_arm_traj: torch.Tensor, first_hand_traj: torch.Tensor, second_hand_traj: torch.Tensor, + *, + resources: _CoordinatedPickResources, ) -> torch.Tensor: n_waypoints = first_arm_traj.shape[1] full = torch.empty( @@ -313,11 +226,10 @@ def _assemble_phase( device=self.device, ) full[:, :, :] = state.last_qpos.to(self.device).unsqueeze(1) - full[:, :, self.dual_arm_joint_ids] = self._compose_dual_arm_trajectory( - first_arm_traj, second_arm_traj - ) - full[:, :, self.first_hand_joint_ids] = first_hand_traj - full[:, :, self.second_hand_joint_ids] = second_hand_traj + full[:, :, list(resources.left_arm.joint_ids)] = first_arm_traj + full[:, :, list(resources.right_arm.joint_ids)] = second_arm_traj + full[:, :, list(resources.left_hand.joint_ids)] = first_hand_traj + full[:, :, list(resources.right_hand.joint_ids)] = second_hand_traj return full @staticmethod @@ -431,21 +343,23 @@ def _interpolate_object_pose( return poses -class CoordinatedPickment(AtomicAction[CoordinatedPickTarget]): +class CoordinatedPickment( + AtomicAction[CoordinatedPickGoal, CoordinatedPickmentOptions] +): """Pick and move a single object pinched by two hands.""" - TargetType: ClassVar[type] = CoordinatedPickTarget + skill_id: ClassVar[str] = "coordinated_pickment" + GoalType: ClassVar[type] = CoordinatedPickGoal + OptionsType: ClassVar[type] = CoordinatedPickmentOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("left", "right") + end_effector_roles: ClassVar[tuple[str, ...]] = ("left", "right") _assemble_phase = _DualArmHelpers._assemble_phase - _compose_dual_arm_trajectory = _DualArmHelpers._compose_dual_arm_trajectory _expand_qpos = _DualArmHelpers._expand_qpos - _fail = _DualArmHelpers._fail - _init_dual_arm_parts = _DualArmHelpers._init_dual_arm_parts _interpolate_keyframe_qpos = _DualArmHelpers._interpolate_keyframe_qpos _interpolate_object_pose = _DualArmHelpers._interpolate_object_pose _interpolate_qpos = _DualArmHelpers._interpolate_qpos _interpolate_qpos_keyframes = _DualArmHelpers._interpolate_qpos_keyframes - _lookup_joint_columns = staticmethod(_DualArmHelpers._lookup_joint_columns) _plan_named_arm_trajectory = _DualArmHelpers._plan_named_arm_trajectory _repeat_qpos = staticmethod(_DualArmHelpers._repeat_qpos) _resolve_dual_arm_start = _DualArmHelpers._resolve_dual_arm_start @@ -453,72 +367,83 @@ class CoordinatedPickment(AtomicAction[CoordinatedPickTarget]): def __init__( self, - motion_generator, - cfg: CoordinatedPickmentCfg | None = None, + default_options: CoordinatedPickmentOptions | None = None, ) -> None: - super().__init__(motion_generator, cfg or CoordinatedPickmentCfg()) - if ( - self.cfg.motion_source == "motion_gen" - and self.motion_generator.planner.cfg.planner_type == "curobo" - ): - logger.log_error( - "Coordinated dual-arm planning is not supported by the cuRobo " - "backend. Use a single-arm action or a dedicated multi-arm " - "planner.", - ValueError, + super().__init__(default_options) + + def _on_bind(self) -> None: + """Resolve engine-wide resources from the owning engine.""" + self.n_envs = self.robot.get_qpos().shape[0] + self.robot_dof = self.robot.dof + + def _resolve_resources( + self, + request: ResolvedActionRequest[CoordinatedPickGoal, CoordinatedPickmentOptions], + ) -> _CoordinatedPickResources: + """Resolve left/right roles from robot control parts.""" + binding = request.binding + left_arm = binding.manipulator("left") + right_arm = binding.manipulator("right") + left_hand = binding.end_effector("left") + right_hand = binding.end_effector("right") + if left_arm.name == right_arm.name: + raise ValueError( + "CoordinatedPickment left and right roles must use different " + "manipulator control parts." ) - self._init_dual_arm_parts( - first_arm_control_part=self.cfg.left_arm_control_part, - second_arm_control_part=self.cfg.right_arm_control_part, - first_hand_control_part=self.cfg.left_hand_control_part, - second_hand_control_part=self.cfg.right_hand_control_part, - ) - self.left_arm_joint_ids = self.first_arm_joint_ids - self.right_arm_joint_ids = self.second_arm_joint_ids - self.left_hand_joint_ids = self.first_hand_joint_ids - self.right_hand_joint_ids = self.second_hand_joint_ids - self.left_arm_dof = self.first_arm_dof - self.right_arm_dof = self.second_arm_dof - self.left_hand_dof = self.first_hand_dof - self.right_hand_dof = self.second_hand_dof - - self._validate_hand_qpos_cfg() - self.left_hand_open_qpos = self._expand_qpos( - self.cfg.left_hand_open_qpos, self.left_hand_dof, "left_hand_open_qpos" - ) - self.left_hand_close_qpos = self._expand_qpos( - self.cfg.left_hand_close_qpos, self.left_hand_dof, "left_hand_close_qpos" - ) - self.right_hand_open_qpos = self._expand_qpos( - self.cfg.right_hand_open_qpos, self.right_hand_dof, "right_hand_open_qpos" - ) - self.right_hand_close_qpos = self._expand_qpos( - self.cfg.right_hand_close_qpos, - self.right_hand_dof, - "right_hand_close_qpos", - ) - - def _validate_hand_qpos_cfg(self) -> None: - for name in ( - "left_hand_open_qpos", - "left_hand_close_qpos", - "right_hand_open_qpos", - "right_hand_close_qpos", - ): - if getattr(self.cfg, name) is None: - logger.log_error( - f"{name} must be specified in CoordinatedPickmentCfg", - ValueError, - ) + if left_hand.name == right_hand.name: + raise ValueError( + "CoordinatedPickment left and right roles must use different " + "end-effector control parts." + ) + return _CoordinatedPickResources( + left_arm=left_arm, + right_arm=right_arm, + left_hand=left_hand, + right_hand=right_hand, + left_hand_open_qpos=left_hand.joint_positions( + OPEN_COMMAND, + n_envs=self.n_envs, + device=self.device, + dtype=torch.float32, + ), + left_hand_close_qpos=left_hand.joint_positions( + GRASP_COMMAND, + n_envs=self.n_envs, + device=self.device, + dtype=torch.float32, + ), + right_hand_open_qpos=right_hand.joint_positions( + OPEN_COMMAND, + n_envs=self.n_envs, + device=self.device, + dtype=torch.float32, + ), + right_hand_close_qpos=right_hand.joint_positions( + GRASP_COMMAND, + n_envs=self.n_envs, + device=self.device, + dtype=torch.float32, + ), + ) def _resolve_object_initial_pose( - self, target: CoordinatedPickTarget + self, + target: CoordinatedPickGoal, + context: PlanningContext, ) -> torch.Tensor: if target.object_initial_pose is not None: - return self._resolve_pose(target.object_initial_pose, "object_initial_pose") + return self._resolve_pose( + resolve_pose_goal( + target.object_initial_pose, + context, + name="object_initial_pose", + ), + "object_initial_pose", + ) if target.semantics.entity is None: logger.log_error( - "CoordinatedPickTarget requires object_initial_pose when " + "CoordinatedPickGoal requires object_initial_pose when " "semantics.entity is not provided.", ValueError, ) @@ -529,7 +454,8 @@ def _resolve_object_initial_pose( def _resolve_target( self, - target: CoordinatedPickTarget, + target: CoordinatedPickGoal, + context: PlanningContext, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -539,9 +465,14 @@ def _resolve_target( torch.Tensor, CoordinatedHeldObjectState, ]: - object_initial_pose = self._resolve_object_initial_pose(target) + object_initial_pose = self._resolve_object_initial_pose(target, context) object_target_pose = self._resolve_pose( - target.object_target_pose, "object_target_pose" + resolve_pose_goal( + target.object_target_pose, + context, + name="object_target_pose", + ), + "object_target_pose", ) left_object_to_eef = self._resolve_pose( target.left_object_to_eef, "left_object_to_eef" @@ -571,17 +502,20 @@ def _resolve_target( held_state, ) - def _compute_segment_lengths(self) -> dict[str, int]: - n_close = max(2, self.cfg.hand_interp_steps) - n_hold = max(0, self.cfg.hold_steps) - n_motion = self.cfg.sample_interval - n_close - n_hold + def _compute_segment_lengths( + self, sample_count: int, options: CoordinatedPickmentOptions + ) -> dict[str, int]: + """Split the invocation sample budget across coordinated-pick phases.""" + n_close = max(2, options.hand_interp_steps) + n_hold = max(0, options.hold_steps) + n_motion = sample_count - n_close - n_hold n_approach = n_motion // 3 n_lift = n_motion // 3 n_move = n_motion - n_approach - n_lift if min(n_approach, n_lift, n_move) < 2: logger.log_error( "Not enough waypoints for coordinated pickment. Please increase " - "sample_interval or decrease hand_interp_steps/hold_steps.", + "sample_count or decrease hand_interp_steps/hold_steps.", ValueError, ) return { @@ -592,17 +526,28 @@ def _compute_segment_lengths(self) -> dict[str, int]: "hold": n_hold, } - def get_segment_lengths(self) -> dict[str, int]: - return self._compute_segment_lengths() + def get_segment_lengths( + self, + sample_count: int, + options: CoordinatedPickmentOptions | None = None, + ) -> dict[str, int]: + """Return phase lengths for an explicit invocation sample budget.""" + return self._compute_segment_lengths( + sample_count, self.default_options if options is None else options + ) - def _compute_pre_grasp_xpos(self, grasp_xpos: torch.Tensor) -> torch.Tensor: + def _compute_pre_grasp_xpos( + self, grasp_xpos: torch.Tensor, options: CoordinatedPickmentOptions + ) -> torch.Tensor: grasp_z = grasp_xpos[:, :3, 2] return self.builder.apply_local_offset( - grasp_xpos, -grasp_z * self.cfg.pre_grasp_distance + grasp_xpos, -grasp_z * options.pre_grasp_distance ) - def _select_motion_keyframe_indices(self, n_waypoints: int) -> torch.Tensor: - n_keyframes = min(max(2, int(self.cfg.object_motion_keyframes)), n_waypoints) + def _select_motion_keyframe_indices( + self, n_waypoints: int, options: CoordinatedPickmentOptions + ) -> torch.Tensor: + n_keyframes = min(max(2, options.object_motion_keyframes), n_waypoints) return ( torch.linspace( 0, @@ -693,9 +638,11 @@ def _plan_synchronized_object_motion( left_object_to_eef: torch.Tensor, right_object_to_eef: torch.Tensor, active_mask: torch.Tensor, + resources: _CoordinatedPickResources, + options: CoordinatedPickmentOptions, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: n_waypoints = object_pose_traj.shape[1] - keyframe_indices = self._select_motion_keyframe_indices(n_waypoints) + keyframe_indices = self._select_motion_keyframe_indices(n_waypoints, options) left_traj = torch.zeros( (self.n_envs, len(keyframe_indices), left_start_qpos.shape[-1]), dtype=torch.float32, @@ -716,23 +663,23 @@ def _plan_synchronized_object_motion( ) left_success, left_qpos = self.robot.compute_ik( pose=left_xpos, - name=self.cfg.left_arm_control_part, + name=resources.left_arm.name, joint_seed=left_qpos_seed, ) right_success, right_qpos = self.robot.compute_ik( pose=right_xpos, - name=self.cfg.right_arm_control_part, + name=resources.right_arm.name, joint_seed=right_qpos_seed, ) left_success = self._as_success_mask(left_success) right_success = self._as_success_mask(right_success) self._log_ik_failures( - self.cfg.left_arm_control_part, + resources.left_arm.name, f"object waypoint {waypoint_idx}", success_mask & ~left_success, ) self._log_ik_failures( - self.cfg.right_arm_control_part, + resources.right_arm.name, f"object waypoint {waypoint_idx}", success_mask & ~right_success, ) @@ -758,7 +705,23 @@ def _plan_synchronized_object_motion( self._interpolate_qpos_keyframes(right_traj, keyframe_indices, n_waypoints), ) - def execute(self, target: CoordinatedPickTarget, state: WorldState) -> ActionResult: + def plan( + self, + request: ResolvedActionRequest[CoordinatedPickGoal, CoordinatedPickmentOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan a coordinated pick without committing the dual attachment.""" + target = self.require_goal(request) + options = request.skill_options + resources = self._resolve_resources(request) + if ( + request.motion_policy.motion_source == "motion_gen" + and self.motion_generator.planner.cfg.planner_type == "curobo" + ): + raise ValueError( + "Coordinated dual-arm planning is not supported by the cuRobo backend." + ) + state = context ( object_initial_pose, object_target_pose, @@ -767,11 +730,15 @@ def execute(self, target: CoordinatedPickTarget, state: WorldState) -> ActionRes left_target_xpos, right_target_xpos, held_state, - ) = self._resolve_target(target) - left_start_qpos, right_start_qpos = self._resolve_dual_arm_start(state) - segments = self._compute_segment_lengths() - left_pre_grasp_xpos = self._compute_pre_grasp_xpos(left_grasp_xpos) - right_pre_grasp_xpos = self._compute_pre_grasp_xpos(right_grasp_xpos) + ) = self._resolve_target(target, context) + left_start_qpos, right_start_qpos = self._resolve_dual_arm_start( + state, resources + ) + segments = self._compute_segment_lengths( + request.motion_policy.sample_count, options + ) + left_pre_grasp_xpos = self._compute_pre_grasp_xpos(left_grasp_xpos, options) + right_pre_grasp_xpos = self._compute_pre_grasp_xpos(right_grasp_xpos, options) left_approach_targets = torch.stack( [left_pre_grasp_xpos, left_grasp_xpos], dim=1 ) @@ -780,14 +747,14 @@ def execute(self, target: CoordinatedPickTarget, state: WorldState) -> ActionRes ) success_mask = torch.ones(self.n_envs, dtype=torch.bool, device=self.device) success_mask, left_approach_traj = self._plan_masked_arm_trajectory( - self.cfg.left_arm_control_part, + resources.left_arm.name, left_start_qpos, left_approach_targets, segments["approach"], success_mask, ) success_mask, right_approach_traj = self._plan_masked_arm_trajectory( - self.cfg.right_arm_control_part, + resources.right_arm.name, right_start_qpos, right_approach_targets, segments["approach"], @@ -800,8 +767,9 @@ def execute(self, target: CoordinatedPickTarget, state: WorldState) -> ActionRes state, left_approach_traj, right_approach_traj, - self._repeat_qpos(self.left_hand_open_qpos, segments["approach"]), - self._repeat_qpos(self.right_hand_open_qpos, segments["approach"]), + self._repeat_qpos(resources.left_hand_open_qpos, segments["approach"]), + self._repeat_qpos(resources.right_hand_open_qpos, segments["approach"]), + resources=resources, ) close_trajectory = self._assemble_phase( @@ -809,20 +777,21 @@ def execute(self, target: CoordinatedPickTarget, state: WorldState) -> ActionRes self._repeat_qpos(left_grasp_qpos, segments["close"]), self._repeat_qpos(right_grasp_qpos, segments["close"]), self._interpolate_qpos( - self.left_hand_open_qpos, - self.left_hand_close_qpos, + resources.left_hand_open_qpos, + resources.left_hand_close_qpos, segments["close"], ), self._interpolate_qpos( - self.right_hand_open_qpos, - self.right_hand_close_qpos, + resources.right_hand_open_qpos, + resources.right_hand_close_qpos, segments["close"], ), + resources=resources, ) lift_object_pose = self.builder.apply_local_offset( object_initial_pose, - torch.tensor([0.0, 0.0, self.cfg.lift_height], device=self.device), + torch.tensor([0.0, 0.0, options.lift_height], device=self.device), ) lift_object_traj = self._interpolate_object_pose( object_initial_pose, @@ -838,6 +807,8 @@ def execute(self, target: CoordinatedPickTarget, state: WorldState) -> ActionRes held_state.left_object_to_eef, held_state.right_object_to_eef, success_mask, + resources, + options, ) ) @@ -847,8 +818,9 @@ def execute(self, target: CoordinatedPickTarget, state: WorldState) -> ActionRes state, left_lift_traj, right_lift_traj, - self._repeat_qpos(self.left_hand_close_qpos, segments["lift"]), - self._repeat_qpos(self.right_hand_close_qpos, segments["lift"]), + self._repeat_qpos(resources.left_hand_close_qpos, segments["lift"]), + self._repeat_qpos(resources.right_hand_close_qpos, segments["lift"]), + resources=resources, ) move_object_traj = self._interpolate_object_pose( @@ -865,6 +837,8 @@ def execute(self, target: CoordinatedPickTarget, state: WorldState) -> ActionRes held_state.left_object_to_eef, held_state.right_object_to_eef, success_mask, + resources, + options, ) ) @@ -874,8 +848,9 @@ def execute(self, target: CoordinatedPickTarget, state: WorldState) -> ActionRes state, left_move_traj, right_move_traj, - self._repeat_qpos(self.left_hand_close_qpos, segments["move"]), - self._repeat_qpos(self.right_hand_close_qpos, segments["move"]), + self._repeat_qpos(resources.left_hand_close_qpos, segments["move"]), + self._repeat_qpos(resources.right_hand_close_qpos, segments["move"]), + resources=resources, ) hold_trajectory = torch.empty( @@ -886,8 +861,9 @@ def execute(self, target: CoordinatedPickTarget, state: WorldState) -> ActionRes state, self._repeat_qpos(left_target_qpos, segments["hold"]), self._repeat_qpos(right_target_qpos, segments["hold"]), - self._repeat_qpos(self.left_hand_close_qpos, segments["hold"]), - self._repeat_qpos(self.right_hand_close_qpos, segments["hold"]), + self._repeat_qpos(resources.left_hand_close_qpos, segments["hold"]), + self._repeat_qpos(resources.right_hand_close_qpos, segments["hold"]), + resources=resources, ) full = torch.cat( @@ -912,36 +888,28 @@ def execute(self, target: CoordinatedPickTarget, state: WorldState) -> ActionRes 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( + return self.build_plan( + request, + context, success=success_mask, trajectory=full, - next_state=state.with_updates( - last_qpos=full[:, -1, :].clone(), - held_objects=held_objects, - coordinated_held_objects=coordinated_held_objects, + expected_effects=StateDelta( + held_object_updates={ + resources.left_arm.name: None, + resources.right_arm.name: None, + }, + coordinated_held_object_updates={ + ( + resources.left_arm.name, + resources.right_arm.name, + ): coordinated_held_object, + }, ), ) __all__ = [ - "CoordinatedPickTarget", + "CoordinatedPickGoal", "CoordinatedPickment", - "CoordinatedPickmentCfg", - "CoordinatedPickmentTarget", + "CoordinatedPickmentOptions", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py index 7f77328a5..1c0a0dcd2 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py @@ -24,29 +24,30 @@ import torch from embodichain.lab.sim.planners import MoveType, PlanState -from embodichain.utils import configclass, logger +from embodichain.utils import logger from ._helpers import resolve_object_target -from ..core import ( - ActionTarget, - ActionCfg, - ActionResult, - AtomicAction, - HeldObjectState, - WorldState, - _validate_pose_tensor, -) -from ..trajectory import TrajectoryBuilder +from ..bindings import ResolvedControlPart +from ..control import GRASP_COMMAND, OPEN_COMMAND +from ..core import AtomicAction +from ..effects import StateDelta +from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal +from ..invocation import ActionOptions, ResolvedActionRequest +from ..plans import ActionPlan +from ..policies import MotionPolicy +from ..state import HeldObjectState, PlanningContext @dataclass(frozen=True, slots=True, eq=False) -class CoordinatedPlacementTarget(ActionTarget): +class CoordinatedPlacementGoal: """Object-centric target for dual-arm coordinated placement.""" - placing_object_target_pose: torch.Tensor + goal_kind: ClassVar[str] = "coordinated_placement" + + placing_object_target_pose: PoseGoalValue """Target pose for the object released by the placing arm.""" - support_object_target_pose: torch.Tensor + support_object_target_pose: PoseGoalValue """Target pose for the object held by the support arm.""" placing_height_offset: float | None = None @@ -56,49 +57,24 @@ class CoordinatedPlacementTarget(ActionTarget): """World-Z offset above the support object target pose.""" release: bool | None = None - """Whether the placing hand releases. ``None`` uses the action config.""" + """Whether the placing hand releases. ``None`` uses invocation options.""" def __post_init__(self) -> None: - _validate_pose_tensor( + validate_pose_goal( self.placing_object_target_pose, "placing_object_target_pose", allow_waypoints=False, ) - _validate_pose_tensor( + validate_pose_goal( self.support_object_target_pose, "support_object_target_pose", allow_waypoints=False, ) -@configclass -class CoordinatedPlacementCfg(ActionCfg): - name: str = "coordinated_placement" - """Name of the action, used for identification and logging.""" - - control_part: str = "dual_arm" - """Robot control part containing both placing and support arms.""" - - placing_arm_control_part: str = "left_arm" - """Arm that places and releases its held object.""" - - support_arm_control_part: str = "right_arm" - """Arm that moves the support object and keeps holding it.""" - - placing_hand_control_part: str = "left_hand" - """Hand attached to the placing arm.""" - - support_hand_control_part: str = "right_hand" - """Hand attached to the support arm.""" - - placing_hand_open_qpos: torch.Tensor | None = None - """Placing-hand qpos for the open state, shape ``[hand_dof,]``.""" - - placing_hand_close_qpos: torch.Tensor | None = None - """Placing-hand qpos for the closed state, shape ``[hand_dof,]``.""" - - support_hand_close_qpos: torch.Tensor | None = None - """Support-hand qpos for the closed state, shape ``[hand_dof,]``.""" +@dataclass(frozen=True, slots=True, eq=False) +class CoordinatedPlacementOptions(ActionOptions): + """Per-invocation coordinated placement behavior.""" release: bool = True """Whether to open the placing hand at the aligned placement pose.""" @@ -112,9 +88,6 @@ class CoordinatedPlacementCfg(ActionCfg): lift_height: float = 0.08 """World-Z lift distance for the placing arm after release.""" - sample_interval: int = 100 - """Number of waypoints for the full coordinated placement trajectory.""" - hand_interp_steps: int = 10 """Number of waypoints for the placing-hand release interpolation.""" @@ -124,113 +97,163 @@ class CoordinatedPlacementCfg(ActionCfg): retreat_steps: int = 16 """Number of waypoints used for the placing-arm lift retreat.""" + def __post_init__(self) -> None: + if self.lift_height < 0.0: + raise ValueError("lift_height must be non-negative.") + for name in ("hand_interp_steps", "hold_steps", "retreat_steps"): + if getattr(self, name) < 0: + raise ValueError(f"{name} must be non-negative.") + + +@dataclass(frozen=True, slots=True, eq=False) +class _CoordinatedPlacementResources: + """Invocation-bound control parts and compatible hand commands.""" + + placing_arm: ResolvedControlPart + support_arm: ResolvedControlPart + placing_hand: ResolvedControlPart + support_hand: ResolvedControlPart + placing_hand_open_qpos: torch.Tensor + placing_hand_close_qpos: torch.Tensor + support_hand_close_qpos: torch.Tensor + -class CoordinatedPlacement(AtomicAction[CoordinatedPlacementTarget]): +class CoordinatedPlacement( + AtomicAction[CoordinatedPlacementGoal, CoordinatedPlacementOptions] +): """Coordinate two held objects: support object below, placing object above.""" - TargetType: ClassVar[type] = CoordinatedPlacementTarget + skill_id: ClassVar[str] = "coordinated_placement" + GoalType: ClassVar[type] = CoordinatedPlacementGoal + OptionsType: ClassVar[type] = CoordinatedPlacementOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("placing", "support") + end_effector_roles: ClassVar[tuple[str, ...]] = ("placing", "support") def __init__( self, - motion_generator, - cfg: CoordinatedPlacementCfg | None = None, + default_options: CoordinatedPlacementOptions | None = None, ) -> None: - super().__init__(motion_generator, cfg or CoordinatedPlacementCfg()) - if ( - self.cfg.motion_source == "motion_gen" - and self.motion_generator.planner.cfg.planner_type == "curobo" - ): - logger.log_error( - "Coordinated dual-arm planning is not supported by the cuRobo " - "backend. Use a single-arm action or a dedicated multi-arm " - "planner.", - ValueError, - ) - self.builder = TrajectoryBuilder(motion_generator) + super().__init__(default_options) + + def _on_bind(self) -> None: + """Resolve engine-wide resources from the owning engine.""" self.n_envs = self.robot.get_qpos().shape[0] self.robot_dof = self.robot.dof - self.dual_arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.placing_arm_joint_ids = self.robot.get_joint_ids( - name=self.cfg.placing_arm_control_part - ) - self.support_arm_joint_ids = self.robot.get_joint_ids( - name=self.cfg.support_arm_control_part - ) - self.placing_hand_joint_ids = self.robot.get_joint_ids( - name=self.cfg.placing_hand_control_part - ) - self.support_hand_joint_ids = self.robot.get_joint_ids( - name=self.cfg.support_hand_control_part - ) - self.joint_ids = ( - self.dual_arm_joint_ids - + self.placing_hand_joint_ids - + self.support_hand_joint_ids - ) - self.placing_arm_dof = len(self.placing_arm_joint_ids) - self.support_arm_dof = len(self.support_arm_joint_ids) - self.placing_hand_dof = len(self.placing_hand_joint_ids) - self.support_hand_dof = len(self.support_hand_joint_ids) - - self._validate_hand_qpos_cfg() - self.placing_hand_open_qpos = self.builder.expand_hand_qpos( - self.cfg.placing_hand_open_qpos, - n_envs=self.n_envs, - hand_dof=self.placing_hand_dof, - ) - self.placing_hand_close_qpos = self.builder.expand_hand_qpos( - self.cfg.placing_hand_close_qpos, - n_envs=self.n_envs, - hand_dof=self.placing_hand_dof, - ) - self.support_hand_close_qpos = self.builder.expand_hand_qpos( - self.cfg.support_hand_close_qpos, - n_envs=self.n_envs, - hand_dof=self.support_hand_dof, + def _resolve_resources( + self, + request: ResolvedActionRequest[ + CoordinatedPlacementGoal, CoordinatedPlacementOptions + ], + ) -> _CoordinatedPlacementResources: + """Resolve placing/support roles from robot control parts.""" + binding = request.binding + placing_arm = binding.manipulator("placing") + support_arm = binding.manipulator("support") + placing_hand = binding.end_effector("placing") + support_hand = binding.end_effector("support") + if placing_arm.name == support_arm.name: + raise ValueError( + "CoordinatedPlacement placing and support roles must use " + "different manipulator control parts." + ) + if placing_hand.name == support_hand.name: + raise ValueError( + "CoordinatedPlacement placing and support roles must use " + "different end-effector control parts." + ) + return _CoordinatedPlacementResources( + placing_arm=placing_arm, + support_arm=support_arm, + placing_hand=placing_hand, + support_hand=support_hand, + placing_hand_open_qpos=placing_hand.joint_positions( + OPEN_COMMAND, + n_envs=self.n_envs, + device=self.device, + dtype=torch.float32, + ), + placing_hand_close_qpos=placing_hand.joint_positions( + GRASP_COMMAND, + n_envs=self.n_envs, + device=self.device, + dtype=torch.float32, + ), + support_hand_close_qpos=support_hand.joint_positions( + GRASP_COMMAND, + n_envs=self.n_envs, + device=self.device, + dtype=torch.float32, + ), ) - def execute( - self, target: CoordinatedPlacementTarget, state: WorldState - ) -> ActionResult: + def plan( + self, + request: ResolvedActionRequest[ + CoordinatedPlacementGoal, CoordinatedPlacementOptions + ], + context: PlanningContext, + ) -> ActionPlan: + """Plan coordinated placement without committing attachment changes.""" + target = self.require_goal(request) + options = request.skill_options + resources = self._resolve_resources(request) + if ( + request.motion_policy.motion_source == "motion_gen" + and self.motion_generator.planner.cfg.planner_type == "curobo" + ): + raise ValueError( + "Coordinated dual-arm planning is not supported by the cuRobo backend." + ) + state = context ( 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) + ) = self._resolve_target(target, state, resources, options) + placing_start_qpos, support_start_qpos = self._resolve_start_qpos( + state, resources + ) + segments = self._compute_segment_lengths( + release, request.motion_policy.sample_count, options + ) placing_lift_xpos = self.builder.apply_local_offset( placing_xpos, torch.tensor( - [0.0, 0.0, self.cfg.lift_height], + [0.0, 0.0, options.lift_height], dtype=torch.float32, device=self.device, ), ) ok, placing_approach_traj = self._plan_named_arm_trajectory( - self.cfg.placing_arm_control_part, + resources.placing_arm.name, placing_start_qpos, torch.stack([placing_lift_xpos, placing_xpos], dim=1), segments["approach"], + request.motion_policy, ) if not ok: logger.log_warning("CoordinatedPlacement failed to plan placing approach.") - return self._fail(state) + return self.failed_plan( + request, context, message="Placing approach failed." + ) ok, support_approach_traj = self._plan_named_arm_trajectory( - self.cfg.support_arm_control_part, + resources.support_arm.name, support_start_qpos, support_xpos.unsqueeze(1), segments["approach"], + request.motion_policy, ) if not ok: logger.log_warning("CoordinatedPlacement failed to plan support approach.") - return self._fail(state) + return self.failed_plan( + request, context, message="Support approach failed." + ) placing_place_qpos = placing_approach_traj[:, -1] support_place_qpos = support_approach_traj[:, -1] @@ -238,8 +261,9 @@ def execute( state.last_qpos, placing_approach_traj, support_approach_traj, - self._repeat_qpos(self.placing_hand_close_qpos, segments["approach"]), - self._repeat_qpos(self.support_hand_close_qpos, segments["approach"]), + self._repeat_qpos(resources.placing_hand_close_qpos, segments["approach"]), + self._repeat_qpos(resources.support_hand_close_qpos, segments["approach"]), + resources=resources, ) hold_trajectory = self._empty_phase() @@ -248,8 +272,9 @@ def execute( state.last_qpos, self._repeat_qpos(placing_place_qpos, segments["hold"]), self._repeat_qpos(support_place_qpos, segments["hold"]), - self._repeat_qpos(self.placing_hand_close_qpos, segments["hold"]), - self._repeat_qpos(self.support_hand_close_qpos, segments["hold"]), + self._repeat_qpos(resources.placing_hand_close_qpos, segments["hold"]), + self._repeat_qpos(resources.support_hand_close_qpos, segments["hold"]), + resources=resources, ) release_trajectory = self._empty_phase() @@ -259,32 +284,39 @@ def execute( self._repeat_qpos(placing_place_qpos, segments["release"]), self._repeat_qpos(support_place_qpos, segments["release"]), self.builder.interpolate_hand_qpos( - self.placing_hand_close_qpos, - self.placing_hand_open_qpos, + resources.placing_hand_close_qpos, + resources.placing_hand_open_qpos, n_waypoints=segments["release"], ), - self._repeat_qpos(self.support_hand_close_qpos, segments["release"]), + self._repeat_qpos( + resources.support_hand_close_qpos, segments["release"] + ), + resources=resources, ) ok, placing_retreat_traj = self._plan_named_arm_trajectory( - self.cfg.placing_arm_control_part, + resources.placing_arm.name, placing_place_qpos, placing_lift_xpos.unsqueeze(1), segments["retreat"], + request.motion_policy, ) if not ok: logger.log_warning("CoordinatedPlacement failed to plan placing retreat.") - return self._fail(state) + return self.failed_plan(request, context, message="Placing retreat failed.") placing_hand_retreat_qpos = ( - self.placing_hand_open_qpos if release else self.placing_hand_close_qpos + resources.placing_hand_open_qpos + if release + else resources.placing_hand_close_qpos ) retreat_trajectory = self._assemble_phase( state.last_qpos, placing_retreat_traj, self._repeat_qpos(support_place_qpos, segments["retreat"]), self._repeat_qpos(placing_hand_retreat_qpos, segments["retreat"]), - self._repeat_qpos(self.support_hand_close_qpos, segments["retreat"]), + self._repeat_qpos(resources.support_hand_close_qpos, segments["retreat"]), + resources=resources, ) full = torch.cat( @@ -296,52 +328,40 @@ 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, + resources.placing_arm.name, + resources.support_arm.name, } - coordinated_held_objects = { - key: value - for key, value in state.coordinated_held_objects.items() - if involved_control_parts.isdisjoint(key) + coordinated_removals = { + key: None + for key in state.coordinated_held_objects + if not involved_control_parts.isdisjoint(key) } - return ActionResult( + return self.build_plan( + request, + context, success=True, trajectory=full, - next_state=state.with_updates( - last_qpos=full[:, -1, :].clone(), - held_objects=held_objects, - coordinated_held_objects=coordinated_held_objects, + expected_effects=StateDelta( + held_object_updates={ + resources.placing_arm.name: ( + None if release else placing_held_object + ), + resources.support_arm.name: support_held_object, + }, + coordinated_held_object_updates=coordinated_removals, ), ) - def _validate_hand_qpos_cfg(self) -> None: - required_names = ( - "placing_hand_open_qpos", - "placing_hand_close_qpos", - "support_hand_close_qpos", - ) - for name in required_names: - if getattr(self.cfg, name) is None: - logger.log_error( - f"{name} must be specified in CoordinatedPlacementCfg", - ValueError, - ) - def _resolve_object_pose( self, - pose: torch.Tensor, + pose: PoseGoalValue, height_offset: float, name: str, + context: PlanningContext, ) -> torch.Tensor: object_pose = resolve_object_target( - pose, + resolve_pose_goal(pose, context, name=name), n_envs=self.n_envs, device=self.device, name=name, @@ -394,8 +414,10 @@ def _resolve_held_state( def _resolve_target( self, - target: CoordinatedPlacementTarget, - state: WorldState, + target: CoordinatedPlacementGoal, + state: PlanningContext, + resources: _CoordinatedPlacementResources, + options: CoordinatedPlacementOptions, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -403,27 +425,29 @@ def _resolve_target( HeldObjectState, HeldObjectState, ]: - placing_held_object = state.get_held_object(self.cfg.placing_arm_control_part) + placing_control_part = resources.placing_arm.name + support_control_part = resources.support_arm.name + placing_held_object = state.get_held_object(placing_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}.", + f"part {placing_control_part!r}.", ValueError, ) - support_held_object = state.get_held_object(self.cfg.support_arm_control_part) + support_held_object = state.get_held_object(support_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}.", + f"part {support_control_part!r}.", ValueError, ) placing_height_offset = ( - self.cfg.placing_height_offset + options.placing_height_offset if target.placing_height_offset is None else target.placing_height_offset ) support_height_offset = ( - self.cfg.support_height_offset + options.support_height_offset if target.support_height_offset is None else target.support_height_offset ) @@ -431,11 +455,13 @@ def _resolve_target( target.placing_object_target_pose, placing_height_offset, "placing_object_target_pose", + state, ) support_object_pose = self._resolve_object_pose( target.support_object_target_pose, support_height_offset, "support_object_target_pose", + state, ) placing_object_to_eef = self._resolve_object_to_eef( placing_held_object, @@ -447,7 +473,7 @@ def _resolve_target( ) placing_xpos = torch.bmm(placing_object_pose, placing_object_to_eef) support_xpos = torch.bmm(support_object_pose, support_object_to_eef) - release = self.cfg.release if target.release is None else target.release + release = options.release if target.release is None else target.release return ( placing_xpos, support_xpos, @@ -465,29 +491,38 @@ def _resolve_target( ) def _resolve_start_qpos( - self, state: WorldState + self, + state: PlanningContext, + resources: _CoordinatedPlacementResources, ) -> tuple[torch.Tensor, torch.Tensor]: if state.last_qpos.shape != (self.n_envs, self.robot_dof): logger.log_error( - f"WorldState.last_qpos must have shape ({self.n_envs}, {self.robot_dof}), " + "PlanningContext.last_qpos must have shape " + f"({self.n_envs}, {self.robot_dof}), " f"but got {state.last_qpos.shape}", ValueError, ) start_qpos = state.last_qpos.to(device=self.device, dtype=torch.float32) return ( - start_qpos[:, self.placing_arm_joint_ids], - start_qpos[:, self.support_arm_joint_ids], + start_qpos[:, list(resources.placing_arm.joint_ids)], + start_qpos[:, list(resources.support_arm.joint_ids)], ) - def _compute_segment_lengths(self, release: bool) -> dict[str, int]: - n_release = max(2, self.cfg.hand_interp_steps) if release else 0 - n_hold = max(0, self.cfg.hold_steps) - n_retreat = max(2, self.cfg.retreat_steps) - n_approach = self.cfg.sample_interval - n_hold - n_release - n_retreat + def _compute_segment_lengths( + self, + release: bool, + sample_count: int, + options: CoordinatedPlacementOptions, + ) -> dict[str, int]: + """Split the invocation sample budget across placement phases.""" + n_release = max(2, options.hand_interp_steps) if release else 0 + n_hold = max(0, options.hold_steps) + n_retreat = max(2, options.retreat_steps) + n_approach = sample_count - n_hold - n_release - n_retreat if n_approach < 2: logger.log_error( "Not enough waypoints for coordinated placement. Increase " - "sample_interval or decrease hold/release/retreat steps.", + "sample_count or decrease hold/release/retreat steps.", ValueError, ) return { @@ -503,6 +538,7 @@ def _plan_named_arm_trajectory( start_qpos: torch.Tensor, target_poses: torch.Tensor, n_waypoints: int, + motion_policy: MotionPolicy, ) -> tuple[bool, torch.Tensor]: target_states_list = [ [ @@ -517,6 +553,7 @@ def _plan_named_arm_trajectory( n_waypoints, control_part=control_part, arm_dof=start_qpos.shape[-1], + cfg=motion_policy, ) return self.builder.all_envs_success(success), trajectory @@ -538,30 +575,21 @@ def _assemble_phase( support_arm_traj: torch.Tensor, placing_hand_traj: torch.Tensor, support_hand_traj: torch.Tensor, + *, + resources: _CoordinatedPlacementResources, ) -> torch.Tensor: n_waypoints = placing_arm_traj.shape[1] full = base_full_qpos.to(device=self.device, dtype=torch.float32) full = full.unsqueeze(1).repeat(1, n_waypoints, 1).clone() - full[:, :, self.placing_arm_joint_ids] = placing_arm_traj - full[:, :, self.support_arm_joint_ids] = support_arm_traj - full[:, :, self.placing_hand_joint_ids] = placing_hand_traj - full[:, :, self.support_hand_joint_ids] = support_hand_traj + full[:, :, list(resources.placing_arm.joint_ids)] = placing_arm_traj + full[:, :, list(resources.support_arm.joint_ids)] = support_arm_traj + full[:, :, list(resources.placing_hand.joint_ids)] = placing_hand_traj + full[:, :, list(resources.support_hand.joint_ids)] = support_hand_traj return full - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=False, - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, - ) - __all__ = [ "CoordinatedPlacement", - "CoordinatedPlacementCfg", - "CoordinatedPlacementTarget", + "CoordinatedPlacementGoal", + "CoordinatedPlacementOptions", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/hand_over.py b/embodichain/lab/sim/atomic_actions/primitives/hand_over.py index a05758ce6..551f84732 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/hand_over.py +++ b/embodichain/lab/sim/atomic_actions/primitives/hand_over.py @@ -18,57 +18,29 @@ from __future__ import annotations +from dataclasses import dataclass from typing import ClassVar import torch from embodichain.lab.sim.planners import MoveType, PlanState -from embodichain.utils import configclass, logger +from embodichain.utils import logger from embodichain.utils.math import pose_inv -from ..core import ( - ActionCfg, - ActionResult, - AtomicAction, - HeldObjectState, - ObjectSemantics, - WorldState, -) -from .pick_up import GraspTarget -from ..trajectory import TrajectoryBuilder +from ..bindings import ResolvedControlPart +from ..control import GRASP_COMMAND, OPEN_COMMAND +from ..core import AtomicAction, ObjectSemantics +from ..effects import StateDelta +from ..invocation import ActionOptions, ResolvedActionRequest +from ..plans import ActionPlan +from ..policies import MotionPolicy +from ..state import HeldObjectState, PlanningContext +from .pick_up import GraspGoal -@configclass -class HandOverCfg(ActionCfg): - name: str = "hand_over" - """Name of the action, used for identification and logging.""" - - control_part: str = "dual_arm" - """Combined control part containing both the transferring and receiving arms.""" - - transfer_arm_control_part: str = "left_arm" - """Arm that already holds the object and hands it over (the 'a' arm).""" - - receive_arm_control_part: str = "right_arm" - """Arm that grasps the object and carries it away (the 'b' arm).""" - - transfer_hand_control_part: str = "left_hand" - """Hand attached to the transferring arm.""" - - receive_hand_control_part: str = "right_hand" - """Hand attached to the receiving arm.""" - - transfer_hand_open_qpos: torch.Tensor | None = None - """Transferring-hand qpos for the open (released) state, shape ``[hand_dof,]``.""" - - transfer_hand_close_qpos: torch.Tensor | None = None - """Transferring-hand qpos for the closed (holding) state, shape ``[hand_dof,]``.""" - - receive_hand_open_qpos: torch.Tensor | None = None - """Receiving-hand qpos for the open state, shape ``[hand_dof,]``.""" - - receive_hand_close_qpos: torch.Tensor | None = None - """Receiving-hand qpos for the closed state, shape ``[hand_dof,]``.""" +@dataclass(frozen=True, slots=True, eq=False) +class HandOverOptions(ActionOptions): + """Per-invocation handover behavior and object-pose targets.""" receive_pick_object_part: str = "bottom" """Object part the receiving arm grasps during the handover @@ -96,9 +68,6 @@ class HandOverCfg(ActionCfg): lift_height: float = 0.08 """World-Z lift distance for the transferring arm after it releases.""" - sample_interval: int = 120 - """Number of waypoints for the full handover trajectory.""" - hand_interp_steps: int = 10 """Number of waypoints used for the receiving-hand close and the transferring-hand release interpolations.""" @@ -109,8 +78,48 @@ class HandOverCfg(ActionCfg): retreat_steps: int = 24 """Number of waypoints used for the final deliver/retreat phase.""" + def __post_init__(self) -> None: + if not isinstance(self.receive_pick_object_part, str) or not ( + self.receive_pick_object_part + ): + raise ValueError("receive_pick_object_part must be non-empty.") + if self.receive_approach_direction.shape != (3,): + raise ValueError("receive_approach_direction must have shape (3,).") + if not torch.isfinite(self.receive_approach_direction).all() or ( + torch.linalg.vector_norm(self.receive_approach_direction) <= 1.0e-6 + ): + raise ValueError("receive_approach_direction must be finite and non-zero.") + if self.pre_grasp_distance < 0.0 or self.lift_height < 0.0: + raise ValueError("pre_grasp_distance and lift_height must be non-negative.") + for name in ("hand_interp_steps", "hold_steps", "retreat_steps"): + if getattr(self, name) < 0: + raise ValueError(f"{name} must be non-negative.") + object.__setattr__( + self, + "receive_approach_direction", + self.receive_approach_direction.clone(), + ) + for name in ("middle_object_pose", "final_object_pose"): + value = getattr(self, name) + if value is not None: + object.__setattr__(self, name, value.clone()) + + +@dataclass(frozen=True, slots=True, eq=False) +class _HandOverResources: + """Invocation-bound control parts and compatible hand commands.""" + + transfer_arm: ResolvedControlPart + receive_arm: ResolvedControlPart + transfer_hand: ResolvedControlPart + receive_hand: ResolvedControlPart + transfer_hand_open_qpos: torch.Tensor + transfer_hand_close_qpos: torch.Tensor + receive_hand_open_qpos: torch.Tensor + receive_hand_close_qpos: torch.Tensor + -class HandOver(AtomicAction[GraspTarget]): +class HandOver(AtomicAction[GraspGoal, HandOverOptions]): """Hand an object from one arm to the other. The transferring arm (already holding the object) moves it to a middle @@ -119,93 +128,115 @@ class HandOver(AtomicAction[GraspTarget]): arm carries the object to a final pose. """ - TargetType: ClassVar[type] = GraspTarget + skill_id: ClassVar[str] = "hand_over" + GoalType: ClassVar[type] = GraspGoal + OptionsType: ClassVar[type] = HandOverOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("source", "destination") + end_effector_roles: ClassVar[tuple[str, ...]] = ("source", "destination") def __init__( self, - motion_generator, - cfg: HandOverCfg | None = None, + default_options: HandOverOptions | None = None, ) -> None: - super().__init__(motion_generator, cfg or HandOverCfg()) - if ( - self.cfg.motion_source == "motion_gen" - and self.motion_generator.planner.cfg.planner_type == "curobo" - ): - logger.log_error( - "Coordinated dual-arm planning is not supported by the cuRobo " - "backend. Use a single-arm action or a dedicated multi-arm " - "planner.", - ValueError, - ) - self.builder = TrajectoryBuilder(motion_generator) + super().__init__(default_options) + + def _on_bind(self) -> None: + """Resolve engine-wide resources from the owning engine.""" self.n_envs = self.robot.get_qpos().shape[0] self.robot_dof = self.robot.dof - self._validate_pose_cfg() - self.middle_object_pose = self._resolve_matrix( - self.cfg.middle_object_pose, "middle_object_pose" - ) - self.final_object_pose = self._resolve_matrix( - self.cfg.final_object_pose, "final_object_pose" - ) - - self.dual_arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.transfer_arm_joint_ids = self.robot.get_joint_ids( - name=self.cfg.transfer_arm_control_part - ) - self.receive_arm_joint_ids = self.robot.get_joint_ids( - name=self.cfg.receive_arm_control_part - ) - self.transfer_hand_joint_ids = self.robot.get_joint_ids( - name=self.cfg.transfer_hand_control_part - ) - self.receive_hand_joint_ids = self.robot.get_joint_ids( - name=self.cfg.receive_hand_control_part - ) - self.transfer_arm_dof = len(self.transfer_arm_joint_ids) - self.receive_arm_dof = len(self.receive_arm_joint_ids) - self.transfer_hand_dof = len(self.transfer_hand_joint_ids) - self.receive_hand_dof = len(self.receive_hand_joint_ids) - - self._validate_hand_qpos_cfg() - self.transfer_hand_open_qpos = self.builder.expand_hand_qpos( - self.cfg.transfer_hand_open_qpos, - n_envs=self.n_envs, - hand_dof=self.transfer_hand_dof, - ) - self.transfer_hand_close_qpos = self.builder.expand_hand_qpos( - self.cfg.transfer_hand_close_qpos, - n_envs=self.n_envs, - hand_dof=self.transfer_hand_dof, - ) - self.receive_hand_open_qpos = self.builder.expand_hand_qpos( - self.cfg.receive_hand_open_qpos, - n_envs=self.n_envs, - hand_dof=self.receive_hand_dof, - ) - self.receive_hand_close_qpos = self.builder.expand_hand_qpos( - self.cfg.receive_hand_close_qpos, - n_envs=self.n_envs, - hand_dof=self.receive_hand_dof, - ) - - approach = self.cfg.receive_approach_direction.to( - device=self.device, dtype=torch.float32 + def _resolve_resources( + self, + request: ResolvedActionRequest[GraspGoal, HandOverOptions], + ) -> _HandOverResources: + """Resolve source/destination roles from robot control parts.""" + binding = request.binding + transfer_arm = binding.manipulator("source") + receive_arm = binding.manipulator("destination") + transfer_hand = binding.end_effector("source") + receive_hand = binding.end_effector("destination") + if transfer_arm.name == receive_arm.name: + raise ValueError( + "HandOver source and destination must use different manipulator " + "control parts." + ) + if transfer_hand.name == receive_hand.name: + raise ValueError( + "HandOver source and destination must use different end-effector " + "control parts." + ) + return _HandOverResources( + transfer_arm=transfer_arm, + receive_arm=receive_arm, + transfer_hand=transfer_hand, + receive_hand=receive_hand, + transfer_hand_open_qpos=transfer_hand.joint_positions( + OPEN_COMMAND, + n_envs=self.n_envs, + device=self.device, + dtype=torch.float32, + ), + transfer_hand_close_qpos=transfer_hand.joint_positions( + GRASP_COMMAND, + n_envs=self.n_envs, + device=self.device, + dtype=torch.float32, + ), + receive_hand_open_qpos=receive_hand.joint_positions( + OPEN_COMMAND, + n_envs=self.n_envs, + device=self.device, + dtype=torch.float32, + ), + receive_hand_close_qpos=receive_hand.joint_positions( + GRASP_COMMAND, + n_envs=self.n_envs, + device=self.device, + dtype=torch.float32, + ), ) - approach_norm = torch.linalg.vector_norm(approach) - if approach_norm <= 1.0e-6: - logger.log_error("receive_approach_direction must be non-zero.", ValueError) - self.receive_approach_direction = approach / approach_norm # ------------------------------------------------------------------ # Public contract # ------------------------------------------------------------------ - def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: + def plan( + self, + request: ResolvedActionRequest[GraspGoal, HandOverOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan a handover without committing the attachment transfer.""" + target = self.require_goal(request) + options = request.skill_options + self._validate_pose_options(options) + resources = self._resolve_resources(request) + if ( + request.motion_policy.motion_source == "motion_gen" + and self.motion_generator.planner.cfg.planner_type == "curobo" + ): + raise ValueError( + "Coordinated dual-arm planning is not supported by the cuRobo backend." + ) + state = context semantics = target.semantics - transfer_object_to_eef = self._resolve_transfer_object_to_eef(state) - middle_object_pose = self.middle_object_pose.clone() - final_object_pose = self.final_object_pose + transfer_object_to_eef = self._resolve_transfer_object_to_eef( + state, resources.transfer_arm.name + ) + assert options.middle_object_pose is not None + assert options.final_object_pose is not None + middle_object_pose = self._resolve_matrix( + options.middle_object_pose, "middle_object_pose" + ) + final_object_pose = self._resolve_matrix( + options.final_object_pose, "final_object_pose" + ) + receive_approach_direction = options.receive_approach_direction.to( + device=self.device, dtype=torch.float32 + ) + receive_approach_direction = ( + receive_approach_direction + / torch.linalg.vector_norm(receive_approach_direction) + ) # force object pose to have the same rotation as the current object pose, so that the handover is feasible. current_object_pose = target.semantics.entity.get_local_pose(to_matrix=True) middle_object_pose[:, :3, :3] = current_object_pose[:, :3, :3] @@ -216,18 +247,21 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: # 2.2 - receiving grasp on the requested object part at the handover pose. receive_grasp_xpos, grasp_success = self._resolve_receive_grasp( - semantics, middle_object_pose, self.cfg.receive_pick_object_part + semantics, + middle_object_pose, + options.receive_pick_object_part, + receive_approach_direction, ) if not self.builder.all_envs_success(grasp_success): logger.log_warning("HandOver failed to resolve a receiving grasp pose.") - return self._fail(state) + return self.failed_plan(request, context, message="No receiving grasp.") receive_object_to_eef = torch.bmm( pose_inv(middle_object_pose), receive_grasp_xpos ) receive_grasp_z = receive_grasp_xpos[..., :3, 2] receive_pre_grasp_eef = self.builder.apply_local_offset( receive_grasp_xpos, - -receive_grasp_z * self.cfg.pre_grasp_distance, + -receive_grasp_z * options.pre_grasp_distance, ) # 2.4 - receiving arm delivers the object to the final pose. receive_final_eef = torch.bmm(final_object_pose, receive_object_to_eef) @@ -235,57 +269,71 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: transfer_retreat_eef = self.builder.apply_local_offset( transfer_middle_eef, torch.tensor( - [0.0, 0.0, self.cfg.lift_height], + [0.0, 0.0, options.lift_height], dtype=torch.float32, device=self.device, ), ) - transfer_start_qpos, receive_start_qpos = self._resolve_start_qpos(state) - segments = self._compute_segment_lengths() + transfer_start_qpos, receive_start_qpos = self._resolve_start_qpos( + state, resources + ) + segments = self._compute_segment_lengths( + request.motion_policy.sample_count, options + ) ok, transfer_move_traj = self._plan_named_arm_trajectory( - self.cfg.transfer_arm_control_part, + resources.transfer_arm.name, transfer_start_qpos, transfer_middle_eef.unsqueeze(1), segments["transfer"], + request.motion_policy, ) if not ok: logger.log_warning("HandOver failed to plan the transfer move.") - return self._fail(state) + return self.failed_plan(request, context, message="Transfer move failed.") ok, receive_approach_traj = self._plan_named_arm_trajectory( - self.cfg.receive_arm_control_part, + resources.receive_arm.name, receive_start_qpos, torch.stack([receive_pre_grasp_eef, receive_grasp_xpos], dim=1), segments["approach"], + request.motion_policy, ) if not ok: logger.log_warning("HandOver failed to plan the receiving approach.") - return self._fail(state) + return self.failed_plan( + request, context, message="Receiving approach failed." + ) transfer_hold_qpos = transfer_move_traj[:, -1] receive_grasp_qpos = receive_approach_traj[:, -1] ok, transfer_retreat_traj = self._plan_named_arm_trajectory( - self.cfg.transfer_arm_control_part, + resources.transfer_arm.name, transfer_hold_qpos, transfer_retreat_eef.unsqueeze(1), segments["deliver"], + request.motion_policy, ) if not ok: logger.log_warning("HandOver failed to plan the transfer retreat.") - return self._fail(state) + return self.failed_plan( + request, context, message="Transfer retreat failed." + ) ok, receive_deliver_traj = self._plan_named_arm_trajectory( - self.cfg.receive_arm_control_part, + resources.receive_arm.name, receive_grasp_qpos, receive_final_eef.unsqueeze(1), segments["deliver"], + request.motion_policy, ) if not ok: logger.log_warning("HandOver failed to plan the receiving delivery.") - return self._fail(state) + return self.failed_plan( + request, context, message="Receiving delivery failed." + ) phases: list[torch.Tensor] = [] # 2.1 transfer: transferring arm carries the object to the middle pose. @@ -294,8 +342,13 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: state, transfer_move_traj, self._repeat_qpos(receive_start_qpos, segments["transfer"]), - self._repeat_qpos(self.transfer_hand_close_qpos, segments["transfer"]), - self._repeat_qpos(self.receive_hand_open_qpos, segments["transfer"]), + self._repeat_qpos( + resources.transfer_hand_close_qpos, segments["transfer"] + ), + self._repeat_qpos( + resources.receive_hand_open_qpos, segments["transfer"] + ), + resources=resources, ) ) # 2.2 approach: receiving arm moves to the grasp pose; transferring arm holds. @@ -304,8 +357,13 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: state, self._repeat_qpos(transfer_hold_qpos, segments["approach"]), receive_approach_traj, - self._repeat_qpos(self.transfer_hand_close_qpos, segments["approach"]), - self._repeat_qpos(self.receive_hand_open_qpos, segments["approach"]), + self._repeat_qpos( + resources.transfer_hand_close_qpos, segments["approach"] + ), + self._repeat_qpos( + resources.receive_hand_open_qpos, segments["approach"] + ), + resources=resources, ) ) # 2.2 close: receiving hand closes; transferring arm keeps holding. @@ -314,12 +372,15 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: state, self._repeat_qpos(transfer_hold_qpos, segments["close"]), self._repeat_qpos(receive_grasp_qpos, segments["close"]), - self._repeat_qpos(self.transfer_hand_close_qpos, segments["close"]), + self._repeat_qpos( + resources.transfer_hand_close_qpos, segments["close"] + ), self.builder.interpolate_hand_qpos( - self.receive_hand_open_qpos, - self.receive_hand_close_qpos, + resources.receive_hand_open_qpos, + resources.receive_hand_close_qpos, n_waypoints=segments["close"], ), + resources=resources, ) ) if segments["hold"] > 0: @@ -328,8 +389,13 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: state, self._repeat_qpos(transfer_hold_qpos, segments["hold"]), self._repeat_qpos(receive_grasp_qpos, segments["hold"]), - self._repeat_qpos(self.transfer_hand_close_qpos, segments["hold"]), - self._repeat_qpos(self.receive_hand_close_qpos, segments["hold"]), + self._repeat_qpos( + resources.transfer_hand_close_qpos, segments["hold"] + ), + self._repeat_qpos( + resources.receive_hand_close_qpos, segments["hold"] + ), + resources=resources, ) ) # 2.3 release: transferring hand opens; receiving arm keeps holding. @@ -339,11 +405,14 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: self._repeat_qpos(transfer_hold_qpos, segments["release"]), self._repeat_qpos(receive_grasp_qpos, segments["release"]), self.builder.interpolate_hand_qpos( - self.transfer_hand_close_qpos, - self.transfer_hand_open_qpos, + resources.transfer_hand_close_qpos, + resources.transfer_hand_open_qpos, n_waypoints=segments["release"], ), - self._repeat_qpos(self.receive_hand_close_qpos, segments["release"]), + self._repeat_qpos( + resources.receive_hand_close_qpos, segments["release"] + ), + resources=resources, ) ) # 2.4 deliver: receiving arm carries the object away; transferring arm retreats. @@ -352,8 +421,13 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: state, transfer_retreat_traj, receive_deliver_traj, - self._repeat_qpos(self.transfer_hand_open_qpos, segments["deliver"]), - self._repeat_qpos(self.receive_hand_close_qpos, segments["deliver"]), + self._repeat_qpos( + resources.transfer_hand_open_qpos, segments["deliver"] + ), + self._repeat_qpos( + resources.receive_hand_close_qpos, segments["deliver"] + ), + resources=resources, ) ) full = torch.cat(phases, dim=1) @@ -362,15 +436,16 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: object_to_eef=receive_object_to_eef, grasp_xpos=receive_grasp_xpos, ) - held_objects = dict(state.held_objects) - held_objects.pop(self.cfg.transfer_arm_control_part, None) - held_objects[self.cfg.receive_arm_control_part] = held_object - return ActionResult( + return self.build_plan( + request, + context, success=True, trajectory=full, - next_state=state.with_updates( - last_qpos=full[:, -1, :].clone(), - held_objects=held_objects, + expected_effects=StateDelta( + held_object_updates={ + resources.transfer_arm.name: None, + resources.receive_arm.name: held_object, + } ), ) @@ -378,21 +453,13 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: # Resolution helpers # ------------------------------------------------------------------ - def _validate_hand_qpos_cfg(self) -> None: - required_names = ( - "transfer_hand_open_qpos", - "transfer_hand_close_qpos", - "receive_hand_open_qpos", - "receive_hand_close_qpos", - ) - for name in required_names: - if getattr(self.cfg, name) is None: - logger.log_error(f"{name} must be specified in HandOverCfg", ValueError) - - def _validate_pose_cfg(self) -> None: + @staticmethod + def _validate_pose_options(options: HandOverOptions) -> None: for name in ("middle_object_pose", "final_object_pose"): - if getattr(self.cfg, name) is None: - logger.log_error(f"{name} must be specified in HandOverCfg", ValueError) + if getattr(options, name) is None: + logger.log_error( + f"{name} must be specified in HandOverOptions", ValueError + ) def _resolve_matrix(self, matrix: torch.Tensor, name: str) -> torch.Tensor: matrix = matrix.to(device=self.device, dtype=torch.float32) @@ -406,12 +473,16 @@ def _resolve_matrix(self, matrix: torch.Tensor, name: str) -> torch.Tensor: ) return matrix - def _resolve_transfer_object_to_eef(self, state: WorldState) -> torch.Tensor: - held = state.get_held_object(self.cfg.transfer_arm_control_part) + def _resolve_transfer_object_to_eef( + self, + state: PlanningContext, + transfer_control_part: str, + ) -> torch.Tensor: + held = state.get_held_object(transfer_control_part) if held is None: logger.log_error( "HandOver requires an object held by transfer control part " - f"{self.cfg.transfer_arm_control_part!r} (run PickUp first).", + f"{transfer_control_part!r} (run PickUp first).", ValueError, ) return self._resolve_matrix(held.object_to_eef, "held_object.object_to_eef") @@ -421,11 +492,12 @@ def _resolve_receive_grasp( semantics: ObjectSemantics, object_pose: torch.Tensor, object_part: str, + approach_direction: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Select the lowest-cost receiving grasp on ``object_part`` at ``object_pose``.""" grasp_poses_result = semantics.affordance.get_valid_grasp_poses( obj_poses=object_pose, - approach_direction=self.receive_approach_direction, + approach_direction=approach_direction, object_part=object_part, ) n_envs = object_pose.shape[0] @@ -449,31 +521,36 @@ def _resolve_receive_grasp( return grasp_xpos, is_success def _resolve_start_qpos( - self, state: WorldState + self, + state: PlanningContext, + resources: _HandOverResources, ) -> tuple[torch.Tensor, torch.Tensor]: if state.last_qpos.shape != (self.n_envs, self.robot_dof): logger.log_error( - f"WorldState.last_qpos must have shape " + f"PlanningContext.last_qpos must have shape " f"({self.n_envs}, {self.robot_dof}), but got {state.last_qpos.shape}", ValueError, ) start_qpos = state.last_qpos.to(device=self.device, dtype=torch.float32) return ( - start_qpos[:, self.transfer_arm_joint_ids], - start_qpos[:, self.receive_arm_joint_ids], - ) - - def _compute_segment_lengths(self) -> dict[str, int]: - n_close = max(2, self.cfg.hand_interp_steps) - n_release = max(2, self.cfg.hand_interp_steps) - n_deliver = max(2, self.cfg.retreat_steps) - n_hold = max(0, self.cfg.hold_steps) + start_qpos[:, list(resources.transfer_arm.joint_ids)], + start_qpos[:, list(resources.receive_arm.joint_ids)], + ) + + def _compute_segment_lengths( + self, sample_count: int, options: HandOverOptions + ) -> dict[str, int]: + """Split the invocation sample budget across handover phases.""" + n_close = max(2, options.hand_interp_steps) + n_release = max(2, options.hand_interp_steps) + n_deliver = max(2, options.retreat_steps) + n_hold = max(0, options.hold_steps) reserved = n_close + n_release + n_deliver + n_hold - n_transfer = max(2, (self.cfg.sample_interval - reserved) // 2) - n_approach = self.cfg.sample_interval - reserved - n_transfer + n_transfer = max(2, (sample_count - reserved) // 2) + n_approach = sample_count - reserved - n_transfer if n_approach < 2: logger.log_error( - "Not enough waypoints for handover. Increase sample_interval or " + "Not enough waypoints for handover. Increase sample_count or " "decrease hand_interp_steps/hold_steps/retreat_steps.", ValueError, ) @@ -496,6 +573,7 @@ def _plan_named_arm_trajectory( start_qpos: torch.Tensor, target_poses: torch.Tensor, n_waypoints: int, + motion_policy: MotionPolicy, ) -> tuple[bool, torch.Tensor]: target_states_list = [ [ @@ -510,7 +588,7 @@ def _plan_named_arm_trajectory( n_waypoints, control_part=control_part, arm_dof=start_qpos.shape[-1], - cfg=self.cfg, + cfg=motion_policy, ) return self.builder.all_envs_success(success), trajectory @@ -520,11 +598,13 @@ def _repeat_qpos(qpos: torch.Tensor, n_waypoints: int) -> torch.Tensor: def _assemble_phase( self, - state: WorldState, + state: PlanningContext, transfer_arm_traj: torch.Tensor, receive_arm_traj: torch.Tensor, transfer_hand_traj: torch.Tensor, receive_hand_traj: torch.Tensor, + *, + resources: _HandOverResources, ) -> torch.Tensor: n_waypoints = transfer_arm_traj.shape[1] base = ( @@ -533,22 +613,11 @@ def _assemble_phase( .repeat(1, n_waypoints, 1) .clone() ) - base[:, :, self.transfer_arm_joint_ids] = transfer_arm_traj - base[:, :, self.receive_arm_joint_ids] = receive_arm_traj - base[:, :, self.transfer_hand_joint_ids] = transfer_hand_traj - base[:, :, self.receive_hand_joint_ids] = receive_hand_traj + base[:, :, list(resources.transfer_arm.joint_ids)] = transfer_arm_traj + base[:, :, list(resources.receive_arm.joint_ids)] = receive_arm_traj + base[:, :, list(resources.transfer_hand.joint_ids)] = transfer_hand_traj + base[:, :, list(resources.receive_hand.joint_ids)] = receive_hand_traj return base - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=torch.zeros(self.n_envs, dtype=torch.bool, device=self.device), - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, - ) - -__all__ = ["HandOver", "HandOverCfg"] +__all__ = ["HandOver", "HandOverOptions"] 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 49e7e7f42..ec8a40f87 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py @@ -23,138 +23,110 @@ import torch -from embodichain.lab.sim.planners import MoveType, PlanOptions, PlanState -from embodichain.utils import configclass - -from ._helpers import arm_qpos_from_state -from ..core import ( - ActionTarget, - ActionCfg, - ActionResult, - AtomicAction, - WorldState, - _validate_pose_tensor, -) -from ..trajectory import TrajectoryBuilder +from embodichain.lab.sim.planners import MoveType, PlanState +from ..core import AtomicAction +from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal +from ..invocation import ActionOptions, ResolvedActionRequest +from ..plans import ActionPlan, CompletionConditionKind +from ..state import PlanningContext @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. +class EndEffectorPoseGoal: + """End-effector pose goal with optional batched intermediate waypoints.""" - Accepts: + goal_kind: ClassVar[str] = "end_effector_pose" - - ``(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. - """ + xpos: PoseGoalValue + """Homogeneous pose with shape ``(4,4)``, ``(B,4,4)`` or ``(B,N,4,4)``.""" def __post_init__(self) -> None: - _validate_pose_tensor(self.xpos, "xpos", allow_waypoints=True) - - -@configclass -class MoveEndEffectorCfg(ActionCfg): - name: str = "move_end_effector" - """Name of the action, used for identification and logging.""" + validate_pose_goal(self.xpos, "xpos", allow_waypoints=True) - sample_interval: int = 50 - """Number of waypoints in the planned trajectory.""" - - plan_opts: PlanOptions | None = None - """Optional planner-specific options copied for each motion-generator call.""" +@dataclass(frozen=True, slots=True, eq=False) +class MoveEndEffectorOptions(ActionOptions): + """Per-invocation behavior for :class:`MoveEndEffector`.""" -class MoveEndEffector(AtomicAction[EndEffectorPoseTarget]): - """Plan a free-space end-effector move to a target pose. - The :class:`EndEffectorPoseTarget` may carry either a single waypoint - ``(n_envs, 4, 4)`` (or a broadcastable ``(4, 4)``) or a multi-waypoint - trajectory ``(n_envs, n_waypoint, 4, 4)``. In the multi-waypoint case the - action plans a single trajectory that visits every waypoint in order, - starting from the inherited ``WorldState.last_qpos``; IK is solved for each - waypoint with the previous waypoint's solution as the seed. - """ +class MoveEndEffector(AtomicAction[EndEffectorPoseGoal, MoveEndEffectorOptions]): + """Plan a free-space move for a bound manipulator.""" - TargetType: ClassVar[type] = EndEffectorPoseTarget + skill_id: ClassVar[str] = "move_end_effector" + GoalType: ClassVar[type] = EndEffectorPoseGoal + OptionsType: ClassVar[type] = MoveEndEffectorOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) def __init__( self, - motion_generator, - cfg: MoveEndEffectorCfg | None = None, + default_options: MoveEndEffectorOptions | None = None, ) -> None: - super().__init__(motion_generator, cfg or MoveEndEffectorCfg()) - self.builder = TrajectoryBuilder(motion_generator) - self.n_envs = self.robot.get_qpos().shape[0] - self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.arm_dof = len(self.arm_joint_ids) - self.robot_dof = self.robot.dof - - def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionResult: - move_xpos = self.builder.resolve_pose_target(target.xpos, n_envs=self.n_envs) + super().__init__(default_options) + + def plan( + self, + request: ResolvedActionRequest[EndEffectorPoseGoal, MoveEndEffectorOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan an end-effector pose goal from the observed joint state.""" + goal = self.require_goal(request) + manipulator = request.binding.manipulator("primary") + control_part = manipulator.name + joint_ids = list(manipulator.joint_ids) + arm_dof = manipulator.dof + move_xpos = self.builder.resolve_pose_target( + resolve_pose_goal(goal.xpos, context, name="xpos"), + n_envs=context.batch_size, + ) start_qpos = self.builder.resolve_start_qpos( - arm_qpos_from_state(state, self.arm_joint_ids), - n_envs=self.n_envs, - arm_dof=self.arm_dof, - control_part=self.cfg.control_part, + context.robot.qpos[:, joint_ids], + n_envs=context.batch_size, + arm_dof=arm_dof, + control_part=control_part, ) - target_states_list = self._build_target_states(move_xpos) - success, arm_traj = self.builder.plan_arm_traj( - target_states_list, + target_states = self._build_target_states(move_xpos, context.batch_size) + result = self.builder.generate_arm_plan( + target_states, start_qpos, - self.cfg.sample_interval, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=self.cfg, + request.motion_policy.sample_count, + control_part=control_part, + arm_dof=arm_dof, + cfg=request.motion_policy, + ) + success, trajectory = self.builder.to_full_robot_trajectory( + result, + base_qpos=context.robot.qpos, + joint_ids=joint_ids, + env_ids=context.env_ids, + control_dt=request.motion_policy.control_dt, ) - full = self._embed(arm_traj, state.last_qpos) - return ActionResult( + return self.build_plan( + request, + context, success=success, - trajectory=full, - next_state=state.with_updates( - last_qpos=full[:, -1, :].clone(), - ), + trajectory=trajectory, + completion_kind=CompletionConditionKind.EEF_GOAL_REACHED, ) - def _build_target_states(self, move_xpos: torch.Tensor) -> list[list[PlanState]]: - """Build per-env PlanState lists from a single- or multi-waypoint target.""" + @staticmethod + def _build_target_states( + move_xpos: torch.Tensor, + batch_size: int, + ) -> list[list[PlanState]]: + """Build per-environment planner states for pose waypoints.""" if move_xpos.dim() == 3: move_xpos = move_xpos.unsqueeze(1) - n_waypoint = move_xpos.shape[1] return [ [ PlanState(xpos=move_xpos[i, j], move_type=MoveType.EEF_MOVE) - for j in range(n_waypoint) + for j in range(move_xpos.shape[1]) ] - for i in range(self.n_envs) + for i in range(batch_size) ] - def _embed( - self, arm_traj: torch.Tensor, last_full_qpos: torch.Tensor - ) -> torch.Tensor: - n_wp = arm_traj.shape[1] - full = torch.empty( - (self.n_envs, n_wp, self.robot_dof), - dtype=torch.float32, - device=self.device, - ) - full[:, :, :] = last_full_qpos.unsqueeze(1) - full[:, :, self.arm_joint_ids] = arm_traj - return full - - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=torch.zeros(self.n_envs, dtype=torch.bool, device=self.device), - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, - ) - -__all__ = ["EndEffectorPoseTarget", "MoveEndEffector", "MoveEndEffectorCfg"] +__all__ = [ + "EndEffectorPoseGoal", + "MoveEndEffector", + "MoveEndEffectorOptions", +] 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 c26683351..98f5f18b6 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py @@ -24,49 +24,38 @@ import torch from embodichain.lab.sim.planners import MoveType, PlanState -from embodichain.utils import configclass, logger +from embodichain.utils import logger from embodichain.utils.math import axis_angle_to_rotation_matrix, get_relative_rotation from ._helpers import arm_qpos_from_state, resolve_object_target -from ..core import ( - ActionTarget, - ActionCfg, - ActionResult, - AtomicAction, - WorldState, - _validate_pose_tensor, -) -from ..trajectory import TrajectoryBuilder +from ..control import GRASP_COMMAND +from ..core import AtomicAction +from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal +from ..invocation import ActionOptions, ResolvedActionRequest +from ..plans import ActionPlan +from ..state import PlanningContext @dataclass(frozen=True, slots=True, eq=False) -class HeldObjectPoseTarget(ActionTarget): +class HeldObjectPoseGoal: """Desired pose for the object held by this action's control part.""" - object_target_pose: torch.Tensor + goal_kind: ClassVar[str] = "held_object_pose" + + object_target_pose: PoseGoalValue """Target object pose, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" def __post_init__(self) -> None: - _validate_pose_tensor( + validate_pose_goal( self.object_target_pose, "object_target_pose", allow_waypoints=False, ) -@configclass -class MoveHeldObjectCfg(ActionCfg): - name: str = "move_held_object" - """Name of the action, used for identification and logging.""" - - sample_interval: int = 50 - """Number of waypoints in the planned trajectory.""" - - hand_control_part: str = "hand" - """Name of the robot part that controls the hand joints.""" - - hand_close_qpos: torch.Tensor | None = None - """Joint positions for the closed hand state, shape ``[hand_dof,]``.""" +@dataclass(frozen=True, slots=True, eq=False) +class MoveHeldObjectOptions(ActionOptions): + """Per-invocation held-object transport behavior.""" obj_upright_direction: torch.Tensor | None = None """Optional object-local direction to align with world up while moving.""" @@ -74,56 +63,92 @@ class MoveHeldObjectCfg(ActionCfg): pick_rotate_upright: float | None = None """Optional rotation in radians used by the legacy upright transport mode.""" + def __post_init__(self) -> None: + if self.obj_upright_direction is not None: + if ( + self.obj_upright_direction.shape != (3,) + or not torch.isfinite(self.obj_upright_direction).all() + ): + raise ValueError( + "obj_upright_direction must be a finite tensor with shape (3,)." + ) + object.__setattr__( + self, "obj_upright_direction", self.obj_upright_direction.clone() + ) + -class MoveHeldObject(AtomicAction[HeldObjectPoseTarget]): +class MoveHeldObject(AtomicAction[HeldObjectPoseGoal, MoveHeldObjectOptions]): """Move the held object to a target object pose; keep the gripper closed.""" - TargetType: ClassVar[type] = HeldObjectPoseTarget + skill_id: ClassVar[str] = "move_held_object" + GoalType: ClassVar[type] = HeldObjectPoseGoal + OptionsType: ClassVar[type] = MoveHeldObjectOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) def __init__( self, - motion_generator, - cfg: MoveHeldObjectCfg | None = None, + default_options: MoveHeldObjectOptions | None = None, ) -> None: - super().__init__(motion_generator, cfg or MoveHeldObjectCfg()) - self.builder = TrajectoryBuilder(motion_generator) + super().__init__(default_options) + + def _on_bind(self) -> None: + """Resolve engine-wide resources from the owning engine.""" self.n_envs = self.robot.get_qpos().shape[0] - self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part) - self.arm_dof = len(self.arm_joint_ids) self.robot_dof = self.robot.dof - if self.cfg.hand_close_qpos is None: - logger.log_error( - "hand_close_qpos must be specified in MoveHeldObjectCfg", ValueError - ) - self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) - - def execute(self, target: HeldObjectPoseTarget, state: WorldState) -> ActionResult: - held_object = state.get_held_object(self.cfg.control_part) + def plan( + self, + request: ResolvedActionRequest[HeldObjectPoseGoal, MoveHeldObjectOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan held-object transport without changing the attachment relation.""" + target = self.require_goal(request) + options = request.skill_options + binding = request.binding + manipulator = binding.manipulator() + end_effector = binding.end_effector() + control_part = manipulator.name + arm_joint_ids = list(manipulator.joint_ids) + hand_joint_ids = list(end_effector.joint_ids) + hand_grasp_qpos = end_effector.joint_positions( + GRASP_COMMAND, + n_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + state = context + held_object = state.get_held_object(control_part) if held_object is None: logger.log_error( "MoveHeldObject requires an object held by control part " - f"{self.cfg.control_part!r} - run PickUp first.", + f"{control_part!r} - run PickUp first.", ValueError, ) object_target_pose = resolve_object_target( - target.object_target_pose, n_envs=self.n_envs, device=self.device + resolve_pose_goal( + target.object_target_pose, + context, + name="object_target_pose", + ), + n_envs=self.n_envs, + device=self.device, ) start_arm_qpos = self.builder.resolve_start_qpos( - arm_qpos_from_state(state, self.arm_joint_ids), + arm_qpos_from_state(state, arm_joint_ids), n_envs=self.n_envs, - arm_dof=self.arm_dof, - control_part=self.cfg.control_part, + arm_dof=manipulator.dof, + control_part=control_part, ) end_arm_xpos = self.robot.compute_fk( - start_arm_qpos, name=self.cfg.control_part, to_matrix=True + start_arm_qpos, name=control_part, to_matrix=True ) - if self.cfg.pick_rotate_upright is not None: + if options.pick_rotate_upright is not None: self._apply_configured_upright_rotation( object_target_pose, end_arm_xpos, held_object.semantics.entity.get_local_pose(to_matrix=True), + options, ) object_to_eef = held_object.object_to_eef.to( device=self.device, dtype=torch.float32 @@ -132,7 +157,7 @@ def execute(self, target: HeldObjectPoseTarget, state: WorldState) -> ActionResu object_to_eef = object_to_eef.unsqueeze(0).repeat(self.n_envs, 1, 1) move_eef_xpos = torch.bmm(object_target_pose, object_to_eef) - if self.cfg.pick_rotate_upright is None: + if options.pick_rotate_upright is None: self._apply_automatic_transport_rotation(move_eef_xpos, end_arm_xpos) target_states_list = [ @@ -142,10 +167,10 @@ def execute(self, target: HeldObjectPoseTarget, state: WorldState) -> ActionResu success, arm_traj = self.builder.plan_arm_traj( target_states_list, start_arm_qpos, - self.cfg.sample_interval, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=self.cfg, + request.motion_policy.sample_count, + control_part=control_part, + arm_dof=manipulator.dof, + cfg=request.motion_policy, ) full = torch.empty( @@ -154,15 +179,15 @@ def execute(self, target: HeldObjectPoseTarget, state: WorldState) -> ActionResu device=self.device, ) full[:, :, :] = state.last_qpos.unsqueeze(1) - full[:, :, self.arm_joint_ids] = arm_traj - full[:, :, self.hand_joint_ids] = self.hand_close_qpos + full[:, :, arm_joint_ids] = arm_traj + full[:, :, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1) - return ActionResult( + return self.build_plan( + request, + context, success=success, trajectory=full, - next_state=state.with_updates( - last_qpos=full[:, -1, :].clone(), - ), + phase_name="transport", ) def _apply_configured_upright_rotation( @@ -170,20 +195,21 @@ def _apply_configured_upright_rotation( object_target_pose: torch.Tensor, end_arm_xpos: torch.Tensor, held_object_xpos: torch.Tensor, + options: MoveHeldObjectOptions, ) -> None: - if self.cfg.obj_upright_direction is None: + if options.obj_upright_direction is None: upright_direction = torch.tensor( [0.0, 0.0, 1.0], device=self.device, dtype=torch.float32 ) else: - upright_direction = self.cfg.obj_upright_direction.to( + upright_direction = options.obj_upright_direction.to( device=self.device, dtype=torch.float32 ) object_upright = torch.matmul(held_object_xpos[:, :3, :3], upright_direction) dot_result = torch.sum(end_arm_xpos[:, :3, 1] * object_upright, dim=-1) revert_flag = torch.where(dot_result < 0, 1.0, -1.0) axis_angle = ( - -float(self.cfg.pick_rotate_upright) + -float(options.pick_rotate_upright) * revert_flag.unsqueeze(-1) * end_arm_xpos[:, :3, 0] ) @@ -244,16 +270,9 @@ def _apply_automatic_transport_rotation( move_eef_xpos[:, :3, :3], ) - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=torch.zeros(self.n_envs, dtype=torch.bool, device=self.device), - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, - ) - -__all__ = ["HeldObjectPoseTarget", "MoveHeldObject", "MoveHeldObjectCfg"] +__all__ = [ + "HeldObjectPoseGoal", + "MoveHeldObject", + "MoveHeldObjectOptions", +] diff --git a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py index e619c113e..1397265a0 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py @@ -23,163 +23,130 @@ import torch -from embodichain.utils import configclass, logger - -from ..core import ( - ActionTarget, - ActionCfg, - ActionResult, - AtomicAction, - WorldState, -) -from ..trajectory import TrajectoryBuilder +from ..core import AtomicAction +from ..invocation import ActionOptions, ResolvedActionRequest +from ..plans import ActionPlan, CompletionConditionKind +from ..state import PlanningContext @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. +class JointPositionGoal: + """Explicit or named joint-space goal for a bound robot resource.""" - Accepts: + goal_kind: ClassVar[str] = "joint_position" - - ``(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. - """ + target: torch.Tensor | str + """Joint qpos/waypoints or a named control-part profile command.""" def __post_init__(self) -> None: - if not isinstance(self.qpos, torch.Tensor): + if isinstance(self.target, str): + if not self.target.strip(): + raise ValueError("Named joint-position target must not be empty.") + return + if not isinstance(self.target, torch.Tensor): raise TypeError( - f"qpos must be a torch.Tensor, got {type(self.qpos).__name__}." + "target must be a torch.Tensor or str, " + f"got {type(self.target).__name__}." ) - if self.qpos.dim() not in (1, 2, 3) or self.qpos.shape[-1] == 0: + if self.target.dim() not in (1, 2, 3) or self.target.shape[-1] == 0: raise ValueError( - "qpos must have shape (control_dof,), (n_envs, control_dof), " + "Tensor target must have shape (control_dof,), " + "(n_envs, control_dof), " "or (n_envs, n_waypoint, control_dof), " - f"got {tuple(self.qpos.shape)}." + f"got {tuple(self.target.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" - """Name of the action, used for identification and logging.""" - - sample_interval: int = 50 - """Number of waypoints in the interpolated joint-space trajectory.""" - - named_joint_positions: dict[str, torch.Tensor] | None = None - """Optional named joint targets resolved by ``NamedJointPositionTarget``.""" +class MoveJointsOptions(ActionOptions): + """Per-invocation behavior for :class:`MoveJoints`.""" -class MoveJoints(AtomicAction[JointPositionTarget | NamedJointPositionTarget]): - """Plan a joint-space move for the configured control part. +class MoveJoints(AtomicAction[JointPositionGoal, MoveJointsOptions]): + """Plan joint motion from the observed state to one or more waypoints.""" - The :class:`JointPositionTarget` may carry either a single waypoint - ``(n_envs, control_dof)`` or a multi-waypoint trajectory - ``(n_envs, n_waypoint, control_dof)``. In the multi-waypoint case the - action plans a single trajectory that visits every waypoint in order, - starting from the inherited ``WorldState.last_qpos``. - """ - - TargetType: ClassVar[tuple[type, ...]] = ( - JointPositionTarget, - NamedJointPositionTarget, - ) + skill_id: ClassVar[str] = "move_joints" + GoalType: ClassVar[type] = JointPositionGoal + OptionsType: ClassVar[type] = MoveJointsOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + agent_visible: ClassVar[bool] = False def __init__( self, - motion_generator, - cfg: MoveJointsCfg | None = None, + default_options: MoveJointsOptions | None = None, ) -> None: - super().__init__(motion_generator, cfg or MoveJointsCfg()) - self.builder = TrajectoryBuilder(motion_generator) - self.n_envs = self.robot.get_qpos().shape[0] - self.joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.joint_dof = len(self.joint_ids) - self.robot_dof = self.robot.dof - self.named_joint_positions = self.cfg.named_joint_positions or {} - - def execute( + super().__init__(default_options) + + def plan( self, - target: JointPositionTarget | NamedJointPositionTarget, - state: WorldState, - ) -> ActionResult: + request: ResolvedActionRequest[JointPositionGoal, MoveJointsOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan a joint-space goal without mutating the robot or task state.""" + goal = self.require_goal(request) + manipulator = request.binding.manipulator("primary") + control_part = manipulator.name + joint_ids = list(manipulator.joint_ids) + joint_dof = manipulator.dof target_qpos = self.builder.resolve_joint_target( - self._resolve_target_qpos(target), - n_envs=self.n_envs, - joint_dof=self.joint_dof, - control_part=self.cfg.control_part, + self._resolve_target_qpos( + goal, + request=request, + context=context, + ), + n_envs=context.batch_size, + joint_dof=joint_dof, + control_part=control_part, ) start_qpos = self.builder.resolve_start_qpos( - state.last_qpos[:, self.joint_ids], - n_envs=self.n_envs, - arm_dof=self.joint_dof, - control_part=self.cfg.control_part, + context.robot.qpos[:, joint_ids], + n_envs=context.batch_size, + arm_dof=joint_dof, + control_part=control_part, ) - success, joint_traj = self.builder.plan_joint_motion( + result = self.builder.generate_joint_plan( start_qpos, target_qpos, - self.cfg.sample_interval, - control_part=self.cfg.control_part, - arm_dof=self.joint_dof, - cfg=self.cfg, + request.motion_policy.sample_count, + control_part=control_part, + arm_dof=joint_dof, + cfg=request.motion_policy, ) - full = self._embed(joint_traj, state.last_qpos) - return ActionResult( + success, trajectory = self.builder.to_full_robot_trajectory( + result, + base_qpos=context.robot.qpos, + joint_ids=joint_ids, + env_ids=context.env_ids, + control_dt=request.motion_policy.control_dt, + ) + return self.build_plan( + request, + context, success=success, - trajectory=full, - next_state=state.with_updates( - last_qpos=full[:, -1, :].clone(), - ), + trajectory=trajectory, + completion_kind=CompletionConditionKind.JOINT_GOAL_REACHED, ) def _resolve_target_qpos( - self, target: JointPositionTarget | NamedJointPositionTarget - ) -> torch.Tensor: - if isinstance(target, JointPositionTarget): - return target.qpos - if target.name not in self.named_joint_positions: - logger.log_error( - f"Unknown named joint-position target '{target.name}' for " - f"MoveJoints. Available targets: {sorted(self.named_joint_positions)}", - KeyError, - ) - return self.named_joint_positions[target.name] - - def _embed( - self, joint_traj: torch.Tensor, last_full_qpos: torch.Tensor + self, + goal: JointPositionGoal, + *, + request: ResolvedActionRequest[JointPositionGoal, MoveJointsOptions], + context: PlanningContext, ) -> torch.Tensor: - n_wp = joint_traj.shape[1] - full = torch.empty( - (self.n_envs, n_wp, self.robot_dof), - dtype=torch.float32, + """Resolve an explicit or named joint goal to a tensor.""" + if isinstance(goal.target, torch.Tensor): + return goal.target + return request.binding.manipulator("primary").joint_positions( + goal.target, + n_envs=context.batch_size, device=self.device, + dtype=context.robot.qpos.dtype, ) - full[:, :, :] = last_full_qpos.unsqueeze(1) - full[:, :, self.joint_ids] = joint_traj - return full __all__ = [ - "JointPositionTarget", + "JointPositionGoal", "MoveJoints", - "MoveJointsCfg", - "NamedJointPositionTarget", + "MoveJointsOptions", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index fbb0c35fe..6320f9a6f 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -25,7 +25,7 @@ import torch from embodichain.lab.sim.planners import MoveType, PlanState -from embodichain.utils import configclass, logger +from embodichain.utils import logger from embodichain.utils.math import ( axis_angle_to_rotation_matrix, pose_inv, @@ -35,23 +35,23 @@ from ._helpers import arm_qpos_from_state from ..affordance import AntipodalAffordance -from ..core import ( - ActionCfg, - ActionResult, - AtomicAction, - HeldObjectState, - ObjectSemantics, - WorldState, - _validate_pose_tensor, -) -from ..targets import ObjectActionTarget -from ..trajectory import TrajectoryBuilder +from ..bindings import ResolvedControlPart +from ..control import GRASP_COMMAND, OPEN_COMMAND +from ..core import AtomicAction, ObjectSemantics +from ..effects import StateDelta +from ..goals import ObjectActionGoal, validate_pose_tensor +from ..invocation import ActionOptions, ResolvedActionRequest +from ..plans import ActionPlan +from ..policies import MotionPolicy +from ..state import HeldObjectState, PlanningContext @dataclass(frozen=True, slots=True, eq=False) -class GraspTarget(ObjectActionTarget): +class GraspGoal(ObjectActionGoal): """Pickup target with an affordance-selected or supplied grasp pose.""" + goal_kind: ClassVar[str] = "grasp" + grasp_xpos: torch.Tensor | None = None """Optional end-effector grasp pose. @@ -61,31 +61,18 @@ class GraspTarget(ObjectActionTarget): """ def __post_init__(self) -> None: - ObjectActionTarget.__post_init__(self) + ObjectActionGoal.__post_init__(self) if self.grasp_xpos is not None: - _validate_pose_tensor(self.grasp_xpos, "grasp_xpos", allow_waypoints=False) + validate_pose_tensor(self.grasp_xpos, "grasp_xpos", allow_waypoints=False) -@configclass -class PickUpCfg(ActionCfg): - name: str = "pick_up" - """Name of the action, used for identification and logging.""" - - sample_interval: int = 80 - """Number of waypoints for the full trajectory (approach + hand + lift).""" +@dataclass(frozen=True, slots=True, eq=False) +class PickUpOptions(ActionOptions): + """Per-invocation pickup behavior.""" hand_interp_steps: int = 5 """Number of waypoints for the gripper close interpolation phase.""" - hand_control_part: str = "hand" - """Name of the robot part that controls the hand joints.""" - - hand_open_qpos: torch.Tensor | None = None - """Joint positions for the open hand state, shape ``[hand_dof,]``.""" - - hand_close_qpos: torch.Tensor | None = None - """Joint positions for the closed hand state, shape ``[hand_dof,]``.""" - pick_object_part: str = "center" """Name of the object part to pick up (used for grasp pose generation). Currently support [center | top | bottom].""" @@ -110,63 +97,82 @@ class PickUpCfg(ActionCfg): rotate_upright: float | None = None """Optional rotation (radians) about the grasp x-axis to apply after grasp selection.""" + def __post_init__(self) -> None: + if self.hand_interp_steps < 1: + raise ValueError("hand_interp_steps must be at least 1.") + if not isinstance(self.pick_object_part, str) or not self.pick_object_part: + raise ValueError("pick_object_part must be a non-empty string.") + if self.lift_height < 0.0: + raise ValueError("lift_height must be non-negative.") + if self.pre_grasp_distance < 0.0: + raise ValueError("pre_grasp_distance must be non-negative.") + if self.approach_direction.shape != (3,): + raise ValueError("approach_direction must have shape (3,).") + if not torch.isfinite(self.approach_direction).all(): + raise ValueError("approach_direction must contain finite values.") + if torch.linalg.vector_norm(self.approach_direction) <= 1.0e-6: + raise ValueError("approach_direction must be non-zero.") + if self.approach_alignment_max_angle is not None and not ( + 0.0 <= self.approach_alignment_max_angle <= math.pi / 2 + ): + raise ValueError("approach_alignment_max_angle must be in [0, pi / 2].") + if self.obj_upright_direction is not None and ( + self.obj_upright_direction.shape != (3,) + or not torch.isfinite(self.obj_upright_direction).all() + ): + raise ValueError("obj_upright_direction must be a finite (3,) tensor.") + object.__setattr__(self, "approach_direction", self.approach_direction.clone()) + object.__setattr__( + self, + "downstream_object_target_poses", + tuple(value.clone() for value in self.downstream_object_target_poses), + ) + if self.obj_upright_direction is not None: + object.__setattr__( + self, "obj_upright_direction", self.obj_upright_direction.clone() + ) -class PickUp(AtomicAction[GraspTarget]): + +class PickUp(AtomicAction[GraspGoal, PickUpOptions]): """Approach a grasp pose, close the gripper, lift.""" - TargetType: ClassVar[type] = GraspTarget + skill_id: ClassVar[str] = "pick_up" + GoalType: ClassVar[type] = GraspGoal + OptionsType: ClassVar[type] = PickUpOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) def __init__( self, - motion_generator, - cfg: PickUpCfg | None = None, + default_options: PickUpOptions | None = None, ) -> None: - super().__init__(motion_generator, cfg or PickUpCfg()) - self.builder = TrajectoryBuilder(motion_generator) + super().__init__(default_options) + + def _on_bind(self) -> None: + """Resolve engine-wide resources from the owning engine.""" self.n_envs = self.robot.get_qpos().shape[0] - self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part) - self.arm_dof = len(self.arm_joint_ids) self.robot_dof = self.robot.dof - if self.cfg.hand_open_qpos is None: - logger.log_error( - "hand_open_qpos must be specified in PickUpCfg", ValueError - ) - if self.cfg.hand_close_qpos is None: - logger.log_error( - "hand_close_qpos must be specified in PickUpCfg", ValueError - ) - self.hand_open_qpos = self.cfg.hand_open_qpos.to(self.device) - self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) - self.approach_direction = self.cfg.approach_direction.to( - device=self.device, dtype=torch.float32 - ) - approach_norm = torch.linalg.vector_norm(self.approach_direction) - if approach_norm <= 1.0e-6: - logger.log_error("approach_direction must be non-zero.", ValueError) - self.approach_direction = self.approach_direction / approach_norm - if self.cfg.approach_alignment_max_angle is not None and not ( - 0.0 <= self.cfg.approach_alignment_max_angle <= math.pi / 2 - ): - logger.log_error( - "approach_alignment_max_angle must be in [0, pi / 2].", - ValueError, - ) - def _get_full_pickup_trajectory( self, grasp_xpos: torch.Tensor, start_arm_qpos: torch.Tensor, last_qpos: torch.Tensor, - ): + motion_policy: MotionPolicy, + options: PickUpOptions, + approach_direction: torch.Tensor, + manipulator: ResolvedControlPart, + end_effector: ResolvedControlPart, + hand_open_qpos: torch.Tensor, + hand_grasp_qpos: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: pre_grasp_xpos = self.builder.apply_local_offset( - grasp_xpos, -self.approach_direction * self.cfg.pre_grasp_distance + grasp_xpos, -approach_direction * options.pre_grasp_distance ) n_approach, n_close, n_lift = self.builder.split_three_phase( - self.cfg.sample_interval, - self.cfg.hand_interp_steps, + motion_policy.sample_count, + options.hand_interp_steps, first_phase_name="approach", third_phase_name="lift", ) @@ -182,15 +188,15 @@ def _get_full_pickup_trajectory( target_states_list, start_arm_qpos, n_approach, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=self.cfg, + control_part=manipulator.name, + arm_dof=manipulator.dof, + cfg=motion_policy, ) grasp_arm_qpos = approach_arm[:, -1, :] lift_xpos = self.builder.apply_local_offset( grasp_xpos, - torch.tensor([0, 0, 1], device=self.device) * self.cfg.lift_height, + torch.tensor([0, 0, 1], device=self.device) * options.lift_height, ) target_states_list = [ [PlanState(xpos=lift_xpos[i], move_type=MoveType.EEF_MOVE)] @@ -200,14 +206,14 @@ def _get_full_pickup_trajectory( target_states_list, grasp_arm_qpos, n_lift, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=self.cfg, + control_part=manipulator.name, + arm_dof=manipulator.dof, + cfg=motion_policy, ) is_success = approach_success & lift_success hand_close_path = self.builder.interpolate_hand_qpos( - self.hand_open_qpos, self.hand_close_qpos, n_waypoints=n_close + hand_open_qpos, hand_grasp_qpos, n_waypoints=n_close ) n_approach_actual = approach_arm.shape[1] n_lift_actual = lift_arm.shape[1] @@ -217,21 +223,53 @@ def _get_full_pickup_trajectory( device=self.device, ) full[:, :, :] = last_qpos.unsqueeze(1) - full[:, :n_approach_actual, self.arm_joint_ids] = approach_arm - full[:, :n_approach_actual, self.hand_joint_ids] = self.hand_open_qpos - full[:, n_approach_actual : n_approach_actual + n_close, self.arm_joint_ids] = ( + arm_joint_ids = list(manipulator.joint_ids) + hand_joint_ids = list(end_effector.joint_ids) + full[:, :n_approach_actual, arm_joint_ids] = approach_arm + full[:, :n_approach_actual, hand_joint_ids] = hand_open_qpos.unsqueeze(1) + full[:, n_approach_actual : n_approach_actual + n_close, arm_joint_ids] = ( grasp_arm_qpos.unsqueeze(1) ) - full[ - :, n_approach_actual : n_approach_actual + n_close, self.hand_joint_ids - ] = hand_close_path - full[:, n_approach_actual + n_close :, self.arm_joint_ids] = lift_arm - full[:, n_approach_actual + n_close :, self.hand_joint_ids] = ( - self.hand_close_qpos + full[:, n_approach_actual : n_approach_actual + n_close, hand_joint_ids] = ( + hand_close_path + ) + full[:, n_approach_actual + n_close :, arm_joint_ids] = lift_arm + full[:, n_approach_actual + n_close :, hand_joint_ids] = ( + hand_grasp_qpos.unsqueeze(1) ) return is_success, full - def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: + def plan( + self, + request: ResolvedActionRequest[GraspGoal, PickUpOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan approach, close, and lift phases without committing attachment.""" + target = self.require_goal(request) + options = request.skill_options + approach_direction = options.approach_direction.to( + device=self.device, dtype=torch.float32 + ) + approach_direction = approach_direction / torch.linalg.vector_norm( + approach_direction + ) + binding = request.binding + manipulator = binding.manipulator() + end_effector = binding.end_effector() + hand_open_qpos = end_effector.joint_positions( + OPEN_COMMAND, + n_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + hand_grasp_qpos = end_effector.joint_positions( + GRASP_COMMAND, + n_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + control_part = manipulator.name + state = context sem = target.semantics if target.grasp_xpos is None and not isinstance( sem.affordance, AntipodalAffordance @@ -245,26 +283,41 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: "PickUp requires an entity on the target semantics.", ValueError ) start_arm_qpos = self.builder.resolve_start_qpos( - arm_qpos_from_state(state, self.arm_joint_ids), + arm_qpos_from_state(state, list(manipulator.joint_ids)), n_envs=self.n_envs, - arm_dof=self.arm_dof, - control_part=self.cfg.control_part, + arm_dof=manipulator.dof, + control_part=control_part, ) if target.grasp_xpos is None: - is_success, grasp_xpos = self._resolve_grasp_pose(sem, start_arm_qpos) + is_success, grasp_xpos = self._resolve_grasp_pose( + sem, start_arm_qpos, manipulator, options, approach_direction + ) else: grasp_xpos = self.builder.resolve_pose_target( target.grasp_xpos, n_envs=self.n_envs ) - if self.cfg.rotate_upright is not None: - self._apply_upright_rotation(sem, grasp_xpos) + if options.rotate_upright is not None: + grasp_xpos = self._upright_adjusted_grasp_poses( + sem, grasp_xpos, options + ) is_success = torch.ones(self.n_envs, dtype=torch.bool, device=self.device) if not self.builder.all_envs_success(is_success): logger.log_warning("PickUp failed to resolve a grasp pose.") - return self._fail(state) + return self.failed_plan( + request, context, message="Failed to resolve a grasp pose." + ) is_success, full = self._get_full_pickup_trajectory( - grasp_xpos, start_arm_qpos, state.last_qpos + grasp_xpos, + start_arm_qpos, + state.last_qpos, + request.motion_policy, + options, + approach_direction, + manipulator, + end_effector, + hand_open_qpos, + hand_grasp_qpos, ) obj_poses = sem.entity.get_local_pose(to_matrix=True) @@ -272,42 +325,34 @@ 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 + coordinated_updates = { + key: None for key in state.coordinated_held_objects if control_part in key } - return ActionResult( + return self.build_plan( + request, + context, success=is_success, trajectory=full, - next_state=state.with_updates( - last_qpos=full[:, -1, :].clone(), - held_objects=held_objects, - coordinated_held_objects=coordinated_held_objects, + expected_effects=StateDelta( + held_object_updates={control_part: held}, + coordinated_held_object_updates=coordinated_updates, ), - ) - - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=torch.zeros(self.n_envs, dtype=torch.bool, device=self.device), - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, + phase_name="pick", ) def _resolve_grasp_pose( - self, semantics: ObjectSemantics, start_qpos: torch.Tensor + self, + semantics: ObjectSemantics, + start_qpos: torch.Tensor, + manipulator: ResolvedControlPart, + options: PickUpOptions, + approach_direction: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: obj_poses = semantics.entity.get_local_pose(to_matrix=True) grasp_poses_result = semantics.affordance.get_valid_grasp_poses( obj_poses=obj_poses, - approach_direction=self.approach_direction, - object_part=self.cfg.pick_object_part, + approach_direction=approach_direction, + object_part=options.pick_object_part, ) n_envs = obj_poses.shape[0] n_max_pose = max(r[0].shape[0] for r in grasp_poses_result) @@ -333,7 +378,13 @@ def _resolve_grasp_pose( grasp_xpos_padding[i, n_pose:] = grasp_poses[0] grasp_cost_padding[i, n_pose:] = grasp_costs[0] grasp_xpos_padding, ik_success = self._select_feasible_grasp_variants( - semantics, grasp_xpos_padding, start_qpos, obj_poses + semantics, + grasp_xpos_padding, + start_qpos, + obj_poses, + manipulator, + options, + approach_direction, ) grasp_cost_masked = torch.where(ik_success, grasp_cost_padding, 10000.0) best_cost, best_idx = grasp_cost_masked.min(dim=1) @@ -349,6 +400,9 @@ def _select_feasible_grasp_variants( grasp_xpos: torch.Tensor, start_qpos: torch.Tensor, object_poses: torch.Tensor, + manipulator: ResolvedControlPart, + options: PickUpOptions, + approach_direction: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Choose a TCP-roll variant with a feasible pickup and transport path.""" n_envs, n_pose = grasp_xpos.shape[:2] @@ -357,29 +411,31 @@ def _select_feasible_grasp_variants( mirrored_grasp_xpos[..., :3, 1] = -mirrored_grasp_xpos[..., :3, 1] selection_variants = torch.stack([grasp_xpos, mirrored_grasp_xpos], dim=2) grasp_variants = self._upright_adjusted_grasp_poses( - semantics, selection_variants + semantics, selection_variants, options ) pre_grasp_variants = grasp_variants.clone() pre_grasp_z = pre_grasp_variants[..., :3, 2] - pre_grasp_variants[..., :3, 3] -= pre_grasp_z * self.cfg.pre_grasp_distance + pre_grasp_variants[..., :3, 3] -= pre_grasp_z * options.pre_grasp_distance lift_variants = grasp_variants.clone() lift_variants[..., :3, 3] += torch.tensor( - [0.0, 0.0, self.cfg.lift_height], + [0.0, 0.0, options.lift_height], dtype=grasp_variants.dtype, device=self.device, ) pre_grasp_success, pre_grasp_qpos = self._compute_batch_candidate_ik( - pre_grasp_variants, start_qpos + pre_grasp_variants, start_qpos, manipulator ) grasp_success, grasp_qpos = self._compute_batch_candidate_ik( - grasp_variants, pre_grasp_qpos + grasp_variants, pre_grasp_qpos, manipulator ) lift_success, lift_qpos = self._compute_batch_candidate_ik( - lift_variants, grasp_qpos + lift_variants, grasp_qpos, manipulator + ) + alignment_success = self._approach_alignment_mask( + grasp_variants, options, approach_direction ) - alignment_success = self._approach_alignment_mask(grasp_variants) pickup_success = ( alignment_success & pre_grasp_success & grasp_success & lift_success ) @@ -390,7 +446,7 @@ def _select_feasible_grasp_variants( # MoveHeldObject begins after the lift, so screen its target from the # same joint state that the execution stream will use. downstream_seed = lift_qpos - for object_target_pose in self.cfg.downstream_object_target_poses: + for object_target_pose in options.downstream_object_target_poses: object_target_pose = object_target_pose.to( device=self.device, dtype=torch.float32 ) @@ -409,7 +465,7 @@ def _select_feasible_grasp_variants( object_target_pose[:, None, None], object_to_eef_variants ) downstream_success, downstream_seed = self._compute_batch_candidate_ik( - downstream_eef_variants, downstream_seed + downstream_eef_variants, downstream_seed, manipulator ) pickup_success &= downstream_success downstream_success_counts.append(pickup_success.sum(dim=(1, 2)).tolist()) @@ -425,7 +481,7 @@ def _select_feasible_grasp_variants( start_xpos = self.robot.compute_fk( qpos=start_qpos, - name=self.cfg.control_part, + name=manipulator.name, to_matrix=True, ) start_quat = quat_from_matrix(start_xpos[:, :3, :3]) @@ -450,55 +506,60 @@ def _select_feasible_grasp_variants( ik_success = pickup_success[env_idx, pose_idx, best_variant_idx] return selected_grasp_xpos, ik_success - def _approach_alignment_mask(self, grasp_poses: torch.Tensor) -> torch.Tensor: + def _approach_alignment_mask( + self, + grasp_poses: torch.Tensor, + options: PickUpOptions, + approach_direction: torch.Tensor, + ) -> torch.Tensor: """Return candidates whose final TCP z-axis follows the approach direction.""" - max_angle = self.cfg.approach_alignment_max_angle - if self.cfg.rotate_upright is not None or max_angle is None: + max_angle = options.approach_alignment_max_angle + if options.rotate_upright is not None or max_angle is None: return torch.ones( grasp_poses.shape[:3], dtype=torch.bool, device=grasp_poses.device ) grasp_z = torch.nn.functional.normalize(grasp_poses[..., :3, 2], dim=-1) - alignment = torch.sum(grasp_z * self.approach_direction, dim=-1) + alignment = torch.sum(grasp_z * approach_direction, dim=-1) return alignment >= math.cos(float(max_angle)) def _compute_batch_candidate_ik( - self, poses: torch.Tensor, joint_seed: torch.Tensor + self, + poses: torch.Tensor, + joint_seed: torch.Tensor, + manipulator: ResolvedControlPart, ) -> tuple[torch.Tensor, torch.Tensor]: """Solve candidate IK poses while preserving the candidate dimensions.""" n_envs, n_pose, n_variant = poses.shape[:3] flat_poses = poses.reshape(n_envs, n_pose * n_variant, 4, 4) if joint_seed.dim() == 2: joint_seed = joint_seed[:, None, None, :].expand(-1, n_pose, n_variant, -1) - flat_seed = joint_seed.reshape(n_envs, n_pose * n_variant, self.arm_dof) + flat_seed = joint_seed.reshape(n_envs, n_pose * n_variant, manipulator.dof) is_success, qpos = self.robot.compute_batch_ik( pose=flat_poses, - name=self.cfg.control_part, + name=manipulator.name, joint_seed=flat_seed, ) return ( is_success.reshape(n_envs, n_pose, n_variant), - qpos.reshape(n_envs, n_pose, n_variant, self.arm_dof), + qpos.reshape(n_envs, n_pose, n_variant, manipulator.dof), ) - def _apply_upright_rotation( - self, semantics: ObjectSemantics, grasp_xpos: torch.Tensor - ) -> None: - """Apply the configured upright-in-place grasp roll adjustment.""" - grasp_xpos.copy_(self._upright_adjusted_grasp_poses(semantics, grasp_xpos)) - def _upright_adjusted_grasp_poses( - self, semantics: ObjectSemantics, grasp_xpos: torch.Tensor + self, + semantics: ObjectSemantics, + grasp_xpos: torch.Tensor, + options: PickUpOptions, ) -> torch.Tensor: """Return grasp poses after the optional upright-in-place roll adjustment.""" - if self.cfg.rotate_upright is None: + if options.rotate_upright is None: return grasp_xpos - if self.cfg.obj_upright_direction is None: + if options.obj_upright_direction is None: upright_direction = torch.tensor( [0, 0, 1], dtype=torch.float32, device=self.device ) else: - upright_direction = self.cfg.obj_upright_direction.to( + upright_direction = options.obj_upright_direction.to( device=self.device, dtype=torch.float32 ) obj_pose = semantics.entity.get_local_pose(to_matrix=True) @@ -511,7 +572,7 @@ def _upright_adjusted_grasp_poses( dot_result = (grasp_ry * object_axes).sum(dim=-1) revert_flag = torch.where(dot_result < 0, -1.0, 1.0) grasp_rx = adjusted_grasp_xpos[..., :3, 0] - rota_axis_angle = self.cfg.rotate_upright * revert_flag[..., None] * grasp_rx + rota_axis_angle = options.rotate_upright * revert_flag[..., None] * grasp_rx rota_offset = axis_angle_to_rotation_matrix( rota_axis_angle.reshape(-1, 3) ).reshape(*rota_axis_angle.shape[:-1], 3, 3) @@ -521,4 +582,4 @@ def _upright_adjusted_grasp_poses( return adjusted_grasp_xpos -__all__ = ["GraspTarget", "PickUp", "PickUpCfg"] +__all__ = ["GraspGoal", "PickUp", "PickUpOptions"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/place.py b/embodichain/lab/sim/atomic_actions/primitives/place.py index 054c83c16..7bfa24253 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/place.py +++ b/embodichain/lab/sim/atomic_actions/primitives/place.py @@ -24,29 +24,29 @@ import torch from embodichain.lab.sim.planners import MoveType, PlanState -from embodichain.utils import configclass, logger +from embodichain.utils import logger from embodichain.utils.math import quat_error_magnitude, quat_from_matrix from ._helpers import arm_qpos_from_state, resolve_object_target from ..affordance import AssembleAffordance -from ..core import ( - ActionTarget, - ActionCfg, - ActionResult, - AtomicAction, - WorldState, - _validate_pose_tensor, -) -from ..trajectory import TrajectoryBuilder +from ..control import GRASP_COMMAND, OPEN_COMMAND +from ..core import AtomicAction +from ..effects import StateDelta +from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal +from ..invocation import ActionOptions, ResolvedActionRequest +from ..plans import ActionPlan +from ..state import PlanningContext TcpSymmetry = Literal["none", "z_roll_180"] @dataclass(frozen=True, slots=True, eq=False) -class PlaceTarget(ActionTarget): +class PlaceGoal: """End-effector release-pose target used by :class:`Place`.""" - xpos: torch.Tensor + goal_kind: ClassVar[str] = "place_pose" + + xpos: PoseGoalValue """Target end-effector release pose. Accepts ``(4, 4)``, ``(n_envs, 4, 4)``, or @@ -62,7 +62,7 @@ class PlaceTarget(ActionTarget): """ def __post_init__(self) -> None: - _validate_pose_tensor(self.xpos, "xpos", allow_waypoints=True) + validate_pose_goal(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', " @@ -71,40 +71,29 @@ def __post_init__(self) -> None: @dataclass(frozen=True, slots=True, eq=False) -class AssembleTarget(ActionTarget): +class AssembleGoal: """Place a held assemble object onto a base object at a relative pose. The base object pose is read at planning time from :attr:`AssembleAffordance.base_object_entity`, and the assemble object's target pose is ``base_pose @ assemble_to_base_pose``. The held-object - transform (``object_to_eef``) is read from :attr:`WorldState.held_objects` + transform (``object_to_eef``) is read from :class:`PlanningContext` for the place control part, which a prior :class:`PickUp` populates. """ + goal_kind: ClassVar[str] = "assemble" + affordance: AssembleAffordance """Assembly affordance anchoring the assemble object to the base object.""" -@configclass -class PlaceCfg(ActionCfg): - name: str = "place" - """Name of the action, used for identification and logging.""" - - sample_interval: int = 80 - """Number of waypoints for the full trajectory (down + hand + back).""" +@dataclass(frozen=True, slots=True, eq=False) +class PlaceOptions(ActionOptions): + """Per-invocation placement behavior.""" hand_interp_steps: int = 5 """Number of waypoints for the gripper open interpolation phase.""" - hand_control_part: str = "hand" - """Name of the robot part that controls the hand joints.""" - - hand_open_qpos: torch.Tensor | None = None - """Joint positions for the open hand state, shape ``[hand_dof,]``.""" - - hand_close_qpos: torch.Tensor | None = None - """Joint positions for the closed hand state, shape ``[hand_dof,]``.""" - lift_height: float = 0.1 """Height (m) to retract the end-effector after opening the gripper.""" @@ -114,87 +103,110 @@ class PlaceCfg(ActionCfg): cartesian_waypoint_count: int = 1 """Number of fixed-orientation Cartesian keyframes per translation segment.""" + def __post_init__(self) -> None: + if self.hand_interp_steps < 1: + raise ValueError("hand_interp_steps must be at least 1.") + if self.lift_height < 0.0: + raise ValueError("lift_height must be non-negative.") + if self.cartesian_waypoint_count < 1: + raise ValueError("cartesian_waypoint_count must be at least 1.") + -class Place(AtomicAction[PlaceTarget | AssembleTarget]): +class Place(AtomicAction[PlaceGoal | AssembleGoal, PlaceOptions]): """Lower the held object to a place pose, open the gripper, retract. - The :class:`PlaceTarget` may carry either a single waypoint + The :class:`PlaceGoal` may carry either a single waypoint ``(n_envs, 4, 4)`` (or a broadcastable ``(4, 4)``) or a multi-waypoint trajectory ``(n_envs, n_waypoint, 4, 4)``. In the multi-waypoint case the down phase visits every waypoint in order; approaching from above the first waypoint, descending through each waypoint, then opening the gripper at the final waypoint and retracting to above the last waypoint. Starting - joint positions are inherited from ``WorldState.last_qpos``. + joint positions are inherited from :class:`PlanningContext`. - An :class:`AssembleTarget` replaces the explicit EEF pose with an assembly + An :class:`AssembleGoal` replaces the explicit EEF pose with an assembly affordance: the place pose is derived from the base object's current pose and ``assemble_to_base_pose``, converted to an EEF pose through the held - object's ``object_to_eef`` (read from ``WorldState.held_objects``). + object's ``object_to_eef`` (read from :class:`PlanningContext`). """ - TargetType: ClassVar[type | tuple[type, ...]] = ( - PlaceTarget, - AssembleTarget, + skill_id: ClassVar[str] = "place" + GoalType: ClassVar[type | tuple[type, ...]] = ( + PlaceGoal, + AssembleGoal, ) + OptionsType: ClassVar[type] = PlaceOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) def __init__( self, - motion_generator, - cfg: PlaceCfg | None = None, + default_options: PlaceOptions | None = None, ) -> None: - super().__init__(motion_generator, cfg or PlaceCfg()) - self.builder = TrajectoryBuilder(motion_generator) + super().__init__(default_options) + + def _on_bind(self) -> None: + """Resolve engine-wide resources from the owning engine.""" self.n_envs = self.robot.get_qpos().shape[0] - self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part) - self.arm_dof = len(self.arm_joint_ids) self.robot_dof = self.robot.dof - if self.cfg.hand_open_qpos is None: - logger.log_error("hand_open_qpos must be specified in PlaceCfg", ValueError) - if self.cfg.hand_close_qpos is None: - logger.log_error( - "hand_close_qpos must be specified in PlaceCfg", ValueError - ) - self.hand_open_qpos = self.cfg.hand_open_qpos.to(self.device) - self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) - if self.cfg.cartesian_waypoint_count < 1: - logger.log_error("cartesian_waypoint_count must be at least 1.", ValueError) - - def execute( - self, target: PlaceTarget | AssembleTarget, state: WorldState - ) -> ActionResult: - place_xpos = self._resolve_place_xpos(target, state) + def plan( + self, + request: ResolvedActionRequest[PlaceGoal | AssembleGoal, PlaceOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan approach, release, and retract without committing detachment.""" + target = self.require_goal(request) + options = request.skill_options + binding = request.binding + manipulator = binding.manipulator() + end_effector = binding.end_effector() + control_part = manipulator.name + arm_joint_ids = list(manipulator.joint_ids) + hand_joint_ids = list(end_effector.joint_ids) + hand_open_qpos = end_effector.joint_positions( + OPEN_COMMAND, + n_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + hand_grasp_qpos = end_effector.joint_positions( + GRASP_COMMAND, + n_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + state = context + place_xpos = self._resolve_place_xpos(target, state, control_part) if place_xpos.dim() == 3: place_xpos = place_xpos.unsqueeze(1) start_arm_qpos = self.builder.resolve_start_qpos( - arm_qpos_from_state(state, self.arm_joint_ids), + arm_qpos_from_state(state, arm_joint_ids), n_envs=self.n_envs, - arm_dof=self.arm_dof, - control_part=self.cfg.control_part, + arm_dof=manipulator.dof, + control_part=control_part, ) - if isinstance(target, PlaceTarget) and target.tcp_symmetry == "z_roll_180": + if isinstance(target, PlaceGoal) and target.tcp_symmetry == "z_roll_180": place_xpos = self._select_tcp_symmetric_place_variant( - place_xpos, start_arm_qpos + place_xpos, start_arm_qpos, control_part ) n_down, n_open, n_back = self.builder.split_three_phase( - self.cfg.sample_interval, - self.cfg.hand_interp_steps, + request.motion_policy.sample_count, + options.hand_interp_steps, first_phase_name="approach", third_phase_name="back", ) - approach_xpos = self._lifted_pose(place_xpos[:, 0]) - retract_xpos = self._lifted_pose(place_xpos[:, -1]) + approach_xpos = self._lifted_pose(place_xpos[:, 0], options) + retract_xpos = self._lifted_pose(place_xpos[:, -1], options) start_xpos = self.robot.compute_fk( qpos=start_arm_qpos, - name=self.cfg.control_part, + name=control_part, to_matrix=True, ) down_xpos = torch.cat([approach_xpos.unsqueeze(1), place_xpos], dim=1) - down_xpos = self._translation_keyframes(start_xpos, down_xpos) + down_xpos = self._translation_keyframes(start_xpos, down_xpos, options) target_states_list = [ [ @@ -207,14 +219,14 @@ def execute( target_states_list, start_arm_qpos, n_down, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=self.cfg, + control_part=control_part, + arm_dof=manipulator.dof, + cfg=request.motion_policy, ) reach_arm_qpos = down_arm[:, -1, :] back_xpos = self._translation_keyframes( - place_xpos[:, -1], retract_xpos.unsqueeze(1) + place_xpos[:, -1], retract_xpos.unsqueeze(1), options ) target_states_list = [ [ @@ -227,14 +239,14 @@ def execute( target_states_list, reach_arm_qpos, n_back, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=self.cfg, + control_part=control_part, + arm_dof=manipulator.dof, + cfg=request.motion_policy, ) success = down_success & back_success hand_open_path = self.builder.interpolate_hand_qpos( - self.hand_close_qpos, self.hand_open_qpos, n_waypoints=n_open + hand_grasp_qpos, hand_open_qpos, n_waypoints=n_open ) # Allocate from the actually-returned phase lengths so collision-aware @@ -247,36 +259,35 @@ def execute( device=self.device, ) full[:, :, :] = state.last_qpos.unsqueeze(1) - full[:, :n_down_actual, self.arm_joint_ids] = down_arm - full[:, :n_down_actual, self.hand_joint_ids] = self.hand_close_qpos - full[:, n_down_actual : n_down_actual + n_open, self.arm_joint_ids] = ( + full[:, :n_down_actual, arm_joint_ids] = down_arm + full[:, :n_down_actual, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1) + full[:, n_down_actual : n_down_actual + n_open, arm_joint_ids] = ( reach_arm_qpos.unsqueeze(1) ) - full[:, n_down_actual : n_down_actual + n_open, self.hand_joint_ids] = ( - hand_open_path - ) - 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 + full[:, n_down_actual : n_down_actual + n_open, hand_joint_ids] = hand_open_path + full[:, n_down_actual + n_open :, arm_joint_ids] = back_arm + full[:, n_down_actual + n_open :, hand_joint_ids] = hand_open_qpos.unsqueeze(1) + + coordinated_updates = { + key: None for key in state.coordinated_held_objects if control_part in key } - return ActionResult( + return self.build_plan( + request, + context, success=success, trajectory=full, - next_state=state.with_updates( - last_qpos=full[:, -1, :].clone(), - held_objects=held_objects, - coordinated_held_objects=coordinated_held_objects, + expected_effects=StateDelta( + held_object_updates={control_part: None}, + coordinated_held_object_updates=coordinated_updates, ), + phase_name="place", ) def _resolve_place_xpos( - self, target: PlaceTarget | AssembleTarget, state: WorldState + self, + target: PlaceGoal | AssembleGoal, + state: PlanningContext, + control_part: str, ) -> torch.Tensor: """Resolve the place EEF poses from a typed target. @@ -288,12 +299,18 @@ def _resolve_place_xpos( Place EEF poses with shape ``(n_envs, 4, 4)`` or ``(n_envs, n_waypoint, 4, 4)``. """ - if isinstance(target, PlaceTarget): - return self.builder.resolve_pose_target(target.xpos, n_envs=self.n_envs) - return self._resolve_assemble_place_xpos(target, state) + if isinstance(target, PlaceGoal): + return self.builder.resolve_pose_target( + resolve_pose_goal(target.xpos, state, name="xpos"), + n_envs=self.n_envs, + ) + return self._resolve_assemble_place_xpos(target, state, control_part) def _resolve_assemble_place_xpos( - self, target: AssembleTarget, state: WorldState + self, + target: AssembleGoal, + state: PlanningContext, + control_part: str, ) -> torch.Tensor: """Derive the place EEF pose from an assembly affordance. @@ -311,11 +328,11 @@ def _resolve_assemble_place_xpos( Raises: ValueError: If no held object or no base object entity is available. """ - held = state.get_held_object(self.cfg.control_part) + held = state.get_held_object(control_part) if held is None: logger.log_error( - "Place with AssembleTarget requires an object held by control " - f"part {self.cfg.control_part!r} (run PickUp first).", + "Place with AssembleGoal requires an object held by control " + f"part {control_part!r} (run PickUp first).", ValueError, ) affordance = target.affordance @@ -337,13 +354,15 @@ def _resolve_assemble_place_xpos( ) return torch.bmm(assemble_object_pose, object_to_eef) - def _lifted_pose(self, release_xpos: torch.Tensor) -> torch.Tensor: + def _lifted_pose( + self, release_xpos: torch.Tensor, options: PlaceOptions + ) -> torch.Tensor: """Build an above-release pose while respecting the optional TCP z cap.""" lifted_xpos = release_xpos.clone() - lifted_z = release_xpos[:, 2, 3] + self.cfg.lift_height - if self.cfg.max_approach_retract_z is not None: + lifted_z = release_xpos[:, 2, 3] + options.lift_height + if options.max_approach_retract_z is not None: max_z = torch.as_tensor( - self.cfg.max_approach_retract_z, + options.max_approach_retract_z, dtype=release_xpos.dtype, device=release_xpos.device, ) @@ -355,10 +374,13 @@ def _lifted_pose(self, release_xpos: torch.Tensor) -> torch.Tensor: return lifted_xpos def _translation_keyframes( - self, start_xpos: torch.Tensor, target_xpos: torch.Tensor + self, + start_xpos: torch.Tensor, + target_xpos: torch.Tensor, + options: PlaceOptions, ) -> torch.Tensor: """Interpolate translations while holding each segment's target rotation.""" - count = self.cfg.cartesian_waypoint_count + count = options.cartesian_waypoint_count if count == 1: return target_xpos @@ -380,19 +402,11 @@ def _translation_keyframes( ) return keyframes.flatten(1, 2) - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=torch.zeros(self.n_envs, dtype=torch.bool, device=self.device), - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, - ) - def _select_tcp_symmetric_place_variant( - self, place_xpos: torch.Tensor, start_qpos: torch.Tensor + self, + place_xpos: torch.Tensor, + start_qpos: torch.Tensor, + control_part: str, ) -> torch.Tensor: """Choose the closest TCP z-roll variant for an opt-in place target.""" mirrored_place_xpos = place_xpos.clone() @@ -402,7 +416,7 @@ def _select_tcp_symmetric_place_variant( start_xpos = self.robot.compute_fk( qpos=start_qpos, - name=self.cfg.control_part, + name=control_part, to_matrix=True, ) start_quat = quat_from_matrix(start_xpos[:, :3, :3]) @@ -423,4 +437,4 @@ def _select_tcp_symmetric_place_variant( ] -__all__ = ["Place", "PlaceCfg", "PlaceTarget", "AssembleTarget"] +__all__ = ["AssembleGoal", "Place", "PlaceGoal", "PlaceOptions"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/press.py b/embodichain/lab/sim/atomic_actions/primitives/press.py index f1606d543..0bb193fb2 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/press.py +++ b/embodichain/lab/sim/atomic_actions/primitives/press.py @@ -24,93 +24,102 @@ import torch from embodichain.lab.sim.planners import MoveType, PlanState -from embodichain.utils import configclass, logger +from embodichain.utils import logger from ._helpers import arm_qpos_from_state -from ..core import ( - ActionTarget, - ActionCfg, - ActionResult, - AtomicAction, - WorldState, - _validate_pose_tensor, -) -from ..trajectory import TrajectoryBuilder +from ..control import GRASP_COMMAND +from ..core import AtomicAction +from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal +from ..invocation import ActionOptions, ResolvedActionRequest +from ..plans import ActionPlan +from ..state import PlanningContext @dataclass(frozen=True, slots=True, eq=False) -class PressTarget(ActionTarget): +class PressGoal: """Single end-effector contact pose used by :class:`Press`.""" - xpos: torch.Tensor + goal_kind: ClassVar[str] = "press_pose" + + xpos: PoseGoalValue """Contact pose, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" def __post_init__(self) -> None: - _validate_pose_tensor(self.xpos, "xpos", allow_waypoints=False) - + validate_pose_goal(self.xpos, "xpos", allow_waypoints=False) -@configclass -class PressCfg(ActionCfg): - name: str = "press" - """Name of the action, used for identification and logging.""" - sample_interval: int = 80 - """Number of waypoints for the full trajectory (hand close + down + back).""" +@dataclass(frozen=True, slots=True, eq=False) +class PressOptions(ActionOptions): + """Per-invocation press behavior.""" hand_interp_steps: int = 5 """Number of waypoints for closing the gripper before pressing.""" - hand_control_part: str = "hand" - """Name of the robot part that controls the hand joints.""" - - hand_close_qpos: torch.Tensor | None = None - """Joint positions for the closed hand state, shape ``[hand_dof,]``.""" + def __post_init__(self) -> None: + if self.hand_interp_steps < 1: + raise ValueError("hand_interp_steps must be at least 1.") -class Press(AtomicAction[PressTarget]): +class Press(AtomicAction[PressGoal, PressOptions]): """Close the gripper, press down to a target pose, then return.""" - TargetType: ClassVar[type] = PressTarget + skill_id: ClassVar[str] = "press" + GoalType: ClassVar[type] = PressGoal + OptionsType: ClassVar[type] = PressOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) def __init__( self, - motion_generator, - cfg: PressCfg | None = None, + default_options: PressOptions | None = None, ) -> None: - super().__init__(motion_generator, cfg or PressCfg()) - self.builder = TrajectoryBuilder(motion_generator) + super().__init__(default_options) + + def _on_bind(self) -> None: + """Resolve engine-wide resources from the owning engine.""" self.n_envs = self.robot.get_qpos().shape[0] - self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part) - self.arm_dof = len(self.arm_joint_ids) - self.hand_dof = len(self.hand_joint_ids) self.robot_dof = self.robot.dof - if self.cfg.hand_close_qpos is None: - logger.log_error( - "hand_close_qpos must be specified in PressCfg", ValueError - ) - self.hand_close_qpos = self.builder.expand_hand_qpos( - self.cfg.hand_close_qpos, + def plan( + self, + request: ResolvedActionRequest[PressGoal, PressOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan a close, press, and retract sequence.""" + target = self.require_goal(request) + options = request.skill_options + binding = request.binding + manipulator = binding.manipulator() + end_effector = binding.end_effector() + control_part = manipulator.name + arm_joint_ids = list(manipulator.joint_ids) + hand_joint_ids = list(end_effector.joint_ids) + hand_close_qpos = end_effector.joint_positions( + GRASP_COMMAND, + n_envs=self.n_envs, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + state = context + press_xpos = self.builder.resolve_pose_target( + resolve_pose_goal(target.xpos, context, name="xpos"), n_envs=self.n_envs, - hand_dof=self.hand_dof, ) - - 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), + arm_qpos_from_state(state, arm_joint_ids), n_envs=self.n_envs, - arm_dof=self.arm_dof, - control_part=self.cfg.control_part, + arm_dof=manipulator.dof, + control_part=control_part, ) - start_hand_qpos = state.last_qpos[:, self.hand_joint_ids] + start_hand_qpos = state.last_qpos[:, hand_joint_ids] - n_close, n_down, n_back = self._compute_phase_waypoints() + n_close, n_down, n_back = self._compute_phase_waypoints( + request.motion_policy.sample_count, options + ) hand_close_path = self.builder.interpolate_hand_qpos( start_hand_qpos, - self.hand_close_qpos, + hand_close_qpos, n_waypoints=n_close, ) @@ -122,9 +131,9 @@ def execute(self, target: PressTarget, state: WorldState) -> ActionResult: target_states_list, start_arm_qpos, n_down, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=self.cfg, + control_part=control_part, + arm_dof=manipulator.dof, + cfg=request.motion_policy, ) press_arm_qpos = down_arm[:, -1, :] @@ -132,9 +141,9 @@ def execute(self, target: PressTarget, state: WorldState) -> ActionResult: press_arm_qpos, start_arm_qpos, n_back, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=self.cfg, + control_part=control_part, + arm_dof=manipulator.dof, + cfg=request.motion_policy, ) success = down_success & back_success @@ -148,53 +157,41 @@ def execute(self, target: PressTarget, state: WorldState) -> ActionResult: device=self.device, ) full[:, :, :] = state.last_qpos.unsqueeze(1) - full[:, :n_close, self.arm_joint_ids] = start_arm_qpos.unsqueeze(1) - full[:, :n_close, self.hand_joint_ids] = hand_close_path - full[:, n_close : n_close + n_down_actual, self.arm_joint_ids] = down_arm - full[:, n_close : n_close + n_down_actual, self.hand_joint_ids] = ( - self.hand_close_qpos.unsqueeze(1) + full[:, :n_close, arm_joint_ids] = start_arm_qpos.unsqueeze(1) + full[:, :n_close, hand_joint_ids] = hand_close_path + full[:, n_close : n_close + n_down_actual, arm_joint_ids] = down_arm + full[:, n_close : n_close + n_down_actual, hand_joint_ids] = ( + hand_close_qpos.unsqueeze(1) ) - full[:, n_close + n_down_actual :, self.arm_joint_ids] = back_arm - full[:, n_close + n_down_actual :, self.hand_joint_ids] = ( - self.hand_close_qpos.unsqueeze(1) + full[:, n_close + n_down_actual :, arm_joint_ids] = back_arm + full[:, n_close + n_down_actual :, hand_joint_ids] = hand_close_qpos.unsqueeze( + 1 ) - return ActionResult( + return self.build_plan( + request, + context, success=success, trajectory=full, - next_state=state.with_updates( - last_qpos=full[:, -1, :].clone(), - ), + phase_name="press", ) - def _compute_phase_waypoints(self) -> tuple[int, int, int]: - n_close = self.cfg.hand_interp_steps - if n_close < 1: - logger.log_error( - "hand_interp_steps must be at least 1 for PressCfg.", ValueError - ) + def _compute_phase_waypoints( + self, sample_count: int, options: PressOptions + ) -> tuple[int, int, int]: + """Split the invocation sample budget across press phases.""" + n_close = options.hand_interp_steps - motion_waypoints = self.cfg.sample_interval - n_close + motion_waypoints = sample_count - n_close n_down = motion_waypoints // 2 n_back = motion_waypoints - n_down if n_down < 2 or n_back < 2: logger.log_error( "Not enough waypoints for press trajectory. Increase " - "sample_interval or decrease hand_interp_steps.", + "MotionPolicy.sample_count or decrease hand_interp_steps.", ValueError, ) return n_close, n_down, n_back - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=torch.zeros(self.n_envs, dtype=torch.bool, device=self.device), - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, - ) - -__all__ = ["Press", "PressCfg", "PressTarget"] +__all__ = ["Press", "PressGoal", "PressOptions"] diff --git a/embodichain/lab/sim/atomic_actions/runtime.py b/embodichain/lab/sim/atomic_actions/runtime.py new file mode 100644 index 000000000..608fc1416 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/runtime.py @@ -0,0 +1,270 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Engine-owned planning resources shared by atomic actions.""" + +from __future__ import annotations + +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING + +import torch + +from .bindings import ActionBinding, ResolvedActionBinding, ResolvedControlPart +from .control import ( + ActionControlOverrides, + ControlCommand, + ControlPartCommandProfile, +) +from .core import resolve_runtime_device +from .trajectory import TrajectoryBuilder + +if TYPE_CHECKING: + from embodichain.lab.sim.objects import Robot + from embodichain.lab.sim.planners import MotionGenerator + + +class ActionPlanningServices: + """Planning resources exclusively owned by one atomic-action engine. + + An action may borrow these resources after the engine binds it, but callers + never pass a motion generator to individual actions. Keeping the generator + and trajectory builder here gives one engine a single planner backend, + robot, device, cache, and collision-world owner. + + Args: + motion_generator: Motion generator owned by the engine. + control_profiles: Semantic command profiles keyed by names from the + owned robot's ``control_parts`` mapping. + """ + + def __init__( + self, + motion_generator: MotionGenerator, + control_profiles: Mapping[str, ControlPartCommandProfile] | None = None, + ) -> None: + self._motion_generator = motion_generator + self._robot: Robot = motion_generator.robot + self._device = resolve_runtime_device(motion_generator.device) + self._trajectory_builder = TrajectoryBuilder(motion_generator) + self._control_profiles = self._snapshot_control_profiles( + {} if control_profiles is None else control_profiles + ) + self._binding_cache: dict[ + tuple[tuple[tuple[str, str], ...], tuple[tuple[str, str], ...]], + ResolvedActionBinding, + ] = {} + + @property + def motion_generator(self) -> MotionGenerator: + """Return the single motion generator owned by the engine.""" + return self._motion_generator + + @property + def robot(self) -> Robot: + """Return the robot planned by this service set.""" + return self._robot + + @property + def device(self) -> torch.device: + """Return the concrete device used for planning.""" + return self._device + + @property + def trajectory_builder(self) -> TrajectoryBuilder: + """Return the shared stateless trajectory builder.""" + return self._trajectory_builder + + @property + def control_profiles(self) -> Mapping[str, ControlPartCommandProfile]: + """Return owned semantic command profiles keyed by control-part name.""" + return MappingProxyType( + { + name: profile.snapshot() + for name, profile in self._control_profiles.items() + } + ) + + @property + def planner_name(self) -> str: + """Return the configured planner backend name.""" + planner_cfg = getattr( + getattr(self._motion_generator, "planner", None), "cfg", None + ) + planner_name = getattr(planner_cfg, "planner_type", None) + return "unknown" if planner_name is None else str(planner_name) + + def resolve_binding( + self, + binding: ActionBinding, + control_overrides: ActionControlOverrides | None = None, + ) -> ResolvedActionBinding: + """Resolve binding names against the owned robot's control parts. + + ``ActionBinding`` deliberately carries stable string references only. + This method establishes that every reference is a key in + ``Robot.control_parts`` and resolves its full-robot joint indices. + + Args: + binding: Semantic-role mapping to validate and resolve. + control_overrides: Optional per-role command replacements for this + invocation revision. + + Returns: + Immutable runtime resources for action planning. + + Raises: + TypeError: If ``binding`` or ``Robot.control_parts`` is invalid. + ValueError: If a referenced control part is unknown or empty. + """ + if not isinstance(binding, ActionBinding): + raise TypeError("binding must be an ActionBinding.") + cache_key = ( + tuple(sorted(binding.manipulators.items())), + tuple(sorted(binding.end_effectors.items())), + ) + resolved = self._binding_cache.get(cache_key) + if resolved is None: + control_parts = getattr(self.robot, "control_parts", None) + if not isinstance(control_parts, Mapping): + if binding.manipulators or binding.end_effectors: + raise TypeError( + "ActionBinding resources must come from " + "Robot.control_parts, but the engine robot does not " + "define a control-parts mapping." + ) + control_parts = {} + + resolved = ResolvedActionBinding( + manipulators=self._resolve_resource_map( + binding.manipulators, + control_parts=control_parts, + resource_kind="manipulator", + ), + end_effectors=self._resolve_resource_map( + binding.end_effectors, + control_parts=control_parts, + resource_kind="end effector", + ), + ) + self._binding_cache[cache_key] = resolved + + if control_overrides is None: + return resolved + if not isinstance(control_overrides, ActionControlOverrides): + raise TypeError("control_overrides must be an ActionControlOverrides.") + if control_overrides.is_empty: + return resolved + return ResolvedActionBinding( + manipulators=self._apply_command_overrides( + resolved.manipulators, + control_overrides.manipulators, + resource_kind="manipulator", + ), + end_effectors=self._apply_command_overrides( + resolved.end_effectors, + control_overrides.end_effectors, + resource_kind="end effector", + ), + ) + + def _resolve_resource_map( + self, + resources: Mapping[str, str], + *, + control_parts: Mapping[str, object], + resource_kind: str, + ) -> dict[str, ResolvedControlPart]: + """Resolve one role map through ``Robot.control_parts``.""" + available = sorted(str(name) for name in control_parts) + resolved: dict[str, ResolvedControlPart] = {} + for role, name in resources.items(): + if name not in control_parts: + raise ValueError( + f"ActionBinding {resource_kind} role {role!r} references " + f"control part {name!r}, but Robot.control_parts contains " + f"{available}." + ) + joint_ids = tuple(self.robot.get_joint_ids(name=name)) + if not joint_ids: + raise ValueError( + f"Robot control part {name!r} bound to {resource_kind} role " + f"{role!r} contains no joints." + ) + profile = self._control_profiles.get(name) + resolved[role] = ResolvedControlPart( + name=name, + joint_ids=joint_ids, + commands={} if profile is None else profile.commands, + ) + return resolved + + def _snapshot_control_profiles( + self, + profiles: Mapping[str, ControlPartCommandProfile], + ) -> Mapping[str, ControlPartCommandProfile]: + """Validate control-part profile ownership and freeze snapshots.""" + if not isinstance(profiles, Mapping): + raise TypeError("control_profiles must be a mapping.") + control_parts = getattr(self.robot, "control_parts", None) + if not isinstance(control_parts, Mapping): + if profiles: + raise TypeError( + "Control-part command profiles require Robot.control_parts." + ) + control_parts = {} + snapshots: dict[str, ControlPartCommandProfile] = {} + available = sorted(str(name) for name in control_parts) + for name, profile in profiles.items(): + if not isinstance(name, str) or not name.strip(): + raise ValueError( + "control_profiles keys must be non-empty control-part names." + ) + if name not in control_parts: + raise ValueError( + f"Control profile references unknown control part {name!r}; " + f"Robot.control_parts contains {available}." + ) + if not isinstance(profile, ControlPartCommandProfile): + raise TypeError( + "control_profiles values must be " + "ControlPartCommandProfile instances." + ) + snapshots[name] = profile.snapshot() + return MappingProxyType(snapshots) + + @staticmethod + def _apply_command_overrides( + resources: Mapping[str, ResolvedControlPart], + overrides: Mapping[str, Mapping[str, ControlCommand]], + *, + resource_kind: str, + ) -> dict[str, ResolvedControlPart]: + """Apply role-scoped commands to already resolved control parts.""" + unknown_roles = sorted(set(overrides) - set(resources)) + if unknown_roles: + raise KeyError( + f"Command overrides reference unbound {resource_kind} roles " + f"{unknown_roles}; bound roles are {sorted(resources)}." + ) + return { + role: resource.with_command_overrides(overrides.get(role, {})) + for role, resource in resources.items() + } + + +__all__ = ["ActionPlanningServices"] diff --git a/embodichain/lab/sim/atomic_actions/state.py b/embodichain/lab/sim/atomic_actions/state.py new file mode 100644 index 000000000..8b2044ab8 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/state.py @@ -0,0 +1,556 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Observed robot state, symbolic task state, and scene snapshots.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Mapping, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from .core import ObjectSemantics + + +def _resolve_runtime_device(device: torch.device | str) -> torch.device: + """Resolve an indexless CUDA device to the active concrete GPU index.""" + resolved = torch.device(device) + if resolved.type == "cuda" and resolved.index is None: + return torch.device(f"cuda:{torch.cuda.current_device()}") + return resolved + + +def _validate_pose(value: torch.Tensor, name: str) -> int | None: + """Validate a homogeneous transform and return its explicit batch size.""" + if not isinstance(value, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor.") + if value.shape == (4, 4): + return None + if value.dim() != 3 or value.shape[-2:] != (4, 4) or value.shape[0] == 0: + raise ValueError( + f"{name} must have shape (4, 4) or (n_envs, 4, 4), " + f"got {tuple(value.shape)}." + ) + return int(value.shape[0]) + + +def _normalize_mask( + value: torch.Tensor | None, + *, + batch_size: int, + device: torch.device, + name: str, +) -> torch.Tensor: + """Return an owned boolean mask with shape ``(batch_size,)``.""" + if value is None: + return torch.ones(batch_size, dtype=torch.bool, device=device) + if not isinstance(value, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor or None.") + if value.dtype != torch.bool: + raise TypeError(f"{name} must have dtype torch.bool, got {value.dtype}.") + if value.shape != (batch_size,): + raise ValueError( + f"{name} must have shape ({batch_size},), got {tuple(value.shape)}." + ) + return value.to(device=device).clone() + + +def _broadcast_pose( + value: torch.Tensor, + *, + batch_size: int, + device: torch.device, + name: str, +) -> torch.Tensor: + """Resolve an optionally batched pose to the task-state batch.""" + pose_batch_size = _validate_pose(value, name) + if value.device != device: + raise ValueError(f"{name} must use task-state device {device}.") + if pose_batch_size is None: + return value.unsqueeze(0).expand(batch_size, -1, -1).clone() + if pose_batch_size != batch_size: + raise ValueError( + f"{name} batch size must be {batch_size}, got {pose_batch_size}." + ) + return value.clone() + + +@dataclass(frozen=True, slots=True, eq=False) +class HeldObjectState: + """Observed or projected relation between an object and one manipulator.""" + + semantics: ObjectSemantics + """Semantics of the held object.""" + + object_to_eef: torch.Tensor + """Object-to-end-effector transform.""" + + grasp_xpos: torch.Tensor + """End-effector grasp pose.""" + + env_mask: torch.Tensor | None = None + """Environments in which the relation is active.""" + + def __post_init__(self) -> None: + from .core import ObjectSemantics + + if not isinstance(self.semantics, ObjectSemantics): + raise TypeError("semantics must be an ObjectSemantics instance.") + object_batch = _validate_pose(self.object_to_eef, "object_to_eef") + grasp_batch = _validate_pose(self.grasp_xpos, "grasp_xpos") + explicit_batches = { + size for size in (object_batch, grasp_batch) if size is not None + } + if len(explicit_batches) > 1: + raise ValueError("Held-object poses must use the same batch size.") + if self.object_to_eef.device != self.grasp_xpos.device: + raise ValueError("Held-object poses must use the same device.") + if self.env_mask is not None: + mask_batch = int(self.env_mask.shape[0]) if self.env_mask.dim() == 1 else -1 + batch_size = next(iter(explicit_batches), mask_batch) + object.__setattr__( + self, + "env_mask", + _normalize_mask( + self.env_mask, + batch_size=batch_size, + device=self.object_to_eef.device, + name="env_mask", + ), + ) + + +@dataclass(frozen=True, slots=True, eq=False) +class CoordinatedHeldObjectState: + """Observed or projected relation for an object held by two manipulators.""" + + semantics: ObjectSemantics + left_object_to_eef: torch.Tensor + right_object_to_eef: torch.Tensor + left_grasp_xpos: torch.Tensor + right_grasp_xpos: torch.Tensor + env_mask: torch.Tensor | None = None + + def __post_init__(self) -> None: + from .core import ObjectSemantics + + if not isinstance(self.semantics, ObjectSemantics): + raise TypeError("semantics must be an ObjectSemantics instance.") + poses = { + "left_object_to_eef": self.left_object_to_eef, + "right_object_to_eef": self.right_object_to_eef, + "left_grasp_xpos": self.left_grasp_xpos, + "right_grasp_xpos": self.right_grasp_xpos, + } + batches = {_validate_pose(value, name) for name, value in poses.items()} + batches.discard(None) + if len(batches) > 1: + raise ValueError("Coordinated held-object poses must share a batch size.") + if len({value.device for value in poses.values()}) != 1: + raise ValueError("Coordinated held-object poses must share a device.") + if self.env_mask is not None: + mask_batch = int(self.env_mask.shape[0]) if self.env_mask.dim() == 1 else -1 + batch_size = next(iter(batches), mask_batch) + object.__setattr__( + self, + "env_mask", + _normalize_mask( + self.env_mask, + batch_size=batch_size, + device=self.left_object_to_eef.device, + name="env_mask", + ), + ) + + +def _normalize_held( + value: HeldObjectState, + *, + batch_size: int, + device: torch.device, +) -> HeldObjectState: + """Normalize a held-object relation to one task-state batch.""" + return HeldObjectState( + semantics=value.semantics, + object_to_eef=_broadcast_pose( + value.object_to_eef, + batch_size=batch_size, + device=device, + name="HeldObjectState.object_to_eef", + ), + grasp_xpos=_broadcast_pose( + value.grasp_xpos, + batch_size=batch_size, + device=device, + name="HeldObjectState.grasp_xpos", + ), + env_mask=_normalize_mask( + value.env_mask, + batch_size=batch_size, + device=device, + name="HeldObjectState.env_mask", + ), + ) + + +def _normalize_coordinated_held( + value: CoordinatedHeldObjectState, + *, + batch_size: int, + device: torch.device, +) -> CoordinatedHeldObjectState: + """Normalize a coordinated relation to one task-state batch.""" + return CoordinatedHeldObjectState( + semantics=value.semantics, + left_object_to_eef=_broadcast_pose( + value.left_object_to_eef, + batch_size=batch_size, + device=device, + name="CoordinatedHeldObjectState.left_object_to_eef", + ), + right_object_to_eef=_broadcast_pose( + value.right_object_to_eef, + batch_size=batch_size, + device=device, + name="CoordinatedHeldObjectState.right_object_to_eef", + ), + left_grasp_xpos=_broadcast_pose( + value.left_grasp_xpos, + batch_size=batch_size, + device=device, + name="CoordinatedHeldObjectState.left_grasp_xpos", + ), + right_grasp_xpos=_broadcast_pose( + value.right_grasp_xpos, + batch_size=batch_size, + device=device, + name="CoordinatedHeldObjectState.right_grasp_xpos", + ), + env_mask=_normalize_mask( + value.env_mask, + batch_size=batch_size, + device=device, + name="CoordinatedHeldObjectState.env_mask", + ), + ) + + +@dataclass(frozen=True, slots=True, eq=False) +class TaskState: + """Symbolic task state, separate from measured robot state.""" + + batch_size: int + """Number of vectorized environments represented by the state.""" + + device: torch.device | str + """Device used by per-environment masks and relation tensors.""" + + held_objects: Mapping[str, HeldObjectState] = field(default_factory=dict) + """Single-manipulator held-object relations keyed by control resource.""" + + coordinated_held_objects: Mapping[tuple[str, str], CoordinatedHeldObjectState] = ( + field(default_factory=dict) + ) + """Two-manipulator held-object relations keyed by ordered resource pairs.""" + + def __post_init__(self) -> None: + if self.batch_size <= 0: + raise ValueError("TaskState.batch_size must be greater than zero.") + device = _resolve_runtime_device(self.device) + normalized_held: dict[str, HeldObjectState] = {} + for resource, value in self.held_objects.items(): + if not isinstance(resource, str) or not resource: + raise TypeError("held_objects keys must be non-empty strings.") + if not isinstance(value, HeldObjectState): + raise TypeError("held_objects values must be HeldObjectState objects.") + normalized_held[resource] = _normalize_held( + value, batch_size=self.batch_size, device=device + ) + + normalized_coordinated: dict[tuple[str, str], CoordinatedHeldObjectState] = {} + for resources, value in self.coordinated_held_objects.items(): + if ( + not isinstance(resources, tuple) + or len(resources) != 2 + or not all(isinstance(item, str) and item for item in resources) + ): + raise TypeError( + "coordinated_held_objects keys must be pairs of non-empty strings." + ) + if not isinstance(value, CoordinatedHeldObjectState): + raise TypeError( + "coordinated_held_objects values must be " + "CoordinatedHeldObjectState objects." + ) + normalized_coordinated[resources] = _normalize_coordinated_held( + value, batch_size=self.batch_size, device=device + ) + + object.__setattr__(self, "device", device) + object.__setattr__(self, "held_objects", MappingProxyType(normalized_held)) + object.__setattr__( + self, + "coordinated_held_objects", + MappingProxyType(normalized_coordinated), + ) + + @classmethod + def empty( + cls, + batch_size: int, + device: torch.device | str, + ) -> TaskState: + """Create an empty symbolic state. + + Args: + batch_size: Number of represented environments. + device: Tensor device used by the state. + + Returns: + Empty task state with explicit batch metadata. + """ + return cls(batch_size=batch_size, device=device) + + def get_held_object(self, resource: str) -> HeldObjectState | None: + """Return the object held by ``resource``, if any.""" + return self.held_objects.get(resource) + + def get_coordinated_held_object( + self, + first_resource: str, + second_resource: str, + ) -> CoordinatedHeldObjectState | None: + """Return the relation for an ordered resource pair, if any.""" + return self.coordinated_held_objects.get((first_resource, second_resource)) + + +@dataclass(frozen=True, slots=True, eq=False) +class RobotObservation: + """Measured robot state used as the start of planning or replanning.""" + + timestamp: float + qpos: torch.Tensor + qvel: torch.Tensor + qeffort: torch.Tensor | None = None + root_pose: torch.Tensor | None = None + root_twist: torch.Tensor | None = None + + def __post_init__(self) -> None: + if self.timestamp < 0.0: + raise ValueError("RobotObservation.timestamp must be non-negative.") + if not isinstance(self.qpos, torch.Tensor) or self.qpos.dim() != 2: + raise ValueError( + "RobotObservation.qpos must have shape (n_envs, robot_dof)." + ) + if self.qpos.shape[0] == 0 or self.qpos.shape[1] == 0: + raise ValueError("RobotObservation.qpos dimensions must be non-zero.") + if not isinstance(self.qvel, torch.Tensor): + raise TypeError("RobotObservation.qvel must be a torch.Tensor.") + if self.qvel.shape != self.qpos.shape: + raise ValueError("RobotObservation.qvel must match qpos shape.") + if self.qvel.device != self.qpos.device: + raise ValueError("RobotObservation.qpos and qvel must share a device.") + if self.qeffort is not None: + if self.qeffort.shape != self.qpos.shape: + raise ValueError("RobotObservation.qeffort must match qpos shape.") + if self.qeffort.device != self.qpos.device: + raise ValueError("RobotObservation.qeffort must share the qpos device.") + object.__setattr__(self, "qpos", self.qpos.clone()) + object.__setattr__(self, "qvel", self.qvel.clone()) + if self.qeffort is not None: + object.__setattr__(self, "qeffort", self.qeffort.clone()) + if self.root_pose is not None: + object.__setattr__(self, "root_pose", self.root_pose.clone()) + if self.root_twist is not None: + object.__setattr__(self, "root_twist", self.root_twist.clone()) + + @property + def batch_size(self) -> int: + """Number of represented vectorized environments.""" + return int(self.qpos.shape[0]) + + @property + def robot_dof(self) -> int: + """Number of robot joint-position columns.""" + return int(self.qpos.shape[1]) + + def with_qpos(self, qpos: torch.Tensor) -> RobotObservation: + """Create a projected observation with a new position and zero velocity. + + Args: + qpos: Projected joint positions with the same shape as this observation. + + Returns: + New observation suitable for compiling the next action. + """ + return RobotObservation( + timestamp=self.timestamp, + qpos=qpos, + qvel=torch.zeros_like(qpos), + qeffort=self.qeffort, + root_pose=self.root_pose, + root_twist=self.root_twist, + ) + + +@dataclass(frozen=True, slots=True, eq=False) +class EntityState: + """Scene entity state addressable by a stable entity identifier.""" + + pose: torch.Tensor + confidence: float = 1.0 + + def __post_init__(self) -> None: + _validate_pose(self.pose, "EntityState.pose") + if not 0.0 <= self.confidence <= 1.0: + raise ValueError("EntityState.confidence must be in [0, 1].") + object.__setattr__(self, "pose", self.pose.clone()) + + +@dataclass(frozen=True, slots=True, eq=False) +class SceneSnapshot: + """Versioned scene state used to ground dynamic goals and obstacles.""" + + timestamp: float + version: int + entities: Mapping[str, EntityState] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.timestamp < 0.0: + raise ValueError("SceneSnapshot.timestamp must be non-negative.") + if self.version < 0: + raise ValueError("SceneSnapshot.version must be non-negative.") + normalized: dict[str, EntityState] = {} + for entity_id, state in self.entities.items(): + if not isinstance(entity_id, str) or not entity_id: + raise ValueError("Scene entity identifiers must be non-empty strings.") + if not isinstance(state, EntityState): + raise TypeError( + "SceneSnapshot entities must contain EntityState values." + ) + normalized[entity_id] = state + object.__setattr__(self, "entities", MappingProxyType(normalized)) + + @classmethod + def empty(cls) -> SceneSnapshot: + """Create an empty initial scene snapshot.""" + return cls(timestamp=0.0, version=0) + + +@dataclass(frozen=True, slots=True, eq=False) +class PlanningContext: + """Complete side-effect-free input to :meth:`AtomicAction.plan`.""" + + robot: RobotObservation + task: TaskState + scene: SceneSnapshot + env_ids: torch.Tensor + + def __post_init__(self) -> None: + if not isinstance(self.robot, RobotObservation): + raise TypeError("robot must be a RobotObservation.") + if not isinstance(self.task, TaskState): + raise TypeError("task must be a TaskState.") + if not isinstance(self.scene, SceneSnapshot): + raise TypeError("scene must be a SceneSnapshot.") + if self.task.batch_size != self.robot.batch_size: + raise ValueError("TaskState and RobotObservation batch sizes must match.") + if self.task.device != self.robot.qpos.device: + raise ValueError("TaskState and RobotObservation must share a device.") + if not isinstance(self.env_ids, torch.Tensor): + raise TypeError("env_ids must be a torch.Tensor.") + if self.env_ids.dtype != torch.long: + raise TypeError("env_ids must have dtype torch.long.") + if self.env_ids.shape != (self.robot.batch_size,): + raise ValueError( + "env_ids must identify every row in the planning batch; expected " + f"shape ({self.robot.batch_size},), got {tuple(self.env_ids.shape)}." + ) + if self.env_ids.device != self.robot.qpos.device: + raise ValueError("env_ids and robot tensors must share a device.") + if torch.unique(self.env_ids).numel() != self.env_ids.numel(): + raise ValueError("env_ids must be unique.") + object.__setattr__(self, "env_ids", self.env_ids.clone()) + + @property + def batch_size(self) -> int: + """Number of environments in this planning request.""" + return self.robot.batch_size + + @property + def last_qpos(self) -> torch.Tensor: + """Measured joint positions used as the planning start state.""" + return self.robot.qpos + + @property + def held_objects(self) -> Mapping[str, HeldObjectState]: + """Single-resource held-object relations.""" + return self.task.held_objects + + @property + def coordinated_held_objects( + self, + ) -> Mapping[tuple[str, str], CoordinatedHeldObjectState]: + """Coordinated held-object relations.""" + return self.task.coordinated_held_objects + + def get_held_object(self, resource: str) -> HeldObjectState | None: + """Return the object held by ``resource``, if any.""" + return self.task.get_held_object(resource) + + def get_coordinated_held_object( + self, + first_resource: str, + second_resource: str, + ) -> CoordinatedHeldObjectState | None: + """Return a coordinated held-object relation, if any.""" + return self.task.get_coordinated_held_object(first_resource, second_resource) + + def project( + self, + *, + qpos: torch.Tensor, + task: TaskState, + ) -> PlanningContext: + """Create the hypothetical context used to compile a following action. + + Args: + qpos: Projected terminal joint positions. + task: Task state after applying expected effects. + + Returns: + New context. No measured state or simulator state is mutated. + """ + return PlanningContext( + robot=self.robot.with_qpos(qpos), + task=task, + scene=self.scene, + env_ids=self.env_ids, + ) + + +__all__ = [ + "CoordinatedHeldObjectState", + "EntityState", + "HeldObjectState", + "PlanningContext", + "RobotObservation", + "SceneSnapshot", + "TaskState", +] diff --git a/embodichain/lab/sim/atomic_actions/targets.py b/embodichain/lab/sim/atomic_actions/targets.py deleted file mode 100644 index 414a432ea..000000000 --- a/embodichain/lab/sim/atomic_actions/targets.py +++ /dev/null @@ -1,47 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""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/embodichain/lab/sim/atomic_actions/trajectory.py b/embodichain/lab/sim/atomic_actions/trajectory.py index e854aa580..7bfd4644b 100644 --- a/embodichain/lab/sim/atomic_actions/trajectory.py +++ b/embodichain/lab/sim/atomic_actions/trajectory.py @@ -31,30 +31,29 @@ from embodichain.lab.sim.utility.action_utils import interpolate_with_distance from embodichain.utils import logger +from .core import resolve_runtime_device +from .plans import TimedTrajectory + if TYPE_CHECKING: from embodichain.lab.sim.planners import MotionGenerator - -def _resolve_runtime_device(device: torch.device | str) -> torch.device: - """Resolve an indexless CUDA device to the active concrete GPU index.""" - resolved = torch.device(device) - if resolved.type == "cuda" and resolved.index is None: - return torch.device(f"cuda:{torch.cuda.current_device()}") - return resolved + from .policies import MotionPolicy class TrajectoryBuilder: """Stateless trajectory utilities shared by every atomic action. - Holds a reference to the motion generator (and through it, the robot and - device) so callers don't have to thread those through each helper call. - All methods are pure: no per-call state is kept on the builder. + :class:`~embodichain.lab.sim.atomic_actions.runtime.ActionPlanningServices` + creates one builder per engine. It holds a reference to the engine's motion + generator (and through it, the robot and device) so actions do not thread + those through each helper call. All methods are pure: no per-call state is + kept on the builder. """ def __init__(self, motion_generator: MotionGenerator) -> None: self.motion_generator = motion_generator self.robot = motion_generator.robot - self.device = _resolve_runtime_device(self.robot.device) + self.device = resolve_runtime_device(self.robot.device) # ------------------------------------------------------------------ # Success / shape helpers @@ -66,6 +65,49 @@ def all_envs_success(self, is_success: bool | torch.Tensor) -> bool: return bool(torch.all(is_success).item()) return bool(is_success) + def _resolve_success_mask( + self, + success: bool | torch.Tensor, + *, + n_envs: int, + name: str, + ) -> torch.Tensor: + """Normalize planner success to a boolean tensor with shape ``(n_envs,)``.""" + if isinstance(success, bool): + return torch.full((n_envs,), success, dtype=torch.bool, device=self.device) + if not isinstance(success, torch.Tensor): + logger.log_error( + f"{name} must be a bool or torch.Tensor, got " + f"{type(success).__name__}.", + TypeError, + ) + success = success.to(self.device) + if success.dtype != torch.bool: + integer_dtypes = { + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + } + if success.dtype not in integer_dtypes or not torch.all( + (success == 0) | (success == 1) + ): + logger.log_error( + f"{name} must be boolean or a binary integer tensor, " + f"got dtype {success.dtype}.", + TypeError, + ) + success = success.to(dtype=torch.bool) + if success.dim() == 0 or success.shape == (1,): + success = success.reshape(1).expand(n_envs) + if success.shape != (n_envs,): + logger.log_error( + f"{name} must have shape ({n_envs},), got {tuple(success.shape)}.", + ValueError, + ) + return success.clone() + def resolve_pose_target(self, target: torch.Tensor, *, n_envs: int) -> torch.Tensor: """Resolve an end-effector pose target into batched homogeneous transforms. @@ -85,7 +127,7 @@ def resolve_pose_target(self, target: torch.Tensor, *, n_envs: int) -> torch.Ten f"or ({n_envs}, n_waypoint, 4, 4)", TypeError, ) - target = target.to(device=self.device, dtype=torch.float32) + target = target.to(device=self.device, dtype=torch.float32).clone() if target.shape == (4, 4): target = target.unsqueeze(0).repeat(n_envs, 1, 1) if target.dim() == 3: @@ -165,7 +207,7 @@ def resolve_joint_target( f"({n_envs}, n_waypoint, {joint_dof})", TypeError, ) - target_qpos = target_qpos.to(device=self.device, dtype=torch.float32) + target_qpos = target_qpos.to(device=self.device, dtype=torch.float32).clone() if target_qpos.shape == (joint_dof,): target_qpos = target_qpos.unsqueeze(0).repeat(n_envs, 1) if target_qpos.dim() == 2: @@ -299,7 +341,7 @@ def split_three_phase( # Arm trajectory planning # ------------------------------------------------------------------ - def plan_arm_traj( + def generate_arm_plan( self, target_states_list: list[list[PlanState]], start_qpos: torch.Tensor, @@ -307,18 +349,26 @@ def plan_arm_traj( *, control_part: str, arm_dof: int, - cfg: "ActionCfg | None" = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Plan batched arm trajectories for all environments. + cfg: "MotionPolicy | None" = None, + ) -> PlanResult: + """Generate a normalized arm plan while retaining backend timing. + + Args: + target_states_list: Cartesian planner states grouped by environment. + start_qpos: Controlled-joint start positions. + n_waypoints: Requested output sample count. + control_part: Bound robot control resource. + arm_dof: Number of controlled joints. + cfg: Motion-generation policy. - Returns ``(success:(B,), trajectory:(B, n_waypoints, arm_dof))``. - ``cfg.motion_source`` selects 'ik_interp' (default) or 'motion_gen'. + Returns: + Normalized planner result. Failed rows contain hold positions. """ motion_source = ( getattr(cfg, "motion_source", "ik_interp") if cfg else "ik_interp" ) if motion_source == "motion_gen": - return self._plan_motion_gen( + return self._generate_motion_gen_plan( target_states_list, start_qpos, n_waypoints, @@ -326,13 +376,47 @@ def plan_arm_traj( arm_dof=arm_dof, cfg=cfg, ) - return self._plan_ik_interp( + success, positions = self._plan_ik_interp( target_states_list, start_qpos, n_waypoints, control_part=control_part, arm_dof=arm_dof, ) + return PlanResult(success=success, positions=positions) + + def plan_arm_traj( + self, + target_states_list: list[list[PlanState]], + start_qpos: torch.Tensor, + n_waypoints: int, + *, + control_part: str, + arm_dof: int, + cfg: "MotionPolicy | None" = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Plan batched arm trajectories for all environments. + + Returns ``(success:(B,), trajectory:(B, n_waypoints, arm_dof))``. + ``cfg.motion_source`` selects 'ik_interp' (default) or 'motion_gen'. + """ + result = self.generate_arm_plan( + target_states_list, + start_qpos, + n_waypoints, + control_part=control_part, + arm_dof=arm_dof, + cfg=cfg, + ) + assert result.positions is not None + return ( + self._resolve_success_mask( + result.success, + n_envs=start_qpos.shape[0], + name="Arm plan success", + ), + result.positions, + ) def _plan_ik_interp( self, @@ -362,6 +446,11 @@ def _plan_ik_interp( is_success, qpos = self.robot.compute_ik( pose=xpos_traj[:, j], name=control_part, joint_seed=qpos_seed ) + is_success = self._resolve_success_mask( + is_success, + n_envs=n_envs, + name=f"IK success for target state {j}", + ) if not self.all_envs_success(is_success): logger.log_warning( f"Failed to compute IK for target state {j} in some environments." @@ -387,9 +476,38 @@ def _plan_motion_gen( *, control_part: str, arm_dof: int, - cfg: "ActionCfg | None", + cfg: "MotionPolicy | None", ) -> tuple[torch.Tensor, torch.Tensor]: """Motion-generator trajectory source for Cartesian (EEF) targets.""" + result = self._generate_motion_gen_plan( + target_states_list, + start_qpos, + n_waypoints, + control_part=control_part, + arm_dof=arm_dof, + cfg=cfg, + ) + assert result.positions is not None + return ( + self._resolve_success_mask( + result.success, + n_envs=start_qpos.shape[0], + name="MotionGenerator PlanResult.success", + ), + result.positions, + ) + + def _generate_motion_gen_plan( + self, + target_states_list: list[list[PlanState]], + start_qpos: torch.Tensor, + n_waypoints: int, + *, + control_part: str, + arm_dof: int, + cfg: "MotionPolicy | None", + ) -> PlanResult: + """Generate and normalize one Cartesian motion-generator result.""" if self.motion_generator is None: logger.log_error( "motion_source='motion_gen' requires a MotionGenerator on the engine", @@ -407,7 +525,9 @@ def _plan_motion_gen( is_interpolate=True, ), ) - return self._process_motion_gen_result(result, start_qpos, n_waypoints, arm_dof) + return self._normalize_motion_gen_result( + result, start_qpos, n_waypoints, arm_dof + ) def _process_motion_gen_result( self, @@ -417,13 +537,34 @@ def _process_motion_gen_result( arm_dof: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Validate a MotionGenerator PlanResult and apply sample/hold policy.""" - success = ( - result.success.to(self.device) - if isinstance(result.success, torch.Tensor) - else torch.tensor(result.success, device=self.device) + normalized = self._normalize_motion_gen_result( + result, start_qpos, n_waypoints, arm_dof ) - positions = result.positions + assert normalized.positions is not None + return ( + self._resolve_success_mask( + normalized.success, + n_envs=start_qpos.shape[0], + name="MotionGenerator PlanResult.success", + ), + normalized.positions, + ) + + def _normalize_motion_gen_result( + self, + result: PlanResult, + start_qpos: torch.Tensor, + n_waypoints: int, + arm_dof: int, + ) -> PlanResult: + """Validate a planner result and retain compatible timing metadata.""" n_envs = start_qpos.shape[0] + success = self._resolve_success_mask( + result.success, + n_envs=n_envs, + name="MotionGenerator PlanResult.success", + ) + positions = result.positions if positions is None or positions.ndim != 3: logger.log_error( "MotionGenerator returned no (B, N, controlled_dof) positions", @@ -441,17 +582,74 @@ def _process_motion_gen_result( "MotionGenerator returned non-finite or wrong-device positions", ValueError, ) + resampled = False if not self.motion_generator.planner.preserve_plan_samples: if positions.shape[1] != n_waypoints: positions = interpolate_with_distance( trajectory=positions, interp_num=n_waypoints, device=self.device ) + resampled = True positions = positions.to(self.device) + + def normalize_derivative( + value: torch.Tensor | None, + name: str, + ) -> torch.Tensor | None: + if value is None or resampled: + return None + if value.shape != positions.shape: + logger.log_error( + f"MotionGenerator {name} must match positions shape, " + f"got {tuple(value.shape)} and {tuple(positions.shape)}.", + ValueError, + ) + value = value.to(self.device) + if not torch.isfinite(value).all(): + logger.log_error( + f"MotionGenerator returned non-finite {name}.", ValueError + ) + return value + + velocities = normalize_derivative(result.velocities, "velocities") + accelerations = normalize_derivative(result.accelerations, "accelerations") + dt = None if resampled else result.dt + if dt is not None: + dt = dt.to(device=self.device, dtype=torch.float32) + if dt.shape != positions.shape[:2]: + logger.log_error( + "MotionGenerator dt must match the positions batch and sample " + f"dimensions, got {tuple(dt.shape)} and {tuple(positions.shape[:2])}.", + ValueError, + ) + if not torch.isfinite(dt).all() or (dt < 0).any(): + logger.log_error( + "MotionGenerator returned invalid time deltas.", ValueError + ) + duration: float | torch.Tensor = dt.sum(dim=1) + else: + duration = result.duration # Failed envs hold start qpos across all waypoints. if not success.all(): held = start_qpos.unsqueeze(1).repeat(1, positions.shape[1], 1) positions = torch.where(success[:, None, None], positions, held) - return success, positions + if velocities is not None: + velocities = torch.where( + success[:, None, None], velocities, torch.zeros_like(velocities) + ) + if accelerations is not None: + accelerations = torch.where( + success[:, None, None], + accelerations, + torch.zeros_like(accelerations), + ) + return PlanResult( + success=success, + positions=positions, + velocities=velocities, + accelerations=accelerations, + dt=dt, + duration=duration, + ) def _to_batched_plan_states( self, target_states_list: list[list[PlanState]], n_envs: int @@ -489,7 +687,11 @@ def _to_batched_plan_states( ) return batched - def _build_plan_opts(self, cfg: "ActionCfg | None", n_waypoints: int): + def _build_plan_opts( + self, + cfg: "MotionPolicy | None", + n_waypoints: int, + ): """Build planner options from action configuration (three-way factory).""" configured_plan_opts = getattr(cfg, "plan_opts", None) if configured_plan_opts is not None: @@ -524,7 +726,7 @@ def _build_plan_opts(self, cfg: "ActionCfg | None", n_waypoints: int): ValueError, ) - def plan_joint_motion( + def generate_joint_plan( self, start_qpos: torch.Tensor, target_qpos: torch.Tensor, @@ -532,9 +734,9 @@ def plan_joint_motion( *, control_part: str, arm_dof: int, - cfg: "ActionCfg | None" = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Plan a joint-space trajectory through one or more target waypoints. + cfg: "MotionPolicy | None" = None, + ) -> PlanResult: + """Generate a joint-space plan while retaining backend timing. For ``motion_source='motion_gen'``, this delegates only when the selected backend includes :attr:`MoveType.JOINT_MOVE` in its supported @@ -544,7 +746,7 @@ def plan_joint_motion( interpolation. Returns: - ``(success:(B,), trajectory:(B, N, arm_dof))``. + Normalized planner result. Failed rows contain hold positions. """ motion_source = ( getattr(cfg, "motion_source", "ik_interp") if cfg else "ik_interp" @@ -572,12 +774,108 @@ def plan_joint_motion( is_interpolate=True, ), ) - return self._process_motion_gen_result( + return self._normalize_motion_gen_result( result, start_qpos, n_waypoints, arm_dof ) success = torch.ones(start_qpos.shape[0], dtype=torch.bool, device=self.device) trajectory = self.plan_joint_traj(start_qpos, target_qpos, n_waypoints) - return success, trajectory + return PlanResult(success=success, positions=trajectory) + + def plan_joint_motion( + self, + start_qpos: torch.Tensor, + target_qpos: torch.Tensor, + n_waypoints: int, + *, + control_part: str, + arm_dof: int, + cfg: "MotionPolicy | None" = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return the position-only view used by phased complex skills.""" + result = self.generate_joint_plan( + start_qpos, + target_qpos, + n_waypoints, + control_part=control_part, + arm_dof=arm_dof, + cfg=cfg, + ) + assert result.positions is not None + return ( + self._resolve_success_mask( + result.success, + n_envs=start_qpos.shape[0], + name="Joint plan success", + ), + result.positions, + ) + + def to_full_robot_trajectory( + self, + result: PlanResult, + *, + base_qpos: torch.Tensor, + joint_ids: list[int], + env_ids: torch.Tensor, + control_dt: float, + ) -> tuple[torch.Tensor, TimedTrajectory]: + """Embed a controlled-joint planner result into a timed robot trajectory. + + Args: + result: Normalized controlled-joint planner result. + base_qpos: Full-robot start positions with shape ``(B, robot_dof)``. + joint_ids: Full-robot columns controlled by the plan. + env_ids: Stable environment identifiers. + control_dt: Fallback command period when timing is absent. + + Returns: + Per-environment success and full-robot timed trajectory. + """ + positions = result.positions + if positions is None or positions.dim() != 3: + raise ValueError( + "PlanResult.positions must have shape (B, N, control_dof)." + ) + if positions.shape[0] != base_qpos.shape[0]: + raise ValueError("PlanResult and base_qpos batch sizes must match.") + if positions.shape[2] != len(joint_ids): + raise ValueError("PlanResult controlled DoF does not match joint_ids.") + full_positions = ( + base_qpos.unsqueeze(1).expand(-1, positions.shape[1], -1).clone() + ) + full_positions[:, :, joint_ids] = positions + + def embed_derivative(value: torch.Tensor | None) -> torch.Tensor | None: + if value is None: + return None + full = torch.zeros_like(full_positions) + full[:, :, joint_ids] = value + return full + + duration: float | torch.Tensor | None = result.duration + if result.dt is None: + duration_tensor = torch.as_tensor( + result.duration, + dtype=torch.float32, + device=base_qpos.device, + ) + if not bool((duration_tensor > 0.0).any().item()): + duration = None + timed = TimedTrajectory.from_positions( + full_positions, + env_ids=env_ids, + control_dt=control_dt, + velocities=embed_derivative(result.velocities), + accelerations=embed_derivative(result.accelerations), + dt=result.dt, + duration=duration, + ) + success = self._resolve_success_mask( + result.success, + n_envs=base_qpos.shape[0], + name="PlanResult.success", + ) + return success, timed def plan_joint_traj( self, diff --git a/embodichain/lab/sim/planners/curobo/curobo_planner.py b/embodichain/lab/sim/planners/curobo/curobo_planner.py index 47b9f5f7b..1aa173696 100644 --- a/embodichain/lab/sim/planners/curobo/curobo_planner.py +++ b/embodichain/lab/sim/planners/curobo/curobo_planner.py @@ -369,7 +369,7 @@ class CuroboPlannerCfg(BasePlannerCfg): When ``False`` (default), :class:`~embodichain.lab.sim.atomic_actions.trajectory.TrajectoryBuilder` resamples the returned trajectory to the atomic action's ``sample_interval`` waypoint count - matching the documented contract of - :class:`~embodichain.lab.sim.atomic_actions.primitives.move_end_effector.MoveEndEffectorCfg.sample_interval` + :attr:`~embodichain.lab.sim.atomic_actions.MotionPolicy.sample_count` and the other planners. The resample is arc-length piecewise-linear along cuRobo's joint-space path, so the collision-free path is preserved; only the sample density changes (cuRobo's own count is derived from diff --git a/examples/sim/planners/curobo_planner.py b/examples/sim/planners/curobo_planner.py index 8e71ee8d1..e12c0e717 100644 --- a/examples/sim/planners/curobo_planner.py +++ b/examples/sim/planners/curobo_planner.py @@ -59,10 +59,11 @@ visualization_cfg_from_args, ) from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - EndEffectorPoseTarget, - MoveEndEffector, - MoveEndEffectorCfg, + EndEffectorPoseGoal, + MotionPolicy, ) from embodichain.data import get_data_path from embodichain.lab.sim.cfg import RenderCfg, RigidBodyAttributesCfg @@ -747,26 +748,16 @@ def main() -> None: ) ) engine = AtomicActionEngine(motion_generator) - engine.register( - MoveEndEffector( - motion_generator, - MoveEndEffectorCfg( - motion_source="motion_gen", - control_part=control_part, - plan_opts=CuroboPlanOptions( - dynamic_obstacle_poses=( - obstacle_poses if use_independent_worlds else None - ), - max_attempts=args.max_attempts, - ), - # sample_interval sets the returned trajectory's waypoint count. - # cuRobo's own collision-checked samples are arc-length resampled - # to this count; set CuroboPlannerCfg.preserve_plan_samples=True - # above to keep cuRobo's raw samples (count from interpolation_dt). - sample_interval=30, + binding = ActionBinding(manipulators={"primary": control_part}) + motion_policy = MotionPolicy( + motion_source="motion_gen", + plan_opts=CuroboPlanOptions( + dynamic_obstacle_poses=( + obstacle_poses if use_independent_worlds else None ), + max_attempts=args.max_attempts, ), - name="move_end_effector", + sample_count=30, ) initial_qpos = robot.get_qpos(name=control_part) @@ -776,9 +767,18 @@ def main() -> None: to_matrix=True, ) plan_start = time.perf_counter() - success, trajectory, _ = engine.run( - [("move_end_effector", EndEffectorPoseTarget(xpos=target_xpos))] + compiled = engine.compile( + ( + ActionInvocation( + "move_end_effector", + EndEffectorPoseGoal(xpos=target_xpos), + binding, + motion_policy, + ), + ) ) + success = compiled.plan_success + trajectory = compiled.trajectory.positions planning_duration = time.perf_counter() - plan_start print(f"cuRobo atomic-action success by environment: {success.tolist()}") @@ -805,9 +805,18 @@ def main() -> None: print(f"maximum final TCP position error: {final_errors.max().item():.4f} m") plan_start = time.perf_counter() - success, trajectory, _ = engine.run( - [("move_end_effector", EndEffectorPoseTarget(xpos=initial_xpos))] + compiled = engine.compile( + ( + ActionInvocation( + "move_end_effector", + EndEffectorPoseGoal(xpos=initial_xpos), + binding, + motion_policy, + ), + ) ) + success = compiled.plan_success + trajectory = compiled.trajectory.positions planning_duration = time.perf_counter() - plan_start print(f"cuRobo return-action success by environment: {success.tolist()}") print(f"full-DoF trajectory shape: {tuple(trajectory.shape)}") diff --git a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py index d4df76ea8..8ad488594 100644 --- a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py +++ b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py @@ -120,17 +120,30 @@ def _run_case( ): """Run one MoveEndEffector case.""" torch = ensure_torch() - from embodichain.lab.sim.atomic_actions import EndEffectorPoseTarget + from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + EndEffectorPoseGoal, + MotionPolicy, + ) reset_robot(robot, initial_qpos) target_pose = _make_pose(sim.device, pose_case.xyz) elapsed, mem_delta, peak_gpu, result = timed_call( - lambda: atomic_engine.run( - steps=[("move_end_effector", EndEffectorPoseTarget(xpos=target_pose))] + lambda: atomic_engine.compile( + ( + ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal(xpos=target_pose), + binding=ActionBinding(manipulators={"primary": "arm"}), + motion_policy=MotionPolicy(sample_count=MOVE_SAMPLE_INTERVAL), + ), + ) ) ) - is_success, traj, _ = result + is_success = result.plan_success + traj = result.trajectory.positions video_path = None if should_record_case(args, recorded_count, bool(is_success)): reset_robot(robot, initial_qpos) @@ -220,8 +233,6 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: ensure_torch() from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, - MoveEndEffector, - MoveEndEffectorCfg, ) from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg from embodichain.lab.sim.planners import ToppraPlannerCfg @@ -251,14 +262,6 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) ) atomic_engine = AtomicActionEngine(motion_generator=motion_gen) - atomic_engine.register( - MoveEndEffector( - motion_gen, - cfg=MoveEndEffectorCfg( - control_part="arm", sample_interval=MOVE_SAMPLE_INTERVAL - ), - ) - ) results: list[dict[str, object]] = [] video_paths: list[str] = [] diff --git a/scripts/benchmark/atomic_action/move_held_object_benchmark.py b/scripts/benchmark/atomic_action/move_held_object_benchmark.py index 97e7d14ea..da623c0f1 100644 --- a/scripts/benchmark/atomic_action/move_held_object_benchmark.py +++ b/scripts/benchmark/atomic_action/move_held_object_benchmark.py @@ -175,13 +175,14 @@ def _prepare_held_state( ): """Run PickUp precondition outside the timed MoveHeldObject block.""" from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - EndEffectorPoseTarget, - GraspTarget, - MoveEndEffector, - MoveEndEffectorCfg, - PickUp, - PickUpCfg, + ControlPartCommandProfile, + EndEffectorPoseGoal, + GraspGoal, + MotionPolicy, + PickUpOptions, ) from scripts.tutorials.atomic_action.move_held_object import ( build_grasp_generator_cfg, @@ -191,33 +192,14 @@ def _prepare_held_state( ) hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) - atomic_engine = AtomicActionEngine(motion_generator=motion_gen) - atomic_engine.register( - MoveEndEffector( - motion_gen, - cfg=MoveEndEffectorCfg( - control_part="arm", - sample_interval=MOVE_SAMPLE_INTERVAL, - ), - ) - ) - atomic_engine.register( - PickUp( - motion_gen, - cfg=PickUpCfg( - control_part="arm", - hand_control_part="hand", - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - approach_direction=resolve_pickup_approach_direction( - pickup_approach, position_case, sim.device - ), - pre_grasp_distance=0.15, - lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) + atomic_engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, ) semantics = create_antipodal_object_semantics( obj=obj, @@ -230,12 +212,37 @@ def _prepare_held_state( move_position = obj_pose[0, :3, 3].clone() move_position[2] = 0.36 move_target = make_pre_pick_eef_pose(robot, move_position) - is_success, traj, state = atomic_engine.run( - steps=[ - ("move_end_effector", EndEffectorPoseTarget(xpos=move_target)), - ("pick_up", GraspTarget(semantics=semantics)), - ] + binding = ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ) + result = atomic_engine.compile( + ( + ActionInvocation( + "move_end_effector", + EndEffectorPoseGoal(xpos=move_target), + binding, + MotionPolicy(sample_count=MOVE_SAMPLE_INTERVAL), + ), + ActionInvocation( + "pick_up", + GraspGoal(semantics=semantics), + binding, + MotionPolicy(sample_count=PICK_SAMPLE_INTERVAL), + skill_options=PickUpOptions( + approach_direction=resolve_pickup_approach_direction( + pickup_approach, position_case, sim.device + ), + pre_grasp_distance=0.15, + lift_height=0.16, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ), + ) ) + is_success = result.plan_success + traj = result.trajectory.positions + state = result.projected_context if not is_success or state.get_held_object("arm") is None: raise RuntimeError( "Failed to prepare held-object state for MoveHeldObject benchmark." @@ -262,10 +269,12 @@ def _run_case( ): """Run one MoveHeldObject benchmark case.""" from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - HeldObjectPoseTarget, - MoveHeldObject, - MoveHeldObjectCfg, + ControlPartCommandProfile, + HeldObjectPoseGoal, + MotionPolicy, ) from scripts.tutorials.atomic_action.move_held_object import ( compute_pick_close_end_step, @@ -301,31 +310,36 @@ def _run_case( pickup_approach_direction_tuple(pickup_approach, position_case) ) precondition_waypoints = int(precondition_traj.shape[1]) - atomic_engine = AtomicActionEngine(motion_generator=motion_gen) - atomic_engine.register( - MoveHeldObject( - motion_gen, - cfg=MoveHeldObjectCfg( - control_part="arm", - hand_control_part="hand", - hand_close_qpos=hand_close, - sample_interval=MOVE_HELD_OBJECT_SAMPLE_INTERVAL, - ), - ) + atomic_engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + grasp=hand_close, + ) + }, ) target_pose = _make_object_target_pose(sim.device, case.xyz) elapsed, mem_delta, peak_gpu, result = timed_call( - lambda: atomic_engine.run( - steps=[ - ( - "move_held_object", - HeldObjectPoseTarget(object_target_pose=target_pose), - ) - ], - state=state, + lambda: atomic_engine.compile( + ( + ActionInvocation( + skill_id="move_held_object", + goal=HeldObjectPoseGoal(object_target_pose=target_pose), + binding=ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ), + motion_policy=MotionPolicy( + sample_count=MOVE_HELD_OBJECT_SAMPLE_INTERVAL + ), + ), + ), + context=state, ) ) - is_success, traj, final_state = result + is_success = result.plan_success + traj = result.trajectory.positions + final_state = result.projected_context torch = ensure_torch() precondition_obj_position = None final_obj_position = None diff --git a/scripts/benchmark/atomic_action/move_joints_benchmark.py b/scripts/benchmark/atomic_action/move_joints_benchmark.py index 1f89def0a..0fc6e6272 100644 --- a/scripts/benchmark/atomic_action/move_joints_benchmark.py +++ b/scripts/benchmark/atomic_action/move_joints_benchmark.py @@ -109,22 +109,29 @@ def _qpos(values, device): def _targets_for_sequence(sequence_case: JointSequenceCase, device): """Build typed MoveJoints targets for a sequence case.""" from embodichain.lab.sim.atomic_actions import ( - JointPositionTarget, - NamedJointPositionTarget, + ActionBinding, + ActionInvocation, + JointPositionGoal, + MotionPolicy, ) targets = [] + binding = ActionBinding(manipulators={"primary": "arm"}) + policy = MotionPolicy(sample_count=MOVE_JOINTS_SAMPLE_INTERVAL) for index, name in enumerate(sequence_case.sequence): if index == 0 and name == "ready": - targets.append(("move_joints", NamedJointPositionTarget(name="ready"))) + goal = JointPositionGoal(target="ready") else: - targets.append( - ( - "move_joints", - JointPositionTarget(qpos=_qpos(JOINT_TARGETS[name], device)), - ) + goal = JointPositionGoal(target=_qpos(JOINT_TARGETS[name], device)) + targets.append( + ActionInvocation( + skill_id="move_joints", + goal=goal, + binding=binding, + motion_policy=policy, ) - return targets + ) + return tuple(targets) def _run_case( @@ -141,8 +148,11 @@ def _run_case( torch = ensure_torch() reset_robot(robot, initial_qpos) steps = _targets_for_sequence(case, sim.device) - elapsed, mem_delta, peak_gpu, result = timed_call(lambda: atomic_engine.run(steps)) - is_success, traj, _ = result + elapsed, mem_delta, peak_gpu, result = timed_call( + lambda: atomic_engine.compile(steps) + ) + is_success = result.plan_success + traj = result.trajectory.positions video_path = None if should_record_case(args, recorded_count, bool(is_success)): reset_robot(robot, initial_qpos) @@ -230,8 +240,7 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: ensure_torch() from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, - MoveJoints, - MoveJointsCfg, + ControlPartCommandProfile, ) from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg from embodichain.lab.sim.planners import ToppraPlannerCfg @@ -261,18 +270,12 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) ) ready_qpos = _qpos(READY_QPOS, sim.device) - atomic_engine = AtomicActionEngine(motion_generator=motion_gen) - atomic_engine.register( - MoveJoints( - motion_gen, - cfg=MoveJointsCfg( - control_part="arm", - sample_interval=MOVE_JOINTS_SAMPLE_INTERVAL, - named_joint_positions={"ready": ready_qpos}, - ), - ) + atomic_engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "arm": ControlPartCommandProfile.joint_positions(ready=ready_qpos), + }, ) - results: list[dict[str, object]] = [] video_paths: list[str] = [] print("\n=== MoveJoints Sequence Sweep ===") diff --git a/scripts/benchmark/atomic_action/pickup_benchmark.py b/scripts/benchmark/atomic_action/pickup_benchmark.py index 9c8f56ff7..8e6686d88 100644 --- a/scripts/benchmark/atomic_action/pickup_benchmark.py +++ b/scripts/benchmark/atomic_action/pickup_benchmark.py @@ -123,10 +123,13 @@ def _run_case( ): """Run one PickUp benchmark case.""" from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - GraspTarget, - PickUp, - PickUpCfg, + ControlPartCommandProfile, + GraspGoal, + PickUpOptions, + MotionPolicy, ) from scripts.tutorials.atomic_action.pickup import ( build_grasp_generator_cfg, @@ -156,22 +159,14 @@ def _run_case( approach_direction_text = format_vector3( pickup_approach_direction_tuple(approach, position_case) ) - atomic_engine = AtomicActionEngine(motion_generator=motion_gen) - atomic_engine.register( - PickUp( - motion_gen, - cfg=PickUpCfg( - control_part="arm", - hand_control_part="hand", - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - approach_direction=approach_direction, - pre_grasp_distance=0.15, - lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) + atomic_engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, ) semantics = create_antipodal_object_semantics( obj=obj, @@ -181,11 +176,29 @@ def _run_case( build_grasp_generator_cfg=build_grasp_generator_cfg, ) elapsed, mem_delta, peak_gpu, result = timed_call( - lambda: atomic_engine.run( - steps=[("pick_up", GraspTarget(semantics=semantics))] + lambda: atomic_engine.compile( + ( + ActionInvocation( + skill_id="pick_up", + goal=GraspGoal(semantics=semantics), + binding=ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ), + motion_policy=MotionPolicy(sample_count=PICK_SAMPLE_INTERVAL), + skill_options=PickUpOptions( + approach_direction=approach_direction, + pre_grasp_distance=0.15, + lift_height=0.16, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ), + ) ) ) - is_success, traj, final_state = result + is_success = result.plan_success + traj = result.trajectory.positions + final_state = result.projected_context final_obj_position = None object_lift_delta_m = None object_xy_drift_m = None diff --git a/scripts/benchmark/atomic_action/place_benchmark.py b/scripts/benchmark/atomic_action/place_benchmark.py index 39d30100a..071be0d51 100644 --- a/scripts/benchmark/atomic_action/place_benchmark.py +++ b/scripts/benchmark/atomic_action/place_benchmark.py @@ -174,10 +174,13 @@ def _prepare_held_state( ): """Run PickUp precondition outside the timed Place block.""" from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - GraspTarget, - PickUp, - PickUpCfg, + ControlPartCommandProfile, + GraspGoal, + PickUpOptions, + MotionPolicy, ) from scripts.tutorials.atomic_action.place import ( build_grasp_generator_cfg, @@ -188,24 +191,14 @@ def _prepare_held_state( hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) initialize_pre_pick_robot_pose(robot, obj, hand_open) - atomic_engine = AtomicActionEngine(motion_generator=motion_gen) - atomic_engine.register( - PickUp( - motion_gen, - cfg=PickUpCfg( - control_part="arm", - hand_control_part="hand", - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - approach_direction=resolve_pickup_approach_direction( - pickup_approach, position_case, sim.device - ), - pre_grasp_distance=0.15, - lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) + atomic_engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, ) semantics = create_antipodal_object_semantics( obj=obj, @@ -214,9 +207,30 @@ def _prepare_held_state( build_gripper_collision_cfg=build_gripper_collision_cfg, build_grasp_generator_cfg=build_grasp_generator_cfg, ) - is_success, traj, state = atomic_engine.run( - steps=[("pick_up", GraspTarget(semantics=semantics))] + result = atomic_engine.compile( + ( + ActionInvocation( + skill_id="pick_up", + goal=GraspGoal(semantics=semantics), + binding=ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ), + motion_policy=MotionPolicy(sample_count=PICK_SAMPLE_INTERVAL), + skill_options=PickUpOptions( + approach_direction=resolve_pickup_approach_direction( + pickup_approach, position_case, sim.device + ), + pre_grasp_distance=0.15, + lift_height=0.16, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ), + ) ) + is_success = result.plan_success + traj = result.trajectory.positions + state = result.projected_context 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) @@ -241,10 +255,13 @@ def _run_case( ): """Run one Place benchmark case.""" from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - Place, - PlaceCfg, - PlaceTarget, + ControlPartCommandProfile, + MotionPolicy, + PlaceGoal, + PlaceOptions, ) from scripts.tutorials.atomic_action.place import ( compute_pick_close_end_step, @@ -281,29 +298,39 @@ def _run_case( pickup_approach_direction_tuple(pickup_approach, position_case) ) precondition_waypoints = int(precondition_traj.shape[1]) - atomic_engine = AtomicActionEngine(motion_generator=motion_gen) - atomic_engine.register( - Place( - motion_gen, - cfg=PlaceCfg( - control_part="arm", - hand_control_part="hand", - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - lift_height=PLACE_LIFT_HEIGHT, - sample_interval=PLACE_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) + atomic_engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, ) place_pose = _make_place_pose(sim.device, case.xyz) elapsed, mem_delta, peak_gpu, result = timed_call( - lambda: atomic_engine.run( - steps=[("place", PlaceTarget(xpos=place_pose))], - state=state, + lambda: atomic_engine.compile( + ( + ActionInvocation( + skill_id="place", + goal=PlaceGoal(xpos=place_pose), + binding=ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ), + motion_policy=MotionPolicy(sample_count=PLACE_SAMPLE_INTERVAL), + skill_options=PlaceOptions( + lift_height=PLACE_LIFT_HEIGHT, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ), + ), + context=state, ) ) - is_success, traj, final_state = result + is_success = result.plan_success + traj = result.trajectory.positions + final_state = result.projected_context torch = ensure_torch() precondition_obj_position = None final_obj_position = None diff --git a/scripts/benchmark/atomic_action/press_benchmark.py b/scripts/benchmark/atomic_action/press_benchmark.py index 549429346..4d408ce41 100644 --- a/scripts/benchmark/atomic_action/press_benchmark.py +++ b/scripts/benchmark/atomic_action/press_benchmark.py @@ -81,13 +81,14 @@ def _ensure_runtime_imports() -> None: import torch as torch_module from embodichain.lab.sim import SimulationManager as simulation_manager_cls from embodichain.lab.sim.atomic_actions import ( + ActionBinding as action_binding_cls, + ActionInvocation as action_invocation_cls, AtomicActionEngine as atomic_action_engine_cls, - EndEffectorPoseTarget as end_effector_pose_target_cls, - MoveEndEffector as move_end_effector_cls, - MoveEndEffectorCfg as move_end_effector_cfg_cls, - Press as press_cls, - PressCfg as press_cfg_cls, - PressTarget as press_target_cls, + ControlPartCommandProfile as control_part_command_profile_cls, + EndEffectorPoseGoal as end_effector_pose_target_cls, + MotionPolicy as motion_policy_cls, + PressGoal as press_target_cls, + PressOptions as press_options_cls, ) from embodichain.lab.sim.cfg import ( RigidBodyAttributesCfg as rigid_body_attributes_cfg_cls, @@ -123,12 +124,13 @@ def _ensure_runtime_imports() -> None: "torch": torch_module, "SimulationManager": simulation_manager_cls, "AtomicActionEngine": atomic_action_engine_cls, - "EndEffectorPoseTarget": end_effector_pose_target_cls, - "MoveEndEffector": move_end_effector_cls, - "MoveEndEffectorCfg": move_end_effector_cfg_cls, - "Press": press_cls, - "PressCfg": press_cfg_cls, - "PressTarget": press_target_cls, + "ControlPartCommandProfile": control_part_command_profile_cls, + "ActionBinding": action_binding_cls, + "ActionInvocation": action_invocation_cls, + "EndEffectorPoseGoal": end_effector_pose_target_cls, + "MotionPolicy": motion_policy_cls, + "PressGoal": press_target_cls, + "PressOptions": press_options_cls, "RigidBodyAttributesCfg": rigid_body_attributes_cfg_cls, "RigidObjectCfg": rigid_object_cfg_cls, "VisualMaterialCfg": visual_material_cfg_cls, @@ -480,27 +482,13 @@ def _build_atomic_engine( ) -> AtomicActionEngine: """Build a Press benchmark engine with MoveEndEffector pre-positioning.""" hand_close = get_hand_close_qpos(robot, device) - atomic_engine = AtomicActionEngine(motion_generator=motion_gen) - atomic_engine.register( - MoveEndEffector( - motion_gen, - cfg=MoveEndEffectorCfg( - control_part="arm", - sample_interval=MOVE_SAMPLE_INTERVAL, - ), - ) - ) - atomic_engine.register( - Press( - motion_gen, - cfg=PressCfg( - control_part="arm", - hand_control_part="hand", - hand_close_qpos=hand_close, - sample_interval=PRESS_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) + atomic_engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + grasp=hand_close, + ) + }, ) return atomic_engine @@ -579,12 +567,31 @@ def _timed_atomic_run( _sync_cuda() start = time.perf_counter() - is_success, traj, _ = atomic_engine.run( - steps=[ - ("move_end_effector", EndEffectorPoseTarget(xpos=move_target)), - ("press", PressTarget(xpos=press_target)), - ] + binding = ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ) + result = atomic_engine.compile( + ( + ActionInvocation( + "move_end_effector", + EndEffectorPoseGoal(xpos=move_target), + binding, + MotionPolicy(sample_count=MOVE_SAMPLE_INTERVAL), + ), + ActionInvocation( + "press", + PressGoal(xpos=press_target), + binding, + MotionPolicy(sample_count=PRESS_SAMPLE_INTERVAL), + skill_options=PressOptions( + hand_interp_steps=HAND_INTERP_STEPS, + ), + ), + ) ) + is_success = result.plan_success + traj = result.trajectory.positions _sync_cuda() elapsed = time.perf_counter() - start diff --git a/scripts/benchmark/planners/neural_planner/BENCHMARK_DESIGN.md b/scripts/benchmark/planners/neural_planner/BENCHMARK_DESIGN.md index 6646c9b86..01d7843d6 100644 --- a/scripts/benchmark/planners/neural_planner/BENCHMARK_DESIGN.md +++ b/scripts/benchmark/planners/neural_planner/BENCHMARK_DESIGN.md @@ -107,14 +107,14 @@ one Markdown report with exactly three tables: ### 2.3 Atomic Action integration -`ActionCfg.motion_source` defaults to `"ik_interp"`. Existing Atomic Action -benchmarks construct a TOPPRA `MotionGenerator`, but cases that do not +`MotionPolicy.motion_source` defaults to `"ik_interp"`. Existing Atomic Action +benchmarks construct a TOPPRA `MotionGenerator`, but invocations that do not explicitly change `motion_source` still use local IK plus interpolation. NMG and cuRobo Atomic Action evaluation must explicitly set: ```python -cfg.motion_source = "motion_gen" +motion_policy = MotionPolicy(motion_source="motion_gen") ``` The NMG checkpoint must remain confined to `NeuralPlannerCfg`. Atomic Action diff --git a/scripts/tutorials/atomic_action/assemble.py b/scripts/tutorials/atomic_action/assemble.py index 1bbd18625..22a57a7a5 100644 --- a/scripts/tutorials/atomic_action/assemble.py +++ b/scripts/tutorials/atomic_action/assemble.py @@ -20,7 +20,7 @@ (object B). The relative pose of the can with respect to the cube is declared on an :class:`~embodichain.lab.sim.atomic_actions.AssembleAffordance` and consumed by the :class:`~embodichain.lab.sim.atomic_actions.Place` action through an -:class:`~embodichain.lab.sim.atomic_actions.AssembleTarget`. +:class:`~embodichain.lab.sim.atomic_actions.AssembleGoal`. """ from __future__ import annotations @@ -34,38 +34,38 @@ if str(_REPO_ROOT) not in sys.path: sys.path.insert(0, str(_REPO_ROOT)) -import numpy as np import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AssembleAffordance, - AssembleTarget, + AssembleGoal, AtomicActionEngine, - GraspTarget, - PickUp, - PickUpCfg, - Place, - PlaceCfg, -) -from embodichain.lab.sim.cfg import ( - JointDrivePropertiesCfg, - RigidBodyAttributesCfg, - RigidObjectCfg, - RobotCfg, - URDFCfg, + ControlPartCommandProfile, + GraspGoal, + PickUpOptions, + PlaceOptions, + MotionPolicy, ) +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.data import get_data_path from embodichain.lab.sim.objects import RigidObject, Robot from embodichain.lab.sim.shapes import CubeCfg, MeshCfg -from embodichain.lab.sim.solvers import URSolverCfg from embodichain.utils import logger +from scripts.tutorials.atomic_action.scenario_utils import ( + add_dual_ur5_robot, + add_support_surface, + make_dual_ur5_solver_cfg, + settle_object, +) from scripts.tutorials.atomic_action.tutorial_utils import ( broadcast_pose_batch, clone_local_pose_from_first_env, create_antipodal_semantics, create_toppra_motion_generator, + create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, get_hand_open_close_qpos, @@ -75,14 +75,8 @@ serve_tutorial_scene, ) -ARM_URDF_PATH = "UniversalRobots/UR5/UR5.urdf" -GRIPPER_URDF_PATH = "DH_PGI_140_80/DH_PGI_140_80.urdf" OBJECT_MESH_PATH = get_data_path("SodaCan/simple_cola_can.obj") GRIPPER_TCP_Z = 0.155 -ROBOT_INIT_POS = (1.95, 0.0, 0.1) -ROBOT_INIT_ROT = (0.0, 0.0, -90.0) -LEFT_ARM_HOME = (0.0, 0.0, -1.57, -1.57, 1.57, 1.57) -RIGHT_ARM_HOME = (-1.57, -1.57, -1.57, -1.57, 0.0, 0.0) SUPPORT_SURFACE_Z = 0.50 SUPPORT_SURFACE_SIZE = (0.8, 1.2, 0.02) SUPPORT_SURFACE_CENTER = ( @@ -130,153 +124,42 @@ def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the assembly demo.""" - parser = argparse.ArgumentParser(description="Dual-arm object assembly demo") - add_env_launcher_args_to_parser(parser) - parser.set_defaults(device="cpu", renderer="hybrid") - parser.add_argument("--n_sample", type=int, default=10000) - parser.add_argument("--force_reannotate", action="store_true") - parser.add_argument( - "--diagnose_plan", - action="store_true", - help="Plan and print diagnostics without playing the trajectory.", - ) - parser.add_argument( - "--auto_play", - action="store_true", - help="Run the viewer demo without waiting for keyboard input.", - ) - parser.add_argument( - "--headless_play", - action="store_true", - help="Execute planned trajectories without opening the viewer window.", - ) - parser.add_argument( - "--no_vis_eef_axis", - action="store_true", - help="Skip drawing the assembly target axis marker.", + parser = create_tutorial_argument_parser( + "Dual-arm object assembly demo", + features=( + "diagnose_plan", + "grasp_sampling", + "headless_play", + "visualize_axes", + ), + default_device="cpu", + default_renderer="hybrid", ) return parser.parse_args() -def rotation_z(yaw: float) -> np.ndarray: - """Build a 3x3 yaw rotation matrix.""" - cos_yaw = math.cos(yaw) - sin_yaw = math.sin(yaw) - return np.array( - [ - [cos_yaw, -sin_yaw, 0.0], - [sin_yaw, cos_yaw, 0.0], - [0.0, 0.0, 1.0], - ], - dtype=np.float32, - ) - - -def make_transform(xyz: tuple[float, float, float], yaw: float) -> np.ndarray: - """Build a homogeneous transform from translation and yaw.""" - transform = np.eye(4, dtype=np.float32) - transform[:3, :3] = rotation_z(yaw) - transform[:3, 3] = np.asarray(xyz, dtype=np.float32) - return transform - - def create_dual_ur5_robot(sim: SimulationManager) -> Robot: """Create a dual-UR5 robot with one PGI gripper on each arm.""" - arm_urdf_path = ARM_URDF_PATH - gripper_urdf_path = GRIPPER_URDF_PATH - tcp = [ - [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, 0.0], - [0.0, 0.0, 1.0, GRIPPER_TCP_Z], - [0.0, 0.0, 0.0, 1.0], - ] - cfg = RobotCfg( + return add_dual_ur5_robot( + sim, uid="DualUR5Assemble", - urdf_cfg=URDFCfg( - components=[ - { - "component_type": "left_arm", - "urdf_path": arm_urdf_path, - "transform": make_transform((-0.3, -1.45, 0.4), np.pi / 2), - }, - { - "component_type": "right_arm", - "urdf_path": arm_urdf_path, - "transform": make_transform((0.3, -1.45, 0.4), np.pi / 2), - }, - {"component_type": "left_hand", "urdf_path": gripper_urdf_path}, - {"component_type": "right_hand", "urdf_path": gripper_urdf_path}, - ], - fname="dual_ur5_assemble", - name_case={"joint": "upper", "link": "lower"}, - ), - drive_pros=JointDrivePropertiesCfg( - stiffness={ - "LEFT_JOINT[0-9]": 1e4, - "RIGHT_JOINT[0-9]": 1e4, - "LEFT_GRIPPER_FINGER[1-2]_JOINT_1": 1e2, - "RIGHT_GRIPPER_FINGER[1-2]_JOINT_1": 1e2, - }, - damping={ - "LEFT_JOINT[0-9]": 1e3, - "RIGHT_JOINT[0-9]": 1e3, - "LEFT_GRIPPER_FINGER[1-2]_JOINT_1": 1e1, - "RIGHT_GRIPPER_FINGER[1-2]_JOINT_1": 1e1, - }, - max_effort={ - "LEFT_JOINT[0-9]": 1e5, - "RIGHT_JOINT[0-9]": 1e5, - "LEFT_GRIPPER_FINGER[1-2]_JOINT_1": 1e3, - "RIGHT_GRIPPER_FINGER[1-2]_JOINT_1": 1e3, - }, - drive_type="force", + urdf_name="dual_ur5_assemble", + solver_cfg=make_dual_ur5_solver_cfg( + GRIPPER_TCP_Z, + ur_ik_nearest_weight=(1.0, 4.0, 1.0, 1.0, 1.0, 1.0), ), - control_parts={ - "left_arm": ["LEFT_JOINT[0-9]"], - "right_arm": ["RIGHT_JOINT[0-9]"], - "dual_arm": ["LEFT_JOINT[0-9]", "RIGHT_JOINT[0-9]"], - "left_hand": ["LEFT_GRIPPER_FINGER1_JOINT_1"], - "right_hand": ["RIGHT_GRIPPER_FINGER1_JOINT_1"], - }, - solver_cfg={ - "left_arm": URSolverCfg( - ur_type="ur5", - tcp=tcp, - end_link_name="left_ee_link", - root_link_name="left_base_link", - ik_nearest_weight=[1.0, 4.0, 1.0, 1.0, 1.0, 1.0], - ), - "right_arm": URSolverCfg( - ur_type="ur5", - tcp=tcp, - end_link_name="right_ee_link", - root_link_name="right_base_link", - ik_nearest_weight=[1.0, 4.0, 1.0, 1.0, 1.0, 1.0], - ), - }, - init_pos=list(ROBOT_INIT_POS), - init_rot=list(ROBOT_INIT_ROT), - init_qpos=list(LEFT_ARM_HOME) + list(RIGHT_ARM_HOME) + [0.0, 0.0, 0.0, 0.0], + hand_stiffness=1e2, + hand_damping=1e1, + hand_max_effort=1e3, ) - return sim.add_robot(cfg=cfg) def create_support_surface(sim: SimulationManager) -> RigidObject: """Create a compact support slab under the staged objects.""" - return sim.add_rigid_object( - cfg=RigidObjectCfg( - uid="support_surface", - shape=CubeCfg(size=list(SUPPORT_SURFACE_SIZE)), - attrs=RigidBodyAttributesCfg( - mass=10.0, - dynamic_friction=0.9, - static_friction=0.95, - restitution=0.01, - ), - body_type="static", - init_pos=list(SUPPORT_SURFACE_CENTER), - init_rot=[0.0, 0.0, 0.0], - ) + return add_support_surface( + sim, + size=SUPPORT_SURFACE_SIZE, + center=SUPPORT_SURFACE_CENTER, ) @@ -334,16 +217,6 @@ def create_base_object(sim: SimulationManager) -> RigidObject: ) -def settle_object(sim: SimulationManager, obj: RigidObject, step: int = 5) -> None: - """Settle an object before planning.""" - if sim.device.type == "cuda": - sim.init_gpu_physics() - obj.reset() - if step > 0: - sim.update(step=step) - obj.clear_dynamics() - - def compute_can_half_height(can: RigidObject) -> float: """Return half the soda-can extent along world Z when laid on its side.""" vertices = can.get_vertices(env_ids=[0], scale=True)[0].to(torch.float32) @@ -405,43 +278,30 @@ def run_assemble_demo( ) # Step 1 - the left arm picks the soda can up by its top part. - pick_up_action = PickUp( - motion_generator=motion_gen, - cfg=PickUpCfg( - name="pick_up", - control_part="left_arm", - hand_control_part="left_hand", - hand_open_qpos=left_open, - hand_close_qpos=left_close, - pick_object_part="top", - pre_grasp_distance=PICKUP_PRE_GRASP_DISTANCE, - lift_height=PICKUP_LIFT_HEIGHT, - sample_interval=PICKUP_SAMPLE_INTERVAL, - hand_interp_steps=PICKUP_HAND_INTERP_STEPS, - approach_direction=torch.as_tensor( - [0.0, -math.sqrt(0.5), -math.sqrt(0.5)], dtype=torch.float32 - ), - downstream_object_target_poses=(assemble_object_target_pose,), + pick_up_options = PickUpOptions( + pick_object_part="top", + pre_grasp_distance=PICKUP_PRE_GRASP_DISTANCE, + lift_height=PICKUP_LIFT_HEIGHT, + hand_interp_steps=PICKUP_HAND_INTERP_STEPS, + approach_direction=torch.as_tensor( + [0.0, -math.sqrt(0.5), -math.sqrt(0.5)], dtype=torch.float32 ), + downstream_object_target_poses=(assemble_object_target_pose,), ) # Step 2 - the left arm places the can directly above the cube. - place_action = Place( + place_options = PlaceOptions( + lift_height=PLACE_LIFT_HEIGHT, + hand_interp_steps=PLACE_HAND_INTERP_STEPS, + ) + engine = AtomicActionEngine( motion_generator=motion_gen, - cfg=PlaceCfg( - name="place", - control_part="left_arm", - hand_control_part="left_hand", - hand_open_qpos=left_open, - hand_close_qpos=left_close, - lift_height=PLACE_LIFT_HEIGHT, - sample_interval=PLACE_SAMPLE_INTERVAL, - hand_interp_steps=PLACE_HAND_INTERP_STEPS, - ), + control_profiles={ + "left_hand": ControlPartCommandProfile.joint_positions( + open=left_open, + grasp=left_close, + ) + }, ) - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register(pick_up_action) - engine.register(place_action) - wait_for_user = prepare_tutorial_scene( sim, args, "Inspect the scene, then press Enter to plan PickUp -> Place..." ) @@ -456,12 +316,30 @@ def run_assemble_demo( assemble_object_entity=can, assemble_to_base_pose=assemble_to_base, ) - success, traj, _ = engine.run( - [ - ("pick_up", GraspTarget(can_semantics)), - ("place", AssembleTarget(affordance=assemble_affordance)), - ] + binding = ActionBinding( + manipulators={"primary": "left_arm"}, + end_effectors={"primary": "left_hand"}, + ) + compiled = engine.compile( + ( + ActionInvocation( + "pick_up", + GraspGoal(can_semantics), + binding, + MotionPolicy(sample_count=PICKUP_SAMPLE_INTERVAL), + skill_options=pick_up_options, + ), + ActionInvocation( + "place", + AssembleGoal(affordance=assemble_affordance), + binding, + MotionPolicy(sample_count=PLACE_SAMPLE_INTERVAL), + skill_options=place_options, + ), + ) ) + success = compiled.plan_success + traj = compiled.trajectory.positions if not success.all(): logger.log_warning("Failed to plan the assemble demo trajectory.") diff --git a/scripts/tutorials/atomic_action/coordinated_pickment.py b/scripts/tutorials/atomic_action/coordinated_pickment.py index a9f1a06a1..7d73c7d31 100644 --- a/scripts/tutorials/atomic_action/coordinated_pickment.py +++ b/scripts/tutorials/atomic_action/coordinated_pickment.py @@ -23,8 +23,6 @@ from __future__ import annotations import argparse -import math -import os import sys import time from dataclasses import dataclass @@ -34,50 +32,58 @@ if str(_REPO_ROOT) not in sys.path: sys.path.insert(0, str(_REPO_ROOT)) -import numpy as np import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.atomic_actions import ( - Affordance, + ActionBinding, + ActionInvocation, AtomicActionEngine, - CoordinatedPickTarget, + ControlPartCommandProfile, + CoordinatedPickGoal, CoordinatedPickment, - CoordinatedPickmentCfg, - ObjectSemantics, + CoordinatedPickmentOptions, + MotionPolicy, ) from embodichain.lab.sim.cfg import ( - JointDrivePropertiesCfg, RigidBodyAttributesCfg, RigidObjectCfg, - RobotCfg, - URDFCfg, ) from embodichain.lab.sim.objects import RigidObject, Robot -from embodichain.lab.sim.shapes import CubeCfg, MeshCfg -from embodichain.lab.sim.solvers import PytorchSolverCfg +from embodichain.lab.sim.shapes import MeshCfg from embodichain.utils import logger from embodichain.utils.math import matrix_from_euler +from scripts.tutorials.atomic_action.scenario_utils import ( + add_dual_ur5_robot, + add_support_surface, + compute_local_bounds, + compute_world_bounds, + create_manual_object_semantics, + get_local_vertices, + invert_pose, + log_action_plan, + make_dual_ur5_solver_cfg, + normalize_vector, + resolve_cached_data_path, + rotate_pose_about_world_z, + settle_object, +) from scripts.tutorials.atomic_action.tutorial_utils import ( broadcast_pose_batch, clone_local_pose_from_first_env, create_toppra_motion_generator, + create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, + format_tensor, + get_hand_open_close_qpos, prepare_tutorial_scene, replay_trajectory, run_tutorial, ) -ARM_URDF_PATH = "UniversalRobots/UR5/UR5.urdf" -GRIPPER_URDF_PATH = "DH_PGI_140_80/DH_PGI_140_80.urdf" PICKMENT_ASSET_ROOT = "CoordinatedPlacementAndPickment" GRIPPER_TCP_Z = 0.121 -ROBOT_INIT_POS = (1.95, 0.0, 0.1) -ROBOT_INIT_ROT = (0.0, 0.0, -90.0) -LEFT_ARM_HOME = (0.0, 0.0, -1.57, -1.57, 1.57, 1.57) -RIGHT_ARM_HOME = (-1.57, -1.57, -1.57, -1.57, 0.0, 0.0) SUPPORT_SURFACE_Z = 0.65 SUPPORT_SURFACE_SIZE = (0.60, 0.60, 0.02) SUPPORT_SURFACE_CENTER = ( @@ -151,28 +157,16 @@ class PickmentObjectPreset: def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the demo.""" - parser = argparse.ArgumentParser(description="Dual-arm coordinated pickment demo") - add_env_launcher_args_to_parser(parser) - parser.set_defaults(device="cpu", renderer="hybrid") - parser.add_argument( - "--diagnose_plan", - action="store_true", - help="Plan and print diagnostics without playing the trajectory.", - ) - parser.add_argument( - "--debug_state", - action="store_true", - help="Log hand targets and object poses during execution.", - ) - parser.add_argument( - "--auto_play", - action="store_true", - help="Run the viewer demo without waiting for keyboard input.", - ) - parser.add_argument( - "--headless_play", - action="store_true", - help="Execute planned trajectories without opening the viewer window.", + parser = create_tutorial_argument_parser( + "Dual-arm coordinated pickment demo", + features=( + "debug_state", + "diagnose_plan", + "headless_play", + "visualize_axes", + ), + default_device="cpu", + default_renderer="hybrid", ) parser.add_argument( "--object", @@ -180,154 +174,27 @@ def parse_arguments() -> argparse.Namespace: default="pencil", help="Object mesh to grasp in the coordinated pickment demo.", ) - parser.add_argument( - "--no_vis_eef_axis", - action="store_true", - help="Do not draw the pickment target/grasp coordinate frames before planning.", - ) return parser.parse_args() -def get_cached_data_path(data_path: str) -> str: - """Resolve an asset path from the local cache before importing data helpers.""" - if os.path.isabs(data_path): - return data_path - - data_root = Path( - os.environ.get( - "EMBODICHAIN_DATA_ROOT", - str(Path.home() / ".cache" / "embodichain_data"), - ) - ) - candidates = ( - data_root / data_path, - data_root / "extract" / data_path, - ) - for candidate in candidates: - if candidate.exists(): - return str(candidate) - - from embodichain.data import get_data_path - - return get_data_path(data_path) - - -def rotation_z(yaw: float) -> np.ndarray: - """Build a 3x3 yaw rotation matrix.""" - cos_yaw = math.cos(yaw) - sin_yaw = math.sin(yaw) - return np.array( - [ - [cos_yaw, -sin_yaw, 0.0], - [sin_yaw, cos_yaw, 0.0], - [0.0, 0.0, 1.0], - ], - dtype=np.float32, - ) - - -def make_transform(xyz: tuple[float, float, float], yaw: float) -> np.ndarray: - """Build a homogeneous transform from translation and yaw.""" - transform = np.eye(4, dtype=np.float32) - transform[:3, :3] = rotation_z(yaw) - transform[:3, 3] = np.asarray(xyz, dtype=np.float32) - return transform - - def create_dual_ur5_robot(sim: SimulationManager) -> Robot: """Create a dual-UR5 robot with one PGI gripper on each arm.""" - arm_urdf_path = get_cached_data_path(ARM_URDF_PATH) - gripper_urdf_path = get_cached_data_path(GRIPPER_URDF_PATH) - tcp = [ - [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, 0.0], - [0.0, 0.0, 1.0, GRIPPER_TCP_Z], - [0.0, 0.0, 0.0, 1.0], - ] - cfg = RobotCfg( + return add_dual_ur5_robot( + sim, uid="DualUR5CoordinatedPickment", - urdf_cfg=URDFCfg( - components=[ - { - "component_type": "left_arm", - "urdf_path": arm_urdf_path, - "transform": make_transform((-0.3, -1.45, 0.4), np.pi / 2), - }, - { - "component_type": "right_arm", - "urdf_path": arm_urdf_path, - "transform": make_transform((0.3, -1.45, 0.4), np.pi / 2), - }, - {"component_type": "left_hand", "urdf_path": gripper_urdf_path}, - {"component_type": "right_hand", "urdf_path": gripper_urdf_path}, - ], - fname="dual_ur5_coordinated_pickment", - name_case={"joint": "upper", "link": "lower"}, - ), - drive_pros=JointDrivePropertiesCfg( - stiffness={ - "LEFT_JOINT[0-9]": 1e4, - "RIGHT_JOINT[0-9]": 1e4, - "LEFT_GRIPPER_FINGER[1-2]_JOINT_1": 1e3, - "RIGHT_GRIPPER_FINGER[1-2]_JOINT_1": 1e3, - }, - damping={ - "LEFT_JOINT[0-9]": 1e3, - "RIGHT_JOINT[0-9]": 1e3, - "LEFT_GRIPPER_FINGER[1-2]_JOINT_1": 1e2, - "RIGHT_GRIPPER_FINGER[1-2]_JOINT_1": 1e2, - }, - max_effort={ - "LEFT_JOINT[0-9]": 1e5, - "RIGHT_JOINT[0-9]": 1e5, - "LEFT_GRIPPER_FINGER[1-2]_JOINT_1": 1e4, - "RIGHT_GRIPPER_FINGER[1-2]_JOINT_1": 1e4, - }, - drive_type="force", - ), - control_parts={ - "left_arm": ["LEFT_JOINT[0-9]"], - "right_arm": ["RIGHT_JOINT[0-9]"], - "dual_arm": ["LEFT_JOINT[0-9]", "RIGHT_JOINT[0-9]"], - "left_hand": ["LEFT_GRIPPER_FINGER1_JOINT_1"], - "right_hand": ["RIGHT_GRIPPER_FINGER1_JOINT_1"], - }, - solver_cfg={ - "left_arm": PytorchSolverCfg( - end_link_name="left_ee_link", - root_link_name="left_base_link", - tcp=tcp, - num_samples=30, - ), - "right_arm": PytorchSolverCfg( - end_link_name="right_ee_link", - root_link_name="right_base_link", - tcp=tcp, - num_samples=30, - ), - }, - init_pos=list(ROBOT_INIT_POS), - init_rot=list(ROBOT_INIT_ROT), - init_qpos=list(LEFT_ARM_HOME) + list(RIGHT_ARM_HOME) + [0.0, 0.0, 0.0, 0.0], + urdf_name="dual_ur5_coordinated_pickment", + arm_urdf_path=resolve_cached_data_path("UniversalRobots/UR5/UR5.urdf"), + gripper_urdf_path=resolve_cached_data_path("DH_PGI_140_80/DH_PGI_140_80.urdf"), + solver_cfg=make_dual_ur5_solver_cfg(GRIPPER_TCP_Z, solver="pytorch"), ) - return sim.add_robot(cfg=cfg) def create_support_surface(sim: SimulationManager) -> RigidObject: """Create a compact support slab under the staged object.""" - return sim.add_rigid_object( - cfg=RigidObjectCfg( - uid="support_surface", - shape=CubeCfg(size=list(SUPPORT_SURFACE_SIZE)), - attrs=RigidBodyAttributesCfg( - mass=10.0, - dynamic_friction=0.9, - static_friction=0.95, - restitution=0.01, - ), - body_type="static", - init_pos=list(SUPPORT_SURFACE_CENTER), - ) + return add_support_surface( + sim, + size=SUPPORT_SURFACE_SIZE, + center=SUPPORT_SURFACE_CENTER, ) @@ -340,7 +207,7 @@ def create_pickment_object( cfg=RigidObjectCfg( uid=preset.label, shape=MeshCfg( - fpath=get_cached_data_path(preset.mesh_path), compute_uv=False + fpath=resolve_cached_data_path(preset.mesh_path), compute_uv=False ), attrs=RigidBodyAttributesCfg( mass=0.01, @@ -366,55 +233,6 @@ def create_pickment_object( return obj -def settle_object(sim: SimulationManager, obj: RigidObject, step: int = 5) -> None: - """Settle an object before planning.""" - if sim.device.type == "cuda": - sim.init_gpu_physics() - obj.reset() - if step > 0: - sim.update(step=step) - obj.clear_dynamics() - - -def create_object_semantics(obj: RigidObject, label: str) -> ObjectSemantics: - """Create minimal object semantics for manually specified grasps.""" - return ObjectSemantics( - label=label, - geometry={}, - affordance=Affordance(object_label=label), - entity=obj, - ) - - -def get_hand_open_close_qpos( - robot: Robot, - hand_control_part: str, - device: torch.device, - close_qpos: float, -) -> tuple[torch.Tensor, torch.Tensor]: - """Get open and close qpos for a PGI gripper control part.""" - limits = robot.get_qpos_limits(name=hand_control_part)[0].to( - device=device, dtype=torch.float32 - ) - hand_open = limits[:, 0] - hand_close = torch.clamp( - torch.full_like(limits[:, 1], close_qpos), - min=limits[:, 0], - max=limits[:, 1], - ) - return hand_open, hand_close - - -def get_local_vertices(obj: RigidObject) -> torch.Tensor: - """Get scaled local mesh vertices.""" - return obj.get_vertices(env_ids=[0], scale=True)[0] - - -def compute_local_bounds(vertices: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Compute local mesh AABB from scaled vertices.""" - return vertices.min(dim=0).values, vertices.max(dim=0).values - - def compute_supported_init_pos( obj: RigidObject, preset: PickmentObjectPreset, @@ -430,50 +248,6 @@ def compute_supported_init_pos( return (preset.init_xy[0], preset.init_xy[1], z) -def invert_pose(pose: torch.Tensor) -> torch.Tensor: - """Invert batched homogeneous transforms.""" - inv_pose = pose.clone() - rot_t = pose[:, :3, :3].transpose(1, 2) - inv_pose[:, :3, :3] = rot_t - inv_pose[:, :3, 3] = -torch.bmm(rot_t, pose[:, :3, 3:4]).squeeze(-1) - return inv_pose - - -def transform_points(pose: torch.Tensor, points: torch.Tensor) -> torch.Tensor: - """Transform local points by a homogeneous pose.""" - return points @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] - - -def compute_world_bounds( - object_pose: torch.Tensor, - local_vertices: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Compute world AABB from transformed local mesh vertices.""" - world_vertices = transform_points(object_pose, local_vertices) - return world_vertices.min(dim=0).values, world_vertices.max(dim=0).values - - -def normalize_vector(vector: torch.Tensor, fallback: torch.Tensor) -> torch.Tensor: - """Normalize a vector with a deterministic fallback for degenerate cases.""" - norm = torch.linalg.norm(vector) - if norm < 1e-6: - return fallback.to(device=vector.device, dtype=vector.dtype) - return vector / norm - - -def rotate_pose_about_world_z(pose: torch.Tensor, yaw_deg: float) -> torch.Tensor: - """Rotate pose orientation about world Z while preserving translation.""" - yaw = math.radians(yaw_deg) - rot = torch.eye(3, dtype=pose.dtype, device=pose.device) - rot[0, 0] = math.cos(yaw) - rot[0, 1] = -math.sin(yaw) - rot[1, 0] = math.sin(yaw) - rot[1, 1] = math.cos(yaw) - rotated_pose = pose.clone() - rotated_pose[:3, :3] = rot @ pose[:3, :3] - return rotated_pose - - def build_object_grasp_poses( object_pose: torch.Tensor, local_vertices: torch.Tensor, @@ -552,28 +326,6 @@ def build_object_target_pose( return pose -def format_tensor(tensor: torch.Tensor) -> str: - """Format tensor values for compact logging.""" - rounded = (tensor.detach().cpu() * 10000.0).round() / 10000.0 - return str(rounded.tolist()) - - -def log_action_plan( - robot: Robot, - action_name: str, - traj: torch.Tensor, - joint_ids: list[int], - segments: dict[str, int] | None = None, -) -> None: - """Log common action plan details.""" - joint_names = [robot.joint_names[joint_id] for joint_id in joint_ids] - logger.log_info(f"{action_name} joint ids: {joint_ids}") - logger.log_info(f"{action_name} joint names: {joint_names}") - logger.log_info(f"{action_name} trajectory shape: {tuple(traj.shape)}") - if segments is not None: - logger.log_info(f"{action_name} trajectory segments: {segments}") - - def log_scene_targets( object_label: str, object_pose: torch.Tensor, @@ -656,37 +408,42 @@ def run_coordinated_pickment_demo( object_pose = object_pose_batch[0].to(device=sim.device, dtype=torch.float32) n_envs = object_pose_batch.shape[0] object_vertices = get_local_vertices(obj) - object_semantics = create_object_semantics(obj, preset.label) + object_semantics = create_manual_object_semantics(obj, preset.label) motion_gen = create_toppra_motion_generator(robot) left_open, left_close = get_hand_open_close_qpos( - robot, "left_hand", sim.device, preset.hand_close_qpos + robot, + hand_control_part="left_hand", + close_qpos=preset.hand_close_qpos, ) right_open, right_close = get_hand_open_close_qpos( - robot, "right_hand", sim.device, preset.hand_close_qpos + robot, + hand_control_part="right_hand", + close_qpos=preset.hand_close_qpos, + ) + pickment_options = CoordinatedPickmentOptions( + pre_grasp_distance=PICKMENT_PRE_GRASP_DISTANCE, + lift_height=PICKMENT_LIFT_HEIGHT, + hand_interp_steps=PICKMENT_HAND_INTERP_STEPS, + hold_steps=PICKMENT_HOLD_STEPS, + object_motion_keyframes=PICKMENT_OBJECT_MOTION_KEYFRAMES, ) - pickment_action = CoordinatedPickment( + engine = AtomicActionEngine( motion_generator=motion_gen, - cfg=CoordinatedPickmentCfg( - control_part="dual_arm", - left_arm_control_part="left_arm", - right_arm_control_part="right_arm", - left_hand_control_part="left_hand", - right_hand_control_part="right_hand", - left_hand_open_qpos=left_open, - left_hand_close_qpos=left_close, - right_hand_open_qpos=right_open, - right_hand_close_qpos=right_close, - pre_grasp_distance=PICKMENT_PRE_GRASP_DISTANCE, - lift_height=PICKMENT_LIFT_HEIGHT, - sample_interval=PICKMENT_SAMPLE_INTERVAL, - hand_interp_steps=PICKMENT_HAND_INTERP_STEPS, - hold_steps=PICKMENT_HOLD_STEPS, - object_motion_keyframes=PICKMENT_OBJECT_MOTION_KEYFRAMES, - ), + control_profiles={ + "left_hand": ControlPartCommandProfile.joint_positions( + open=left_open, + grasp=left_close, + ), + "right_hand": ControlPartCommandProfile.joint_positions( + open=right_open, + grasp=right_close, + ), + }, ) - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register(pickment_action) + pickment_action = engine.actions["coordinated_pickment"] + if not isinstance(pickment_action, CoordinatedPickment): + raise RuntimeError("Unexpected coordinated_pickment implementation.") left_grasp_pose, right_grasp_pose = build_object_grasp_poses( object_pose, @@ -724,7 +481,7 @@ 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 = CoordinatedPickTarget( + pickment_target = CoordinatedPickGoal( semantics=object_semantics, object_target_pose=broadcast_pose_batch(target_pose, num_envs=n_envs), left_object_to_eef=left_object_to_eef, @@ -737,7 +494,22 @@ def run_coordinated_pickment_demo( ) start_time = time.time() - success, traj, _ = engine.run([("coordinated_pickment", pickment_target)]) + compiled = engine.compile( + ( + ActionInvocation( + "coordinated_pickment", + pickment_target, + ActionBinding( + manipulators={"left": "left_arm", "right": "right_arm"}, + end_effectors={"left": "left_hand", "right": "right_hand"}, + ), + MotionPolicy(sample_count=PICKMENT_SAMPLE_INTERVAL), + skill_options=pickment_options, + ), + ) + ) + success = compiled.plan_success + traj = compiled.trajectory.positions logger.log_info( f"Plan coordinated pickment cost time: {time.time() - start_time:.2f} seconds" ) @@ -750,7 +522,10 @@ def run_coordinated_pickment_demo( "coordinated_pickment", traj, joint_ids, - pickment_action.get_segment_lengths(), + pickment_action.get_segment_lengths( + PICKMENT_SAMPLE_INTERVAL, + pickment_options, + ), ) if args.diagnose_plan: diff --git a/scripts/tutorials/atomic_action/coordinated_placement.py b/scripts/tutorials/atomic_action/coordinated_placement.py index 27f46d47a..32a4c73e1 100644 --- a/scripts/tutorials/atomic_action/coordinated_placement.py +++ b/scripts/tutorials/atomic_action/coordinated_placement.py @@ -24,8 +24,6 @@ from __future__ import annotations import argparse -import math -import os import sys import time from pathlib import Path @@ -38,39 +36,53 @@ import torch from scipy.spatial.transform import Rotation as SciRotation -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.atomic_actions import ( - Affordance, + ActionBinding, + ActionInvocation, AtomicActionEngine, + ControlPartCommandProfile, CoordinatedPlacement, - CoordinatedPlacementCfg, - CoordinatedPlacementTarget, - GraspTarget, + CoordinatedPlacementOptions, + CoordinatedPlacementGoal, + GraspGoal, HeldObjectState, ObjectSemantics, - PickUp, - PickUpCfg, - WorldState, + PickUpOptions, + MotionPolicy, + TaskState, ) from embodichain.lab.sim.cfg import ( - JointDrivePropertiesCfg, RigidBodyAttributesCfg, RigidObjectCfg, - RobotCfg, - URDFCfg, ) from embodichain.lab.sim.objects import RigidObject, Robot from embodichain.lab.sim.shapes import MeshCfg -from embodichain.lab.sim.solvers import URSolverCfg from embodichain.utils import logger +from scripts.tutorials.atomic_action.scenario_utils import ( + add_dual_ur5_robot, + compute_local_bounds, + compute_world_bounds, + create_manual_object_semantics, + get_local_vertices, + invert_pose, + log_action_plan, + make_dual_ur5_solver_cfg, + normalize_vector, + resolve_cached_data_path, + rotate_pose_about_world_z, + settle_object, + transform_points, +) from scripts.tutorials.atomic_action.tutorial_utils import ( broadcast_pose_batch, clone_local_pose_from_first_env, create_toppra_motion_generator, + create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, - make_ur5_solver_cfg, + format_tensor, + get_hand_open_close_qpos, prepare_tutorial_scene, replay_trajectory, run_tutorial, @@ -104,8 +116,6 @@ def transform_baseline_pose( return tuple(float(value) for value in pos), tuple(float(value) for value in rot) -ARM_URDF_PATH = "UniversalRobots/UR5/UR5.urdf" -GRIPPER_URDF_PATH = "DH_PGI_140_80/DH_PGI_140_80.urdf" PLACEMENT_ASSET_ROOT = "CoordinatedPlacementAndPickment" TABLE_MESH_PATH = f"{PLACEMENT_ASSET_ROOT}/table.glb" BREAD_MESH_PATH = f"{PLACEMENT_ASSET_ROOT}/bread.glb" @@ -116,9 +126,6 @@ def transform_baseline_pose( PICK_SAMPLE_INTERVAL = 100 COORDINATED_SAMPLE_INTERVAL = 120 ROBOT_INIT_POS = (1.85, 0.0, 0.1) -ROBOT_INIT_ROT = (0.0, 0.0, -90.0) -LEFT_ARM_HOME = (0.0, 0.0, -1.57, -1.57, 1.57, 1.57) -RIGHT_ARM_HOME = (-1.57, -1.57, -1.57, -1.57, 0.0, 0.0) TABLE_TOP_Z = 0.65 BASELINE_TABLE_TOP_Z = 0.3621708124799265 SCENE_Z_OFFSET = TABLE_TOP_Z - BASELINE_TABLE_TOP_Z @@ -192,151 +199,36 @@ def transform_baseline_pose( def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the demo.""" - parser = argparse.ArgumentParser(description="Dual-arm coordinated placement demo") - add_env_launcher_args_to_parser(parser) - parser.set_defaults(device="cuda", renderer="hybrid") - parser.add_argument( - "--diagnose_plan", - action="store_true", - help="Plan and print diagnostics without playing the trajectory.", - ) - parser.add_argument( - "--debug_state", - action="store_true", - help="Log hand targets and object poses during execution.", - ) - parser.add_argument( - "--auto_play", - action="store_true", - help="Run the viewer demo without waiting for keyboard input.", - ) - parser.add_argument( - "--no_vis_eef_axis", - action="store_true", - help="Do not draw coordinated placement target coordinate frames.", - ) - parser.add_argument( - "--headless_play", - action="store_true", - help="Execute planned trajectories without opening the viewer window.", + parser = create_tutorial_argument_parser( + "Dual-arm coordinated placement demo", + features=( + "debug_state", + "diagnose_plan", + "headless_play", + "visualize_axes", + ), + default_device="cuda", + default_renderer="hybrid", ) return parser.parse_args() -def get_cached_data_path(data_path: str) -> str: - """Resolve an asset path from the local cache before importing data helpers.""" - if os.path.isabs(data_path): - return data_path - - data_root = Path( - os.environ.get( - "EMBODICHAIN_DATA_ROOT", - str(Path.home() / ".cache" / "embodichain_data"), - ) - ) - candidates = ( - data_root / data_path, - data_root / "extract" / data_path, - ) - for candidate in candidates: - if candidate.exists(): - return str(candidate) - - from embodichain.data import get_data_path - - return get_data_path(data_path) - - -def rotation_z(yaw: float) -> np.ndarray: - """Build a 3x3 yaw rotation matrix.""" - cos_yaw = math.cos(yaw) - sin_yaw = math.sin(yaw) - return np.array( - [ - [cos_yaw, -sin_yaw, 0.0], - [sin_yaw, cos_yaw, 0.0], - [0.0, 0.0, 1.0], - ], - dtype=np.float32, - ) - - -def make_transform(xyz: tuple[float, float, float], yaw: float) -> np.ndarray: - """Build a homogeneous transform from translation and yaw.""" - transform = np.eye(4, dtype=np.float32) - transform[:3, :3] = rotation_z(yaw) - transform[:3, 3] = np.asarray(xyz, dtype=np.float32) - return transform - - -def make_prefixed_ur5_solver_cfg(prefix: str) -> URSolverCfg: - """Create a UR5 solver cfg for a prefixed arm in the assembled robot.""" - cfg = make_ur5_solver_cfg(GRIPPER_TCP_Z) - cfg.root_link_name = f"{prefix}_base_link" - cfg.end_link_name = f"{prefix}_ee_link" - return cfg - - def create_dual_ur5_robot(sim: SimulationManager) -> Robot: """Create a dual-UR5 robot with one PGI gripper on each arm.""" - arm_urdf_path = get_cached_data_path(ARM_URDF_PATH) - gripper_urdf_path = get_cached_data_path(GRIPPER_URDF_PATH) - cfg = RobotCfg( + return add_dual_ur5_robot( + sim, uid="DualUR5CoordinatedPlacement", - urdf_cfg=URDFCfg( - components=[ - { - "component_type": "left_arm", - "urdf_path": arm_urdf_path, - "transform": make_transform((-0.3, -1.45, 0.4), np.pi / 2), - }, - { - "component_type": "right_arm", - "urdf_path": arm_urdf_path, - "transform": make_transform((0.3, -1.45, 0.4), np.pi / 2), - }, - {"component_type": "left_hand", "urdf_path": gripper_urdf_path}, - {"component_type": "right_hand", "urdf_path": gripper_urdf_path}, - ], - fname="dual_ur5_coordinated_placement", + urdf_name="dual_ur5_coordinated_placement", + arm_urdf_path=resolve_cached_data_path("UniversalRobots/UR5/UR5.urdf"), + gripper_urdf_path=resolve_cached_data_path("DH_PGI_140_80/DH_PGI_140_80.urdf"), + solver_cfg=make_dual_ur5_solver_cfg( + GRIPPER_TCP_Z, + clear_urdf_path=True, ), - drive_pros=JointDrivePropertiesCfg( - stiffness={ - "left_joint[0-9]": 1e4, - "right_joint[0-9]": 1e4, - "left_gripper_finger[1-2]_joint_1": 1e3, - "right_gripper_finger[1-2]_joint_1": 1e3, - }, - damping={ - "left_joint[0-9]": 1e3, - "right_joint[0-9]": 1e3, - "left_gripper_finger[1-2]_joint_1": 1e2, - "right_gripper_finger[1-2]_joint_1": 1e2, - }, - max_effort={ - "left_joint[0-9]": 1e5, - "right_joint[0-9]": 1e5, - "left_gripper_finger[1-2]_joint_1": 1e4, - "right_gripper_finger[1-2]_joint_1": 1e4, - }, - drive_type="force", - ), - control_parts={ - "left_arm": ["left_joint[0-9]"], - "right_arm": ["right_joint[0-9]"], - "dual_arm": ["left_joint[0-9]", "right_joint[0-9]"], - "left_hand": ["left_gripper_finger1_joint_1"], - "right_hand": ["right_gripper_finger1_joint_1"], - }, - solver_cfg={ - "left_arm": make_prefixed_ur5_solver_cfg("left"), - "right_arm": make_prefixed_ur5_solver_cfg("right"), - }, - init_pos=list(ROBOT_INIT_POS), - init_rot=list(ROBOT_INIT_ROT), - init_qpos=list(LEFT_ARM_HOME) + list(RIGHT_ARM_HOME) + [0.0, 0.0, 0.0, 0.0], + init_pos=ROBOT_INIT_POS, + joint_name_case="lower", + set_urdf_name_case=False, ) - return sim.add_robot(cfg=cfg) def create_table(sim: SimulationManager) -> RigidObject: @@ -344,7 +236,7 @@ def create_table(sim: SimulationManager) -> RigidObject: return sim.add_rigid_object( cfg=RigidObjectCfg( uid="table", - shape=MeshCfg(fpath=get_cached_data_path(TABLE_MESH_PATH)), + shape=MeshCfg(fpath=resolve_cached_data_path(TABLE_MESH_PATH)), attrs=RigidBodyAttributesCfg( mass=10.0, dynamic_friction=0.9, @@ -364,7 +256,7 @@ def create_bread(sim: SimulationManager) -> RigidObject: cfg=RigidObjectCfg( uid="bread", shape=MeshCfg( - fpath=get_cached_data_path(BREAD_MESH_PATH), compute_uv=False + fpath=resolve_cached_data_path(BREAD_MESH_PATH), compute_uv=False ), attrs=RigidBodyAttributesCfg( mass=0.01, @@ -388,7 +280,9 @@ def create_pan(sim: SimulationManager) -> RigidObject: return sim.add_rigid_object( cfg=RigidObjectCfg( uid="pan", - shape=MeshCfg(fpath=get_cached_data_path(PAN_MESH_PATH), compute_uv=False), + shape=MeshCfg( + fpath=resolve_cached_data_path(PAN_MESH_PATH), compute_uv=False + ), attrs=RigidBodyAttributesCfg( mass=0.01, dynamic_friction=0.97, @@ -410,100 +304,6 @@ def create_pan(sim: SimulationManager) -> RigidObject: ) -def settle_object(sim: SimulationManager, obj: RigidObject, step: int = 5) -> None: - """Settle an object before planning.""" - if sim.device.type == "cuda": - sim.init_gpu_physics() - obj.reset() - if step > 0: - sim.update(step=step) - obj.clear_dynamics() - - -def create_object_semantics(obj: RigidObject, label: str) -> ObjectSemantics: - """Create minimal object semantics for manually specified grasps.""" - return ObjectSemantics( - label=label, - geometry={}, - affordance=Affordance(object_label=label), - entity=obj, - ) - - -def get_hand_open_close_qpos( - robot: Robot, hand_control_part: str, device: torch.device -) -> tuple[torch.Tensor, torch.Tensor]: - """Get open and close qpos for a PGI gripper control part.""" - limits = robot.get_qpos_limits(name=hand_control_part)[0].to( - device=device, dtype=torch.float32 - ) - hand_open = limits[:, 0] - hand_close = torch.minimum(limits[:, 1], torch.full_like(limits[:, 1], 0.030)) - return hand_open, hand_close - - -def get_pan_handle_open_close_qpos( - robot: Robot, hand_control_part: str, device: torch.device -) -> tuple[torch.Tensor, torch.Tensor]: - """Get right hand qpos tuned for holding the thin pan handle.""" - limits = robot.get_qpos_limits(name=hand_control_part)[0].to( - device=device, dtype=torch.float32 - ) - hand_open = limits[:, 0] - hand_close = torch.clamp( - torch.full_like(limits[:, 1], PAN_HANDLE_CLOSE_QPOS), - min=limits[:, 0], - max=limits[:, 1], - ) - return hand_open, hand_close - - -def invert_pose(pose: torch.Tensor) -> torch.Tensor: - """Invert batched homogeneous transforms.""" - inv_pose = pose.clone() - rot_t = pose[:, :3, :3].transpose(1, 2) - inv_pose[:, :3, :3] = rot_t - inv_pose[:, :3, 3] = -torch.bmm(rot_t, pose[:, :3, 3:4]).squeeze(-1) - return inv_pose - - -def get_local_vertices(obj: RigidObject) -> torch.Tensor: - """Get scaled local mesh vertices.""" - return obj.get_vertices(env_ids=[0], scale=True)[0] - - -def compute_local_bounds(vertices: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Compute local mesh AABB from scaled vertices.""" - return vertices.min(dim=0).values, vertices.max(dim=0).values - - -def transform_points(pose: torch.Tensor, points: torch.Tensor) -> torch.Tensor: - """Transform local points by a homogeneous pose.""" - return points @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] - - -def compute_world_bounds( - object_pose: torch.Tensor, - local_vertices: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Compute world AABB from transformed local mesh vertices.""" - world_vertices = transform_points(object_pose, local_vertices) - return world_vertices.min(dim=0).values, world_vertices.max(dim=0).values - - -def rotate_pose_about_world_z(pose: torch.Tensor, yaw_deg: float) -> torch.Tensor: - """Rotate pose orientation about world Z while preserving translation.""" - yaw = math.radians(yaw_deg) - rot = torch.eye(3, dtype=pose.dtype, device=pose.device) - rot[0, 0] = math.cos(yaw) - rot[0, 1] = -math.sin(yaw) - rot[1, 0] = math.sin(yaw) - rot[1, 1] = math.cos(yaw) - rotated_pose = pose.clone() - rotated_pose[:3, :3] = rot @ pose[:3, :3] - return rotated_pose - - def get_pan_basin_vertices(pan_vertices: torch.Tensor) -> torch.Tensor: """Select pan vertices that belong to the basin instead of the long handle.""" basin_vertices = pan_vertices[pan_vertices[:, 2] <= PAN_BASIN_LOCAL_Z_MAX] @@ -552,14 +352,6 @@ def build_flat_object_grasp_pose( return build_top_down_tcp_pose(grasp_position, device) -def normalize_vector(vector: torch.Tensor, fallback: torch.Tensor) -> torch.Tensor: - """Normalize a vector with a deterministic fallback for degenerate cases.""" - norm = torch.linalg.norm(vector) - if norm < 1e-6: - return fallback.to(device=vector.device, dtype=vector.dtype) - return vector / norm - - def build_pan_handle_grasp_pose( pan_pose: torch.Tensor, pan_vertices: torch.Tensor, @@ -652,12 +444,6 @@ def build_placing_object_target_pose( return pose -def format_tensor(tensor: torch.Tensor) -> str: - """Format tensor values for compact logging.""" - rounded = (tensor.detach().cpu() * 10000.0).round() / 10000.0 - return str(rounded.tolist()) - - def compute_actual_held_state( robot: Robot, semantics: ObjectSemantics, @@ -727,22 +513,6 @@ def draw_coordinated_axes( ) -def log_action_plan( - robot: Robot, - action_name: str, - traj: torch.Tensor, - joint_ids: list[int], - segments: dict[str, int] | None = None, -) -> None: - """Log common action plan details.""" - joint_names = [robot.joint_names[joint_id] for joint_id in joint_ids] - logger.log_info(f"{action_name} joint ids: {joint_ids}") - logger.log_info(f"{action_name} joint names: {joint_names}") - logger.log_info(f"{action_name} trajectory shape: {tuple(traj.shape)}") - if segments is not None: - logger.log_info(f"{action_name} trajectory segments: {segments}") - - def log_execution_state( robot: Robot, bread: RigidObject, @@ -784,68 +554,60 @@ def run_coordinated_placement_demo( pan_vertices = get_local_vertices(pan) bread_local_min, bread_local_max = compute_local_bounds(bread_vertices) log_scene_targets(bread_pose, pan_pose) - bread_semantics = create_object_semantics(bread, BREAD_LABEL) - pan_semantics = create_object_semantics(pan, PAN_LABEL) + bread_semantics = create_manual_object_semantics(bread, BREAD_LABEL) + pan_semantics = create_manual_object_semantics(pan, PAN_LABEL) motion_gen = create_toppra_motion_generator(robot) - right_open, right_close = get_pan_handle_open_close_qpos( - robot, "right_hand", sim.device + right_open, right_close = get_hand_open_close_qpos( + robot, + hand_control_part="right_hand", + close_qpos=PAN_HANDLE_CLOSE_QPOS, ) - left_open, left_close = get_hand_open_close_qpos(robot, "left_hand", sim.device) - left_pick_action = PickUp( - motion_generator=motion_gen, - cfg=PickUpCfg( - control_part="left_arm", - hand_control_part="left_hand", - hand_open_qpos=left_open, - hand_close_qpos=left_close, - pre_grasp_distance=PICK_APPROACH_DISTANCE, - lift_height=0.12, - sample_interval=PICK_SAMPLE_INTERVAL, - hand_interp_steps=10, - ), + left_open, left_close = get_hand_open_close_qpos( + robot, + hand_control_part="left_hand", + close_qpos=0.030, ) - right_pick_action = PickUp( - motion_generator=motion_gen, - cfg=PickUpCfg( - control_part="right_arm", - hand_control_part="right_hand", - hand_open_qpos=right_open, - hand_close_qpos=right_close, - pre_grasp_distance=PICK_APPROACH_DISTANCE, - lift_height=0.10, - sample_interval=PAN_PICK_SAMPLE_INTERVAL, - hand_interp_steps=PAN_PICK_HAND_INTERP_STEPS, - ), + left_pick_options = PickUpOptions( + pre_grasp_distance=PICK_APPROACH_DISTANCE, + lift_height=0.12, + hand_interp_steps=10, + ) + right_pick_options = PickUpOptions( + pre_grasp_distance=PICK_APPROACH_DISTANCE, + lift_height=0.10, + hand_interp_steps=PAN_PICK_HAND_INTERP_STEPS, + ) + coordinated_options = CoordinatedPlacementOptions( + release=True, + placing_height_offset=BREAD_TARGET_HEIGHT_OFFSET, + support_height_offset=SUPPORT_TARGET_HEIGHT_OFFSET, + lift_height=PLACE_LIFT_HEIGHT, + hand_interp_steps=10, + hold_steps=6, + retreat_steps=18, ) - coordinated_action = CoordinatedPlacement( + engine = AtomicActionEngine( motion_generator=motion_gen, - cfg=CoordinatedPlacementCfg( - control_part="dual_arm", - placing_arm_control_part="left_arm", - support_arm_control_part="right_arm", - placing_hand_control_part="left_hand", - support_hand_control_part="right_hand", - placing_hand_open_qpos=left_open, - placing_hand_close_qpos=left_close, - support_hand_close_qpos=right_close, - release=True, - placing_height_offset=BREAD_TARGET_HEIGHT_OFFSET, - support_height_offset=SUPPORT_TARGET_HEIGHT_OFFSET, - lift_height=PLACE_LIFT_HEIGHT, - sample_interval=COORDINATED_SAMPLE_INTERVAL, - hand_interp_steps=10, - hold_steps=6, - retreat_steps=18, - ), + control_profiles={ + "left_hand": ControlPartCommandProfile.joint_positions( + open=left_open, + grasp=left_close, + ), + "right_hand": ControlPartCommandProfile.joint_positions( + open=right_open, + grasp=right_close, + ), + }, ) - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register(coordinated_action) + coordinated_action = engine.actions["coordinated_placement"] + if not isinstance(coordinated_action, CoordinatedPlacement): + raise RuntimeError("Unexpected coordinated_placement implementation.") full_joint_ids = list(range(robot.dof)) - state = WorldState(last_qpos=robot.get_qpos().clone()) + state = engine.initial_context() wait_for_user = prepare_tutorial_scene( - sim, args, "Inspect the scene, then press Enter to plan left pick-up..." + sim, args, "Inspect the scene, then press Enter to compile both pick-ups..." ) bread_grasp_pose = build_flat_object_grasp_pose( @@ -856,52 +618,65 @@ def run_coordinated_placement_demo( sim.device, z_clearance=BREAD_GRASP_Z_CLEARANCE, ) - start_time = time.time() - left_pick_result = left_pick_action.execute( - GraspTarget( - semantics=bread_semantics, - grasp_xpos=broadcast_pose_batch(bread_grasp_pose, num_envs=n_envs), - ), - state, - ) - logger.log_info( - f"Plan left bread pick-up cost time: {time.time() - start_time:.2f} seconds" - ) - if not left_pick_result.success.all(): - logger.log_warning("Failed to plan left bread pick-up trajectory.") - return - left_pick_traj = left_pick_result.trajectory - state = left_pick_result.next_state - bread_held_state = state.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) - pan_grasp_pose = build_pan_handle_grasp_pose( pan_pose, pan_vertices, sim.device, z_clearance=PAN_GRASP_Z_CLEARANCE, ) - start_time = time.time() - right_pick_result = right_pick_action.execute( - GraspTarget( - semantics=pan_semantics, - grasp_xpos=broadcast_pose_batch(pan_grasp_pose, num_envs=n_envs), + pick_invocations = ( + ActionInvocation( + skill_id="pick_up", + goal=GraspGoal( + semantics=bread_semantics, + grasp_xpos=broadcast_pose_batch(bread_grasp_pose, num_envs=n_envs), + ), + binding=ActionBinding( + manipulators={"primary": "left_arm"}, + end_effectors={"primary": "left_hand"}, + ), + motion_policy=MotionPolicy(sample_count=PICK_SAMPLE_INTERVAL), + skill_options=left_pick_options, + ), + ActionInvocation( + skill_id="pick_up", + goal=GraspGoal( + semantics=pan_semantics, + grasp_xpos=broadcast_pose_batch(pan_grasp_pose, num_envs=n_envs), + ), + binding=ActionBinding( + manipulators={"primary": "right_arm"}, + end_effectors={"primary": "right_hand"}, + ), + motion_policy=MotionPolicy(sample_count=PAN_PICK_SAMPLE_INTERVAL), + skill_options=right_pick_options, ), - state, ) + start_time = time.time() + pick_compiled = engine.compile(pick_invocations, state) logger.log_info( - f"Plan right pan pick-up cost time: {time.time() - start_time:.2f} seconds" + f"Compile both pick-ups cost time: {time.time() - start_time:.2f} seconds" ) - if not right_pick_result.success.all(): + if len(pick_compiled.action_plans) != len(pick_invocations): + logger.log_warning("Failed to compile both pick-up trajectories.") + return + left_pick_result, right_pick_result = pick_compiled.action_plans + if not left_pick_result.plan_success.all(): + logger.log_warning("Failed to plan left bread pick-up trajectory.") + return + if not pick_compiled.plan_success.all(): logger.log_warning("Failed to plan right pan pick-up trajectory.") return - right_pick_traj = right_pick_result.trajectory - state = right_pick_result.next_state + left_pick_traj = left_pick_result.trajectory.positions + right_pick_traj = right_pick_result.trajectory.positions + state = pick_compiled.projected_context + 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.") 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, "left_pick_up", left_pick_traj, full_joint_ids) log_action_plan(robot, "right_pick_up", right_pick_traj, full_joint_ids) if args.diagnose_plan: @@ -910,7 +685,7 @@ def run_coordinated_placement_demo( if wait_for_user: input("Press Enter to execute both pick-up trajectories...") - def log_pick_execution(step_idx: int, total_steps: int) -> None: + def log_trajectory_execution(step_idx: int, total_steps: int) -> None: if args.debug_state and ( step_idx % max(1, total_steps // 10) == 0 or step_idx == total_steps - 1 ): @@ -925,7 +700,7 @@ def log_pick_execution(step_idx: int, total_steps: int) -> None: hold_steps=0, trajectory_sim_steps=TRAJECTORY_SIM_STEPS, joint_ids=full_joint_ids, - on_trajectory_step=log_pick_execution, + on_trajectory_step=log_trajectory_execution, record=False, ) bread.clear_dynamics() @@ -938,10 +713,12 @@ def log_pick_execution(step_idx: int, total_steps: int) -> None: hold_steps=0, trajectory_sim_steps=TRAJECTORY_SIM_STEPS, joint_ids=full_joint_ids, - on_trajectory_step=log_pick_execution, + on_trajectory_step=log_trajectory_execution, record=False, ) pan.clear_dynamics() + # Reconcile the projected task state with measurements before compiling + # placement. This keeps the second phase robust to pick-up execution error. bread_pose_batch = clone_local_pose_from_first_env(bread).to( device=sim.device, dtype=torch.float32 ) @@ -966,10 +743,18 @@ def log_pick_execution(step_idx: int, total_steps: int) -> None: "right_arm", sim.device, ) - held_objects = dict(state.held_objects) + held_objects = dict(state.task.held_objects) held_objects["left_arm"] = bread_held_state held_objects["right_arm"] = pan_held_state - state = state.with_updates(held_objects=held_objects) + state = state.project( + qpos=robot.get_qpos().clone(), + task=TaskState( + batch_size=state.batch_size, + device=state.robot.qpos.device, + held_objects=held_objects, + coordinated_held_objects=state.task.coordinated_held_objects, + ), + ) support_target_pose = build_support_object_target_pose(pan_pose, sim.device) placing_target_pose = build_placing_object_target_pose( @@ -992,7 +777,7 @@ def log_pick_execution(step_idx: int, total_steps: int) -> None: placing_target_pose, num_envs=n_envs, ) - coordinated_target = CoordinatedPlacementTarget( + coordinated_target = CoordinatedPlacementGoal( placing_object_target_pose=broadcast_pose_batch( placing_target_pose, num_envs=n_envs ), @@ -1004,9 +789,30 @@ def log_pick_execution(step_idx: int, total_steps: int) -> None: release=True, ) start_time = time.time() - coordinated_success, coordinated_traj, state = engine.run( - [("coordinated_placement", coordinated_target)], state + placement_compiled = engine.compile( + ( + ActionInvocation( + skill_id="coordinated_placement", + goal=coordinated_target, + binding=ActionBinding( + manipulators={ + "placing": "left_arm", + "support": "right_arm", + }, + end_effectors={ + "placing": "left_hand", + "support": "right_hand", + }, + ), + motion_policy=MotionPolicy(sample_count=COORDINATED_SAMPLE_INTERVAL), + skill_options=coordinated_options, + ), + ), + state, ) + coordinated_success = placement_compiled.plan_success + coordinated_traj = placement_compiled.trajectory.positions + state = placement_compiled.projected_context logger.log_info( "Plan coordinated placement cost time: " f"{time.time() - start_time:.2f} seconds" @@ -1019,7 +825,11 @@ def log_pick_execution(step_idx: int, total_steps: int) -> None: "coordinated_placement", coordinated_traj, full_joint_ids, - coordinated_action._compute_segment_lengths(coordinated_action.cfg.release), + coordinated_action._compute_segment_lengths( + coordinated_options.release, + COORDINATED_SAMPLE_INTERVAL, + coordinated_options, + ), ) if args.diagnose_plan: @@ -1035,12 +845,6 @@ def log_pick_execution(step_idx: int, total_steps: int) -> None: if wait_for_user: input("Press Enter to execute coordinated placement...") - def log_execution(step_idx: int, total_steps: int) -> None: - if args.debug_state and ( - step_idx % max(1, total_steps // 10) == 0 or step_idx == total_steps - 1 - ): - log_execution_state(robot, bread, pan, step_idx, total_steps) - replay_trajectory( sim, robot, @@ -1050,7 +854,7 @@ def log_execution(step_idx: int, total_steps: int) -> None: hold_steps=80, trajectory_sim_steps=TRAJECTORY_SIM_STEPS, joint_ids=full_joint_ids, - on_trajectory_step=log_execution, + on_trajectory_step=log_trajectory_execution, look_at=( (-0.25, 0.0, 2.5), (-0.05, 0.0, 0.72), diff --git a/scripts/tutorials/atomic_action/hand_over.py b/scripts/tutorials/atomic_action/hand_over.py index 33a509299..35864027b 100644 --- a/scripts/tutorials/atomic_action/hand_over.py +++ b/scripts/tutorials/atomic_action/hand_over.py @@ -24,7 +24,6 @@ from __future__ import annotations import argparse -import math import sys from pathlib import Path @@ -32,34 +31,34 @@ if str(_REPO_ROOT) not in sys.path: sys.path.insert(0, str(_REPO_ROOT)) -import numpy as np import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.atomic_actions import ( - GraspTarget, + ActionBinding, + ActionInvocation, + GraspGoal, AtomicActionEngine, - HandOver, - HandOverCfg, - PickUp, - PickUpCfg, -) -from embodichain.lab.sim.cfg import ( - JointDrivePropertiesCfg, - RigidBodyAttributesCfg, - RigidObjectCfg, - RobotCfg, - URDFCfg, + ControlPartCommandProfile, + HandOverOptions, + PickUpOptions, + MotionPolicy, ) +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.data import get_data_path from embodichain.lab.sim.objects import RigidObject, Robot -from embodichain.lab.sim.shapes import CubeCfg, MeshCfg -from embodichain.lab.sim.solvers import URSolverCfg +from embodichain.lab.sim.shapes import MeshCfg from embodichain.utils import logger +from scripts.tutorials.atomic_action.scenario_utils import ( + add_dual_ur5_robot, + add_support_surface, + make_dual_ur5_solver_cfg, + settle_object, +) from scripts.tutorials.atomic_action.tutorial_utils import ( create_antipodal_semantics, create_toppra_motion_generator, + create_tutorial_argument_parser, create_tutorial_simulation, get_hand_open_close_qpos, clone_local_pose_from_first_env, @@ -69,14 +68,8 @@ serve_tutorial_scene, ) -ARM_URDF_PATH = "UniversalRobots/UR5/UR5.urdf" -GRIPPER_URDF_PATH = "DH_PGI_140_80/DH_PGI_140_80.urdf" OBJECT_MESH_PATH = get_data_path("SodaCan/simple_cola_can.obj") GRIPPER_TCP_Z = 0.155 -ROBOT_INIT_POS = (1.95, 0.0, 0.1) -ROBOT_INIT_ROT = (0.0, 0.0, -90.0) -LEFT_ARM_HOME = (0.0, 0.0, -1.57, -1.57, 1.57, 1.57) -RIGHT_ARM_HOME = (-1.57, -1.57, -1.57, -1.57, 0.0, 0.0) SUPPORT_SURFACE_Z = 0.50 SUPPORT_SURFACE_SIZE = (0.8, 1.2, 0.02) SUPPORT_SURFACE_CENTER = ( @@ -117,146 +110,37 @@ def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the demo.""" - parser = argparse.ArgumentParser(description="Dual-arm handover demo") - add_env_launcher_args_to_parser(parser) - parser.set_defaults(device="cpu", renderer="hybrid") - parser.add_argument( - "--diagnose_plan", - action="store_true", - help="Plan and print diagnostics without playing the trajectory.", - ) - parser.add_argument( - "--auto_play", - action="store_true", - help="Run the viewer demo without waiting for keyboard input.", - ) - parser.add_argument( - "--headless_play", - action="store_true", - help="Execute planned trajectories without opening the viewer window.", + parser = create_tutorial_argument_parser( + "Dual-arm handover demo", + features=("diagnose_plan", "headless_play"), + default_device="cpu", + default_renderer="hybrid", ) return parser.parse_args() -def rotation_z(yaw: float) -> np.ndarray: - """Build a 3x3 yaw rotation matrix.""" - cos_yaw = math.cos(yaw) - sin_yaw = math.sin(yaw) - return np.array( - [ - [cos_yaw, -sin_yaw, 0.0], - [sin_yaw, cos_yaw, 0.0], - [0.0, 0.0, 1.0], - ], - dtype=np.float32, - ) - - -def make_transform(xyz: tuple[float, float, float], yaw: float) -> np.ndarray: - """Build a homogeneous transform from translation and yaw.""" - transform = np.eye(4, dtype=np.float32) - transform[:3, :3] = rotation_z(yaw) - transform[:3, 3] = np.asarray(xyz, dtype=np.float32) - return transform - - def create_dual_ur5_robot(sim: SimulationManager) -> Robot: """Create a dual-UR5 robot with one PGI gripper on each arm.""" - arm_urdf_path = ARM_URDF_PATH - gripper_urdf_path = GRIPPER_URDF_PATH - tcp = [ - [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, 0.0], - [0.0, 0.0, 1.0, GRIPPER_TCP_Z], - [0.0, 0.0, 0.0, 1.0], - ] - cfg = RobotCfg( + return add_dual_ur5_robot( + sim, uid="DualUR5HandOver", - urdf_cfg=URDFCfg( - components=[ - { - "component_type": "left_arm", - "urdf_path": arm_urdf_path, - "transform": make_transform((-0.3, -1.45, 0.4), np.pi / 2), - }, - { - "component_type": "right_arm", - "urdf_path": arm_urdf_path, - "transform": make_transform((0.3, -1.45, 0.4), np.pi / 2), - }, - {"component_type": "left_hand", "urdf_path": gripper_urdf_path}, - {"component_type": "right_hand", "urdf_path": gripper_urdf_path}, - ], - fname="dual_ur5_hand_over", - name_case={"joint": "upper", "link": "lower"}, - ), - drive_pros=JointDrivePropertiesCfg( - stiffness={ - "LEFT_JOINT[0-9]": 1e4, - "RIGHT_JOINT[0-9]": 1e4, - "LEFT_GRIPPER_FINGER[1-2]_JOINT_1": 1e2, - "RIGHT_GRIPPER_FINGER[1-2]_JOINT_1": 1e2, - }, - damping={ - "LEFT_JOINT[0-9]": 1e3, - "RIGHT_JOINT[0-9]": 1e3, - "LEFT_GRIPPER_FINGER[1-2]_JOINT_1": 1e1, - "RIGHT_GRIPPER_FINGER[1-2]_JOINT_1": 1e1, - }, - max_effort={ - "LEFT_JOINT[0-9]": 1e5, - "RIGHT_JOINT[0-9]": 1e5, - "LEFT_GRIPPER_FINGER[1-2]_JOINT_1": 1e3, - "RIGHT_GRIPPER_FINGER[1-2]_JOINT_1": 1e3, - }, - drive_type="force", + urdf_name="dual_ur5_hand_over", + solver_cfg=make_dual_ur5_solver_cfg( + GRIPPER_TCP_Z, + ur_ik_nearest_weight=(1.0, 4.0, 1.0, 1.0, 1.0, 1.0), ), - control_parts={ - "left_arm": ["LEFT_JOINT[0-9]"], - "right_arm": ["RIGHT_JOINT[0-9]"], - "dual_arm": ["LEFT_JOINT[0-9]", "RIGHT_JOINT[0-9]"], - "left_hand": ["LEFT_GRIPPER_FINGER1_JOINT_1"], - "right_hand": ["RIGHT_GRIPPER_FINGER1_JOINT_1"], - }, - solver_cfg={ - "left_arm": URSolverCfg( - ur_type="ur5", - tcp=tcp, - end_link_name="left_ee_link", - root_link_name="left_base_link", - ik_nearest_weight=[1.0, 4.0, 1.0, 1.0, 1.0, 1.0], - ), - "right_arm": URSolverCfg( - ur_type="ur5", - tcp=tcp, - end_link_name="right_ee_link", - root_link_name="right_base_link", - ik_nearest_weight=[1.0, 4.0, 1.0, 1.0, 1.0, 1.0], - ), - }, - init_pos=list(ROBOT_INIT_POS), - init_rot=list(ROBOT_INIT_ROT), - init_qpos=list(LEFT_ARM_HOME) + list(RIGHT_ARM_HOME) + [0.0, 0.0, 0.0, 0.0], + hand_stiffness=1e2, + hand_damping=1e1, + hand_max_effort=1e3, ) - return sim.add_robot(cfg=cfg) def create_support_surface(sim: SimulationManager) -> RigidObject: """Create a compact support slab under the staged object.""" - return sim.add_rigid_object( - cfg=RigidObjectCfg( - uid="support_surface", - shape=CubeCfg(size=list(SUPPORT_SURFACE_SIZE)), - attrs=RigidBodyAttributesCfg( - mass=10.0, - dynamic_friction=0.9, - static_friction=0.95, - restitution=0.01, - ), - body_type="static", - init_pos=list(SUPPORT_SURFACE_CENTER), - init_rot=[0.0, 0.0, 0.0], - ) + return add_support_surface( + sim, + size=SUPPORT_SURFACE_SIZE, + center=SUPPORT_SURFACE_CENTER, ) @@ -287,16 +171,6 @@ def create_handover_object(sim: SimulationManager) -> RigidObject: ) -def settle_object(sim: SimulationManager, obj: RigidObject, step: int = 5) -> None: - """Settle an object before planning.""" - if sim.device.type == "cuda": - sim.init_gpu_physics() - obj.reset() - if step > 0: - sim.update(step=step) - obj.clear_dynamics() - - def run_handover_demo( args: argparse.Namespace, sim: SimulationManager, @@ -342,68 +216,80 @@ def run_handover_demo( ) # Step 1 - the left arm picks the object up by its top part. - pick_up_action = PickUp( - motion_generator=motion_gen, - cfg=PickUpCfg( - name="pick_up", - control_part="left_arm", - hand_control_part="left_hand", - hand_open_qpos=left_open, - hand_close_qpos=left_close, - pick_object_part="top", - pre_grasp_distance=PICKUP_PRE_GRASP_DISTANCE, - lift_height=PICKUP_LIFT_HEIGHT, - sample_interval=PICKUP_SAMPLE_INTERVAL, - hand_interp_steps=PICKUP_HAND_INTERP_STEPS, - approach_direction=torch.as_tensor( - [0.0, -707106781, -707106781], dtype=torch.float32 - ), + pick_up_options = PickUpOptions( + pick_object_part="top", + pre_grasp_distance=PICKUP_PRE_GRASP_DISTANCE, + lift_height=PICKUP_LIFT_HEIGHT, + hand_interp_steps=PICKUP_HAND_INTERP_STEPS, + approach_direction=torch.as_tensor( + [0.0, -707106781, -707106781], dtype=torch.float32 ), ) # Step 2 - hand the object from the left arm to the right arm. - handover_action = HandOver( + handover_options = HandOverOptions( + receive_pick_object_part="bottom", + middle_object_pose=middle_pose, + final_object_pose=final_pose, + pre_grasp_distance=HANDOVER_PRE_GRASP_DISTANCE, + lift_height=HANDOVER_LIFT_HEIGHT, + hand_interp_steps=HANDOVER_HAND_INTERP_STEPS, + hold_steps=HANDOVER_HOLD_STEPS, + retreat_steps=HANDOVER_RETREAT_STEPS, + receive_approach_direction=torch.as_tensor( + [0.0, 707106781, -707106781], dtype=torch.float32 + ), + ) + engine = AtomicActionEngine( motion_generator=motion_gen, - cfg=HandOverCfg( - name="hand_over", - control_part="dual_arm", - transfer_arm_control_part="left_arm", - receive_arm_control_part="right_arm", - transfer_hand_control_part="left_hand", - receive_hand_control_part="right_hand", - transfer_hand_open_qpos=left_open, - transfer_hand_close_qpos=left_close, - receive_hand_open_qpos=right_open, - receive_hand_close_qpos=right_close, - receive_pick_object_part="bottom", - middle_object_pose=middle_pose, - final_object_pose=final_pose, - pre_grasp_distance=HANDOVER_PRE_GRASP_DISTANCE, - lift_height=HANDOVER_LIFT_HEIGHT, - sample_interval=HANDOVER_SAMPLE_INTERVAL, - hand_interp_steps=HANDOVER_HAND_INTERP_STEPS, - hold_steps=HANDOVER_HOLD_STEPS, - retreat_steps=HANDOVER_RETREAT_STEPS, - receive_approach_direction=torch.as_tensor( - [0.0, 707106781, -707106781], dtype=torch.float32 + control_profiles={ + "left_hand": ControlPartCommandProfile.joint_positions( + open=left_open, + grasp=left_close, ), - ), + "right_hand": ControlPartCommandProfile.joint_positions( + open=right_open, + grasp=right_close, + ), + }, ) - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register(pick_up_action) - engine.register(handover_action) - wait_for_user = prepare_tutorial_scene( sim, args, "Inspect the scene, then press Enter to plan the handover..." ) # wait for object to drop for _ in range(20): sim.update(step=10) - success, traj, _ = engine.run( - [ - ("pick_up", GraspTarget(object_semantics)), - ("hand_over", GraspTarget(object_semantics)), - ] + compiled = engine.compile( + ( + ActionInvocation( + "pick_up", + GraspGoal(object_semantics), + ActionBinding( + manipulators={"primary": "left_arm"}, + end_effectors={"primary": "left_hand"}, + ), + MotionPolicy(sample_count=PICKUP_SAMPLE_INTERVAL), + skill_options=pick_up_options, + ), + ActionInvocation( + "hand_over", + GraspGoal(object_semantics), + ActionBinding( + manipulators={ + "source": "left_arm", + "destination": "right_arm", + }, + end_effectors={ + "source": "left_hand", + "destination": "right_hand", + }, + ), + MotionPolicy(sample_count=HANDOVER_SAMPLE_INTERVAL), + skill_options=handover_options, + ), + ) ) + success = compiled.plan_success + traj = compiled.trajectory.positions if not success.all(): logger.log_warning("Failed to plan the full pick-up + handover trajectory.") diff --git a/scripts/tutorials/atomic_action/move_end_effector.py b/scripts/tutorials/atomic_action/move_end_effector.py index dff45def5..27cef7267 100644 --- a/scripts/tutorials/atomic_action/move_end_effector.py +++ b/scripts/tutorials/atomic_action/move_end_effector.py @@ -28,12 +28,12 @@ import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - EndEffectorPoseTarget, - MoveEndEffector, - MoveEndEffectorCfg, + EndEffectorPoseGoal, + MotionPolicy, ) from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( @@ -41,6 +41,7 @@ broadcast_pose_batch, broadcast_waypoint_pose_batch, create_toppra_motion_generator, + create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, make_top_down_eef_pose, @@ -55,12 +56,10 @@ def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the MoveEndEffector tutorial.""" - parser = argparse.ArgumentParser( - description="Demonstrate MoveEndEffector with a multi-waypoint pose trajectory." + parser = create_tutorial_argument_parser( + "Demonstrate MoveEndEffector with a multi-waypoint pose trajectory.", + features=("visualize_axes",), ) - add_env_launcher_args_to_parser(parser) - parser.add_argument("--auto_play", action="store_true") - parser.add_argument("--no_vis_eef_axis", action="store_true") return parser.parse_args() @@ -72,12 +71,6 @@ def main() -> None: motion_gen = create_toppra_motion_generator(robot) engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register( - MoveEndEffector( - motion_gen, - cfg=MoveEndEffectorCfg(sample_interval=MOVE_SAMPLE_INTERVAL), - ) - ) poses = torch.stack( [ @@ -99,15 +92,17 @@ def main() -> None: sim, args, "Inspect the robot, then press Enter to plan MoveEndEffector..." ) - success, trajectory, _ = engine.run( - [ - ( - "move_end_effector", - EndEffectorPoseTarget(broadcast_waypoint_pose_batch(poses, n_envs)), - ) - ] + compiled = engine.compile( + ( + ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal(broadcast_waypoint_pose_batch(poses, n_envs)), + binding=ActionBinding(manipulators={"primary": "arm"}), + motion_policy=MotionPolicy(sample_count=MOVE_SAMPLE_INTERVAL), + ), + ) ) - if not success.all(): + if not compiled.plan_success.all(): logger.log_warning("Failed to plan MoveEndEffector demo trajectory.") return @@ -116,7 +111,7 @@ def main() -> None: replay_trajectory( sim, robot, - trajectory, + compiled.trajectory.positions, args, video_prefix="move_end_effector_auto_play", hold_steps=POST_TRAJECTORY_STEPS, diff --git a/scripts/tutorials/atomic_action/move_held_object.py b/scripts/tutorials/atomic_action/move_held_object.py index d30b8c7d1..22d4873e6 100644 --- a/scripts/tutorials/atomic_action/move_held_object.py +++ b/scripts/tutorials/atomic_action/move_held_object.py @@ -29,18 +29,16 @@ import torch from embodichain.data import get_data_path -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - EndEffectorPoseTarget, - GraspTarget, - HeldObjectPoseTarget, - MoveEndEffector, - MoveEndEffectorCfg, - MoveHeldObject, - MoveHeldObjectCfg, - PickUp, - PickUpCfg, + ControlPartCommandProfile, + EndEffectorPoseGoal, + GraspGoal, + HeldObjectPoseGoal, + PickUpOptions, + MotionPolicy, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.lab.sim.objects import RigidObject @@ -52,9 +50,11 @@ clone_local_pose_from_first_env, create_antipodal_semantics, create_toppra_motion_generator, + create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, get_hand_open_close_qpos, + make_clear_dynamics_callback, make_eef_pose_at, prepare_tutorial_scene, replay_trajectory, @@ -72,14 +72,10 @@ def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the MoveHeldObject tutorial.""" - parser = argparse.ArgumentParser( - description="Pick up a paper cup and move it by object pose." + parser = create_tutorial_argument_parser( + "Pick up a paper cup and move it by object pose.", + features=("grasp_sampling", "visualize_axes"), ) - add_env_launcher_args_to_parser(parser) - parser.add_argument("--n_sample", type=int, default=10000) - parser.add_argument("--force_reannotate", action="store_true") - parser.add_argument("--auto_play", action="store_true") - parser.add_argument("--no_vis_eef_axis", action="store_true") return parser.parse_args() @@ -126,35 +122,15 @@ def main() -> None: motion_gen = create_toppra_motion_generator(robot) hand_open, hand_close = get_hand_open_close_qpos(robot) - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register( - MoveEndEffector( - motion_gen, MoveEndEffectorCfg(sample_interval=MOVE_SAMPLE_INTERVAL) - ) - ) - engine.register( - PickUp( - motion_gen, - PickUpCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - pre_grasp_distance=0.15, - lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, ) - engine.register( - MoveHeldObject( - motion_gen, - MoveHeldObjectCfg( - hand_close_qpos=hand_close, - sample_interval=MOVE_HELD_OBJECT_SAMPLE_INTERVAL, - ), - ) - ) - semantics = create_antipodal_semantics( obj, label="paper_cup", @@ -172,14 +148,38 @@ def main() -> None: sim, args, "Inspect the paper cup, then press Enter to plan..." ) - success, trajectory, _ = engine.run( - [ - ("move_end_effector", EndEffectorPoseTarget(move_target)), - ("pick_up", GraspTarget(semantics)), - ("move_held_object", HeldObjectPoseTarget(object_target)), - ] + binding = ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, ) - if not success.all(): + compiled = engine.compile( + ( + ActionInvocation( + "move_end_effector", + EndEffectorPoseGoal(move_target), + binding, + MotionPolicy(sample_count=MOVE_SAMPLE_INTERVAL), + ), + ActionInvocation( + "pick_up", + GraspGoal(semantics), + binding, + MotionPolicy(sample_count=PICK_SAMPLE_INTERVAL), + skill_options=PickUpOptions( + pre_grasp_distance=0.15, + lift_height=0.16, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ), + ActionInvocation( + "move_held_object", + HeldObjectPoseGoal(object_target), + binding, + MotionPolicy(sample_count=MOVE_HELD_OBJECT_SAMPLE_INTERVAL), + ), + ) + ) + if not compiled.plan_success.all(): logger.log_warning("Failed to plan MoveHeldObject demo trajectory.") return @@ -190,22 +190,14 @@ def main() -> None: + round((PICK_SAMPLE_INTERVAL - HAND_INTERP_STEPS) * 0.6) + HAND_INTERP_STEPS ) - dynamics_cleared = False - - def clear_object_dynamics(step_idx: int, _: int) -> None: - nonlocal dynamics_cleared - if not dynamics_cleared and step_idx + 1 >= clear_after_step: - obj.clear_dynamics() - dynamics_cleared = True - replay_trajectory( sim, robot, - trajectory, + compiled.trajectory.positions, args, video_prefix="move_held_object_auto_play", hold_steps=POST_TRAJECTORY_STEPS, - on_trajectory_step=clear_object_dynamics, + on_trajectory_step=make_clear_dynamics_callback(obj, clear_after_step), ) if wait_for_user: input("Press Enter to exit the simulation...") diff --git a/scripts/tutorials/atomic_action/move_joints.py b/scripts/tutorials/atomic_action/move_joints.py index 2e0addb13..0439bf876 100644 --- a/scripts/tutorials/atomic_action/move_joints.py +++ b/scripts/tutorials/atomic_action/move_joints.py @@ -28,18 +28,19 @@ import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - JointPositionTarget, - MoveJoints, - MoveJointsCfg, - NamedJointPositionTarget, + ControlPartCommandProfile, + JointPositionGoal, + MotionPolicy, ) from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( add_ur5_gripper_robot, create_toppra_motion_generator, + create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, prepare_tutorial_scene, @@ -53,12 +54,10 @@ def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the MoveJoints tutorial.""" - parser = argparse.ArgumentParser( - description="Demonstrate MoveJoints with named and explicit qpos targets." + parser = create_tutorial_argument_parser( + "Demonstrate MoveJoints with named and explicit qpos targets.", + features=("visualize_axes",), ) - add_env_launcher_args_to_parser(parser) - parser.add_argument("--auto_play", action="store_true") - parser.add_argument("--no_vis_eef_axis", action="store_true") return parser.parse_args() @@ -77,17 +76,12 @@ def main() -> None: [0.0, -1.57, 1.57, -1.57, -1.57, 0.0], ) ) - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register( - MoveJoints( - motion_gen, - cfg=MoveJointsCfg( - sample_interval=MOVE_JOINTS_SAMPLE_INTERVAL, - named_joint_positions={"ready": ready}, - ), - ) + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "arm": ControlPartCommandProfile.joint_positions(ready=ready), + }, ) - if not args.no_vis_eef_axis: draw_axis_marker( sim, @@ -101,13 +95,19 @@ def main() -> None: waypoints = ( torch.stack([mid, home]).unsqueeze(0).repeat(robot.get_qpos().shape[0], 1, 1) ) - success, trajectory, _ = engine.run( - [ - ("move_joints", NamedJointPositionTarget("ready")), - ("move_joints", JointPositionTarget(waypoints)), - ] + binding = ActionBinding(manipulators={"primary": "arm"}) + policy = MotionPolicy(sample_count=MOVE_JOINTS_SAMPLE_INTERVAL) + compiled = engine.compile( + ( + ActionInvocation( + "move_joints", JointPositionGoal("ready"), binding, policy + ), + ActionInvocation( + "move_joints", JointPositionGoal(waypoints), binding, policy + ), + ) ) - if not success.all(): + if not compiled.plan_success.all(): logger.log_warning("Failed to plan MoveJoints demo trajectory.") return @@ -116,7 +116,7 @@ def main() -> None: replay_trajectory( sim, robot, - trajectory, + compiled.trajectory.positions, args, video_prefix="move_joints_auto_play", hold_steps=POST_TRAJECTORY_STEPS, diff --git a/scripts/tutorials/atomic_action/pickup.py b/scripts/tutorials/atomic_action/pickup.py index cf0acef4c..403061292 100644 --- a/scripts/tutorials/atomic_action/pickup.py +++ b/scripts/tutorials/atomic_action/pickup.py @@ -28,12 +28,14 @@ import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - GraspTarget, - PickUp, - PickUpCfg, + ControlPartCommandProfile, + GraspGoal, + PickUpOptions, + MotionPolicy, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.lab.sim.objects import RigidObject @@ -44,10 +46,12 @@ clone_local_pose_from_first_env, create_antipodal_semantics, create_toppra_motion_generator, + create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, get_hand_open_close_qpos, initialize_pre_pick_robot_pose, + make_clear_dynamics_callback, prepare_tutorial_scene, replay_trajectory, run_tutorial, @@ -67,16 +71,14 @@ def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the PickUp tutorial.""" - parser = argparse.ArgumentParser(description="Demonstrate PickUp on a cube.") - add_env_launcher_args_to_parser(parser) - parser.add_argument("--n_sample", type=int, default=10000) - parser.add_argument("--force_reannotate", action="store_true") - parser.add_argument("--auto_play", action="store_true") + parser = create_tutorial_argument_parser( + "Demonstrate PickUp on a cube.", + features=("grasp_sampling", "visualize_axes"), + ) parser.add_argument( "--approach", choices=[*APPROACH_DIRECTIONS, "custom"], default="top" ) parser.add_argument("--custom_approach_direction", type=float, nargs=3) - parser.add_argument("--no_vis_eef_axis", action="store_true") return parser.parse_args() @@ -130,20 +132,14 @@ def main() -> None: initialize_pre_pick_robot_pose(robot, obj, hand_open) motion_gen = create_toppra_motion_generator(robot) - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register( - PickUp( - motion_gen, - cfg=PickUpCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - approach_direction=resolve_approach_direction(args, sim.device), - pre_grasp_distance=0.15, - lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, ) semantics = create_antipodal_semantics( obj, @@ -157,8 +153,26 @@ def main() -> None: sim, args, "Inspect the cube, then press Enter to plan PickUp..." ) - success, trajectory, _ = engine.run([("pick_up", GraspTarget(semantics))]) - if not success.all(): + compiled = engine.compile( + ( + ActionInvocation( + skill_id="pick_up", + goal=GraspGoal(semantics), + binding=ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ), + motion_policy=MotionPolicy(sample_count=PICK_SAMPLE_INTERVAL), + skill_options=PickUpOptions( + approach_direction=resolve_approach_direction(args, sim.device), + pre_grasp_distance=0.15, + lift_height=0.16, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ), + ) + ) + if not compiled.plan_success.all(): logger.log_warning("Failed to plan PickUp demo trajectory.") return @@ -167,22 +181,14 @@ def main() -> None: clear_after_step = ( round((PICK_SAMPLE_INTERVAL - HAND_INTERP_STEPS) * 0.6) + HAND_INTERP_STEPS ) - dynamics_cleared = False - - def clear_object_dynamics(step_idx: int, _: int) -> None: - nonlocal dynamics_cleared - if not dynamics_cleared and step_idx + 1 >= clear_after_step: - obj.clear_dynamics() - dynamics_cleared = True - replay_trajectory( sim, robot, - trajectory, + compiled.trajectory.positions, args, video_prefix="pickup_cube_auto_play", hold_steps=POST_TRAJECTORY_STEPS, - on_trajectory_step=clear_object_dynamics, + on_trajectory_step=make_clear_dynamics_callback(obj, clear_after_step), ) if wait_for_user: input("Press Enter to exit the simulation...") diff --git a/scripts/tutorials/atomic_action/place.py b/scripts/tutorials/atomic_action/place.py index fb3ae0efe..6531e59cc 100644 --- a/scripts/tutorials/atomic_action/place.py +++ b/scripts/tutorials/atomic_action/place.py @@ -28,15 +28,16 @@ import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - GraspTarget, - PickUp, - PickUpCfg, - Place, - PlaceCfg, - PlaceTarget, + ControlPartCommandProfile, + GraspGoal, + PickUpOptions, + PlaceGoal, + PlaceOptions, + MotionPolicy, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.lab.sim.objects import RigidObject @@ -49,10 +50,12 @@ clone_local_pose_from_first_env, create_antipodal_semantics, create_toppra_motion_generator, + create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, get_hand_open_close_qpos, initialize_pre_pick_robot_pose, + make_clear_dynamics_callback, prepare_tutorial_scene, replay_trajectory, run_tutorial, @@ -69,14 +72,10 @@ def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the Place tutorial.""" - parser = argparse.ArgumentParser( - description="Pick up a cube and place it at a target pose." + parser = create_tutorial_argument_parser( + "Pick up a cube and place it at a target pose.", + features=("grasp_sampling", "visualize_axes"), ) - add_env_launcher_args_to_parser(parser) - parser.add_argument("--n_sample", type=int, default=10000) - parser.add_argument("--force_reannotate", action="store_true") - parser.add_argument("--auto_play", action="store_true") - parser.add_argument("--no_vis_eef_axis", action="store_true") return parser.parse_args() @@ -131,33 +130,15 @@ def main() -> None: hand_open, hand_close = get_hand_open_close_qpos(robot) initialize_pre_pick_robot_pose(robot, obj, hand_open) - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register( - PickUp( - motion_gen, - cfg=PickUpCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - pre_grasp_distance=0.15, - lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, ) - engine.register( - Place( - motion_gen, - cfg=PlaceCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, - lift_height=PLACE_LIFT_HEIGHT, - sample_interval=PLACE_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) - ) - semantics = create_antipodal_semantics( obj, label="cube", @@ -175,20 +156,40 @@ def main() -> None: sim, args, "Inspect the cube, then press Enter to plan PickUp -> Place..." ) - success, trajectory, _ = engine.run( - [ - ("pick_up", GraspTarget(semantics)), - ( + binding = ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ) + compiled = engine.compile( + ( + ActionInvocation( + "pick_up", + GraspGoal(semantics), + binding, + MotionPolicy(sample_count=PICK_SAMPLE_INTERVAL), + skill_options=PickUpOptions( + pre_grasp_distance=0.15, + lift_height=0.16, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ), + ActionInvocation( "place", - PlaceTarget( + PlaceGoal( broadcast_waypoint_pose_batch( place_poses, robot.get_qpos().shape[0] ) ), + binding, + MotionPolicy(sample_count=PLACE_SAMPLE_INTERVAL), + skill_options=PlaceOptions( + lift_height=PLACE_LIFT_HEIGHT, + hand_interp_steps=HAND_INTERP_STEPS, + ), ), - ] + ) ) - if not success.all(): + if not compiled.plan_success.all(): logger.log_warning("Failed to plan Place demo trajectory.") return @@ -197,22 +198,14 @@ def main() -> None: clear_after_step = ( round((PICK_SAMPLE_INTERVAL - HAND_INTERP_STEPS) * 0.6) + HAND_INTERP_STEPS ) - dynamics_cleared = False - - def clear_object_dynamics(step_idx: int, _: int) -> None: - nonlocal dynamics_cleared - if not dynamics_cleared and step_idx + 1 >= clear_after_step: - obj.clear_dynamics() - dynamics_cleared = True - replay_trajectory( sim, robot, - trajectory, + compiled.trajectory.positions, args, video_prefix="place_auto_play", hold_steps=POST_TRAJECTORY_STEPS, - on_trajectory_step=clear_object_dynamics, + on_trajectory_step=make_clear_dynamics_callback(obj, clear_after_step), ) if wait_for_user: input("Press Enter to exit the simulation...") diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py index 39cd829dd..83a8dc0ff 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -28,15 +28,15 @@ import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - EndEffectorPoseTarget, - MoveEndEffector, - MoveEndEffectorCfg, - Press, - PressCfg, - PressTarget, + ControlPartCommandProfile, + EndEffectorPoseGoal, + PressOptions, + PressGoal, + MotionPolicy, ) from embodichain.lab.sim.cfg import ( RigidBodyAttributesCfg, @@ -50,6 +50,7 @@ add_ur5_gripper_robot, broadcast_pose_batch, create_toppra_motion_generator, + create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, format_tensor, @@ -72,15 +73,14 @@ def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the Press tutorial.""" - parser = argparse.ArgumentParser(description="Demonstrate Press on a wooden block.") - add_env_launcher_args_to_parser(parser) - parser.add_argument("--auto_play", action="store_true") - parser.add_argument("--debug_state", action="store_true") + parser = create_tutorial_argument_parser( + "Demonstrate Press on a wooden block.", + features=("debug_state", "visualize_axes"), + ) parser.add_argument( "--press_tolerance", type=float, default=DEFAULT_PRESS_TOLERANCE ) parser.add_argument("--block_pos", type=float, nargs=2, default=(-0.30, -0.12)) - parser.add_argument("--no_vis_eef_axis", action="store_true") return parser.parse_args() @@ -159,24 +159,16 @@ def main() -> None: block.clear_dynamics() motion_gen = create_toppra_motion_generator(robot) - hand_close = get_hand_open_close_qpos(robot)[1] - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register( - MoveEndEffector( - motion_gen, MoveEndEffectorCfg(sample_interval=MOVE_SAMPLE_INTERVAL) - ) - ) - engine.register( - Press( - motion_gen, - PressCfg( - hand_close_qpos=hand_close, - sample_interval=PRESS_SAMPLE_INTERVAL, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) + hand_open, hand_close = get_hand_open_close_qpos(robot) + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, ) - block_center = block.get_local_pose(to_matrix=True)[0, :3, 3] press_position = block_center.clone() press_position[2] += 0.5 * BLOCK_SIZE[2] + PRESS_SURFACE_OFFSET @@ -191,15 +183,33 @@ def main() -> None: sim, args, "Inspect the wooden block, then press Enter to plan..." ) - success, trajectory, _ = engine.run( - [ - ("move_end_effector", EndEffectorPoseTarget(move_target)), - ("press", PressTarget(press_target)), - ] + binding = ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ) + compiled = engine.compile( + ( + ActionInvocation( + "move_end_effector", + EndEffectorPoseGoal(move_target), + binding, + MotionPolicy(sample_count=MOVE_SAMPLE_INTERVAL), + ), + ActionInvocation( + "press", + PressGoal(press_target), + binding, + MotionPolicy(sample_count=PRESS_SAMPLE_INTERVAL), + skill_options=PressOptions( + hand_interp_steps=HAND_INTERP_STEPS, + ), + ), + ) ) - if not success.all(): + if not compiled.plan_success.all(): logger.log_warning("Failed to plan Press demo trajectory.") return + trajectory = compiled.trajectory.positions is_center_hit, center_error, hit_step, hit_pos, expected_pos = ( compute_press_center_check(robot, trajectory, block, args.press_tolerance) ) diff --git a/scripts/tutorials/atomic_action/scenario_utils.py b/scripts/tutorials/atomic_action/scenario_utils.py new file mode 100644 index 000000000..b51aefb91 --- /dev/null +++ b/scripts/tutorials/atomic_action/scenario_utils.py @@ -0,0 +1,351 @@ +# ---------------------------------------------------------------------------- +# 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 scene and dual-arm helpers for atomic-action tutorials.""" + +from __future__ import annotations + +import math +import os +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Literal + +import numpy as np +import torch + +from embodichain.data import get_data_path +from embodichain.lab.sim import SimulationManager +from embodichain.lab.sim.atomic_actions import Affordance, ObjectSemantics +from embodichain.lab.sim.cfg import ( + JointDrivePropertiesCfg, + RigidBodyAttributesCfg, + RigidObjectCfg, + RobotCfg, + URDFCfg, +) +from embodichain.lab.sim.objects import RigidObject, Robot +from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.sim.solvers import PytorchSolverCfg, SolverCfg, URSolverCfg +from embodichain.utils import logger + +ARM_URDF_PATH = "UniversalRobots/UR5/UR5.urdf" +GRIPPER_URDF_PATH = "DH_PGI_140_80/DH_PGI_140_80.urdf" +LEFT_ARM_HOME = (0.0, 0.0, -1.57, -1.57, 1.57, 1.57) +RIGHT_ARM_HOME = (-1.57, -1.57, -1.57, -1.57, 0.0, 0.0) +DUAL_UR5_INIT_POS = (1.95, 0.0, 0.1) +DUAL_UR5_INIT_ROT = (0.0, 0.0, -90.0) + + +def resolve_cached_data_path(data_path: str) -> str: + """Resolve an asset from the local cache, falling back to project data.""" + if os.path.isabs(data_path): + return data_path + + data_root = Path( + os.environ.get( + "EMBODICHAIN_DATA_ROOT", + str(Path.home() / ".cache" / "embodichain_data"), + ) + ) + for candidate in (data_root / data_path, data_root / "extract" / data_path): + if candidate.exists(): + return str(candidate) + return get_data_path(data_path) + + +def make_yaw_transform(xyz: tuple[float, float, float], yaw: float) -> np.ndarray: + """Build a homogeneous transform from translation and world yaw.""" + cos_yaw = math.cos(yaw) + sin_yaw = math.sin(yaw) + transform = np.eye(4, dtype=np.float32) + transform[:3, :3] = np.array( + ( + (cos_yaw, -sin_yaw, 0.0), + (sin_yaw, cos_yaw, 0.0), + (0.0, 0.0, 1.0), + ), + dtype=np.float32, + ) + transform[:3, 3] = np.asarray(xyz, dtype=np.float32) + return transform + + +def make_dual_ur5_solver_cfg( + tcp_z: float, + *, + solver: Literal["pytorch", "ur"] = "ur", + ur_ik_nearest_weight: Sequence[float] | None = None, + clear_urdf_path: bool = False, + pytorch_num_samples: int = 30, +) -> dict[str, SolverCfg]: + """Build matching left/right solver configs for dual-UR5 tutorials.""" + tcp = [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, tcp_z], + [0.0, 0.0, 0.0, 1.0], + ] + configs: dict[str, SolverCfg] = {} + for prefix in ("left", "right"): + if solver == "pytorch": + config = PytorchSolverCfg( + end_link_name=f"{prefix}_ee_link", + root_link_name=f"{prefix}_base_link", + tcp=tcp, + num_samples=pytorch_num_samples, + ) + else: + config = URSolverCfg( + ur_type="ur5", + end_link_name=f"{prefix}_ee_link", + root_link_name=f"{prefix}_base_link", + tcp=tcp, + ik_nearest_weight=ur_ik_nearest_weight, + ) + if clear_urdf_path: + config.urdf_path = None + configs[f"{prefix}_arm"] = config + return configs + + +def add_dual_ur5_robot( + sim: SimulationManager, + *, + uid: str, + urdf_name: str, + solver_cfg: Mapping[str, SolverCfg], + init_pos: Sequence[float] = DUAL_UR5_INIT_POS, + init_rot: Sequence[float] = DUAL_UR5_INIT_ROT, + left_arm_home: Sequence[float] = LEFT_ARM_HOME, + right_arm_home: Sequence[float] = RIGHT_ARM_HOME, + arm_urdf_path: str = ARM_URDF_PATH, + gripper_urdf_path: str = GRIPPER_URDF_PATH, + joint_name_case: str = "upper", + set_urdf_name_case: bool = True, + hand_stiffness: float = 1e3, + hand_damping: float = 1e2, + hand_max_effort: float = 1e4, +) -> Robot: + """Add the common dual-UR5 and dual-gripper tutorial embodiment.""" + if joint_name_case not in {"lower", "upper"}: + raise ValueError("joint_name_case must be 'lower' or 'upper'.") + prefix = str.upper if joint_name_case == "upper" else str.lower + left_joint = prefix("left_joint[0-9]") + right_joint = prefix("right_joint[0-9]") + left_hand = prefix("left_gripper_finger[1-2]_joint_1") + right_hand = prefix("right_gripper_finger[1-2]_joint_1") + left_hand_control = prefix("left_gripper_finger1_joint_1") + right_hand_control = prefix("right_gripper_finger1_joint_1") + + urdf_cfg = URDFCfg( + components=[ + { + "component_type": "left_arm", + "urdf_path": arm_urdf_path, + "transform": make_yaw_transform((-0.3, -1.45, 0.4), np.pi / 2), + }, + { + "component_type": "right_arm", + "urdf_path": arm_urdf_path, + "transform": make_yaw_transform((0.3, -1.45, 0.4), np.pi / 2), + }, + {"component_type": "left_hand", "urdf_path": gripper_urdf_path}, + {"component_type": "right_hand", "urdf_path": gripper_urdf_path}, + ], + fname=urdf_name, + ) + if set_urdf_name_case: + urdf_cfg.name_case = {"joint": joint_name_case, "link": "lower"} + + cfg = RobotCfg( + uid=uid, + urdf_cfg=urdf_cfg, + drive_pros=JointDrivePropertiesCfg( + stiffness={ + left_joint: 1e4, + right_joint: 1e4, + left_hand: hand_stiffness, + right_hand: hand_stiffness, + }, + damping={ + left_joint: 1e3, + right_joint: 1e3, + left_hand: hand_damping, + right_hand: hand_damping, + }, + max_effort={ + left_joint: 1e5, + right_joint: 1e5, + left_hand: hand_max_effort, + right_hand: hand_max_effort, + }, + drive_type="force", + ), + control_parts={ + "left_arm": [left_joint], + "right_arm": [right_joint], + "dual_arm": [left_joint, right_joint], + "left_hand": [left_hand_control], + "right_hand": [right_hand_control], + }, + solver_cfg=dict(solver_cfg), + init_pos=list(init_pos), + init_rot=list(init_rot), + init_qpos=list(left_arm_home) + list(right_arm_home) + [0.0, 0.0, 0.0, 0.0], + ) + return sim.add_robot(cfg=cfg) + + +def add_support_surface( + sim: SimulationManager, + *, + size: Sequence[float], + center: Sequence[float], +) -> RigidObject: + """Add the standard static support slab used by dual-arm tutorials.""" + return sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="support_surface", + shape=CubeCfg(size=list(size)), + attrs=RigidBodyAttributesCfg( + mass=10.0, + dynamic_friction=0.9, + static_friction=0.95, + restitution=0.01, + ), + body_type="static", + init_pos=list(center), + init_rot=[0.0, 0.0, 0.0], + ) + ) + + +def settle_object(sim: SimulationManager, obj: RigidObject, step: int = 5) -> None: + """Reset, settle, and freeze an object before tutorial planning.""" + if sim.device.type == "cuda": + sim.init_gpu_physics() + obj.reset() + if step > 0: + sim.update(step=step) + obj.clear_dynamics() + + +def create_manual_object_semantics(obj: RigidObject, label: str) -> ObjectSemantics: + """Create minimal semantics for a caller-provided grasp pose.""" + return ObjectSemantics( + label=label, + geometry={}, + affordance=Affordance(object_label=label), + entity=obj, + ) + + +def get_local_vertices(obj: RigidObject) -> torch.Tensor: + """Return scaled local vertices from the first environment.""" + return obj.get_vertices(env_ids=[0], scale=True)[0] + + +def compute_local_bounds( + vertices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute a local mesh axis-aligned bounding box.""" + return vertices.min(dim=0).values, vertices.max(dim=0).values + + +def invert_pose(pose: torch.Tensor) -> torch.Tensor: + """Invert a batch of homogeneous transforms.""" + inv_pose = pose.clone() + rot_t = pose[:, :3, :3].transpose(1, 2) + inv_pose[:, :3, :3] = rot_t + inv_pose[:, :3, 3] = -torch.bmm(rot_t, pose[:, :3, 3:4]).squeeze(-1) + return inv_pose + + +def transform_points(pose: torch.Tensor, points: torch.Tensor) -> torch.Tensor: + """Transform local points by a homogeneous pose.""" + return points @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + + +def compute_world_bounds( + object_pose: torch.Tensor, + local_vertices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute a world-space AABB for local mesh vertices.""" + world_vertices = transform_points(object_pose, local_vertices) + return world_vertices.min(dim=0).values, world_vertices.max(dim=0).values + + +def normalize_vector(vector: torch.Tensor, fallback: torch.Tensor) -> torch.Tensor: + """Normalize a vector with a deterministic degenerate fallback.""" + norm = torch.linalg.norm(vector) + if norm < 1e-6: + return fallback.to(device=vector.device, dtype=vector.dtype) + return vector / norm + + +def rotate_pose_about_world_z(pose: torch.Tensor, yaw_deg: float) -> torch.Tensor: + """Rotate pose orientation about world Z while preserving translation.""" + yaw = math.radians(yaw_deg) + rot = torch.eye(3, dtype=pose.dtype, device=pose.device) + rot[0, 0] = math.cos(yaw) + rot[0, 1] = -math.sin(yaw) + rot[1, 0] = math.sin(yaw) + rot[1, 1] = math.cos(yaw) + rotated_pose = pose.clone() + rotated_pose[:3, :3] = rot @ pose[:3, :3] + return rotated_pose + + +def log_action_plan( + robot: Robot, + action_name: str, + trajectory: torch.Tensor, + joint_ids: list[int], + segments: Mapping[str, int] | None = None, +) -> None: + """Log joint and segment details for a planned tutorial action.""" + joint_names = [robot.joint_names[joint_id] for joint_id in joint_ids] + logger.log_info(f"{action_name} joint ids: {joint_ids}") + logger.log_info(f"{action_name} joint names: {joint_names}") + logger.log_info(f"{action_name} trajectory shape: {tuple(trajectory.shape)}") + if segments is not None: + logger.log_info(f"{action_name} trajectory segments: {dict(segments)}") + + +__all__ = [ + "ARM_URDF_PATH", + "DUAL_UR5_INIT_POS", + "DUAL_UR5_INIT_ROT", + "GRIPPER_URDF_PATH", + "LEFT_ARM_HOME", + "RIGHT_ARM_HOME", + "add_dual_ur5_robot", + "add_support_surface", + "compute_local_bounds", + "compute_world_bounds", + "create_manual_object_semantics", + "get_local_vertices", + "invert_pose", + "log_action_plan", + "make_dual_ur5_solver_cfg", + "make_yaw_transform", + "normalize_vector", + "resolve_cached_data_path", + "rotate_pose_about_world_z", + "settle_object", + "transform_points", +] diff --git a/scripts/tutorials/atomic_action/tutorial_utils.py b/scripts/tutorials/atomic_action/tutorial_utils.py index 48be356df..7edc483c1 100644 --- a/scripts/tutorials/atomic_action/tutorial_utils.py +++ b/scripts/tutorials/atomic_action/tutorial_utils.py @@ -20,11 +20,12 @@ import argparse import time -from collections.abc import Callable, Sequence +from collections.abc import Callable, Collection, Sequence +from typing import Literal import torch -from embodichain.data import get_data_path +from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.atomic_actions import ( @@ -74,6 +75,67 @@ (0.0401, 0.0000, -0.9992), ) +TutorialCliFeature = Literal[ + "debug_state", + "diagnose_plan", + "grasp_sampling", + "headless_play", + "visualize_axes", +] + + +def create_tutorial_argument_parser( + description: str, + *, + features: Collection[TutorialCliFeature] = (), + default_device: str | None = None, + default_renderer: str | None = None, +) -> argparse.ArgumentParser: + """Create a launcher parser with the shared atomic-tutorial switches.""" + parser = argparse.ArgumentParser(description=description) + add_env_launcher_args_to_parser(parser) + defaults = {} + if default_device is not None: + defaults["device"] = default_device + if default_renderer is not None: + defaults["renderer"] = default_renderer + if defaults: + parser.set_defaults(**defaults) + + parser.add_argument( + "--auto_play", + action="store_true", + help="Run the demo without waiting for keyboard input.", + ) + if "debug_state" in features: + parser.add_argument( + "--debug_state", + action="store_true", + help="Log robot and object state during execution.", + ) + if "diagnose_plan" in features: + parser.add_argument( + "--diagnose_plan", + action="store_true", + help="Plan and print diagnostics without playing the trajectory.", + ) + if "grasp_sampling" in features: + parser.add_argument("--n_sample", type=int, default=10000) + parser.add_argument("--force_reannotate", action="store_true") + if "headless_play" in features: + parser.add_argument( + "--headless_play", + action="store_true", + help="Execute trajectories without opening the viewer window.", + ) + if "visualize_axes" in features: + parser.add_argument( + "--no_vis_eef_axis", + action="store_true", + help="Skip drawing target coordinate-frame markers.", + ) + return parser + def make_ur5_solver_cfg(tcp_z: float) -> URSolverCfg: """Create the UR5 arm solver cfg used by atomic-action tutorials.""" @@ -206,8 +268,10 @@ def get_hand_open_close_qpos( hand_limits = robot.get_qpos_limits(name=hand_control_part)[0].to( device=robot.device, dtype=torch.float32 ) - return hand_limits[:, 0], torch.minimum( - hand_limits[:, 1], torch.full_like(hand_limits[:, 1], close_qpos) + return hand_limits[:, 0], torch.clamp( + torch.full_like(hand_limits[:, 1], close_qpos), + min=hand_limits[:, 0], + max=hand_limits[:, 1], ) @@ -483,6 +547,22 @@ def replay_trajectory( stop_auto_play_recording(sim, recording_started) +def make_clear_dynamics_callback( + obj: RigidObject, + clear_after_step: int, +) -> Callable[[int, int], None]: + """Create a one-shot replay callback that freezes an attached object.""" + dynamics_cleared = False + + def clear_dynamics(step_idx: int, _: int) -> None: + nonlocal dynamics_cleared + if not dynamics_cleared and step_idx + 1 >= clear_after_step: + obj.clear_dynamics() + dynamics_cleared = True + + return clear_dynamics + + def get_tutorial_window_size(args: argparse.Namespace) -> tuple[int, int]: """Return the viewer window size used by atomic-action tutorials.""" return VIEWER_WIDTH, VIEWER_HEIGHT @@ -712,12 +792,14 @@ def create_ur5_gripper_robot_cfg( "GRIPPER_TCP_Z", "GRIPPER_URDF_PATH", "TOP_DOWN_EEF_ROTATION", + "TutorialCliFeature", "add_ur5_gripper_robot", "broadcast_pose_batch", "broadcast_waypoint_pose_batch", "clone_local_pose_from_first_env", "create_antipodal_semantics", "create_toppra_motion_generator", + "create_tutorial_argument_parser", "create_tutorial_simulation", "create_ur5_gripper_robot_cfg", "format_tensor", @@ -725,6 +807,7 @@ def create_ur5_gripper_robot_cfg( "initialize_pre_pick_robot_pose", "make_ur5_solver_cfg", "make_eef_pose_at", + "make_clear_dynamics_callback", "make_top_down_eef_pose", "get_tutorial_window_size", "prepare_tutorial_scene", diff --git a/tests/sim/atomic_actions/test_action_result_success.py b/tests/sim/atomic_actions/test_action_result_success.py deleted file mode 100644 index 5a5d684b8..000000000 --- a/tests/sim/atomic_actions/test_action_result_success.py +++ /dev/null @@ -1,49 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Tests for ActionResult.success tensor and ActionCfg.motion_source.""" - -from __future__ import annotations - -import pytest -import torch - -from embodichain.lab.sim.atomic_actions.core import ActionCfg, ActionResult, WorldState - - -class TestActionResultSuccess: - def test_success_all_tensor(self): - r = ActionResult( - success=torch.tensor([True, False]), - trajectory=torch.zeros(2, 0, 3), - next_state=WorldState(last_qpos=torch.zeros(2, 3)), - ) - assert r.success_all is False - - def test_bool_deprecation(self): - r = ActionResult( - success=torch.tensor([True, True]), - trajectory=torch.zeros(2, 0, 3), - next_state=WorldState(last_qpos=torch.zeros(2, 3)), - ) - with pytest.warns(DeprecationWarning): - assert bool(r) is True - - -class TestActionCfgMotionSource: - def test_defaults(self): - cfg = ActionCfg() - assert cfg.motion_source == "ik_interp" diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index 808d2c4fb..d0e8a2d6d 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -14,170 +14,265 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Tests for the concrete atomic action classes.""" +"""Tests for built-in atomic actions under the plan/invocation contract.""" from __future__ import annotations +from typing import TypeVar +from unittest.mock import Mock + import pytest import torch -from unittest.mock import Mock, patch -from embodichain.lab.sim.atomic_actions.affordance import ( - AntipodalAffordance, -) from embodichain.lab.sim.atomic_actions import ( - AssembleTarget, - 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, + ActionBinding, + ActionInvocation, + Affordance, + AntipodalAffordance, + AssembleGoal, AtomicAction, + AtomicActionEngine, + ControlPartCommandProfile, CoordinatedHeldObjectState, - HeldObjectState, - ObjectSemantics, - WorldState, -) -from embodichain.lab.sim.atomic_actions.actions import ( + CoordinatedPickGoal, CoordinatedPickment, - CoordinatedPickmentCfg, + CoordinatedPickmentOptions, CoordinatedPlacement, - CoordinatedPlacementCfg, + CoordinatedPlacementGoal, + CoordinatedPlacementOptions, + EndEffectorPoseGoal, + GraspGoal, + HandOver, + HandOverOptions, + HeldObjectPoseGoal, + HeldObjectState, + JointPositionGoal, + MotionPolicy, MoveEndEffector, - MoveEndEffectorCfg, - MoveJoints, - MoveJointsCfg, + MoveEndEffectorOptions, MoveHeldObject, - MoveHeldObjectCfg, + MoveHeldObjectOptions, + MoveJoints, + MoveJointsOptions, + ObjectSemantics, PickUp, - PickUpCfg, + PickUpOptions, Place, - PlaceCfg, + PlaceGoal, + PlaceOptions, + PlanningContext, Press, - PressCfg, + PressGoal, + PressOptions, + RobotObservation, + SceneSnapshot, + TaskState, ) +from embodichain.lab.sim.planners import MoveType NUM_ENVS = 2 ARM_DOF = 6 HAND_DOF = 2 -TOTAL_DOF = ARM_DOF + HAND_DOF -DUAL_ARM_DOF = 12 -DUAL_TOTAL_DOF = DUAL_ARM_DOF + 2 * HAND_DOF +ROBOT_DOF = ARM_DOF + HAND_DOF +DUAL_ARM_DOF = 2 * ARM_DOF +DUAL_ROBOT_DOF = DUAL_ARM_DOF + 2 * HAND_DOF + +ActionT = TypeVar("ActionT", bound=AtomicAction) + + +@pytest.fixture(autouse=True) +def _torch_interpolation(monkeypatch: pytest.MonkeyPatch) -> None: + """Use a small torch interpolation stand-in without initializing Warp.""" + + def interpolate( + trajectory: torch.Tensor, + interp_num: int, + device: torch.device, + ) -> torch.Tensor: + indices = torch.linspace( + 0, + trajectory.shape[1] - 1, + interp_num, + device=device, + ) + lower = indices.floor().to(torch.long) + upper = indices.ceil().to(torch.long) + weights = (indices - lower).view(1, -1, 1) + return torch.lerp(trajectory[:, lower], trajectory[:, upper], weights) + + monkeypatch.setattr( + "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", + interpolate, + ) -def _make_mock_robot(): +def _robot() -> Mock: robot = Mock() robot.device = torch.device("cpu") - robot.dof = TOTAL_DOF - - def get_qpos(name=None): - if name == "arm": + robot.dof = ROBOT_DOF + robot.control_parts = { + "arm": object(), + "hand": object(), + "alternate_arm": object(), + "alternate_hand": object(), + } + + def get_qpos(name: str | None = None) -> torch.Tensor: + if name in {"arm", "alternate_arm"}: return torch.zeros(NUM_ENVS, ARM_DOF) - if name == "hand": + if name in {"hand", "alternate_hand"}: return torch.zeros(NUM_ENVS, HAND_DOF) - return torch.zeros(NUM_ENVS, TOTAL_DOF) + return torch.zeros(NUM_ENVS, ROBOT_DOF) - robot.get_qpos = get_qpos - - def get_joint_ids(name=None): - if name == "arm": + def get_joint_ids(name: str | None = None) -> list[int]: + if name in {"arm", "alternate_arm"}: return list(range(ARM_DOF)) - if name == "hand": - return list(range(ARM_DOF, TOTAL_DOF)) - return list(range(TOTAL_DOF)) + if name in {"hand", "alternate_hand"}: + return list(range(ARM_DOF, ROBOT_DOF)) + return list(range(ROBOT_DOF)) + + def compute_ik( + pose: torch.Tensor | None = None, + name: str | None = None, + joint_seed: torch.Tensor | None = None, + **_: object, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert joint_seed is not None + return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed.clone() + + def compute_fk( + qpos: torch.Tensor | None = None, + name: str | None = None, + to_matrix: bool = True, + ) -> torch.Tensor: + count = NUM_ENVS if qpos is None else qpos.shape[0] + return torch.eye(4).repeat(count, 1, 1) + + robot.get_qpos.side_effect = get_qpos + robot.get_joint_ids.side_effect = get_joint_ids + robot.compute_ik.side_effect = compute_ik + robot.compute_fk.side_effect = compute_fk + return robot - robot.get_joint_ids = get_joint_ids - def compute_ik(pose=None, qpos_seed=None, name=None, joint_seed=None): - seed = joint_seed if joint_seed is not None else qpos_seed - if seed is None: - seed = torch.zeros(NUM_ENVS, ARM_DOF) - return torch.ones(NUM_ENVS, dtype=torch.bool), seed.clone() +def _motion_generator() -> Mock: + generator = Mock() + generator.robot = _robot() + generator.device = torch.device("cpu") + generator.planner.cfg.planner_type = "stub" + generator.planner.preserve_plan_samples = False + generator.planner.supports_move_type.return_value = False + return generator + + +def _bind_action( + generator: Mock, + action: ActionT, + control_profiles: dict[str, ControlPartCommandProfile] | None = None, +) -> ActionT: + """Bind one configured action to an engine-owned test backend.""" + profiles = { + name: ControlPartCommandProfile.joint_positions( + open=torch.zeros(len(generator.robot.get_joint_ids(name=name))), + grasp=torch.ones(len(generator.robot.get_joint_ids(name=name))), + ) + for name in generator.robot.control_parts + if "hand" in name + } + profiles.update({} if control_profiles is None else control_profiles) + engine = AtomicActionEngine( + generator, + control_profiles=profiles, + load_builtins=False, + ) + engine.register(action) + return action - robot.compute_ik = compute_ik - def compute_batch_ik(pose=None, name=None, joint_seed=None): - if joint_seed is not None: - return ( - torch.ones(joint_seed.shape[:2], dtype=torch.bool), - joint_seed.clone(), - ) - return torch.ones(NUM_ENVS, dtype=torch.bool), torch.zeros(NUM_ENVS, ARM_DOF) +def _plan_action( + action: AtomicAction, + invocation: ActionInvocation, + context: PlanningContext, +): + """Resolve a caller-owned invocation before calling the action planner.""" + return action.plan(action.resolve_request(invocation), context) - robot.compute_batch_ik = compute_batch_ik - def compute_fk(qpos=None, name=None, to_matrix=True): - n = qpos.shape[0] if qpos is not None else NUM_ENVS - return torch.eye(4).unsqueeze(0).repeat(n, 1, 1) +def _context(task: TaskState | None = None) -> PlanningContext: + qpos = torch.zeros(NUM_ENVS, ROBOT_DOF) + return PlanningContext( + robot=RobotObservation(timestamp=0.0, qpos=qpos, qvel=torch.zeros_like(qpos)), + task=task or TaskState.empty(batch_size=NUM_ENVS, device="cpu"), + scene=SceneSnapshot.empty(), + env_ids=torch.arange(NUM_ENVS), + ) - robot.compute_fk = compute_fk - return robot +def _binding() -> ActionBinding: + return ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ) -def _make_mock_motion_generator(): - mg = Mock() - mg.robot = _make_mock_robot() - mg.device = torch.device("cpu") - return mg +def _invocation( + skill_id: str, + goal, + *, + sample_count: int = 20, +) -> ActionInvocation: + return ActionInvocation( + skill_id=skill_id, + goal=goal, + binding=_binding(), + motion_policy=MotionPolicy(sample_count=sample_count), + ) -def _make_curobo_mock_motion_generator( - result_positions, success=None, preserve_plan_samples=True -): - """Mock MotionGenerator whose planner is a cuRobo backend. - - ``result_positions`` is ``(B, N, ARM_DOF)``. The planner preserves samples - and accepts both EEF and joint targets directly, matching the real cuRobo - capabilities. - ``preserve_plan_samples`` defaults to ``True`` to exercise the opt-in raw - path; pass ``False`` to exercise the default resample-to-sample_interval - path. - """ - mg = _make_mock_motion_generator() - planner = Mock() - planner.cfg.planner_type = "curobo" - planner.supported_move_types = frozenset({MoveType.EEF_MOVE, MoveType.JOINT_MOVE}) - planner.supports_move_type.side_effect = ( - lambda move_type: move_type in planner.supported_move_types + +def _semantics() -> ObjectSemantics: + entity = Mock() + entity.get_local_pose.return_value = torch.eye(4).repeat(NUM_ENVS, 1, 1) + return ObjectSemantics( + affordance=Affordance(), + geometry={}, + label="test_object", + entity=entity, ) - planner.preserve_plan_samples = preserve_plan_samples - mg.planner = planner - B = result_positions.shape[0] - if success is None: - success = torch.ones(B, dtype=torch.bool) - mg.generate.return_value = PlanResult(success=success, positions=result_positions) - return mg -def _make_dual_arm_mock_robot(): +def _held(semantics: ObjectSemantics | None = None) -> HeldObjectState: + poses = torch.eye(4).repeat(NUM_ENVS, 1, 1) + return HeldObjectState( + semantics=semantics or _semantics(), + object_to_eef=poses, + grasp_xpos=poses, + ) + + +def _dual_motion_generator() -> Mock: robot = Mock() robot.device = torch.device("cpu") - robot.dof = DUAL_TOTAL_DOF - - def get_qpos(name=None): - if name == "left_arm": - return torch.zeros(NUM_ENVS, ARM_DOF) - if name == "right_arm": + robot.dof = DUAL_ROBOT_DOF + robot.control_parts = { + "dual_arm": object(), + "left_arm": object(), + "right_arm": object(), + "left_hand": object(), + "right_hand": object(), + } + + def get_qpos(name: str | None = None) -> torch.Tensor: + if name in {"left_arm", "right_arm"}: return torch.zeros(NUM_ENVS, ARM_DOF) if name == "dual_arm": return torch.zeros(NUM_ENVS, DUAL_ARM_DOF) - if name in ("left_hand", "right_hand"): + if name in {"left_hand", "right_hand"}: return torch.zeros(NUM_ENVS, HAND_DOF) - return torch.zeros(NUM_ENVS, DUAL_TOTAL_DOF) - - robot.get_qpos = get_qpos + return torch.zeros(NUM_ENVS, DUAL_ROBOT_DOF) - def get_joint_ids(name=None): + def get_joint_ids(name: str | None = None) -> list[int]: if name == "left_arm": - return list(range(0, ARM_DOF)) + return list(range(ARM_DOF)) if name == "right_arm": return list(range(ARM_DOF, DUAL_ARM_DOF)) if name == "dual_arm": @@ -185,1268 +280,690 @@ def get_joint_ids(name=None): if name == "left_hand": return list(range(DUAL_ARM_DOF, DUAL_ARM_DOF + HAND_DOF)) if name == "right_hand": - return list(range(DUAL_ARM_DOF + HAND_DOF, DUAL_TOTAL_DOF)) - return list(range(DUAL_TOTAL_DOF)) - - robot.get_joint_ids = get_joint_ids - - def compute_ik(pose=None, name=None, joint_seed=None, qpos_seed=None): + return list(range(DUAL_ARM_DOF + HAND_DOF, DUAL_ROBOT_DOF)) + return list(range(DUAL_ROBOT_DOF)) + + def compute_ik( + pose: torch.Tensor | None = None, + name: str | None = None, + joint_seed: torch.Tensor | None = None, + qpos_seed: torch.Tensor | None = None, + **_: object, + ) -> tuple[torch.Tensor, torch.Tensor]: seed = joint_seed if joint_seed is not None else qpos_seed - if seed is None: - seed = torch.zeros(NUM_ENVS, ARM_DOF) + assert seed is not None offset = 0.1 if name == "left_arm" else 0.2 return torch.ones(seed.shape[0], dtype=torch.bool), seed + offset - robot.compute_ik = compute_ik - return robot - - -def _make_dual_arm_mock_motion_generator(): - mg = Mock() - mg.robot = _make_dual_arm_mock_robot() - mg.device = torch.device("cpu") - return mg - + def compute_fk( + qpos: torch.Tensor | None = None, + name: str | None = None, + to_matrix: bool = True, + ) -> torch.Tensor: + count = NUM_ENVS if qpos is None else qpos.shape[0] + return torch.eye(4).repeat(count, 1, 1) + + robot.get_qpos.side_effect = get_qpos + robot.get_joint_ids.side_effect = get_joint_ids + robot.compute_ik.side_effect = compute_ik + robot.compute_fk.side_effect = compute_fk + + generator = Mock() + generator.robot = robot + generator.device = torch.device("cpu") + generator.planner.cfg.planner_type = "stub" + generator.planner.preserve_plan_samples = False + generator.planner.supports_move_type.return_value = False + return generator + + +def _dual_context(task: TaskState | None = None) -> PlanningContext: + qpos = torch.zeros(NUM_ENVS, DUAL_ROBOT_DOF) + return PlanningContext( + robot=RobotObservation(0.0, qpos, torch.zeros_like(qpos)), + task=task or TaskState.empty(NUM_ENVS, "cpu"), + scene=SceneSnapshot.empty(), + env_ids=torch.arange(NUM_ENVS), + ) -def _hand_open(): - return torch.zeros(HAND_DOF, dtype=torch.float32) +def _dual_binding( + first_role: str, + second_role: str, +) -> ActionBinding: + return ActionBinding( + manipulators={ + first_role: "left_arm", + second_role: "right_arm", + }, + end_effectors={ + first_role: "left_hand", + second_role: "right_hand", + }, + ) -def _hand_close(): - return torch.full((HAND_DOF,), 0.025, dtype=torch.float32) +def test_builtin_descriptors_expose_goals_not_legacy_targets() -> None: + assert MoveEndEffector.GoalType is EndEffectorPoseGoal + assert MoveJoints.GoalType is JointPositionGoal + assert PickUp.GoalType is GraspGoal + assert MoveHeldObject.GoalType is HeldObjectPoseGoal + assert Place.GoalType == (PlaceGoal, AssembleGoal) + assert Press.GoalType is PressGoal + assert CoordinatedPickment.GoalType is CoordinatedPickGoal + assert CoordinatedPlacement.GoalType is CoordinatedPlacementGoal + assert HandOver.GoalType is GraspGoal + + +@pytest.mark.parametrize( + "options", + ( + PickUpOptions(), + MoveHeldObjectOptions(), + PlaceOptions(), + PressOptions(), + CoordinatedPickmentOptions(), + CoordinatedPlacementOptions(), + HandOverOptions(), + ), +) +def test_action_options_do_not_contain_embodiment_resources(options: object) -> None: + field_names = getattr(options, "__dataclass_fields__") -# --------------------------------------------------------------------------- -# MoveEndEffector -# --------------------------------------------------------------------------- + assert "control_part" not in field_names + assert not any(name.endswith("_control_part") for name in field_names) + assert not any(name.endswith("_qpos") for name in field_names) -class TestMoveEndEffectorAction: - def setup_method(self): - self.mg = _make_mock_motion_generator() +def test_joint_position_goal_rejects_unsupported_target_type() -> None: + with pytest.raises(TypeError, match="torch.Tensor or str"): + JointPositionGoal(target=1.0) # type: ignore[arg-type] - def test_target_type_is_pose_target(self): - assert MoveEndEffector.TargetType is EndEffectorPoseTarget - def test_default_name_is_explicit(self): - assert MoveEndEffectorCfg().name == "move_end_effector" +def test_joint_position_goal_rejects_empty_named_target() -> None: + with pytest.raises(ValueError, match="must not be empty"): + JointPositionGoal(target=" ") - def test_execute_returns_full_dof_trajectory(self): - action = MoveEndEffector(self.mg, MoveEndEffectorCfg(sample_interval=10)) - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - return_value=torch.zeros(NUM_ENVS, 10, ARM_DOF), - ): - state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)) - result = action.execute(EndEffectorPoseTarget(xpos=torch.eye(4)), state) - assert isinstance(result, ActionResult) - assert result.success.all() - assert result.success.shape == (NUM_ENVS,) - assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) - # 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)) - pose0 = torch.eye(4) - pose1 = torch.eye(4) - pose1[0, 3] = 1.0 - # (n_envs, n_waypoint, 4, 4) trajectory target - multi_xpos = ( - torch.stack([pose0, pose1], dim=0).unsqueeze(0).repeat(NUM_ENVS, 1, 1, 1) - ) - seen_poses = [] - - def compute_ik(pose=None, name=None, joint_seed=None, **kwargs): - seen_poses.append(pose.clone()) - return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed.clone() - - self.mg.robot.compute_ik = Mock(side_effect=compute_ik) - - captured = {} - - def interpolate(trajectory, interp_num, device): - captured["keyframes"] = trajectory - return trajectory[:, -1:, :].repeat(1, interp_num, 1) - - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - side_effect=interpolate, - ): - result = action.execute( - EndEffectorPoseTarget(xpos=multi_xpos), - WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), - ) - - assert result.success.all() - assert result.success.shape == (NUM_ENVS,) - assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) - # Two waypoints -> two IK calls, in order. - assert len(seen_poses) == 2 - assert torch.allclose(seen_poses[0], pose0.unsqueeze(0).repeat(NUM_ENVS, 1, 1)) - assert torch.allclose(seen_poses[1], pose1.unsqueeze(0).repeat(NUM_ENVS, 1, 1)) - # start prepended to the two IK solutions -> 3 keyframes. - assert captured["keyframes"].shape == (NUM_ENVS, 3, ARM_DOF) - - -# --------------------------------------------------------------------------- -# MoveJoints -# --------------------------------------------------------------------------- - - -class TestMoveJointsAction: - def setup_method(self): - self.mg = _make_mock_motion_generator() - - def test_target_type_accepts_explicit_and_named_joint_targets(self): - assert MoveJoints.TargetType == (JointPositionTarget, NamedJointPositionTarget) - - def test_default_name_is_explicit(self): - assert MoveJointsCfg().name == "move_joints" - - def test_execute_with_explicit_qpos_returns_full_dof_trajectory(self): - action = MoveJoints(self.mg, MoveJointsCfg(sample_interval=10)) - target_qpos = torch.full((ARM_DOF,), 0.5) - hand_qpos = torch.full((NUM_ENVS, HAND_DOF), 0.25) - last_qpos = torch.cat([torch.zeros(NUM_ENVS, ARM_DOF), hand_qpos], dim=1) - sem = ObjectSemantics( - affordance=AntipodalAffordance(), geometry={}, label="mug" - ) - held = HeldObjectState( - semantics=sem, - object_to_eef=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1), - grasp_xpos=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1), - ) +@pytest.mark.parametrize( + "target", + ( + torch.tensor(1.0), + torch.empty(0), + torch.zeros(1, 1, 1, 1), + ), +) +def test_joint_position_goal_rejects_invalid_tensor_shape( + target: torch.Tensor, +) -> None: + with pytest.raises(ValueError, match="Tensor target must have shape"): + JointPositionGoal(target=target) + + +def test_move_end_effector_returns_full_robot_timed_plan() -> None: + generator = _motion_generator() + action = _bind_action(generator, MoveEndEffector()) + context = _context() + + plan = _plan_action( + action, + _invocation( + "move_end_effector", + EndEffectorPoseGoal(torch.eye(4)), + sample_count=10, + ), + context, + ) - def interpolate(trajectory, interp_num, device): - assert trajectory.shape == (NUM_ENVS, 2, ARM_DOF) - return trajectory[:, -1:, :].repeat(1, interp_num, 1) - - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - side_effect=interpolate, - ): - result = action.execute( - JointPositionTarget(qpos=target_qpos), - WorldState(last_qpos=last_qpos, held_objects={"arm": held}), - ) - - assert result.success.all() - assert result.success.shape == (NUM_ENVS,) - assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) - assert torch.allclose(result.trajectory[:, -1, :ARM_DOF], target_qpos) - assert torch.allclose(result.trajectory[:, -1, ARM_DOF:], hand_qpos) - assert result.next_state.get_held_object("arm") is held - - def test_execute_with_named_qpos_resolves_cfg_target(self): - action = MoveJoints( - self.mg, - MoveJointsCfg( - sample_interval=8, - named_joint_positions={"home": torch.full((ARM_DOF,), 0.2)}, - ), - ) - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - side_effect=lambda trajectory, interp_num, device: trajectory[ - :, -1:, : - ].repeat(1, interp_num, 1), - ): - result = action.execute( - NamedJointPositionTarget(name="home"), - WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), - ) - assert result.success.all() - assert result.success.shape == (NUM_ENVS,) - assert torch.allclose( - result.next_state.last_qpos[:, :ARM_DOF], - torch.full((NUM_ENVS, ARM_DOF), 0.2), - ) + assert plan.plan_success.tolist() == [True, True] + assert plan.trajectory.positions.shape == (NUM_ENVS, 10, ROBOT_DOF) + assert plan.trajectory.duration.tolist() == pytest.approx([0.15, 0.15]) + assert plan.expected_effects.is_empty + + +def test_move_joints_uses_binding_and_preserves_uncontrolled_joints() -> None: + generator = _motion_generator() + named = {"ready": torch.full((ARM_DOF,), 0.4)} + action = _bind_action( + generator, + MoveJoints(), + control_profiles={ + "arm": ControlPartCommandProfile.joint_positions(**named), + }, + ) + qpos = torch.zeros(NUM_ENVS, ROBOT_DOF) + qpos[:, ARM_DOF:] = 0.7 + context = PlanningContext( + robot=RobotObservation(0.0, qpos, torch.zeros_like(qpos)), + task=TaskState.empty(NUM_ENVS, "cpu"), + scene=SceneSnapshot.empty(), + env_ids=torch.arange(NUM_ENVS), + ) - def test_execute_with_multi_waypoint_qpos_visits_each_waypoint(self): - action = MoveJoints(self.mg, MoveJointsCfg(sample_interval=10)) - # (n_envs, n_waypoint, control_dof) trajectory target - waypoint_qpos = ( - torch.stack( - [ - torch.full((ARM_DOF,), 0.3), - torch.full((ARM_DOF,), 0.7), - ], - dim=0, - ) - .unsqueeze(0) - .repeat(NUM_ENVS, 1, 1) - ) - last_qpos = torch.zeros(NUM_ENVS, TOTAL_DOF) - - captured = {} - - def interpolate(trajectory, interp_num, device): - captured["keyframes"] = trajectory - return trajectory[:, -1:, :].repeat(1, interp_num, 1) - - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - side_effect=interpolate, - ): - result = action.execute( - JointPositionTarget(qpos=waypoint_qpos), - WorldState(last_qpos=last_qpos), - ) - - assert result.success.all() - assert result.success.shape == (NUM_ENVS,) - assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) - # start prepended to the two waypoints -> 3 keyframes - keyframes = captured["keyframes"] - assert keyframes.shape == (NUM_ENVS, 3, ARM_DOF) - assert torch.allclose(keyframes[:, 0, :], torch.zeros(NUM_ENVS, ARM_DOF)) - assert torch.allclose(keyframes[:, 1, :], torch.full((NUM_ENVS, ARM_DOF), 0.3)) - assert torch.allclose(keyframes[:, 2, :], torch.full((NUM_ENVS, ARM_DOF), 0.7)) - # final state lands on the last waypoint - assert torch.allclose( - result.next_state.last_qpos[:, :ARM_DOF], - torch.full((NUM_ENVS, ARM_DOF), 0.7), - ) + plan = _plan_action( + action, + _invocation("move_joints", JointPositionGoal("ready"), sample_count=8), + context, + ) - def test_unknown_named_qpos_raises(self): - action = MoveJoints(self.mg, MoveJointsCfg()) - with pytest.raises(KeyError, match="missing"): - action.execute( - NamedJointPositionTarget(name="missing"), - WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), - ) + assert torch.allclose(plan.trajectory.positions[:, -1, :ARM_DOF], named["ready"]) + assert torch.all(plan.trajectory.positions[:, :, ARM_DOF:] == 0.7) -# --------------------------------------------------------------------------- -# PickUp -# --------------------------------------------------------------------------- +def test_pick_and_place_declare_effects_without_mutating_context() -> None: + generator = _motion_generator() + pick = _bind_action(generator, PickUp()) + initial = _context() + semantics = _semantics() + grasp = torch.eye(4).repeat(NUM_ENVS, 1, 1) + pick_plan = _plan_action( + pick, + _invocation("pick_up", GraspGoal(semantics=semantics, grasp_xpos=grasp)), + initial, + ) + picked_task = pick_plan.expected_effects.apply(initial.task, pick_plan.plan_success) -class TestPickUpAction: - def setup_method(self): - self.mg = _make_mock_motion_generator() + assert initial.task.get_held_object("arm") is None + assert picked_task.get_held_object("arm") is not None - def test_target_type_is_grasp_target(self): - assert PickUp.TargetType is GraspTarget + place = _bind_action(generator, Place()) + picked_context = PlanningContext( + robot=initial.robot, + task=picked_task, + scene=initial.scene, + env_ids=initial.env_ids, + ) + place_plan = _plan_action( + place, + _invocation("place", PlaceGoal(torch.eye(4))), + picked_context, + ) + placed_task = place_plan.expected_effects.apply( + picked_task, place_plan.plan_success + ) - def test_approach_alignment_filter_is_opt_in(self): - assert PickUpCfg().approach_alignment_max_angle is None + assert picked_task.get_held_object("arm") is not None + assert placed_task.get_held_object("arm") is None - def test_execute_populates_held_object_state(self): - cfg = PickUpCfg( - hand_open_qpos=_hand_open(), - hand_close_qpos=_hand_close(), - sample_interval=20, - hand_interp_steps=4, - ) - action = PickUp(self.mg, cfg) - - # Fake affordance that returns a single identity grasp pose. - affordance = AntipodalAffordance() - affordance.get_valid_grasp_poses = Mock( - return_value=[ - (torch.eye(4).unsqueeze(0), torch.tensor([0.5])) - for _ in range(NUM_ENVS) - ] - ) - entity = Mock() - entity.get_local_pose = Mock( - return_value=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1) - ) +def test_move_held_object_requires_projected_attachment() -> None: + generator = _motion_generator() + action = _bind_action(generator, MoveHeldObject()) + invocation = _invocation( + "move_held_object", + HeldObjectPoseGoal(torch.eye(4)), + sample_count=10, + ) - sem = ObjectSemantics( - affordance=affordance, - geometry={}, - label="mug", - entity=entity, - ) + with pytest.raises(ValueError, match="requires an object held"): + _plan_action(action, invocation, _context()) - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - side_effect=lambda trajectory, interp_num, device: torch.zeros( - NUM_ENVS, interp_num, ARM_DOF - ), - ): - state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)) - result = action.execute(GraspTarget(semantics=sem), state) - assert result.success.all() - assert result.success.shape == (NUM_ENVS,) - assert result.trajectory.shape[0] == NUM_ENVS - assert result.trajectory.shape[2] == TOTAL_DOF - 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( - self.mg, - PickUpCfg( - hand_open_qpos=_hand_open(), - hand_close_qpos=_hand_close(), - sample_interval=12, - hand_interp_steps=4, - ), - ) - entity = Mock() - entity.get_local_pose.return_value = ( - torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1) - ) - affordance = Mock() - grasp_xpos = torch.eye(4) - - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - return_value=torch.zeros(NUM_ENVS, 4, ARM_DOF), - ): - result = action.execute( - GraspTarget( - semantics=ObjectSemantics( - affordance=affordance, - geometry={}, - entity=entity, - ), - grasp_xpos=grasp_xpos, - ), - WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), - ) - - assert result.success.all() - held_object = result.next_state.get_held_object("arm") - assert held_object is not None - assert torch.allclose( - held_object.grasp_xpos, - grasp_xpos.unsqueeze(0).repeat(NUM_ENVS, 1, 1), - ) - affordance.get_valid_grasp_poses.assert_not_called() - - def test_execute_chooses_symmetric_grasp_variant_closest_to_start_pose(self): - cfg = PickUpCfg( - hand_open_qpos=_hand_open(), - hand_close_qpos=_hand_close(), - sample_interval=20, - hand_interp_steps=4, - ) - action = PickUp(self.mg, cfg) - compute_batch_ik = self.mg.robot.compute_batch_ik - self.mg.robot.compute_batch_ik = Mock(side_effect=compute_batch_ik) - - rz_pi_grasp = torch.eye(4) - rz_pi_grasp[:3, :3] = torch.diag(torch.tensor([-1.0, -1.0, 1.0])) - affordance = AntipodalAffordance() - affordance.get_valid_grasp_poses = Mock( - return_value=[ - (rz_pi_grasp.unsqueeze(0), torch.tensor([0.5])) for _ in range(NUM_ENVS) - ] - ) + held = _held() + task = TaskState( + batch_size=NUM_ENVS, + device="cpu", + held_objects={"arm": held}, + ) + plan = _plan_action(action, invocation, _context(task)) + assert plan.plan_success.all() + assert plan.expected_effects.is_empty - entity = Mock() - entity.get_local_pose = Mock( - return_value=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1) - ) - sem = ObjectSemantics( - affordance=affordance, - geometry={}, - label="mug", - entity=entity, - ) - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - side_effect=lambda trajectory, interp_num, device: torch.zeros( - NUM_ENVS, interp_num, ARM_DOF - ), - ): - state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)) - result = action.execute(GraspTarget(semantics=sem), state) +def test_press_uses_invocation_sample_budget() -> None: + generator = _motion_generator() + action = _bind_action(generator, Press()) - assert result.success.all() - assert result.success.shape == (NUM_ENVS,) - 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(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) - assert call.kwargs["joint_seed"].shape == (NUM_ENVS, 2, ARM_DOF) + plan = _plan_action( + action, + _invocation("press", PressGoal(torch.eye(4)), sample_count=12), + _context(), + ) + assert plan.plan_success.tolist() == [True, True] + assert plan.trajectory.waypoint_count == 12 + assert plan.expected_effects.is_empty -# --------------------------------------------------------------------------- -# MoveHeldObject -# --------------------------------------------------------------------------- +def test_motion_source_and_sample_count_are_not_action_config_fields() -> None: + with pytest.raises(TypeError): + MoveEndEffectorOptions(motion_source="motion_gen") # type: ignore[call-arg] + with pytest.raises(TypeError): + MoveJointsOptions(sample_interval=10) # type: ignore[call-arg] -class TestMoveHeldObjectAction: - def setup_method(self): - self.mg = _make_mock_motion_generator() - def test_target_type_is_held_object_target(self): - assert MoveHeldObject.TargetType is HeldObjectPoseTarget +def test_move_joints_rejects_binding_with_wrong_goal_skill() -> None: + generator = _motion_generator() + action = _bind_action(generator, MoveJoints()) + invocation = ActionInvocation( + skill_id="move_end_effector", + goal=JointPositionGoal(torch.zeros(ARM_DOF)), + binding=ActionBinding(manipulators={"primary": "arm"}), + ) + with pytest.raises(ValueError, match="skill_id"): + action.resolve_request(invocation) - def test_default_name_is_explicit(self): - assert ( - MoveHeldObjectCfg(hand_close_qpos=_hand_close()).name == "move_held_object" - ) - def test_legacy_upright_configuration_is_retained(self): - upright_direction = torch.tensor([0.0, 0.0, 1.0]) - cfg = MoveHeldObjectCfg( - hand_close_qpos=_hand_close(), - obj_upright_direction=upright_direction, - pick_rotate_upright=torch.pi / 2, - ) +def test_planner_timing_is_preserved_in_simple_action() -> None: + generator = _motion_generator() + generator.planner.cfg.planner_type = "toppra" + generator.planner.supports_move_type.side_effect = ( + lambda move_type: move_type is MoveType.JOINT_MOVE + ) + generator.generate.return_value.success = torch.ones(NUM_ENVS, dtype=torch.bool) + generator.generate.return_value.positions = torch.ones(NUM_ENVS, 3, ARM_DOF) + generator.generate.return_value.velocities = torch.full((NUM_ENVS, 3, ARM_DOF), 0.5) + generator.generate.return_value.accelerations = torch.zeros(NUM_ENVS, 3, ARM_DOF) + generator.generate.return_value.dt = torch.tensor([[0.0, 0.1, 0.2]]).repeat( + NUM_ENVS, 1 + ) + generator.generate.return_value.duration = torch.full((NUM_ENVS,), 0.3) + action = _bind_action(generator, MoveJoints()) + invocation = ActionInvocation( + skill_id="move_joints", + goal=JointPositionGoal(torch.ones(ARM_DOF)), + binding=ActionBinding(manipulators={"primary": "arm"}), + motion_policy=MotionPolicy(motion_source="motion_gen", sample_count=3), + ) - assert torch.equal(cfg.obj_upright_direction, upright_direction) - assert cfg.pick_rotate_upright == torch.pi / 2 + plan = _plan_action(action, invocation, _context()) + + assert plan.trajectory.duration.tolist() == pytest.approx([0.3, 0.3]) + assert plan.trajectory.velocities is not None + assert torch.all(plan.trajectory.velocities[:, :, :ARM_DOF] == 0.5) + assert torch.all(plan.trajectory.velocities[:, :, ARM_DOF:] == 0.0) + + +def test_move_end_effector_visits_batched_waypoints_in_order() -> None: + generator = _motion_generator() + solved_poses: list[torch.Tensor] = [] + + def compute_ik( + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + **_: object, + ) -> tuple[torch.Tensor, torch.Tensor]: + solved_poses.append(pose.clone()) + return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed + 0.1 + + generator.robot.compute_ik.side_effect = compute_ik + waypoints = torch.eye(4).reshape(1, 1, 4, 4).repeat(NUM_ENVS, 2, 1, 1) + waypoints[:, 0, 0, 3] = 0.1 + waypoints[:, 1, 0, 3] = 0.3 + + action = _bind_action(generator, MoveEndEffector()) + plan = _plan_action( + action, + _invocation( + "move_end_effector", + EndEffectorPoseGoal(waypoints), + sample_count=9, + ), + _context(), + ) - def test_requires_held_object_in_state(self): - cfg = MoveHeldObjectCfg( - hand_close_qpos=_hand_close(), - sample_interval=10, - ) - action = MoveHeldObject(self.mg, cfg) - state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)) - with pytest.raises(Exception): - action.execute(HeldObjectPoseTarget(object_target_pose=torch.eye(4)), state) - - def test_preserves_held_object(self): - cfg = MoveHeldObjectCfg( - hand_close_qpos=_hand_close(), - sample_interval=10, - ) - action = MoveHeldObject(self.mg, cfg) - sem = ObjectSemantics( - affordance=AntipodalAffordance(), geometry={}, label="mug" - ) - held = HeldObjectState( - semantics=sem, - object_to_eef=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1), - grasp_xpos=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1), - ) - 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), - ): - result = action.execute( - HeldObjectPoseTarget(object_target_pose=torch.eye(4)), state - ) - assert result.success.all() - assert result.success.shape == (NUM_ENVS,) - assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) - assert result.next_state.get_held_object("arm") is held - - def test_automatic_rotation_adjustment_is_isolated_per_environment(self): - action = MoveHeldObject( - self.mg, - MoveHeldObjectCfg( - hand_close_qpos=_hand_close(), - sample_interval=10, - ), - ) - upward_tcp = torch.eye(4) - downward_tcp = torch.eye(4) - downward_tcp[:3, :3] = torch.diag(torch.tensor([1.0, -1.0, -1.0])) - action.robot.compute_fk = Mock( - return_value=torch.stack([upward_tcp, downward_tcp]) - ) - action.builder.plan_arm_traj = Mock( - return_value=( - torch.ones(NUM_ENVS, dtype=torch.bool), - torch.zeros(NUM_ENVS, 10, ARM_DOF), - ) - ) - semantics = ObjectSemantics( - affordance=AntipodalAffordance(), geometry={}, label="mug" - ) - held = HeldObjectState( - semantics=semantics, - 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), - ) + assert plan.plan_success.all() + assert len(solved_poses) == 2 + assert solved_poses[0][:, 0, 3].tolist() == pytest.approx([0.1, 0.1]) + assert solved_poses[1][:, 0, 3].tolist() == pytest.approx([0.3, 0.3]) - result = action.execute( - HeldObjectPoseTarget(object_target_pose=torch.eye(4)), - WorldState( - last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF), - held_objects={"arm": held}, - ), - ) - assert result.success.all() - target_states = action.builder.plan_arm_traj.call_args.args[0] - assert not torch.allclose(target_states[0][0].xpos[:3, :3], torch.eye(3)) - assert torch.allclose(target_states[1][0].xpos[:3, :3], torch.eye(3)) - - -# --------------------------------------------------------------------------- -# Place -# --------------------------------------------------------------------------- - - -class TestPlaceAction: - def setup_method(self): - self.mg = _make_mock_motion_generator() - - def test_target_type_is_pose_target(self): - assert PlaceTarget in Place.TargetType - assert AssembleTarget in Place.TargetType - - def test_rejects_non_positive_cartesian_waypoint_count(self): - with pytest.raises(Exception, match="cartesian_waypoint_count"): - Place( - self.mg, - PlaceCfg( - hand_open_qpos=_hand_open(), - hand_close_qpos=_hand_close(), - cartesian_waypoint_count=0, - ), - ) - - def test_execute_clears_held_object(self): - cfg = PlaceCfg( - hand_open_qpos=_hand_open(), - hand_close_qpos=_hand_close(), - sample_interval=20, - hand_interp_steps=4, - ) - action = Place(self.mg, cfg) - sem = ObjectSemantics( - affordance=AntipodalAffordance(), geometry={}, label="mug" - ) - held = HeldObjectState( - semantics=sem, - object_to_eef=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1), - grasp_xpos=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1), - ) - 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(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.get_held_object("arm") is None - - def test_execute_with_multi_waypoint_visits_each_waypoint(self): - cfg = PlaceCfg( - hand_open_qpos=_hand_open(), - hand_close_qpos=_hand_close(), - sample_interval=20, - hand_interp_steps=4, - lift_height=0.1, - ) - action = Place(self.mg, cfg) - sem = ObjectSemantics( - affordance=AntipodalAffordance(), geometry={}, label="mug" - ) - held = HeldObjectState( - semantics=sem, - object_to_eef=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1), - grasp_xpos=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1), - ) - state = WorldState( - last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF), - held_objects={"arm": held}, - ) - - pose0 = torch.eye(4) - pose1 = torch.eye(4) - pose1[0, 3] = 1.0 - # (n_envs, n_waypoint, 4, 4) trajectory target - multi_xpos = ( - torch.stack([pose0, pose1], dim=0).unsqueeze(0).repeat(NUM_ENVS, 1, 1, 1) - ) - seen_poses = [] - - def compute_ik(pose=None, name=None, joint_seed=None, **kwargs): - seen_poses.append(pose.clone()) - return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed.clone() - - self.mg.robot.compute_ik = Mock(side_effect=compute_ik) - - captured = {} - - def interpolate(trajectory, interp_num, device): - # Only the down phase carries more than 2 keyframes; capture it. - if trajectory.shape[1] > 2: - captured["down_keyframes"] = trajectory - return trajectory[:, -1:, :].repeat(1, interp_num, 1) - - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - side_effect=interpolate, - ): - result = action.execute(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.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 - approach = pose0.clone() - approach[2, 3] += lift_height - retract = pose1.clone() - retract[2, 3] += lift_height - assert torch.allclose( - seen_poses[0], approach.unsqueeze(0).repeat(NUM_ENVS, 1, 1) - ) - assert torch.allclose(seen_poses[1], pose0.unsqueeze(0).repeat(NUM_ENVS, 1, 1)) - assert torch.allclose(seen_poses[2], pose1.unsqueeze(0).repeat(NUM_ENVS, 1, 1)) - assert torch.allclose( - seen_poses[3], retract.unsqueeze(0).repeat(NUM_ENVS, 1, 1) - ) - # start prepended to the 3 down-phase IK solutions -> 4 keyframes. - assert captured["down_keyframes"].shape == (NUM_ENVS, 4, ARM_DOF) - - @pytest.mark.parametrize( - ("release_z", "expected_lifted_z"), - [(0.7, 0.8), (0.85, 0.85)], +def test_move_joints_visits_waypoints_and_rejects_unknown_names() -> None: + generator = _motion_generator() + action = _bind_action( + generator, + MoveJoints(), + control_profiles={ + "arm": ControlPartCommandProfile.joint_positions(ready=torch.zeros(ARM_DOF)) + }, + ) + waypoints = torch.stack( + [ + torch.full((NUM_ENVS, ARM_DOF), 0.3), + torch.full((NUM_ENVS, ARM_DOF), 0.7), + ], + dim=1, ) - def test_caps_approach_and_retract_world_z_without_descending( - self, - release_z, - expected_lifted_z, - ): - action = Place( - self.mg, - PlaceCfg( - hand_open_qpos=_hand_open(), - hand_close_qpos=_hand_close(), - sample_interval=20, - hand_interp_steps=4, - lift_height=0.15, - max_approach_retract_z=0.8, - ), - ) - release_pose = torch.eye(4) - release_pose[2, 3] = release_z - seen_poses = [] - - def compute_ik(pose=None, name=None, joint_seed=None, **kwargs): - seen_poses.append(pose.clone()) - return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed.clone() - - self.mg.robot.compute_ik = Mock(side_effect=compute_ik) - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - side_effect=lambda trajectory, interp_num, device: trajectory[ - :, -1:, : - ].repeat(1, interp_num, 1), - ): - result = action.execute( - PlaceTarget(xpos=release_pose), - WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), - ) - - assert result.success.all() - assert [pose[0, 2, 3].item() for pose in seen_poses] == pytest.approx( - [expected_lifted_z, release_z, expected_lifted_z] - ) - - def test_cartesian_waypoints_hold_target_rotation_during_translation(self): - waypoint_count = 3 - action = Place( - self.mg, - PlaceCfg( - hand_open_qpos=_hand_open(), - hand_close_qpos=_hand_close(), - sample_interval=24, - hand_interp_steps=4, - lift_height=0.1, - cartesian_waypoint_count=waypoint_count, - ), - ) - target = torch.eye(4) - target[:3, :3] = torch.tensor( - [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]] - ) - target[:3, 3] = torch.tensor([0.3, 0.0, 0.2]) - seen_poses = [] - - def compute_ik(pose=None, name=None, joint_seed=None, **kwargs): - seen_poses.append(pose.clone()) - return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed.clone() - - self.mg.robot.compute_ik = Mock(side_effect=compute_ik) - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - side_effect=lambda trajectory, interp_num, device: trajectory[ - :, -1:, : - ].repeat(1, interp_num, 1), - ): - result = action.execute( - PlaceTarget(xpos=target), - WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), - ) - - assert result.success.all() - assert len(seen_poses) == 3 * waypoint_count - expected_rotation = target[:3, :3].unsqueeze(0).repeat(NUM_ENVS, 1, 1) - for pose in seen_poses: - assert torch.allclose(pose[:, :3, :3], expected_rotation) - - def test_execute_preserves_release_pose_without_tcp_symmetry(self): - cfg = PlaceCfg( - hand_open_qpos=_hand_open(), - hand_close_qpos=_hand_close(), - sample_interval=20, - hand_interp_steps=4, - ) - action = Place(self.mg, cfg) - state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)) - rz_pi_pose = torch.eye(4) - rz_pi_pose[:3, :3] = torch.diag(torch.tensor([-1.0, -1.0, 1.0])) - seen_poses = [] - - def compute_ik(pose=None, name=None, joint_seed=None, **kwargs): - seen_poses.append(pose.clone()) - return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed.clone() - - def repeat_last_keyframe(trajectory, interp_num, device): - return trajectory[:, -1:, :].repeat(1, interp_num, 1) - - self.mg.robot.compute_ik = Mock(side_effect=compute_ik) - - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - side_effect=repeat_last_keyframe, - ): - result = action.execute(PlaceTarget(xpos=rz_pi_pose), state) - - assert result.success.all() - assert result.success.shape == (NUM_ENVS,) - assert len(seen_poses) == 3 - assert torch.allclose( - seen_poses[1], rz_pi_pose.unsqueeze(0).repeat(NUM_ENVS, 1, 1) - ) - def test_execute_with_tcp_symmetry_selects_closest_release_variant(self): - cfg = PlaceCfg( - hand_open_qpos=_hand_open(), - hand_close_qpos=_hand_close(), - sample_interval=20, - hand_interp_steps=4, - lift_height=0.1, - ) - action = Place(self.mg, cfg) - state = WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)) - rz_pi_pose = torch.eye(4) - rz_pi_pose[:3, :3] = torch.diag(torch.tensor([-1.0, -1.0, 1.0])) - seen_poses = [] - - def compute_ik(pose=None, name=None, joint_seed=None, **kwargs): - seen_poses.append(pose.clone()) - return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed.clone() - - def repeat_last_keyframe(trajectory, interp_num, device): - return trajectory[:, -1:, :].repeat(1, interp_num, 1) - - self.mg.robot.compute_ik = Mock(side_effect=compute_ik) - - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - side_effect=repeat_last_keyframe, - ): - result = action.execute( - PlaceTarget( - xpos=rz_pi_pose, - tcp_symmetry="z_roll_180", - ), - state, - ) - - assert result.success.all() - assert result.success.shape == (NUM_ENVS,) - assert len(seen_poses) == 3 - expected_release = torch.eye(4) - expected_retract = torch.eye(4) - expected_retract[2, 3] += cfg.lift_height - assert torch.allclose( - seen_poses[0], expected_retract.unsqueeze(0).repeat(NUM_ENVS, 1, 1) - ) - assert torch.allclose( - seen_poses[1], expected_release.unsqueeze(0).repeat(NUM_ENVS, 1, 1) - ) - assert torch.allclose( - seen_poses[2], expected_retract.unsqueeze(0).repeat(NUM_ENVS, 1, 1) - ) + plan = _plan_action( + action, + _invocation( + "move_joints", + JointPositionGoal(waypoints), + sample_count=7, + ), + _context(), + ) + assert torch.allclose(plan.trajectory.positions[:, 3, :ARM_DOF], waypoints[:, 0]) + assert torch.allclose(plan.trajectory.positions[:, -1, :ARM_DOF], waypoints[:, 1]) + with pytest.raises(KeyError, match="has no command"): + _plan_action( + action, + _invocation("move_joints", JointPositionGoal("missing")), + _context(), + ) + + +def test_pick_explicit_grasp_bypasses_sampling_and_records_grasp() -> None: + generator = _motion_generator() + affordance = AntipodalAffordance() + affordance.get_valid_grasp_poses = Mock() + entity = Mock() + entity.get_local_pose.return_value = torch.eye(4).repeat(NUM_ENVS, 1, 1) + semantics = ObjectSemantics( + affordance=affordance, + geometry={}, + label="explicit-grasp-object", + entity=entity, + ) + grasp = torch.eye(4).repeat(NUM_ENVS, 1, 1) + grasp[:, 0, 3] = torch.tensor([0.1, 0.2]) + action = _bind_action(generator, PickUp()) + + plan = _plan_action( + action, + _invocation( + "pick_up", + GraspGoal(semantics=semantics, grasp_xpos=grasp), + sample_count=20, + ), + _context(), + ) + projected = plan.expected_effects.apply(_context().task, plan.plan_success) + + affordance.get_valid_grasp_poses.assert_not_called() + held = projected.get_held_object("arm") + assert held is not None + assert torch.allclose(held.grasp_xpos, grasp) + + +def test_pick_uses_binding_control_part_as_effect_resource() -> None: + generator = _motion_generator() + action = _bind_action(generator, PickUp()) + invocation = ActionInvocation( + skill_id="pick_up", + goal=GraspGoal(semantics=_semantics(), grasp_xpos=torch.eye(4)), + binding=ActionBinding( + manipulators={"primary": "alternate_arm"}, + end_effectors={"primary": "alternate_hand"}, + ), + motion_policy=MotionPolicy(sample_count=20), + ) + context = _context() -# --------------------------------------------------------------------------- -# Press -# --------------------------------------------------------------------------- + plan = _plan_action(action, invocation, context) + projected = plan.expected_effects.apply(context.task, plan.plan_success) + assert projected.get_held_object("alternate_arm") is not None + assert projected.get_held_object("arm") is None -class TestPressAction: - def setup_method(self): - self.mg = _make_mock_motion_generator() - def test_target_type_is_press_target(self): - assert Press.TargetType is PressTarget +def test_press_closes_hand_without_changing_projected_attachment() -> None: + held = _held() + task = TaskState( + batch_size=NUM_ENVS, + device="cpu", + held_objects={"arm": held}, + ) + generator = _motion_generator() + action = _bind_action( + generator, + Press(default_options=PressOptions(hand_interp_steps=4)), + ) - def test_default_name_is_explicit(self): - assert PressCfg(hand_close_qpos=_hand_close()).name == "press" + plan = _plan_action( + action, + _invocation("press", PressGoal(torch.eye(4)), sample_count=12), + _context(task), + ) + projected = plan.expected_effects.apply(task, plan.plan_success) + + assert torch.all(plan.trajectory.positions[:, -1, ARM_DOF:] == 1.0) + projected_held = projected.get_held_object("arm") + assert projected_held is not None + assert projected_held.semantics is held.semantics + assert torch.equal(projected_held.object_to_eef, held.object_to_eef) + + +def test_handover_does_not_mutate_cached_final_pose() -> None: + generator = _dual_motion_generator() + handover_options = HandOverOptions( + middle_object_pose=torch.eye(4), + final_object_pose=torch.eye(4), + hand_interp_steps=4, + hold_steps=2, + retreat_steps=5, + ) + action = _bind_action( + generator, + HandOver(default_options=handover_options), + ) + assert handover_options.final_object_pose is not None + original_final_pose = handover_options.final_object_pose.clone() + current_object_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + current_object_pose[:, :3, :3] = torch.diag(torch.tensor([-1.0, -1.0, 1.0])) + semantics = _semantics() + semantics.entity.get_local_pose.return_value = current_object_pose + task = TaskState( + batch_size=NUM_ENVS, + device="cpu", + held_objects={"left_arm": _held(semantics)}, + ) + receive_grasp = torch.eye(4).repeat(NUM_ENVS, 1, 1) + action._resolve_receive_grasp = Mock( + return_value=(receive_grasp, torch.ones(NUM_ENVS, dtype=torch.bool)) + ) - def test_execute_closes_hand_and_preserves_held_object(self): - cfg = PressCfg( - hand_close_qpos=_hand_close(), - sample_interval=12, - hand_interp_steps=4, - ) - action = Press(self.mg, cfg) - sem = ObjectSemantics( - affordance=AntipodalAffordance(), geometry={}, label="mug" - ) - held = HeldObjectState( - semantics=sem, - object_to_eef=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1), - grasp_xpos=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1), - ) - start_hand_qpos = torch.full((NUM_ENVS, HAND_DOF), 0.01) - last_qpos = torch.cat([torch.zeros(NUM_ENVS, ARM_DOF), start_hand_qpos], dim=1) - state = WorldState(last_qpos=last_qpos, held_objects={"arm": held}) - - def interpolate(trajectory, interp_num, device): - return trajectory[:, -1:, :].repeat(1, interp_num, 1) - - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - side_effect=interpolate, - ): - result = action.execute(PressTarget(xpos=torch.eye(4)), state) - - assert result.success.all() - assert result.success.shape == (NUM_ENVS,) - assert result.trajectory.shape == (NUM_ENVS, 12, TOTAL_DOF) - expected_hand_qpos = _hand_close().unsqueeze(0).repeat(NUM_ENVS, 1) - assert torch.allclose(result.trajectory[:, -1, ARM_DOF:], expected_hand_qpos) - assert torch.allclose( - result.next_state.last_qpos[:, :ARM_DOF], - last_qpos[:, :ARM_DOF], - ) - assert result.next_state.get_held_object("arm") is held + def plan_from_start( + control_part: str, + start_qpos: torch.Tensor, + target_poses: torch.Tensor, + n_waypoints: int, + motion_policy: MotionPolicy, + ) -> tuple[bool, torch.Tensor]: + return True, start_qpos.unsqueeze(1).repeat(1, n_waypoints, 1) + + action._plan_named_arm_trajectory = Mock(side_effect=plan_from_start) + invocation = ActionInvocation( + skill_id="hand_over", + goal=GraspGoal(semantics=semantics), + binding=_dual_binding("source", "destination"), + motion_policy=MotionPolicy(sample_count=30), + ) + plan = _plan_action(action, invocation, _dual_context(task)) -# --------------------------------------------------------------------------- -# CoordinatedPickment -# --------------------------------------------------------------------------- + assert plan.plan_success.all() + assert torch.equal(handover_options.final_object_pose, original_final_pose) -class TestCoordinatedPickmentAction: - def setup_method(self): - self.mg = _make_dual_arm_mock_motion_generator() +def test_coordinated_pick_returns_full_dof_plan_and_projected_relation() -> None: + generator = _dual_motion_generator() + action = _bind_action( + generator, + CoordinatedPickment( + default_options=CoordinatedPickmentOptions( + hand_interp_steps=4, + hold_steps=2, + object_motion_keyframes=3, + ), + ), + ) + semantics = ObjectSemantics( + affordance=AntipodalAffordance(), geometry={}, label="coordinated-object" + ) + invocation = ActionInvocation( + skill_id="coordinated_pickment", + goal=CoordinatedPickGoal( + semantics=semantics, + object_target_pose=torch.eye(4), + left_object_to_eef=torch.eye(4), + right_object_to_eef=torch.eye(4), + object_initial_pose=torch.eye(4), + ), + binding=_dual_binding("left", "right"), + motion_policy=MotionPolicy(sample_count=30), + ) + context = _dual_context() + + plan = _plan_action(action, invocation, context) + projected = plan.expected_effects.apply(context.task, plan.plan_success) + + assert plan.plan_success.tolist() == [True, True] + assert plan.trajectory.positions.shape == (NUM_ENVS, 30, DUAL_ROBOT_DOF) + assert projected.get_held_object("left_arm") is None + assert projected.get_held_object("right_arm") is None + assert isinstance( + projected.get_coordinated_held_object("left_arm", "right_arm"), + CoordinatedHeldObjectState, + ) - def test_target_type_is_coordinated_pickment_target(self): - assert CoordinatedPickment.TargetType is CoordinatedPickTarget - assert CoordinatedPickment.__bases__ == (AtomicAction,) - def test_execute_returns_full_dof_trajectory_and_dual_held_state(self): - cfg = CoordinatedPickmentCfg( - left_hand_open_qpos=_hand_open(), - left_hand_close_qpos=_hand_close(), - right_hand_open_qpos=_hand_open(), - right_hand_close_qpos=_hand_close(), - sample_interval=30, - hand_interp_steps=4, - hold_steps=2, - object_motion_keyframes=3, - ) - action = CoordinatedPickment(self.mg, cfg) - sem = ObjectSemantics( - affordance=AntipodalAffordance(), geometry={}, label="pencil" - ) - state = WorldState(last_qpos=torch.zeros(NUM_ENVS, DUAL_TOTAL_DOF)) - result = action.execute( - CoordinatedPickTarget( - semantics=sem, - object_target_pose=torch.eye(4), - left_object_to_eef=torch.eye(4), - right_object_to_eef=torch.eye(4), - object_initial_pose=torch.eye(4), - ), - state, - ) - assert result.success.all() - assert result.success.shape == (NUM_ENVS,) - assert result.trajectory.shape == (NUM_ENVS, 30, DUAL_TOTAL_DOF) - assert torch.allclose( - result.trajectory[:, -1, action.left_hand_joint_ids], - _hand_close().unsqueeze(0).repeat(NUM_ENVS, 1), - ) - assert torch.allclose( - result.trajectory[:, -1, action.right_hand_joint_ids], - _hand_close().unsqueeze(0).repeat(NUM_ENVS, 1), - ) - held_object = result.next_state.get_coordinated_held_object( - "left_arm", "right_arm" - ) - assert isinstance(held_object, CoordinatedHeldObjectState) - assert result.next_state.held_objects == {} - - def test_execute_freezes_only_environment_with_partial_ik_failure(self): - action = CoordinatedPickment( - self.mg, - CoordinatedPickmentCfg( - left_hand_open_qpos=_hand_open(), - left_hand_close_qpos=_hand_close(), - right_hand_open_qpos=_hand_open(), - right_hand_close_qpos=_hand_close(), - sample_interval=30, +def test_coordinated_pick_holds_only_environment_with_ik_failure() -> None: + generator = _dual_motion_generator() + original_compute_ik = generator.robot.compute_ik.side_effect + + def fail_second_environment( + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + **kwargs: object, + ) -> tuple[torch.Tensor, torch.Tensor]: + success, qpos = original_compute_ik( + pose=pose, + name=name, + joint_seed=joint_seed, + **kwargs, + ) + if name == "right_arm" and float(pose[1, 0, 3]) > 0.15: + success = success.clone() + success[1] = False + return success, qpos + + generator.robot.compute_ik.side_effect = fail_second_environment + action = _bind_action( + generator, + CoordinatedPickment( + default_options=CoordinatedPickmentOptions( hand_interp_steps=4, hold_steps=2, object_motion_keyframes=3, ), - ) - original_compute_ik = self.mg.robot.compute_ik - - def fail_second_env_during_move( - pose=None, name=None, joint_seed=None, qpos_seed=None - ): - success, qpos = original_compute_ik( - pose=pose, - name=name, - joint_seed=joint_seed, - qpos_seed=qpos_seed, - ) - if name == "right_arm" and float(pose[1, 0, 3]) > 0.15: - success = success.clone() - success[1] = False - return success, qpos - - self.mg.robot.compute_ik = fail_second_env_during_move - target_pose = torch.eye(4) - target_pose[0, 3] = 0.3 - state = WorldState(last_qpos=torch.zeros(NUM_ENVS, DUAL_TOTAL_DOF)) - semantics = ObjectSemantics( - affordance=AntipodalAffordance(), geometry={}, label="tray" - ) - - result = action.execute( - CoordinatedPickTarget( - semantics=semantics, - object_target_pose=target_pose, - left_object_to_eef=torch.eye(4), - right_object_to_eef=torch.eye(4), - object_initial_pose=torch.eye(4), + ), + ) + target_pose = torch.eye(4) + target_pose[0, 3] = 0.3 + invocation = ActionInvocation( + skill_id="coordinated_pickment", + goal=CoordinatedPickGoal( + semantics=ObjectSemantics( + affordance=AntipodalAffordance(), geometry={}, label="tray" ), - state, - ) - - assert result.success.tolist() == [True, False] - assert not torch.allclose(result.trajectory[0], state.last_qpos[0]) - assert torch.allclose( - result.trajectory[1], - state.last_qpos[1].unsqueeze(0).repeat(30, 1), - ) - assert torch.allclose(result.next_state.last_qpos[1], state.last_qpos[1]) - - -# --------------------------------------------------------------------------- -# CoordinatedPlacement -# --------------------------------------------------------------------------- - - -class TestCoordinatedPlacementAction: - def setup_method(self): - self.mg = _make_dual_arm_mock_motion_generator() - self.cfg = CoordinatedPlacementCfg( - placing_hand_open_qpos=_hand_open(), - placing_hand_close_qpos=_hand_close(), - support_hand_close_qpos=_hand_close(), - sample_interval=30, - hand_interp_steps=4, - hold_steps=3, - retreat_steps=5, - lift_height=0.08, - ) - self.action = CoordinatedPlacement(self.mg, cfg=self.cfg) - - def _make_target_and_state( - self, - ) -> tuple[ - CoordinatedPlacementTarget, - WorldState, - HeldObjectState, - HeldObjectState, - ]: - placing_pose = torch.eye(4) - placing_pose[0, 3] = 0.2 - support_pose = torch.eye(4) - support_pose[0, 3] = 0.2 - support_pose[2, 3] = -0.05 - - placing_object_to_eef = torch.eye(4) - placing_object_to_eef[2, 3] = 0.12 - support_object_to_eef = torch.eye(4) - support_object_to_eef[2, 3] = 0.10 - - placing_semantics = ObjectSemantics( - affordance=AntipodalAffordance(), geometry={}, label="placing" - ) - support_semantics = ObjectSemantics( - affordance=AntipodalAffordance(), geometry={}, label="support" - ) - 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, - ) - 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 + object_target_pose=target_pose, + left_object_to_eef=torch.eye(4), + right_object_to_eef=torch.eye(4), + object_initial_pose=torch.eye(4), + ), + binding=_dual_binding("left", "right"), + motion_policy=MotionPolicy(sample_count=30), + ) + context = _dual_context() - def test_target_type_is_coordinated_placement_target(self): - assert CoordinatedPlacement.TargetType is CoordinatedPlacementTarget + plan = _plan_action(action, invocation, context) + projected = plan.expected_effects.apply(context.task, plan.plan_success) - def test_init_sets_dual_arm_and_hand_joint_ids(self): - assert self.action.dual_arm_joint_ids == list(range(DUAL_ARM_DOF)) - assert self.action.placing_arm_joint_ids == list(range(ARM_DOF)) - assert self.action.support_arm_joint_ids == list(range(ARM_DOF, DUAL_ARM_DOF)) - assert self.action.placing_hand_joint_ids == list( - range(DUAL_ARM_DOF, DUAL_ARM_DOF + HAND_DOF) - ) - assert self.action.support_hand_joint_ids == list( - range(DUAL_ARM_DOF + HAND_DOF, DUAL_TOTAL_DOF) - ) - assert self.action.joint_ids == list(range(DUAL_TOTAL_DOF)) - - def test_resolve_target_composes_object_and_tcp_poses(self): - target, 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 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( - 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 - assert segments["approach"] >= 2 - assert segments["release"] == self.cfg.hand_interp_steps - assert segments["retreat"] == self.cfg.retreat_steps - - def test_execute_returns_full_dof_and_final_hand_states(self): - target, state, _, support_source = self._make_target_and_state() - - def interpolate(trajectory, interp_num, device): - weights = torch.linspace( - 0.0, - 1.0, - steps=interp_num, - dtype=trajectory.dtype, - device=trajectory.device, - ) - return torch.lerp( - trajectory[:, :1], - trajectory[:, -1:], - weights.view(1, -1, 1), - ) - - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - side_effect=interpolate, - ): - result = self.action.execute(target, state) - - assert result.success is True - assert result.trajectory.shape == ( - NUM_ENVS, - self.cfg.sample_interval, - DUAL_TOTAL_DOF, - ) - assert torch.allclose( - result.trajectory[:, -1, self.action.placing_hand_joint_ids], - _hand_open().unsqueeze(0).repeat(NUM_ENVS, 1), - ) - assert torch.allclose( - result.trajectory[:, -1, self.action.support_hand_joint_ids], - _hand_close().unsqueeze(0).repeat(NUM_ENVS, 1), - ) - assert result.next_state.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) + assert plan.plan_success.tolist() == [True, False] + assert not torch.allclose(plan.trajectory.positions[0], context.robot.qpos[0]) + assert torch.allclose( + plan.trajectory.positions[1], + context.robot.qpos[1].unsqueeze(0).repeat(30, 1), + ) + held = projected.get_coordinated_held_object("left_arm", "right_arm") + assert held is not None + assert held.env_mask.tolist() == [True, False] -# --------------------------------------------------------------------------- -# MoveJoints + cuRobo motion_gen routing -# --------------------------------------------------------------------------- +def test_coordinated_placement_projects_release_and_support_attachment() -> None: + generator = _dual_motion_generator() + action = _bind_action( + generator, + CoordinatedPlacement( + default_options=CoordinatedPlacementOptions( + hand_interp_steps=4, + hold_steps=3, + retreat_steps=5, + ), + ), + ) + placing = _held( + ObjectSemantics(affordance=AntipodalAffordance(), geometry={}, label="placing") + ) + support = _held( + ObjectSemantics(affordance=AntipodalAffordance(), geometry={}, label="support") + ) + task = TaskState( + batch_size=NUM_ENVS, + device="cpu", + held_objects={"left_arm": placing, "right_arm": support}, + ) + invocation = ActionInvocation( + skill_id="coordinated_placement", + goal=CoordinatedPlacementGoal( + placing_object_target_pose=torch.eye(4), + support_object_target_pose=torch.eye(4), + ), + binding=_dual_binding("placing", "support"), + motion_policy=MotionPolicy(sample_count=30), + ) + context = _dual_context(task) + plan = _plan_action(action, invocation, context) + projected = plan.expected_effects.apply(context.task, plan.plan_success) -class TestMoveJointsCurobo: - def setup_method(self): - # shared per-test mg is created in each test (result shapes differ). - pass + assert plan.plan_success.tolist() == [True, True] + assert plan.trajectory.positions.shape == (NUM_ENVS, 30, DUAL_ROBOT_DOF) + assert projected.get_held_object("left_arm") is None + assert projected.get_held_object("right_arm") is not None + assert projected.get_held_object("right_arm").semantics is support.semantics - def _action(self, mg, **cfg_kw): - return MoveJoints( - mg, - MoveJointsCfg(motion_source="motion_gen", **cfg_kw), - ) - def test_one_waypoint_routes_joint_move_to_motion_gen(self): - mg = _make_curobo_mock_motion_generator( - result_positions=torch.zeros(NUM_ENVS, 5, ARM_DOF) - ) - action = self._action(mg, sample_interval=10) - result = action.execute( - JointPositionTarget(qpos=torch.full((ARM_DOF,), 0.5)), - WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), - ) - assert result.success.tolist() == [True, True] - # With preserve_plan_samples=True (opt-in), cuRobo's raw length (5) is - # returned unchanged rather than resampled to sample_interval (10). - assert result.trajectory.shape == (NUM_ENVS, 5, TOTAL_DOF) - # Full-DoF preservation: hand joints stay at the inherited state (zeros). - assert torch.allclose( - result.trajectory[:, :, ARM_DOF:], torch.zeros(NUM_ENVS, 5, HAND_DOF) - ) - plan_states = mg.generate.call_args.args[0] - assert all(s.move_type is MoveType.JOINT_MOVE for s in plan_states) - # The builder requests target preparation; MotionGenerator skips it - # because cuRobo declares native JOINT_MOVE support. - assert mg.generate.call_args.kwargs["options"].is_interpolate is True - - def test_default_resamples_to_sample_interval(self): - mg = _make_curobo_mock_motion_generator( - result_positions=torch.zeros(NUM_ENVS, 5, ARM_DOF), - preserve_plan_samples=False, - ) - action = self._action(mg, sample_interval=10) - with patch( - "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", - return_value=torch.zeros(NUM_ENVS, 10, ARM_DOF), - ) as interp: - result = action.execute( - JointPositionTarget(qpos=torch.full((ARM_DOF,), 0.5)), - WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), - ) - assert result.success.tolist() == [True, True] - # Default preserve_plan_samples=False resamples cuRobo's raw length (5) - # up to the action's sample_interval (10). - assert interp.call_count == 1 - assert interp.call_args.kwargs["interp_num"] == 10 - assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) - - def test_multi_waypoint_routes_ordered_joint_states(self): - mg = _make_curobo_mock_motion_generator( - result_positions=torch.zeros(NUM_ENVS, 5, ARM_DOF) - ) - action = self._action(mg, sample_interval=10) - waypoint_qpos = ( - torch.stack( - [torch.full((ARM_DOF,), 0.3), torch.full((ARM_DOF,), 0.7)], dim=0 - ) - .unsqueeze(0) - .repeat(NUM_ENVS, 1, 1) - ) - result = action.execute( - JointPositionTarget(qpos=waypoint_qpos), - WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), - ) - assert result.success.tolist() == [True, True] - assert result.trajectory.shape == (NUM_ENVS, 5, TOTAL_DOF) - plan_states = mg.generate.call_args.args[0] - assert len(plan_states) == 2 - assert all(s.move_type is MoveType.JOINT_MOVE for s in plan_states) - # Ordered: first waypoint, then second. - assert torch.allclose(plan_states[0].qpos, torch.full((NUM_ENVS, ARM_DOF), 0.3)) - assert torch.allclose(plan_states[1].qpos, torch.full((NUM_ENVS, ARM_DOF), 0.7)) - - def test_failure_holds_start_qpos(self): - positions = torch.zeros(NUM_ENVS, 5, ARM_DOF) - positions[1] = 1.0 # env 1 "would move" but is marked failed - mg = _make_curobo_mock_motion_generator( - result_positions=positions, success=torch.tensor([True, False]) - ) - action = self._action(mg, sample_interval=10) - last_qpos = torch.zeros(NUM_ENVS, TOTAL_DOF) - last_qpos[1, :ARM_DOF] = 0.7 # env 1 start - result = action.execute( - JointPositionTarget(qpos=torch.full((ARM_DOF,), 0.5)), - WorldState(last_qpos=last_qpos), - ) - assert result.success.tolist() == [True, False] - # Failed env held at its start arm qpos across all samples. - assert torch.allclose( - result.trajectory[1, :, :ARM_DOF], torch.full((5, ARM_DOF), 0.7) - ) +def test_coordinated_actions_reject_curobo_motion_generation() -> None: + generator = _dual_motion_generator() + generator.planner.cfg.planner_type = "curobo" + policy = MotionPolicy(motion_source="motion_gen", sample_count=30) + pick = _bind_action( + generator, + CoordinatedPickment(), + ) + pick_invocation = ActionInvocation( + skill_id="coordinated_pickment", + goal=CoordinatedPickGoal( + semantics=ObjectSemantics( + affordance=AntipodalAffordance(), geometry={}, label="object" + ), + object_target_pose=torch.eye(4), + left_object_to_eef=torch.eye(4), + right_object_to_eef=torch.eye(4), + object_initial_pose=torch.eye(4), + ), + binding=_dual_binding("left", "right"), + motion_policy=policy, + ) + with pytest.raises(ValueError, match="not supported"): + _plan_action(pick, pick_invocation, _dual_context()) -class TestCoordinatedRejectsCurobo: - def test_coordinated_pickment_rejects_curobo(self): - mg = _make_dual_arm_mock_motion_generator() - mg.planner.cfg.planner_type = "curobo" - cfg = CoordinatedPickmentCfg( - left_hand_open_qpos=_hand_open(), - left_hand_close_qpos=_hand_close(), - right_hand_open_qpos=_hand_open(), - right_hand_close_qpos=_hand_close(), - motion_source="motion_gen", - ) - with pytest.raises(ValueError, match="not supported"): - CoordinatedPickment(mg, cfg) - - def test_coordinated_placement_rejects_curobo(self): - mg = _make_dual_arm_mock_motion_generator() - mg.planner.cfg.planner_type = "curobo" - cfg = CoordinatedPlacementCfg( - placing_hand_open_qpos=_hand_open(), - placing_hand_close_qpos=_hand_close(), - support_hand_close_qpos=_hand_close(), - motion_source="motion_gen", - ) - with pytest.raises(ValueError, match="not supported"): - CoordinatedPlacement(mg, cfg) + placement = _bind_action( + generator, + CoordinatedPlacement(), + ) + placement_invocation = ActionInvocation( + skill_id="coordinated_placement", + goal=CoordinatedPlacementGoal(torch.eye(4), torch.eye(4)), + binding=_dual_binding("placing", "support"), + motion_policy=policy, + ) + with pytest.raises(ValueError, match="not supported"): + _plan_action(placement, placement_invocation, _dual_context()) diff --git a/tests/sim/atomic_actions/test_control.py b/tests/sim/atomic_actions/test_control.py new file mode 100644 index 000000000..ee056aeaa --- /dev/null +++ b/tests/sim/atomic_actions/test_control.py @@ -0,0 +1,137 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for control-part semantic command profiles and invocation overrides.""" + +from __future__ import annotations + +from unittest.mock import Mock + +import pytest +import torch + +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionControlOverrides, + ActionPlanningServices, + ControlPartCommandProfile, + JointPositionCommand, +) + + +def _services() -> ActionPlanningServices: + robot = Mock() + robot.device = torch.device("cpu") + robot.control_parts = {"arm": object(), "hand": object()} + robot.get_joint_ids.side_effect = lambda name: ( + [0, 1, 2] if name == "arm" else [3, 4] + ) + generator = Mock() + generator.robot = robot + generator.device = torch.device("cpu") + generator.planner.cfg.planner_type = "stub" + return ActionPlanningServices( + generator, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=torch.zeros(2), + grasp=torch.ones(2), + ) + }, + ) + + +def test_joint_position_command_broadcasts_owned_batch() -> None: + source = torch.tensor([0.1, 0.2]) + command = JointPositionCommand(source) + source.fill_(9.0) + + resolved = command.resolve(n_envs=3, control_dof=2, device="cpu") + resolved[0].fill_(7.0) + + assert torch.allclose(resolved[1:], torch.tensor([[0.1, 0.2], [0.1, 0.2]])) + assert torch.allclose(command.positions, torch.tensor([0.1, 0.2])) + + +def test_joint_position_command_rejects_incompatible_control_part() -> None: + command = JointPositionCommand(torch.zeros(2)) + + with pytest.raises(ValueError, match="resolved control part has 3"): + command.resolve(n_envs=1, control_dof=3, device="cpu") + + +def test_control_profile_is_resolved_from_robot_control_part() -> None: + resolved = _services().resolve_binding( + ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ) + ) + + grasp = resolved.end_effector().joint_positions( + "grasp", + n_envs=2, + device="cpu", + ) + + assert grasp.tolist() == [[1.0, 1.0], [1.0, 1.0]] + with pytest.raises(KeyError, match="Available commands"): + resolved.end_effector().joint_positions( + "pinch", + n_envs=2, + device="cpu", + ) + + +def test_invocation_override_replaces_only_resolved_role_snapshot() -> None: + services = _services() + override_source = torch.full((2,), 0.4) + overrides = ActionControlOverrides( + end_effectors={ + "primary": {"grasp": JointPositionCommand(override_source)}, + } + ) + override_source.fill_(8.0) + binding = ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ) + + overridden = services.resolve_binding(binding, overrides) + base = services.resolve_binding(binding) + overrides.end_effectors["primary"]["grasp"].positions.fill_(6.0) # type: ignore[attr-defined] + + assert torch.allclose( + overridden.end_effector().joint_positions("grasp", n_envs=1, device="cpu"), + torch.full((1, 2), 0.4), + ) + assert torch.equal( + base.end_effector().joint_positions("grasp", n_envs=1, device="cpu"), + torch.ones(1, 2), + ) + + +def test_override_rejects_role_not_present_in_binding() -> None: + services = _services() + binding = ActionBinding(end_effectors={"primary": "hand"}) + overrides = ActionControlOverrides( + end_effectors={ + "destination": {"open": JointPositionCommand(torch.zeros(2))}, + } + ) + + with pytest.raises(KeyError, match="unbound end effector roles"): + services.resolve_binding(binding, overrides) diff --git a/tests/sim/atomic_actions/test_core.py b/tests/sim/atomic_actions/test_core.py index 9fabb3699..574b49b8e 100644 --- a/tests/sim/atomic_actions/test_core.py +++ b/tests/sim/atomic_actions/test_core.py @@ -14,321 +14,197 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Tests for atomic-action target contracts and shared core state.""" +"""Tests for atomic-action goals, state, policies, effects, and plans.""" from __future__ import annotations -import dataclasses -from typing import get_args +from dataclasses import FrozenInstanceError 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 import ( - BuiltinTarget, - CoordinatedPickTarget, - CoordinatedPickmentTarget, - CoordinatedPlacementTarget, - EndEffectorPoseTarget, - GraspTarget, - HeldObjectPoseTarget, - JointPositionTarget, - NamedJointPositionTarget, - ObjectActionTarget, - PlaceTarget, - PressTarget, -) -from embodichain.lab.sim.atomic_actions.core import ( - ActionTarget, - ActionCfg, - ActionResult, - CoordinatedHeldObjectState, + ActionBinding, + ActionInvocation, + Affordance, + EndEffectorPoseGoal, + EntityState, HeldObjectState, + MotionPolicy, ObjectSemantics, - WorldState, + PlanningContext, + RecoveryPolicy, + RobotObservation, + SceneEntityPose, + SceneSnapshot, + StateDelta, + TaskState, + TimedTrajectory, +) +from embodichain.lab.sim.atomic_actions.goals import ( + collect_scene_dependencies, + resolve_pose_goal, ) -class TestTypedTargets: - def test_core_does_not_own_concrete_target_types(self): - assert not hasattr(core_module, "GraspTarget") +def _semantics(label: str = "object") -> ObjectSemantics: + return ObjectSemantics(affordance=Affordance(), geometry={}, label=label) - 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 _held(batch_size: int = 2) -> HeldObjectState: + pose = torch.eye(4).repeat(batch_size, 1, 1) + return HeldObjectState( + semantics=_semantics(), + object_to_eef=pose, + grasp_xpos=pose, + ) - 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 _context(scene: SceneSnapshot | None = None) -> PlanningContext: + qpos = torch.zeros(2, 4) + return PlanningContext( + robot=RobotObservation(timestamp=1.0, qpos=qpos, qvel=torch.zeros_like(qpos)), + task=TaskState.empty(batch_size=2, device="cpu"), + scene=scene or SceneSnapshot.empty(), + env_ids=torch.tensor([4, 7], dtype=torch.long), ) - 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 - - 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_place_target_rejects_unknown_tcp_symmetry(self): - with pytest.raises(ValueError, match="tcp_symmetry"): - 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): - t.xpos = torch.zeros(4, 4) # type: ignore[misc] - - def test_joint_position_target_holds_qpos(self): - qpos = torch.zeros(6) - assert JointPositionTarget(qpos=qpos).qpos is qpos - - def test_joint_position_target_is_frozen(self): - t = JointPositionTarget(qpos=torch.zeros(6)) - with pytest.raises(dataclasses.FrozenInstanceError): - t.qpos = torch.ones(6) # type: ignore[misc] - - def test_named_joint_position_target_holds_name(self): - assert NamedJointPositionTarget(name="home").name == "home" - - def test_named_joint_position_target_is_frozen(self): - t = NamedJointPositionTarget(name="home") - with pytest.raises(dataclasses.FrozenInstanceError): - t.name = "ready" # type: ignore[misc] - - def test_grasp_target_holds_semantics(self): - sem = ObjectSemantics(affordance=Affordance(), geometry={}, label="mug") - 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") - t = GraspTarget(semantics=sem) - with pytest.raises(dataclasses.FrozenInstanceError): - t.semantics = ObjectSemantics( # type: ignore[misc] - affordance=Affordance(), geometry={}, label="other" - ) - - def test_held_object_target_holds_pose(self): - x = torch.eye(4) - assert HeldObjectPoseTarget(object_target_pose=x).object_target_pose is x - - def test_held_object_target_is_frozen(self): - t = HeldObjectPoseTarget(object_target_pose=torch.eye(4)) - with pytest.raises(dataclasses.FrozenInstanceError): - t.object_target_pose = torch.zeros(4, 4) # type: ignore[misc] - - def test_coordinated_pick_target_holds_object_offsets(self): - sem = ObjectSemantics(affordance=Affordance(), geometry={}, label="pencil") - target = CoordinatedPickTarget( - semantics=sem, - object_target_pose=torch.eye(4), - left_object_to_eef=torch.eye(4), - right_object_to_eef=torch.eye(4), - ) - 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_only_holds_desired_state(self): - target = CoordinatedPlacementTarget( - placing_object_target_pose=torch.eye(4), - support_object_target_pose=torch.eye(4), - ) - assert isinstance(target, ActionTarget) - assert not hasattr(target, "placing_held_object") - assert target.support_object_target_pose.shape == (4, 4) - - -class TestObjectSemantics: - def test_does_not_mutate_affordance_geometry(self): - # The redesign removes the __post_init__ aliasing footgun. - aff = Affordance() - geometry = {"bounding_box": [0.1, 0.1, 0.1]} - ObjectSemantics(affordance=aff, geometry=geometry, label="mug") - # affordance should not have a geometry attribute, or if it does it should - # NOT be the same object as the semantics' geometry dict. - assert getattr(aff, "geometry", None) is not geometry - - def test_sets_object_label_on_affordance(self): - aff = Affordance() - ObjectSemantics(affordance=aff, geometry={}, label="mug") - assert aff.object_label == "mug" - - def test_default_optional_fields(self): - sem = ObjectSemantics(affordance=Affordance(), geometry={}) - assert sem.label == "none" - assert sem.properties == {} - assert sem.entity is None - - -class TestHeldObjectState: - def test_required_fields(self): - sem = ObjectSemantics(affordance=Affordance(), geometry={}) - s = HeldObjectState( - semantics=sem, - object_to_eef=torch.eye(4).unsqueeze(0), - grasp_xpos=torch.eye(4).unsqueeze(0), - ) - assert s.semantics is sem - assert s.object_to_eef.shape == (1, 4, 4) - assert s.grasp_xpos.shape == (1, 4, 4) - - -class TestCoordinatedHeldObjectState: - def test_required_fields(self): - sem = ObjectSemantics(affordance=Affordance(), geometry={}) - s = CoordinatedHeldObjectState( - semantics=sem, - left_object_to_eef=torch.eye(4).unsqueeze(0), - right_object_to_eef=torch.eye(4).unsqueeze(0), - left_grasp_xpos=torch.eye(4).unsqueeze(0), - right_grasp_xpos=torch.eye(4).unsqueeze(0), - ) - assert s.semantics is sem - assert s.left_object_to_eef.shape == (1, 4, 4) - assert s.right_grasp_xpos.shape == (1, 4, 4) - - -class TestWorldState: - def test_constructs_with_last_qpos_only(self): - qpos = torch.zeros(2, 6) - ws = WorldState(last_qpos=qpos) - assert ws.last_qpos is qpos - assert ws.held_objects == {} - assert ws.coordinated_held_objects == {} - - def test_carries_held_object(self): - sem = ObjectSemantics(affordance=Affordance(), geometry={}) - held = HeldObjectState( - semantics=sem, - object_to_eef=torch.eye(4).unsqueeze(0), - grasp_xpos=torch.eye(4).unsqueeze(0), - ) - ws = WorldState( - last_qpos=torch.zeros(1, 6), - held_objects={"left_arm": held}, + + +def test_action_binding_is_role_based_and_immutable() -> None: + binding = ActionBinding( + manipulators={"primary": "left_arm"}, + end_effectors={"primary": "left_hand"}, + ) + + assert binding.manipulator() == "left_arm" + assert binding.end_effector() == "left_hand" + with pytest.raises(TypeError): + binding.manipulators["primary"] = "right_arm" + with pytest.raises(KeyError, match="destination"): + binding.manipulator("destination") + + +def test_invocation_rejects_values_without_goal_contract() -> None: + with pytest.raises(TypeError, match="goal_kind"): + ActionInvocation( + skill_id="move_end_effector", + goal=object(), # type: ignore[arg-type] + binding=ActionBinding(manipulators={"primary": "arm"}), ) - 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={}) - held = CoordinatedHeldObjectState( - semantics=sem, - left_object_to_eef=torch.eye(4).unsqueeze(0), - right_object_to_eef=torch.eye(4).unsqueeze(0), - left_grasp_xpos=torch.eye(4).unsqueeze(0), - right_grasp_xpos=torch.eye(4).unsqueeze(0), + + +def test_motion_and_recovery_policy_validate_shared_parameters() -> None: + policy = MotionPolicy(sample_count=24, control_dt=0.01) + assert policy.sample_count == 24 + assert policy.control_dt == 0.01 + with pytest.raises(ValueError, match="sample_count"): + MotionPolicy(sample_count=1) + with pytest.raises(ValueError, match="max_replans"): + RecoveryPolicy(max_replans=-1) + + +def test_task_state_normalizes_held_relations_and_masks_updates() -> None: + held = _held() + state = TaskState( + batch_size=2, + device="cpu", + held_objects={"left_arm": held}, + ) + replacement = _held() + updated = StateDelta( + held_object_updates={ + "left_arm": None, + "right_arm": replacement, + } + ).apply(state, torch.tensor([True, False])) + + left = updated.get_held_object("left_arm") + right = updated.get_held_object("right_arm") + assert left is not None and left.env_mask.tolist() == [False, True] + assert right is not None and right.env_mask.tolist() == [True, False] + assert state.get_held_object("right_arm") is None + + +def test_robot_observation_owns_input_tensors() -> None: + qpos = torch.zeros(2, 4) + observation = RobotObservation( + timestamp=0.0, + qpos=qpos, + qvel=torch.zeros_like(qpos), + ) + qpos.fill_(1.0) + assert torch.count_nonzero(observation.qpos) == 0 + with pytest.raises(FrozenInstanceError): + observation.timestamp = 2.0 + + +def test_scene_entity_pose_is_resolved_late_from_snapshot() -> None: + entity_pose = torch.eye(4).repeat(2, 1, 1) + entity_pose[:, 0, 3] = torch.tensor([0.2, 0.4]) + offset = torch.eye(4) + offset[2, 3] = 0.1 + reference = SceneEntityPose("cup", relative_pose=offset) + context = _context( + SceneSnapshot( + timestamp=1.0, + version=3, + entities={"cup": EntityState(entity_pose, confidence=0.9)}, ) - ws = WorldState( - last_qpos=torch.zeros(1, 14), - coordinated_held_objects={("left_arm", "right_arm"): held}, + ) + + resolved = resolve_pose_goal(reference, context, name="xpos") + + assert resolved[:, 0, 3].tolist() == pytest.approx([0.2, 0.4]) + assert resolved[:, 2, 3].tolist() == pytest.approx([0.1, 0.1]) + assert collect_scene_dependencies(EndEffectorPoseGoal(reference)) == ("cup",) + + +def test_scene_entity_pose_enforces_confidence() -> None: + context = _context( + SceneSnapshot( + timestamp=1.0, + version=1, + entities={"cup": EntityState(torch.eye(4), confidence=0.2)}, ) - 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), + ) + with pytest.raises(ValueError, match="confidence"): + resolve_pose_goal( + SceneEntityPose("cup", minimum_confidence=0.8), + context, + name="xpos", ) - assert ws.held_objects == {} - - -class TestActionResult: - def test_shape_contract(self): - traj = torch.zeros(2, 10, 8) - ws = WorldState(last_qpos=torch.zeros(2, 8)) - res = ActionResult(success=True, trajectory=traj, next_state=ws) - assert res.success is True - assert res.trajectory.shape == (2, 10, 8) - assert res.next_state is ws - - -class TestActionCfg: - def test_defaults(self): - cfg = ActionCfg() - assert cfg.name == "default" - assert cfg.control_part == "arm" - assert cfg.interpolation_type == "linear" - assert cfg.velocity_limit is None - assert cfg.acceleration_limit is None + + +def test_timed_trajectory_synthesizes_timing_and_holds_selected_rows() -> None: + positions = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4) + trajectory = TimedTrajectory.from_positions( + positions, + env_ids=torch.tensor([4, 7]), + control_dt=0.02, + ) + held = trajectory.hold_rows( + torch.tensor([True, False]), + torch.full((2, 4), -1.0), + ) + + assert trajectory.duration.tolist() == pytest.approx([0.04, 0.04]) + assert torch.equal(held.positions[0], positions[0]) + assert torch.all(held.positions[1] == -1.0) + + +def test_timed_trajectory_concatenates_metadata() -> None: + first = TimedTrajectory.from_positions( + torch.zeros(2, 2, 4), + env_ids=torch.tensor([0, 1]), + control_dt=0.1, + ) + second = TimedTrajectory.from_positions( + torch.ones(2, 3, 4), + env_ids=torch.tensor([0, 1]), + control_dt=0.2, + ) + + result = TimedTrajectory.concatenate((first, second)) + + assert result.positions.shape == (2, 5, 4) + assert result.duration.tolist() == pytest.approx([0.5, 0.5]) 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 bb12e25c5..67c0a15dc 100644 --- a/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py +++ b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py @@ -46,12 +46,11 @@ CuroboWorldCfg, ) from embodichain.lab.sim.atomic_actions import ( # noqa: E402 + ActionBinding, + ActionInvocation, AtomicActionEngine, - EndEffectorPoseTarget, -) -from embodichain.lab.sim.atomic_actions.actions import ( # noqa: E402 - MoveEndEffector, - MoveEndEffectorCfg, + EndEffectorPoseGoal, + MotionPolicy, ) ROBOT_UID = "curobo_franka" @@ -88,17 +87,6 @@ def _make_franka_curobo_engine(): ) ) engine = AtomicActionEngine(mg) - engine.register( - MoveEndEffector( - mg, - MoveEndEffectorCfg( - motion_source="motion_gen", - control_part=CONTROL_PART, - sample_interval=SAMPLE_INTERVAL, - ), - ), - name="move_end_effector", - ) return sim, robot, engine @@ -140,11 +128,22 @@ def test_atomic_move_end_effector_uses_curobo_v2(): sim, robot, engine = _make_franka_curobo_engine() try: target = _reachable_target_beyond_demo_block(robot) - success, trajectory, _ = engine.run( - [("move_end_effector", EndEffectorPoseTarget(xpos=target))] + result = engine.compile( + ( + ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal(xpos=target), + binding=ActionBinding(manipulators={"primary": CONTROL_PART}), + motion_policy=MotionPolicy( + motion_source="motion_gen", + sample_count=SAMPLE_INTERVAL, + ), + ), + ) ) - assert success.shape == (1,) - assert bool(success.item()) + trajectory = result.trajectory.positions + assert result.plan_success.shape == (1,) + assert bool(result.plan_success.item()) assert trajectory.shape[2] == robot.dof # Default preserve_plan_samples=False resamples cuRobo's raw samples to # the action's sample_interval waypoint count. diff --git a/tests/sim/atomic_actions/test_engine.py b/tests/sim/atomic_actions/test_engine.py index 1d25fdb6e..c30ca1d1a 100644 --- a/tests/sim/atomic_actions/test_engine.py +++ b/tests/sim/atomic_actions/test_engine.py @@ -14,237 +14,352 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Tests for atomic_actions.engine.""" +"""Tests for the atomic-action registry and static compiler.""" from __future__ import annotations +from dataclasses import replace +from typing import ClassVar +from unittest.mock import Mock + import pytest import torch -from unittest.mock import Mock -from embodichain.lab.sim.atomic_actions.affordance import Affordance from embodichain.lab.sim.atomic_actions import ( - EndEffectorPoseTarget, - GraspTarget, - HeldObjectPoseTarget, - JointPositionTarget, - NamedJointPositionTarget, - PlaceTarget, -) -from embodichain.lab.sim.atomic_actions.core import ( - ActionTarget, - ActionResult, + ActionBinding, + ActionControlOverrides, + ActionInvocation, + ActionOptions, + ActionPlan, AtomicAction, - HeldObjectState, - ObjectSemantics, - WorldState, -) -from embodichain.lab.sim.atomic_actions.engine import ( AtomicActionEngine, - get_registered_actions, + BUILTIN_ACTION_TYPES, + ControlPartCommandProfile, + JointPositionCommand, + JointPositionGoal, + MotionPolicy, + PlanningContext, + PressGoal, + PressOptions, + ResolvedActionRequest, register_action, + get_registered_actions, unregister_action, ) -# --------------------------------------------------------------------------- -# Global registry (kept from old design) -# --------------------------------------------------------------------------- +class StubAction(AtomicAction[JointPositionGoal, ActionOptions]): + """Deterministic test action that commands every robot joint.""" + + skill_id: ClassVar[str] = "stub" + GoalType: ClassVar[type] = JointPositionGoal + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + + def plan( + self, + request: ResolvedActionRequest[JointPositionGoal, ActionOptions], + context: PlanningContext, + ) -> ActionPlan: + goal = self.require_goal(request) + assert isinstance(goal.target, torch.Tensor) + target = goal.target.to(context.robot.qpos) + if target.dim() == 1: + target = target.unsqueeze(0).expand(context.batch_size, -1) + success = torch.ones( + context.batch_size, dtype=torch.bool, device=context.robot.qpos.device + ) + if torch.isnan(target).any(dim=1).any(): + success &= ~torch.isnan(target).any(dim=1) + target = torch.nan_to_num(target) + trajectory = torch.stack([context.robot.qpos, target], dim=1) + return self.build_plan( + request, + context, + success=success, + trajectory=trajectory, + ) -class TestGlobalRegistry: - def teardown_method(self): - unregister_action("_test_dummy") - def test_register_and_retrieve(self): - cls = Mock() - register_action("_test_dummy", cls) - assert get_registered_actions()["_test_dummy"] is cls +class OtherStubAction(StubAction): + """Second configured skill used to verify shared engine resources.""" - def test_unregister(self): - register_action("_test_dummy", Mock()) - unregister_action("_test_dummy") - assert "_test_dummy" not in get_registered_actions() + skill_id: ClassVar[str] = "other_stub" - def test_unregister_nonexistent_is_noop(self): - unregister_action("_does_not_exist") - def test_get_registered_actions_returns_copy(self): - out = get_registered_actions() - out["_should_not_persist"] = Mock() - assert "_should_not_persist" not in get_registered_actions() +def _motion_generator( + batch_size: int = 2, + robot_dof: int = 3, +) -> Mock: + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = robot_dof + robot.control_parts = {"all": object()} + robot.get_qpos.return_value = torch.zeros(batch_size, robot_dof) + robot.get_qvel.return_value = torch.zeros(batch_size, robot_dof) + robot.get_joint_ids.return_value = list(range(robot_dof)) + generator = Mock() + generator.robot = robot + generator.device = torch.device("cpu") + generator.planner.cfg.planner_type = "stub_planner" + return generator + + +def _engine( + batch_size: int = 2, + robot_dof: int = 3, + control_profiles: dict[str, ControlPartCommandProfile] | None = None, + *, + load_builtins: bool = False, +) -> AtomicActionEngine: + return AtomicActionEngine( + _motion_generator(batch_size, robot_dof), + control_profiles=control_profiles, + load_builtins=load_builtins, + ) + + +def _invocation( + qpos: torch.Tensor, +) -> ActionInvocation[JointPositionGoal, ActionOptions]: + return ActionInvocation( + skill_id="stub", + goal=JointPositionGoal(qpos), + binding=ActionBinding(manipulators={"primary": "all"}), + motion_policy=MotionPolicy(sample_count=2), + ) + + +def test_global_registry_uses_stable_skill_id() -> None: + unregister_action("stub") + register_action(StubAction) + try: + assert get_registered_actions()["stub"] is StubAction + register_action(StubAction) + finally: + unregister_action("stub") + + +def test_engine_loads_fresh_builtin_instances_by_default() -> None: + first = AtomicActionEngine(_motion_generator()) + second = AtomicActionEngine(_motion_generator()) + expected_ids = tuple(action_type.skill_id for action_type in BUILTIN_ACTION_TYPES) + + assert tuple(first.actions) == expected_ids + assert all(action.is_bound for action in first.actions.values()) + assert all( + first.actions[skill_id] is not second.actions[skill_id] + for skill_id in expected_ids + ) -# --------------------------------------------------------------------------- -# Engine run() semantics -# --------------------------------------------------------------------------- +def test_engine_can_disable_builtin_loading() -> None: + assert _engine(load_builtins=False).actions == {} + + +def test_auto_registered_builtin_accepts_per_invocation_options() -> None: + engine = _engine(load_builtins=True) + options = PressOptions(hand_interp_steps=7) + invocation = ActionInvocation( + skill_id="press", + goal=PressGoal(torch.eye(4)), + binding=ActionBinding( + manipulators={"primary": "all"}, + end_effectors={"primary": "all"}, + ), + motion_policy=MotionPolicy(sample_count=20), + skill_options=options, + ) + + request = engine.resolve(invocation) + assert request.skill_options.hand_interp_steps == 7 + assert request.skill_options is not options -NUM_ENVS = 2 -TOTAL_DOF = 8 +def test_engine_compile_projects_terminal_state_between_actions() -> None: + engine = _engine() + engine.register(StubAction()) + first = torch.ones(2, 3) + second = torch.full((2, 3), 2.0) + + compiled = engine.compile((_invocation(first), _invocation(second))) + + assert compiled.plan_success.tolist() == [True, True] + assert compiled.trajectory.positions.shape == (2, 4, 3) + assert torch.equal(compiled.action_plans[1].trajectory.positions[:, 0], first) + assert torch.equal(compiled.projected_context.robot.qpos, second) + assert torch.count_nonzero(engine.robot.get_qpos()) == 0 + + +def test_engine_compile_holds_failed_rows_for_remaining_actions() -> None: + engine = _engine() + engine.register(StubAction()) + first = torch.tensor([[1.0, 1.0, 1.0], [float("nan"), 2.0, 2.0]]) + second = torch.full((2, 3), 4.0) + + compiled = engine.compile((_invocation(first), _invocation(second))) + + assert compiled.plan_success.tolist() == [True, False] + assert torch.all(compiled.projected_context.robot.qpos[0] == 4.0) + assert torch.all(compiled.projected_context.robot.qpos[1] == 0.0) + assert torch.all(compiled.trajectory.positions[1] == 0.0) + + +def test_engine_compile_empty_sequence_is_successful_noop() -> None: + engine = _engine() + context = engine.initial_context() + + compiled = engine.compile((), context) + + assert compiled.plan_success.tolist() == [True, True] + assert compiled.trajectory.positions.shape == (2, 0, 3) + assert compiled.projected_context is context -def _make_mg(): - robot = Mock() - robot.device = torch.device("cpu") - robot.dof = TOTAL_DOF - robot.get_qpos.return_value = torch.zeros(NUM_ENVS, TOTAL_DOF) - - mg = Mock() - mg.robot = robot - mg.device = torch.device("cpu") - return mg - - -def _fake_action(name, target_type, *, sets_held=False, clears_held=False, fails=False): - action = Mock(spec=AtomicAction) - action.TargetType = target_type - action.cfg = Mock() - action.cfg.name = name - - def execute(target, state): - if fails: - return ActionResult( - success=False, - trajectory=torch.empty(NUM_ENVS, 0, TOTAL_DOF), - next_state=state, - ) - held_objects = dict(state.held_objects) - if sets_held: - sem = ObjectSemantics(affordance=Affordance(), geometry={}, label="x") - 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_objects.pop("arm", None) - traj = torch.zeros(NUM_ENVS, 5, TOTAL_DOF) - return ActionResult( - success=True, - trajectory=traj, - next_state=state.with_updates( - last_qpos=traj[:, -1, :].clone(), - held_objects=held_objects, - ), - ) - action.execute = Mock(side_effect=execute) - return action - - -class TestEngineRun: - def setup_method(self): - self.mg = _make_mg() - self.engine = AtomicActionEngine(self.mg) - - def test_register_and_lookup(self): - action = _fake_action("pick_up", GraspTarget, sets_held=True) - self.engine.register(action) - assert "pick_up" in self.engine.actions - - def test_register_with_explicit_name_overrides_cfg(self): - action = _fake_action("pick_up", GraspTarget, sets_held=True) - self.engine.register(action, name="custom") - assert "custom" in self.engine.actions - - def test_run_concatenates_trajectories(self): - a = _fake_action("a", EndEffectorPoseTarget) - b = _fake_action("b", EndEffectorPoseTarget) - self.engine.register(a, name="a") - self.engine.register(b, name="b") - success, traj, _ = self.engine.run( - [ - ("a", EndEffectorPoseTarget(torch.eye(4))), - ("b", EndEffectorPoseTarget(torch.eye(4))), - ] - ) - assert success.all().item() - assert traj.shape == (NUM_ENVS, 10, TOTAL_DOF) - - def test_run_threads_world_state(self): - pick = _fake_action("pick", GraspTarget, sets_held=True) - move = _fake_action("move", HeldObjectPoseTarget) - 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") - sem = ObjectSemantics(affordance=Affordance(), geometry={}, label="x") - success, _, final_state = self.engine.run( - [ - ("pick", GraspTarget(sem)), - ("move", HeldObjectPoseTarget(torch.eye(4))), - ("place", PlaceTarget(torch.eye(4))), - ] - ) - assert success.all().item() - # 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.get_held_object("arm") is not None - # Final state cleared by place. - assert final_state.get_held_object("arm") is None - - def test_run_stops_on_first_failure(self): - a = _fake_action("a", EndEffectorPoseTarget) - b = _fake_action("b", EndEffectorPoseTarget, fails=True) - c = _fake_action("c", EndEffectorPoseTarget) - self.engine.register(a, name="a") - self.engine.register(b, name="b") - self.engine.register(c, name="c") - success, traj, _ = self.engine.run( - [ - ("a", EndEffectorPoseTarget(torch.eye(4))), - ("b", EndEffectorPoseTarget(torch.eye(4))), - ("c", EndEffectorPoseTarget(torch.eye(4))), - ] - ) - assert not success.any().item() - # `c` should not have been called. - c.execute.assert_not_called() - # We get only the trajectory from executed steps (all-dead breaks the loop). - assert traj.shape == (NUM_ENVS, 5, TOTAL_DOF) - - def test_run_raises_on_unknown_action_name(self): - with pytest.raises(KeyError, match="ghost"): - self.engine.run([("ghost", EndEffectorPoseTarget(torch.eye(4)))]) - - def test_run_raises_on_target_type_mismatch(self): - a = _fake_action("a", EndEffectorPoseTarget) - self.engine.register(a, name="a") - with pytest.raises(TypeError, match="target"): - self.engine.run([("a", HeldObjectPoseTarget(torch.eye(4)))]) - - def test_run_raises_on_tuple_target_type_mismatch(self): - a = _fake_action("a", (JointPositionTarget, NamedJointPositionTarget)) - self.engine.register(a, name="a") - with pytest.raises( - TypeError, match="JointPositionTarget.*NamedJointPositionTarget" - ): - self.engine.run([("a", EndEffectorPoseTarget(torch.eye(4)))]) - - def test_run_seeds_state_from_robot_when_none_provided(self): - a = _fake_action("a", EndEffectorPoseTarget) - self.engine.register(a, name="a") - seed_qpos = self.mg.robot.get_qpos.return_value - self.engine.run([("a", EndEffectorPoseTarget(torch.eye(4)))]) - # 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_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) +def test_engine_rejects_unknown_skill() -> None: + engine = _engine() + with pytest.raises(KeyError, match="stub"): + engine.compile((_invocation(torch.zeros(2, 3)),)) + + +def test_engine_rejects_duplicate_instance_registration() -> None: + engine = _engine() + first = StubAction() + second = StubAction() + engine.register(first) + with pytest.raises(ValueError, match="already registered"): + engine.register(second) + + +def test_engine_replaces_registered_action_only_when_explicit() -> None: + engine = _engine() + first = StubAction() + replacement = StubAction() + engine.register(first) + + engine.register(replacement, replace=True) + + assert engine.actions["stub"] is replacement + assert replacement.is_bound + + +def test_engine_binds_one_planning_service_to_every_action() -> None: + engine = _engine() + first = StubAction() + second = OtherStubAction() + + engine.register(first) + engine.register(second) + + assert first.motion_generator is engine.motion_generator + assert second.motion_generator is engine.motion_generator + assert first.builder is engine.planning_services.trajectory_builder + assert second.builder is first.builder + + +def test_engine_resolves_action_binding_from_robot_control_parts() -> None: + engine = _engine(robot_dof=3) + + resolved = engine.planning_services.resolve_binding( + ActionBinding(manipulators={"primary": "all"}) + ) + + assert resolved.manipulator().name == "all" + assert resolved.manipulator().joint_ids == (0, 1, 2) + assert resolved.manipulator().dof == 3 + + +def test_engine_resolves_invocation_control_override_into_request() -> None: + engine = _engine( + robot_dof=3, + control_profiles={ + "all": ControlPartCommandProfile.joint_positions(ready=torch.zeros(3)) + }, + ) + engine.register(StubAction()) + invocation = replace( + _invocation(torch.ones(2, 3)), + control_overrides=ActionControlOverrides( + manipulators={ + "primary": {"ready": JointPositionCommand(torch.full((3,), 0.4))} + } + ), + revision=2, + ) + + request = engine.resolve(invocation) + + assert request.revision == 2 + assert torch.allclose( + request.binding.manipulator().joint_positions("ready", n_envs=2, device="cpu"), + torch.full((2, 3), 0.4), + ) + + +def test_engine_rejects_binding_outside_robot_control_parts() -> None: + engine = _engine() + engine.register(StubAction()) + invocation = ActionInvocation( + skill_id="stub", + goal=JointPositionGoal(torch.zeros(2, 3)), + binding=ActionBinding(manipulators={"primary": "missing_arm"}), + motion_policy=MotionPolicy(sample_count=2), + ) + + with pytest.raises(ValueError, match="Robot.control_parts"): + engine.plan(invocation) + + +def test_engine_motion_generator_is_read_only() -> None: + engine = _engine() + + with pytest.raises(AttributeError): + engine.motion_generator = Mock() # type: ignore[misc] + + +def test_engine_plan_action_supports_unregistered_configured_instance() -> None: + engine = _engine() + action = StubAction() + + plan = engine.plan_action( + action, + _invocation(torch.ones(2, 3)), + engine.initial_context(), + ) + + assert plan.plan_success.tolist() == [True, True] + assert action.is_bound + assert engine.actions == {} + + +def test_action_cannot_be_rebound_to_another_engine() -> None: + action = StubAction() + _engine().register(action) + + with pytest.raises(ValueError, match="another AtomicActionEngine"): + _engine().register(action) + + +def test_unbound_action_rejects_direct_planning() -> None: + action = StubAction() + + with pytest.raises(RuntimeError, match="not bound"): + action.resolve_request(_invocation(torch.ones(2, 3))) + + +def test_engine_rejects_plan_for_a_different_skill() -> None: + engine = _engine() + action = StubAction() + original_plan = action.plan + + def wrong_skill_plan( + request: ResolvedActionRequest[JointPositionGoal, ActionOptions], + context: PlanningContext, + ) -> ActionPlan: + return replace(original_plan(request, context), skill_id="other") + + action.plan = wrong_skill_plan # type: ignore[method-assign] + engine.register(action) + + with pytest.raises(ValueError, match="must match its request"): + engine.compile((_invocation(torch.zeros(2, 3)),)) diff --git a/tests/sim/atomic_actions/test_engine_per_env.py b/tests/sim/atomic_actions/test_engine_per_env.py index 622d93e3f..38179d517 100644 --- a/tests/sim/atomic_actions/test_engine_per_env.py +++ b/tests/sim/atomic_actions/test_engine_per_env.py @@ -14,64 +14,371 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Per-environment failure propagation tests for AtomicActionEngine.run().""" +"""Tests for dynamic goals and closed-loop execution recovery.""" from __future__ import annotations -import torch -import pytest +from typing import ClassVar from unittest.mock import Mock -from embodichain.lab.sim.atomic_actions import EndEffectorPoseTarget -from embodichain.lab.sim.atomic_actions.core import ( - ActionCfg, - ActionResult, +import pytest +import torch + +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + ActionOptions, + ActionPlan, + Affordance, AtomicAction, - WorldState, + AtomicActionEngine, + EndEffectorPoseGoal, + EntityState, + ExecutionEventKind, + ExecutionStatus, + HeldObjectState, + MotionPolicy, + ObjectSemantics, + PlanningContext, + RecoveryPolicy, + ResolvedActionRequest, + RobotObservation, + SceneEntityPose, + SceneSnapshot, + StateDelta, + TaskState, ) -from embodichain.lab.sim.atomic_actions.engine import AtomicActionEngine - - -class _StubAction(AtomicAction): - TargetType = EndEffectorPoseTarget - - def __init__(self, mg, success_vec, traj_len=4, dof=3): - super().__init__(mg, ActionCfg()) - self._success = torch.tensor(success_vec) - self._traj_len = traj_len - self._dof = dof - - def execute(self, target, state): - n = state.last_qpos.shape[0] - traj = torch.zeros(n, self._traj_len, self._dof) - traj[:] = state.last_qpos.unsqueeze(1) - return ActionResult( - success=self._success.clone(), - trajectory=traj, - next_state=WorldState(last_qpos=traj[:, -1, :].clone()), +from embodichain.lab.sim.atomic_actions.goals import resolve_pose_goal + + +class DynamicAction(AtomicAction[EndEffectorPoseGoal, ActionOptions]): + """Test action whose terminal joint command follows a scene entity x pose.""" + + skill_id: ClassVar[str] = "dynamic" + GoalType: ClassVar[type] = EndEffectorPoseGoal + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + + def __init__(self) -> None: + super().__init__() + self.plan_count = 0 + self.requests: list[ResolvedActionRequest] = [] + + def plan( + self, + request: ResolvedActionRequest[EndEffectorPoseGoal, ActionOptions], + context: PlanningContext, + ) -> ActionPlan: + goal = self.require_goal(request) + self.plan_count += 1 + self.requests.append(request) + pose = resolve_pose_goal(goal.xpos, context, name="xpos") + target = pose[:, 0, 3].unsqueeze(1).expand_as(context.robot.qpos) + return self.build_plan( + request, + context, + success=True, + trajectory=torch.stack([context.robot.qpos, target], dim=1), ) -class TestRunPerEnv: - def test_failed_env_holds(self): - mg = Mock() - mg.robot.get_qpos = lambda: torch.zeros(3, 3) - mg.robot.dof = 3 - mg.device = torch.device("cpu") - eng = AtomicActionEngine(mg) - # env 1 fails step 2 - eng.register(_StubAction(mg, [True, True, True]), name="a") - eng.register(_StubAction(mg, [True, False, True]), name="b") - eng.register(_StubAction(mg, [True, True, True]), name="c") - success, traj, state = eng.run( - steps=[ - ("a", EndEffectorPoseTarget(xpos=torch.eye(4))), - ("b", EndEffectorPoseTarget(xpos=torch.eye(4))), - ("c", EndEffectorPoseTarget(xpos=torch.eye(4))), - ] +class EffectAction(DynamicAction): + """Dynamic test action that declares an attachment effect.""" + + skill_id: ClassVar[str] = "effect" + + def plan( + self, + request: ResolvedActionRequest[EndEffectorPoseGoal, ActionOptions], + context: PlanningContext, + ) -> ActionPlan: + goal = self.require_goal(request) + pose = resolve_pose_goal(goal.xpos, context, name="xpos") + target = pose[:, 0, 3].unsqueeze(1).expand_as(context.robot.qpos) + semantics = ObjectSemantics( + affordance=Affordance(), geometry={}, label="object" ) - assert success.tolist() == [True, False, True] - assert traj.shape[1] == 12 # 3 steps * 4 waypoints - # env 1's rows after its failure should equal its pre-failure qpos (held) - # all zeros here, so just check shape and that env 0/2 advanced - assert state.last_qpos.shape == (3, 3) + held = HeldObjectState( + semantics=semantics, + object_to_eef=torch.eye(4), + grasp_xpos=torch.eye(4), + ) + return self.build_plan( + request, + context, + success=True, + trajectory=torch.stack([context.robot.qpos, target], dim=1), + expected_effects=StateDelta(held_object_updates={"arm": held}), + ) + + +def _engine() -> tuple[AtomicActionEngine, DynamicAction]: + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = 2 + robot.control_parts = {"arm": object()} + robot.get_qpos.return_value = torch.zeros(1, 2) + robot.get_qvel.return_value = torch.zeros(1, 2) + robot.get_joint_ids.return_value = [0, 1] + generator = Mock() + generator.robot = robot + generator.device = torch.device("cpu") + generator.planner.cfg.planner_type = "stub" + engine = AtomicActionEngine(generator, load_builtins=False) + action = DynamicAction() + engine.register(action) + return engine, action + + +def _context( + timestamp: float, qpos: float, entity_x: float, version: int +) -> PlanningContext: + positions = torch.full((1, 2), qpos) + pose = torch.eye(4).unsqueeze(0) + pose[:, 0, 3] = entity_x + return PlanningContext( + robot=RobotObservation( + timestamp=timestamp, + qpos=positions, + qvel=torch.zeros_like(positions), + ), + task=TaskState.empty(batch_size=1, device="cpu"), + scene=SceneSnapshot( + timestamp=timestamp, + version=version, + entities={"target": EntityState(pose)}, + ), + env_ids=torch.tensor([0], dtype=torch.long), + ) + + +def _invocation( + *, + max_replans: int = 2, + max_phase_retries: int = 2, + phase_timeout: float = 30.0, +) -> ActionInvocation[EndEffectorPoseGoal]: + return ActionInvocation( + skill_id="dynamic", + goal=EndEffectorPoseGoal(SceneEntityPose("target")), + binding=ActionBinding(manipulators={"primary": "arm"}), + motion_policy=MotionPolicy(sample_count=2), + recovery_policy=RecoveryPolicy( + max_replans=max_replans, + max_phase_retries=max_phase_retries, + tracking_error_threshold=0.05, + goal_translation_threshold=0.02, + phase_timeout=phase_timeout, + ), + invocation_id="dynamic-call", + ) + + +def test_session_completes_incremental_command_sequence() -> None: + engine, _ = _engine() + session = engine.start((_invocation(),), _context(0.0, 0.0, 0.2, 0)) + + first = session.tick(_context(0.0, 0.0, 0.2, 0)) + second = session.tick(_context(0.1, 0.0, 0.2, 0)) + final = session.tick(_context(0.2, 0.2, 0.2, 0)) + + assert first.command is not None and torch.all(first.command.positions == 0.0) + assert all(event.invocation_id == "dynamic-call" for event in first.events) + assert second.command is not None and torch.all(second.command.positions == 0.2) + assert final.status is ExecutionStatus.COMPLETED + assert final.eligible_mask.tolist() == [True] + + +def test_scene_motion_replans_late_bound_goal() -> None: + engine, action = _engine() + session = engine.start((_invocation(),), _context(0.0, 0.0, 0.1, 0)) + session.tick(_context(0.0, 0.0, 0.1, 0)) + + tick = session.tick(_context(0.1, 0.0, 0.3, 1)) + + kinds = {event.kind for event in tick.events} + assert ExecutionEventKind.DYNAMIC_GOAL_CHANGED in kinds + assert ExecutionEventKind.REPLANNED in kinds + assert action.plan_count == 2 + assert action.requests[0] is action.requests[1] + assert tick.command is not None + + +def test_session_revision_replans_from_latest_context() -> None: + engine, action = _engine() + original = _invocation() + session = engine.start((original,), _context(0.0, 0.0, 0.1, 0)) + revised_pose = torch.eye(4).unsqueeze(0) + revised_pose[:, 0, 3] = 0.8 + revised = ActionInvocation( + skill_id=original.skill_id, + goal=EndEffectorPoseGoal(revised_pose), + binding=original.binding, + motion_policy=original.motion_policy, + recovery_policy=original.recovery_policy, + invocation_id=original.invocation_id, + revision=1, + ) + + session.revise_current(revised) + first = session.tick(_context(0.0, 0.0, 0.1, 0)) + second = session.tick(_context(0.1, 0.0, 0.1, 0)) + + assert action.plan_count == 2 + assert action.requests[0] is not action.requests[1] + assert [request.revision for request in action.requests] == [0, 1] + assert any( + event.kind is ExecutionEventKind.INVOCATION_REVISED + and event.invocation_revision == 1 + for event in first.events + ) + assert second.command is not None + assert torch.all(second.command.positions == 0.8) + + +def test_session_revision_must_advance_same_invocation() -> None: + engine, _ = _engine() + original = _invocation() + session = engine.start((original,), _context(0.0, 0.0, 0.1, 0)) + + with pytest.raises(ValueError, match="must advance"): + session.revise_current(original) + + with pytest.raises(ValueError, match="invocation_id"): + session.revise_current( + ActionInvocation( + skill_id=original.skill_id, + goal=original.goal, + binding=original.binding, + motion_policy=original.motion_policy, + recovery_policy=original.recovery_policy, + invocation_id="another-call", + revision=1, + ) + ) + + +def test_tracking_error_fails_when_replan_budget_is_zero() -> None: + engine, _ = _engine() + session = engine.start( + (_invocation(max_replans=0),), + _context(0.0, 0.0, 0.2, 0), + ) + session.tick(_context(0.0, 0.0, 0.2, 0)) + + tick = session.tick(_context(0.1, 1.0, 0.2, 0)) + + kinds = {event.kind for event in tick.events} + assert ExecutionEventKind.TRACKING_ERROR in kinds + assert ExecutionEventKind.RECOVERY_EXHAUSTED in kinds + assert tick.status is ExecutionStatus.FAILED + assert tick.eligible_mask.tolist() == [False] + + +def test_phase_timeout_retry_budget_is_bounded() -> None: + engine, action = _engine() + session = engine.start( + ( + _invocation( + max_phase_retries=1, + phase_timeout=0.05, + ), + ), + _context(0.0, 0.0, 0.2, 0), + ) + session.tick(_context(0.0, 0.0, 0.2, 0)) + + retry = session.tick(_context(0.1, 0.0, 0.2, 0)) + exhausted = session.tick(_context(0.2, 0.0, 0.2, 0)) + + retry_kinds = {event.kind for event in retry.events} + assert ExecutionEventKind.PHASE_TIMEOUT in retry_kinds + assert ExecutionEventKind.ACTION_RETRY in retry_kinds + assert ExecutionEventKind.REPLANNED in retry_kinds + assert action.plan_count == 2 + exhausted_kinds = {event.kind for event in exhausted.events} + assert ExecutionEventKind.PHASE_TIMEOUT in exhausted_kinds + assert ExecutionEventKind.RECOVERY_EXHAUSTED in exhausted_kinds + assert exhausted.status is ExecutionStatus.FAILED + assert exhausted.eligible_mask.tolist() == [False] + + +def test_session_rejects_changed_environment_identity() -> None: + engine, _ = _engine() + initial = _context(0.0, 0.0, 0.2, 0) + session = engine.start((_invocation(),), initial) + changed = PlanningContext( + robot=initial.robot, + task=initial.task, + scene=initial.scene, + env_ids=torch.tensor([7], dtype=torch.long), + ) + + with pytest.raises(ValueError, match="env_ids must remain stable"): + session.tick(changed) + + +def test_session_rejects_regressing_scene_snapshot() -> None: + engine, _ = _engine() + session = engine.start((_invocation(),), _context(1.0, 0.0, 0.2, 2)) + + with pytest.raises(ValueError, match="versions must be monotonic"): + session.tick(_context(1.0, 0.0, 0.2, 1)) + + +def test_nonempty_effect_is_committed_only_after_external_verification() -> None: + engine, _ = _engine() + effect = EffectAction() + engine.register(effect) + invocation = _invocation() + invocation = ActionInvocation( + skill_id="effect", + goal=invocation.goal, + binding=invocation.binding, + motion_policy=invocation.motion_policy, + recovery_policy=invocation.recovery_policy, + ) + session = engine.start((invocation,), _context(0.0, 0.0, 0.2, 0)) + session.tick(_context(0.0, 0.0, 0.2, 0)) + session.tick(_context(0.1, 0.0, 0.2, 0)) + + waiting = session.tick(_context(0.2, 0.2, 0.2, 0)) + completed = session.tick( + _context(0.3, 0.2, 0.2, 0), + effect_success=torch.tensor([True]), + ) + + assert waiting.status is ExecutionStatus.RUNNING + assert waiting.task_state.get_held_object("arm") is None + assert any( + event.kind is ExecutionEventKind.EFFECT_VERIFICATION_REQUIRED + for event in waiting.events + ) + assert completed.status is ExecutionStatus.COMPLETED + assert completed.task_state.get_held_object("arm") is not None + + +def test_effect_failure_does_not_commit_and_exhausts_retry_budget() -> None: + engine, _ = _engine() + engine.register(EffectAction()) + base = _invocation(max_phase_retries=0) + invocation = ActionInvocation( + skill_id="effect", + goal=base.goal, + binding=base.binding, + motion_policy=base.motion_policy, + recovery_policy=base.recovery_policy, + ) + session = engine.start((invocation,), _context(0.0, 0.0, 0.2, 0)) + session.tick(_context(0.0, 0.0, 0.2, 0)) + session.tick(_context(0.1, 0.0, 0.2, 0)) + + failed = session.tick( + _context(0.2, 0.2, 0.2, 0), + effect_success=torch.tensor([False]), + ) + + assert failed.status is ExecutionStatus.FAILED + assert failed.task_state.get_held_object("arm") is None + assert any( + event.kind is ExecutionEventKind.RECOVERY_EXHAUSTED for event in failed.events + ) diff --git a/tests/sim/atomic_actions/test_motion_source_e2e.py b/tests/sim/atomic_actions/test_motion_source_e2e.py index 93f6535d4..66778af01 100644 --- a/tests/sim/atomic_actions/test_motion_source_e2e.py +++ b/tests/sim/atomic_actions/test_motion_source_e2e.py @@ -25,12 +25,11 @@ from embodichain.lab.sim.robots import CobotMagicCfg from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - EndEffectorPoseTarget, -) -from embodichain.lab.sim.atomic_actions.actions import ( - MoveEndEffector, - MoveEndEffectorCfg, + EndEffectorPoseGoal, + MotionPolicy, ) @@ -59,12 +58,6 @@ def _setup(self, motion_source: str): MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=self.ROBOT_UID)) ) engine = AtomicActionEngine(mg) - cfg = MoveEndEffectorCfg( - motion_source=motion_source, - control_part=self.CONTROL_PART, - sample_interval=self.SAMPLE_INTERVAL, - ) - engine.register(MoveEndEffector(mg, cfg), name="move_end_effector") return sim, robot, engine def _teardown(self, sim): @@ -86,11 +79,23 @@ def _run_reach_test(self, motion_source: str): sim, robot, engine = self._setup(motion_source) try: target, arm_ids = self._reachable_target(robot) - success, traj, _ = engine.run( - [("move_end_effector", EndEffectorPoseTarget(xpos=target))] + result = engine.compile( + ( + ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal(xpos=target), + binding=ActionBinding( + manipulators={"primary": self.CONTROL_PART} + ), + motion_policy=MotionPolicy( + motion_source=motion_source, + sample_count=self.SAMPLE_INTERVAL, + ), + ), + ) ) - assert success.all().item(), f"{motion_source} reported failure" - final_q = traj[0, -1, arm_ids] + assert result.plan_success.all().item(), f"{motion_source} reported failure" + final_q = result.trajectory.positions[0, -1, arm_ids] fk = robot.compute_fk( qpos=final_q[None], name=self.CONTROL_PART, to_matrix=True )[0] diff --git a/tests/sim/atomic_actions/test_trajectory.py b/tests/sim/atomic_actions/test_trajectory.py index 006a713ca..461a6dbac 100644 --- a/tests/sim/atomic_actions/test_trajectory.py +++ b/tests/sim/atomic_actions/test_trajectory.py @@ -23,6 +23,9 @@ from unittest.mock import Mock, patch from embodichain.lab.sim.atomic_actions.trajectory import TrajectoryBuilder +from embodichain.lab.sim.atomic_actions.primitives._helpers import ( + resolve_object_target, +) def _make_mock_motion_generator(num_envs: int = 2, arm_dof: int = 6) -> Mock: @@ -71,6 +74,24 @@ def test_tensor_all_true(self): def test_tensor_any_false(self): assert self.builder.all_envs_success(torch.tensor([True, False])) is False + def test_binary_integer_success_is_normalized_at_planner_boundary(self): + success = self.builder._resolve_success_mask( + torch.tensor([1, 0], dtype=torch.int32), + n_envs=2, + name="IK success", + ) + + assert success.dtype == torch.bool + assert success.tolist() == [True, False] + + def test_non_binary_integer_success_is_rejected(self): + with pytest.raises(TypeError, match="binary integer"): + self.builder._resolve_success_mask( + torch.tensor([1, 2], dtype=torch.int32), + n_envs=2, + name="IK success", + ) + class TestResolvePoseTarget: def setup_method(self): @@ -86,6 +107,12 @@ def test_batched_pose_passes_through(self): out = self.builder.resolve_pose_target(pose, n_envs=2) assert torch.equal(out, pose) + def test_batched_pose_returns_owned_tensor(self): + pose = torch.eye(4).unsqueeze(0).repeat(2, 1, 1) + out = self.builder.resolve_pose_target(pose, n_envs=2) + out[:, 0, 0] = 2.0 + assert torch.equal(pose, torch.eye(4).unsqueeze(0).repeat(2, 1, 1)) + def test_pose_converts_to_builder_dtype_and_device(self): pose = torch.eye(4, dtype=torch.float64) out = self.builder.resolve_pose_target(pose, n_envs=2) @@ -115,6 +142,20 @@ def test_multi_waypoint_empty_raises(self): self.builder.resolve_pose_target(empty, n_envs=2) +class TestResolveObjectTarget: + def test_batched_pose_returns_owned_tensor(self): + pose = torch.eye(4).unsqueeze(0).repeat(2, 1, 1) + out = resolve_object_target( + pose, + n_envs=2, + device=torch.device("cpu"), + ) + + out[:, 0, 0] = 2.0 + + assert torch.equal(pose, torch.eye(4).unsqueeze(0).repeat(2, 1, 1)) + + class TestResolveJointTarget: def setup_method(self): self.builder = TrajectoryBuilder(_make_mock_motion_generator()) @@ -135,6 +176,15 @@ def test_batched_qpos_passes_through(self): ) assert torch.equal(out, qpos) + def test_batched_qpos_returns_owned_tensor(self): + qpos = torch.arange(12, dtype=torch.float32).reshape(2, 6) + expected = qpos.clone() + out = self.builder.resolve_joint_target( + qpos, n_envs=2, joint_dof=6, control_part="arm" + ) + out.zero_() + assert torch.equal(qpos, expected) + def test_wrong_shape_raises(self): with pytest.raises(Exception): self.builder.resolve_joint_target( diff --git a/tests/sim/atomic_actions/test_trajectory_motion_source.py b/tests/sim/atomic_actions/test_trajectory_motion_source.py index 4e7e4e6ac..1edcdd0f2 100644 --- a/tests/sim/atomic_actions/test_trajectory_motion_source.py +++ b/tests/sim/atomic_actions/test_trajectory_motion_source.py @@ -23,8 +23,9 @@ from unittest.mock import Mock, patch from embodichain.lab.sim.atomic_actions.trajectory import TrajectoryBuilder +from embodichain.lab.sim.planners import PlanOptions from embodichain.lab.sim.planners.utils import PlanState, PlanResult, MoveType -from embodichain.lab.sim.atomic_actions.core import ActionCfg +from embodichain.lab.sim.atomic_actions.policies import MotionPolicy def _mock_mg(num_envs=2, arm_dof=6, planner_type="toppra"): @@ -74,6 +75,31 @@ def _pose_targets_for_two_envs(): class TestPlanArmTrajMotionGen: + def test_shared_action_cfg_forwards_copied_plan_options(self): + mg = _mock_mg(num_envs=2, arm_dof=6) + mg.generate.return_value = PlanResult( + success=torch.ones(2, dtype=torch.bool), + positions=torch.zeros(2, 6, 6), + ) + builder = TrajectoryBuilder(mg) + cfg = MotionPolicy( + motion_source="motion_gen", + plan_opts=PlanOptions(), + ) + + builder.plan_arm_traj( + _pose_targets_for_two_envs(), + torch.zeros(2, 6), + 6, + control_part="arm", + arm_dof=6, + cfg=cfg, + ) + + forwarded = mg.generate.call_args.kwargs["options"].plan_opts + assert type(forwarded) is PlanOptions + assert forwarded is not cfg.plan_opts + def test_motion_gen_path_delegates_to_generate(self): mg = _mock_mg(num_envs=3, arm_dof=6) from embodichain.lab.sim.planners.utils import PlanResult @@ -83,7 +109,7 @@ def test_motion_gen_path_delegates_to_generate(self): positions=torch.zeros(3, 12, 6), ) builder = TrajectoryBuilder(mg) - cfg = ActionCfg(motion_source="motion_gen", control_part="arm") + cfg = MotionPolicy(motion_source="motion_gen") start_qpos = torch.zeros(3, 6) # per-env list[list[PlanState]] with single-env PlanStates (action contract) target_states_list = [ @@ -113,7 +139,7 @@ def test_motion_gen_path_delegates_to_generate(self): def test_ik_interp_path_unchanged(self): mg = _mock_mg(num_envs=2, arm_dof=6) builder = TrajectoryBuilder(mg) - cfg = ActionCfg(motion_source="ik_interp", control_part="arm") + cfg = MotionPolicy(motion_source="ik_interp") start_qpos = torch.zeros(2, 6) target_states_list = [ [PlanState(xpos=torch.eye(4), move_type=MoveType.EEF_MOVE)] @@ -146,9 +172,8 @@ def test_curobo_builder_preserves_cartesian_targets_and_samples(self): n_waypoints=20, control_part="arm", arm_dof=6, - cfg=ActionCfg( + cfg=MotionPolicy( motion_source="motion_gen", - control_part="arm", ), ) assert success.tolist() == [True, True] @@ -165,7 +190,7 @@ def test_curobo_builder_preserves_cartesian_targets_and_samples(self): def test_invalid_motion_source_raises(self): with pytest.raises(ValueError, match="motion_source"): - ActionCfg(motion_source="bogus") + MotionPolicy(motion_source="bogus") def test_nan_positions_rejected(self): positions = torch.zeros(2, 5, 6) @@ -179,9 +204,8 @@ def test_nan_positions_rejected(self): n_waypoints=10, control_part="arm", arm_dof=6, - cfg=ActionCfg( + cfg=MotionPolicy( motion_source="motion_gen", - control_part="arm", ), ) @@ -198,9 +222,8 @@ def test_none_positions_rejected(self): n_waypoints=10, control_part="arm", arm_dof=6, - cfg=ActionCfg( + cfg=MotionPolicy( motion_source="motion_gen", - control_part="arm", ), ) @@ -220,9 +243,8 @@ def test_failed_row_holds_start_qpos(self): n_waypoints=10, control_part="arm", arm_dof=6, - cfg=ActionCfg( + cfg=MotionPolicy( motion_source="motion_gen", - control_part="arm", ), ) assert success.tolist() == [True, False] @@ -243,9 +265,8 @@ def test_indexless_cuda_robot_accepts_indexed_curobo_result(self): n_waypoints=10, control_part="arm", arm_dof=6, - cfg=ActionCfg( + cfg=MotionPolicy( motion_source="motion_gen", - control_part="arm", ), ) @@ -262,7 +283,7 @@ def test_joint_capable_backend_delegates_through_motion_generator(self): positions=torch.zeros(2, 8, 6), ) builder = TrajectoryBuilder(mg) - cfg = ActionCfg(motion_source="motion_gen", control_part="arm") + cfg = MotionPolicy(motion_source="motion_gen") success, trajectory = builder.plan_joint_motion( torch.zeros(2, 6), @@ -283,7 +304,7 @@ def test_neural_joint_motion_falls_back_to_local_interpolation(self): builder = TrajectoryBuilder(mg) start = torch.zeros(2, 6) target = torch.ones(2, 6) - cfg = ActionCfg(motion_source="motion_gen", control_part="arm") + cfg = MotionPolicy(motion_source="motion_gen") with patch( "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", diff --git a/tests/sim/planners/test_curobo_planner.py b/tests/sim/planners/test_curobo_planner.py index 931aeecbd..3d952c443 100644 --- a/tests/sim/planners/test_curobo_planner.py +++ b/tests/sim/planners/test_curobo_planner.py @@ -630,8 +630,6 @@ def _make_curobo_engine( ) -> object: from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, - MoveEndEffector, - MoveEndEffectorCfg, ) from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator @@ -645,17 +643,6 @@ def _make_curobo_engine( ) ) engine = AtomicActionEngine(motion_generator) - engine.register( - MoveEndEffector( - motion_generator, - MoveEndEffectorCfg( - motion_source="motion_gen", - control_part=_SIM_CONTROL_PART, - sample_interval=80, - ), - ), - name="move_end_effector", - ) return engine @@ -663,7 +650,12 @@ def _make_curobo_engine( @pytest.mark.slow def test_curobo_reuses_non_graph_backend(): from embodichain.lab.sim import SimulationManager - from embodichain.lab.sim.atomic_actions import EndEffectorPoseTarget + from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + EndEffectorPoseGoal, + MotionPolicy, + ) pytest.importorskip("curobo", reason="cuRobo V2 not installed.") sim, robot, block = _build_curobo_scene() @@ -671,9 +663,18 @@ def test_curobo_reuses_non_graph_backend(): engine = _make_curobo_engine(block) target = _target_beyond_block(robot) - success, trajectory, _ = engine.run( - [("move_end_effector", EndEffectorPoseTarget(xpos=target))] + result = engine.compile( + ( + ActionInvocation( + "move_end_effector", + EndEffectorPoseGoal(xpos=target), + ActionBinding(manipulators={"primary": _SIM_CONTROL_PART}), + MotionPolicy(motion_source="motion_gen", sample_count=80), + ), + ) ) + success = result.plan_success + trajectory = result.trajectory.positions assert bool(success.item()), "first plan failed" assert trajectory.shape[0] == 1 @@ -681,9 +682,17 @@ def test_curobo_reuses_non_graph_backend(): assert planner.cfg.use_cuda_graph is False assert len(planner._backend_cache) == 1 - success, _, _ = engine.run( - [("move_end_effector", EndEffectorPoseTarget(xpos=target))] + result = engine.compile( + ( + ActionInvocation( + "move_end_effector", + EndEffectorPoseGoal(xpos=target), + ActionBinding(manipulators={"primary": _SIM_CONTROL_PART}), + MotionPolicy(motion_source="motion_gen", sample_count=80), + ), + ) ) + success = result.plan_success assert bool(success.item()), "second plan failed" assert len(planner._backend_cache) == 1 finally: @@ -695,7 +704,12 @@ def test_curobo_reuses_non_graph_backend(): @pytest.mark.slow def test_curobo_uses_accelerator_with_cpu_physics(): from embodichain.lab.sim import SimulationManager - from embodichain.lab.sim.atomic_actions import EndEffectorPoseTarget + from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + EndEffectorPoseGoal, + MotionPolicy, + ) pytest.importorskip("curobo", reason="cuRobo V2 not installed.") sim, robot, block = _build_curobo_scene(sim_device="cpu") @@ -703,9 +717,18 @@ def test_curobo_uses_accelerator_with_cpu_physics(): engine = _make_curobo_engine(block, use_cuda_graph=True) target = _target_beyond_block(robot) - success, trajectory, _ = engine.run( - [("move_end_effector", EndEffectorPoseTarget(xpos=target))] + result = engine.compile( + ( + ActionInvocation( + "move_end_effector", + EndEffectorPoseGoal(xpos=target), + ActionBinding(manipulators={"primary": _SIM_CONTROL_PART}), + MotionPolicy(motion_source="motion_gen", sample_count=80), + ), + ) ) + success = result.plan_success + trajectory = result.trajectory.positions planner = engine.motion_generator.planner backend = next(iter(planner._backend_cache.values()))