From c58ce08a03576033c9c048dae2e49dcced0a2627 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 2 Aug 2026 08:02:13 +0000 Subject: [PATCH 1/5] refactor(atomic-actions): enforce per-environment contracts --- embodichain/lab/sim/atomic_actions/core.py | 596 +++++++++++++++++- embodichain/lab/sim/atomic_actions/engine.py | 49 +- .../sim/atomic_actions/primitives/_helpers.py | 2 +- .../primitives/coordinated_placement.py | 1 + .../atomic_actions/primitives/hand_over.py | 2 +- .../primitives/move_end_effector.py | 5 +- .../sim/atomic_actions/primitives/pick_up.py | 8 +- .../lab/sim/atomic_actions/trajectory.py | 72 ++- tests/sim/atomic_actions/test_actions.py | 82 ++- tests/sim/atomic_actions/test_core.py | 168 ++++- tests/sim/atomic_actions/test_engine.py | 4 + .../sim/atomic_actions/test_engine_per_env.py | 65 ++ tests/sim/atomic_actions/test_trajectory.py | 50 ++ .../test_trajectory_motion_source.py | 27 + 14 files changed, 1074 insertions(+), 57 deletions(-) diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index fccb39be8..a58c43735 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -27,7 +27,15 @@ from .affordance import Affordance if TYPE_CHECKING: - from embodichain.lab.sim.planners import MotionGenerator + from embodichain.lab.sim.planners import MotionGenerator, PlanOptions + + +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 # ============================================================================= @@ -118,10 +126,38 @@ class HeldObjectState: """Semantics of the held object.""" object_to_eef: torch.Tensor - """Batched transform from object frame to end-effector frame, shape [n_envs, 4, 4].""" + """Object-to-end-effector transform, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" grasp_xpos: torch.Tensor - """Batched end-effector pose used to grasp the object, shape [n_envs, 4, 4].""" + """Grasp pose, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" + + env_mask: torch.Tensor | None = None + """Environments in which the held-object relation is active, shape ``(n_envs,)``.""" + + def __post_init__(self) -> None: + object_batch_size = _validate_held_pose( + self.object_to_eef, "HeldObjectState.object_to_eef" + ) + grasp_batch_size = _validate_held_pose( + self.grasp_xpos, "HeldObjectState.grasp_xpos" + ) + known_batch_sizes = { + size for size in (object_batch_size, grasp_batch_size) if size is not None + } + if len(known_batch_sizes) > 1: + raise ValueError( + "HeldObjectState pose tensors must use the same batch size, " + f"got {object_batch_size} and {grasp_batch_size}." + ) + if self.grasp_xpos.device != self.object_to_eef.device: + raise ValueError("HeldObjectState pose tensors must use the same device.") + batch_size = next(iter(known_batch_sizes), None) + self.env_mask = _normalize_optional_env_mask( + self.env_mask, + batch_size=batch_size, + device=self.object_to_eef.device, + name="HeldObjectState.env_mask", + ) @dataclass(slots=True, eq=False) @@ -132,16 +168,361 @@ class CoordinatedHeldObjectState: """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]``.""" + """Left object-to-end-effector transform, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" right_object_to_eef: torch.Tensor - """Transform from object frame to right end-effector frame, shape ``[n_envs, 4, 4]``.""" + """Right object-to-end-effector transform, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" left_grasp_xpos: torch.Tensor - """Left end-effector grasp pose for the shared object, shape ``[n_envs, 4, 4]``.""" + """Left grasp pose, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" right_grasp_xpos: torch.Tensor - """Right end-effector grasp pose for the shared object, shape ``[n_envs, 4, 4]``.""" + """Right grasp pose, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" + + env_mask: torch.Tensor | None = None + """Environments in which the coordinated hold is active, shape ``(n_envs,)``.""" + + def __post_init__(self) -> None: + pose_fields = { + "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, + } + batch_sizes = { + name: _validate_held_pose(value, f"CoordinatedHeldObjectState.{name}") + for name, value in pose_fields.items() + } + known_batch_sizes = {size for size in batch_sizes.values() if size is not None} + if len(known_batch_sizes) > 1: + raise ValueError( + "CoordinatedHeldObjectState pose tensors must use the same batch " + f"size, got {batch_sizes}." + ) + devices = {value.device for value in pose_fields.values()} + if len(devices) != 1: + raise ValueError( + "CoordinatedHeldObjectState pose tensors must use the same device." + ) + batch_size = next(iter(known_batch_sizes), None) + self.env_mask = _normalize_optional_env_mask( + self.env_mask, + batch_size=batch_size, + device=self.left_object_to_eef.device, + name="CoordinatedHeldObjectState.env_mask", + ) + + +def _validate_held_pose(value: torch.Tensor, name: str) -> int | None: + """Validate a held-state pose and return its explicit batch size, if any.""" + 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) with n_envs > 0, " + f"got {tuple(value.shape)}." + ) + return int(value.shape[0]) + + +def _normalize_optional_env_mask( + value: torch.Tensor | None, + *, + batch_size: int | None, + device: torch.device, + name: str, +) -> torch.Tensor | None: + """Normalize a mask when a held-state batch can already be inferred.""" + if batch_size is None and value is None: + return None + if batch_size is None: + 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.dim() != 1 or value.shape[0] == 0: + raise ValueError( + f"{name} must have shape (n_envs,) with n_envs > 0, " + f"got {tuple(value.shape)}." + ) + batch_size = int(value.shape[0]) + return _normalize_env_mask( + value, + batch_size=batch_size, + device=device, + name=name, + ) + + +def _normalize_env_mask( + value: torch.Tensor | None, + *, + batch_size: int, + device: torch.device, + name: str, +) -> torch.Tensor: + """Return an owned boolean environment 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_held_pose( + value: torch.Tensor, + *, + batch_size: int, + device: torch.device, + name: str, +) -> torch.Tensor: + """Resolve an optionally batched held-state pose to a world-state batch.""" + pose_batch_size = _validate_held_pose(value, name) + if value.device != device: + raise ValueError(f"{name} must use the same device as WorldState.last_qpos.") + if pose_batch_size is None: + return value.unsqueeze(0).expand(batch_size, -1, -1).clone() + if pose_batch_size != batch_size: + raise ValueError( + "Held-object state batch size must match WorldState.last_qpos; " + f"expected {batch_size}, got {pose_batch_size}." + ) + return value + + +def _normalize_held_object_state( + value: HeldObjectState, + *, + batch_size: int, + device: torch.device, +) -> HeldObjectState: + """Return a held-object state normalized to a world-state batch.""" + object_batch_size = _validate_held_pose( + value.object_to_eef, "HeldObjectState.object_to_eef" + ) + grasp_batch_size = _validate_held_pose( + value.grasp_xpos, "HeldObjectState.grasp_xpos" + ) + if ( + object_batch_size == batch_size + and grasp_batch_size == batch_size + and value.object_to_eef.device == device + and value.grasp_xpos.device == device + and isinstance(value.env_mask, torch.Tensor) + and value.env_mask.dtype == torch.bool + and value.env_mask.shape == (batch_size,) + and value.env_mask.device == device + ): + return value + return HeldObjectState( + semantics=value.semantics, + object_to_eef=_broadcast_held_pose( + value.object_to_eef, + batch_size=batch_size, + device=device, + name="HeldObjectState.object_to_eef", + ), + grasp_xpos=_broadcast_held_pose( + value.grasp_xpos, + batch_size=batch_size, + device=device, + name="HeldObjectState.grasp_xpos", + ), + env_mask=_normalize_env_mask( + value.env_mask, + batch_size=batch_size, + device=device, + name="HeldObjectState.env_mask", + ), + ) + + +def _normalize_coordinated_held_object_state( + value: CoordinatedHeldObjectState, + *, + batch_size: int, + device: torch.device, +) -> CoordinatedHeldObjectState: + """Return a coordinated-held state normalized to a world-state batch.""" + pose_fields = { + "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, + } + pose_batch_sizes = { + name: _validate_held_pose( + pose, + f"CoordinatedHeldObjectState.{name}", + ) + for name, pose in pose_fields.items() + } + if ( + all(size == batch_size for size in pose_batch_sizes.values()) + and all(pose.device == device for pose in pose_fields.values()) + and isinstance(value.env_mask, torch.Tensor) + and value.env_mask.dtype == torch.bool + and value.env_mask.shape == (batch_size,) + and value.env_mask.device == device + ): + return value + normalized_poses = { + name: _broadcast_held_pose( + pose, + batch_size=batch_size, + device=device, + name=f"CoordinatedHeldObjectState.{name}", + ) + for name, pose in pose_fields.items() + } + return CoordinatedHeldObjectState( + semantics=value.semantics, + **normalized_poses, + env_mask=_normalize_env_mask( + value.env_mask, + batch_size=batch_size, + device=device, + name="CoordinatedHeldObjectState.env_mask", + ), + ) + + +def _merge_held_object_state( + previous: HeldObjectState | None, + candidate: HeldObjectState | None, + update_mask: torch.Tensor, +) -> HeldObjectState | None: + """Merge one held-object entry using a per-environment update mask.""" + if previous is not None: + assert previous.env_mask is not None + if candidate is not None: + assert candidate.env_mask is not None + if previous is None and candidate is None: + return None + if previous is None: + assert candidate is not None + env_mask = candidate.env_mask & update_mask + if not env_mask.any(): + return None + return HeldObjectState( + semantics=candidate.semantics, + object_to_eef=candidate.object_to_eef, + grasp_xpos=candidate.grasp_xpos, + env_mask=env_mask, + ) + if candidate is None: + env_mask = previous.env_mask & ~update_mask + if not env_mask.any(): + return None + return HeldObjectState( + semantics=previous.semantics, + object_to_eef=previous.object_to_eef, + grasp_xpos=previous.grasp_xpos, + env_mask=env_mask, + ) + + 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 control part " + "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_held_object_state( + previous: CoordinatedHeldObjectState | None, + candidate: CoordinatedHeldObjectState | None, + update_mask: torch.Tensor, +) -> CoordinatedHeldObjectState | None: + """Merge one coordinated-held entry using a per-environment update mask.""" + if previous is not None: + assert previous.env_mask is not None + if candidate is not None: + assert candidate.env_mask is not None + if previous is None and candidate is None: + return None + if previous is None: + assert candidate is not None + env_mask = candidate.env_mask & update_mask + if not env_mask.any(): + return None + return CoordinatedHeldObjectState( + semantics=candidate.semantics, + left_object_to_eef=candidate.left_object_to_eef, + right_object_to_eef=candidate.right_object_to_eef, + left_grasp_xpos=candidate.left_grasp_xpos, + right_grasp_xpos=candidate.right_grasp_xpos, + env_mask=env_mask, + ) + if candidate is None: + env_mask = previous.env_mask & ~update_mask + if not env_mask.any(): + return None + return CoordinatedHeldObjectState( + semantics=previous.semantics, + left_object_to_eef=previous.left_object_to_eef, + right_object_to_eef=previous.right_object_to_eef, + left_grasp_xpos=previous.left_grasp_xpos, + right_grasp_xpos=previous.right_grasp_xpos, + env_mask=env_mask, + ) + + 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 semantics for one control-part " + "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(slots=True, eq=False) @@ -159,6 +540,69 @@ class WorldState: ) """Objects jointly held by two control parts, keyed by their ordered pair.""" + def __post_init__(self) -> None: + if not isinstance(self.last_qpos, torch.Tensor): + raise TypeError("WorldState.last_qpos must be a torch.Tensor.") + if ( + self.last_qpos.dim() != 2 + or self.last_qpos.shape[0] == 0 + or self.last_qpos.shape[1] == 0 + ): + raise ValueError( + "WorldState.last_qpos must have shape (n_envs, robot_dof) with " + f"both dimensions non-zero, got {tuple(self.last_qpos.shape)}." + ) + held_objects: dict[str, HeldObjectState] = {} + for control_part, held in self.held_objects.items(): + if not isinstance(control_part, str) or not control_part: + raise TypeError( + "WorldState.held_objects keys must be non-empty strings." + ) + if not isinstance(held, HeldObjectState): + raise TypeError( + "WorldState.held_objects values must be HeldObjectState instances." + ) + held_objects[control_part] = _normalize_held_object_state( + held, + batch_size=self.batch_size, + device=self.last_qpos.device, + ) + coordinated_held_objects: dict[tuple[str, str], CoordinatedHeldObjectState] = {} + for control_parts, held in self.coordinated_held_objects.items(): + if ( + not isinstance(control_parts, tuple) + or len(control_parts) != 2 + or not all(isinstance(part, str) and part for part in control_parts) + ): + raise TypeError( + "WorldState.coordinated_held_objects keys must be pairs of " + "non-empty control-part names." + ) + if not isinstance(held, CoordinatedHeldObjectState): + raise TypeError( + "WorldState.coordinated_held_objects values must be " + "CoordinatedHeldObjectState instances." + ) + coordinated_held_objects[control_parts] = ( + _normalize_coordinated_held_object_state( + held, + batch_size=self.batch_size, + device=self.last_qpos.device, + ) + ) + self.held_objects = held_objects + self.coordinated_held_objects = coordinated_held_objects + + @property + def batch_size(self) -> int: + """Number of vectorized environments represented by this state.""" + return int(self.last_qpos.shape[0]) + + @property + def robot_dof(self) -> int: + """Number of robot joint-position columns represented by this state.""" + return int(self.last_qpos.shape[1]) + 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) @@ -195,14 +639,79 @@ def with_updates( ), ) + def masked_merge( + self, + candidate: WorldState, + update_mask: torch.Tensor, + ) -> WorldState: + """Merge a candidate successor for selected environments. + + Args: + candidate: Candidate successor returned by an atomic action. + update_mask: Boolean tensor of shape ``(n_envs,)``. Candidate robot + and held-object state is committed only where this mask is true. + + Returns: + A new state that preserves the current values in unselected rows. + + Raises: + TypeError: If ``candidate`` or ``update_mask`` has an invalid type. + ValueError: If state shapes, devices, or semantics are incompatible. + """ + if not isinstance(candidate, WorldState): + raise TypeError("candidate must be a WorldState instance.") + if candidate.last_qpos.shape != self.last_qpos.shape: + raise ValueError( + "Candidate WorldState.last_qpos must match the current shape, " + f"got {tuple(candidate.last_qpos.shape)} and " + f"{tuple(self.last_qpos.shape)}." + ) + if candidate.last_qpos.device != self.last_qpos.device: + raise ValueError("WorldState values being merged must use the same device.") + update_mask = _normalize_env_mask( + update_mask, + batch_size=self.batch_size, + device=self.last_qpos.device, + name="update_mask", + ) + + held_objects: dict[str, HeldObjectState] = {} + held_keys = dict.fromkeys((*self.held_objects, *candidate.held_objects)) + for key in held_keys: + merged = _merge_held_object_state( + self.held_objects.get(key), candidate.held_objects.get(key), update_mask + ) + if merged is not None: + held_objects[key] = merged + + coordinated_held_objects: dict[tuple[str, str], CoordinatedHeldObjectState] = {} + coordinated_keys = dict.fromkeys( + (*self.coordinated_held_objects, *candidate.coordinated_held_objects) + ) + for key in coordinated_keys: + merged = _merge_coordinated_held_object_state( + self.coordinated_held_objects.get(key), + candidate.coordinated_held_objects.get(key), + update_mask, + ) + if merged is not None: + coordinated_held_objects[key] = merged + + return WorldState( + last_qpos=torch.where( + update_mask[:, None], candidate.last_qpos, self.last_qpos + ), + held_objects=held_objects, + coordinated_held_objects=coordinated_held_objects, + ) + @dataclass(slots=True, eq=False) class ActionResult: """Return value of every AtomicAction.execute call.""" - 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,).""" + success: torch.Tensor + """Per-environment planning success, normalized to shape ``(n_envs,)``.""" trajectory: torch.Tensor """Full-robot trajectory, shape (n_envs, n_waypoints, robot.dof).""" @@ -210,12 +719,63 @@ class ActionResult: next_state: WorldState """World state to feed into the next action.""" + def __post_init__(self) -> None: + if not isinstance(self.trajectory, torch.Tensor): + raise TypeError("ActionResult.trajectory must be a torch.Tensor.") + if self.trajectory.dim() != 3: + raise ValueError( + "ActionResult.trajectory must have shape " + f"(n_envs, n_waypoints, robot_dof), got {tuple(self.trajectory.shape)}." + ) + if not isinstance(self.next_state, WorldState): + raise TypeError("ActionResult.next_state must be a WorldState instance.") + expected_shape = ( + self.next_state.batch_size, + self.next_state.robot_dof, + ) + if (self.trajectory.shape[0], self.trajectory.shape[2]) != expected_shape: + raise ValueError( + "ActionResult trajectory batch/DoF must match next_state.last_qpos; " + f"got trajectory {tuple(self.trajectory.shape)} and state " + f"{tuple(self.next_state.last_qpos.shape)}." + ) + if self.trajectory.device != self.next_state.last_qpos.device: + raise ValueError( + "ActionResult trajectory and next_state.last_qpos must use the " + "same device." + ) + + batch_size = self.next_state.batch_size + if isinstance(self.success, bool): + success = torch.full( + (batch_size,), + self.success, + dtype=torch.bool, + device=self.trajectory.device, + ) + elif isinstance(self.success, torch.Tensor): + if self.success.dtype != torch.bool: + raise TypeError( + "ActionResult.success must have dtype torch.bool, " + f"got {self.success.dtype}." + ) + success = self.success.to(device=self.trajectory.device) + if success.dim() == 0 or success.shape == (1,): + success = success.reshape(1).expand(batch_size) + if success.shape != (batch_size,): + raise ValueError( + f"ActionResult.success must have shape ({batch_size},), " + f"got {tuple(success.shape)}." + ) + success = success.clone() + else: + raise TypeError("ActionResult.success must be a bool or torch.Tensor.") + self.success = success + @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) + return bool(torch.all(self.success).item()) def __bool__(self) -> bool: import warnings as _w @@ -240,8 +800,13 @@ class ActionCfg: name: str = "default" control_part: str = "arm" interpolation_type: str = "linear" + """Interpolation policy. Only ``"linear"`` is currently implemented.""" + velocity_limit: float | None = None acceleration_limit: float | None = None + plan_opts: PlanOptions | None = None + """Optional planner-specific options copied for each motion-generator call.""" + motion_source: str = "ik_interp" """Trajectory source: 'ik_interp' (default, batched IK + linear interp) or 'motion_gen' (batched MotionGenerator).""" @@ -253,6 +818,11 @@ def __post_init__(self) -> None: f"motion_source must be one of {sorted(valid_sources)}, " f"but got {self.motion_source!r}." ) + if self.interpolation_type != "linear": + raise ValueError( + "interpolation_type currently supports only 'linear', " + f"but got {self.interpolation_type!r}." + ) # ============================================================================= diff --git a/embodichain/lab/sim/atomic_actions/engine.py b/embodichain/lab/sim/atomic_actions/engine.py index b1b0b3915..0d2e284ed 100644 --- a/embodichain/lab/sim/atomic_actions/engine.py +++ b/embodichain/lab/sim/atomic_actions/engine.py @@ -26,6 +26,7 @@ ActionResult, AtomicAction, WorldState, + _resolve_runtime_device, ) if TYPE_CHECKING: @@ -73,7 +74,7 @@ class AtomicActionEngine: def __init__(self, motion_generator: MotionGenerator) -> None: self.motion_generator = motion_generator self.robot = motion_generator.robot - self.device = motion_generator.device + self.device = _resolve_runtime_device(motion_generator.device) self._actions: dict[str, AtomicAction] = {} @property @@ -125,7 +126,23 @@ def run( if state is None: state = WorldState(last_qpos=self.robot.get_qpos().clone()) - b = state.last_qpos.shape[0] + if state.robot_dof != self.robot.dof: + raise ValueError( + "Initial WorldState DoF must match the engine robot, " + f"got {state.robot_dof} and {self.robot.dof}." + ) + robot_batch_size = int(self.robot.get_qpos().shape[0]) + if state.batch_size != robot_batch_size: + raise ValueError( + "Initial WorldState batch size must match the engine robot, " + f"got {state.batch_size} and {robot_batch_size}." + ) + if state.last_qpos.device != self.device: + raise ValueError( + "Initial WorldState and AtomicActionEngine must use the same device." + ) + + b = state.batch_size full_traj = torch.empty( (b, 0, self.robot.dof), dtype=torch.float32, @@ -148,23 +165,29 @@ def run( 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) - ) - step_success = step_success.to(self.device) + if result.trajectory.shape[0] != b: + raise ValueError( + f"Action '{name}' returned batch {result.trajectory.shape[0]}, " + f"but the engine state batch is {b}." + ) + if result.trajectory.shape[2] != self.robot.dof: + raise ValueError( + f"Action '{name}' returned {result.trajectory.shape[2]} DoF, " + f"but the engine robot has {self.robot.dof}." + ) + if result.trajectory.device != self.device: + raise ValueError( + f"Action '{name}' returned a trajectory on " + f"{result.trajectory.device}, expected {self.device}." + ) + step_success = result.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 - ) - ) + state = state.masked_merge(result.next_state, alive) return alive, full_traj, state diff --git a/embodichain/lab/sim/atomic_actions/primitives/_helpers.py b/embodichain/lab/sim/atomic_actions/primitives/_helpers.py index 2aecb8cad..8aec84027 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/_helpers.py +++ b/embodichain/lab/sim/atomic_actions/primitives/_helpers.py @@ -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): diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py index 7f77328a5..35e8c9fc8 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py @@ -517,6 +517,7 @@ def _plan_named_arm_trajectory( n_waypoints, control_part=control_part, arm_dof=start_qpos.shape[-1], + cfg=self.cfg, ) return self.builder.all_envs_success(success), trajectory diff --git a/embodichain/lab/sim/atomic_actions/primitives/hand_over.py b/embodichain/lab/sim/atomic_actions/primitives/hand_over.py index a05758ce6..ba0298d01 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/hand_over.py +++ b/embodichain/lab/sim/atomic_actions/primitives/hand_over.py @@ -205,7 +205,7 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: 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 + final_object_pose = self.final_object_pose.clone() # 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] 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..682e18fa8 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py @@ -23,7 +23,7 @@ import torch -from embodichain.lab.sim.planners import MoveType, PlanOptions, PlanState +from embodichain.lab.sim.planners import MoveType, PlanState from embodichain.utils import configclass from ._helpers import arm_qpos_from_state @@ -64,9 +64,6 @@ class MoveEndEffectorCfg(ActionCfg): 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.""" - class MoveEndEffector(AtomicAction[EndEffectorPoseTarget]): """Plan a free-space end-effector move to a target pose. diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index fbb0c35fe..85cbe8940 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -257,7 +257,7 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: target.grasp_xpos, n_envs=self.n_envs ) if self.cfg.rotate_upright is not None: - self._apply_upright_rotation(sem, grasp_xpos) + grasp_xpos = self._upright_adjusted_grasp_poses(sem, grasp_xpos) 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.") @@ -480,12 +480,6 @@ def _compute_batch_candidate_ik( qpos.reshape(n_envs, n_pose, n_variant, self.arm_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 ) -> torch.Tensor: diff --git a/embodichain/lab/sim/atomic_actions/trajectory.py b/embodichain/lab/sim/atomic_actions/trajectory.py index e854aa580..0a4953568 100644 --- a/embodichain/lab/sim/atomic_actions/trajectory.py +++ b/embodichain/lab/sim/atomic_actions/trajectory.py @@ -31,18 +31,12 @@ from embodichain.lab.sim.utility.action_utils import interpolate_with_distance from embodichain.utils import logger +from .core import _resolve_runtime_device + 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 - - class TrajectoryBuilder: """Stateless trajectory utilities shared by every atomic action. @@ -66,6 +60,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 +122,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 +202,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: @@ -362,6 +399,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." @@ -417,13 +459,13 @@ 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) + n_envs = start_qpos.shape[0] + success = self._resolve_success_mask( + result.success, + n_envs=n_envs, + name="MotionGenerator PlanResult.success", ) positions = result.positions - n_envs = start_qpos.shape[0] if positions is None or positions.ndim != 3: logger.log_error( "MotionGenerator returned no (B, N, controlled_dof) positions", diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index 808d2c4fb..16a6b45ab 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -31,6 +31,8 @@ CoordinatedPlacementTarget, EndEffectorPoseTarget, GraspTarget, + HandOver, + HandOverCfg, HeldObjectPoseTarget, JointPositionTarget, NamedJointPositionTarget, @@ -1042,6 +1044,67 @@ def interpolate(trajectory, interp_num, device): assert result.next_state.get_held_object("arm") is held +# --------------------------------------------------------------------------- +# HandOver +# --------------------------------------------------------------------------- + + +class TestHandOverAction: + def test_execute_does_not_mutate_cached_final_pose(self): + motion_generator = _make_dual_arm_mock_motion_generator() + action = HandOver( + motion_generator, + HandOverCfg( + transfer_hand_open_qpos=_hand_open(), + transfer_hand_close_qpos=_hand_close(), + receive_hand_open_qpos=_hand_open(), + receive_hand_close_qpos=_hand_close(), + middle_object_pose=torch.eye(4), + final_object_pose=torch.eye(4), + sample_interval=30, + hand_interp_steps=4, + hold_steps=2, + retreat_steps=5, + ), + ) + original_final_pose = action.final_object_pose.clone() + current_object_pose = torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1) + current_object_pose[:, :3, :3] = torch.diag(torch.tensor([-1.0, -1.0, 1.0])) + entity = Mock() + entity.get_local_pose.return_value = current_object_pose + semantics = ObjectSemantics( + affordance=AntipodalAffordance(), + geometry={}, + label="handover-object", + entity=entity, + ) + 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), + ) + state = WorldState( + last_qpos=torch.zeros(NUM_ENVS, DUAL_TOTAL_DOF), + held_objects={"left_arm": held}, + ) + receive_grasp = torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1) + action._resolve_receive_grasp = Mock( + return_value=( + receive_grasp, + torch.ones(NUM_ENVS, dtype=torch.bool), + ) + ) + + def plan_from_start(control_part, start_qpos, target_poses, n_waypoints): + return True, start_qpos.unsqueeze(1).repeat(1, n_waypoints, 1) + + action._plan_named_arm_trajectory = Mock(side_effect=plan_from_start) + + action.execute(GraspTarget(semantics=semantics), state) + + assert torch.equal(action.final_object_pose, original_final_pose) + + # --------------------------------------------------------------------------- # CoordinatedPickment # --------------------------------------------------------------------------- @@ -1176,6 +1239,23 @@ def setup_method(self): ) self.action = CoordinatedPlacement(self.mg, cfg=self.cfg) + def test_named_arm_planning_forwards_action_configuration(self): + self.action.builder.plan_arm_traj = Mock( + return_value=( + torch.ones(NUM_ENVS, dtype=torch.bool), + torch.zeros(NUM_ENVS, 4, ARM_DOF), + ) + ) + + self.action._plan_named_arm_trajectory( + "left_arm", + torch.zeros(NUM_ENVS, ARM_DOF), + torch.eye(4).reshape(1, 1, 4, 4).repeat(NUM_ENVS, 1, 1, 1), + 4, + ) + + assert self.action.builder.plan_arm_traj.call_args.kwargs["cfg"] is self.cfg + def _make_target_and_state( self, ) -> tuple[ @@ -1298,7 +1378,7 @@ def interpolate(trajectory, interp_num, device): ): result = self.action.execute(target, state) - assert result.success is True + assert result.success.tolist() == [True] * NUM_ENVS assert result.trajectory.shape == ( NUM_ENVS, self.cfg.sample_interval, diff --git a/tests/sim/atomic_actions/test_core.py b/tests/sim/atomic_actions/test_core.py index 9fabb3699..5c7133ac4 100644 --- a/tests/sim/atomic_actions/test_core.py +++ b/tests/sim/atomic_actions/test_core.py @@ -249,6 +249,26 @@ def test_required_fields(self): assert s.semantics is sem assert s.object_to_eef.shape == (1, 4, 4) assert s.grasp_xpos.shape == (1, 4, 4) + assert s.env_mask.tolist() == [True] + + def test_rejects_mismatched_pose_batches(self): + sem = ObjectSemantics(affordance=Affordance(), geometry={}) + with pytest.raises(ValueError, match="same batch size"): + HeldObjectState( + semantics=sem, + object_to_eef=torch.eye(4).unsqueeze(0), + grasp_xpos=torch.eye(4).unsqueeze(0).repeat(2, 1, 1), + ) + + def test_rejects_invalid_env_mask(self): + sem = ObjectSemantics(affordance=Affordance(), geometry={}) + with pytest.raises(ValueError, match="env_mask"): + HeldObjectState( + semantics=sem, + object_to_eef=torch.eye(4).unsqueeze(0), + grasp_xpos=torch.eye(4).unsqueeze(0), + env_mask=torch.ones(2, dtype=torch.bool), + ) class TestCoordinatedHeldObjectState: @@ -313,16 +333,155 @@ def test_with_updates_does_not_alias_held_state_dictionaries(self): ) assert ws.held_objects == {} + def test_rejects_non_batched_last_qpos(self): + with pytest.raises(ValueError, match="last_qpos"): + WorldState(last_qpos=torch.zeros(6)) + + def test_rejects_held_state_with_different_batch(self): + held = 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="batch size"): + WorldState( + last_qpos=torch.zeros(2, 6), + held_objects={"arm": held}, + ) + + def test_broadcasts_unbatched_held_state_at_world_boundary(self): + held = HeldObjectState( + semantics=ObjectSemantics(affordance=Affordance(), geometry={}), + object_to_eef=torch.eye(4), + grasp_xpos=torch.eye(4), + ) + + world = WorldState( + last_qpos=torch.zeros(2, 6), + held_objects={"arm": held}, + ) + + normalized = world.get_held_object("arm") + assert normalized is not None + assert normalized is not held + assert normalized.object_to_eef.shape == (2, 4, 4) + assert normalized.grasp_xpos.shape == (2, 4, 4) + assert normalized.env_mask.tolist() == [True, True] + + def test_masked_merge_applies_new_hold_only_to_successful_envs(self): + batch_size = 2 + previous = WorldState(last_qpos=torch.zeros(batch_size, 6)) + held = HeldObjectState( + semantics=ObjectSemantics(affordance=Affordance(), geometry={}), + object_to_eef=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), + grasp_xpos=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), + ) + candidate = WorldState( + last_qpos=torch.ones(batch_size, 6), + held_objects={"arm": held}, + ) + + merged = previous.masked_merge( + candidate, torch.tensor([True, False], dtype=torch.bool) + ) + + assert merged.get_held_object("arm").env_mask.tolist() == [True, False] + + def test_masked_merge_preserves_removed_hold_in_failed_envs(self): + batch_size = 2 + held = HeldObjectState( + semantics=ObjectSemantics(affordance=Affordance(), geometry={}), + object_to_eef=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), + grasp_xpos=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), + ) + previous = WorldState( + last_qpos=torch.zeros(batch_size, 6), + held_objects={"arm": held}, + ) + candidate = WorldState(last_qpos=torch.ones(batch_size, 6)) + + merged = previous.masked_merge( + candidate, torch.tensor([True, False], dtype=torch.bool) + ) + + assert merged.get_held_object("arm").env_mask.tolist() == [False, True] + + def test_masked_merge_preserves_coordinated_hold_in_failed_envs(self): + batch_size = 2 + held = CoordinatedHeldObjectState( + semantics=ObjectSemantics(affordance=Affordance(), geometry={}), + left_object_to_eef=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), + right_object_to_eef=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), + left_grasp_xpos=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), + right_grasp_xpos=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), + ) + previous = WorldState( + last_qpos=torch.zeros(batch_size, 12), + coordinated_held_objects={("left_arm", "right_arm"): held}, + ) + candidate = WorldState(last_qpos=torch.ones(batch_size, 12)) + + merged = previous.masked_merge( + candidate, torch.tensor([True, False], dtype=torch.bool) + ) + + coordinated = merged.get_coordinated_held_object("left_arm", "right_arm") + assert coordinated is not None + assert coordinated.env_mask.tolist() == [False, True] + + def test_masked_merge_updates_qpos_per_environment(self): + previous = WorldState(last_qpos=torch.zeros(2, 3)) + candidate = WorldState(last_qpos=torch.ones(2, 3)) + + merged = previous.masked_merge( + candidate, torch.tensor([True, False], dtype=torch.bool) + ) + + assert torch.equal(merged.last_qpos[0], torch.ones(3)) + assert torch.equal(merged.last_qpos[1], torch.zeros(3)) + class TestActionResult: - def test_shape_contract(self): + def test_bool_success_is_normalized_to_per_environment_tensor(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.success.tolist() == [True, True] assert res.trajectory.shape == (2, 10, 8) assert res.next_state is ws + def test_scalar_tensor_success_is_normalized(self): + res = ActionResult( + success=torch.tensor(False), + trajectory=torch.zeros(2, 0, 3), + next_state=WorldState(last_qpos=torch.zeros(2, 3)), + ) + assert res.success.tolist() == [False, False] + + def test_rejects_non_boolean_success_tensor(self): + with pytest.raises(TypeError, match="torch.bool"): + ActionResult( + success=torch.ones(2), + trajectory=torch.zeros(2, 0, 3), + next_state=WorldState(last_qpos=torch.zeros(2, 3)), + ) + + def test_rejects_wrong_success_shape(self): + with pytest.raises(ValueError, match="success"): + ActionResult( + success=torch.ones(3, dtype=torch.bool), + trajectory=torch.zeros(2, 0, 3), + next_state=WorldState(last_qpos=torch.zeros(2, 3)), + ) + + def test_rejects_trajectory_state_dof_mismatch(self): + with pytest.raises(ValueError, match="batch/DoF"): + ActionResult( + success=torch.ones(2, dtype=torch.bool), + trajectory=torch.zeros(2, 4, 4), + next_state=WorldState(last_qpos=torch.zeros(2, 3)), + ) + class TestActionCfg: def test_defaults(self): @@ -332,3 +491,8 @@ def test_defaults(self): assert cfg.interpolation_type == "linear" assert cfg.velocity_limit is None assert cfg.acceleration_limit is None + assert cfg.plan_opts is None + + def test_rejects_unsupported_interpolation_type(self): + with pytest.raises(ValueError, match="interpolation_type"): + ActionCfg(interpolation_type="cubic") diff --git a/tests/sim/atomic_actions/test_engine.py b/tests/sim/atomic_actions/test_engine.py index 1d25fdb6e..128dd190f 100644 --- a/tests/sim/atomic_actions/test_engine.py +++ b/tests/sim/atomic_actions/test_engine.py @@ -147,6 +147,10 @@ def test_register_with_explicit_name_overrides_cfg(self): self.engine.register(action, name="custom") assert "custom" in self.engine.actions + def test_rejects_initial_state_with_wrong_batch_size(self): + with pytest.raises(ValueError, match="batch size"): + self.engine.run([], state=WorldState(last_qpos=torch.zeros(1, TOTAL_DOF))) + def test_run_concatenates_trajectories(self): a = _fake_action("a", EndEffectorPoseTarget) b = _fake_action("b", EndEffectorPoseTarget) diff --git a/tests/sim/atomic_actions/test_engine_per_env.py b/tests/sim/atomic_actions/test_engine_per_env.py index 622d93e3f..186dbddfa 100644 --- a/tests/sim/atomic_actions/test_engine_per_env.py +++ b/tests/sim/atomic_actions/test_engine_per_env.py @@ -22,11 +22,14 @@ import pytest from unittest.mock import Mock +from embodichain.lab.sim.atomic_actions.affordance import Affordance from embodichain.lab.sim.atomic_actions import EndEffectorPoseTarget from embodichain.lab.sim.atomic_actions.core import ( ActionCfg, ActionResult, AtomicAction, + HeldObjectState, + ObjectSemantics, WorldState, ) from embodichain.lab.sim.atomic_actions.engine import AtomicActionEngine @@ -52,6 +55,30 @@ def execute(self, target, state): ) +class _HeldStateAction(_StubAction): + def __init__(self, mg, success_vec, *, set_held): + super().__init__(mg, success_vec) + self._set_held = set_held + + def execute(self, target, state): + result = super().execute(target, state) + held_objects = {} + if self._set_held: + batch_size = state.batch_size + held_objects["arm"] = HeldObjectState( + semantics=ObjectSemantics( + affordance=Affordance(), geometry={}, label="test-object" + ), + object_to_eef=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), + grasp_xpos=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), + ) + return ActionResult( + success=result.success, + trajectory=result.trajectory, + next_state=result.next_state.with_updates(held_objects=held_objects), + ) + + class TestRunPerEnv: def test_failed_env_holds(self): mg = Mock() @@ -75,3 +102,41 @@ def test_failed_env_holds(self): # 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) + + def test_failed_env_does_not_acquire_successful_env_held_state(self): + mg = Mock() + mg.robot.get_qpos = lambda: torch.zeros(3, 3) + mg.robot.dof = 3 + mg.device = torch.device("cpu") + engine = AtomicActionEngine(mg) + engine.register( + _HeldStateAction(mg, [True, False, True], set_held=True), name="pick" + ) + + _, _, state = engine.run([("pick", EndEffectorPoseTarget(xpos=torch.eye(4)))]) + + assert state.get_held_object("arm").env_mask.tolist() == [True, False, True] + + def test_failed_env_preserves_held_state_when_successful_envs_release(self): + mg = Mock() + mg.robot.get_qpos = lambda: torch.zeros(3, 3) + mg.robot.dof = 3 + mg.device = torch.device("cpu") + engine = AtomicActionEngine(mg) + engine.register( + _HeldStateAction(mg, [True, False, True], set_held=False), name="place" + ) + held = HeldObjectState( + semantics=ObjectSemantics( + affordance=Affordance(), geometry={}, label="test-object" + ), + object_to_eef=torch.eye(4).unsqueeze(0).repeat(3, 1, 1), + grasp_xpos=torch.eye(4).unsqueeze(0).repeat(3, 1, 1), + ) + initial = WorldState(last_qpos=torch.zeros(3, 3), held_objects={"arm": held}) + + _, _, state = engine.run( + [("place", EndEffectorPoseTarget(xpos=torch.eye(4)))], state=initial + ) + + assert state.get_held_object("arm").env_mask.tolist() == [False, True, False] 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..3d742682b 100644 --- a/tests/sim/atomic_actions/test_trajectory_motion_source.py +++ b/tests/sim/atomic_actions/test_trajectory_motion_source.py @@ -23,6 +23,7 @@ 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 @@ -74,6 +75,32 @@ 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 = ActionCfg( + motion_source="motion_gen", + control_part="arm", + 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 From 26d3419d3100a39d4045406f53507c8b46b4bf2c Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 2 Aug 2026 09:59:30 +0000 Subject: [PATCH 2/5] refactor(atomic-actions): add typed planning and recovery runtime --- .agents/skills/add-atomic-action/SKILL.md | 439 ++-- agent_context/MAP.yaml | 26 +- .../topics/atomic-actions/atomic-actions.md | 243 +- ...hain.lab.sim.atomic_actions.primitives.rst | 43 +- .../embodichain.lab.sim.atomic_actions.rst | 300 +-- ...dichain.lab.sim.atomic_actions.targets.rst | 21 - .../sim/atomic_actions/builtin_actions.md | 358 +-- .../overview/sim/atomic_actions/index.md | 409 +--- docs/source/tutorial/atomic_actions.rst | 355 +-- .../lab/sim/atomic_actions/__init__.py | 191 +- embodichain/lab/sim/atomic_actions/actions.py | 61 - .../lab/sim/atomic_actions/bindings.py | 106 + embodichain/lab/sim/atomic_actions/core.py | 1038 +++----- embodichain/lab/sim/atomic_actions/effects.py | 284 +++ embodichain/lab/sim/atomic_actions/engine.py | 343 ++- .../lab/sim/atomic_actions/execution.py | 685 ++++++ embodichain/lab/sim/atomic_actions/goals.py | 215 ++ .../lab/sim/atomic_actions/invocation.py | 79 + embodichain/lab/sim/atomic_actions/plans.py | 471 ++++ .../lab/sim/atomic_actions/policies.py | 120 + .../sim/atomic_actions/primitives/__init__.py | 40 +- .../sim/atomic_actions/primitives/_helpers.py | 11 +- .../primitives/coordinated_pickment.py | 190 +- .../primitives/coordinated_placement.py | 175 +- .../atomic_actions/primitives/hand_over.py | 136 +- .../primitives/move_end_effector.py | 152 +- .../primitives/move_held_object.py | 77 +- .../atomic_actions/primitives/move_joints.py | 157 +- .../sim/atomic_actions/primitives/pick_up.py | 96 +- .../sim/atomic_actions/primitives/place.py | 117 +- .../sim/atomic_actions/primitives/press.py | 81 +- embodichain/lab/sim/atomic_actions/state.py | 556 +++++ embodichain/lab/sim/atomic_actions/targets.py | 47 - .../lab/sim/atomic_actions/trajectory.py | 296 ++- examples/sim/planners/curobo_planner.py | 58 +- .../move_end_effector_benchmark.py | 30 +- .../move_held_object_benchmark.py | 78 +- .../atomic_action/move_joints_benchmark.py | 35 +- .../atomic_action/pickup_benchmark.py | 24 +- .../atomic_action/place_benchmark.py | 49 +- .../atomic_action/press_benchmark.py | 51 +- scripts/tutorials/atomic_action/assemble.py | 37 +- .../atomic_action/coordinated_pickment.py | 26 +- .../atomic_action/coordinated_placement.py | 111 +- scripts/tutorials/atomic_action/hand_over.py | 41 +- .../atomic_action/move_end_effector.py | 32 +- .../atomic_action/move_held_object.py | 52 +- .../tutorials/atomic_action/move_joints.py | 28 +- scripts/tutorials/atomic_action/pickup.py | 24 +- scripts/tutorials/atomic_action/place.py | 36 +- scripts/tutorials/atomic_action/press.py | 41 +- .../test_action_result_success.py | 49 - tests/sim/atomic_actions/test_actions.py | 2096 ++++++----------- tests/sim/atomic_actions/test_core.py | 600 ++--- .../test_curobo_motion_source_e2e.py | 38 +- tests/sim/atomic_actions/test_engine.py | 347 ++- .../sim/atomic_actions/test_engine_per_env.py | 388 ++- .../atomic_actions/test_motion_source_e2e.py | 34 +- .../test_trajectory_motion_source.py | 30 +- tests/sim/planners/test_curobo_planner.py | 64 +- 60 files changed, 6540 insertions(+), 5777 deletions(-) delete mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.targets.rst delete mode 100644 embodichain/lab/sim/atomic_actions/actions.py create mode 100644 embodichain/lab/sim/atomic_actions/bindings.py create mode 100644 embodichain/lab/sim/atomic_actions/effects.py create mode 100644 embodichain/lab/sim/atomic_actions/execution.py create mode 100644 embodichain/lab/sim/atomic_actions/goals.py create mode 100644 embodichain/lab/sim/atomic_actions/invocation.py create mode 100644 embodichain/lab/sim/atomic_actions/plans.py create mode 100644 embodichain/lab/sim/atomic_actions/policies.py create mode 100644 embodichain/lab/sim/atomic_actions/state.py delete mode 100644 embodichain/lab/sim/atomic_actions/targets.py delete mode 100644 tests/sim/atomic_actions/test_action_result_success.py diff --git a/.agents/skills/add-atomic-action/SKILL.md b/.agents/skills/add-atomic-action/SKILL.md index 20b13b301..d2e3b00a5 100644 --- a/.agents/skills/add-atomic-action/SKILL.md +++ b/.agents/skills/add-atomic-action/SKILL.md @@ -1,341 +1,232 @@ --- 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. 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 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` | +| 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 = action.plan(invocation: ActionInvocation[Goal], context: PlanningContext) +``` -import torch +Do not add compatibility code for `ActionTarget`, `WorldState`, `ActionResult`, +`execute()`, or `AtomicActionEngine.run()`. -from embodichain.utils import configclass -from embodichain.lab.sim.atomic_actions.core import ActionCfg +## 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. -@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" +```python +from dataclasses import dataclass +from typing import ClassVar + +import torch + + +@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 implementation configuration -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`. +Extend `ActionCfg` directly with `@configclass`. Keep only implementation-owned +behavior such as distances, gripper positions, grasp constraints, and phase +split counts. + +Do not put `motion_source`, planner choice, sample count, control period, +velocity limits, collision policy, or recovery thresholds in the action config; +those belong to `MotionPolicy` or `RecoveryPolicy` on the invocation. ```python -from dataclasses import dataclass +from embodichain.lab.sim.atomic_actions import ActionCfg +from embodichain.utils import configclass + -@dataclass(frozen=True) -class PushTarget: - contact_pose: torch.Tensor # (4, 4) or (n_envs, 4, 4) EEF contact pose +@configclass +class PushCfg(ActionCfg): + name: str = "push" + push_distance: float = 0.05 ``` -### 3. Implement the action class +## 3. Implement the planner -Subclass `AtomicAction` directly, declare `TargetType`, compose a `TrajectoryBuilder`, -and implement `execute`. +Inherit `AtomicAction[PushGoal]` directly. Declare stable metadata and resolve +resources from semantic binding roles. ```python -from __future__ import annotations - -import torch 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 ( + ActionInvocation, + ActionPlan, AtomicAction, - WorldState, + PlanningContext, + StateDelta, + TrajectoryBuilder, ) -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 +class Push(AtomicAction[PushGoal]): + skill_id: ClassVar[str] = "push" + GoalType: ClassVar[type] = PushGoal + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) - def __init__(self, motion_generator, cfg: PushCfg | None = None): + def __init__(self, motion_generator, cfg: PushCfg | None = 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( + def plan( + self, + invocation: ActionInvocation[PushGoal], + context: PlanningContext, + ) -> ActionPlan: + goal = self.require_goal(invocation) + control_part = invocation.binding.manipulator("primary") + joint_ids = self.robot.get_joint_ids(name=control_part) + start_qpos = context.robot.qpos[:, joint_ids] + + # Build planner states and generate controlled-joint motion using + # invocation.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, + invocation.motion_policy.sample_count, + control_part=control_part, + arm_dof=len(joint_ids), + cfg=invocation.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=invocation.motion_policy.control_dt, ) - full[:, :, :] = state.last_qpos.unsqueeze(1) - full[:, :, self.arm_joint_ids] = arm_traj - - return ActionResult( + return self.build_plan( + invocation, + 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 +- 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(invocation, 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(motion_generator, PushCfg())) ``` -### 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", -] +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,)) ``` -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. +Use `engine.start(...).tick(...)` instead when dynamic scene updates or online +error recovery are required. -### 6. Update the supported actions table +## 5. Export and document -Add a row to the table in `docs/source/overview/sim/atomic_actions/builtin_actions.md`: +Export the goal, config, and action from: -```markdown -| `Push` | Single | `PushTarget` — contact pose | Approach → push forward | Add a demo asset or `N/A` | -``` +1. `embodichain/lab/sim/atomic_actions/primitives/__init__.py` +2. `embodichain/lab/sim/atomic_actions/__init__.py` -### 7. Write a test +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. -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. +## 6. Test behavior -```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 -``` +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 -## 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) | +| 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 planner/recovery knobs in action config | Move them to invocation policies. | +| 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..81dbd04bf 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,32 @@ topics: - AtomicAction - AtomicActionEngine - TrajectoryBuilder - - ActionResult - - WorldState + - ActionPlan + - PlanningContext + - ExecutionSession + - StateDelta - held_objects - - ObjectActionTarget - - PlaceTarget - - PressTarget - - target ownership + - ActionBinding + - SceneEntityPose + - dynamic goal + - error recovery - ActionCfg - - motion_source + - MotionPolicy + - RecoveryPolicy - plan_arm_traj - register_action paths: - topics/atomic-actions/atomic-actions.md source_of_truth: - embodichain/lab/sim/atomic_actions/core.py - - embodichain/lab/sim/atomic_actions/targets.py + - embodichain/lab/sim/atomic_actions/goals.py + - embodichain/lab/sim/atomic_actions/bindings.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..9441ee8a0 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -1,180 +1,99 @@ -# 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 = action.plan(invocation: ActionInvocation, context: PlanningContext) +``` -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. +There is no `ActionTarget`, `WorldState`, `ActionResult`, `execute()`, or +`AtomicActionEngine.run()` compatibility surface. -``` -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 -``` +`ActionInvocation` separates: -All tensor shapes carry a leading batch dim `B = n_envs`. +- an action-owned typed goal (`goal_kind` is its stable discriminator); +- `ActionBinding`, which maps semantic roles to concrete robot resources; +- reusable `MotionPolicy` planner/timing choices; +- bounded `RecoveryPolicy` thresholds and retry budgets. -## Core Types +`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`. -### Typed Targets +## Static compilation -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. +Register configured action instances by their class-level stable `skill_id`, +then call: -`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 +compiled = engine.compile(invocations, context=None) +``` -| 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 +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. + +## 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 -engine = AtomicActionEngine(motion_generator) -engine.register(MoveEndEffector(motion_generator, cfg=MoveEndEffectorCfg())) -success, traj, final_state = engine.run(steps=[("move_end_effector", target)]) +session = engine.start(invocations, initial_context) +tick = session.tick(latest_context, effect_success=None) ``` -`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())`. +An `ExecutionSession` emits at most one `JointCommand` per tick and monitors: + +- joint tracking error against the previous command; +- translation/rotation drift of referenced scene entities; +- phase timeout; +- planner and semantic-effect failure. + +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. + +## 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. Action configs retain +only implementation-specific behavior such as gripper poses, phase splits, +lift distances, and grasp constraints. -## 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`, `NamedJointPositionGoal` | 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. Declare `skill_id`, `GoalType`, and required semantic roles on the action. +3. Validate with `require_goal(invocation)`. +4. Plan from `context.robot.qpos`; never read an implicit live start state. +5. Return full-robot positions or a `TimedTrajectory` through `build_plan()`. +6. Declare symbolic changes with `StateDelta`; do not mutate context or commit + physical effects during planning. +7. 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..c245c4ada 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 @@ -8,13 +8,10 @@ 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 an :class:`~embodichain.lab.sim.atomic_actions.ActionInvocation` 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 @@ -36,21 +33,23 @@ resulting trajectories along the time axis. CoordinatedPickment CoordinatedPlacementCfg CoordinatedPlacement + HandOverCfg + HandOver - .. rubric:: Built-in Target Contracts + .. rubric:: Built-in Goal Contracts .. autosummary:: - EndEffectorPoseTarget - JointPositionTarget - NamedJointPositionTarget - GraspTarget - HeldObjectPoseTarget - PlaceTarget - PressTarget - CoordinatedPickTarget - CoordinatedPickmentTarget - CoordinatedPlacementTarget + EndEffectorPoseGoal + JointPositionGoal + NamedJointPositionGoal + GraspGoal + HeldObjectPoseGoal + PlaceGoal + AssembleGoal + PressGoal + CoordinatedPickGoal + CoordinatedPlacementGoal .. currentmodule:: embodichain.lab.sim.atomic_actions.primitives @@ -117,3 +116,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..5e7990410 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,143 @@ embodichain.lab.sim.atomic_actions .. automodule:: embodichain.lab.sim.atomic_actions - .. rubric:: Classes + .. rubric:: Planning contracts + + .. autosummary:: + + ActionGoal + ActionBinding + ActionInvocation + 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 + NamedJointPositionGoal + 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:: ActionInvocation + :members: -.. autoclass:: Affordance - :members: - :show-inheritance: +.. autoclass:: MotionPolicy + :members: + :exclude-members: __init__, copy, replace, to_dict -.. autoclass:: AntipodalAffordance - :members: - :show-inheritance: +.. autoclass:: RecoveryPolicy + :members: + :exclude-members: __init__, copy, replace, to_dict -.. autoclass:: InteractionPoints - :members: - :show-inheritance: +.. autoclass:: PlanningContext + :members: -.. autoclass:: ObjectSemantics - :members: - :show-inheritance: +.. autoclass:: RobotObservation + :members: -.. autoclass:: ActionTarget - :members: - :show-inheritance: +.. autoclass:: TaskState + :members: -.. autoclass:: ObjectActionTarget - :members: - :show-inheritance: +.. autoclass:: SceneSnapshot + :members: -.. autoclass:: EndEffectorPoseTarget - :members: - :show-inheritance: +.. autoclass:: SceneEntityPose + :members: -.. autoclass:: PlaceTarget - :members: - :show-inheritance: +.. autoclass:: StateDelta + :members: -.. autoclass:: PressTarget - :members: - :show-inheritance: +.. autoclass:: TimedTrajectory + :members: -.. autoclass:: JointPositionTarget - :members: - :show-inheritance: +.. autoclass:: ActionPlan + :members: -.. autoclass:: NamedJointPositionTarget - :members: - :show-inheritance: +Engine and execution +-------------------- -.. autoclass:: GraspTarget - :members: - :show-inheritance: +.. autoclass:: AtomicAction + :members: -.. autoclass:: HeldObjectPoseTarget - :members: - :show-inheritance: +.. autoclass:: AtomicActionEngine + :members: -.. autoclass:: CoordinatedPickTarget - :members: - :show-inheritance: +.. autoclass:: ExecutionSession + :members: -.. autodata:: CoordinatedPickmentTarget +.. autoclass:: ExecutionTick + :members: -.. autoclass:: CoordinatedPlacementTarget - :members: - :show-inheritance: +.. autoclass:: JointCommand + :members: -.. autodata:: Target +.. autoclass:: ExecutionEvent + :members: -.. autodata:: BuiltinTarget +Semantic objects and helpers +---------------------------- + +.. autoclass:: ObjectSemantics + :members: .. autoclass:: HeldObjectState - :members: - :show-inheritance: + :members: .. autoclass:: CoordinatedHeldObjectState - :members: - :show-inheritance: - -.. autoclass:: WorldState - :members: - :show-inheritance: - -.. autoclass:: ActionResult - :members: - :show-inheritance: - -.. autoclass:: ActionCfg - :members: - :exclude-members: __init__, copy, replace, to_dict - -.. autoclass:: AtomicAction - :members: - :show-inheritance: - -Trajectory helpers ------------------- + :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 ------------------ - -.. autoclass:: AtomicActionEngine - :members: - :show-inheritance: - -.. autofunction:: register_action - -.. autofunction:: unregister_action - -.. autofunction:: get_registered_actions + :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..8372c6d08 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -1,314 +1,44 @@ -(builtin_actions)= - -# Built-in Actions - -```{currentmodule} embodichain.lab.sim.atomic_actions -``` - -The following actions are available out of the box: - -```{note} -The built-in atomic actions currently support gripper-based manipulation only. Dexterous-hand manipulation is not supported yet. -``` - -| Action | Arm | Target type | Motion phases | Demo | -|---|---|---|---|---| -| `MoveEndEffector` | Single | `EndEffectorPoseTarget` — EEF pose | Move end-effector to pose | MoveEndEffector | -| `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 | - ---- - -## `MoveEndEffector` - -Moves the end-effector to a target pose in free space. - -| 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 | - -**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](../../../_static/atomic_actions/move_end_effector.gif) - ---- - -## `MoveJoints` - -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. - -| Config field | Default | Description | -|---|---|---| -| `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 | - -**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`. - -![MoveJoints demo](../../../_static/atomic_actions/move_joints.gif) - ---- - -## `PickUp` - -Three-phase grasp motion: *approach → close gripper → lift*. - -| 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) - ---- - -## `MoveHeldObject` - -Moves a held object to an object-centric target pose while preserving the grasp. It requires -the `HeldObjectState` populated by a prior `PickUp` (read from -`WorldState.held_objects[control_part]`) -and preserves it in its successor state. - -`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. - -| 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 | - -**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. - -![MoveHeldObject demo](../../../_static/atomic_actions/move_held_object.gif) - ---- - -## `Place` - -Three-phase release motion: *lower → open gripper → retract*. Mirrors `PickUp`. - -`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`. - -| 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)` | - -**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`. - -**Tutorial:** `scripts/tutorials/atomic_action/assemble.py` - -![Assemble demo](../../../_static/atomic_actions/assemble.gif) - ---- - -## `Press` - -Three-phase contact motion: *close gripper → press down → return*. This is useful -for button-like or contact-based interactions where the end-effector should reach a -target pose and then return to the pre-press arm pose. - -`Press` does not create or clear `WorldState.held_objects`; it preserves the state -threaded into it. - -| Config field | Default | Description | -|---|---|---| -| `hand_close_qpos` | `None` | **Required.** Gripper closed joint positions | -| `hand_control_part` | `"hand"` | Robot control part for the gripper | -| `hand_interp_steps` | `5` | Waypoints for the gripper close phase | -| `sample_interval` | `80` | Total waypoints across all three phases | - -**Target:** `PressTarget(xpos=...)` — the EEF pose to press, a `torch.Tensor` -of shape `(4, 4)` or `(n_envs, 4, 4)`. - -![Press demo](../../../_static/atomic_actions/press.gif) - ---- - -## `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. - -| Config field | Default | Description | -|---|---|---| -| `control_part` | `"dual_arm"` | Combined arm control part | -| `left_arm_control_part` / `right_arm_control_part` | `"left_arm"` / `"right_arm"` | Arm control parts for each grasp | -| `left_hand_control_part` / `right_hand_control_part` | `"left_hand"` / `"right_hand"` | Hand control parts for each gripper | -| `pre_grasp_distance` | `0.10` | Distance to back away from each grasp TCP | -| `lift_height` | `0.08` | World-Z lift distance before moving to the target pose | -| `object_motion_keyframes` | `6` | Sparse object-pose IK keyframes for synchronized motion | -| `sample_interval` | `120` | Total waypoints across all phases | - -**Target:** `CoordinatedPickTarget(...)` with a target object pose, object -semantics, and left/right object-to-EEF transforms. -It inherits the same `ObjectActionTarget(semantics)` base as `GraspTarget`, but -keeps the dual-arm pose fields in its own action-specific contract. - -`CoordinatedPickmentTarget` remains a compatibility alias. - -**Tutorial:** `scripts/tutorials/atomic_action/coordinated_pickment.py` - -![CoordinatedPickment demo](../../../_static/atomic_actions/coordinated_pickment.gif) - ---- - -## `CoordinatedPlacement` - -Dual-arm object-centric placement. The support arm moves its held object to a lower -target pose and keeps its gripper closed. The placing arm moves its held object to -the aligned upper target pose, optionally opens the placing hand, then lifts away. - -`CoordinatedPlacement` reads both held objects from -`WorldState.held_objects`, keyed by `placing_arm_control_part` and -`support_arm_control_part`. The target contains desired poses and per-call -overrides only. - -| Config field | Default | Description | -|---|---|---| -| `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) - ---- - -## `HandOver` - -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. - -`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. - -| 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) +# Built-in atomic actions + +All built-ins implement `plan(invocation, context) -> ActionPlan`. Motion and +recovery settings belong to the invocation rather than each action config. + +| Skill ID | Goal | Semantic roles | Expected task effect | +|---|---|---|---| +| `move_end_effector` | `EndEffectorPoseGoal` | manipulator `primary` | none | +| `move_joints` | `JointPositionGoal` or `NamedJointPositionGoal` | manipulator `primary` | none | +| `pick_up` | `GraspGoal` | manipulator/end effector `primary` | attach object | +| `move_held_object` | `HeldObjectPoseGoal` | manipulator/end effector `primary` | none | +| `place` | `PlaceGoal` or `AssembleGoal` | manipulator/end effector `primary` | detach object | +| `press` | `PressGoal` | manipulator/end effector `primary` | none | +| `coordinated_pickment` | `CoordinatedPickGoal` | `left`, `right` | create coordinated attachment | +| `coordinated_placement` | `CoordinatedPlacementGoal` | `placing`, `support` | update/remove attachments | +| `hand_over` | `GraspGoal` | `source`, `destination` | transfer attachment | + +## Pose goals + +Pose-valued goals accept explicit tensors. Selected goals also accept a +`SceneEntityPose(entity_id, relative_pose=...)`, which is resolved from every +new `SceneSnapshot` during planning or replanning. Explicit tensors may use +`(4, 4)` or `(B, 4, 4)` shapes; waypoint-capable goals additionally accept +`(B, N, 4, 4)`. + +## Configuration boundaries + +Action configs contain implementation-owned behavior such as gripper open/close +positions, phase split counts, lift distance, and grasp-selection constraints. +They do not contain planner choice, motion source, trajectory sample count, +velocity limits, recovery budgets, or dynamic-goal thresholds. Those reusable +choices live in `MotionPolicy` and `RecoveryPolicy`. + +`MoveJoints` and `MoveEndEffector` resolve the concrete manipulator entirely from +`ActionBinding`. Complex manipulation skills currently validate their semantic +bindings against the configured hardware resources used by their phase-specific +hand parameters. + +## Planning versus physical success + +`ActionPlan.plan_success` reports whether a valid motion was produced per +environment. It does not prove that a grasp, release, contact, or handover +occurred. Such actions declare a `StateDelta`; an `ExecutionSession` commits it +only after the caller supplies a successful semantic-effect verification mask. diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index b80ef1bb1..286023b90 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -1,369 +1,74 @@ -# Atomic Actions +# Atomic 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. - -```{note} -Atomic actions currently support gripper-based manipulation only. Dexterous-hand manipulation is not supported yet. -``` - -## Design Overview - -The module is organized into three layers: - -``` -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) -``` - -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: - -``` -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)) -``` - -### Core Concepts - -**`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 - -**`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. - -**Typed targets** describe *where* an action should go. Each one is a small frozen, -identity-equality dataclass inheriting the open `ActionTarget` marker. Every action -declares the target type, or tuple of target types, it accepts via its `TargetType` -class variable. Object-centric actions can share the semantic-object contract -provided by `ObjectActionTarget`, while keeping their pose roles action-specific: - -| Target | Constructor | Used by | -|---|---|---| -| `EndEffectorPoseTarget` | `EndEffectorPoseTarget(xpos)` | `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 | -|---|---|---| -| `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 - -```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, -) - -# 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)}, -) - -# 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, - ), - geometry={}, - label="mug", - entity=obj, -) +```{toctree} +:hidden: -# 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")), - ] -) -# traj: (n_envs, n_waypoints, robot.dof) +builtin_actions ``` ---- - -## How to Extend: Adding a Custom Action - -You can add any motion primitive by subclassing `AtomicAction`, composing a -`TrajectoryBuilder` for the shared planning math, and registering an instance with the engine. -Built-in primitives live one action per module under -`embodichain/lab/sim/atomic_actions/primitives/`, while -`embodichain.lab.sim.atomic_actions` remains the public import surface and -`embodichain.lab.sim.atomic_actions.actions` stays as a compatibility re-export. - -### Step 1 — Define the config - -```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 +Atomic actions turn typed, grounded skill requests into full-robot timed motion. +Planning is side-effect free and execution is incremental. + +```text +Action Agent / task graph + | + | semantic skill call + v +grounder + capability binder + | + | ActionInvocation + v +AtomicAction.plan(invocation, PlanningContext) + | + | ActionPlan + StateDelta + +------------------------------+ + | | + v v +AtomicActionEngine.compile ExecutionSession.tick +(fixed-scene/offline) (dynamic/closed-loop) ``` -### Step 2 — Implement the action - -```python -from dataclasses import dataclass -import torch -from typing import ClassVar -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 +## Contracts - 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)") +- `ActionGoal`: structural protocol implemented by action-owned frozen goal + dataclasses. There is no common target object or closed union. +- `ActionBinding`: semantic role to robot-resource mapping. Goals do not carry + arm or hand names. +- `MotionPolicy`: planner, interpolation, sample count, timing, and limits. +- `RecoveryPolicy`: replan/retry budgets, tracking and dynamic-goal thresholds, + and phase timeout. +- `PlanningContext`: measured `RobotObservation`, verified `TaskState`, versioned + `SceneSnapshot`, and stable environment IDs. +- `ActionPlan`: one or more scene-bound phases, timed trajectories, completion + conditions, diagnostics, and uncommitted `StateDelta` effects. -class Push(AtomicAction[PushTarget]): - TargetType: ClassVar[type] = PushTarget +## Static and dynamic use - 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] +`AtomicActionEngine.compile()` plans a fixed sequence and returns a +`CompiledTrajectory`. It projects terminal qpos and expected task effects only +inside the returned context; it never changes simulator state. - 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(), - ), - ) -``` +`AtomicActionEngine.start()` creates an `ExecutionSession`. Each `tick()` takes +the latest context and emits at most one `JointCommand`. The session detects +tracking error, phase timeout, and movement of entities referenced by +`SceneEntityPose`, then replans from the latest observation within the configured +budget. Non-empty symbolic effects require external verification before commit. -### Step 3 — Register and use - -Register an instance with the engine so it can be referenced by name in `run()`: +## Example ```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))] +binding = ActionBinding(manipulators={"primary": "left_arm"}) +invocation = ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal(target_pose), + binding=binding, + motion_policy=MotionPolicy(sample_count=80), ) -``` -To publish an action class for third-party discovery (independent of any engine instance), -use the global registry: - -```python -from embodichain.lab.sim.atomic_actions import register_action, unregister_action, get_registered_actions - -register_action("push", Push) # registers the class under "push" -unregister_action("push") # removes it -all_actions = get_registered_actions() # dict[str, type[AtomicAction]] +engine = AtomicActionEngine(motion_generator) +engine.register(MoveEndEffector(motion_generator, MoveEndEffectorCfg())) +compiled = engine.compile((invocation,)) ``` -> **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. - ---- - -```{toctree} -:maxdepth: 1 - -builtin_actions -``` - -## 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 -``` +See [Built-in actions](builtin_actions.md) for the shipped skill catalog and +[the tutorial](../../../tutorial/atomic_actions.rst) for closed-loop usage. diff --git a/docs/source/tutorial/atomic_actions.rst b/docs/source/tutorial/atomic_actions.rst index f5d6f53c9..c2777124e 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -1,293 +1,148 @@ -.. _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: +Atomic actions are typed, side-effect-free motion planners. An action receives a +grounded :class:`~embodichain.lab.sim.atomic_actions.ActionInvocation` and the +latest :class:`~embodichain.lab.sim.atomic_actions.PlanningContext`, then returns +an :class:`~embodichain.lab.sim.atomic_actions.ActionPlan`. -- ``move_end_effector.py`` -- ``move_joints.py`` -- ``pickup.py`` -- ``move_held_object.py`` -- ``place.py`` -- ``press.py`` -- ``coordinated_pickment.py`` -- ``coordinated_placement.py`` +The contracts deliberately separate four concerns: -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``: +* a **goal** describes what should happen; +* an **ActionBinding** maps semantic roles such as ``primary`` or ``source`` to + robot control resources; +* 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. -.. code-block:: bash +Static compilation +------------------ - 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 - -The concrete implementations are organized one primitive per module under -``embodichain/lab/sim/atomic_actions/primitives``. Public imports from -``embodichain.lab.sim.atomic_actions`` remain the recommended API, and -``embodichain.lab.sim.atomic_actions.actions`` stays as a compatibility -re-export surface. - -Typical Usage -------------- - -Setting up the engine -~~~~~~~~~~~~~~~~~~~~~ +Use :meth:`~embodichain.lab.sim.atomic_actions.AtomicActionEngine.compile` when +the scene is treated as fixed during planning: .. code-block:: python - import torch - from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - PickUp, PickUpCfg, - Place, PlaceCfg, - MoveEndEffector, MoveEndEffectorCfg, - MoveJoints, MoveJointsCfg, + EndEffectorPoseGoal, + MotionPolicy, + MoveEndEffector, + MoveEndEffectorCfg, ) - motion_gen = MotionGenerator(cfg=MotionGenCfg(...)) + engine = AtomicActionEngine(motion_generator) + engine.register(MoveEndEffector(motion_generator, MoveEndEffectorCfg())) - 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)}, + 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), ) + compiled = engine.compile((invocation,)) + trajectory = compiled.trajectory.positions - # 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)) +``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. -Defining object semantics -~~~~~~~~~~~~~~~~~~~~~~~~~ +Dynamic goals and closed-loop execution +--------------------------------------- + +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 ( - ObjectSemantics, - AntipodalAffordance, + EndEffectorPoseGoal, + RecoveryPolicy, + SceneEntityPose, ) - 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 + invocation = ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal( + xpos=SceneEntityPose("moving_tray", relative_pose=tray_to_tcp) ), - generator_cfg=GraspGeneratorCfg( - antipodal_sampler_cfg=AntipodalSamplerCfg( - n_sample=20000, max_length=0.088, min_length=0.003 - ) + binding=ActionBinding(manipulators={"primary": "left_arm"}), + recovery_policy=RecoveryPolicy( + max_replans=3, + tracking_error_threshold=0.05, + goal_translation_threshold=0.02, ), - force_reannotate=False, - ) - - 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 - ) - -Executing a pick-place-end-effector sequence -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. code-block:: python - - from embodichain.lab.sim.atomic_actions import ( - GraspTarget, - EndEffectorPoseTarget, - PlaceTarget, - ) - - place_xpos = ... # torch.Tensor [4, 4] — target placement pose - rest_xpos = ... # torch.Tensor [4, 4] — resting pose after placing - - is_success, trajectory, _ = engine.run( - steps=[ - ("pick_up", GraspTarget(semantics=semantics)), - ("place", PlaceTarget(xpos=place_xpos)), - ("move_end_effector", EndEffectorPoseTarget(xpos=rest_xpos)), - ] - ) - # trajectory: [n_envs, n_waypoints, robot_dof] - - for i in range(trajectory.shape[1]): - robot.set_qpos(trajectory[:, i]) - sim.update(step=4) - -Moving in joint space -~~~~~~~~~~~~~~~~~~~~~ - -``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``. - -.. code-block:: python - - from embodichain.lab.sim.atomic_actions import ( - JointPositionTarget, - NamedJointPositionTarget, ) - home_qpos = torch.zeros(6, dtype=torch.float32, device=device) + session = engine.start((invocation,), initial_context) + while session.status.value == "running": + tick = session.tick(latest_context) + if tick.command is not None: + send_joint_command(tick.command) - is_success, trajectory, _ = engine.run( - steps=[ - ("move_joints", NamedJointPositionTarget(name="home")), - ("move_joints", JointPositionTarget(qpos=home_qpos)), - ] - ) +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. -Moving a held object -~~~~~~~~~~~~~~~~~~~~ +Task-state effects +------------------ -``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. +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 - from embodichain.lab.sim.atomic_actions import ( - MoveHeldObject, MoveHeldObjectCfg, - GraspTarget, HeldObjectPoseTarget, - ) - - 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)) + 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) - 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) +This prevents a successful trajectory plan from being mistaken for a successful +physical grasp or release. - is_success, trajectory, _ = engine.run( - steps=[ - ("pick_up", GraspTarget(semantics=semantics)), - ("move_held_object", HeldObjectPoseTarget(object_target_pose=object_target_pose)), - ] - ) +Adding an action +---------------- -Registering custom actions -~~~~~~~~~~~~~~~~~~~~~~~~~~ +Define an action-owned frozen goal dataclass with a stable ``goal_kind``. Then +implement ``plan(invocation, 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 ----------------------- + @dataclass(frozen=True, slots=True) + class PushGoal: + goal_kind: ClassVar[str] = "push" + contact_pose: torch.Tensor + + class Push(AtomicAction[PushGoal]): + skill_id: ClassVar[str] = "push" + GoalType: ClassVar[type] = PushGoal + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + + def plan( + self, + invocation: ActionInvocation[PushGoal], + context: PlanningContext, + ) -> ActionPlan: + goal = self.require_goal(invocation) + # Resolve the bound resource, plan from context.robot.qpos, and + # return a full-robot TimedTrajectory or position tensor. + return self.build_plan( + invocation, + context, + success=success_mask, + trajectory=full_robot_positions, + ) -- ``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`. +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..e196d8664 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -14,13 +14,13 @@ # 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. +An action consumes an :class:`ActionInvocation` and 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 +31,147 @@ AssembleAffordance, InteractionPoints, ) -from .core import ( - ActionTarget, - ActionCfg, - ActionResult, - AtomicAction, - CoordinatedHeldObjectState, - HeldObjectState, - ObjectSemantics, - Target, - WorldState, -) +from .bindings import ActionBinding +from .core import ActionCfg, 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 +from .plans import ( + ActionPlan, + CompiledTrajectory, + CompletionCondition, + CompletionConditionKind, + PhaseSpec, + PlannedPhase, + PlannerDiagnostics, + TimedTrajectory, +) +from .policies import MotionPolicy, RecoveryPolicy from .primitives import ( - AssembleTarget, - CoordinatedPickTarget, + AssembleGoal, + CoordinatedPickGoal, CoordinatedPickment, CoordinatedPickmentCfg, - CoordinatedPickmentTarget, CoordinatedPlacement, CoordinatedPlacementCfg, - CoordinatedPlacementTarget, - EndEffectorPoseTarget, - GraspTarget, + CoordinatedPlacementGoal, + EndEffectorPoseGoal, + GraspGoal, HandOver, HandOverCfg, - HeldObjectPoseTarget, - JointPositionTarget, + HeldObjectPoseGoal, + JointPositionGoal, MoveEndEffector, MoveEndEffectorCfg, MoveHeldObject, MoveHeldObjectCfg, MoveJoints, MoveJointsCfg, - NamedJointPositionTarget, + NamedJointPositionGoal, PickUp, PickUpCfg, Place, PlaceCfg, - PlaceTarget, + PlaceGoal, Press, PressCfg, - PressTarget, + PressGoal, ) -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", + "ActionCfg", + "ActionGoal", + "ActionInvocation", + "ActionPlan", "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", + "CompiledTrajectory", + "CompletionCondition", + "CompletionConditionKind", + "CoordinatedHeldObjectState", + "CoordinatedPickGoal", "CoordinatedPickment", - "CoordinatedPlacement", - "HandOver", - "MoveEndEffector", - "MoveJoints", - "MoveHeldObject", - "PickUp", - "Place", - "Press", "CoordinatedPickmentCfg", + "CoordinatedPlacement", "CoordinatedPlacementCfg", + "CoordinatedPlacementGoal", + "EndEffectorPoseGoal", + "EntityState", + "ExecutionEvent", + "ExecutionEventKind", + "ExecutionSession", + "ExecutionStatus", + "ExecutionTick", + "GraspGoal", + "HandOver", "HandOverCfg", + "HeldObjectPoseGoal", + "HeldObjectState", + "InteractionPoints", + "JointPositionGoal", + "JointCommand", + "MotionPolicy", + "MoveEndEffector", "MoveEndEffectorCfg", - "MoveJointsCfg", + "MoveHeldObject", "MoveHeldObjectCfg", + "MoveJoints", + "MoveJointsCfg", + "NamedJointPositionGoal", + "ObjectActionGoal", + "ObjectSemantics", + "PhaseSpec", + "PickUp", "PickUpCfg", + "Place", "PlaceCfg", + "PlaceGoal", + "PlannedPhase", + "PlannerDiagnostics", + "PlanningContext", + "PoseGoalValue", + "Press", "PressCfg", - # Engine - "AtomicActionEngine", + "PressGoal", + "RecoveryPolicy", + "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..c17e536f0 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/bindings.py @@ -0,0 +1,106 @@ +# ---------------------------------------------------------------------------- +# 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-resource bindings for atomic actions.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Mapping + + +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 embodiment-specific control resources. + + The action and an agent-facing request refer to roles such as ``primary``, + ``source`` and ``destination``. Only the compiler or application binding + layer needs to know concrete robot resources such as ``left_arm``. + """ + + manipulators: Mapping[str, str] = field(default_factory=dict) + """Manipulator resources keyed by semantic role.""" + + end_effectors: Mapping[str, str] = field(default_factory=dict) + """End-effector resources 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 resource bound to ``role``. + + Args: + role: Semantic manipulator role. + + Returns: + Concrete robot control-resource name. + + 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 end-effector resource bound to ``role``. + + Args: + role: Semantic end-effector role. + + Returns: + Concrete robot control-resource name. + + 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 + + +__all__ = ["ActionBinding"] diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index a58c43735..6f49f52fa 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -14,866 +14,370 @@ # 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 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, GoalT +from .plans import ( + ActionPlan, + CompletionCondition, + CompletionConditionKind, + PhaseSpec, + PlannedPhase, + PlannerDiagnostics, + TimedTrajectory, +) if TYPE_CHECKING: - from embodichain.lab.sim.planners import MotionGenerator, PlanOptions + from embodichain.lab.sim.planners import MotionGenerator + + from .state import PlanningContext + +def resolve_runtime_device(device: torch.device | str) -> torch.device: + """Resolve an indexless CUDA device to the active concrete GPU index. -def _resolve_runtime_device(device: torch.device | str) -> torch.device: - """Resolve an indexless CUDA device to the active concrete GPU index.""" + 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 -# ============================================================================= -# ObjectSemantics -# ============================================================================= - - @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) - - -def _validate_pose_tensor( - value: torch.Tensor, - name: str, - *, - allow_waypoints: bool, -) -> None: - """Validate the environment-independent part of a pose tensor contract.""" - if not isinstance(value, torch.Tensor): - raise TypeError(f"{name} must be a torch.Tensor, got {type(value).__name__}.") - valid_dims = {2, 3, 4} if allow_waypoints else {2, 3} - if value.dim() not in valid_dims or value.shape[-2:] != (4, 4): - supported = "(4, 4), (n_envs, 4, 4)" - if allow_waypoints: - supported += ", or (n_envs, n_waypoint, 4, 4)" - raise ValueError( - f"{name} must have shape {supported}, got {tuple(value.shape)}." - ) - - -# ``Target`` 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 +@dataclass(frozen=True, slots=True) +class SkillDescriptor: + """Machine-readable metadata for one registered atomic skill.""" - -# ============================================================================= -# 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 - """Object-to-end-effector transform, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" - - grasp_xpos: torch.Tensor - """Grasp pose, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" - - env_mask: torch.Tensor | None = None - """Environments in which the held-object relation is active, shape ``(n_envs,)``.""" + skill_id: str + goal_type: type[Any] | tuple[type[Any], ...] + manipulator_roles: tuple[str, ...] = () + end_effector_roles: tuple[str, ...] = () + agent_visible: bool = True def __post_init__(self) -> None: - object_batch_size = _validate_held_pose( - self.object_to_eef, "HeldObjectState.object_to_eef" - ) - grasp_batch_size = _validate_held_pose( - self.grasp_xpos, "HeldObjectState.grasp_xpos" - ) - known_batch_sizes = { - size for size in (object_batch_size, grasp_batch_size) if size is not None - } - if len(known_batch_sizes) > 1: - raise ValueError( - "HeldObjectState pose tensors must use the same batch size, " - f"got {object_batch_size} and {grasp_batch_size}." - ) - if self.grasp_xpos.device != self.object_to_eef.device: - raise ValueError("HeldObjectState pose tensors must use the same device.") - batch_size = next(iter(known_batch_sizes), None) - self.env_mask = _normalize_optional_env_mask( - self.env_mask, - batch_size=batch_size, - device=self.object_to_eef.device, - name="HeldObjectState.env_mask", + 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.") + 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) -@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 - """Left object-to-end-effector transform, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" - - right_object_to_eef: torch.Tensor - """Right object-to-end-effector transform, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" - - left_grasp_xpos: torch.Tensor - """Left grasp pose, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" - - right_grasp_xpos: torch.Tensor - """Right grasp pose, shape ``(4, 4)`` or ``(n_envs, 4, 4)``.""" +@configclass +class ActionCfg: + """Base configuration for implementation-owned skill behavior.""" - env_mask: torch.Tensor | None = None - """Environments in which the coordinated hold is active, shape ``(n_envs,)``.""" + name: str = "default" def __post_init__(self) -> None: - pose_fields = { - "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, - } - batch_sizes = { - name: _validate_held_pose(value, f"CoordinatedHeldObjectState.{name}") - for name, value in pose_fields.items() - } - known_batch_sizes = {size for size in batch_sizes.values() if size is not None} - if len(known_batch_sizes) > 1: - raise ValueError( - "CoordinatedHeldObjectState pose tensors must use the same batch " - f"size, got {batch_sizes}." - ) - devices = {value.device for value in pose_fields.values()} - if len(devices) != 1: - raise ValueError( - "CoordinatedHeldObjectState pose tensors must use the same device." - ) - batch_size = next(iter(known_batch_sizes), None) - self.env_mask = _normalize_optional_env_mask( - self.env_mask, - batch_size=batch_size, - device=self.left_object_to_eef.device, - name="CoordinatedHeldObjectState.env_mask", - ) - + if not isinstance(self.name, str) or not self.name: + raise ValueError("name must be a non-empty string.") -def _validate_held_pose(value: torch.Tensor, name: str) -> int | None: - """Validate a held-state pose and return its explicit batch size, if any.""" - 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) with n_envs > 0, " - f"got {tuple(value.shape)}." - ) - return int(value.shape[0]) - - -def _normalize_optional_env_mask( - value: torch.Tensor | None, - *, - batch_size: int | None, - device: torch.device, - name: str, -) -> torch.Tensor | None: - """Normalize a mask when a held-state batch can already be inferred.""" - if batch_size is None and value is None: - return None - if batch_size is None: - 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.dim() != 1 or value.shape[0] == 0: - raise ValueError( - f"{name} must have shape (n_envs,) with n_envs > 0, " - f"got {tuple(value.shape)}." - ) - batch_size = int(value.shape[0]) - return _normalize_env_mask( - value, - batch_size=batch_size, - device=device, - name=name, - ) - - -def _normalize_env_mask( - value: torch.Tensor | None, - *, - batch_size: int, - device: torch.device, - name: str, -) -> torch.Tensor: - """Return an owned boolean environment 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_held_pose( - value: torch.Tensor, - *, - batch_size: int, - device: torch.device, - name: str, -) -> torch.Tensor: - """Resolve an optionally batched held-state pose to a world-state batch.""" - pose_batch_size = _validate_held_pose(value, name) - if value.device != device: - raise ValueError(f"{name} must use the same device as WorldState.last_qpos.") - if pose_batch_size is None: - return value.unsqueeze(0).expand(batch_size, -1, -1).clone() - if pose_batch_size != batch_size: - raise ValueError( - "Held-object state batch size must match WorldState.last_qpos; " - f"expected {batch_size}, got {pose_batch_size}." - ) - return value - - -def _normalize_held_object_state( - value: HeldObjectState, - *, - batch_size: int, - device: torch.device, -) -> HeldObjectState: - """Return a held-object state normalized to a world-state batch.""" - object_batch_size = _validate_held_pose( - value.object_to_eef, "HeldObjectState.object_to_eef" - ) - grasp_batch_size = _validate_held_pose( - value.grasp_xpos, "HeldObjectState.grasp_xpos" - ) - if ( - object_batch_size == batch_size - and grasp_batch_size == batch_size - and value.object_to_eef.device == device - and value.grasp_xpos.device == device - and isinstance(value.env_mask, torch.Tensor) - and value.env_mask.dtype == torch.bool - and value.env_mask.shape == (batch_size,) - and value.env_mask.device == device - ): - return value - return HeldObjectState( - semantics=value.semantics, - object_to_eef=_broadcast_held_pose( - value.object_to_eef, - batch_size=batch_size, - device=device, - name="HeldObjectState.object_to_eef", - ), - grasp_xpos=_broadcast_held_pose( - value.grasp_xpos, - batch_size=batch_size, - device=device, - name="HeldObjectState.grasp_xpos", - ), - env_mask=_normalize_env_mask( - value.env_mask, - batch_size=batch_size, - device=device, - name="HeldObjectState.env_mask", - ), - ) - - -def _normalize_coordinated_held_object_state( - value: CoordinatedHeldObjectState, - *, - batch_size: int, - device: torch.device, -) -> CoordinatedHeldObjectState: - """Return a coordinated-held state normalized to a world-state batch.""" - pose_fields = { - "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, - } - pose_batch_sizes = { - name: _validate_held_pose( - pose, - f"CoordinatedHeldObjectState.{name}", - ) - for name, pose in pose_fields.items() - } - if ( - all(size == batch_size for size in pose_batch_sizes.values()) - and all(pose.device == device for pose in pose_fields.values()) - and isinstance(value.env_mask, torch.Tensor) - and value.env_mask.dtype == torch.bool - and value.env_mask.shape == (batch_size,) - and value.env_mask.device == device - ): - return value - normalized_poses = { - name: _broadcast_held_pose( - pose, - batch_size=batch_size, - device=device, - name=f"CoordinatedHeldObjectState.{name}", - ) - for name, pose in pose_fields.items() - } - return CoordinatedHeldObjectState( - semantics=value.semantics, - **normalized_poses, - env_mask=_normalize_env_mask( - value.env_mask, - batch_size=batch_size, - device=device, - name="CoordinatedHeldObjectState.env_mask", - ), - ) - - -def _merge_held_object_state( - previous: HeldObjectState | None, - candidate: HeldObjectState | None, - update_mask: torch.Tensor, -) -> HeldObjectState | None: - """Merge one held-object entry using a per-environment update mask.""" - if previous is not None: - assert previous.env_mask is not None - if candidate is not None: - assert candidate.env_mask is not None - if previous is None and candidate is None: - return None - if previous is None: - assert candidate is not None - env_mask = candidate.env_mask & update_mask - if not env_mask.any(): - return None - return HeldObjectState( - semantics=candidate.semantics, - object_to_eef=candidate.object_to_eef, - grasp_xpos=candidate.grasp_xpos, - env_mask=env_mask, - ) - if candidate is None: - env_mask = previous.env_mask & ~update_mask - if not env_mask.any(): - return None - return HeldObjectState( - semantics=previous.semantics, - object_to_eef=previous.object_to_eef, - grasp_xpos=previous.grasp_xpos, - env_mask=env_mask, - ) - 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 control part " - "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_held_object_state( - previous: CoordinatedHeldObjectState | None, - candidate: CoordinatedHeldObjectState | None, - update_mask: torch.Tensor, -) -> CoordinatedHeldObjectState | None: - """Merge one coordinated-held entry using a per-environment update mask.""" - if previous is not None: - assert previous.env_mask is not None - if candidate is not None: - assert candidate.env_mask is not None - if previous is None and candidate is None: - return None - if previous is None: - assert candidate is not None - env_mask = candidate.env_mask & update_mask - if not env_mask.any(): - return None - return CoordinatedHeldObjectState( - semantics=candidate.semantics, - left_object_to_eef=candidate.left_object_to_eef, - right_object_to_eef=candidate.right_object_to_eef, - left_grasp_xpos=candidate.left_grasp_xpos, - right_grasp_xpos=candidate.right_grasp_xpos, - env_mask=env_mask, - ) - if candidate is None: - env_mask = previous.env_mask & ~update_mask - if not env_mask.any(): - return None - return CoordinatedHeldObjectState( - semantics=previous.semantics, - left_object_to_eef=previous.left_object_to_eef, - right_object_to_eef=previous.right_object_to_eef, - left_grasp_xpos=previous.left_grasp_xpos, - right_grasp_xpos=previous.right_grasp_xpos, - env_mask=env_mask, - ) +class AtomicAction(Generic[GoalT], ABC): + """Side-effect-free planner for one semantically meaningful robot skill.""" - 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 semantics for one control-part " - "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(slots=True, eq=False) -class WorldState: - """State the engine threads through a sequence of actions.""" - - last_qpos: torch.Tensor - """Robot joint positions at the start of the next action, shape [n_envs, robot.dof].""" - - held_objects: dict[str, HeldObjectState] = field(default_factory=dict) - """Objects held by individual control parts, keyed by control-part name.""" - - coordinated_held_objects: dict[tuple[str, str], CoordinatedHeldObjectState] = field( - default_factory=dict - ) - """Objects jointly held by two control parts, keyed by their ordered pair.""" - - def __post_init__(self) -> None: - if not isinstance(self.last_qpos, torch.Tensor): - raise TypeError("WorldState.last_qpos must be a torch.Tensor.") - if ( - self.last_qpos.dim() != 2 - or self.last_qpos.shape[0] == 0 - or self.last_qpos.shape[1] == 0 - ): - raise ValueError( - "WorldState.last_qpos must have shape (n_envs, robot_dof) with " - f"both dimensions non-zero, got {tuple(self.last_qpos.shape)}." - ) - held_objects: dict[str, HeldObjectState] = {} - for control_part, held in self.held_objects.items(): - if not isinstance(control_part, str) or not control_part: - raise TypeError( - "WorldState.held_objects keys must be non-empty strings." - ) - if not isinstance(held, HeldObjectState): - raise TypeError( - "WorldState.held_objects values must be HeldObjectState instances." - ) - held_objects[control_part] = _normalize_held_object_state( - held, - batch_size=self.batch_size, - device=self.last_qpos.device, - ) - coordinated_held_objects: dict[tuple[str, str], CoordinatedHeldObjectState] = {} - for control_parts, held in self.coordinated_held_objects.items(): - if ( - not isinstance(control_parts, tuple) - or len(control_parts) != 2 - or not all(isinstance(part, str) and part for part in control_parts) - ): - raise TypeError( - "WorldState.coordinated_held_objects keys must be pairs of " - "non-empty control-part names." - ) - if not isinstance(held, CoordinatedHeldObjectState): - raise TypeError( - "WorldState.coordinated_held_objects values must be " - "CoordinatedHeldObjectState instances." - ) - coordinated_held_objects[control_parts] = ( - _normalize_coordinated_held_object_state( - held, - batch_size=self.batch_size, - device=self.last_qpos.device, - ) - ) - self.held_objects = held_objects - self.coordinated_held_objects = coordinated_held_objects + skill_id: ClassVar[str] + """Stable registry identifier for this skill.""" - @property - def batch_size(self) -> int: - """Number of vectorized environments represented by this state.""" - return int(self.last_qpos.shape[0]) + GoalType: ClassVar[type[Any] | tuple[type[Any], ...]] + """Concrete goal dataclass or dataclasses accepted by this skill.""" - @property - def robot_dof(self) -> int: - """Number of robot joint-position columns represented by this state.""" - return int(self.last_qpos.shape[1]) + 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 - ), + 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 = resolve_runtime_device(self.robot.device) + + @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, + manipulator_roles=cls.manipulator_roles, + end_effector_roles=cls.end_effector_roles, + agent_visible=cls.agent_visible, ) - def masked_merge( - self, - candidate: WorldState, - update_mask: torch.Tensor, - ) -> WorldState: - """Merge a candidate successor for selected environments. + def require_goal(self, invocation: ActionInvocation[GoalT]) -> GoalT: + """Validate an invocation and return its concrete goal. Args: - candidate: Candidate successor returned by an atomic action. - update_mask: Boolean tensor of shape ``(n_envs,)``. Candidate robot - and held-object state is committed only where this mask is true. + invocation: Grounded invocation to validate. Returns: - A new state that preserves the current values in unselected rows. + Invocation goal narrowed to this action's declared type. Raises: - TypeError: If ``candidate`` or ``update_mask`` has an invalid type. - ValueError: If state shapes, devices, or semantics are incompatible. + ValueError: If the stable skill identifier does not match. + TypeError: If the goal type is incompatible. + KeyError: If a required binding role is missing. """ - if not isinstance(candidate, WorldState): - raise TypeError("candidate must be a WorldState instance.") - if candidate.last_qpos.shape != self.last_qpos.shape: + if invocation.skill_id != self.skill_id: raise ValueError( - "Candidate WorldState.last_qpos must match the current shape, " - f"got {tuple(candidate.last_qpos.shape)} and " - f"{tuple(self.last_qpos.shape)}." + f"Invocation skill_id {invocation.skill_id!r} does not match " + f"{self.skill_id!r}." ) - if candidate.last_qpos.device != self.last_qpos.device: - raise ValueError("WorldState values being merged must use the same device.") - update_mask = _normalize_env_mask( - update_mask, - batch_size=self.batch_size, - device=self.last_qpos.device, - name="update_mask", - ) - - held_objects: dict[str, HeldObjectState] = {} - held_keys = dict.fromkeys((*self.held_objects, *candidate.held_objects)) - for key in held_keys: - merged = _merge_held_object_state( - self.held_objects.get(key), candidate.held_objects.get(key), update_mask + 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__ ) - if merged is not None: - held_objects[key] = merged - - coordinated_held_objects: dict[tuple[str, str], CoordinatedHeldObjectState] = {} - coordinated_keys = dict.fromkeys( - (*self.coordinated_held_objects, *candidate.coordinated_held_objects) - ) - for key in coordinated_keys: - merged = _merge_coordinated_held_object_state( - self.coordinated_held_objects.get(key), - candidate.coordinated_held_objects.get(key), - update_mask, + raise TypeError( + f"Skill {self.skill_id!r} expects goal {expected}, got " + f"{type(invocation.goal).__name__}." ) - if merged is not None: - coordinated_held_objects[key] = merged - - return WorldState( - last_qpos=torch.where( - update_mask[:, None], candidate.last_qpos, self.last_qpos - ), - held_objects=held_objects, - coordinated_held_objects=coordinated_held_objects, + for role in self.manipulator_roles: + invocation.binding.manipulator(role) + for role in self.end_effector_roles: + invocation.binding.end_effector(role) + required_planner = invocation.motion_policy.planner + configured_planner = getattr( + getattr(self.motion_generator, "planner", None), "cfg", None ) - - -@dataclass(slots=True, eq=False) -class ActionResult: - """Return value of every AtomicAction.execute call.""" - - success: torch.Tensor - """Per-environment planning success, normalized to shape ``(n_envs,)``.""" - - trajectory: torch.Tensor - """Full-robot trajectory, shape (n_envs, n_waypoints, robot.dof).""" - - next_state: WorldState - """World state to feed into the next action.""" - - def __post_init__(self) -> None: - if not isinstance(self.trajectory, torch.Tensor): - raise TypeError("ActionResult.trajectory must be a torch.Tensor.") - if self.trajectory.dim() != 3: + configured_planner_name = getattr(configured_planner, "planner_type", None) + if required_planner is not None and required_planner != configured_planner_name: raise ValueError( - "ActionResult.trajectory must have shape " - f"(n_envs, n_waypoints, robot_dof), got {tuple(self.trajectory.shape)}." - ) - if not isinstance(self.next_state, WorldState): - raise TypeError("ActionResult.next_state must be a WorldState instance.") - expected_shape = ( - self.next_state.batch_size, - self.next_state.robot_dof, - ) - if (self.trajectory.shape[0], self.trajectory.shape[2]) != expected_shape: - raise ValueError( - "ActionResult trajectory batch/DoF must match next_state.last_qpos; " - f"got trajectory {tuple(self.trajectory.shape)} and state " - f"{tuple(self.next_state.last_qpos.shape)}." - ) - if self.trajectory.device != self.next_state.last_qpos.device: - raise ValueError( - "ActionResult trajectory and next_state.last_qpos must use the " - "same device." + f"Motion policy requires planner {required_planner!r}, but this " + f"action uses {configured_planner_name!r}." ) + return invocation.goal - batch_size = self.next_state.batch_size - if isinstance(self.success, bool): - success = torch.full( - (batch_size,), - self.success, + def build_plan( + self, + invocation: ActionInvocation[GoalT], + 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: + invocation: Grounded invocation 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 action config name. + 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. + + Returns: + Side-effect-free action plan. + """ + self.require_goal(invocation) + if isinstance(success, bool): + success_mask = torch.full( + (context.batch_size,), + success, dtype=torch.bool, - device=self.trajectory.device, + device=self.device, ) - elif isinstance(self.success, torch.Tensor): - if self.success.dtype != torch.bool: - raise TypeError( - "ActionResult.success must have dtype torch.bool, " - f"got {self.success.dtype}." - ) - success = self.success.to(device=self.trajectory.device) - if success.dim() == 0 or success.shape == (1,): - success = success.reshape(1).expand(batch_size) - if success.shape != (batch_size,): + 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( - f"ActionResult.success must have shape ({batch_size},), " - f"got {tuple(success.shape)}." + "Planning success must have shape " + f"({context.batch_size},), got {tuple(success_mask.shape)}." ) - success = success.clone() + success_mask = success_mask.clone() else: - raise TypeError("ActionResult.success must be a bool or torch.Tensor.") - self.success = success - - @property - def success_all(self) -> bool: - """True only if all environments succeeded.""" - return bool(torch.all(self.success).item()) - - def __bool__(self) -> bool: - import warnings as _w + raise TypeError("Planning success must be bool or torch.Tensor.") - _w.warn( - "ActionResult bool() is deprecated; use .success_all", - DeprecationWarning, - stacklevel=2, - ) - return self.success_all - - -# ============================================================================= -# Configuration base -# ============================================================================= - - -@configclass -class ActionCfg: - """Configuration shared by all atomic actions.""" - - name: str = "default" - control_part: str = "arm" - interpolation_type: str = "linear" - """Interpolation policy. Only ``"linear"`` is currently implemented.""" - - velocity_limit: float | None = None - acceleration_limit: float | None = None - plan_opts: PlanOptions | None = None - """Optional planner-specific options copied for each motion-generator call.""" - - motion_source: str = "ik_interp" - """Trajectory source: 'ik_interp' (default, batched IK + linear interp) - or 'motion_gen' (batched MotionGenerator).""" - - 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"but got {self.motion_source!r}." + if isinstance(trajectory, torch.Tensor): + timed = TimedTrajectory.from_positions( + trajectory, + env_ids=context.env_ids, + control_dt=invocation.motion_policy.control_dt, ) - if self.interpolation_type != "linear": - raise ValueError( - "interpolation_type currently supports only 'linear', " - f"but got {self.interpolation_type!r}." + 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: + backend = getattr( + getattr(getattr(self.motion_generator, "planner", None), "cfg", None), + "planner_type", + invocation.motion_policy.motion_source, ) + diagnostics = PlannerDiagnostics(backend=str(backend)) + phase = PlannedPhase( + spec=PhaseSpec( + name=phase_name or self.cfg.name, + goal=invocation.goal, + replannable=replannable, + completion_condition=CompletionCondition( + kind=completion_kind, + tolerance=completion_tolerance, + ), + recovery_policy=invocation.recovery_policy, + scene_dependencies=collect_scene_dependencies(invocation.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=invocation.invocation_id, + ) + def failed_plan( + self, + invocation: ActionInvocation[GoalT], + context: PlanningContext, + *, + message: str | None = None, + ) -> ActionPlan: + """Build a failed empty plan without changing task state. -# ============================================================================= -# AtomicAction ABC (slim) -# ============================================================================= - - -class AtomicAction(Generic[TargetT], ABC): - """Abstract base for atomic actions. - - 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. - """ - - TargetType: ClassVar[type[ActionTarget] | tuple[type[ActionTarget], ...]] - """Concrete target dataclass or dataclasses accepted by ``execute``.""" + Args: + invocation: Grounded 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. + """ + backend = getattr( + getattr(getattr(self.motion_generator, "planner", None), "cfg", None), + "planner_type", + invocation.motion_policy.motion_source, + ) + return self.build_plan( + invocation, + 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=str(backend), 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, + invocation: ActionInvocation[GoalT], + 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). + invocation: Fully typed and embodiment-bound 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 0d2e284ed..369870604 100644 --- a/embodichain/lab/sim/atomic_actions/engine.py +++ b/embodichain/lab/sim/atomic_actions/engine.py @@ -14,183 +14,274 @@ # 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 embodichain.utils import logger +import torch -from .core import ( - ActionTarget, - ActionResult, - AtomicAction, - WorldState, - _resolve_runtime_device, -) +from .core import AtomicAction, resolve_runtime_device +from .invocation import ActionInvocation +from .plans import ActionPlan, CompiledTrajectory, TimedTrajectory +from .state import PlanningContext, RobotObservation, SceneSnapshot, TaskState if TYPE_CHECKING: from embodichain.lab.sim.planners import MotionGenerator - -# ============================================================================= -# Global action registry (kept for third-party extensions) -# ============================================================================= + from .execution import ExecutionSession _global_action_registry: dict[str, type[AtomicAction]] = {} -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__ - +def register_action(action_class: type[AtomicAction]) -> None: + """Register an atomic action class under its stable skill identifier. + + Args: + action_class: Concrete :class:`AtomicAction` subclass. + + 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_action_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_action_registry[descriptor.skill_id] = action_class -def register_action(name: str, action_class: type[AtomicAction]) -> None: - """Register a custom AtomicAction subclass globally under ``name``.""" - _global_action_registry[name] = action_class +def unregister_action(skill_id: str) -> None: + """Remove a globally registered skill class if present. -def unregister_action(name: str) -> None: - """Remove a previously-registered action class. No-op if absent.""" - _global_action_registry.pop(name, None) + Args: + skill_id: Stable registered skill identifier. + """ + _global_action_registry.pop(skill_id, None) def get_registered_actions() -> dict[str, type[AtomicAction]]: - """Return a copy of the global action-class registry.""" - return _global_action_registry.copy() - - -# ============================================================================= -# AtomicActionEngine -# ============================================================================= + """Return a copy of the global skill-class registry.""" + return dict(_global_action_registry) class AtomicActionEngine: - """Sequences typed atomic actions while threading WorldState through them.""" + """Compile grounded atomic invocations without stepping the environment.""" def __init__(self, motion_generator: MotionGenerator) -> None: self.motion_generator = motion_generator self.robot = motion_generator.robot - self.device = _resolve_runtime_device(motion_generator.device) + self.device = resolve_runtime_device(motion_generator.device) self._actions: dict[str, AtomicAction] = {} @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) -> None: + """Register one action instance using its descriptor. + + Args: + action: Configured action instance. + + Raises: + TypeError: If ``action`` is not an AtomicAction. + ValueError: If its robot or skill identifier is incompatible. + """ + if not isinstance(action, AtomicAction): + raise TypeError("action must be an AtomicAction instance.") + if action.robot is not self.robot: + raise ValueError("Registered actions must use the engine robot.") + descriptor = action.descriptor() + existing = self._actions.get(descriptor.skill_id) + if existing is not None and existing is not action: + 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 + self._actions[descriptor.skill_id] = action - def run( + def initial_context( 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. + *, + task: TaskState | None = None, + scene: SceneSnapshot | None = None, + timestamp: float = 0.0, + ) -> PlanningContext: + """Capture the robot state needed to start offline compilation. Args: - steps: Iterable of ``(action_name, typed_target)`` pairs. - state: Initial world state. If None, seeded from ``robot.get_qpos()``. + 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: - ``(success, concatenated_full_dof_trajectory, final_state)``. + 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), + ) + + 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: + action = self._actions.get(invocation.skill_id) + if action is None: + raise KeyError( + f"No atomic action registered for skill {invocation.skill_id!r}." + ) + if not alive.any(): + break + previous_qpos = projected.robot.qpos + plan = action.plan(invocation, projected) + self._validate_plan(plan, projected, invocation) + 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, + ) - ``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. + def start( + self, + invocations: Iterable[ActionInvocation], + context: PlanningContext | None = None, + ) -> ExecutionSession: + """Start closed-loop execution for a grounded invocation sequence. - An empty ``steps`` iterable is a successful no-op returning an - empty trajectory and the seed state. + 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(...)``. """ - if state is None: - state = WorldState(last_qpos=self.robot.get_qpos().clone()) + 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}." + ) + 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.") - if state.robot_dof != self.robot.dof: + def _validate_plan( + self, + plan: ActionPlan, + context: PlanningContext, + invocation: ActionInvocation, + ) -> None: + """Validate one action result before it is composed.""" + if plan.skill_id != invocation.skill_id: raise ValueError( - "Initial WorldState DoF must match the engine robot, " - f"got {state.robot_dof} and {self.robot.dof}." + "ActionPlan.skill_id must match its invocation, " + f"got {plan.skill_id!r} and {invocation.skill_id!r}." ) - robot_batch_size = int(self.robot.get_qpos().shape[0]) - if state.batch_size != robot_batch_size: + if plan.invocation_id != invocation.invocation_id: raise ValueError( - "Initial WorldState batch size must match the engine robot, " - f"got {state.batch_size} and {robot_batch_size}." + "ActionPlan.invocation_id must preserve the invocation correlation id." ) - if state.last_qpos.device != self.device: + 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( - "Initial WorldState and AtomicActionEngine must use the same device." + "Every action phase must record the planning scene version." ) - b = state.batch_size - full_traj = torch.empty( - (b, 0, self.robot.dof), - dtype=torch.float32, - device=self.device, - ) - alive = torch.ones(b, dtype=torch.bool, device=self.device) - - for name, target in steps: - if name not in self._actions: - 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, - ) - if not alive.any(): - # All envs dead: no further motion to plan. - break - prev_last_qpos = state.last_qpos.clone() - result: ActionResult = action.execute(target, state) - if result.trajectory.shape[0] != b: - raise ValueError( - f"Action '{name}' returned batch {result.trajectory.shape[0]}, " - f"but the engine state batch is {b}." - ) - if result.trajectory.shape[2] != self.robot.dof: - raise ValueError( - f"Action '{name}' returned {result.trajectory.shape[2]} DoF, " - f"but the engine robot has {self.robot.dof}." - ) - if result.trajectory.device != self.device: - raise ValueError( - f"Action '{name}' returned a trajectory on " - f"{result.trajectory.device}, expected {self.device}." - ) - step_success = result.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 = state.masked_merge(result.next_state, alive) - - return alive, full_traj, state - __all__ = [ "AtomicActionEngine", diff --git a/embodichain/lab/sim/atomic_actions/execution.py b/embodichain/lab/sim/atomic_actions/execution.py new file mode 100644 index 000000000..4def86424 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/execution.py @@ -0,0 +1,685 @@ +# ---------------------------------------------------------------------------- +# 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 +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" + 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_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.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._invocations = 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 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.""" + invocation = self._invocations[self._invocation_index] + action = self._engine.actions.get(invocation.skill_id) + if action is None: + raise KeyError( + f"No atomic action registered for skill {invocation.skill_id!r}." + ) + plan = action.plan(invocation, context) + self._engine._validate_plan(plan, context, invocation) + 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._invocations): + 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._invocations[self._invocation_index].skill_id + if self._invocation_index < len(self._invocations) + else None + ) + invocation_id = ( + self._invocations[self._invocation_index].invocation_id + if self._invocation_index < len(self._invocations) + else None + ) + return ExecutionEvent( + kind=kind, + timestamp=self._context.robot.timestamp, + skill_id=skill_id, + invocation_id=invocation_id, + invocation_index=min(self._invocation_index, len(self._invocations) - 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..56ca04531 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/invocation.py @@ -0,0 +1,79 @@ +# ---------------------------------------------------------------------------- +# 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 dataclasses import dataclass, field +from typing import Generic, TypeVar + +from .bindings import ActionBinding +from .goals import ActionGoal +from .policies import MotionPolicy, RecoveryPolicy + +GoalT = TypeVar("GoalT", bound=ActionGoal) + + +@dataclass(frozen=True, slots=True) +class ActionInvocation(Generic[GoalT]): + """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 for the selected robot embodiment.""" + + motion_policy: MotionPolicy = field(default_factory=MotionPolicy) + """Reusable motion-generation settings.""" + + recovery_policy: RecoveryPolicy = field(default_factory=RecoveryPolicy) + """Bounded local execution recovery settings.""" + + invocation_id: str | None = None + """Optional correlation identifier propagated into execution traces.""" + + 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.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.") + + +__all__ = ["ActionInvocation", "GoalT"] diff --git a/embodichain/lab/sim/atomic_actions/plans.py b/embodichain/lab/sim/atomic_actions/plans.py new file mode 100644 index 000000000..5da8dda35 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/plans.py @@ -0,0 +1,471 @@ +# ---------------------------------------------------------------------------- +# 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 + + 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.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..1adebcb4f --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/policies.py @@ -0,0 +1,120 @@ +# ---------------------------------------------------------------------------- +# 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 typing import TYPE_CHECKING + +from embodichain.utils import configclass + +if TYPE_CHECKING: + from embodichain.lab.sim.planners import PlanOptions + + +@configclass +class MotionPolicy: + """Reusable motion-generation policy supplied with an action invocation.""" + + 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.") + + +@configclass +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..21c109442 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/__init__.py +++ b/embodichain/lab/sim/atomic_actions/primitives/__init__.py @@ -19,65 +19,63 @@ from __future__ import annotations from .coordinated_pickment import ( - CoordinatedPickTarget, + CoordinatedPickGoal, CoordinatedPickment, CoordinatedPickmentCfg, - CoordinatedPickmentTarget, ) from .coordinated_placement import ( CoordinatedPlacement, CoordinatedPlacementCfg, - CoordinatedPlacementTarget, + CoordinatedPlacementGoal, ) from .hand_over import HandOver, HandOverCfg from .move_end_effector import ( - EndEffectorPoseTarget, + EndEffectorPoseGoal, MoveEndEffector, MoveEndEffectorCfg, ) from .move_held_object import ( - HeldObjectPoseTarget, + HeldObjectPoseGoal, MoveHeldObject, MoveHeldObjectCfg, ) from .move_joints import ( - JointPositionTarget, + JointPositionGoal, MoveJoints, MoveJointsCfg, - NamedJointPositionTarget, + NamedJointPositionGoal, ) -from .pick_up import GraspTarget, PickUp, PickUpCfg -from .place import AssembleTarget, Place, PlaceCfg, PlaceTarget -from .press import Press, PressCfg, PressTarget +from .pick_up import GraspGoal, PickUp, PickUpCfg +from .place import AssembleGoal, Place, PlaceCfg, PlaceGoal +from .press import Press, PressCfg, PressGoal __all__ = [ - "AssembleTarget", - "CoordinatedPickTarget", + "AssembleGoal", + "CoordinatedPickGoal", "CoordinatedPickment", "CoordinatedPickmentCfg", - "CoordinatedPickmentTarget", "CoordinatedPlacement", "CoordinatedPlacementCfg", - "CoordinatedPlacementTarget", - "EndEffectorPoseTarget", - "GraspTarget", + "CoordinatedPlacementGoal", + "EndEffectorPoseGoal", + "GraspGoal", "HandOver", "HandOverCfg", - "HeldObjectPoseTarget", - "JointPositionTarget", + "HeldObjectPoseGoal", + "JointPositionGoal", "MoveEndEffector", "MoveEndEffectorCfg", "MoveHeldObject", "MoveHeldObjectCfg", "MoveJoints", "MoveJointsCfg", - "NamedJointPositionTarget", + "NamedJointPositionGoal", "PickUp", "PickUpCfg", "Place", "PlaceCfg", - "PlaceTarget", + "PlaceGoal", "Press", "PressCfg", - "PressTarget", + "PressGoal", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/_helpers.py b/embodichain/lab/sim/atomic_actions/primitives/_helpers.py index 8aec84027..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( @@ -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..73a744b6c 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py @@ -28,21 +28,29 @@ from ..core import ( ActionCfg, - ActionResult, AtomicAction, - CoordinatedHeldObjectState, - WorldState, - _validate_pose_tensor, ) -from ..targets import ObjectActionTarget +from ..effects import StateDelta +from ..goals import ( + ObjectActionGoal, + PoseGoalValue, + resolve_pose_goal, + validate_pose_goal, + validate_pose_tensor, +) +from ..invocation import ActionInvocation +from ..plans import ActionPlan +from ..state import CoordinatedHeldObjectState, PlanningContext from ..trajectory import TrajectoryBuilder @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,38 +59,34 @@ 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" @@ -124,9 +128,6 @@ 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.""" @@ -195,17 +196,6 @@ def _lookup_joint_columns( ) return [joint_id_to_col[joint_id] for joint_id in joint_ids] - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=False, - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, - ) - def _expand_qpos(self, qpos: torch.Tensor, dof: int, name: str) -> torch.Tensor: qpos = qpos.to(device=self.device, dtype=torch.float32) if qpos.shape == (dof,): @@ -233,7 +223,7 @@ def _resolve_pose(self, pose: torch.Tensor, name: str) -> torch.Tensor: def _resolve_dual_arm_start( self, - state: WorldState, + state: PlanningContext, ) -> tuple[torch.Tensor, torch.Tensor]: dual_start = state.last_qpos[:, self.dual_arm_joint_ids].to( device=self.device, dtype=torch.float32 @@ -300,7 +290,7 @@ def _compose_dual_arm_trajectory( def _assemble_phase( self, - state: WorldState, + state: PlanningContext, first_arm_traj: torch.Tensor, second_arm_traj: torch.Tensor, first_hand_traj: torch.Tensor, @@ -431,15 +421,17 @@ def _interpolate_object_pose( return poses -class CoordinatedPickment(AtomicAction[CoordinatedPickTarget]): +class CoordinatedPickment(AtomicAction[CoordinatedPickGoal]): """Pick and move a single object pinched by two hands.""" - TargetType: ClassVar[type] = CoordinatedPickTarget + skill_id: ClassVar[str] = "coordinated_pickment" + GoalType: ClassVar[type] = CoordinatedPickGoal + 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 @@ -457,16 +449,6 @@ def __init__( cfg: CoordinatedPickmentCfg | 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, - ) 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, @@ -512,13 +494,22 @@ def _validate_hand_qpos_cfg(self) -> None: ) 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 +520,8 @@ def _resolve_object_initial_pose( def _resolve_target( self, - target: CoordinatedPickTarget, + target: CoordinatedPickGoal, + context: PlanningContext, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -539,9 +531,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 +568,18 @@ def _resolve_target( held_state, ) - def _compute_segment_lengths(self) -> dict[str, int]: + def _compute_segment_lengths(self, sample_count: int) -> dict[str, int]: + """Split the invocation sample budget across coordinated-pick phases.""" n_close = max(2, self.cfg.hand_interp_steps) n_hold = max(0, self.cfg.hold_steps) - n_motion = self.cfg.sample_interval - n_close - n_hold + n_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,8 +590,9 @@ 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) -> dict[str, int]: + """Return phase lengths for an explicit invocation sample budget.""" + return self._compute_segment_lengths(sample_count) def _compute_pre_grasp_xpos(self, grasp_xpos: torch.Tensor) -> torch.Tensor: grasp_z = grasp_xpos[:, :3, 2] @@ -758,7 +757,40 @@ 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, + invocation: ActionInvocation[CoordinatedPickGoal], + context: PlanningContext, + ) -> ActionPlan: + """Plan a coordinated pick without committing the dual attachment.""" + target = self.require_goal(invocation) + bindings = ( + (invocation.binding.manipulator("left"), self.cfg.left_arm_control_part), + ( + invocation.binding.manipulator("right"), + self.cfg.right_arm_control_part, + ), + ( + invocation.binding.end_effector("left"), + self.cfg.left_hand_control_part, + ), + ( + invocation.binding.end_effector("right"), + self.cfg.right_hand_control_part, + ), + ) + if any(actual != expected for actual, expected in bindings): + raise ValueError( + "CoordinatedPickment bindings do not match its configured resources." + ) + if ( + invocation.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,9 +799,9 @@ def execute(self, target: CoordinatedPickTarget, state: WorldState) -> ActionRes left_target_xpos, right_target_xpos, held_state, - ) = self._resolve_target(target) + ) = self._resolve_target(target, context) left_start_qpos, right_start_qpos = self._resolve_dual_arm_start(state) - segments = self._compute_segment_lengths() + segments = self._compute_segment_lengths(invocation.motion_policy.sample_count) left_pre_grasp_xpos = self._compute_pre_grasp_xpos(left_grasp_xpos) right_pre_grasp_xpos = self._compute_pre_grasp_xpos(right_grasp_xpos) left_approach_targets = torch.stack( @@ -912,36 +944,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( + invocation, + 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={ + self.cfg.left_arm_control_part: None, + self.cfg.right_arm_control_part: None, + }, + coordinated_held_object_updates={ + ( + self.cfg.left_arm_control_part, + self.cfg.right_arm_control_part, + ): coordinated_held_object, + }, ), ) __all__ = [ - "CoordinatedPickTarget", + "CoordinatedPickGoal", "CoordinatedPickment", "CoordinatedPickmentCfg", - "CoordinatedPickmentTarget", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py index 35e8c9fc8..d743e4675 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py @@ -28,25 +28,28 @@ from ._helpers import resolve_object_target from ..core import ( - ActionTarget, ActionCfg, - ActionResult, AtomicAction, - HeldObjectState, - WorldState, - _validate_pose_tensor, ) +from ..effects import StateDelta +from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal +from ..invocation import ActionInvocation +from ..plans import ActionPlan +from ..policies import MotionPolicy +from ..state import HeldObjectState, PlanningContext from ..trajectory import TrajectoryBuilder @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 @@ -59,12 +62,12 @@ class CoordinatedPlacementTarget(ActionTarget): """Whether the placing hand releases. ``None`` uses the action config.""" 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, @@ -112,9 +115,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.""" @@ -125,10 +125,13 @@ class CoordinatedPlacementCfg(ActionCfg): """Number of waypoints used for the placing-arm lift retreat.""" -class CoordinatedPlacement(AtomicAction[CoordinatedPlacementTarget]): +class CoordinatedPlacement(AtomicAction[CoordinatedPlacementGoal]): """Coordinate two held objects: support object below, placing object above.""" - TargetType: ClassVar[type] = CoordinatedPlacementTarget + skill_id: ClassVar[str] = "coordinated_placement" + GoalType: ClassVar[type] = CoordinatedPlacementGoal + manipulator_roles: ClassVar[tuple[str, ...]] = ("placing", "support") + end_effector_roles: ClassVar[tuple[str, ...]] = ("placing", "support") def __init__( self, @@ -136,16 +139,6 @@ def __init__( cfg: CoordinatedPlacementCfg | 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) self.n_envs = self.robot.get_qpos().shape[0] self.robot_dof = self.robot.dof @@ -190,9 +183,43 @@ def __init__( hand_dof=self.support_hand_dof, ) - def execute( - self, target: CoordinatedPlacementTarget, state: WorldState - ) -> ActionResult: + def plan( + self, + invocation: ActionInvocation[CoordinatedPlacementGoal], + context: PlanningContext, + ) -> ActionPlan: + """Plan coordinated placement without committing attachment changes.""" + target = self.require_goal(invocation) + bindings = ( + ( + invocation.binding.manipulator("placing"), + self.cfg.placing_arm_control_part, + ), + ( + invocation.binding.manipulator("support"), + self.cfg.support_arm_control_part, + ), + ( + invocation.binding.end_effector("placing"), + self.cfg.placing_hand_control_part, + ), + ( + invocation.binding.end_effector("support"), + self.cfg.support_hand_control_part, + ), + ) + if any(actual != expected for actual, expected in bindings): + raise ValueError( + "CoordinatedPlacement bindings do not match its configured resources." + ) + if ( + invocation.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, @@ -201,7 +228,9 @@ def execute( 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) + segments = self._compute_segment_lengths( + release, invocation.motion_policy.sample_count + ) placing_lift_xpos = self.builder.apply_local_offset( placing_xpos, @@ -217,20 +246,26 @@ def execute( placing_start_qpos, torch.stack([placing_lift_xpos, placing_xpos], dim=1), segments["approach"], + invocation.motion_policy, ) if not ok: logger.log_warning("CoordinatedPlacement failed to plan placing approach.") - return self._fail(state) + return self.failed_plan( + invocation, context, message="Placing approach failed." + ) ok, support_approach_traj = self._plan_named_arm_trajectory( self.cfg.support_arm_control_part, support_start_qpos, support_xpos.unsqueeze(1), segments["approach"], + invocation.motion_policy, ) if not ok: logger.log_warning("CoordinatedPlacement failed to plan support approach.") - return self._fail(state) + return self.failed_plan( + invocation, context, message="Support approach failed." + ) placing_place_qpos = placing_approach_traj[:, -1] support_place_qpos = support_approach_traj[:, -1] @@ -271,10 +306,13 @@ def execute( placing_place_qpos, placing_lift_xpos.unsqueeze(1), segments["retreat"], + invocation.motion_policy, ) if not ok: logger.log_warning("CoordinatedPlacement failed to plan placing retreat.") - return self._fail(state) + return self.failed_plan( + invocation, context, message="Placing retreat failed." + ) placing_hand_retreat_qpos = ( self.placing_hand_open_qpos if release else self.placing_hand_close_qpos @@ -296,28 +334,28 @@ def execute( ], dim=1, ) - held_objects = dict(state.held_objects) - if release: - held_objects.pop(self.cfg.placing_arm_control_part, None) - else: - held_objects[self.cfg.placing_arm_control_part] = placing_held_object - held_objects[self.cfg.support_arm_control_part] = support_held_object involved_control_parts = { self.cfg.placing_arm_control_part, self.cfg.support_arm_control_part, } - coordinated_held_objects = { - key: value - for key, value in state.coordinated_held_objects.items() - if involved_control_parts.isdisjoint(key) + coordinated_removals = { + key: None + for key in state.coordinated_held_objects + if not involved_control_parts.isdisjoint(key) } - return ActionResult( + return self.build_plan( + invocation, + 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={ + self.cfg.placing_arm_control_part: ( + None if release else placing_held_object + ), + self.cfg.support_arm_control_part: support_held_object, + }, + coordinated_held_object_updates=coordinated_removals, ), ) @@ -336,12 +374,13 @@ def _validate_hand_qpos_cfg(self) -> None: 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 +433,8 @@ def _resolve_held_state( def _resolve_target( self, - target: CoordinatedPlacementTarget, - state: WorldState, + target: CoordinatedPlacementGoal, + state: PlanningContext, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -431,11 +470,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, @@ -465,11 +506,12 @@ def _resolve_target( ) def _resolve_start_qpos( - self, state: WorldState + self, state: PlanningContext ) -> 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, ) @@ -479,15 +521,20 @@ def _resolve_start_qpos( start_qpos[:, self.support_arm_joint_ids], ) - def _compute_segment_lengths(self, release: bool) -> dict[str, int]: + def _compute_segment_lengths( + self, + release: bool, + sample_count: int, + ) -> dict[str, int]: + """Split the invocation sample budget across placement phases.""" 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 + 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 +550,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,7 +565,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 @@ -549,20 +597,9 @@ def _assemble_phase( full[:, :, self.support_hand_joint_ids] = support_hand_traj return full - def _fail(self, state: WorldState) -> ActionResult: - return ActionResult( - success=False, - trajectory=torch.empty( - (self.n_envs, 0, self.robot_dof), - dtype=torch.float32, - device=self.device, - ), - next_state=state, - ) - __all__ = [ "CoordinatedPlacement", "CoordinatedPlacementCfg", - "CoordinatedPlacementTarget", + "CoordinatedPlacementGoal", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/hand_over.py b/embodichain/lab/sim/atomic_actions/primitives/hand_over.py index ba0298d01..8c251ff0f 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/hand_over.py +++ b/embodichain/lab/sim/atomic_actions/primitives/hand_over.py @@ -28,13 +28,15 @@ from ..core import ( ActionCfg, - ActionResult, AtomicAction, - HeldObjectState, ObjectSemantics, - WorldState, ) -from .pick_up import GraspTarget +from ..effects import StateDelta +from ..invocation import ActionInvocation +from ..plans import ActionPlan +from ..policies import MotionPolicy +from ..state import HeldObjectState, PlanningContext +from .pick_up import GraspGoal from ..trajectory import TrajectoryBuilder @@ -96,9 +98,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.""" @@ -110,7 +109,7 @@ class HandOverCfg(ActionCfg): """Number of waypoints used for the final deliver/retreat phase.""" -class HandOver(AtomicAction[GraspTarget]): +class HandOver(AtomicAction[GraspGoal]): """Hand an object from one arm to the other. The transferring arm (already holding the object) moves it to a middle @@ -119,7 +118,10 @@ 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 + manipulator_roles: ClassVar[tuple[str, ...]] = ("source", "destination") + end_effector_roles: ClassVar[tuple[str, ...]] = ("source", "destination") def __init__( self, @@ -127,16 +129,6 @@ def __init__( cfg: HandOverCfg | 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) self.n_envs = self.robot.get_qpos().shape[0] self.robot_dof = self.robot.dof @@ -201,7 +193,41 @@ def __init__( # Public contract # ------------------------------------------------------------------ - def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: + def plan( + self, + invocation: ActionInvocation[GraspGoal], + context: PlanningContext, + ) -> ActionPlan: + """Plan a handover without committing the attachment transfer.""" + target = self.require_goal(invocation) + expected_bindings = ( + ( + invocation.binding.manipulator("source"), + self.cfg.transfer_arm_control_part, + ), + ( + invocation.binding.manipulator("destination"), + self.cfg.receive_arm_control_part, + ), + ( + invocation.binding.end_effector("source"), + self.cfg.transfer_hand_control_part, + ), + ( + invocation.binding.end_effector("destination"), + self.cfg.receive_hand_control_part, + ), + ) + if any(actual != expected for actual, expected in expected_bindings): + raise ValueError("HandOver bindings do not match its configured resources.") + if ( + invocation.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() @@ -220,7 +246,7 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: ) 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(invocation, context, message="No receiving grasp.") receive_object_to_eef = torch.bmm( pose_inv(middle_object_pose), receive_grasp_xpos ) @@ -242,27 +268,33 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: ) transfer_start_qpos, receive_start_qpos = self._resolve_start_qpos(state) - segments = self._compute_segment_lengths() + segments = self._compute_segment_lengths(invocation.motion_policy.sample_count) ok, transfer_move_traj = self._plan_named_arm_trajectory( self.cfg.transfer_arm_control_part, transfer_start_qpos, transfer_middle_eef.unsqueeze(1), segments["transfer"], + invocation.motion_policy, ) if not ok: logger.log_warning("HandOver failed to plan the transfer move.") - return self._fail(state) + return self.failed_plan( + invocation, context, message="Transfer move failed." + ) ok, receive_approach_traj = self._plan_named_arm_trajectory( self.cfg.receive_arm_control_part, receive_start_qpos, torch.stack([receive_pre_grasp_eef, receive_grasp_xpos], dim=1), segments["approach"], + invocation.motion_policy, ) if not ok: logger.log_warning("HandOver failed to plan the receiving approach.") - return self._fail(state) + return self.failed_plan( + invocation, context, message="Receiving approach failed." + ) transfer_hold_qpos = transfer_move_traj[:, -1] receive_grasp_qpos = receive_approach_traj[:, -1] @@ -272,20 +304,26 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: transfer_hold_qpos, transfer_retreat_eef.unsqueeze(1), segments["deliver"], + invocation.motion_policy, ) if not ok: logger.log_warning("HandOver failed to plan the transfer retreat.") - return self._fail(state) + return self.failed_plan( + invocation, context, message="Transfer retreat failed." + ) ok, receive_deliver_traj = self._plan_named_arm_trajectory( self.cfg.receive_arm_control_part, receive_grasp_qpos, receive_final_eef.unsqueeze(1), segments["deliver"], + invocation.motion_policy, ) if not ok: logger.log_warning("HandOver failed to plan the receiving delivery.") - return self._fail(state) + return self.failed_plan( + invocation, context, message="Receiving delivery failed." + ) phases: list[torch.Tensor] = [] # 2.1 transfer: transferring arm carries the object to the middle pose. @@ -362,15 +400,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( + invocation, + 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={ + self.cfg.transfer_arm_control_part: None, + self.cfg.receive_arm_control_part: held_object, + } ), ) @@ -406,7 +445,7 @@ def _resolve_matrix(self, matrix: torch.Tensor, name: str) -> torch.Tensor: ) return matrix - def _resolve_transfer_object_to_eef(self, state: WorldState) -> torch.Tensor: + def _resolve_transfer_object_to_eef(self, state: PlanningContext) -> torch.Tensor: held = state.get_held_object(self.cfg.transfer_arm_control_part) if held is None: logger.log_error( @@ -449,11 +488,11 @@ def _resolve_receive_grasp( return grasp_xpos, is_success def _resolve_start_qpos( - self, state: WorldState + self, state: PlanningContext ) -> 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, ) @@ -463,17 +502,18 @@ def _resolve_start_qpos( start_qpos[:, self.receive_arm_joint_ids], ) - def _compute_segment_lengths(self) -> dict[str, int]: + def _compute_segment_lengths(self, sample_count: int) -> dict[str, int]: + """Split the invocation sample budget across handover phases.""" 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) 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 +536,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 +551,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,7 +561,7 @@ 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, @@ -539,16 +580,5 @@ def _assemble_phase( base[:, :, self.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"] 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 682e18fa8..29c38db26 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py @@ -26,57 +26,40 @@ from embodichain.lab.sim.planners import MoveType, 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 ..core import ActionCfg, AtomicAction +from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal +from ..invocation import ActionInvocation +from ..plans import ActionPlan, CompletionConditionKind +from ..state import PlanningContext from ..trajectory import TrajectoryBuilder @dataclass(frozen=True, slots=True, eq=False) -class EndEffectorPoseTarget(ActionTarget): - """End-effector pose target used by :class:`MoveEndEffector`.""" +class EndEffectorPoseGoal: + """End-effector pose goal with optional batched intermediate waypoints.""" - xpos: torch.Tensor - """Target end-effector homogeneous transform. + goal_kind: ClassVar[str] = "end_effector_pose" - Accepts: - - - ``(4, 4)`` or ``(n_envs, 4, 4)`` — a single waypoint. - - ``(n_envs, n_waypoint, 4, 4)`` — a multi-waypoint trajectory whose - waypoints are visited in order. - """ + 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) + validate_pose_goal(self.xpos, "xpos", allow_waypoints=True) @configclass class MoveEndEffectorCfg(ActionCfg): - name: str = "move_end_effector" - """Name of the action, used for identification and logging.""" - - sample_interval: int = 50 - """Number of waypoints in the planned trajectory.""" + """Skill-specific MoveEndEffector configuration.""" + name: str = "move_end_effector" -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]): + """Plan a free-space move for a bound manipulator.""" - TargetType: ClassVar[type] = EndEffectorPoseTarget + skill_id: ClassVar[str] = "move_end_effector" + GoalType: ClassVar[type] = EndEffectorPoseGoal + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) def __init__( self, @@ -85,73 +68,66 @@ def __init__( ) -> 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) + def plan( + self, + invocation: ActionInvocation[EndEffectorPoseGoal], + context: PlanningContext, + ) -> ActionPlan: + """Plan an end-effector pose goal from the observed joint state.""" + goal = self.require_goal(invocation) + control_part = invocation.binding.manipulator("primary") + joint_ids = self.robot.get_joint_ids(name=control_part) + arm_dof = len(joint_ids) + 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, + invocation.motion_policy.sample_count, + control_part=control_part, + arm_dof=arm_dof, + cfg=invocation.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=invocation.motion_policy.control_dt, ) - full = self._embed(arm_traj, state.last_qpos) - return ActionResult( + return self.build_plan( + invocation, + 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", "MoveEndEffectorCfg"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py b/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py index c26683351..0705e6840 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py @@ -29,25 +29,27 @@ from ._helpers import arm_qpos_from_state, resolve_object_target from ..core import ( - ActionTarget, ActionCfg, - ActionResult, AtomicAction, - WorldState, - _validate_pose_tensor, ) +from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal +from ..invocation import ActionInvocation +from ..plans import ActionPlan +from ..state import PlanningContext from ..trajectory import TrajectoryBuilder @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, @@ -59,8 +61,8 @@ 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.""" + control_part: str = "arm" + """Manipulator resource used by this configured action instance.""" hand_control_part: str = "hand" """Name of the robot part that controls the hand joints.""" @@ -75,10 +77,13 @@ class MoveHeldObjectCfg(ActionCfg): """Optional rotation in radians used by the legacy upright transport mode.""" -class MoveHeldObject(AtomicAction[HeldObjectPoseTarget]): +class MoveHeldObject(AtomicAction[HeldObjectPoseGoal]): """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 + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) def __init__( self, @@ -99,7 +104,22 @@ def __init__( ) self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) - def execute(self, target: HeldObjectPoseTarget, state: WorldState) -> ActionResult: + def plan( + self, + invocation: ActionInvocation[HeldObjectPoseGoal], + context: PlanningContext, + ) -> ActionPlan: + """Plan held-object transport without changing the attachment relation.""" + target = self.require_goal(invocation) + if invocation.binding.manipulator() != self.cfg.control_part: + raise ValueError( + "MoveHeldObject manipulator binding does not match its config." + ) + if invocation.binding.end_effector() != self.cfg.hand_control_part: + raise ValueError( + "MoveHeldObject end-effector binding does not match its config." + ) + state = context held_object = state.get_held_object(self.cfg.control_part) if held_object is None: logger.log_error( @@ -108,7 +128,13 @@ def execute(self, target: HeldObjectPoseTarget, state: WorldState) -> ActionResu 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), @@ -142,10 +168,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, + invocation.motion_policy.sample_count, control_part=self.cfg.control_part, arm_dof=self.arm_dof, - cfg=self.cfg, + cfg=invocation.motion_policy, ) full = torch.empty( @@ -157,12 +183,12 @@ def execute(self, target: HeldObjectPoseTarget, state: WorldState) -> ActionResu full[:, :, self.arm_joint_ids] = arm_traj full[:, :, self.hand_joint_ids] = self.hand_close_qpos - return ActionResult( + return self.build_plan( + invocation, + context, success=success, trajectory=full, - next_state=state.with_updates( - last_qpos=full[:, -1, :].clone(), - ), + phase_name="transport", ) def _apply_configured_upright_rotation( @@ -244,16 +270,5 @@ 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", "MoveHeldObjectCfg"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py index e619c113e..3fc9fcc83 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py @@ -25,29 +25,21 @@ from embodichain.utils import configclass, logger -from ..core import ( - ActionTarget, - ActionCfg, - ActionResult, - AtomicAction, - WorldState, -) +from ..core import ActionCfg, AtomicAction +from ..invocation import ActionInvocation +from ..plans import ActionPlan, CompletionConditionKind +from ..state import PlanningContext from ..trajectory import TrajectoryBuilder @dataclass(frozen=True, slots=True, eq=False) -class JointPositionTarget(ActionTarget): - """Joint-space target for a configured robot control part.""" +class JointPositionGoal: + """Joint-space goal for a bound robot control resource.""" - qpos: torch.Tensor - """Target joint positions. - - 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. - """ + qpos: torch.Tensor + """One joint waypoint or a batched sequence of joint waypoints.""" def __post_init__(self) -> None: if not isinstance(self.qpos, torch.Tensor): @@ -63,11 +55,13 @@ def __post_init__(self) -> None: @dataclass(frozen=True, slots=True, eq=False) -class NamedJointPositionTarget(ActionTarget): - """Named joint-space target resolved from :class:`MoveJointsCfg`.""" +class NamedJointPositionGoal: + """Named joint-space goal resolved from :class:`MoveJointsCfg`.""" + + goal_kind: ClassVar[str] = "named_joint_position" name: str - """Name of a joint-position target in ``MoveJointsCfg.named_joint_positions``.""" + """Name in ``MoveJointsCfg.named_joint_positions``.""" def __post_init__(self) -> None: if not isinstance(self.name, str): @@ -78,30 +72,23 @@ def __post_init__(self) -> None: @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.""" + """Skill-specific MoveJoints configuration.""" + name: str = "move_joints" named_joint_positions: dict[str, torch.Tensor] | None = None - """Optional named joint targets resolved by ``NamedJointPositionTarget``.""" + """Optional named goals. Motion settings belong to ``MotionPolicy``.""" -class MoveJoints(AtomicAction[JointPositionTarget | NamedJointPositionTarget]): - """Plan a joint-space move for the configured control part. +class MoveJoints(AtomicAction[JointPositionGoal | NamedJointPositionGoal]): + """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[tuple[type, ...]] = ( + JointPositionGoal, + NamedJointPositionGoal, ) + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + agent_visible: ClassVar[bool] = False def __init__( self, @@ -110,76 +97,72 @@ def __init__( ) -> 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( + def plan( self, - target: JointPositionTarget | NamedJointPositionTarget, - state: WorldState, - ) -> ActionResult: + invocation: ActionInvocation[JointPositionGoal | NamedJointPositionGoal], + context: PlanningContext, + ) -> ActionPlan: + """Plan a joint-space goal without mutating the robot or task state.""" + goal = self.require_goal(invocation) + control_part = invocation.binding.manipulator("primary") + joint_ids = self.robot.get_joint_ids(name=control_part) + joint_dof = len(joint_ids) 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), + 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, + invocation.motion_policy.sample_count, + control_part=control_part, + arm_dof=joint_dof, + cfg=invocation.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=invocation.motion_policy.control_dt, ) - full = self._embed(joint_traj, state.last_qpos) - return ActionResult( + return self.build_plan( + invocation, + 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 + self, + goal: JointPositionGoal | NamedJointPositionGoal, ) -> torch.Tensor: - if isinstance(target, JointPositionTarget): - return target.qpos - if target.name not in self.named_joint_positions: + """Resolve an explicit or named joint goal to a tensor.""" + if isinstance(goal, JointPositionGoal): + return goal.qpos + if goal.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)}", + f"Unknown named joint-position goal {goal.name!r}. Available " + f"goals: {sorted(self.named_joint_positions)}", KeyError, ) - return self.named_joint_positions[target.name] - - def _embed( - self, joint_traj: torch.Tensor, last_full_qpos: torch.Tensor - ) -> torch.Tensor: - n_wp = joint_traj.shape[1] - full = torch.empty( - (self.n_envs, n_wp, self.robot_dof), - dtype=torch.float32, - device=self.device, - ) - full[:, :, :] = last_full_qpos.unsqueeze(1) - full[:, :, self.joint_ids] = joint_traj - return full + return self.named_joint_positions[goal.name] __all__ = [ - "JointPositionTarget", + "JointPositionGoal", "MoveJoints", "MoveJointsCfg", - "NamedJointPositionTarget", + "NamedJointPositionGoal", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index 85cbe8940..4a812802d 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -37,21 +37,24 @@ from ..affordance import AntipodalAffordance from ..core import ( ActionCfg, - ActionResult, AtomicAction, - HeldObjectState, ObjectSemantics, - WorldState, - _validate_pose_tensor, ) -from ..targets import ObjectActionTarget +from ..effects import StateDelta +from ..goals import ObjectActionGoal, validate_pose_tensor +from ..invocation import ActionInvocation +from ..plans import ActionPlan +from ..policies import MotionPolicy +from ..state import HeldObjectState, PlanningContext from ..trajectory import TrajectoryBuilder @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,9 +64,9 @@ 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 @@ -71,8 +74,8 @@ 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).""" + control_part: str = "arm" + """Manipulator resource used by this configured action instance.""" hand_interp_steps: int = 5 """Number of waypoints for the gripper close interpolation phase.""" @@ -111,10 +114,13 @@ class PickUpCfg(ActionCfg): """Optional rotation (radians) about the grasp x-axis to apply after grasp selection.""" -class PickUp(AtomicAction[GraspTarget]): +class PickUp(AtomicAction[GraspGoal]): """Approach a grasp pose, close the gripper, lift.""" - TargetType: ClassVar[type] = GraspTarget + skill_id: ClassVar[str] = "pick_up" + GoalType: ClassVar[type] = GraspGoal + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) def __init__( self, @@ -159,13 +165,14 @@ def _get_full_pickup_trajectory( grasp_xpos: torch.Tensor, start_arm_qpos: torch.Tensor, last_qpos: torch.Tensor, - ): + motion_policy: MotionPolicy, + ) -> tuple[torch.Tensor, torch.Tensor]: pre_grasp_xpos = self.builder.apply_local_offset( grasp_xpos, -self.approach_direction * self.cfg.pre_grasp_distance ) n_approach, n_close, n_lift = self.builder.split_three_phase( - self.cfg.sample_interval, + motion_policy.sample_count, self.cfg.hand_interp_steps, first_phase_name="approach", third_phase_name="lift", @@ -184,7 +191,7 @@ def _get_full_pickup_trajectory( n_approach, control_part=self.cfg.control_part, arm_dof=self.arm_dof, - cfg=self.cfg, + cfg=motion_policy, ) grasp_arm_qpos = approach_arm[:, -1, :] @@ -202,7 +209,7 @@ def _get_full_pickup_trajectory( n_lift, control_part=self.cfg.control_part, arm_dof=self.arm_dof, - cfg=self.cfg, + cfg=motion_policy, ) is_success = approach_success & lift_success @@ -231,7 +238,18 @@ def _get_full_pickup_trajectory( ) return is_success, full - def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: + def plan( + self, + invocation: ActionInvocation[GraspGoal], + context: PlanningContext, + ) -> ActionPlan: + """Plan approach, close, and lift phases without committing attachment.""" + target = self.require_goal(invocation) + if invocation.binding.manipulator() != self.cfg.control_part: + raise ValueError("PickUp manipulator binding does not match its config.") + if invocation.binding.end_effector() != self.cfg.hand_control_part: + raise ValueError("PickUp end-effector binding does not match its config.") + state = context sem = target.semantics if target.grasp_xpos is None and not isinstance( sem.affordance, AntipodalAffordance @@ -261,10 +279,15 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: 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( + invocation, 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, + invocation.motion_policy, ) obj_poses = sem.entity.get_local_pose(to_matrix=True) @@ -272,32 +295,21 @@ 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 self.cfg.control_part in key } - return ActionResult( + return self.build_plan( + invocation, + 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, - ), - ) - - 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, + expected_effects=StateDelta( + held_object_updates={self.cfg.control_part: held}, + coordinated_held_object_updates=coordinated_updates, ), - next_state=state, + phase_name="pick", ) def _resolve_grasp_pose( @@ -515,4 +527,4 @@ def _upright_adjusted_grasp_poses( return adjusted_grasp_xpos -__all__ = ["GraspTarget", "PickUp", "PickUpCfg"] +__all__ = ["GraspGoal", "PickUp", "PickUpCfg"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/place.py b/embodichain/lab/sim/atomic_actions/primitives/place.py index 054c83c16..80b4c47c7 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/place.py +++ b/embodichain/lab/sim/atomic_actions/primitives/place.py @@ -30,23 +30,26 @@ 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 ..effects import StateDelta +from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal +from ..invocation import ActionInvocation +from ..plans import ActionPlan +from ..state import PlanningContext from ..trajectory import TrajectoryBuilder 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 +65,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,16 +74,18 @@ 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.""" @@ -90,8 +95,8 @@ 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).""" + control_part: str = "arm" + """Manipulator resource used by this configured action instance.""" hand_interp_steps: int = 5 """Number of waypoints for the gripper open interpolation phase.""" @@ -115,27 +120,30 @@ class PlaceCfg(ActionCfg): """Number of fixed-orientation Cartesian keyframes per translation segment.""" -class Place(AtomicAction[PlaceTarget | AssembleTarget]): +class Place(AtomicAction[PlaceGoal | AssembleGoal]): """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, ) + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) def __init__( self, @@ -161,9 +169,18 @@ def __init__( if self.cfg.cartesian_waypoint_count < 1: logger.log_error("cartesian_waypoint_count must be at least 1.", ValueError) - def execute( - self, target: PlaceTarget | AssembleTarget, state: WorldState - ) -> ActionResult: + def plan( + self, + invocation: ActionInvocation[PlaceGoal | AssembleGoal], + context: PlanningContext, + ) -> ActionPlan: + """Plan approach, release, and retract without committing detachment.""" + target = self.require_goal(invocation) + if invocation.binding.manipulator() != self.cfg.control_part: + raise ValueError("Place manipulator binding does not match its config.") + if invocation.binding.end_effector() != self.cfg.hand_control_part: + raise ValueError("Place end-effector binding does not match its config.") + state = context place_xpos = self._resolve_place_xpos(target, state) if place_xpos.dim() == 3: place_xpos = place_xpos.unsqueeze(1) @@ -174,12 +191,12 @@ def execute( arm_dof=self.arm_dof, control_part=self.cfg.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 ) n_down, n_open, n_back = self.builder.split_three_phase( - self.cfg.sample_interval, + invocation.motion_policy.sample_count, self.cfg.hand_interp_steps, first_phase_name="approach", third_phase_name="back", @@ -209,7 +226,7 @@ def execute( n_down, control_part=self.cfg.control_part, arm_dof=self.arm_dof, - cfg=self.cfg, + cfg=invocation.motion_policy, ) reach_arm_qpos = down_arm[:, -1, :] @@ -229,7 +246,7 @@ def execute( n_back, control_part=self.cfg.control_part, arm_dof=self.arm_dof, - cfg=self.cfg, + cfg=invocation.motion_policy, ) success = down_success & back_success @@ -258,25 +275,25 @@ def execute( 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 + coordinated_updates = { + key: None + for key in state.coordinated_held_objects + if self.cfg.control_part in key } - return ActionResult( + return self.build_plan( + invocation, + 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={self.cfg.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 ) -> torch.Tensor: """Resolve the place EEF poses from a typed target. @@ -288,12 +305,15 @@ 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) + 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) def _resolve_assemble_place_xpos( - self, target: AssembleTarget, state: WorldState + self, target: AssembleGoal, state: PlanningContext ) -> torch.Tensor: """Derive the place EEF pose from an assembly affordance. @@ -314,7 +334,7 @@ def _resolve_assemble_place_xpos( held = state.get_held_object(self.cfg.control_part) if held is None: logger.log_error( - "Place with AssembleTarget requires an object held by control " + "Place with AssembleGoal requires an object held by control " f"part {self.cfg.control_part!r} (run PickUp first).", ValueError, ) @@ -380,17 +400,6 @@ 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 ) -> torch.Tensor: @@ -423,4 +432,4 @@ def _select_tcp_symmetric_place_variant( ] -__all__ = ["Place", "PlaceCfg", "PlaceTarget", "AssembleTarget"] +__all__ = ["AssembleGoal", "Place", "PlaceCfg", "PlaceGoal"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/press.py b/embodichain/lab/sim/atomic_actions/primitives/press.py index f1606d543..ddd0ea9c7 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/press.py +++ b/embodichain/lab/sim/atomic_actions/primitives/press.py @@ -28,25 +28,27 @@ from ._helpers import arm_qpos_from_state from ..core import ( - ActionTarget, ActionCfg, - ActionResult, AtomicAction, - WorldState, - _validate_pose_tensor, ) +from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal +from ..invocation import ActionInvocation +from ..plans import ActionPlan +from ..state import PlanningContext from ..trajectory import TrajectoryBuilder @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 @@ -54,8 +56,8 @@ 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).""" + control_part: str = "arm" + """Manipulator resource used by this configured action instance.""" hand_interp_steps: int = 5 """Number of waypoints for closing the gripper before pressing.""" @@ -67,10 +69,13 @@ class PressCfg(ActionCfg): """Joint positions for the closed hand state, shape ``[hand_dof,]``.""" -class Press(AtomicAction[PressTarget]): +class Press(AtomicAction[PressGoal]): """Close the gripper, press down to a target pose, then return.""" - TargetType: ClassVar[type] = PressTarget + skill_id: ClassVar[str] = "press" + GoalType: ClassVar[type] = PressGoal + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) def __init__( self, @@ -96,8 +101,22 @@ def __init__( 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) + def plan( + self, + invocation: ActionInvocation[PressGoal], + context: PlanningContext, + ) -> ActionPlan: + """Plan a close, press, and retract sequence.""" + target = self.require_goal(invocation) + if invocation.binding.manipulator() != self.cfg.control_part: + raise ValueError("Press manipulator binding does not match its config.") + if invocation.binding.end_effector() != self.cfg.hand_control_part: + raise ValueError("Press end-effector binding does not match its config.") + state = context + press_xpos = self.builder.resolve_pose_target( + resolve_pose_goal(target.xpos, context, name="xpos"), + n_envs=self.n_envs, + ) start_arm_qpos = self.builder.resolve_start_qpos( arm_qpos_from_state(state, self.arm_joint_ids), n_envs=self.n_envs, @@ -106,7 +125,9 @@ def execute(self, target: PressTarget, state: WorldState) -> ActionResult: ) start_hand_qpos = state.last_qpos[:, self.hand_joint_ids] - n_close, n_down, n_back = self._compute_phase_waypoints() + n_close, n_down, n_back = self._compute_phase_waypoints( + invocation.motion_policy.sample_count + ) hand_close_path = self.builder.interpolate_hand_qpos( start_hand_qpos, @@ -124,7 +145,7 @@ def execute(self, target: PressTarget, state: WorldState) -> ActionResult: n_down, control_part=self.cfg.control_part, arm_dof=self.arm_dof, - cfg=self.cfg, + cfg=invocation.motion_policy, ) press_arm_qpos = down_arm[:, -1, :] @@ -134,7 +155,7 @@ def execute(self, target: PressTarget, state: WorldState) -> ActionResult: n_back, control_part=self.cfg.control_part, arm_dof=self.arm_dof, - cfg=self.cfg, + cfg=invocation.motion_policy, ) success = down_success & back_success @@ -159,42 +180,32 @@ def execute(self, target: PressTarget, state: WorldState) -> ActionResult: self.hand_close_qpos.unsqueeze(1) ) - return ActionResult( + return self.build_plan( + invocation, + 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]: + def _compute_phase_waypoints(self, sample_count: int) -> tuple[int, int, int]: + """Split the invocation sample budget across press phases.""" n_close = self.cfg.hand_interp_steps if n_close < 1: logger.log_error( "hand_interp_steps must be at least 1 for PressCfg.", ValueError ) - motion_waypoints = self.cfg.sample_interval - n_close + 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", "PressCfg", "PressGoal"] 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 0a4953568..17165008c 100644 --- a/embodichain/lab/sim/atomic_actions/trajectory.py +++ b/embodichain/lab/sim/atomic_actions/trajectory.py @@ -31,11 +31,14 @@ from embodichain.lab.sim.utility.action_utils import interpolate_with_distance from embodichain.utils import logger -from .core import _resolve_runtime_device +from .core import resolve_runtime_device +from .plans import TimedTrajectory if TYPE_CHECKING: from embodichain.lab.sim.planners import MotionGenerator + from .policies import MotionPolicy + class TrajectoryBuilder: """Stateless trajectory utilities shared by every atomic action. @@ -48,7 +51,7 @@ class TrajectoryBuilder: 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 @@ -336,7 +339,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, @@ -344,18 +347,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, @@ -363,12 +374,46 @@ 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( @@ -429,9 +474,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", @@ -449,7 +523,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, @@ -459,6 +535,27 @@ def _process_motion_gen_result( arm_dof: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Validate a MotionGenerator PlanResult and apply sample/hold policy.""" + normalized = self._normalize_motion_gen_result( + result, start_qpos, n_waypoints, arm_dof + ) + 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, @@ -483,17 +580,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 @@ -531,7 +685,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: @@ -566,7 +724,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, @@ -574,9 +732,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 @@ -586,7 +744,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" @@ -614,12 +772,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/examples/sim/planners/curobo_planner.py b/examples/sim/planners/curobo_planner.py index 8e71ee8d1..cd19f86cc 100644 --- a/examples/sim/planners/curobo_planner.py +++ b/examples/sim/planners/curobo_planner.py @@ -59,10 +59,13 @@ visualization_cfg_from_args, ) from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - EndEffectorPoseTarget, + EndEffectorPoseGoal, MoveEndEffector, MoveEndEffectorCfg, + MotionPolicy, ) from embodichain.data import get_data_path from embodichain.lab.sim.cfg import RenderCfg, RigidBodyAttributesCfg @@ -747,26 +750,17 @@ 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, + engine.register(MoveEndEffector(motion_generator, MoveEndEffectorCfg())) + 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 +770,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 +808,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..d1381653d 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) @@ -251,14 +264,7 @@ 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 - ), - ) - ) + atomic_engine.register(MoveEndEffector(motion_gen, cfg=MoveEndEffectorCfg())) 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..9ee113330 100644 --- a/scripts/benchmark/atomic_action/move_held_object_benchmark.py +++ b/scripts/benchmark/atomic_action/move_held_object_benchmark.py @@ -175,11 +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, + EndEffectorPoseGoal, + GraspGoal, MoveEndEffector, MoveEndEffectorCfg, + MotionPolicy, PickUp, PickUpCfg, ) @@ -192,15 +195,7 @@ 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(MoveEndEffector(motion_gen, cfg=MoveEndEffectorCfg())) atomic_engine.register( PickUp( motion_gen, @@ -214,7 +209,6 @@ def _prepare_held_state( ), pre_grasp_distance=0.15, lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, hand_interp_steps=HAND_INTERP_STEPS, ), ) @@ -230,12 +224,29 @@ 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), + ), + ) ) + 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 +273,13 @@ def _run_case( ): """Run one MoveHeldObject benchmark case.""" from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, - HeldObjectPoseTarget, + HeldObjectPoseGoal, MoveHeldObject, MoveHeldObjectCfg, + MotionPolicy, ) from scripts.tutorials.atomic_action.move_held_object import ( compute_pick_close_end_step, @@ -309,23 +323,31 @@ def _run_case( control_part="arm", hand_control_part="hand", hand_close_qpos=hand_close, - sample_interval=MOVE_HELD_OBJECT_SAMPLE_INTERVAL, ), ) ) target_pose = _make_object_target_pose(sim.device, case.xyz) elapsed, mem_delta, peak_gpu, result = timed_call( - lambda: atomic_engine.run( - steps=[ - ( - "move_held_object", - HeldObjectPoseTarget(object_target_pose=target_pose), - ) - ], - state=state, + 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..e0f4b9f62 100644 --- a/scripts/benchmark/atomic_action/move_joints_benchmark.py +++ b/scripts/benchmark/atomic_action/move_joints_benchmark.py @@ -109,22 +109,30 @@ 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, + NamedJointPositionGoal, ) 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 = NamedJointPositionGoal(name="ready") else: - targets.append( - ( - "move_joints", - JointPositionTarget(qpos=_qpos(JOINT_TARGETS[name], device)), - ) + goal = JointPositionGoal(qpos=_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 +149,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) @@ -266,8 +277,6 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: MoveJoints( motion_gen, cfg=MoveJointsCfg( - control_part="arm", - sample_interval=MOVE_JOINTS_SAMPLE_INTERVAL, named_joint_positions={"ready": ready_qpos}, ), ) diff --git a/scripts/benchmark/atomic_action/pickup_benchmark.py b/scripts/benchmark/atomic_action/pickup_benchmark.py index 9c8f56ff7..6bc6acabf 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, + GraspGoal, PickUp, PickUpCfg, + MotionPolicy, ) from scripts.tutorials.atomic_action.pickup import ( build_grasp_generator_cfg, @@ -168,7 +171,6 @@ def _run_case( approach_direction=approach_direction, pre_grasp_distance=0.15, lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, hand_interp_steps=HAND_INTERP_STEPS, ), ) @@ -181,11 +183,23 @@ 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), + ), + ) ) ) - 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..f35df7a46 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, + GraspGoal, PickUp, PickUpCfg, + MotionPolicy, ) from scripts.tutorials.atomic_action.place import ( build_grasp_generator_cfg, @@ -202,7 +205,6 @@ def _prepare_held_state( ), pre_grasp_distance=0.15, lift_height=0.16, - sample_interval=PICK_SAMPLE_INTERVAL, hand_interp_steps=HAND_INTERP_STEPS, ), ) @@ -214,9 +216,22 @@ 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), + ), + ) ) + 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 +256,13 @@ def _run_case( ): """Run one Place benchmark case.""" from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, AtomicActionEngine, + MotionPolicy, Place, PlaceCfg, - PlaceTarget, + PlaceGoal, ) from scripts.tutorials.atomic_action.place import ( compute_pick_close_end_step, @@ -291,19 +309,30 @@ def _run_case( hand_open_qpos=hand_open, hand_close_qpos=hand_close, lift_height=PLACE_LIFT_HEIGHT, - sample_interval=PLACE_SAMPLE_INTERVAL, hand_interp_steps=HAND_INTERP_STEPS, ), ) ) place_pose = _make_place_pose(sim.device, case.xyz) elapsed, mem_delta, peak_gpu, result = timed_call( - lambda: atomic_engine.run( - steps=[("place", 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), + ), + ), + 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..d96dd7f37 100644 --- a/scripts/benchmark/atomic_action/press_benchmark.py +++ b/scripts/benchmark/atomic_action/press_benchmark.py @@ -81,13 +81,16 @@ 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, + EndEffectorPoseGoal as end_effector_pose_target_cls, MoveEndEffector as move_end_effector_cls, MoveEndEffectorCfg as move_end_effector_cfg_cls, + MotionPolicy as motion_policy_cls, Press as press_cls, PressCfg as press_cfg_cls, - PressTarget as press_target_cls, + PressGoal as press_target_cls, ) from embodichain.lab.sim.cfg import ( RigidBodyAttributesCfg as rigid_body_attributes_cfg_cls, @@ -123,12 +126,15 @@ def _ensure_runtime_imports() -> None: "torch": torch_module, "SimulationManager": simulation_manager_cls, "AtomicActionEngine": atomic_action_engine_cls, - "EndEffectorPoseTarget": end_effector_pose_target_cls, + "ActionBinding": action_binding_cls, + "ActionInvocation": action_invocation_cls, + "EndEffectorPoseGoal": end_effector_pose_target_cls, "MoveEndEffector": move_end_effector_cls, "MoveEndEffectorCfg": move_end_effector_cfg_cls, + "MotionPolicy": motion_policy_cls, "Press": press_cls, "PressCfg": press_cfg_cls, - "PressTarget": press_target_cls, + "PressGoal": press_target_cls, "RigidBodyAttributesCfg": rigid_body_attributes_cfg_cls, "RigidObjectCfg": rigid_object_cfg_cls, "VisualMaterialCfg": visual_material_cfg_cls, @@ -481,15 +487,7 @@ def _build_atomic_engine( """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(MoveEndEffector(motion_gen, cfg=MoveEndEffectorCfg())) atomic_engine.register( Press( motion_gen, @@ -497,7 +495,6 @@ def _build_atomic_engine( control_part="arm", hand_control_part="hand", hand_close_qpos=hand_close, - sample_interval=PRESS_SAMPLE_INTERVAL, hand_interp_steps=HAND_INTERP_STEPS, ), ) @@ -579,12 +576,28 @@ 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), + ), + ) ) + is_success = result.plan_success + traj = result.trajectory.positions _sync_cuda() elapsed = time.perf_counter() - start diff --git a/scripts/tutorials/atomic_action/assemble.py b/scripts/tutorials/atomic_action/assemble.py index 1bbd18625..016cecb6e 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 @@ -40,14 +40,17 @@ 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, + GraspGoal, PickUp, PickUpCfg, Place, PlaceCfg, + MotionPolicy, ) from embodichain.lab.sim.cfg import ( JointDrivePropertiesCfg, @@ -416,7 +419,6 @@ def run_assemble_demo( 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 @@ -434,7 +436,6 @@ def run_assemble_demo( 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, ), ) @@ -456,12 +457,28 @@ 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), + ), + ActionInvocation( + "place", + AssembleGoal(affordance=assemble_affordance), + binding, + MotionPolicy(sample_count=PLACE_SAMPLE_INTERVAL), + ), + ) ) + 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..ae72f85b4 100644 --- a/scripts/tutorials/atomic_action/coordinated_pickment.py +++ b/scripts/tutorials/atomic_action/coordinated_pickment.py @@ -40,12 +40,15 @@ 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, Affordance, AtomicActionEngine, - CoordinatedPickTarget, + CoordinatedPickGoal, CoordinatedPickment, CoordinatedPickmentCfg, ObjectSemantics, + MotionPolicy, ) from embodichain.lab.sim.cfg import ( JointDrivePropertiesCfg, @@ -679,7 +682,6 @@ def run_coordinated_pickment_demo( 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, @@ -724,7 +726,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 +739,21 @@ 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), + ), + ) + ) + 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 +766,7 @@ def run_coordinated_pickment_demo( "coordinated_pickment", traj, joint_ids, - pickment_action.get_segment_lengths(), + pickment_action.get_segment_lengths(PICKMENT_SAMPLE_INTERVAL), ) if args.diagnose_plan: diff --git a/scripts/tutorials/atomic_action/coordinated_placement.py b/scripts/tutorials/atomic_action/coordinated_placement.py index 27f46d47a..487a14dd9 100644 --- a/scripts/tutorials/atomic_action/coordinated_placement.py +++ b/scripts/tutorials/atomic_action/coordinated_placement.py @@ -41,17 +41,20 @@ 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, Affordance, AtomicActionEngine, CoordinatedPlacement, CoordinatedPlacementCfg, - CoordinatedPlacementTarget, - GraspTarget, + CoordinatedPlacementGoal, + GraspGoal, HeldObjectState, ObjectSemantics, PickUp, PickUpCfg, - WorldState, + MotionPolicy, + TaskState, ) from embodichain.lab.sim.cfg import ( JointDrivePropertiesCfg, @@ -801,7 +804,6 @@ def run_coordinated_placement_demo( hand_close_qpos=left_close, pre_grasp_distance=PICK_APPROACH_DISTANCE, lift_height=0.12, - sample_interval=PICK_SAMPLE_INTERVAL, hand_interp_steps=10, ), ) @@ -814,7 +816,6 @@ def run_coordinated_placement_demo( 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, ), ) @@ -833,7 +834,6 @@ def run_coordinated_placement_demo( 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, @@ -842,7 +842,7 @@ def run_coordinated_placement_demo( engine = AtomicActionEngine(motion_generator=motion_gen) engine.register(coordinated_action) full_joint_ids = list(range(robot.dof)) - state = WorldState(last_qpos=robot.get_qpos().clone()) + state = engine.initial_context() wait_for_user = prepare_tutorial_scene( sim, args, "Inspect the scene, then press Enter to plan left pick-up..." @@ -857,21 +857,34 @@ def run_coordinated_placement_demo( 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), + left_pick_result = left_pick_action.plan( + 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), ), 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(): + if not left_pick_result.plan_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 + left_pick_traj = left_pick_result.trajectory.positions + state = state.project( + qpos=left_pick_traj[:, -1], + task=left_pick_result.expected_effects.apply( + state.task, left_pick_result.plan_success + ), + ) 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.") @@ -884,21 +897,34 @@ def run_coordinated_placement_demo( 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), + right_pick_result = right_pick_action.plan( + 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), ), state, ) logger.log_info( f"Plan right pan pick-up cost time: {time.time() - start_time:.2f} seconds" ) - if not right_pick_result.success.all(): + if not right_pick_result.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 + right_pick_traj = right_pick_result.trajectory.positions + state = state.project( + qpos=right_pick_traj[:, -1], + task=right_pick_result.expected_effects.apply( + state.task, right_pick_result.plan_success + ), + ) 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.") @@ -966,10 +992,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 +1026,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 +1038,29 @@ 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 + 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), + ), + ), + state, ) + coordinated_success = compiled.plan_success + coordinated_traj = compiled.trajectory.positions + state = compiled.projected_context logger.log_info( "Plan coordinated placement cost time: " f"{time.time() - start_time:.2f} seconds" @@ -1019,7 +1073,10 @@ 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_action.cfg.release, + COORDINATED_SAMPLE_INTERVAL, + ), ) if args.diagnose_plan: diff --git a/scripts/tutorials/atomic_action/hand_over.py b/scripts/tutorials/atomic_action/hand_over.py index 33a509299..bd34ac280 100644 --- a/scripts/tutorials/atomic_action/hand_over.py +++ b/scripts/tutorials/atomic_action/hand_over.py @@ -38,12 +38,15 @@ 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, + MotionPolicy, ) from embodichain.lab.sim.cfg import ( JointDrivePropertiesCfg, @@ -353,7 +356,6 @@ def run_handover_demo( 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 @@ -379,7 +381,6 @@ def run_handover_demo( 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, @@ -398,12 +399,36 @@ def run_handover_demo( # 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), + ), + 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), + ), + ) ) + 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..9b937b1d3 100644 --- a/scripts/tutorials/atomic_action/move_end_effector.py +++ b/scripts/tutorials/atomic_action/move_end_effector.py @@ -30,10 +30,13 @@ 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, + EndEffectorPoseGoal, MoveEndEffector, MoveEndEffectorCfg, + MotionPolicy, ) from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( @@ -72,12 +75,7 @@ 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), - ) - ) + engine.register(MoveEndEffector(motion_gen, cfg=MoveEndEffectorCfg())) poses = torch.stack( [ @@ -99,15 +97,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 +116,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..b0b7d67d8 100644 --- a/scripts/tutorials/atomic_action/move_held_object.py +++ b/scripts/tutorials/atomic_action/move_held_object.py @@ -31,16 +31,19 @@ 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, + EndEffectorPoseGoal, + GraspGoal, + HeldObjectPoseGoal, MoveEndEffector, MoveEndEffectorCfg, MoveHeldObject, MoveHeldObjectCfg, PickUp, PickUpCfg, + MotionPolicy, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.lab.sim.objects import RigidObject @@ -127,11 +130,7 @@ def main() -> None: 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(MoveEndEffector(motion_gen, MoveEndEffectorCfg())) engine.register( PickUp( motion_gen, @@ -140,7 +139,6 @@ def main() -> None: 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, ), ) @@ -150,7 +148,6 @@ def main() -> None: motion_gen, MoveHeldObjectCfg( hand_close_qpos=hand_close, - sample_interval=MOVE_HELD_OBJECT_SAMPLE_INTERVAL, ), ) ) @@ -172,14 +169,33 @@ 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"}, + ) + 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), + ), + ActionInvocation( + "move_held_object", + HeldObjectPoseGoal(object_target), + binding, + MotionPolicy(sample_count=MOVE_HELD_OBJECT_SAMPLE_INTERVAL), + ), + ) ) - if not success.all(): + if not compiled.plan_success.all(): logger.log_warning("Failed to plan MoveHeldObject demo trajectory.") return @@ -201,7 +217,7 @@ def clear_object_dynamics(step_idx: int, _: int) -> None: replay_trajectory( sim, robot, - trajectory, + compiled.trajectory.positions, args, video_prefix="move_held_object_auto_play", hold_steps=POST_TRAJECTORY_STEPS, diff --git a/scripts/tutorials/atomic_action/move_joints.py b/scripts/tutorials/atomic_action/move_joints.py index 2e0addb13..915e4754b 100644 --- a/scripts/tutorials/atomic_action/move_joints.py +++ b/scripts/tutorials/atomic_action/move_joints.py @@ -30,11 +30,14 @@ 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, + JointPositionGoal, MoveJoints, MoveJointsCfg, - NamedJointPositionTarget, + NamedJointPositionGoal, + MotionPolicy, ) from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( @@ -82,7 +85,6 @@ def main() -> None: MoveJoints( motion_gen, cfg=MoveJointsCfg( - sample_interval=MOVE_JOINTS_SAMPLE_INTERVAL, named_joint_positions={"ready": ready}, ), ) @@ -101,13 +103,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", NamedJointPositionGoal("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 +124,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..a8556d357 100644 --- a/scripts/tutorials/atomic_action/pickup.py +++ b/scripts/tutorials/atomic_action/pickup.py @@ -30,10 +30,13 @@ 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, + GraspGoal, PickUp, PickUpCfg, + MotionPolicy, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.lab.sim.objects import RigidObject @@ -140,7 +143,6 @@ def main() -> None: 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, ), ) @@ -157,8 +159,20 @@ 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), + ), + ) + ) + if not compiled.plan_success.all(): logger.log_warning("Failed to plan PickUp demo trajectory.") return @@ -178,7 +192,7 @@ def clear_object_dynamics(step_idx: int, _: int) -> None: replay_trajectory( sim, robot, - trajectory, + compiled.trajectory.positions, args, video_prefix="pickup_cube_auto_play", hold_steps=POST_TRAJECTORY_STEPS, diff --git a/scripts/tutorials/atomic_action/place.py b/scripts/tutorials/atomic_action/place.py index fb3ae0efe..85f2164c5 100644 --- a/scripts/tutorials/atomic_action/place.py +++ b/scripts/tutorials/atomic_action/place.py @@ -30,13 +30,16 @@ 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, + GraspGoal, PickUp, PickUpCfg, Place, PlaceCfg, - PlaceTarget, + PlaceGoal, + MotionPolicy, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.lab.sim.objects import RigidObject @@ -140,7 +143,6 @@ def main() -> None: 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, ), ) @@ -152,7 +154,6 @@ def main() -> None: 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, ), ) @@ -175,20 +176,31 @@ 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), + ), + ActionInvocation( "place", - PlaceTarget( + PlaceGoal( broadcast_waypoint_pose_batch( place_poses, robot.get_qpos().shape[0] ) ), + binding, + MotionPolicy(sample_count=PLACE_SAMPLE_INTERVAL), ), - ] + ) ) - if not success.all(): + if not compiled.plan_success.all(): logger.log_warning("Failed to plan Place demo trajectory.") return @@ -208,7 +220,7 @@ def clear_object_dynamics(step_idx: int, _: int) -> None: replay_trajectory( sim, robot, - trajectory, + compiled.trajectory.positions, args, video_prefix="place_auto_play", hold_steps=POST_TRAJECTORY_STEPS, diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py index 39cd829dd..dc51668a7 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -30,13 +30,16 @@ 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, + EndEffectorPoseGoal, MoveEndEffector, MoveEndEffectorCfg, Press, PressCfg, - PressTarget, + PressGoal, + MotionPolicy, ) from embodichain.lab.sim.cfg import ( RigidBodyAttributesCfg, @@ -161,17 +164,12 @@ def main() -> None: 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(MoveEndEffector(motion_gen, MoveEndEffectorCfg())) engine.register( Press( motion_gen, PressCfg( hand_close_qpos=hand_close, - sample_interval=PRESS_SAMPLE_INTERVAL, hand_interp_steps=HAND_INTERP_STEPS, ), ) @@ -191,15 +189,30 @@ 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), + ), + ) ) - 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/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 16a6b45ab..b79fe21d9 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -14,172 +14,214 @@ # 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 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, - HandOver, - HandOverCfg, - HeldObjectPoseTarget, - JointPositionTarget, - NamedJointPositionTarget, - PlaceTarget, - PressTarget, -) -from embodichain.lab.sim.planners.utils import MoveType, PlanResult -from embodichain.lab.sim.atomic_actions.core import ( - ActionResult, - AtomicAction, + ActionBinding, + ActionInvocation, + Affordance, + AntipodalAffordance, + AssembleGoal, CoordinatedHeldObjectState, - HeldObjectState, - ObjectSemantics, - WorldState, -) -from embodichain.lab.sim.atomic_actions.actions import ( + CoordinatedPickGoal, CoordinatedPickment, CoordinatedPickmentCfg, CoordinatedPlacement, CoordinatedPlacementCfg, + CoordinatedPlacementGoal, + EndEffectorPoseGoal, + GraspGoal, + HandOver, + HandOverCfg, + HeldObjectPoseGoal, + HeldObjectState, + JointPositionGoal, + MotionPolicy, MoveEndEffector, MoveEndEffectorCfg, - MoveJoints, - MoveJointsCfg, MoveHeldObject, MoveHeldObjectCfg, + MoveJoints, + MoveJointsCfg, + NamedJointPositionGoal, + ObjectSemantics, PickUp, PickUpCfg, Place, PlaceCfg, + PlaceGoal, + PlanningContext, Press, PressCfg, + PressGoal, + 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 + + +@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 + robot.dof = ROBOT_DOF - def get_qpos(name=None): + def get_qpos(name: str | None = None) -> torch.Tensor: if name == "arm": return torch.zeros(NUM_ENVS, ARM_DOF) if name == "hand": return torch.zeros(NUM_ENVS, HAND_DOF) - return torch.zeros(NUM_ENVS, TOTAL_DOF) - - robot.get_qpos = get_qpos + return torch.zeros(NUM_ENVS, ROBOT_DOF) - def get_joint_ids(name=None): + def get_joint_ids(name: str | None = None) -> list[int]: if name == "arm": return list(range(ARM_DOF)) if name == "hand": - return list(range(ARM_DOF, TOTAL_DOF)) - return list(range(TOTAL_DOF)) + 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 - 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 _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_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 _binding() -> ActionBinding: + return ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ) - robot.compute_fk = compute_fk - return robot + +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_mock_motion_generator(): - mg = Mock() - mg.robot = _make_mock_robot() - mg.device = torch.device("cpu") - return mg - - -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 + robot.dof = DUAL_ROBOT_DOF - def get_qpos(name=None): - if name == "left_arm": - return torch.zeros(NUM_ENVS, ARM_DOF) - if name == "right_arm": + 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) + return torch.zeros(NUM_ENVS, DUAL_ROBOT_DOF) - robot.get_qpos = get_qpos - - 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": @@ -187,1346 +229,632 @@ 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 _hand_open(): - return torch.zeros(HAND_DOF, dtype=torch.float32) - + 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_close(): - return torch.full((HAND_DOF,), 0.025, 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", + }, + ) -# --------------------------------------------------------------------------- -# MoveEndEffector -# --------------------------------------------------------------------------- +def test_builtin_descriptors_expose_goals_not_legacy_targets() -> None: + assert MoveEndEffector.GoalType is EndEffectorPoseGoal + assert MoveJoints.GoalType == (JointPositionGoal, NamedJointPositionGoal) + 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 + + +def test_move_end_effector_returns_full_robot_timed_plan() -> None: + action = MoveEndEffector(_motion_generator(), MoveEndEffectorCfg()) + context = _context() + + plan = action.plan( + _invocation( + "move_end_effector", + EndEffectorPoseGoal(torch.eye(4)), + sample_count=10, + ), + context, + ) -class TestMoveEndEffectorAction: - def setup_method(self): - self.mg = _make_mock_motion_generator() + 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 = MoveJoints(generator, MoveJointsCfg(named_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_target_type_is_pose_target(self): - assert MoveEndEffector.TargetType is EndEffectorPoseTarget + plan = action.plan( + _invocation("move_joints", NamedJointPositionGoal("ready"), sample_count=8), + context, + ) - def test_default_name_is_explicit(self): - assert MoveEndEffectorCfg().name == "move_end_effector" + assert torch.allclose(plan.trajectory.positions[:, -1, :ARM_DOF], named["ready"]) + assert torch.all(plan.trajectory.positions[:, :, ARM_DOF:] == 0.7) - 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), - ) +def test_pick_and_place_declare_effects_without_mutating_context() -> None: + generator = _motion_generator() + hand_open = torch.zeros(HAND_DOF) + hand_close = torch.ones(HAND_DOF) + pick = PickUp( + generator, + PickUpCfg(hand_open_qpos=hand_open, hand_close_qpos=hand_close), + ) + initial = _context() + semantics = _semantics() + grasp = torch.eye(4).repeat(NUM_ENVS, 1, 1) - 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), - ) + pick_plan = pick.plan( + _invocation("pick_up", GraspGoal(semantics=semantics, grasp_xpos=grasp)), + initial, + ) + picked_task = pick_plan.expected_effects.apply(initial.task, pick_plan.plan_success) - 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), - ) + assert initial.task.get_held_object("arm") is None + assert picked_task.get_held_object("arm") is not None - 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)), - ) + place = Place( + generator, + PlaceCfg(hand_open_qpos=hand_open, hand_close_qpos=hand_close), + ) + picked_context = PlanningContext( + robot=initial.robot, + task=picked_task, + scene=initial.scene, + env_ids=initial.env_ids, + ) + place_plan = place.plan( + _invocation("place", PlaceGoal(torch.eye(4))), + picked_context, + ) + placed_task = place_plan.expected_effects.apply( + picked_task, place_plan.plan_success + ) + assert picked_task.get_held_object("arm") is not None + assert placed_task.get_held_object("arm") is None -# --------------------------------------------------------------------------- -# PickUp -# --------------------------------------------------------------------------- +def test_move_held_object_requires_projected_attachment() -> None: + generator = _motion_generator() + action = MoveHeldObject( + generator, + MoveHeldObjectCfg(hand_close_qpos=torch.ones(HAND_DOF)), + ) + invocation = _invocation( + "move_held_object", + HeldObjectPoseGoal(torch.eye(4)), + sample_count=10, + ) -class TestPickUpAction: - def setup_method(self): - self.mg = _make_mock_motion_generator() + with pytest.raises(ValueError, match="requires an object held"): + action.plan(invocation, _context()) - def test_target_type_is_grasp_target(self): - assert PickUp.TargetType is GraspTarget + held = _held() + task = TaskState( + batch_size=NUM_ENVS, + device="cpu", + held_objects={"arm": held}, + ) + plan = action.plan(invocation, _context(task)) + assert plan.plan_success.all() + assert plan.expected_effects.is_empty - def test_approach_alignment_filter_is_opt_in(self): - assert PickUpCfg().approach_alignment_max_angle 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) - ] - ) +def test_press_uses_invocation_sample_budget() -> None: + generator = _motion_generator() + action = Press(generator, PressCfg(hand_close_qpos=torch.ones(HAND_DOF))) - entity = Mock() - entity.get_local_pose = Mock( - return_value=torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1) - ) + plan = action.plan( + _invocation("press", PressGoal(torch.eye(4)), sample_count=12), + _context(), + ) - sem = ObjectSemantics( - affordance=affordance, - geometry={}, - label="mug", - entity=entity, - ) + assert plan.plan_success.tolist() == [True, True] + assert plan.trajectory.waypoint_count == 12 + assert plan.expected_effects.is_empty - 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) - ] - ) +def test_motion_source_and_sample_count_are_not_action_config_fields() -> None: + with pytest.raises(TypeError): + MoveEndEffectorCfg(motion_source="motion_gen") + with pytest.raises(TypeError): + MoveJointsCfg(sample_interval=10) - 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_move_joints_rejects_binding_with_wrong_goal_skill() -> None: + action = MoveJoints(_motion_generator(), MoveJointsCfg()) + 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.plan(invocation, _context()) - 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) +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 = MoveJoints(generator, MoveJointsCfg()) + 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), + ) -# --------------------------------------------------------------------------- -# MoveHeldObject -# --------------------------------------------------------------------------- + plan = action.plan(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 + + plan = MoveEndEffector(generator).plan( + _invocation( + "move_end_effector", + EndEffectorPoseGoal(waypoints), + sample_count=9, + ), + _context(), + ) + 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]) -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_visits_waypoints_and_rejects_unknown_names() -> None: + action = MoveJoints( + _motion_generator(), + MoveJointsCfg(named_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_default_name_is_explicit(self): - assert ( - MoveHeldObjectCfg(hand_close_qpos=_hand_close()).name == "move_held_object" - ) + plan = action.plan( + _invocation( + "move_joints", + JointPositionGoal(waypoints), + sample_count=7, + ), + _context(), + ) - 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, - ) + 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="Unknown named joint-position goal"): + action.plan( + _invocation("move_joints", NamedJointPositionGoal("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 = PickUp( + generator, + PickUpCfg( + hand_open_qpos=torch.zeros(HAND_DOF), + hand_close_qpos=torch.ones(HAND_DOF), + ), + ) - assert torch.equal(cfg.obj_upright_direction, upright_direction) - assert cfg.pick_rotate_upright == torch.pi / 2 + plan = action.plan( + _invocation( + "pick_up", + GraspGoal(semantics=semantics, grasp_xpos=grasp), + sample_count=20, + ), + _context(), + ) + projected = plan.expected_effects.apply(_context().task, plan.plan_success) - 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), - ) + 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) - 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_press_closes_hand_without_changing_projected_attachment() -> None: + held = _held() + task = TaskState( + batch_size=NUM_ENVS, + device="cpu", + held_objects={"arm": held}, + ) + action = Press( + _motion_generator(), + PressCfg(hand_close_qpos=torch.ones(HAND_DOF), hand_interp_steps=4), ) - 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, + plan = action.plan( + _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() + action = HandOver( + generator, + HandOverCfg( + transfer_hand_open_qpos=torch.zeros(HAND_DOF), + transfer_hand_close_qpos=torch.ones(HAND_DOF), + receive_hand_open_qpos=torch.zeros(HAND_DOF), + receive_hand_close_qpos=torch.ones(HAND_DOF), + middle_object_pose=torch.eye(4), + final_object_pose=torch.eye(4), 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) - ) - - -# --------------------------------------------------------------------------- -# Press -# --------------------------------------------------------------------------- + hold_steps=2, + retreat_steps=5, + ), + ) + original_final_pose = action.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 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), + ) -class TestPressAction: - def setup_method(self): - self.mg = _make_mock_motion_generator() + plan = action.plan(invocation, _dual_context(task)) - def test_target_type_is_press_target(self): - assert Press.TargetType is PressTarget + assert plan.plan_success.all() + assert torch.equal(action.final_object_pose, original_final_pose) - def test_default_name_is_explicit(self): - assert PressCfg(hand_close_qpos=_hand_close()).name == "press" - def test_execute_closes_hand_and_preserves_held_object(self): - cfg = PressCfg( - hand_close_qpos=_hand_close(), - sample_interval=12, +def test_coordinated_pick_returns_full_dof_plan_and_projected_relation() -> None: + action = CoordinatedPickment( + _dual_motion_generator(), + CoordinatedPickmentCfg( + left_hand_open_qpos=torch.zeros(HAND_DOF), + left_hand_close_qpos=torch.ones(HAND_DOF), + right_hand_open_qpos=torch.zeros(HAND_DOF), + right_hand_close_qpos=torch.ones(HAND_DOF), 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 - - -# --------------------------------------------------------------------------- -# HandOver -# --------------------------------------------------------------------------- - - -class TestHandOverAction: - def test_execute_does_not_mutate_cached_final_pose(self): - motion_generator = _make_dual_arm_mock_motion_generator() - action = HandOver( - motion_generator, - HandOverCfg( - transfer_hand_open_qpos=_hand_open(), - transfer_hand_close_qpos=_hand_close(), - receive_hand_open_qpos=_hand_open(), - receive_hand_close_qpos=_hand_close(), - middle_object_pose=torch.eye(4), - final_object_pose=torch.eye(4), - sample_interval=30, - hand_interp_steps=4, - hold_steps=2, - retreat_steps=5, - ), - ) - original_final_pose = action.final_object_pose.clone() - current_object_pose = torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1) - current_object_pose[:, :3, :3] = torch.diag(torch.tensor([-1.0, -1.0, 1.0])) - entity = Mock() - entity.get_local_pose.return_value = current_object_pose - semantics = ObjectSemantics( - affordance=AntipodalAffordance(), - geometry={}, - label="handover-object", - entity=entity, - ) - held = HeldObjectState( + 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_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, DUAL_TOTAL_DOF), - held_objects={"left_arm": held}, - ) - receive_grasp = torch.eye(4).unsqueeze(0).repeat(NUM_ENVS, 1, 1) - action._resolve_receive_grasp = Mock( - return_value=( - receive_grasp, - torch.ones(NUM_ENVS, dtype=torch.bool), - ) - ) - - def plan_from_start(control_part, start_qpos, target_poses, n_waypoints): - return True, start_qpos.unsqueeze(1).repeat(1, n_waypoints, 1) - - action._plan_named_arm_trajectory = Mock(side_effect=plan_from_start) - - action.execute(GraspTarget(semantics=semantics), state) - - assert torch.equal(action.final_object_pose, original_final_pose) - - -# --------------------------------------------------------------------------- -# CoordinatedPickment -# --------------------------------------------------------------------------- - - -class TestCoordinatedPickmentAction: - def setup_method(self): - self.mg = _make_dual_arm_mock_motion_generator() + 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 = action.plan(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, +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 = CoordinatedPickment( + generator, + CoordinatedPickmentCfg( + left_hand_open_qpos=torch.zeros(HAND_DOF), + left_hand_close_qpos=torch.ones(HAND_DOF), + right_hand_open_qpos=torch.zeros(HAND_DOF), + right_hand_close_qpos=torch.ones(HAND_DOF), 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, - 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 -# --------------------------------------------------------------------------- + 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() + plan = action.plan(invocation, context) + projected = plan.expected_effects.apply(context.task, plan.plan_success) -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, + 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] + + +def test_coordinated_placement_projects_release_and_support_attachment() -> None: + generator = _dual_motion_generator() + action = CoordinatedPlacement( + generator, + CoordinatedPlacementCfg( + placing_hand_open_qpos=torch.zeros(HAND_DOF), + placing_hand_close_qpos=torch.ones(HAND_DOF), + support_hand_close_qpos=torch.ones(HAND_DOF), hand_interp_steps=4, hold_steps=3, retreat_steps=5, - lift_height=0.08, - ) - self.action = CoordinatedPlacement(self.mg, cfg=self.cfg) - - def test_named_arm_planning_forwards_action_configuration(self): - self.action.builder.plan_arm_traj = Mock( - return_value=( - torch.ones(NUM_ENVS, dtype=torch.bool), - torch.zeros(NUM_ENVS, 4, ARM_DOF), - ) - ) - - self.action._plan_named_arm_trajectory( - "left_arm", - torch.zeros(NUM_ENVS, ARM_DOF), - torch.eye(4).reshape(1, 1, 4, 4).repeat(NUM_ENVS, 1, 1, 1), - 4, - ) - - assert self.action.builder.plan_arm_traj.call_args.kwargs["cfg"] is 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 - - def test_target_type_is_coordinated_placement_target(self): - assert CoordinatedPlacement.TargetType is CoordinatedPlacementTarget - - def test_init_sets_dual_arm_and_hand_joint_ids(self): - assert self.action.dual_arm_joint_ids == list(range(DUAL_ARM_DOF)) - assert self.action.placing_arm_joint_ids == list(range(ARM_DOF)) - assert self.action.support_arm_joint_ids == list(range(ARM_DOF, DUAL_ARM_DOF)) - assert self.action.placing_hand_joint_ids == list( - range(DUAL_ARM_DOF, DUAL_ARM_DOF + HAND_DOF) - ) - assert self.action.support_hand_joint_ids == list( - range(DUAL_ARM_DOF + HAND_DOF, DUAL_TOTAL_DOF) - ) - assert self.action.joint_ids == list(range(DUAL_TOTAL_DOF)) - - def test_resolve_target_composes_object_and_tcp_poses(self): - target, 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.tolist() == [True] * NUM_ENVS - 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) - - -# --------------------------------------------------------------------------- -# MoveJoints + cuRobo motion_gen routing -# --------------------------------------------------------------------------- - - -class TestMoveJointsCurobo: - def setup_method(self): - # shared per-test mg is created in each test (result shapes differ). - pass - - 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) - ) + ), + ) + 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 = action.plan(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 not None + assert projected.get_held_object("right_arm").semantics is support.semantics + + +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 = CoordinatedPickment( + generator, + CoordinatedPickmentCfg( + left_hand_open_qpos=torch.zeros(HAND_DOF), + left_hand_close_qpos=torch.ones(HAND_DOF), + right_hand_open_qpos=torch.zeros(HAND_DOF), + right_hand_close_qpos=torch.ones(HAND_DOF), + ), + ) + 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"): + pick.plan(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 = CoordinatedPlacement( + generator, + CoordinatedPlacementCfg( + placing_hand_open_qpos=torch.zeros(HAND_DOF), + placing_hand_close_qpos=torch.ones(HAND_DOF), + support_hand_close_qpos=torch.ones(HAND_DOF), + ), + ) + 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"): + placement.plan(placement_invocation, _dual_context()) diff --git a/tests/sim/atomic_actions/test_core.py b/tests/sim/atomic_actions/test_core.py index 5c7133ac4..574b49b8e 100644 --- a/tests/sim/atomic_actions/test_core.py +++ b/tests/sim/atomic_actions/test_core.py @@ -14,485 +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) - assert s.env_mask.tolist() == [True] - - def test_rejects_mismatched_pose_batches(self): - sem = ObjectSemantics(affordance=Affordance(), geometry={}) - with pytest.raises(ValueError, match="same batch size"): - HeldObjectState( - semantics=sem, - object_to_eef=torch.eye(4).unsqueeze(0), - grasp_xpos=torch.eye(4).unsqueeze(0).repeat(2, 1, 1), - ) - - def test_rejects_invalid_env_mask(self): - sem = ObjectSemantics(affordance=Affordance(), geometry={}) - with pytest.raises(ValueError, match="env_mask"): - HeldObjectState( - semantics=sem, - object_to_eef=torch.eye(4).unsqueeze(0), - grasp_xpos=torch.eye(4).unsqueeze(0), - env_mask=torch.ones(2, dtype=torch.bool), - ) - - -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}, - ) - 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), - ) - ws = WorldState( - last_qpos=torch.zeros(1, 14), - coordinated_held_objects={("left_arm", "right_arm"): held}, - ) - assert ws.get_coordinated_held_object("left_arm", "right_arm") is held - - def test_with_updates_does_not_alias_held_state_dictionaries(self): - ws = WorldState(last_qpos=torch.zeros(1, 6)) - successor = ws.with_updates(last_qpos=torch.ones(1, 6)) - successor.held_objects["arm"] = HeldObjectState( - semantics=ObjectSemantics(affordance=Affordance(), geometry={}), - object_to_eef=torch.eye(4).unsqueeze(0), - grasp_xpos=torch.eye(4).unsqueeze(0), - ) - assert ws.held_objects == {} - def test_rejects_non_batched_last_qpos(self): - with pytest.raises(ValueError, match="last_qpos"): - WorldState(last_qpos=torch.zeros(6)) - def test_rejects_held_state_with_different_batch(self): - held = 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="batch size"): - WorldState( - last_qpos=torch.zeros(2, 6), - held_objects={"arm": held}, - ) - - def test_broadcasts_unbatched_held_state_at_world_boundary(self): - held = HeldObjectState( - semantics=ObjectSemantics(affordance=Affordance(), geometry={}), - object_to_eef=torch.eye(4), - grasp_xpos=torch.eye(4), - ) +def test_action_binding_is_role_based_and_immutable() -> None: + binding = ActionBinding( + manipulators={"primary": "left_arm"}, + end_effectors={"primary": "left_hand"}, + ) - world = WorldState( - last_qpos=torch.zeros(2, 6), - held_objects={"arm": held}, - ) + 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") - normalized = world.get_held_object("arm") - assert normalized is not None - assert normalized is not held - assert normalized.object_to_eef.shape == (2, 4, 4) - assert normalized.grasp_xpos.shape == (2, 4, 4) - assert normalized.env_mask.tolist() == [True, True] - - def test_masked_merge_applies_new_hold_only_to_successful_envs(self): - batch_size = 2 - previous = WorldState(last_qpos=torch.zeros(batch_size, 6)) - held = HeldObjectState( - semantics=ObjectSemantics(affordance=Affordance(), geometry={}), - object_to_eef=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), - grasp_xpos=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), - ) - candidate = WorldState( - last_qpos=torch.ones(batch_size, 6), - held_objects={"arm": held}, - ) - merged = previous.masked_merge( - candidate, torch.tensor([True, False], dtype=torch.bool) +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 merged.get_held_object("arm").env_mask.tolist() == [True, False] - def test_masked_merge_preserves_removed_hold_in_failed_envs(self): - batch_size = 2 - held = HeldObjectState( - semantics=ObjectSemantics(affordance=Affordance(), geometry={}), - object_to_eef=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), - grasp_xpos=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), - ) - previous = WorldState( - last_qpos=torch.zeros(batch_size, 6), - held_objects={"arm": held}, - ) - candidate = WorldState(last_qpos=torch.ones(batch_size, 6)) +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) + - merged = previous.masked_merge( - candidate, torch.tensor([True, False], dtype=torch.bool) +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)}, ) + ) - assert merged.get_held_object("arm").env_mask.tolist() == [False, True] + resolved = resolve_pose_goal(reference, context, name="xpos") - def test_masked_merge_preserves_coordinated_hold_in_failed_envs(self): - batch_size = 2 - held = CoordinatedHeldObjectState( - semantics=ObjectSemantics(affordance=Affordance(), geometry={}), - left_object_to_eef=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), - right_object_to_eef=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), - left_grasp_xpos=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), - right_grasp_xpos=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), - ) - previous = WorldState( - last_qpos=torch.zeros(batch_size, 12), - coordinated_held_objects={("left_arm", "right_arm"): held}, - ) - candidate = WorldState(last_qpos=torch.ones(batch_size, 12)) + 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",) - merged = previous.masked_merge( - candidate, torch.tensor([True, False], dtype=torch.bool) + +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)}, + ) + ) + with pytest.raises(ValueError, match="confidence"): + resolve_pose_goal( + SceneEntityPose("cup", minimum_confidence=0.8), + context, + name="xpos", ) - coordinated = merged.get_coordinated_held_object("left_arm", "right_arm") - assert coordinated is not None - assert coordinated.env_mask.tolist() == [False, True] - def test_masked_merge_updates_qpos_per_environment(self): - previous = WorldState(last_qpos=torch.zeros(2, 3)) - candidate = WorldState(last_qpos=torch.ones(2, 3)) +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), + ) - merged = previous.masked_merge( - candidate, torch.tensor([True, False], dtype=torch.bool) - ) + 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) - assert torch.equal(merged.last_qpos[0], torch.ones(3)) - assert torch.equal(merged.last_qpos[1], torch.zeros(3)) +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, + ) -class TestActionResult: - def test_bool_success_is_normalized_to_per_environment_tensor(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.tolist() == [True, True] - assert res.trajectory.shape == (2, 10, 8) - assert res.next_state is ws + result = TimedTrajectory.concatenate((first, second)) - def test_scalar_tensor_success_is_normalized(self): - res = ActionResult( - success=torch.tensor(False), - trajectory=torch.zeros(2, 0, 3), - next_state=WorldState(last_qpos=torch.zeros(2, 3)), - ) - assert res.success.tolist() == [False, False] - - def test_rejects_non_boolean_success_tensor(self): - with pytest.raises(TypeError, match="torch.bool"): - ActionResult( - success=torch.ones(2), - trajectory=torch.zeros(2, 0, 3), - next_state=WorldState(last_qpos=torch.zeros(2, 3)), - ) - - def test_rejects_wrong_success_shape(self): - with pytest.raises(ValueError, match="success"): - ActionResult( - success=torch.ones(3, dtype=torch.bool), - trajectory=torch.zeros(2, 0, 3), - next_state=WorldState(last_qpos=torch.zeros(2, 3)), - ) - - def test_rejects_trajectory_state_dof_mismatch(self): - with pytest.raises(ValueError, match="batch/DoF"): - ActionResult( - success=torch.ones(2, dtype=torch.bool), - trajectory=torch.zeros(2, 4, 4), - next_state=WorldState(last_qpos=torch.zeros(2, 3)), - ) - - -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 - assert cfg.plan_opts is None - - def test_rejects_unsupported_interpolation_type(self): - with pytest.raises(ValueError, match="interpolation_type"): - ActionCfg(interpolation_type="cubic") + 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..2ae980638 100644 --- a/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py +++ b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py @@ -46,10 +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 + EndEffectorPoseGoal, + MotionPolicy, MoveEndEffector, MoveEndEffectorCfg, ) @@ -88,17 +89,7 @@ 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", - ) + engine.register(MoveEndEffector(mg, MoveEndEffectorCfg())) return sim, robot, engine @@ -140,11 +131,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 128dd190f..9ae563a20 100644 --- a/tests/sim/atomic_actions/test_engine.py +++ b/tests/sim/atomic_actions/test_engine.py @@ -14,241 +14,164 @@ # 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, + ActionCfg, + ActionInvocation, + ActionPlan, AtomicAction, - HeldObjectState, - ObjectSemantics, - WorldState, -) -from embodichain.lab.sim.atomic_actions.engine import ( AtomicActionEngine, - get_registered_actions, + JointPositionGoal, + MotionPolicy, + PlanningContext, register_action, + get_registered_actions, unregister_action, ) -# --------------------------------------------------------------------------- -# Global registry (kept from old design) -# --------------------------------------------------------------------------- +class StubAction(AtomicAction[JointPositionGoal]): + """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, + invocation: ActionInvocation[JointPositionGoal], + context: PlanningContext, + ) -> ActionPlan: + goal = self.require_goal(invocation) + target = goal.qpos.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( + invocation, + 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 +def _engine(batch_size: int = 2, robot_dof: int = 3) -> AtomicActionEngine: + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = robot_dof + robot.get_qpos.return_value = torch.zeros(batch_size, robot_dof) + robot.get_qvel.return_value = torch.zeros(batch_size, robot_dof) + generator = Mock() + generator.robot = robot + generator.device = torch.device("cpu") + generator.planner.cfg.planner_type = "stub_planner" + return AtomicActionEngine(generator) - def test_unregister(self): - register_action("_test_dummy", Mock()) - unregister_action("_test_dummy") - assert "_test_dummy" not in get_registered_actions() - def test_unregister_nonexistent_is_noop(self): - unregister_action("_does_not_exist") +def _invocation(qpos: torch.Tensor) -> ActionInvocation[JointPositionGoal]: + return ActionInvocation( + skill_id="stub", + goal=JointPositionGoal(qpos), + binding=ActionBinding(manipulators={"primary": "all"}), + motion_policy=MotionPolicy(sample_count=2), + ) - 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 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") -# --------------------------------------------------------------------------- -# Engine run() semantics -# --------------------------------------------------------------------------- +def test_engine_compile_projects_terminal_state_between_actions() -> None: + engine = _engine() + engine.register(StubAction(engine.motion_generator, ActionCfg(name="stub"))) + first = torch.ones(2, 3) + second = torch.full((2, 3), 2.0) -NUM_ENVS = 2 -TOTAL_DOF = 8 + 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 _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_rejects_initial_state_with_wrong_batch_size(self): - with pytest.raises(ValueError, match="batch size"): - self.engine.run([], state=WorldState(last_qpos=torch.zeros(1, TOTAL_DOF))) - - 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_compile_holds_failed_rows_for_remaining_actions() -> None: + engine = _engine() + engine.register(StubAction(engine.motion_generator, ActionCfg(name="stub"))) + 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 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(engine.motion_generator, ActionCfg(name="first")) + second = StubAction(engine.motion_generator, ActionCfg(name="second")) + engine.register(first) + with pytest.raises(ValueError, match="already registered"): + engine.register(second) + + +def test_engine_rejects_plan_for_a_different_skill() -> None: + engine = _engine() + action = StubAction(engine.motion_generator, ActionCfg(name="stub")) + original_plan = action.plan + + def wrong_skill_plan( + invocation: ActionInvocation[JointPositionGoal], + context: PlanningContext, + ) -> ActionPlan: + return replace(original_plan(invocation, context), skill_id="other") + + action.plan = wrong_skill_plan # type: ignore[method-assign] + engine.register(action) + + with pytest.raises(ValueError, match="must match its invocation"): + 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 186dbddfa..12318ce36 100644 --- a/tests/sim/atomic_actions/test_engine_per_env.py +++ b/tests/sim/atomic_actions/test_engine_per_env.py @@ -14,129 +14,311 @@ # 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.affordance import Affordance -from embodichain.lab.sim.atomic_actions import EndEffectorPoseTarget -from embodichain.lab.sim.atomic_actions.core import ( +import pytest +import torch + +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, ActionCfg, - ActionResult, + ActionInvocation, + ActionPlan, + Affordance, AtomicAction, + AtomicActionEngine, + EndEffectorPoseGoal, + EntityState, + ExecutionEventKind, + ExecutionStatus, HeldObjectState, + MotionPolicy, ObjectSemantics, - WorldState, + PlanningContext, + RecoveryPolicy, + 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 _HeldStateAction(_StubAction): - def __init__(self, mg, success_vec, *, set_held): - super().__init__(mg, success_vec) - self._set_held = set_held - - def execute(self, target, state): - result = super().execute(target, state) - held_objects = {} - if self._set_held: - batch_size = state.batch_size - held_objects["arm"] = HeldObjectState( - semantics=ObjectSemantics( - affordance=Affordance(), geometry={}, label="test-object" - ), - object_to_eef=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), - grasp_xpos=torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1), - ) - return ActionResult( - success=result.success, - trajectory=result.trajectory, - next_state=result.next_state.with_updates(held_objects=held_objects), - ) +class DynamicAction(AtomicAction[EndEffectorPoseGoal]): + """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",) -class TestRunPerEnv: - def test_failed_env_holds(self): - mg = Mock() - mg.robot.get_qpos = lambda: torch.zeros(3, 3) - mg.robot.dof = 3 - mg.device = torch.device("cpu") - eng = AtomicActionEngine(mg) - # env 1 fails step 2 - eng.register(_StubAction(mg, [True, True, True]), name="a") - eng.register(_StubAction(mg, [True, False, True]), name="b") - eng.register(_StubAction(mg, [True, True, True]), name="c") - success, traj, state = eng.run( - steps=[ - ("a", EndEffectorPoseTarget(xpos=torch.eye(4))), - ("b", EndEffectorPoseTarget(xpos=torch.eye(4))), - ("c", EndEffectorPoseTarget(xpos=torch.eye(4))), - ] - ) - assert success.tolist() == [True, False, True] - assert traj.shape[1] == 12 # 3 steps * 4 waypoints - # env 1's rows after its failure should equal its pre-failure qpos (held) - # all zeros here, so just check shape and that env 0/2 advanced - assert state.last_qpos.shape == (3, 3) - - def test_failed_env_does_not_acquire_successful_env_held_state(self): - mg = Mock() - mg.robot.get_qpos = lambda: torch.zeros(3, 3) - mg.robot.dof = 3 - mg.device = torch.device("cpu") - engine = AtomicActionEngine(mg) - engine.register( - _HeldStateAction(mg, [True, False, True], set_held=True), name="pick" + def __init__(self, motion_generator) -> None: + super().__init__(motion_generator, ActionCfg(name="dynamic")) + self.plan_count = 0 + + def plan( + self, + invocation: ActionInvocation[EndEffectorPoseGoal], + context: PlanningContext, + ) -> ActionPlan: + goal = self.require_goal(invocation) + self.plan_count += 1 + pose = resolve_pose_goal(goal.xpos, context, name="xpos") + target = pose[:, 0, 3].unsqueeze(1).expand_as(context.robot.qpos) + return self.build_plan( + invocation, + context, + success=True, + trajectory=torch.stack([context.robot.qpos, target], dim=1), ) - _, _, state = engine.run([("pick", EndEffectorPoseTarget(xpos=torch.eye(4)))]) - assert state.get_held_object("arm").env_mask.tolist() == [True, False, True] +class EffectAction(DynamicAction): + """Dynamic test action that declares an attachment effect.""" - def test_failed_env_preserves_held_state_when_successful_envs_release(self): - mg = Mock() - mg.robot.get_qpos = lambda: torch.zeros(3, 3) - mg.robot.dof = 3 - mg.device = torch.device("cpu") - engine = AtomicActionEngine(mg) - engine.register( - _HeldStateAction(mg, [True, False, True], set_held=False), name="place" + skill_id: ClassVar[str] = "effect" + + def plan( + self, + invocation: ActionInvocation[EndEffectorPoseGoal], + context: PlanningContext, + ) -> ActionPlan: + goal = self.require_goal(invocation) + 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" ) held = HeldObjectState( - semantics=ObjectSemantics( - affordance=Affordance(), geometry={}, label="test-object" - ), - object_to_eef=torch.eye(4).unsqueeze(0).repeat(3, 1, 1), - grasp_xpos=torch.eye(4).unsqueeze(0).repeat(3, 1, 1), + semantics=semantics, + object_to_eef=torch.eye(4), + grasp_xpos=torch.eye(4), ) - initial = WorldState(last_qpos=torch.zeros(3, 3), held_objects={"arm": held}) - - _, _, state = engine.run( - [("place", EndEffectorPoseTarget(xpos=torch.eye(4)))], state=initial + return self.build_plan( + invocation, + context, + success=True, + trajectory=torch.stack([context.robot.qpos, target], dim=1), + expected_effects=StateDelta(held_object_updates={"arm": held}), ) - assert state.get_held_object("arm").env_mask.tolist() == [False, True, False] + +def _engine() -> tuple[AtomicActionEngine, DynamicAction]: + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = 2 + robot.get_qpos.return_value = torch.zeros(1, 2) + robot.get_qvel.return_value = torch.zeros(1, 2) + generator = Mock() + generator.robot = robot + generator.device = torch.device("cpu") + generator.planner.cfg.planner_type = "stub" + engine = AtomicActionEngine(generator) + action = DynamicAction(generator) + 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 tick.command is not None + + +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.motion_generator) + 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(engine.motion_generator)) + 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..2c949ce95 100644 --- a/tests/sim/atomic_actions/test_motion_source_e2e.py +++ b/tests/sim/atomic_actions/test_motion_source_e2e.py @@ -25,10 +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 ( + EndEffectorPoseGoal, + MotionPolicy, MoveEndEffector, MoveEndEffectorCfg, ) @@ -59,12 +60,7 @@ 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") + engine.register(MoveEndEffector(mg, MoveEndEffectorCfg())) return sim, robot, engine def _teardown(self, sim): @@ -86,11 +82,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_motion_source.py b/tests/sim/atomic_actions/test_trajectory_motion_source.py index 3d742682b..1edcdd0f2 100644 --- a/tests/sim/atomic_actions/test_trajectory_motion_source.py +++ b/tests/sim/atomic_actions/test_trajectory_motion_source.py @@ -25,7 +25,7 @@ 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"): @@ -82,9 +82,8 @@ def test_shared_action_cfg_forwards_copied_plan_options(self): positions=torch.zeros(2, 6, 6), ) builder = TrajectoryBuilder(mg) - cfg = ActionCfg( + cfg = MotionPolicy( motion_source="motion_gen", - control_part="arm", plan_opts=PlanOptions(), ) @@ -110,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 = [ @@ -140,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)] @@ -173,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] @@ -192,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) @@ -206,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", ), ) @@ -225,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", ), ) @@ -247,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] @@ -270,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", ), ) @@ -289,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), @@ -310,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..81edea96e 100644 --- a/tests/sim/planners/test_curobo_planner.py +++ b/tests/sim/planners/test_curobo_planner.py @@ -645,17 +645,7 @@ 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", - ) + engine.register(MoveEndEffector(motion_generator, MoveEndEffectorCfg())) return engine @@ -663,7 +653,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 +666,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 +685,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 +707,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 +720,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())) From 1c76efcbfc68b986ce5e746320f1654d0ec9092e Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 2 Aug 2026 15:40:13 +0000 Subject: [PATCH 3/5] wip --- .agents/skills/add-atomic-action/SKILL.md | 24 +- .../topics/atomic-actions/atomic-actions.md | 15 +- ...hain.lab.sim.atomic_actions.primitives.rst | 1 - .../embodichain.lab.sim.atomic_actions.rst | 1 - .../sim/atomic_actions/builtin_actions.md | 559 ++++++++++++++++-- .../overview/sim/atomic_actions/index.md | 327 ++++++++-- docs/source/tutorial/atomic_actions.rst | 56 +- .../lab/sim/atomic_actions/__init__.py | 4 +- embodichain/lab/sim/atomic_actions/core.py | 95 ++- embodichain/lab/sim/atomic_actions/engine.py | 102 +++- .../lab/sim/atomic_actions/execution.py | 8 +- .../sim/atomic_actions/primitives/__init__.py | 2 - .../primitives/coordinated_pickment.py | 14 +- .../primitives/coordinated_placement.py | 13 +- .../atomic_actions/primitives/hand_over.py | 19 +- .../primitives/move_end_effector.py | 5 +- .../primitives/move_held_object.py | 18 +- .../atomic_actions/primitives/move_joints.py | 67 +-- .../sim/atomic_actions/primitives/pick_up.py | 35 +- .../sim/atomic_actions/primitives/place.py | 25 +- .../sim/atomic_actions/primitives/press.py | 17 +- embodichain/lab/sim/atomic_actions/runtime.py | 81 +++ .../lab/sim/atomic_actions/trajectory.py | 8 +- examples/sim/planners/curobo_planner.py | 2 +- .../move_end_effector_benchmark.py | 2 +- .../move_held_object_benchmark.py | 4 +- .../atomic_action/move_joints_benchmark.py | 6 +- .../atomic_action/pickup_benchmark.py | 1 - .../atomic_action/place_benchmark.py | 2 - .../atomic_action/press_benchmark.py | 3 +- scripts/tutorials/atomic_action/assemble.py | 2 - .../atomic_action/coordinated_pickment.py | 1 - .../atomic_action/coordinated_placement.py | 9 +- scripts/tutorials/atomic_action/hand_over.py | 2 - .../atomic_action/move_end_effector.py | 2 +- .../atomic_action/move_held_object.py | 4 +- .../tutorials/atomic_action/move_joints.py | 4 +- scripts/tutorials/atomic_action/pickup.py | 1 - scripts/tutorials/atomic_action/place.py | 2 - scripts/tutorials/atomic_action/press.py | 3 +- tests/sim/atomic_actions/test_actions.py | 214 ++++--- .../test_curobo_motion_source_e2e.py | 2 +- tests/sim/atomic_actions/test_engine.py | 73 ++- .../sim/atomic_actions/test_engine_per_env.py | 10 +- .../atomic_actions/test_motion_source_e2e.py | 2 +- tests/sim/planners/test_curobo_planner.py | 2 +- 46 files changed, 1468 insertions(+), 381 deletions(-) create mode 100644 embodichain/lab/sim/atomic_actions/runtime.py diff --git a/.agents/skills/add-atomic-action/SKILL.md b/.agents/skills/add-atomic-action/SKILL.md index d2e3b00a5..dbf587540 100644 --- a/.agents/skills/add-atomic-action/SKILL.md +++ b/.agents/skills/add-atomic-action/SKILL.md @@ -6,8 +6,10 @@ description: Add a new simulation atomic action or motion primitive to EmbodiCha # Add Atomic Action Add an action-owned goal and a side-effect-free `AtomicAction.plan()` -implementation. Keep task-graph/MLLM logic, simulator stepping, controller I/O, -and physical-effect commits outside the action. +implementation. The engine owns all motion-planning resources; action +constructors accept only implementation configuration. Keep task-graph/MLLM +logic, simulator stepping, controller I/O, and physical-effect commits outside +the action. ## Read the current contracts @@ -22,15 +24,20 @@ Inspect only the files relevant to the requested skill: | 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 -plan = action.plan(invocation: ActionInvocation[Goal], context: PlanningContext) +plan = engine.plan(invocation: ActionInvocation[Goal], context: PlanningContext) ``` +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()`. @@ -94,7 +101,6 @@ from embodichain.lab.sim.atomic_actions import ( AtomicAction, PlanningContext, StateDelta, - TrajectoryBuilder, ) @@ -103,9 +109,8 @@ class Push(AtomicAction[PushGoal]): GoalType: ClassVar[type] = PushGoal manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) - def __init__(self, motion_generator, cfg: PushCfg | None = None) -> None: - super().__init__(motion_generator, cfg or PushCfg()) - self.builder = TrajectoryBuilder(motion_generator) + def __init__(self, cfg: PushCfg | None = None) -> None: + super().__init__(cfg or PushCfg()) def plan( self, @@ -145,6 +150,8 @@ class Push(AtomicAction[PushGoal]): Follow these invariants: +- 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 @@ -164,7 +171,7 @@ Follow these invariants: Register an instance by its class-level `skill_id`: ```python -engine.register(Push(motion_generator, PushCfg())) +engine.register(Push(PushCfg())) ``` Use the global registry only for discoverable third-party classes: @@ -225,6 +232,7 @@ then use the `pre-commit-check` skill before committing. | 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 planner/recovery knobs in action config | Move them to invocation policies. | +| Pass a motion generator to each action | Pass it once to `AtomicActionEngine`; construct actions from config 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`. | diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index 9441ee8a0..a0a4c1e8c 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -5,7 +5,7 @@ Atomic actions are side-effect-free, environment-batched planners: ```python -plan = action.plan(invocation: ActionInvocation, context: PlanningContext) +plan = engine.plan(invocation: ActionInvocation, context: PlanningContext) ``` There is no `ActionTarget`, `WorldState`, `ActionResult`, `execute()`, or @@ -24,6 +24,13 @@ 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, and +shared `TrajectoryBuilder`. Actions accept only implementation configuration +and borrow those services after `engine.register(action)` or +`engine.plan_action(action, invocation, context)`. A bound action cannot be +reused by another engine. + ## Static compilation Register configured action instances by their class-level stable `skill_id`, @@ -38,6 +45,10 @@ 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 `engine.plan_action(...)` for an unregistered configured instance when an +application needs multiple variants with the same stable `skill_id`; the action +still uses the engine's single motion generator. + ## Dynamic execution and recovery `SceneEntityPose(entity_id, relative_pose)` is resolved from the latest scene @@ -77,7 +88,7 @@ lift distances, and grasp constraints. | Skill ID | Goal type | Roles | |---|---|---| | `move_end_effector` | `EndEffectorPoseGoal` | manipulator `primary` | -| `move_joints` | `JointPositionGoal`, `NamedJointPositionGoal` | manipulator `primary` | +| `move_joints` | `JointPositionGoal` (`target` is explicit qpos or a configured 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` | 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 c245c4ada..8692f3194 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 @@ -42,7 +42,6 @@ full-robot timed trajectory and uncommitted expected effects. EndEffectorPoseGoal JointPositionGoal - NamedJointPositionGoal GraspGoal HeldObjectPoseGoal PlaceGoal 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 5e7990410..d38028b53 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 @@ -41,7 +41,6 @@ embodichain.lab.sim.atomic_actions EndEffectorPoseGoal JointPositionGoal - NamedJointPositionGoal GraspGoal HeldObjectPoseGoal PlaceGoal diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index 8372c6d08..62afb8841 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -1,44 +1,519 @@ +(builtin-actions)= + # Built-in atomic actions -All built-ins implement `plan(invocation, context) -> ActionPlan`. Motion and -recovery settings belong to the invocation rather than each action config. - -| Skill ID | Goal | Semantic roles | Expected task effect | -|---|---|---|---| -| `move_end_effector` | `EndEffectorPoseGoal` | manipulator `primary` | none | -| `move_joints` | `JointPositionGoal` or `NamedJointPositionGoal` | manipulator `primary` | none | -| `pick_up` | `GraspGoal` | manipulator/end effector `primary` | attach object | -| `move_held_object` | `HeldObjectPoseGoal` | manipulator/end effector `primary` | none | -| `place` | `PlaceGoal` or `AssembleGoal` | manipulator/end effector `primary` | detach object | -| `press` | `PressGoal` | manipulator/end effector `primary` | none | -| `coordinated_pickment` | `CoordinatedPickGoal` | `left`, `right` | create coordinated attachment | -| `coordinated_placement` | `CoordinatedPlacementGoal` | `placing`, `support` | update/remove attachments | -| `hand_over` | `GraspGoal` | `source`, `destination` | transfer attachment | - -## Pose goals - -Pose-valued goals accept explicit tensors. Selected goals also accept a -`SceneEntityPose(entity_id, relative_pose=...)`, which is resolved from every -new `SceneSnapshot` during planning or replanning. Explicit tensors may use -`(4, 4)` or `(B, 4, 4)` shapes; waypoint-capable goals additionally accept -`(B, N, 4, 4)`. - -## Configuration boundaries - -Action configs contain implementation-owned behavior such as gripper open/close -positions, phase split counts, lift distance, and grasp-selection constraints. -They do not contain planner choice, motion source, trajectory sample count, -velocity limits, recovery budgets, or dynamic-goal thresholds. Those reusable -choices live in `MotionPolicy` and `RecoveryPolicy`. - -`MoveJoints` and `MoveEndEffector` resolve the concrete manipulator entirely from -`ActionBinding`. Complex manipulation skills currently validate their semantic -bindings against the configured hardware resources used by their phase-specific -hand parameters. - -## Planning versus physical success - -`ActionPlan.plan_success` reports whether a valid motion was produced per -environment. It does not prove that a grasp, release, contact, or handover -occurred. Such actions declare a `StateDelta`; an `ExecutionSession` commits it -only after the caller supplies a successful semantic-effect verification mask. +```{currentmodule} embodichain.lab.sim.atomic_actions +``` + +EmbodiChain ships nine built-in action implementations with stable skill IDs; +applications register the configured instances they need with an engine. +`Place` additionally accepts an `AssembleGoal`, so assembly reuses the same +release primitive instead of introducing a tenth skill ID. + +All built-ins implement +`plan(invocation, context) -> ActionPlan`. Their constructors accept only +action configuration; the owning `AtomicActionEngine` supplies one shared +motion generator and trajectory builder during `register()` or +`plan_action()`. Generic motion and recovery choices belong to the invocation, +not the action config. + +```{note} +The current manipulation primitives use gripper open/close joint positions. +Replacing a gripper with a dexterous hand requires a hand-command abstraction +or new hand-specific phases; it is not yet a drop-in config change. +``` + +## Visual catalog + +The animations below are the focused simulator demos under +`scripts/tutorials/atomic_action/`. + +::::{grid} 1 2 2 2 +:gutter: 2 + +:::{grid-item-card} `MoveEndEffector` +:link: builtin-move-end-effector +:link-type: ref + +`move_end_effector` · free-space EEF pose motion + +MoveEndEffector demo +::: + +:::{grid-item-card} `MoveJoints` +:link: builtin-move-joints +:link-type: ref + +`move_joints` · explicit or named joint-space motion + +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 + +`coordinated_pickment` · dual-arm shared-object pick + +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 task state | Expected task effect | +|---|---|---|---|---| +| `move_end_effector` | `EndEffectorPoseGoal` | manipulator `primary` | none | none | +| `move_joints` | `JointPositionGoal` | manipulator `primary` | none | none | +| `pick_up` | `GraspGoal` | manipulator + end effector `primary` | semantic object/entity | attach object to `primary` manipulator | +| `move_held_object` | `HeldObjectPoseGoal` | manipulator + end effector `primary` | object held by `primary` | preserve attachment | +| `place` | `PlaceGoal`, `AssembleGoal` | manipulator + end effector `primary` | `AssembleGoal` requires an object held by `primary`; ordinary `PlaceGoal` has no planner-enforced attachment precondition | detach object | +| `press` | `PressGoal` | manipulator + end effector `primary` | none | none | +| `coordinated_pickment` | `CoordinatedPickGoal` | manipulator + end effector `left`, `right` | semantic object/entity | create coordinated attachment; clear individual attachments | +| `coordinated_placement` | `CoordinatedPlacementGoal` | manipulator + end effector `placing`, `support` | one individually held object per arm | optionally detach placing object; preserve support attachment | +| `hand_over` | `GraspGoal` | manipulator + end effector `source`, `destination` | object held by source arm | transfer attachment to destination arm | + +`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) +) +``` + +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 concrete robot resources selected for this call; +- the **action config** carries hardware constants and phase-specific behavior; +- `MotionPolicy` carries sample count, timing, motion source, limits, + collision choice, and planner options; +- `RecoveryPolicy` carries all replan/retry thresholds and budgets. + +Complex manipulation actions currently validate that their semantic bindings +match concrete resources configured for their hand qpos and multi-part phase +assembly. `MoveEndEffector` and `MoveJoints` resolve the manipulator entirely +from `ActionBinding`. + +### 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` + +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 | +| Action config | only the inherited action name; 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 `MoveJointsCfg.named_joint_positions`. Named poses are +implementation/hardware knowledge and remain in the action config rather than +becoming separate goal types: + +```python +move_joints = MoveJoints( + MoveJointsCfg(named_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` + +Plans **approach -> close hand -> lift** and declares the object attached to the +bound manipulator. + +| 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 configured manipulator and clear overlapping coordinated attachment state | +| Verification | the attachment effect must be verified during closed-loop execution | + +`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. + +Important `PickUpCfg` fields: + +| Field | Purpose | +|---|---| +| `control_part`, `hand_control_part` | Concrete resources that the `primary` binding must match | +| `hand_open_qpos`, `hand_close_qpos` | Required hardware-specific hand states | +| `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 | + +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. + +**Example:** `scripts/tutorials/atomic_action/pickup.py` + +(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 configured 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` | + +`MoveHeldObjectCfg` holds the concrete arm/hand names, required +`hand_close_qpos`, and optional upright-transport settings. Generic timing and +trajectory sampling remain in `MotionPolicy`. + +**Example:** `scripts/tutorials/atomic_action/move_held_object.py` + +(builtin-place)= + +## `Place` + +Plans **approach/descend -> open hand -> retract**. A multi-waypoint +`PlaceGoal` visits all supplied release waypoints in order and opens at the last +one. + +| Contract | Value | +|---|---| +| Skill ID | `place` | +| Goal | `PlaceGoal(xpos=..., tcp_symmetry="none")` | +| Binding | manipulator + end effector role `primary` | +| State | consumes the configured 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` | + +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. + +Important `PlaceCfg` fields: + +| Field | Purpose | +|---|---| +| `control_part`, `hand_control_part` | Concrete resources that the `primary` binding must match | +| `hand_open_qpos`, `hand_close_qpos` | Required release/holding hand states | +| `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 | + +**Example:** `scripts/tutorials/atomic_action/place.py` + +(builtin-assemble)= + +### Assembly through `Place` + +`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. + +```text +base_object_pose @ assemble_to_base_pose = assemble_object_target_pose +assemble_object_target_pose @ held.object_to_eef = release_eef_pose +``` + +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. + +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. + +**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` | + +`PressCfg` pins the concrete arm/hand resources, required `hand_close_qpos`, and +`hand_interp_steps`. 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` + +Coordinates two arms around one shared object: **approach both grasps -> close +both hands -> lift -> move object -> hold**. + +| 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. + +Important `CoordinatedPickmentCfg` fields group into: + +- combined, left/right arm, and left/right hand control-part names; +- required open/close qpos for both hands; +- `pre_grasp_distance` and `lift_height`; +- `object_motion_keyframes`, `hand_interp_steps`, and `hold_steps`. + +The semantic binding must match those configured resources. 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. + +**Example:** `scripts/tutorials/atomic_action/coordinated_pickment.py` + +(builtin-coordinated-placement)= + +## `CoordinatedPlacement` + +Moves a support object and a placing object together: **align both objects -> +hold -> optionally release the placing hand -> retreat the placing arm**. + +| Contract | Value | +|---|---| +| Skill ID | `coordinated_placement` | +| Goal | `CoordinatedPlacementGoal` | +| Binding | manipulator + end effector roles `placing` and `support` | +| Precondition | separate `HeldObjectState` entries exist for both configured 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 | + +Both object targets may use `SceneEntityPose`, so either can participate in +dynamic-goal invalidation. Goal-level height/release values override defaults in +the action config for that invocation. + +Important `CoordinatedPlacementCfg` fields group into: + +- combined, placing/support arm, and placing/support hand control-part names; +- required placing-hand open/close and support-hand close qpos; +- default `release`, placing/support height offsets, and `lift_height`; +- `hand_interp_steps`, `hold_steps`, and `retreat_steps`. + +The semantic binding must match those configured resources. The same cuRobo +restriction as coordinated pickment applies to dual-arm +`motion_source="motion_gen"` planning. + +**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 | + +`HandOverCfg` currently owns the concrete source/destination arm and hand names, +all four open/close hand qpos values, destination grasp region and approach +direction, middle/final object poses, and phase distances/counts. The semantic +binding must match these configured resources. + +The middle and final poses are currently fixed configuration tensors rather +than `SceneEntityPose` goal fields. Consequently, handover supports tracking- +error and timeout recovery, but does not yet provide automatic moving-handover- +point invalidation. The action queries the semantic object's live orientation +when replanning and preserves that orientation at the configured 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 286023b90..476bbba55 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -1,3 +1,5 @@ +(atomic-actions)= + # Atomic actions ```{toctree} @@ -6,69 +8,310 @@ builtin_actions ``` -Atomic actions turn typed, grounded skill requests into full-robot timed motion. -Planning is side-effect free and execution is incremental. +```{currentmodule} embodichain.lab.sim.atomic_actions +``` + +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} +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. +``` + +## Architecture and responsibility boundary ```text Action Agent / task graph | - | semantic skill call + | semantic skill call (object, destination, constraints) v grounder + capability binder | - | ActionInvocation + | ActionInvocation + PlanningContext v -AtomicAction.plan(invocation, PlanningContext) - | - | ActionPlan + StateDelta - +------------------------------+ - | | - v v -AtomicActionEngine.compile ExecutionSession.tick -(fixed-scene/offline) (dynamic/closed-loop) ++-------------------------------------------------------------+ +| AtomicActionEngine | +| | +| owns exactly one ActionPlanningServices | +| +-- Robot | +| +-- MotionGenerator / planner backend | +| +-- device and shared TrajectoryBuilder | +| | +| registered AtomicAction.plan(...) -> ActionPlan | ++--------------------------+----------------------------------+ + | + +------------+-------------+ + | | + v v + compile(...) start(...) / tick(...) + fixed projection observed closed loop + | | + v v + CompiledTrajectory JointCommand + events +``` + +The boundary is deliberate: + +| Concern | Owner | Contract | +|---|---|---| +| Task intent and sequencing | Action Agent or task graph | Selects a skill and semantic goal | +| Perception and grounding | Application adapter | Builds scene snapshots, object semantics, and resource bindings | +| 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, and trajectory builder | +| 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 | + +## Core contracts + +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. + +| Contract | Contains | Does not contain | +|---|---|---| +| `ActionGoal` | Action-specific desired outcome, such as an EEF pose or object pose | Arm names, planner instances, recovery counters | +| `ActionBinding` | Semantic-role mappings such as `primary -> left_arm` and `primary -> left_hand` | Motion settings or task geometry | +| `ActionCfg` | Implementation and hardware constants: hand qpos, grasp-selection rules, phase structure | Per-call goal, motion generator, generic recovery settings | +| `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 + +Bindings make an invocation portable across embodiments: + +```python +binding = ActionBinding( + manipulators={"primary": "left_arm"}, + end_effectors={"primary": "left_hand"}, +) +``` + +Single-resource skills use `primary`; handover uses `source` and `destination`; +coordinated pick uses `left` and `right`; coordinated placement uses `placing` +and `support`. The action descriptor declares which roles are required, so a +grounder can validate a call before planning. + +Simple motion skills resolve their concrete control part entirely from the +binding. Some current multi-phase manipulation implementations also keep +concrete arm/hand names in their action config because their preconfigured hand +qpos and phase assembly are hardware-specific. For those actions, the binding +must match the configured resources. This is an implementation constraint, not +a reason to put resource names back into goals. + +### Engine-owned planning resources + +One engine owns one motion generator. Actions borrow its planning services only +after `register()` or `plan_action()` binds them: + +```python +engine = AtomicActionEngine(motion_generator) +engine.register(MoveEndEffector(MoveEndEffectorCfg())) +engine.register(MoveJoints(MoveJointsCfg())) ``` -## Contracts +Consequences of this ownership model: -- `ActionGoal`: structural protocol implemented by action-owned frozen goal - dataclasses. There is no common target object or closed union. -- `ActionBinding`: semantic role to robot-resource mapping. Goals do not carry - arm or hand names. -- `MotionPolicy`: planner, interpolation, sample count, timing, and limits. -- `RecoveryPolicy`: replan/retry budgets, tracking and dynamic-goal thresholds, - and phase timeout. -- `PlanningContext`: measured `RobotObservation`, verified `TaskState`, versioned - `SceneSnapshot`, and stable environment IDs. -- `ActionPlan`: one or more scene-bound phases, timed trajectories, completion - conditions, diagnostics, and uncommitted `StateDelta` effects. +- action constructors contain only skill configuration; +- 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. -## Static and dynamic use +When two differently configured instances share the same stable skill ID, keep +one or both outside the registry and call `engine.plan_action(...)` explicitly: -`AtomicActionEngine.compile()` plans a fixed sequence and returns a -`CompiledTrajectory`. It projects terminal qpos and expected task effects only -inside the returned context; it never changes simulator state. +```python +left_pick = PickUp(left_pick_cfg) +right_pick = PickUp(right_pick_cfg) + +left_plan = engine.plan_action(left_pick, left_invocation, latest_context) +right_plan = engine.plan_action(right_pick, right_invocation, latest_context) +``` + +Both instances still borrow the same engine-owned motion generator. + +## Which planning API to use -`AtomicActionEngine.start()` creates an `ExecutionSession`. Each `tick()` takes -the latest context and emits at most one `JointCommand`. The session detects -tracking error, phase timeout, and movement of entities referenced by -`SceneEntityPose`, then replans from the latest observation within the configured -budget. Non-empty symbolic effects require external verification before commit. +| API | Use it for | Result / behavior | +|---|---|---| +| `AtomicAction.plan(invocation, context)` | Implementing a skill | Action-owned side-effect-free planning hook; application code normally calls it through the engine | +| `engine.plan(invocation, context)` | Planning one registered skill | Resolves the registered action, binds shared resources, and validates its plan | +| `engine.plan_action(action, invocation, context)` | Planning an unregistered configured instance | Supports multiple configurations with one `skill_id` and one engine backend | +| `engine.compile(invocations, context)` | Fixed-scene/offline sequence planning | Returns one concatenated `CompiledTrajectory` and a hypothetical projected context | +| `engine.start(invocations, context)` | Observed incremental execution | Returns an `ExecutionSession`; each `tick()` emits at most one command and recovery events | -## Example +`AtomicAction.plan()` is therefore not a second execution API. It is the +polymorphic implementation point used by the engine. Neither it nor the engine +mutates the simulator. + +## 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 -binding = ActionBinding(manipulators={"primary": "left_arm"}) +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + AtomicActionEngine, + EndEffectorPoseGoal, + ExecutionEventKind, + ExecutionStatus, + MotionPolicy, + MoveEndEffector, + RecoveryPolicy, + SceneEntityPose, +) + +engine = AtomicActionEngine(motion_generator) +engine.register(MoveEndEffector()) + invocation = ActionInvocation( skill_id="move_end_effector", - goal=EndEffectorPoseGoal(target_pose), - binding=binding, - motion_policy=MotionPolicy(sample_count=80), + goal=EndEffectorPoseGoal(xpos=target_pose), + binding=ActionBinding(manipulators={"primary": "left_arm"}), + motion_policy=MotionPolicy(sample_count=80, control_dt=1.0 / 60.0), ) -engine = AtomicActionEngine(motion_generator) -engine.register(MoveEndEffector(motion_generator, MoveEndEffectorCfg())) compiled = engine.compile((invocation,)) +if compiled.plan_success.all(): + positions = compiled.trajectory.positions # (B, N, robot_dof) +``` + +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. + +## Dynamic goals and closed-loop recovery + +Pose-valued goals can use `SceneEntityPose` instead of freezing an object pose +at invocation creation time: + +```python +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() ``` -See [Built-in actions](builtin_actions.md) for the shipped skill catalog and -[the tutorial](../../../tutorial/atomic_actions.rst) for closed-loop usage. +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. + +```{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. +``` + +## 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. + +At the terminal waypoint, an `ExecutionSession` requests an external +per-environment verification mask before committing a non-empty effect: + +```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) +``` + +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 + -> ActionInvocation + -> AtomicActionEngine + -> ActionPlan / execution events +``` + +This keeps learned reasoning independent of planner instances and concrete +joint groups, while invocation IDs, planner diagnostics, and execution events +provide structured feedback for the agent's next decision. + +## Extending the module + +A new primitive should: + +1. define a frozen, action-owned goal dataclass with a stable `goal_kind`; +2. declare `skill_id`, `GoalType`, required semantic roles, and agent visibility; +3. keep embodiment constants in its `ActionCfg` and reusable motion/recovery + choices in invocation policies; +4. implement side-effect-free `plan(invocation, context)` using the + engine-owned planning services; +5. return full-robot timed motion, per-environment planning success, + diagnostics, and uncommitted effects; +6. 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/tutorial/atomic_actions.rst b/docs/source/tutorial/atomic_actions.rst index c2777124e..59c20471b 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -6,6 +6,11 @@ grounded :class:`~embodichain.lab.sim.atomic_actions.ActionInvocation` and the latest :class:`~embodichain.lab.sim.atomic_actions.PlanningContext`, then 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 four concerns: * a **goal** describes what should happen; @@ -16,6 +21,43 @@ The contracts deliberately separate four concerns: * a **PlanningContext** contains measured robot state, verified task state, and a versioned scene snapshot. +The engine exclusively owns the ``MotionGenerator`` and a shared trajectory +builder. Atomic action constructors accept only implementation configuration; +``register()`` binds each action to the engine resources. Use +``engine.plan_action(action, invocation, context)`` for an unregistered, +configuration-specific action 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/pickup.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 ``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. + Static compilation ------------------ @@ -35,7 +77,7 @@ the scene is treated as fixed during planning: ) engine = AtomicActionEngine(motion_generator) - engine.register(MoveEndEffector(motion_generator, MoveEndEffectorCfg())) + engine.register(MoveEndEffector(MoveEndEffectorCfg())) invocation = ActionInvocation( skill_id="move_end_effector", @@ -78,17 +120,24 @@ must be resolved from the latest scene snapshot: ), ) - session = engine.start((invocation,), initial_context) + 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() 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. +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. + Task-state effects ------------------ @@ -128,6 +177,9 @@ implement ``plan(invocation, context)`` and declare the stable skill metadata: GoalType: ClassVar[type] = PushGoal manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + def __init__(self, cfg: PushCfg | None = None) -> None: + super().__init__(cfg or PushCfg()) + def plan( self, invocation: ActionInvocation[PushGoal], diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index e196d8664..103e64c09 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -61,6 +61,7 @@ TimedTrajectory, ) from .policies import MotionPolicy, RecoveryPolicy +from .runtime import ActionPlanningServices from .primitives import ( AssembleGoal, CoordinatedPickGoal, @@ -81,7 +82,6 @@ MoveHeldObjectCfg, MoveJoints, MoveJointsCfg, - NamedJointPositionGoal, PickUp, PickUpCfg, Place, @@ -108,6 +108,7 @@ "ActionGoal", "ActionInvocation", "ActionPlan", + "ActionPlanningServices", "Affordance", "AntipodalAffordance", "AssembleAffordance", @@ -146,7 +147,6 @@ "MoveHeldObjectCfg", "MoveJoints", "MoveJointsCfg", - "NamedJointPositionGoal", "ObjectActionGoal", "ObjectSemantics", "PhaseSpec", diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index 6f49f52fa..70579f650 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -42,9 +42,12 @@ ) 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: @@ -132,7 +135,12 @@ def __post_init__(self) -> None: class AtomicAction(Generic[GoalT], ABC): - """Side-effect-free planner for one semantically meaningful robot skill.""" + """Side-effect-free planner for one semantically meaningful robot skill. + + Actions own only skill configuration. An + :class:`~embodichain.lab.sim.atomic_actions.engine.AtomicActionEngine` binds + its shared planning services before an action is invoked. + """ skill_id: ClassVar[str] """Stable registry identifier for this skill.""" @@ -151,13 +159,70 @@ class AtomicAction(Generic[GoalT], ABC): 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 = resolve_runtime_device(self.robot.device) + self._planning_services: ActionPlanningServices | None = None + + @property + def is_bound(self) -> bool: + """Whether an engine has supplied this action's planning resources.""" + return self._planning_services is not None + + @property + def planning_services(self) -> ActionPlanningServices: + """Return the engine-owned services borrowed by this action. + + 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 + + @property + def motion_generator(self) -> MotionGenerator: + """Return the engine-owned motion generator borrowed by this action.""" + return self.planning_services.motion_generator + + @property + 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 + + @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: @@ -204,10 +269,7 @@ def require_goal(self, invocation: ActionInvocation[GoalT]) -> GoalT: for role in self.end_effector_roles: invocation.binding.end_effector(role) required_planner = invocation.motion_policy.planner - configured_planner = getattr( - getattr(self.motion_generator, "planner", None), "cfg", None - ) - configured_planner_name = getattr(configured_planner, "planner_type", None) + configured_planner_name = self.planning_services.planner_name if required_planner is not None and required_planner != configured_planner_name: raise ValueError( f"Motion policy requires planner {required_planner!r}, but this " @@ -287,12 +349,9 @@ def build_plan( raise ValueError("Trajectory robot_dof must match the planning context.") if diagnostics is None: - backend = getattr( - getattr(getattr(self.motion_generator, "planner", None), "cfg", None), - "planner_type", - invocation.motion_policy.motion_source, + diagnostics = PlannerDiagnostics( + backend=self.planning_services.planner_name ) - diagnostics = PlannerDiagnostics(backend=str(backend)) phase = PlannedPhase( spec=PhaseSpec( name=phase_name or self.cfg.name, @@ -334,11 +393,6 @@ def failed_plan( Returns: Failed action plan with an empty phase trajectory. """ - backend = getattr( - getattr(getattr(self.motion_generator, "planner", None), "cfg", None), - "planner_type", - invocation.motion_policy.motion_source, - ) return self.build_plan( invocation, context, @@ -353,7 +407,8 @@ def failed_plan( ), replannable=True, diagnostics=PlannerDiagnostics( - backend=str(backend), messages=(() if message is None else (message,)) + backend=self.planning_services.planner_name, + messages=(() if message is None else (message,)), ), ) diff --git a/embodichain/lab/sim/atomic_actions/engine.py b/embodichain/lab/sim/atomic_actions/engine.py index 369870604..a97d9a258 100644 --- a/embodichain/lab/sim/atomic_actions/engine.py +++ b/embodichain/lab/sim/atomic_actions/engine.py @@ -22,12 +22,14 @@ import torch -from .core import AtomicAction, resolve_runtime_device +from .core import AtomicAction from .invocation import ActionInvocation 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 @@ -73,14 +75,32 @@ def get_registered_actions() -> dict[str, type[AtomicAction]]: class AtomicActionEngine: - """Compile grounded atomic invocations without stepping the environment.""" + """Own planning resources and coordinate side-effect-free atomic actions.""" def __init__(self, motion_generator: MotionGenerator) -> None: - self.motion_generator = motion_generator - self.robot = motion_generator.robot - self.device = resolve_runtime_device(motion_generator.device) + self._planning_services = ActionPlanningServices(motion_generator) self._actions: dict[str, AtomicAction] = {} + @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 actions(self) -> dict[str, AtomicAction]: """Registered action instances keyed by stable skill identifier.""" @@ -94,20 +114,78 @@ def register(self, action: AtomicAction) -> None: Raises: TypeError: If ``action`` is not an AtomicAction. - ValueError: If its robot or skill identifier is incompatible. + 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.") - if action.robot is not self.robot: - raise ValueError("Registered actions must use the engine robot.") descriptor = action.descriptor() existing = self._actions.get(descriptor.skill_id) if existing is not None and existing is not action: raise ValueError( f"Skill id {descriptor.skill_id!r} is already registered in this engine." ) + action._bind(self._planning_services) self._actions[descriptor.skill_id] = action + def plan_action( + self, + 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 supports multiple configured instances with the + same stable skill identifier while preserving one engine-owned motion + generator. + + Args: + action: Configured action implementation to invoke. + invocation: Grounded request matching the action's skill identifier. + context: Latest measured planning state. + + Returns: + 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) + plan = action.plan(invocation, context) + self._validate_plan(plan, context, invocation) + 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. + """ + action = self._actions.get(invocation.skill_id) + if action is None: + raise KeyError( + f"No atomic action registered for skill {invocation.skill_id!r}." + ) + current = self.initial_context() if context is None else context + return self.plan_action(action, invocation, current) + def initial_context( self, *, @@ -178,16 +256,10 @@ def compile( projected = context for invocation in invocations: - action = self._actions.get(invocation.skill_id) - if action is None: - raise KeyError( - f"No atomic action registered for skill {invocation.skill_id!r}." - ) if not alive.any(): break previous_qpos = projected.robot.qpos - plan = action.plan(invocation, projected) - self._validate_plan(plan, projected, invocation) + 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) diff --git a/embodichain/lab/sim/atomic_actions/execution.py b/embodichain/lab/sim/atomic_actions/execution.py index 4def86424..e0ee840e8 100644 --- a/embodichain/lab/sim/atomic_actions/execution.py +++ b/embodichain/lab/sim/atomic_actions/execution.py @@ -304,13 +304,7 @@ def _plan_current( ) -> None: """Plan the current invocation from the latest observation.""" invocation = self._invocations[self._invocation_index] - action = self._engine.actions.get(invocation.skill_id) - if action is None: - raise KeyError( - f"No atomic action registered for skill {invocation.skill_id!r}." - ) - plan = action.plan(invocation, context) - self._engine._validate_plan(plan, context, invocation) + plan = self._engine.plan(invocation, context) self._plan = plan self._phase_index = min(self._phase_index, len(plan.phases) - 1) self._waypoint_index = 0 diff --git a/embodichain/lab/sim/atomic_actions/primitives/__init__.py b/embodichain/lab/sim/atomic_actions/primitives/__init__.py index 21c109442..e34a7b20a 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/__init__.py +++ b/embodichain/lab/sim/atomic_actions/primitives/__init__.py @@ -43,7 +43,6 @@ JointPositionGoal, MoveJoints, MoveJointsCfg, - NamedJointPositionGoal, ) from .pick_up import GraspGoal, PickUp, PickUpCfg from .place import AssembleGoal, Place, PlaceCfg, PlaceGoal @@ -69,7 +68,6 @@ "MoveHeldObjectCfg", "MoveJoints", "MoveJointsCfg", - "NamedJointPositionGoal", "PickUp", "PickUpCfg", "Place", diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py index 73a744b6c..dc7c76a22 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py @@ -41,7 +41,6 @@ from ..invocation import ActionInvocation from ..plans import ActionPlan from ..state import CoordinatedHeldObjectState, PlanningContext -from ..trajectory import TrajectoryBuilder @dataclass(frozen=True, slots=True, eq=False) @@ -146,7 +145,6 @@ def _init_dual_arm_parts( 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) @@ -445,10 +443,13 @@ class CoordinatedPickment(AtomicAction[CoordinatedPickGoal]): def __init__( self, - motion_generator, cfg: CoordinatedPickmentCfg | None = None, ) -> None: - super().__init__(motion_generator, cfg or CoordinatedPickmentCfg()) + super().__init__(cfg or CoordinatedPickmentCfg()) + self._validate_hand_qpos_cfg() + + def _on_bind(self) -> None: + """Resolve robot-dependent resources from the owning engine.""" 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, @@ -464,7 +465,10 @@ def __init__( self.left_hand_dof = self.first_hand_dof self.right_hand_dof = self.second_hand_dof - self._validate_hand_qpos_cfg() + assert self.cfg.left_hand_open_qpos is not None + assert self.cfg.left_hand_close_qpos is not None + assert self.cfg.right_hand_open_qpos is not None + assert self.cfg.right_hand_close_qpos is not None self.left_hand_open_qpos = self._expand_qpos( self.cfg.left_hand_open_qpos, self.left_hand_dof, "left_hand_open_qpos" ) diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py index d743e4675..768c0bf69 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py @@ -37,7 +37,6 @@ from ..plans import ActionPlan from ..policies import MotionPolicy from ..state import HeldObjectState, PlanningContext -from ..trajectory import TrajectoryBuilder @dataclass(frozen=True, slots=True, eq=False) @@ -135,11 +134,13 @@ class CoordinatedPlacement(AtomicAction[CoordinatedPlacementGoal]): def __init__( self, - motion_generator, cfg: CoordinatedPlacementCfg | None = None, ) -> None: - super().__init__(motion_generator, cfg or CoordinatedPlacementCfg()) - self.builder = TrajectoryBuilder(motion_generator) + super().__init__(cfg or CoordinatedPlacementCfg()) + self._validate_hand_qpos_cfg() + + def _on_bind(self) -> None: + """Resolve robot-dependent resources from the owning engine.""" self.n_envs = self.robot.get_qpos().shape[0] self.robot_dof = self.robot.dof @@ -166,7 +167,9 @@ def __init__( 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() + assert self.cfg.placing_hand_open_qpos is not None + assert self.cfg.placing_hand_close_qpos is not None + assert self.cfg.support_hand_close_qpos is not None self.placing_hand_open_qpos = self.builder.expand_hand_qpos( self.cfg.placing_hand_open_qpos, n_envs=self.n_envs, diff --git a/embodichain/lab/sim/atomic_actions/primitives/hand_over.py b/embodichain/lab/sim/atomic_actions/primitives/hand_over.py index 8c251ff0f..aaca0495d 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/hand_over.py +++ b/embodichain/lab/sim/atomic_actions/primitives/hand_over.py @@ -37,7 +37,6 @@ from ..policies import MotionPolicy from ..state import HeldObjectState, PlanningContext from .pick_up import GraspGoal -from ..trajectory import TrajectoryBuilder @configclass @@ -125,15 +124,18 @@ class HandOver(AtomicAction[GraspGoal]): def __init__( self, - motion_generator, cfg: HandOverCfg | None = None, ) -> None: - super().__init__(motion_generator, cfg or HandOverCfg()) - self.builder = TrajectoryBuilder(motion_generator) + super().__init__(cfg or HandOverCfg()) + self._validate_pose_cfg() + self._validate_hand_qpos_cfg() + + def _on_bind(self) -> None: + """Resolve robot-dependent resources from the owning engine.""" self.n_envs = self.robot.get_qpos().shape[0] self.robot_dof = self.robot.dof - - self._validate_pose_cfg() + assert self.cfg.middle_object_pose is not None + assert self.cfg.final_object_pose is not None self.middle_object_pose = self._resolve_matrix( self.cfg.middle_object_pose, "middle_object_pose" ) @@ -159,7 +161,10 @@ def __init__( 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() + assert self.cfg.transfer_hand_open_qpos is not None + assert self.cfg.transfer_hand_close_qpos is not None + assert self.cfg.receive_hand_open_qpos is not None + assert self.cfg.receive_hand_close_qpos is not None self.transfer_hand_open_qpos = self.builder.expand_hand_qpos( self.cfg.transfer_hand_open_qpos, n_envs=self.n_envs, 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 29c38db26..d35f1ee09 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py @@ -31,7 +31,6 @@ from ..invocation import ActionInvocation from ..plans import ActionPlan, CompletionConditionKind from ..state import PlanningContext -from ..trajectory import TrajectoryBuilder @dataclass(frozen=True, slots=True, eq=False) @@ -63,11 +62,9 @@ class MoveEndEffector(AtomicAction[EndEffectorPoseGoal]): def __init__( self, - motion_generator, cfg: MoveEndEffectorCfg | None = None, ) -> None: - super().__init__(motion_generator, cfg or MoveEndEffectorCfg()) - self.builder = TrajectoryBuilder(motion_generator) + super().__init__(cfg or MoveEndEffectorCfg()) def plan( self, 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 0705e6840..c0a94f2f6 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py @@ -36,7 +36,6 @@ from ..invocation import ActionInvocation from ..plans import ActionPlan from ..state import PlanningContext -from ..trajectory import TrajectoryBuilder @dataclass(frozen=True, slots=True, eq=False) @@ -87,21 +86,22 @@ class MoveHeldObject(AtomicAction[HeldObjectPoseGoal]): def __init__( self, - motion_generator, cfg: MoveHeldObjectCfg | None = None, ) -> None: - super().__init__(motion_generator, cfg or MoveHeldObjectCfg()) - self.builder = TrajectoryBuilder(motion_generator) + super().__init__(cfg or MoveHeldObjectCfg()) + if self.cfg.hand_close_qpos is None: + logger.log_error( + "hand_close_qpos must be specified in MoveHeldObjectCfg", ValueError + ) + + def _on_bind(self) -> None: + """Resolve robot-dependent 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 - ) + assert self.cfg.hand_close_qpos is not None self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) def plan( diff --git a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py index 3fc9fcc83..577888617 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py @@ -29,79 +29,63 @@ from ..invocation import ActionInvocation from ..plans import ActionPlan, CompletionConditionKind from ..state import PlanningContext -from ..trajectory import TrajectoryBuilder @dataclass(frozen=True, slots=True, eq=False) class JointPositionGoal: - """Joint-space goal for a bound robot control resource.""" + """Explicit or named joint-space goal for a bound robot resource.""" goal_kind: ClassVar[str] = "joint_position" - qpos: torch.Tensor - """One joint waypoint or a batched sequence of joint waypoints.""" + target: torch.Tensor | str + """Joint qpos/waypoints or a name in ``MoveJointsCfg.named_joint_positions``.""" 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 NamedJointPositionGoal: - """Named joint-space goal resolved from :class:`MoveJointsCfg`.""" - - goal_kind: ClassVar[str] = "named_joint_position" - - name: str - """Name 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): """Skill-specific MoveJoints configuration.""" name: str = "move_joints" named_joint_positions: dict[str, torch.Tensor] | None = None - """Optional named goals. Motion settings belong to ``MotionPolicy``.""" + """Optional named joint-position targets. Motion settings belong to ``MotionPolicy``.""" -class MoveJoints(AtomicAction[JointPositionGoal | NamedJointPositionGoal]): +class MoveJoints(AtomicAction[JointPositionGoal]): """Plan joint motion from the observed state to one or more waypoints.""" skill_id: ClassVar[str] = "move_joints" - GoalType: ClassVar[tuple[type, ...]] = ( - JointPositionGoal, - NamedJointPositionGoal, - ) + GoalType: ClassVar[type] = JointPositionGoal manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) agent_visible: ClassVar[bool] = False def __init__( self, - motion_generator, cfg: MoveJointsCfg | None = None, ) -> None: - super().__init__(motion_generator, cfg or MoveJointsCfg()) - self.builder = TrajectoryBuilder(motion_generator) + super().__init__(cfg or MoveJointsCfg()) self.named_joint_positions = self.cfg.named_joint_positions or {} def plan( self, - invocation: ActionInvocation[JointPositionGoal | NamedJointPositionGoal], + invocation: ActionInvocation[JointPositionGoal], context: PlanningContext, ) -> ActionPlan: """Plan a joint-space goal without mutating the robot or task state.""" @@ -146,23 +130,22 @@ def plan( def _resolve_target_qpos( self, - goal: JointPositionGoal | NamedJointPositionGoal, + goal: JointPositionGoal, ) -> torch.Tensor: """Resolve an explicit or named joint goal to a tensor.""" - if isinstance(goal, JointPositionGoal): - return goal.qpos - if goal.name not in self.named_joint_positions: + if isinstance(goal.target, torch.Tensor): + return goal.target + if goal.target not in self.named_joint_positions: logger.log_error( - f"Unknown named joint-position goal {goal.name!r}. Available " + f"Unknown named joint-position goal {goal.target!r}. Available " f"goals: {sorted(self.named_joint_positions)}", KeyError, ) - return self.named_joint_positions[goal.name] + return self.named_joint_positions[goal.target] __all__ = [ "JointPositionGoal", "MoveJoints", "MoveJointsCfg", - "NamedJointPositionGoal", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index 4a812802d..8d145b51d 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -46,7 +46,6 @@ from ..plans import ActionPlan from ..policies import MotionPolicy from ..state import HeldObjectState, PlanningContext -from ..trajectory import TrajectoryBuilder @dataclass(frozen=True, slots=True, eq=False) @@ -124,17 +123,9 @@ class PickUp(AtomicAction[GraspGoal]): def __init__( self, - motion_generator, cfg: PickUpCfg | None = None, ) -> None: - super().__init__(motion_generator, cfg or PickUpCfg()) - self.builder = TrajectoryBuilder(motion_generator) - self.n_envs = self.robot.get_qpos().shape[0] - self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part) - self.arm_dof = len(self.arm_joint_ids) - self.robot_dof = self.robot.dof - + super().__init__(cfg or PickUpCfg()) if self.cfg.hand_open_qpos is None: logger.log_error( "hand_open_qpos must be specified in PickUpCfg", ValueError @@ -143,6 +134,23 @@ def __init__( logger.log_error( "hand_close_qpos must be specified in PickUpCfg", ValueError ) + 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 _on_bind(self) -> None: + """Resolve robot-dependent 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 + assert self.cfg.hand_open_qpos is not None + assert self.cfg.hand_close_qpos is not None 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( @@ -152,13 +160,6 @@ def __init__( 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, diff --git a/embodichain/lab/sim/atomic_actions/primitives/place.py b/embodichain/lab/sim/atomic_actions/primitives/place.py index 80b4c47c7..12fdf6c57 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/place.py +++ b/embodichain/lab/sim/atomic_actions/primitives/place.py @@ -38,7 +38,6 @@ from ..invocation import ActionInvocation from ..plans import ActionPlan from ..state import PlanningContext -from ..trajectory import TrajectoryBuilder TcpSymmetry = Literal["none", "z_roll_180"] @@ -147,28 +146,30 @@ class Place(AtomicAction[PlaceGoal | AssembleGoal]): def __init__( self, - motion_generator, cfg: PlaceCfg | None = None, ) -> None: - super().__init__(motion_generator, cfg or PlaceCfg()) - self.builder = TrajectoryBuilder(motion_generator) - self.n_envs = self.robot.get_qpos().shape[0] - self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) - self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part) - self.arm_dof = len(self.arm_joint_ids) - self.robot_dof = self.robot.dof - + super().__init__(cfg or PlaceCfg()) 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 _on_bind(self) -> None: + """Resolve robot-dependent 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 + assert self.cfg.hand_open_qpos is not None + assert self.cfg.hand_close_qpos is not None + self.hand_open_qpos = self.cfg.hand_open_qpos.to(self.device) + self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) + def plan( self, invocation: ActionInvocation[PlaceGoal | AssembleGoal], diff --git a/embodichain/lab/sim/atomic_actions/primitives/press.py b/embodichain/lab/sim/atomic_actions/primitives/press.py index ddd0ea9c7..d9a6c73d2 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/press.py +++ b/embodichain/lab/sim/atomic_actions/primitives/press.py @@ -35,7 +35,6 @@ from ..invocation import ActionInvocation from ..plans import ActionPlan from ..state import PlanningContext -from ..trajectory import TrajectoryBuilder @dataclass(frozen=True, slots=True, eq=False) @@ -79,11 +78,16 @@ class Press(AtomicAction[PressGoal]): def __init__( self, - motion_generator, cfg: PressCfg | None = None, ) -> None: - super().__init__(motion_generator, cfg or PressCfg()) - self.builder = TrajectoryBuilder(motion_generator) + super().__init__(cfg or PressCfg()) + if self.cfg.hand_close_qpos is None: + logger.log_error( + "hand_close_qpos must be specified in PressCfg", ValueError + ) + + def _on_bind(self) -> None: + """Resolve robot-dependent 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) @@ -91,10 +95,7 @@ def __init__( 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 - ) + assert self.cfg.hand_close_qpos is not None self.hand_close_qpos = self.builder.expand_hand_qpos( self.cfg.hand_close_qpos, n_envs=self.n_envs, diff --git a/embodichain/lab/sim/atomic_actions/runtime.py b/embodichain/lab/sim/atomic_actions/runtime.py new file mode 100644 index 000000000..426613ff2 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/runtime.py @@ -0,0 +1,81 @@ +# ---------------------------------------------------------------------------- +# 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 typing import TYPE_CHECKING + +import torch + +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. + """ + + def __init__(self, motion_generator: MotionGenerator) -> 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) + + @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 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) + + +__all__ = ["ActionPlanningServices"] diff --git a/embodichain/lab/sim/atomic_actions/trajectory.py b/embodichain/lab/sim/atomic_actions/trajectory.py index 17165008c..7bfd4644b 100644 --- a/embodichain/lab/sim/atomic_actions/trajectory.py +++ b/embodichain/lab/sim/atomic_actions/trajectory.py @@ -43,9 +43,11 @@ 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: diff --git a/examples/sim/planners/curobo_planner.py b/examples/sim/planners/curobo_planner.py index cd19f86cc..64f16e692 100644 --- a/examples/sim/planners/curobo_planner.py +++ b/examples/sim/planners/curobo_planner.py @@ -750,7 +750,7 @@ def main() -> None: ) ) engine = AtomicActionEngine(motion_generator) - engine.register(MoveEndEffector(motion_generator, MoveEndEffectorCfg())) + engine.register(MoveEndEffector(MoveEndEffectorCfg())) binding = ActionBinding(manipulators={"primary": control_part}) motion_policy = MotionPolicy( motion_source="motion_gen", diff --git a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py index d1381653d..6b079b572 100644 --- a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py +++ b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py @@ -264,7 +264,7 @@ 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())) + atomic_engine.register(MoveEndEffector(cfg=MoveEndEffectorCfg())) 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 9ee113330..797975e0b 100644 --- a/scripts/benchmark/atomic_action/move_held_object_benchmark.py +++ b/scripts/benchmark/atomic_action/move_held_object_benchmark.py @@ -195,10 +195,9 @@ 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())) + atomic_engine.register(MoveEndEffector(cfg=MoveEndEffectorCfg())) atomic_engine.register( PickUp( - motion_gen, cfg=PickUpCfg( control_part="arm", hand_control_part="hand", @@ -318,7 +317,6 @@ def _run_case( atomic_engine = AtomicActionEngine(motion_generator=motion_gen) atomic_engine.register( MoveHeldObject( - motion_gen, cfg=MoveHeldObjectCfg( control_part="arm", hand_control_part="hand", diff --git a/scripts/benchmark/atomic_action/move_joints_benchmark.py b/scripts/benchmark/atomic_action/move_joints_benchmark.py index e0f4b9f62..8ee335184 100644 --- a/scripts/benchmark/atomic_action/move_joints_benchmark.py +++ b/scripts/benchmark/atomic_action/move_joints_benchmark.py @@ -113,7 +113,6 @@ def _targets_for_sequence(sequence_case: JointSequenceCase, device): ActionInvocation, JointPositionGoal, MotionPolicy, - NamedJointPositionGoal, ) targets = [] @@ -121,9 +120,9 @@ def _targets_for_sequence(sequence_case: JointSequenceCase, device): policy = MotionPolicy(sample_count=MOVE_JOINTS_SAMPLE_INTERVAL) for index, name in enumerate(sequence_case.sequence): if index == 0 and name == "ready": - goal = NamedJointPositionGoal(name="ready") + goal = JointPositionGoal(target="ready") else: - goal = JointPositionGoal(qpos=_qpos(JOINT_TARGETS[name], device)) + goal = JointPositionGoal(target=_qpos(JOINT_TARGETS[name], device)) targets.append( ActionInvocation( skill_id="move_joints", @@ -275,7 +274,6 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: atomic_engine = AtomicActionEngine(motion_generator=motion_gen) atomic_engine.register( MoveJoints( - motion_gen, cfg=MoveJointsCfg( named_joint_positions={"ready": ready_qpos}, ), diff --git a/scripts/benchmark/atomic_action/pickup_benchmark.py b/scripts/benchmark/atomic_action/pickup_benchmark.py index 6bc6acabf..3bb3ec7cb 100644 --- a/scripts/benchmark/atomic_action/pickup_benchmark.py +++ b/scripts/benchmark/atomic_action/pickup_benchmark.py @@ -162,7 +162,6 @@ def _run_case( atomic_engine = AtomicActionEngine(motion_generator=motion_gen) atomic_engine.register( PickUp( - motion_gen, cfg=PickUpCfg( control_part="arm", hand_control_part="hand", diff --git a/scripts/benchmark/atomic_action/place_benchmark.py b/scripts/benchmark/atomic_action/place_benchmark.py index f35df7a46..31f93c6c2 100644 --- a/scripts/benchmark/atomic_action/place_benchmark.py +++ b/scripts/benchmark/atomic_action/place_benchmark.py @@ -194,7 +194,6 @@ def _prepare_held_state( atomic_engine = AtomicActionEngine(motion_generator=motion_gen) atomic_engine.register( PickUp( - motion_gen, cfg=PickUpCfg( control_part="arm", hand_control_part="hand", @@ -302,7 +301,6 @@ def _run_case( atomic_engine = AtomicActionEngine(motion_generator=motion_gen) atomic_engine.register( Place( - motion_gen, cfg=PlaceCfg( control_part="arm", hand_control_part="hand", diff --git a/scripts/benchmark/atomic_action/press_benchmark.py b/scripts/benchmark/atomic_action/press_benchmark.py index d96dd7f37..3cd5a2c90 100644 --- a/scripts/benchmark/atomic_action/press_benchmark.py +++ b/scripts/benchmark/atomic_action/press_benchmark.py @@ -487,10 +487,9 @@ def _build_atomic_engine( """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())) + atomic_engine.register(MoveEndEffector(cfg=MoveEndEffectorCfg())) atomic_engine.register( Press( - motion_gen, cfg=PressCfg( control_part="arm", hand_control_part="hand", diff --git a/scripts/tutorials/atomic_action/assemble.py b/scripts/tutorials/atomic_action/assemble.py index 016cecb6e..c05bd2ded 100644 --- a/scripts/tutorials/atomic_action/assemble.py +++ b/scripts/tutorials/atomic_action/assemble.py @@ -409,7 +409,6 @@ 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", @@ -428,7 +427,6 @@ def run_assemble_demo( ) # Step 2 - the left arm places the can directly above the cube. place_action = Place( - motion_generator=motion_gen, cfg=PlaceCfg( name="place", control_part="left_arm", diff --git a/scripts/tutorials/atomic_action/coordinated_pickment.py b/scripts/tutorials/atomic_action/coordinated_pickment.py index ae72f85b4..093128544 100644 --- a/scripts/tutorials/atomic_action/coordinated_pickment.py +++ b/scripts/tutorials/atomic_action/coordinated_pickment.py @@ -669,7 +669,6 @@ def run_coordinated_pickment_demo( robot, "right_hand", sim.device, preset.hand_close_qpos ) pickment_action = CoordinatedPickment( - motion_generator=motion_gen, cfg=CoordinatedPickmentCfg( control_part="dual_arm", left_arm_control_part="left_arm", diff --git a/scripts/tutorials/atomic_action/coordinated_placement.py b/scripts/tutorials/atomic_action/coordinated_placement.py index 487a14dd9..893d7b73d 100644 --- a/scripts/tutorials/atomic_action/coordinated_placement.py +++ b/scripts/tutorials/atomic_action/coordinated_placement.py @@ -796,7 +796,6 @@ def run_coordinated_placement_demo( ) 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", @@ -808,7 +807,6 @@ def run_coordinated_placement_demo( ), ) right_pick_action = PickUp( - motion_generator=motion_gen, cfg=PickUpCfg( control_part="right_arm", hand_control_part="right_hand", @@ -820,7 +818,6 @@ def run_coordinated_placement_demo( ), ) coordinated_action = CoordinatedPlacement( - motion_generator=motion_gen, cfg=CoordinatedPlacementCfg( control_part="dual_arm", placing_arm_control_part="left_arm", @@ -857,7 +854,8 @@ def run_coordinated_placement_demo( z_clearance=BREAD_GRASP_Z_CLEARANCE, ) start_time = time.time() - left_pick_result = left_pick_action.plan( + left_pick_result = engine.plan_action( + left_pick_action, ActionInvocation( skill_id="pick_up", goal=GraspGoal( @@ -897,7 +895,8 @@ def run_coordinated_placement_demo( z_clearance=PAN_GRASP_Z_CLEARANCE, ) start_time = time.time() - right_pick_result = right_pick_action.plan( + right_pick_result = engine.plan_action( + right_pick_action, ActionInvocation( skill_id="pick_up", goal=GraspGoal( diff --git a/scripts/tutorials/atomic_action/hand_over.py b/scripts/tutorials/atomic_action/hand_over.py index bd34ac280..a5366864a 100644 --- a/scripts/tutorials/atomic_action/hand_over.py +++ b/scripts/tutorials/atomic_action/hand_over.py @@ -346,7 +346,6 @@ 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", @@ -364,7 +363,6 @@ def run_handover_demo( ) # Step 2 - hand the object from the left arm to the right arm. handover_action = HandOver( - motion_generator=motion_gen, cfg=HandOverCfg( name="hand_over", control_part="dual_arm", diff --git a/scripts/tutorials/atomic_action/move_end_effector.py b/scripts/tutorials/atomic_action/move_end_effector.py index 9b937b1d3..8e2fa6446 100644 --- a/scripts/tutorials/atomic_action/move_end_effector.py +++ b/scripts/tutorials/atomic_action/move_end_effector.py @@ -75,7 +75,7 @@ def main() -> None: motion_gen = create_toppra_motion_generator(robot) engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register(MoveEndEffector(motion_gen, cfg=MoveEndEffectorCfg())) + engine.register(MoveEndEffector(cfg=MoveEndEffectorCfg())) poses = torch.stack( [ diff --git a/scripts/tutorials/atomic_action/move_held_object.py b/scripts/tutorials/atomic_action/move_held_object.py index b0b7d67d8..81e6975da 100644 --- a/scripts/tutorials/atomic_action/move_held_object.py +++ b/scripts/tutorials/atomic_action/move_held_object.py @@ -130,10 +130,9 @@ def main() -> None: hand_open, hand_close = get_hand_open_close_qpos(robot) engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register(MoveEndEffector(motion_gen, MoveEndEffectorCfg())) + engine.register(MoveEndEffector(MoveEndEffectorCfg())) engine.register( PickUp( - motion_gen, PickUpCfg( hand_open_qpos=hand_open, hand_close_qpos=hand_close, @@ -145,7 +144,6 @@ def main() -> None: ) engine.register( MoveHeldObject( - motion_gen, MoveHeldObjectCfg( hand_close_qpos=hand_close, ), diff --git a/scripts/tutorials/atomic_action/move_joints.py b/scripts/tutorials/atomic_action/move_joints.py index 915e4754b..3ee97b9ad 100644 --- a/scripts/tutorials/atomic_action/move_joints.py +++ b/scripts/tutorials/atomic_action/move_joints.py @@ -36,7 +36,6 @@ JointPositionGoal, MoveJoints, MoveJointsCfg, - NamedJointPositionGoal, MotionPolicy, ) from embodichain.utils import logger @@ -83,7 +82,6 @@ def main() -> None: engine = AtomicActionEngine(motion_generator=motion_gen) engine.register( MoveJoints( - motion_gen, cfg=MoveJointsCfg( named_joint_positions={"ready": ready}, ), @@ -108,7 +106,7 @@ def main() -> None: compiled = engine.compile( ( ActionInvocation( - "move_joints", NamedJointPositionGoal("ready"), binding, policy + "move_joints", JointPositionGoal("ready"), binding, policy ), ActionInvocation( "move_joints", JointPositionGoal(waypoints), binding, policy diff --git a/scripts/tutorials/atomic_action/pickup.py b/scripts/tutorials/atomic_action/pickup.py index a8556d357..84b9fbb88 100644 --- a/scripts/tutorials/atomic_action/pickup.py +++ b/scripts/tutorials/atomic_action/pickup.py @@ -136,7 +136,6 @@ def main() -> None: engine = AtomicActionEngine(motion_generator=motion_gen) engine.register( PickUp( - motion_gen, cfg=PickUpCfg( hand_open_qpos=hand_open, hand_close_qpos=hand_close, diff --git a/scripts/tutorials/atomic_action/place.py b/scripts/tutorials/atomic_action/place.py index 85f2164c5..9fe690367 100644 --- a/scripts/tutorials/atomic_action/place.py +++ b/scripts/tutorials/atomic_action/place.py @@ -137,7 +137,6 @@ def main() -> None: engine = AtomicActionEngine(motion_generator=motion_gen) engine.register( PickUp( - motion_gen, cfg=PickUpCfg( hand_open_qpos=hand_open, hand_close_qpos=hand_close, @@ -149,7 +148,6 @@ def main() -> None: ) engine.register( Place( - motion_gen, cfg=PlaceCfg( hand_open_qpos=hand_open, hand_close_qpos=hand_close, diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py index dc51668a7..2df787461 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -164,10 +164,9 @@ def main() -> None: 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())) + engine.register(MoveEndEffector(MoveEndEffectorCfg())) engine.register( Press( - motion_gen, PressCfg( hand_close_qpos=hand_close, hand_interp_steps=HAND_INTERP_STEPS, diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index b79fe21d9..814e782b6 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -18,6 +18,7 @@ from __future__ import annotations +from typing import TypeVar from unittest.mock import Mock import pytest @@ -29,6 +30,8 @@ Affordance, AntipodalAffordance, AssembleGoal, + AtomicAction, + AtomicActionEngine, CoordinatedHeldObjectState, CoordinatedPickGoal, CoordinatedPickment, @@ -50,7 +53,6 @@ MoveHeldObjectCfg, MoveJoints, MoveJointsCfg, - NamedJointPositionGoal, ObjectSemantics, PickUp, PickUpCfg, @@ -74,6 +76,8 @@ 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: @@ -154,6 +158,13 @@ def _motion_generator() -> Mock: return generator +def _bind_action(generator: Mock, action: ActionT) -> ActionT: + """Bind one configured action to an engine-owned test backend.""" + engine = AtomicActionEngine(generator) + engine.register(action) + return action + + def _context(task: TaskState | None = None) -> PlanningContext: qpos = torch.zeros(NUM_ENVS, ROBOT_DOF) return PlanningContext( @@ -294,7 +305,7 @@ def _dual_binding( def test_builtin_descriptors_expose_goals_not_legacy_targets() -> None: assert MoveEndEffector.GoalType is EndEffectorPoseGoal - assert MoveJoints.GoalType == (JointPositionGoal, NamedJointPositionGoal) + assert MoveJoints.GoalType is JointPositionGoal assert PickUp.GoalType is GraspGoal assert MoveHeldObject.GoalType is HeldObjectPoseGoal assert Place.GoalType == (PlaceGoal, AssembleGoal) @@ -304,8 +315,34 @@ def test_builtin_descriptors_expose_goals_not_legacy_targets() -> None: assert HandOver.GoalType is GraspGoal +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_joint_position_goal_rejects_empty_named_target() -> None: + with pytest.raises(ValueError, match="must not be empty"): + JointPositionGoal(target=" ") + + +@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: - action = MoveEndEffector(_motion_generator(), MoveEndEffectorCfg()) + generator = _motion_generator() + action = _bind_action(generator, MoveEndEffector(MoveEndEffectorCfg())) context = _context() plan = action.plan( @@ -326,7 +363,10 @@ def test_move_end_effector_returns_full_robot_timed_plan() -> None: def test_move_joints_uses_binding_and_preserves_uncontrolled_joints() -> None: generator = _motion_generator() named = {"ready": torch.full((ARM_DOF,), 0.4)} - action = MoveJoints(generator, MoveJointsCfg(named_joint_positions=named)) + action = _bind_action( + generator, + MoveJoints(MoveJointsCfg(named_joint_positions=named)), + ) qpos = torch.zeros(NUM_ENVS, ROBOT_DOF) qpos[:, ARM_DOF:] = 0.7 context = PlanningContext( @@ -337,7 +377,7 @@ def test_move_joints_uses_binding_and_preserves_uncontrolled_joints() -> None: ) plan = action.plan( - _invocation("move_joints", NamedJointPositionGoal("ready"), sample_count=8), + _invocation("move_joints", JointPositionGoal("ready"), sample_count=8), context, ) @@ -349,9 +389,9 @@ def test_pick_and_place_declare_effects_without_mutating_context() -> None: generator = _motion_generator() hand_open = torch.zeros(HAND_DOF) hand_close = torch.ones(HAND_DOF) - pick = PickUp( + pick = _bind_action( generator, - PickUpCfg(hand_open_qpos=hand_open, hand_close_qpos=hand_close), + PickUp(PickUpCfg(hand_open_qpos=hand_open, hand_close_qpos=hand_close)), ) initial = _context() semantics = _semantics() @@ -366,9 +406,9 @@ def test_pick_and_place_declare_effects_without_mutating_context() -> None: assert initial.task.get_held_object("arm") is None assert picked_task.get_held_object("arm") is not None - place = Place( + place = _bind_action( generator, - PlaceCfg(hand_open_qpos=hand_open, hand_close_qpos=hand_close), + Place(PlaceCfg(hand_open_qpos=hand_open, hand_close_qpos=hand_close)), ) picked_context = PlanningContext( robot=initial.robot, @@ -390,9 +430,9 @@ def test_pick_and_place_declare_effects_without_mutating_context() -> None: def test_move_held_object_requires_projected_attachment() -> None: generator = _motion_generator() - action = MoveHeldObject( + action = _bind_action( generator, - MoveHeldObjectCfg(hand_close_qpos=torch.ones(HAND_DOF)), + MoveHeldObject(MoveHeldObjectCfg(hand_close_qpos=torch.ones(HAND_DOF))), ) invocation = _invocation( "move_held_object", @@ -416,7 +456,10 @@ def test_move_held_object_requires_projected_attachment() -> None: def test_press_uses_invocation_sample_budget() -> None: generator = _motion_generator() - action = Press(generator, PressCfg(hand_close_qpos=torch.ones(HAND_DOF))) + action = _bind_action( + generator, + Press(PressCfg(hand_close_qpos=torch.ones(HAND_DOF))), + ) plan = action.plan( _invocation("press", PressGoal(torch.eye(4)), sample_count=12), @@ -436,7 +479,8 @@ def test_motion_source_and_sample_count_are_not_action_config_fields() -> None: def test_move_joints_rejects_binding_with_wrong_goal_skill() -> None: - action = MoveJoints(_motion_generator(), MoveJointsCfg()) + generator = _motion_generator() + action = _bind_action(generator, MoveJoints(MoveJointsCfg())) invocation = ActionInvocation( skill_id="move_end_effector", goal=JointPositionGoal(torch.zeros(ARM_DOF)), @@ -460,7 +504,7 @@ def test_planner_timing_is_preserved_in_simple_action() -> None: NUM_ENVS, 1 ) generator.generate.return_value.duration = torch.full((NUM_ENVS,), 0.3) - action = MoveJoints(generator, MoveJointsCfg()) + action = _bind_action(generator, MoveJoints(MoveJointsCfg())) invocation = ActionInvocation( skill_id="move_joints", goal=JointPositionGoal(torch.ones(ARM_DOF)), @@ -494,7 +538,8 @@ def compute_ik( waypoints[:, 0, 0, 3] = 0.1 waypoints[:, 1, 0, 3] = 0.3 - plan = MoveEndEffector(generator).plan( + action = _bind_action(generator, MoveEndEffector()) + plan = action.plan( _invocation( "move_end_effector", EndEffectorPoseGoal(waypoints), @@ -510,9 +555,12 @@ def compute_ik( def test_move_joints_visits_waypoints_and_rejects_unknown_names() -> None: - action = MoveJoints( - _motion_generator(), - MoveJointsCfg(named_joint_positions={"ready": torch.zeros(ARM_DOF)}), + generator = _motion_generator() + action = _bind_action( + generator, + MoveJoints( + MoveJointsCfg(named_joint_positions={"ready": torch.zeros(ARM_DOF)}) + ), ) waypoints = torch.stack( [ @@ -535,7 +583,7 @@ def test_move_joints_visits_waypoints_and_rejects_unknown_names() -> None: assert torch.allclose(plan.trajectory.positions[:, -1, :ARM_DOF], waypoints[:, 1]) with pytest.raises(KeyError, match="Unknown named joint-position goal"): action.plan( - _invocation("move_joints", NamedJointPositionGoal("missing")), + _invocation("move_joints", JointPositionGoal("missing")), _context(), ) @@ -554,11 +602,13 @@ def test_pick_explicit_grasp_bypasses_sampling_and_records_grasp() -> None: ) grasp = torch.eye(4).repeat(NUM_ENVS, 1, 1) grasp[:, 0, 3] = torch.tensor([0.1, 0.2]) - action = PickUp( + action = _bind_action( generator, - PickUpCfg( - hand_open_qpos=torch.zeros(HAND_DOF), - hand_close_qpos=torch.ones(HAND_DOF), + PickUp( + PickUpCfg( + hand_open_qpos=torch.zeros(HAND_DOF), + hand_close_qpos=torch.ones(HAND_DOF), + ) ), ) @@ -585,9 +635,10 @@ def test_press_closes_hand_without_changing_projected_attachment() -> None: device="cpu", held_objects={"arm": held}, ) - action = Press( - _motion_generator(), - PressCfg(hand_close_qpos=torch.ones(HAND_DOF), hand_interp_steps=4), + generator = _motion_generator() + action = _bind_action( + generator, + Press(PressCfg(hand_close_qpos=torch.ones(HAND_DOF), hand_interp_steps=4)), ) plan = action.plan( @@ -605,18 +656,20 @@ def test_press_closes_hand_without_changing_projected_attachment() -> None: def test_handover_does_not_mutate_cached_final_pose() -> None: generator = _dual_motion_generator() - action = HandOver( + action = _bind_action( generator, - HandOverCfg( - transfer_hand_open_qpos=torch.zeros(HAND_DOF), - transfer_hand_close_qpos=torch.ones(HAND_DOF), - receive_hand_open_qpos=torch.zeros(HAND_DOF), - receive_hand_close_qpos=torch.ones(HAND_DOF), - middle_object_pose=torch.eye(4), - final_object_pose=torch.eye(4), - hand_interp_steps=4, - hold_steps=2, - retreat_steps=5, + HandOver( + HandOverCfg( + transfer_hand_open_qpos=torch.zeros(HAND_DOF), + transfer_hand_close_qpos=torch.ones(HAND_DOF), + receive_hand_open_qpos=torch.zeros(HAND_DOF), + receive_hand_close_qpos=torch.ones(HAND_DOF), + middle_object_pose=torch.eye(4), + final_object_pose=torch.eye(4), + hand_interp_steps=4, + hold_steps=2, + retreat_steps=5, + ) ), ) original_final_pose = action.final_object_pose.clone() @@ -658,16 +711,19 @@ def plan_from_start( def test_coordinated_pick_returns_full_dof_plan_and_projected_relation() -> None: - action = CoordinatedPickment( - _dual_motion_generator(), - CoordinatedPickmentCfg( - left_hand_open_qpos=torch.zeros(HAND_DOF), - left_hand_close_qpos=torch.ones(HAND_DOF), - right_hand_open_qpos=torch.zeros(HAND_DOF), - right_hand_close_qpos=torch.ones(HAND_DOF), - hand_interp_steps=4, - hold_steps=2, - object_motion_keyframes=3, + generator = _dual_motion_generator() + action = _bind_action( + generator, + CoordinatedPickment( + CoordinatedPickmentCfg( + left_hand_open_qpos=torch.zeros(HAND_DOF), + left_hand_close_qpos=torch.ones(HAND_DOF), + right_hand_open_qpos=torch.zeros(HAND_DOF), + right_hand_close_qpos=torch.ones(HAND_DOF), + hand_interp_steps=4, + hold_steps=2, + object_motion_keyframes=3, + ) ), ) semantics = ObjectSemantics( @@ -722,16 +778,18 @@ def fail_second_environment( return success, qpos generator.robot.compute_ik.side_effect = fail_second_environment - action = CoordinatedPickment( + action = _bind_action( generator, - CoordinatedPickmentCfg( - left_hand_open_qpos=torch.zeros(HAND_DOF), - left_hand_close_qpos=torch.ones(HAND_DOF), - right_hand_open_qpos=torch.zeros(HAND_DOF), - right_hand_close_qpos=torch.ones(HAND_DOF), - hand_interp_steps=4, - hold_steps=2, - object_motion_keyframes=3, + CoordinatedPickment( + CoordinatedPickmentCfg( + left_hand_open_qpos=torch.zeros(HAND_DOF), + left_hand_close_qpos=torch.ones(HAND_DOF), + right_hand_open_qpos=torch.zeros(HAND_DOF), + right_hand_close_qpos=torch.ones(HAND_DOF), + hand_interp_steps=4, + hold_steps=2, + object_motion_keyframes=3, + ) ), ) target_pose = torch.eye(4) @@ -768,15 +826,17 @@ def fail_second_environment( def test_coordinated_placement_projects_release_and_support_attachment() -> None: generator = _dual_motion_generator() - action = CoordinatedPlacement( + action = _bind_action( generator, - CoordinatedPlacementCfg( - placing_hand_open_qpos=torch.zeros(HAND_DOF), - placing_hand_close_qpos=torch.ones(HAND_DOF), - support_hand_close_qpos=torch.ones(HAND_DOF), - hand_interp_steps=4, - hold_steps=3, - retreat_steps=5, + CoordinatedPlacement( + CoordinatedPlacementCfg( + placing_hand_open_qpos=torch.zeros(HAND_DOF), + placing_hand_close_qpos=torch.ones(HAND_DOF), + support_hand_close_qpos=torch.ones(HAND_DOF), + hand_interp_steps=4, + hold_steps=3, + retreat_steps=5, + ) ), ) placing = _held( @@ -815,13 +875,15 @@ 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 = CoordinatedPickment( + pick = _bind_action( generator, - CoordinatedPickmentCfg( - left_hand_open_qpos=torch.zeros(HAND_DOF), - left_hand_close_qpos=torch.ones(HAND_DOF), - right_hand_open_qpos=torch.zeros(HAND_DOF), - right_hand_close_qpos=torch.ones(HAND_DOF), + CoordinatedPickment( + CoordinatedPickmentCfg( + left_hand_open_qpos=torch.zeros(HAND_DOF), + left_hand_close_qpos=torch.ones(HAND_DOF), + right_hand_open_qpos=torch.zeros(HAND_DOF), + right_hand_close_qpos=torch.ones(HAND_DOF), + ) ), ) pick_invocation = ActionInvocation( @@ -842,12 +904,14 @@ def test_coordinated_actions_reject_curobo_motion_generation() -> None: with pytest.raises(ValueError, match="not supported"): pick.plan(pick_invocation, _dual_context()) - placement = CoordinatedPlacement( + placement = _bind_action( generator, - CoordinatedPlacementCfg( - placing_hand_open_qpos=torch.zeros(HAND_DOF), - placing_hand_close_qpos=torch.ones(HAND_DOF), - support_hand_close_qpos=torch.ones(HAND_DOF), + CoordinatedPlacement( + CoordinatedPlacementCfg( + placing_hand_open_qpos=torch.zeros(HAND_DOF), + placing_hand_close_qpos=torch.ones(HAND_DOF), + support_hand_close_qpos=torch.ones(HAND_DOF), + ) ), ) placement_invocation = ActionInvocation( 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 2ae980638..1dde97ea8 100644 --- a/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py +++ b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py @@ -89,7 +89,7 @@ def _make_franka_curobo_engine(): ) ) engine = AtomicActionEngine(mg) - engine.register(MoveEndEffector(mg, MoveEndEffectorCfg())) + engine.register(MoveEndEffector(MoveEndEffectorCfg())) return sim, robot, engine diff --git a/tests/sim/atomic_actions/test_engine.py b/tests/sim/atomic_actions/test_engine.py index 9ae563a20..615146b82 100644 --- a/tests/sim/atomic_actions/test_engine.py +++ b/tests/sim/atomic_actions/test_engine.py @@ -54,7 +54,8 @@ def plan( context: PlanningContext, ) -> ActionPlan: goal = self.require_goal(invocation) - target = goal.qpos.to(context.robot.qpos) + 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( @@ -72,6 +73,12 @@ def plan( ) +class OtherStubAction(StubAction): + """Second configured skill used to verify shared engine resources.""" + + skill_id: ClassVar[str] = "other_stub" + + def _engine(batch_size: int = 2, robot_dof: int = 3) -> AtomicActionEngine: robot = Mock() robot.device = torch.device("cpu") @@ -106,7 +113,7 @@ def test_global_registry_uses_stable_skill_id() -> None: def test_engine_compile_projects_terminal_state_between_actions() -> None: engine = _engine() - engine.register(StubAction(engine.motion_generator, ActionCfg(name="stub"))) + engine.register(StubAction(ActionCfg(name="stub"))) first = torch.ones(2, 3) second = torch.full((2, 3), 2.0) @@ -121,7 +128,7 @@ def test_engine_compile_projects_terminal_state_between_actions() -> None: def test_engine_compile_holds_failed_rows_for_remaining_actions() -> None: engine = _engine() - engine.register(StubAction(engine.motion_generator, ActionCfg(name="stub"))) + engine.register(StubAction(ActionCfg(name="stub"))) first = torch.tensor([[1.0, 1.0, 1.0], [float("nan"), 2.0, 2.0]]) second = torch.full((2, 3), 4.0) @@ -152,16 +159,70 @@ def test_engine_rejects_unknown_skill() -> None: def test_engine_rejects_duplicate_instance_registration() -> None: engine = _engine() - first = StubAction(engine.motion_generator, ActionCfg(name="first")) - second = StubAction(engine.motion_generator, ActionCfg(name="second")) + first = StubAction(ActionCfg(name="first")) + second = StubAction(ActionCfg(name="second")) engine.register(first) with pytest.raises(ValueError, match="already registered"): engine.register(second) +def test_engine_binds_one_planning_service_to_every_action() -> None: + engine = _engine() + first = StubAction(ActionCfg(name="first")) + second = OtherStubAction(ActionCfg(name="second")) + + 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_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(ActionCfg(name="temporary")) + + 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(ActionCfg(name="stub")) + _engine().register(action) + + with pytest.raises(ValueError, match="another AtomicActionEngine"): + _engine().register(action) + + +def test_unbound_action_rejects_direct_planning() -> None: + action = StubAction(ActionCfg(name="stub")) + + with pytest.raises(RuntimeError, match="not bound"): + action.plan( + _invocation(torch.ones(2, 3)), + _engine().initial_context(), + ) + + def test_engine_rejects_plan_for_a_different_skill() -> None: engine = _engine() - action = StubAction(engine.motion_generator, ActionCfg(name="stub")) + action = StubAction(ActionCfg(name="stub")) original_plan = action.plan def wrong_skill_plan( diff --git a/tests/sim/atomic_actions/test_engine_per_env.py b/tests/sim/atomic_actions/test_engine_per_env.py index 12318ce36..5b9848248 100644 --- a/tests/sim/atomic_actions/test_engine_per_env.py +++ b/tests/sim/atomic_actions/test_engine_per_env.py @@ -57,8 +57,8 @@ class DynamicAction(AtomicAction[EndEffectorPoseGoal]): GoalType: ClassVar[type] = EndEffectorPoseGoal manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) - def __init__(self, motion_generator) -> None: - super().__init__(motion_generator, ActionCfg(name="dynamic")) + def __init__(self) -> None: + super().__init__(ActionCfg(name="dynamic")) self.plan_count = 0 def plan( @@ -119,7 +119,7 @@ def _engine() -> tuple[AtomicActionEngine, DynamicAction]: generator.device = torch.device("cpu") generator.planner.cfg.planner_type = "stub" engine = AtomicActionEngine(generator) - action = DynamicAction(generator) + action = DynamicAction() engine.register(action) return engine, action @@ -267,7 +267,7 @@ def test_session_rejects_regressing_scene_snapshot() -> None: def test_nonempty_effect_is_committed_only_after_external_verification() -> None: engine, _ = _engine() - effect = EffectAction(engine.motion_generator) + effect = EffectAction() engine.register(effect) invocation = _invocation() invocation = ActionInvocation( @@ -299,7 +299,7 @@ def test_nonempty_effect_is_committed_only_after_external_verification() -> None def test_effect_failure_does_not_commit_and_exhausts_retry_budget() -> None: engine, _ = _engine() - engine.register(EffectAction(engine.motion_generator)) + engine.register(EffectAction()) base = _invocation(max_phase_retries=0) invocation = ActionInvocation( skill_id="effect", diff --git a/tests/sim/atomic_actions/test_motion_source_e2e.py b/tests/sim/atomic_actions/test_motion_source_e2e.py index 2c949ce95..6b434a24a 100644 --- a/tests/sim/atomic_actions/test_motion_source_e2e.py +++ b/tests/sim/atomic_actions/test_motion_source_e2e.py @@ -60,7 +60,7 @@ def _setup(self, motion_source: str): MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=self.ROBOT_UID)) ) engine = AtomicActionEngine(mg) - engine.register(MoveEndEffector(mg, MoveEndEffectorCfg())) + engine.register(MoveEndEffector(MoveEndEffectorCfg())) return sim, robot, engine def _teardown(self, sim): diff --git a/tests/sim/planners/test_curobo_planner.py b/tests/sim/planners/test_curobo_planner.py index 81edea96e..a4e01cb56 100644 --- a/tests/sim/planners/test_curobo_planner.py +++ b/tests/sim/planners/test_curobo_planner.py @@ -645,7 +645,7 @@ def _make_curobo_engine( ) ) engine = AtomicActionEngine(motion_generator) - engine.register(MoveEndEffector(motion_generator, MoveEndEffectorCfg())) + engine.register(MoveEndEffector(MoveEndEffectorCfg())) return engine From dcaf45ea38cabb9b5b5498a51dda2bb90b1cd026 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 2 Aug 2026 17:58:51 +0000 Subject: [PATCH 4/5] refactor(atomic-actions): replace configs with runtime options --- .agents/skills/add-atomic-action/SKILL.md | 77 ++-- agent_context/MAP.yaml | 11 +- .../topics/atomic-actions/atomic-actions.md | 77 +++- ...hain.lab.sim.atomic_actions.primitives.rst | 22 +- .../embodichain.lab.sim.atomic_actions.rst | 32 ++ .../sim/atomic_actions/builtin_actions.md | 177 +++++--- .../overview/sim/atomic_actions/index.md | 252 +++++++++-- .../overview/sim/planners/curobo_planner.md | 6 +- docs/source/tutorial/atomic_actions.rst | 118 ++++- .../lab/sim/atomic_actions/__init__.py | 73 +-- .../lab/sim/atomic_actions/bindings.py | 210 ++++++++- embodichain/lab/sim/atomic_actions/control.py | 247 ++++++++++ embodichain/lab/sim/atomic_actions/core.py | 149 ++++-- embodichain/lab/sim/atomic_actions/engine.py | 105 ++++- .../lab/sim/atomic_actions/execution.py | 97 +++- .../lab/sim/atomic_actions/invocation.py | 88 +++- embodichain/lab/sim/atomic_actions/plans.py | 6 + .../lab/sim/atomic_actions/policies.py | 16 +- .../sim/atomic_actions/primitives/__init__.py | 40 +- .../primitives/coordinated_pickment.py | 408 +++++++---------- .../primitives/coordinated_placement.py | 317 +++++++------ .../atomic_actions/primitives/hand_over.py | 424 ++++++++++-------- .../primitives/move_end_effector.py | 46 +- .../primitives/move_held_object.py | 126 +++--- .../atomic_actions/primitives/move_joints.py | 69 +-- .../sim/atomic_actions/primitives/pick_up.py | 294 +++++++----- .../sim/atomic_actions/primitives/place.py | 190 ++++---- .../sim/atomic_actions/primitives/press.py | 127 +++--- embodichain/lab/sim/atomic_actions/runtime.py | 191 +++++++- .../lab/sim/planners/curobo/curobo_planner.py | 2 +- examples/sim/planners/curobo_planner.py | 3 +- .../move_end_effector_benchmark.py | 3 +- .../move_held_object_benchmark.py | 41 +- .../atomic_action/move_joints_benchmark.py | 15 +- .../atomic_action/pickup_benchmark.py | 19 +- .../atomic_action/place_benchmark.py | 38 +- .../atomic_action/press_benchmark.py | 24 +- .../neural_planner/BENCHMARK_DESIGN.md | 6 +- scripts/tutorials/atomic_action/assemble.py | 29 +- .../atomic_action/coordinated_pickment.py | 28 +- .../atomic_action/coordinated_placement.py | 41 +- scripts/tutorials/atomic_action/hand_over.py | 38 +- .../atomic_action/move_end_effector.py | 3 +- .../atomic_action/move_held_object.py | 29 +- .../tutorials/atomic_action/move_joints.py | 15 +- scripts/tutorials/atomic_action/pickup.py | 17 +- scripts/tutorials/atomic_action/place.py | 23 +- scripts/tutorials/atomic_action/press.py | 21 +- tests/sim/atomic_actions/test_actions.py | 281 +++++++----- tests/sim/atomic_actions/test_control.py | 137 ++++++ .../test_curobo_motion_source_e2e.py | 3 +- tests/sim/atomic_actions/test_engine.py | 112 +++-- .../sim/atomic_actions/test_engine_per_env.py | 78 +++- .../atomic_actions/test_motion_source_e2e.py | 3 +- tests/sim/planners/test_curobo_planner.py | 3 +- 55 files changed, 3340 insertions(+), 1667 deletions(-) create mode 100644 embodichain/lab/sim/atomic_actions/control.py create mode 100644 tests/sim/atomic_actions/test_control.py diff --git a/.agents/skills/add-atomic-action/SKILL.md b/.agents/skills/add-atomic-action/SKILL.md index dbf587540..a007497c4 100644 --- a/.agents/skills/add-atomic-action/SKILL.md +++ b/.agents/skills/add-atomic-action/SKILL.md @@ -7,7 +7,7 @@ description: Add a new simulation atomic action or motion primitive to EmbodiCha Add an action-owned goal and a side-effect-free `AtomicAction.plan()` implementation. The engine owns all motion-planning resources; action -constructors accept only implementation configuration. Keep task-graph/MLLM +constructors accept only optional typed default options. Keep task-graph/MLLM logic, simulator stepping, controller I/O, and physical-effect commits outside the action. @@ -20,6 +20,8 @@ Inspect only the files relevant to the requested skill: | 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` | @@ -66,37 +68,42 @@ options, retry counts, live state, or a generic optional field bag. Use Use `ObjectActionGoal` only when the shared `semantics` field is genuinely required. -## 2. Define implementation configuration +## 2. Define runtime options and control commands -Extend `ActionCfg` directly with `@configclass`. Keep only implementation-owned -behavior such as distances, gripper positions, grasp constraints, and phase -split counts. +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 the action config; -those belong to `MotionPolicy` or `RecoveryPolicy` on the invocation. +velocity limits, collision policy, or recovery thresholds in skill options; +those belong to `MotionPolicy` or `RecoveryPolicy`. ```python -from embodichain.lab.sim.atomic_actions import ActionCfg -from embodichain.utils import configclass +from dataclasses import dataclass + +from embodichain.lab.sim.atomic_actions import ActionOptions -@configclass -class PushCfg(ActionCfg): - name: str = "push" +@dataclass(frozen=True, slots=True, eq=False) +class PushOptions(ActionOptions): push_distance: float = 0.05 ``` +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. + ## 3. Implement the planner -Inherit `AtomicAction[PushGoal]` directly. Declare stable metadata and resolve +Inherit `AtomicAction[PushGoal, PushOptions]` directly. Declare stable metadata and resolve resources from semantic binding roles. ```python from typing import ClassVar from embodichain.lab.sim.atomic_actions import ( - ActionInvocation, + ResolvedActionRequest, ActionPlan, AtomicAction, PlanningContext, @@ -104,43 +111,46 @@ from embodichain.lab.sim.atomic_actions import ( ) -class Push(AtomicAction[PushGoal]): +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, cfg: PushCfg | None = None) -> None: - super().__init__(cfg or PushCfg()) + def __init__(self, default_options: PushOptions | None = None) -> None: + super().__init__(default_options) def plan( self, - invocation: ActionInvocation[PushGoal], + request: ResolvedActionRequest[PushGoal, PushOptions], context: PlanningContext, ) -> ActionPlan: - goal = self.require_goal(invocation) - control_part = invocation.binding.manipulator("primary") - joint_ids = self.robot.get_joint_ids(name=control_part) + 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 - # invocation.motion_policy. Embed it into full robot DoF. + # request.motion_policy. Embed it into full robot DoF. result = self.builder.generate_arm_plan( target_states, start_qpos, - invocation.motion_policy.sample_count, + request.motion_policy.sample_count, control_part=control_part, - arm_dof=len(joint_ids), - cfg=invocation.motion_policy, + arm_dof=manipulator.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=invocation.motion_policy.control_dt, + control_dt=request.motion_policy.control_dt, ) return self.build_plan( - invocation, + request, context, success=success, trajectory=trajectory, @@ -157,7 +167,7 @@ Follow these invariants: - 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(invocation, context, message=...)` for an expected soft +- 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. @@ -171,7 +181,7 @@ Follow these invariants: Register an instance by its class-level `skill_id`: ```python -engine.register(Push(PushCfg())) +engine.register(Push()) ``` Use the global registry only for discoverable third-party classes: @@ -198,7 +208,7 @@ error recovery are required. ## 5. Export and document -Export the goal, config, and action from: +Export the goal, options, and action from: 1. `embodichain/lab/sim/atomic_actions/primitives/__init__.py` 2. `embodichain/lab/sim/atomic_actions/__init__.py` @@ -231,8 +241,11 @@ then use the `pre-commit-check` skill before committing. | 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 planner/recovery knobs in action config | Move them to invocation policies. | -| Pass a motion generator to each action | Pass it once to `AtomicActionEngine`; construct actions from config only. | +| 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`. | diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 81dbd04bf..c597adca7 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -441,7 +441,12 @@ topics: - SceneEntityPose - dynamic goal - error recovery - - ActionCfg + - ActionOptions + - ResolvedActionRequest + - ControlPartCommandProfile + - ActionControlOverrides + - JointPositionCommand + - invocation revision - MotionPolicy - RecoveryPolicy - plan_arm_traj @@ -452,6 +457,10 @@ topics: - embodichain/lab/sim/atomic_actions/core.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 diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index a0a4c1e8c..c6f69d65a 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -14,9 +14,12 @@ There is no `ActionTarget`, `WorldState`, `ActionResult`, `execute()`, or `ActionInvocation` separates: - an action-owned typed goal (`goal_kind` is its stable discriminator); -- `ActionBinding`, which maps semantic roles to concrete robot resources; +- `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. +- 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` @@ -25,9 +28,10 @@ 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, and -shared `TrajectoryBuilder`. Actions accept only implementation configuration -and borrow those services after `engine.register(action)` or +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 after +`engine.register(action)` or `engine.plan_action(action, invocation, context)`. A bound action cannot be reused by another engine. @@ -72,6 +76,18 @@ 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. +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: + +```python +session.revise_current(revised_invocation) +``` + +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 @@ -79,16 +95,46 @@ 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. Action configs retain -only implementation-specific behavior such as gripper poses, phase splits, -lift distances, and grasp constraints. +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. + +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, + control_profiles={ + "left_hand": ControlPartCommandProfile.joint_positions( + open=left_open_qpos, + grasp=left_grasp_qpos, + ), + "left_arm": ControlPartCommandProfile.joint_positions(ready=ready_qpos), + }, +) +``` + +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-ins | Skill ID | Goal type | Roles | |---|---|---| | `move_end_effector` | `EndEffectorPoseGoal` | manipulator `primary` | -| `move_joints` | `JointPositionGoal` (`target` is explicit qpos or a configured name) | 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` | @@ -100,11 +146,12 @@ lift distances, and grasp constraints. ## Extension rules 1. Define a frozen action-owned goal dataclass with `goal_kind`. -2. Declare `skill_id`, `GoalType`, and required semantic roles on the action. -3. Validate with `require_goal(invocation)`. -4. Plan from `context.robot.qpos`; never read an implicit live start state. -5. Return full-robot positions or a `TimedTrajectory` through `build_plan()`. -6. Declare symbolic changes with `StateDelta`; do not mutate context or commit +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. -7. Keep scene stepping, controller I/O, and task-graph/MLLM logic outside the +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 8692f3194..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,7 +8,7 @@ Overview Concrete implementations of the built-in atomic-action primitives. Each primitive is an :class:`~embodichain.lab.sim.atomic_actions.AtomicAction` that -accepts an :class:`~embodichain.lab.sim.atomic_actions.ActionInvocation` and a +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. @@ -17,24 +17,24 @@ full-robot timed trajectory and uncommitted expected effects. .. autosummary:: - MoveEndEffectorCfg MoveEndEffector - MoveJointsCfg + MoveEndEffectorOptions MoveJoints - PickUpCfg + MoveJointsOptions PickUp - MoveHeldObjectCfg + PickUpOptions MoveHeldObject - PlaceCfg + MoveHeldObjectOptions Place - PressCfg + PlaceOptions Press - CoordinatedPickmentCfg + PressOptions CoordinatedPickment - CoordinatedPlacementCfg + CoordinatedPickmentOptions CoordinatedPlacement - HandOverCfg + CoordinatedPlacementOptions HandOver + HandOverOptions .. rubric:: Built-in Goal Contracts 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 d38028b53..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 @@ -9,7 +9,15 @@ embodichain.lab.sim.atomic_actions ActionGoal ActionBinding + ResolvedActionBinding + ResolvedControlPart + ControlCommand + JointPositionCommand + ControlPartCommandProfile + ActionControlOverrides ActionInvocation + ResolvedActionRequest + ActionOptions MotionPolicy RecoveryPolicy RobotObservation @@ -72,9 +80,33 @@ Planning and state .. autoclass:: ActionBinding :members: +.. autoclass:: ResolvedActionBinding + :members: + +.. autoclass:: ResolvedControlPart + :members: + +.. autoclass:: ControlCommand + :members: + +.. autoclass:: JointPositionCommand + :members: + +.. autoclass:: ControlPartCommandProfile + :members: + +.. autoclass:: ActionControlOverrides + :members: + .. autoclass:: ActionInvocation :members: +.. autoclass:: ResolvedActionRequest + :members: + +.. autoclass:: ActionOptions + :members: + .. autoclass:: MotionPolicy :members: :exclude-members: __init__, copy, replace, to_dict diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index 62afb8841..624c5c28c 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -11,16 +11,19 @@ applications register the configured instances they need with an engine. release primitive instead of introducing a tenth skill ID. All built-ins implement -`plan(invocation, context) -> ActionPlan`. Their constructors accept only -action configuration; the owning `AtomicActionEngine` supplies one shared -motion generator and trajectory builder during `register()` or -`plan_action()`. Generic motion and recovery choices belong to the invocation, -not the action config. +`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 during +`register()` or `plan_action()`. Generic motion and recovery choices belong to +the invocation. ```{note} -The current manipulation primitives use gripper open/close joint positions. -Replacing a gripper with a dexterous hand requires a hand-command abstraction -or new hand-specific phases; it is not yet a drop-in config change. +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. ``` ## Visual catalog @@ -125,17 +128,39 @@ The animations below are the focused simulator demos under ## Capability matrix -| Skill ID | Accepted goal | Required binding roles | Required task state | Expected task effect | -|---|---|---|---|---| -| `move_end_effector` | `EndEffectorPoseGoal` | manipulator `primary` | none | none | -| `move_joints` | `JointPositionGoal` | manipulator `primary` | none | none | -| `pick_up` | `GraspGoal` | manipulator + end effector `primary` | semantic object/entity | attach object to `primary` manipulator | -| `move_held_object` | `HeldObjectPoseGoal` | manipulator + end effector `primary` | object held by `primary` | preserve attachment | -| `place` | `PlaceGoal`, `AssembleGoal` | manipulator + end effector `primary` | `AssembleGoal` requires an object held by `primary`; ordinary `PlaceGoal` has no planner-enforced attachment precondition | detach object | -| `press` | `PressGoal` | manipulator + end effector `primary` | none | none | -| `coordinated_pickment` | `CoordinatedPickGoal` | manipulator + end effector `left`, `right` | semantic object/entity | create coordinated attachment; clear individual attachments | -| `coordinated_placement` | `CoordinatedPlacementGoal` | manipulator + end effector `placing`, `support` | one individually held object per arm | optionally detach placing object; preserve support attachment | -| `hand_over` | `GraspGoal` | manipulator + end effector `source`, `destination` | object held by source arm | transfer attachment to destination arm | +| 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 | +|---|---|---| +| `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 @@ -177,16 +202,22 @@ entity as a recovery dependency. Use this rule when configuring a built-in or adding a new one: - the **goal** carries only the requested outcome; -- the **binding** carries concrete robot resources selected for this call; -- the **action config** carries hardware constants and phase-specific behavior; +- 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. -Complex manipulation actions currently validate that their semantic bindings -match concrete resources configured for their hand qpos and multi-part phase -assembly. `MoveEndEffector` and `MoveJoints` resolve the manipulator entirely -from `ActionBinding`. +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 @@ -211,7 +242,7 @@ ordered set of pose waypoints. | Motion | EEF planning from observed arm qpos; output expanded to full robot DoF | | Completion | `EEF_GOAL_REACHED` | | Effect | none | -| Action config | only the inherited action name; reusable motion choices are in `MotionPolicy` | +| 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 @@ -239,14 +270,17 @@ than an EEF pose. `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 `MoveJointsCfg.named_joint_positions`. Named poses are -implementation/hardware knowledge and remain in the action config rather than -becoming separate goal types: +from the bound manipulator's `ControlPartCommandProfile`. Named poses remain +embodiment knowledge without becoming separate goal types: ```python -move_joints = MoveJoints( - MoveJointsCfg(named_joint_positions={"home": home_qpos}) +engine = AtomicActionEngine( + motion_generator, + control_profiles={ + "left_arm": ControlPartCommandProfile.joint_positions(home=home_qpos), + }, ) +engine.register(MoveJoints()) explicit_goal = JointPositionGoal(target=home_qpos) named_goal = JointPositionGoal(target="home") @@ -267,7 +301,7 @@ bound manipulator. | 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 configured manipulator and clear overlapping coordinated attachment state | +| Effect | write `HeldObjectState` for the bound manipulator and clear overlapping coordinated attachment state | | Verification | the attachment effect must be verified during closed-loop execution | `grasp_xpos` may be `(4, 4)` or `(B, 4, 4)`. When omitted, the action samples @@ -275,12 +309,11 @@ 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. -Important `PickUpCfg` fields: +`PickUp` requires `open` and `grasp` commands on the bound end-effector profile. +Important `PickUpOptions` fields: | Field | Purpose | |---|---| -| `control_part`, `hand_control_part` | Concrete resources that the `primary` binding must match | -| `hand_open_qpos`, `hand_close_qpos` | Required hardware-specific hand states | | `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` | @@ -307,14 +340,15 @@ the action derives `target_object_pose @ object_to_eef` from verified task state | Skill ID | `move_held_object` | | Goal | `HeldObjectPoseGoal(object_target_pose=...)` | | Binding | manipulator + end effector role `primary` | -| Precondition | a `HeldObjectState` exists for the configured manipulator, normally from `PickUp` | +| 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` | -`MoveHeldObjectCfg` holds the concrete arm/hand names, required -`hand_close_qpos`, and optional upright-transport settings. Generic timing and -trajectory sampling remain in `MotionPolicy`. +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` @@ -331,7 +365,7 @@ one. | Skill ID | `place` | | Goal | `PlaceGoal(xpos=..., tcp_symmetry="none")` | | Binding | manipulator + end effector role `primary` | -| State | consumes the configured manipulator's attachment when present | +| 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` | @@ -341,12 +375,11 @@ translation remain physically equivalent. The action selects the closer orientation variant from the observed starting state and uses it consistently across all waypoints. -Important `PlaceCfg` fields: +The bound end-effector profile must provide `open` and `grasp`. Important +`PlaceOptions` fields: | Field | Purpose | |---|---| -| `control_part`, `hand_control_part` | Concrete resources that the `primary` binding must match | -| `hand_open_qpos`, `hand_close_qpos` | Required release/holding hand states | | `lift_height` | Approach and retract height | | `hand_interp_steps` | Open-phase discretization | | `max_approach_retract_z` | Optional world-Z ceiling for approach/retract poses | @@ -396,10 +429,11 @@ arm should retreat along its planned path after reaching the target. | Effect | none; existing attachment state is unchanged | | Dynamic target | explicit pose or `SceneEntityPose` | -`PressCfg` pins the concrete arm/hand resources, required `hand_close_qpos`, and -`hand_interp_steps`. Contact detection is not itself a symbolic effect in the -current action; applications that require force/contact confirmation should -verify it externally. +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` @@ -423,16 +457,16 @@ The object target and optional initial pose may use `SceneEntityPose`. When no initial pose is supplied, `ObjectSemantics.entity` provides the object's current pose. -Important `CoordinatedPickmentCfg` fields group into: +Both bound end-effector profiles must provide `open` and `grasp`. Important +`CoordinatedPickmentOptions` fields group into: -- combined, left/right arm, and left/right hand control-part names; -- required open/close qpos for both hands; - `pre_grasp_distance` and `lift_height`; - `object_motion_keyframes`, `hand_interp_steps`, and `hold_steps`. -The semantic binding must match those configured resources. 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. +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. **Example:** `scripts/tutorials/atomic_action/coordinated_pickment.py` @@ -448,23 +482,22 @@ hold -> optionally release the placing hand -> retreat the placing arm**. | Skill ID | `coordinated_placement` | | Goal | `CoordinatedPlacementGoal` | | Binding | manipulator + end effector roles `placing` and `support` | -| Precondition | separate `HeldObjectState` entries exist for both configured arms | +| 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 | Both object targets may use `SceneEntityPose`, so either can participate in -dynamic-goal invalidation. Goal-level height/release values override defaults in -the action config for that invocation. +dynamic-goal invalidation. Goal-level height/release values override +`CoordinatedPlacementOptions` for that invocation. -Important `CoordinatedPlacementCfg` fields group into: +The placing profile must provide `open` and `grasp`; the support profile must +provide `grasp`. Important `CoordinatedPlacementOptions` fields group into: -- combined, placing/support arm, and placing/support hand control-part names; -- required placing-hand open/close and support-hand close qpos; - default `release`, placing/support height offsets, and `lift_height`; - `hand_interp_steps`, `hold_steps`, and `retreat_steps`. -The semantic binding must match those configured resources. The same cuRobo -restriction as coordinated pickment applies to dual-arm +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. **Example:** `scripts/tutorials/atomic_action/coordinated_placement.py` @@ -486,16 +519,18 @@ retreats -> destination delivers**. | Effect | remove source attachment and create destination `HeldObjectState` | | Verification | attachment transfer must be externally verified | -`HandOverCfg` currently owns the concrete source/destination arm and hand names, -all four open/close hand qpos values, destination grasp region and approach -direction, middle/final object poses, and phase distances/counts. The semantic -binding must match these configured resources. - -The middle and final poses are currently fixed configuration tensors rather -than `SceneEntityPose` goal fields. Consequently, handover supports tracking- -error and timeout recovery, but does not yet provide automatic moving-handover- -point invalidation. The action queries the semantic object's live orientation -when replanning and preserves that orientation at the configured middle/final +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 diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index 476bbba55..b6f3bef89 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -27,14 +27,20 @@ and whole-body control are not implemented by this module yet. ## Architecture and responsibility boundary ```text -Action Agent / task graph - | - | semantic skill call (object, destination, constraints) - v -grounder + capability binder - | - | ActionInvocation + PlanningContext - v ++--------------------------------+ +--------------------------------+ +| 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 | | | @@ -60,13 +66,48 @@ The boundary is deliberate: | Concern | Owner | Contract | |---|---|---| -| Task intent and sequencing | Action Agent or task graph | Selects a skill and semantic goal | -| Perception and grounding | Application adapter | Builds scene snapshots, object semantics, and resource bindings | +| 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, and trajectory builder | +| 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 + +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), +) + +# Inspect a single plan, compile a fixed sequence, or execute with recovery. +plan = engine.plan(manual_invocation, latest_context) +compiled = engine.compile((manual_invocation,), latest_context) +session = engine.start((manual_invocation,), latest_context) +``` + +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. + ## Core contracts The public contracts separate values with different owners and lifetimes. This @@ -76,8 +117,10 @@ from leaking into an Action Agent schema. | Contract | Contains | Does not contain | |---|---|---| | `ActionGoal` | Action-specific desired outcome, such as an EEF pose or object pose | Arm names, planner instances, recovery counters | -| `ActionBinding` | Semantic-role mappings such as `primary -> left_arm` and `primary -> left_hand` | Motion settings or task geometry | -| `ActionCfg` | Implementation and hardware constants: hand qpos, grasp-selection rules, phase structure | Per-call goal, motion generator, generic recovery settings | +| `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 | @@ -89,7 +132,17 @@ base class and no closed union that must change whenever a skill is added. ### Semantic resource binding -Bindings make an invocation portable across embodiments: +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 binding = ActionBinding( @@ -98,17 +151,90 @@ binding = ActionBinding( ) ``` -Single-resource skills use `primary`; handover uses `source` and `destination`; -coordinated pick uses `left` and `right`; coordinated placement uses `placing` -and `support`. The action descriptor declares which roles are required, so a -grounder can validate a call before planning. +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, + ), + }, +) +``` -Simple motion skills resolve their concrete control part entirely from the -binding. Some current multi-phase manipulation implementations also keep -concrete arm/hand names in their action config because their preconfigured hand -qpos and phase assembly are hardware-specific. For those actions, the binding -must match the configured resources. This is an implementation constraint, not -a reason to put resource names back into goals. +`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. + +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), + } + } + ), + 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 @@ -116,25 +242,27 @@ One engine owns one motion generator. Actions borrow its planning services only after `register()` or `plan_action()` binds them: ```python -engine = AtomicActionEngine(motion_generator) -engine.register(MoveEndEffector(MoveEndEffectorCfg())) -engine.register(MoveJoints(MoveJointsCfg())) +engine = AtomicActionEngine(motion_generator, control_profiles=profiles) +engine.register(MoveEndEffector()) +engine.register(MoveJoints()) ``` Consequences of this ownership model: -- action constructors contain only skill configuration; +- action constructors optionally contain only typed default 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. -When two differently configured instances share the same stable skill ID, keep -one or both outside the registry and call `engine.plan_action(...)` explicitly: +Prefer invocation `skill_options` when behavior varies per call. If an +application still needs two instances with different default options and the +same stable skill ID, keep one or both outside the registry and call +`engine.plan_action(...)` explicitly: ```python -left_pick = PickUp(left_pick_cfg) -right_pick = PickUp(right_pick_cfg) +left_pick = PickUp(default_options=left_pick_options) +right_pick = PickUp(default_options=right_pick_options) left_plan = engine.plan_action(left_pick, left_invocation, latest_context) right_plan = engine.plan_action(right_pick, right_invocation, latest_context) @@ -146,11 +274,12 @@ Both instances still borrow the same engine-owned motion generator. | API | Use it for | Result / behavior | |---|---|---| -| `AtomicAction.plan(invocation, context)` | Implementing a skill | Action-owned side-effect-free planning hook; application code normally calls it through the engine | +| `AtomicAction.plan(request, context)` | Implementing a skill | Consumes an engine-resolved immutable request; application code normally calls it through the engine | | `engine.plan(invocation, context)` | Planning one registered skill | Resolves the registered action, binds shared resources, and validates its plan | | `engine.plan_action(action, invocation, context)` | Planning an unregistered configured instance | Supports multiple configurations with one `skill_id` and one engine backend | | `engine.compile(invocations, context)` | Fixed-scene/offline sequence planning | Returns one concatenated `CompiledTrajectory` and a hypothetical projected context | | `engine.start(invocations, context)` | Observed incremental execution | Returns an `ExecutionSession`; each `tick()` emits at most one command and recovery events | +| `session.revise_current(invocation)` | Explicit runtime parameter/goal update | Requires a newer revision of the active logical invocation and replans from the latest context | `AtomicAction.plan()` is therefore not a second execution API. It is the polymorphic implementation point used by the engine. Neither it nor the engine @@ -243,6 +372,35 @@ 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 +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) +``` + +`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 @@ -283,28 +441,38 @@ 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 ``` -This keeps learned reasoning independent of planner instances and concrete -joint groups, while invocation IDs, planner diagnostics, and execution events -provide structured feedback for the agent's next decision. +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. declare `skill_id`, `GoalType`, required semantic roles, and agent visibility; -3. keep embodiment constants in its `ActionCfg` and reusable motion/recovery - choices in invocation policies; -4. implement side-effect-free `plan(invocation, context)` using the +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; -5. return full-robot timed motion, per-environment planning success, +6. return full-robot timed motion, per-environment planning success, diagnostics, and uncommitted effects; -6. add registration coverage, contract tests, execution/recovery tests, a +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 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 59c20471b..a5635cfc1 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -1,31 +1,50 @@ Atomic actions ============== -Atomic actions are typed, side-effect-free motion planners. An action receives a -grounded :class:`~embodichain.lab.sim.atomic_actions.ActionInvocation` and the -latest :class:`~embodichain.lab.sim.atomic_actions.PlanningContext`, then returns -an :class:`~embodichain.lab.sim.atomic_actions.ActionPlan`. +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 four concerns: +The contracts deliberately separate six concerns: * a **goal** describes what should happen; * an **ActionBinding** maps semantic roles such as ``primary`` or ``source`` to - robot control resources; + 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. -The engine exclusively owns the ``MotionGenerator`` and a shared trajectory -builder. Atomic action constructors accept only implementation configuration; -``register()`` binds each action to the engine resources. Use +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. Atomic action constructors accept only optional +typed default options; ``register()`` binds each action to the engine resources. Use ``engine.plan_action(action, invocation, context)`` for an unregistered, -configuration-specific action instance. +default-option-specific action instance. Runnable examples ----------------- @@ -58,6 +77,38 @@ 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. +Control-part commands +--------------------- + +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 + + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + ControlPartCommandProfile, + ) + + 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, + ), + }, + ) + +``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. + Static compilation ------------------ @@ -73,11 +124,10 @@ the scene is treated as fixed during planning: EndEffectorPoseGoal, MotionPolicy, MoveEndEffector, - MoveEndEffectorCfg, ) engine = AtomicActionEngine(motion_generator) - engine.register(MoveEndEffector(MoveEndEffectorCfg())) + engine.register(MoveEndEffector()) invocation = ActionInvocation( skill_id="move_end_effector", @@ -133,6 +183,29 @@ 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 + + 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) + +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. + 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 @@ -160,7 +233,8 @@ Adding an action ---------------- Define an action-owned frozen goal dataclass with a stable ``goal_kind``. Then -implement ``plan(invocation, context)`` and declare the stable skill metadata: +define typed runtime options when needed, implement ``plan(request, context)``, +and declare the stable skill metadata: .. code-block:: python @@ -172,24 +246,30 @@ implement ``plan(invocation, context)`` and declare the stable skill metadata: goal_kind: ClassVar[str] = "push" contact_pose: torch.Tensor - class Push(AtomicAction[PushGoal]): + @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, cfg: PushCfg | None = None) -> None: - super().__init__(cfg or PushCfg()) + def __init__(self, default_options: PushOptions | None = None) -> None: + super().__init__(default_options) def plan( self, - invocation: ActionInvocation[PushGoal], + request: ResolvedActionRequest[PushGoal, PushOptions], context: PlanningContext, ) -> ActionPlan: - goal = self.require_goal(invocation) + 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( - invocation, + request, context, success=success_mask, trajectory=full_robot_positions, diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index 103e64c09..f4ca7c71f 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -16,11 +16,13 @@ """Typed planning contracts and built-in atomic actions. -An action consumes an :class:`ActionInvocation` and 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. +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,8 +33,16 @@ AssembleAffordance, InteractionPoints, ) -from .bindings import ActionBinding -from .core import ActionCfg, AtomicAction, ObjectSemantics, SkillDescriptor +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, @@ -49,7 +59,7 @@ JointCommand, ) from .goals import ActionGoal, ObjectActionGoal, PoseGoalValue, SceneEntityPose -from .invocation import ActionInvocation +from .invocation import ActionInvocation, ActionOptions, ResolvedActionRequest from .plans import ( ActionPlan, CompiledTrajectory, @@ -66,30 +76,30 @@ AssembleGoal, CoordinatedPickGoal, CoordinatedPickment, - CoordinatedPickmentCfg, + CoordinatedPickmentOptions, CoordinatedPlacement, - CoordinatedPlacementCfg, CoordinatedPlacementGoal, + CoordinatedPlacementOptions, EndEffectorPoseGoal, GraspGoal, HandOver, - HandOverCfg, + HandOverOptions, HeldObjectPoseGoal, JointPositionGoal, MoveEndEffector, - MoveEndEffectorCfg, + MoveEndEffectorOptions, MoveHeldObject, - MoveHeldObjectCfg, + MoveHeldObjectOptions, MoveJoints, - MoveJointsCfg, + MoveJointsOptions, PickUp, - PickUpCfg, + PickUpOptions, Place, - PlaceCfg, PlaceGoal, + PlaceOptions, Press, - PressCfg, PressGoal, + PressOptions, ) from .state import ( CoordinatedHeldObjectState, @@ -104,9 +114,10 @@ __all__ = [ "ActionBinding", - "ActionCfg", + "ActionControlOverrides", "ActionGoal", "ActionInvocation", + "ActionOptions", "ActionPlan", "ActionPlanningServices", "Affordance", @@ -118,13 +129,15 @@ "CompiledTrajectory", "CompletionCondition", "CompletionConditionKind", + "ControlCommand", + "ControlPartCommandProfile", "CoordinatedHeldObjectState", "CoordinatedPickGoal", "CoordinatedPickment", - "CoordinatedPickmentCfg", + "CoordinatedPickmentOptions", "CoordinatedPlacement", - "CoordinatedPlacementCfg", "CoordinatedPlacementGoal", + "CoordinatedPlacementOptions", "EndEffectorPoseGoal", "EntityState", "ExecutionEvent", @@ -132,37 +145,43 @@ "ExecutionSession", "ExecutionStatus", "ExecutionTick", + "GRASP_COMMAND", "GraspGoal", "HandOver", - "HandOverCfg", + "HandOverOptions", "HeldObjectPoseGoal", "HeldObjectState", "InteractionPoints", "JointPositionGoal", "JointCommand", + "JointPositionCommand", "MotionPolicy", "MoveEndEffector", - "MoveEndEffectorCfg", + "MoveEndEffectorOptions", "MoveHeldObject", - "MoveHeldObjectCfg", + "MoveHeldObjectOptions", "MoveJoints", - "MoveJointsCfg", + "MoveJointsOptions", "ObjectActionGoal", "ObjectSemantics", + "OPEN_COMMAND", "PhaseSpec", "PickUp", - "PickUpCfg", + "PickUpOptions", "Place", - "PlaceCfg", "PlaceGoal", + "PlaceOptions", "PlannedPhase", "PlannerDiagnostics", "PlanningContext", "PoseGoalValue", "Press", - "PressCfg", "PressGoal", + "PressOptions", "RecoveryPolicy", + "ResolvedActionRequest", + "ResolvedActionBinding", + "ResolvedControlPart", "RobotObservation", "SceneSnapshot", "SceneEntityPose", diff --git a/embodichain/lab/sim/atomic_actions/bindings.py b/embodichain/lab/sim/atomic_actions/bindings.py index c17e536f0..5257c5035 100644 --- a/embodichain/lab/sim/atomic_actions/bindings.py +++ b/embodichain/lab/sim/atomic_actions/bindings.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Semantic-role to robot-resource bindings for atomic actions.""" +"""Semantic-role to robot control-part bindings for atomic actions.""" from __future__ import annotations @@ -22,6 +22,10 @@ from types import MappingProxyType from typing import Mapping +import torch + +from .control import ControlCommand, JointPositionCommand + def _normalize_resource_map( values: Mapping[str, str], @@ -43,18 +47,32 @@ def _normalize_resource_map( @dataclass(frozen=True, slots=True) class ActionBinding: - """Bind semantic action roles to embodiment-specific control resources. + """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``. - The action and an agent-facing request refer to roles such as ``primary``, - ``source`` and ``destination``. Only the compiler or application binding - layer needs to know concrete robot resources such as ``left_arm``. + 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 resources keyed by semantic role.""" + """Manipulator control-part names keyed by semantic role.""" end_effectors: Mapping[str, str] = field(default_factory=dict) - """End-effector resources keyed by semantic role.""" + """Tool or hand control-part names keyed by semantic role.""" def __post_init__(self) -> None: object.__setattr__( @@ -69,13 +87,13 @@ def __post_init__(self) -> None: ) def manipulator(self, role: str = "primary") -> str: - """Return the manipulator resource bound to ``role``. + """Return the manipulator control-part name bound to ``role``. Args: role: Semantic manipulator role. Returns: - Concrete robot control-resource name. + Key from the current robot's ``control_parts`` mapping. Raises: KeyError: If the requested role is not bound. @@ -86,13 +104,13 @@ def manipulator(self, role: str = "primary") -> str: raise KeyError(f"No manipulator is bound to role {role!r}.") from exc def end_effector(self, role: str = "primary") -> str: - """Return the end-effector resource bound to ``role``. + """Return the tool/hand control-part name bound to ``role``. Args: role: Semantic end-effector role. Returns: - Concrete robot control-resource name. + Key from the current robot's ``control_parts`` mapping. Raises: KeyError: If the requested role is not bound. @@ -103,4 +121,172 @@ def end_effector(self, role: str = "primary") -> str: raise KeyError(f"No end effector is bound to role {role!r}.") from exc -__all__ = ["ActionBinding"] +@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 70579f650..f9143d41e 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -19,18 +19,24 @@ from __future__ import annotations from abc import ABC, abstractmethod +from copy import deepcopy from dataclasses import dataclass, field 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, GoalT +from .invocation import ( + ActionInvocation, + ActionOptions, + GoalT, + OptionsT, + ResolvedActionRequest, +) from .plans import ( ActionPlan, CompletionCondition, @@ -102,6 +108,7 @@ class SkillDescriptor: 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 @@ -114,6 +121,12 @@ def __post_init__(self) -> None: ) 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( @@ -123,21 +136,10 @@ def __post_init__(self) -> None: object.__setattr__(self, field_name, roles) -@configclass -class ActionCfg: - """Base configuration for implementation-owned skill behavior.""" - - name: str = "default" - - def __post_init__(self) -> None: - if not isinstance(self.name, str) or not self.name: - raise ValueError("name must be a non-empty string.") - - -class AtomicAction(Generic[GoalT], ABC): +class AtomicAction(Generic[GoalT, OptionsT], ABC): """Side-effect-free planner for one semantically meaningful robot skill. - Actions own only skill configuration. An + 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. """ @@ -148,6 +150,9 @@ class AtomicAction(Generic[GoalT], ABC): GoalType: ClassVar[type[Any] | tuple[type[Any], ...]] """Concrete goal dataclass or dataclasses accepted by this skill.""" + OptionsType: ClassVar[type[ActionOptions]] = ActionOptions + """Concrete per-invocation runtime options accepted by this skill.""" + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) """Required semantic manipulator roles.""" @@ -159,11 +164,25 @@ class AtomicAction(Generic[GoalT], ABC): def __init__( self, - cfg: ActionCfg | None = None, + default_options: OptionsT | None = None, ) -> None: - self.cfg = cfg if cfg is not None else ActionCfg() + 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) + @property def is_bound(self) -> bool: """Whether an engine has supplied this action's planning resources.""" @@ -230,23 +249,27 @@ def descriptor(cls) -> SkillDescriptor: 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 require_goal(self, invocation: ActionInvocation[GoalT]) -> GoalT: - """Validate an invocation and return its concrete goal. + def resolve_request( + self, + invocation: ActionInvocation[GoalT, OptionsT], + ) -> ResolvedActionRequest[GoalT, OptionsT]: + """Validate and snapshot an invocation through engine-owned resources. Args: - invocation: Grounded invocation to validate. + invocation: Caller-owned invocation to resolve. Returns: - Invocation goal narrowed to this action's declared type. + Immutable request reused by planning and recovery replans. Raises: ValueError: If the stable skill identifier does not match. - TypeError: If the goal type is incompatible. + TypeError: If the goal or options type is incompatible. KeyError: If a required binding role is missing. """ if invocation.skill_id != self.skill_id: @@ -268,6 +291,16 @@ def require_goal(self, invocation: ActionInvocation[GoalT]) -> GoalT: 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: @@ -275,11 +308,49 @@ def require_goal(self, invocation: ActionInvocation[GoalT]) -> GoalT: f"Motion policy requires planner {required_planner!r}, but this " f"action uses {configured_planner_name!r}." ) - return invocation.goal + 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 def build_plan( self, - invocation: ActionInvocation[GoalT], + request: ResolvedActionRequest[GoalT, OptionsT], context: PlanningContext, *, success: bool | torch.Tensor, @@ -296,12 +367,12 @@ def build_plan( """Build a validated single-phase plan for a primitive implementation. Args: - invocation: Grounded invocation being planned. + 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 action config name. + 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. @@ -310,7 +381,7 @@ def build_plan( Returns: Side-effect-free action plan. """ - self.require_goal(invocation) + self.require_goal(request) if isinstance(success, bool): success_mask = torch.full( (context.batch_size,), @@ -337,7 +408,7 @@ def build_plan( timed = TimedTrajectory.from_positions( trajectory, env_ids=context.env_ids, - control_dt=invocation.motion_policy.control_dt, + control_dt=request.motion_policy.control_dt, ) elif isinstance(trajectory, TimedTrajectory): timed = trajectory @@ -354,15 +425,15 @@ def build_plan( ) phase = PlannedPhase( spec=PhaseSpec( - name=phase_name or self.cfg.name, - goal=invocation.goal, + name=phase_name or self.skill_id, + goal=request.goal, replannable=replannable, completion_condition=CompletionCondition( kind=completion_kind, tolerance=completion_tolerance, ), - recovery_policy=invocation.recovery_policy, - scene_dependencies=collect_scene_dependencies(invocation.goal), + recovery_policy=request.recovery_policy, + scene_dependencies=collect_scene_dependencies(request.goal), ), trajectory=timed, planned_scene_version=context.scene.version, @@ -373,12 +444,13 @@ def build_plan( plan_success=success_mask, phases=(phase,), expected_effects=expected_effects or StateDelta(), - invocation_id=invocation.invocation_id, + invocation_id=request.invocation_id, + invocation_revision=request.revision, ) def failed_plan( self, - invocation: ActionInvocation[GoalT], + request: ResolvedActionRequest[GoalT, OptionsT], context: PlanningContext, *, message: str | None = None, @@ -386,7 +458,7 @@ def failed_plan( """Build a failed empty plan without changing task state. Args: - invocation: Grounded invocation that failed to plan. + request: Resolved invocation that failed to plan. context: Planning input used for the attempt. message: Optional diagnostic message. @@ -394,7 +466,7 @@ def failed_plan( Failed action plan with an empty phase trajectory. """ return self.build_plan( - invocation, + request, context, success=torch.zeros( context.batch_size, dtype=torch.bool, device=self.device @@ -415,13 +487,13 @@ def failed_plan( @abstractmethod def plan( self, - invocation: ActionInvocation[GoalT], + request: ResolvedActionRequest[GoalT, OptionsT], context: PlanningContext, ) -> ActionPlan: """Plan one invocation without stepping simulation or committing state. Args: - invocation: Fully typed and embodiment-bound action request. + request: Immutable, typed, and embodiment-resolved action request. context: Latest observed robot, task, and scene state. Returns: @@ -430,7 +502,6 @@ def plan( __all__ = [ - "ActionCfg", "AtomicAction", "ObjectSemantics", "SkillDescriptor", diff --git a/embodichain/lab/sim/atomic_actions/engine.py b/embodichain/lab/sim/atomic_actions/engine.py index a97d9a258..f79ef2774 100644 --- a/embodichain/lab/sim/atomic_actions/engine.py +++ b/embodichain/lab/sim/atomic_actions/engine.py @@ -18,12 +18,13 @@ from __future__ import annotations -from typing import Iterable, TYPE_CHECKING +from typing import Iterable, Mapping, TYPE_CHECKING import torch from .core import AtomicAction -from .invocation import ActionInvocation +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 @@ -77,8 +78,15 @@ def get_registered_actions() -> dict[str, type[AtomicAction]]: class AtomicActionEngine: """Own planning resources and coordinate side-effect-free atomic actions.""" - def __init__(self, motion_generator: MotionGenerator) -> None: - self._planning_services = ActionPlanningServices(motion_generator) + def __init__( + self, + motion_generator: MotionGenerator, + control_profiles: Mapping[str, ControlPartCommandProfile] | None = None, + ) -> None: + self._planning_services = ActionPlanningServices( + motion_generator, + control_profiles=control_profiles, + ) self._actions: dict[str, AtomicAction] = {} @property @@ -101,6 +109,11 @@ 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 action instances keyed by stable skill identifier.""" @@ -157,8 +170,66 @@ def plan_action( raise TypeError("action must be an AtomicAction instance.") self._validate_context(context) action._bind(self._planning_services) - plan = action.plan(invocation, context) - self._validate_plan(plan, context, invocation) + 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. + + Args: + invocation: Grounded request for a registered skill. + + 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 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( @@ -178,13 +249,9 @@ def plan( 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}." - ) current = self.initial_context() if context is None else context - return self.plan_action(action, invocation, current) + request = self.resolve(invocation) + return self.plan_request(request, current) def initial_context( self, @@ -325,18 +392,22 @@ def _validate_plan( self, plan: ActionPlan, context: PlanningContext, - invocation: ActionInvocation, + request: ResolvedActionRequest, ) -> None: """Validate one action result before it is composed.""" - if plan.skill_id != invocation.skill_id: + if plan.skill_id != request.skill_id: raise ValueError( - "ActionPlan.skill_id must match its invocation, " - f"got {plan.skill_id!r} and {invocation.skill_id!r}." + "ActionPlan.skill_id must match its request, " + f"got {plan.skill_id!r} and {request.skill_id!r}." ) - if plan.invocation_id != invocation.invocation_id: + 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.") diff --git a/embodichain/lab/sim/atomic_actions/execution.py b/embodichain/lab/sim/atomic_actions/execution.py index e0ee840e8..748cebed8 100644 --- a/embodichain/lab/sim/atomic_actions/execution.py +++ b/embodichain/lab/sim/atomic_actions/execution.py @@ -24,7 +24,7 @@ import torch -from .invocation import ActionInvocation +from .invocation import ActionInvocation, ResolvedActionRequest from .plans import ActionPlan, PlannedPhase from .state import EntityState, PlanningContext, SceneSnapshot, TaskState @@ -44,6 +44,7 @@ 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" @@ -64,6 +65,7 @@ class ExecutionEvent: timestamp: float skill_id: str | None invocation_id: str | None + invocation_revision: int invocation_index: int env_mask: torch.Tensor message: str = "" @@ -73,6 +75,8 @@ def __post_init__(self) -> None: 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()) @@ -150,7 +154,9 @@ def __init__( raise ValueError("ExecutionSession requires at least one invocation.") engine._validate_context(context) self._engine = engine - self._invocations = invocations + self._requests: tuple[ResolvedActionRequest, ...] = tuple( + engine.resolve(invocation) for invocation in invocations + ) self._task_state = context.task self._context = context self._invocation_index = 0 @@ -194,6 +200,62 @@ 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, @@ -303,8 +365,17 @@ def _plan_current( event_kind: ExecutionEventKind, ) -> None: """Plan the current invocation from the latest observation.""" - invocation = self._invocations[self._invocation_index] - plan = self._engine.plan(invocation, context) + 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 @@ -495,7 +566,7 @@ def _finish_action( ) ) self._invocation_index += 1 - if self._invocation_index >= len(self._invocations): + if self._invocation_index >= len(self._requests): self._status = ( ExecutionStatus.COMPLETED if self._eligible.any() @@ -623,21 +694,27 @@ def _event( ) -> ExecutionEvent: """Create an event correlated with the current invocation.""" skill_id = ( - self._invocations[self._invocation_index].skill_id - if self._invocation_index < len(self._invocations) + self._requests[self._invocation_index].skill_id + if self._invocation_index < len(self._requests) else None ) invocation_id = ( - self._invocations[self._invocation_index].invocation_id - if self._invocation_index < len(self._invocations) + 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_index=min(self._invocation_index, len(self._invocations) - 1), + invocation_revision=invocation_revision, + invocation_index=min(self._invocation_index, len(self._requests) - 1), env_mask=env_mask, message=message, ) diff --git a/embodichain/lab/sim/atomic_actions/invocation.py b/embodichain/lab/sim/atomic_actions/invocation.py index 56ca04531..5e5ad06eb 100644 --- a/embodichain/lab/sim/atomic_actions/invocation.py +++ b/embodichain/lab/sim/atomic_actions/invocation.py @@ -18,18 +18,33 @@ from __future__ import annotations +from copy import deepcopy from dataclasses import dataclass, field from typing import Generic, TypeVar -from .bindings import ActionBinding +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]): +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. @@ -44,7 +59,7 @@ class ActionInvocation(Generic[GoalT]): """Action-specific goal value object.""" binding: ActionBinding - """Semantic-role bindings for the selected robot embodiment.""" + """Semantic-role bindings to keys in the selected robot's control parts.""" motion_policy: MotionPolicy = field(default_factory=MotionPolicy) """Reusable motion-generation settings.""" @@ -52,9 +67,20 @@ class ActionInvocation(Generic[GoalT]): 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.") @@ -70,10 +96,64 @@ def __post_init__(self) -> None: 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.") -__all__ = ["ActionInvocation", "GoalT"] +@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 index 5da8dda35..f1f9dc3e6 100644 --- a/embodichain/lab/sim/atomic_actions/plans.py +++ b/embodichain/lab/sim/atomic_actions/plans.py @@ -397,10 +397,16 @@ class ActionPlan: 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: diff --git a/embodichain/lab/sim/atomic_actions/policies.py b/embodichain/lab/sim/atomic_actions/policies.py index 1adebcb4f..8d56ab333 100644 --- a/embodichain/lab/sim/atomic_actions/policies.py +++ b/embodichain/lab/sim/atomic_actions/policies.py @@ -18,17 +18,22 @@ from __future__ import annotations +from copy import deepcopy +from dataclasses import dataclass from typing import TYPE_CHECKING -from embodichain.utils import configclass - if TYPE_CHECKING: from embodichain.lab.sim.planners import PlanOptions -@configclass +@dataclass(frozen=True, slots=True) class MotionPolicy: - """Reusable motion-generation policy supplied with an action invocation.""" + """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.""" @@ -77,9 +82,10 @@ def __post_init__(self) -> None: 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)) -@configclass +@dataclass(frozen=True, slots=True) class RecoveryPolicy: """Bounded local recovery policy used by the execution runtime.""" diff --git a/embodichain/lab/sim/atomic_actions/primitives/__init__.py b/embodichain/lab/sim/atomic_actions/primitives/__init__.py index e34a7b20a..9b85035e8 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/__init__.py +++ b/embodichain/lab/sim/atomic_actions/primitives/__init__.py @@ -21,59 +21,55 @@ from .coordinated_pickment import ( CoordinatedPickGoal, CoordinatedPickment, - CoordinatedPickmentCfg, + CoordinatedPickmentOptions, ) from .coordinated_placement import ( CoordinatedPlacement, - CoordinatedPlacementCfg, CoordinatedPlacementGoal, + CoordinatedPlacementOptions, ) -from .hand_over import HandOver, HandOverCfg +from .hand_over import HandOver, HandOverOptions from .move_end_effector import ( EndEffectorPoseGoal, MoveEndEffector, - MoveEndEffectorCfg, + MoveEndEffectorOptions, ) from .move_held_object import ( HeldObjectPoseGoal, MoveHeldObject, - MoveHeldObjectCfg, + MoveHeldObjectOptions, ) -from .move_joints import ( - JointPositionGoal, - MoveJoints, - MoveJointsCfg, -) -from .pick_up import GraspGoal, PickUp, PickUpCfg -from .place import AssembleGoal, Place, PlaceCfg, PlaceGoal -from .press import Press, PressCfg, PressGoal +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 __all__ = [ "AssembleGoal", "CoordinatedPickGoal", "CoordinatedPickment", - "CoordinatedPickmentCfg", + "CoordinatedPickmentOptions", "CoordinatedPlacement", - "CoordinatedPlacementCfg", "CoordinatedPlacementGoal", + "CoordinatedPlacementOptions", "EndEffectorPoseGoal", "GraspGoal", "HandOver", - "HandOverCfg", + "HandOverOptions", "HeldObjectPoseGoal", "JointPositionGoal", "MoveEndEffector", - "MoveEndEffectorCfg", + "MoveEndEffectorOptions", "MoveHeldObject", - "MoveHeldObjectCfg", + "MoveHeldObjectOptions", "MoveJoints", - "MoveJointsCfg", + "MoveJointsOptions", "PickUp", - "PickUpCfg", + "PickUpOptions", "Place", - "PlaceCfg", "PlaceGoal", + "PlaceOptions", "Press", - "PressCfg", "PressGoal", + "PressOptions", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py index dc7c76a22..f8765a595 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py @@ -23,13 +23,12 @@ 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, - AtomicAction, -) +from ..bindings import ResolvedControlPart +from ..control import GRASP_COMMAND, OPEN_COMMAND +from ..core import AtomicAction, ObjectSemantics from ..effects import StateDelta from ..goals import ( ObjectActionGoal, @@ -38,7 +37,7 @@ validate_pose_goal, validate_pose_tensor, ) -from ..invocation import ActionInvocation +from ..invocation import ActionOptions, ResolvedActionRequest from ..plans import ActionPlan from ..state import CoordinatedHeldObjectState, PlanningContext @@ -86,37 +85,9 @@ def __post_init__(self) -> None: ) -@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.""" @@ -133,66 +104,34 @@ class CoordinatedPickmentCfg(ActionCfg): 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.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 + + +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) @@ -222,13 +161,12 @@ def _resolve_pose(self, pose: torch.Tensor, name: str) -> torch.Tensor: def _resolve_dual_arm_start( self, 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( @@ -271,21 +209,6 @@ 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: PlanningContext, @@ -293,6 +216,8 @@ def _assemble_phase( 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( @@ -301,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 @@ -419,23 +343,23 @@ def _interpolate_object_pose( return poses -class CoordinatedPickment(AtomicAction[CoordinatedPickGoal]): +class CoordinatedPickment( + AtomicAction[CoordinatedPickGoal, CoordinatedPickmentOptions] +): """Pick and move a single object pinched by two hands.""" 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 - _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 @@ -443,59 +367,65 @@ class CoordinatedPickment(AtomicAction[CoordinatedPickGoal]): def __init__( self, - cfg: CoordinatedPickmentCfg | None = None, + default_options: CoordinatedPickmentOptions | None = None, ) -> None: - super().__init__(cfg or CoordinatedPickmentCfg()) - self._validate_hand_qpos_cfg() + super().__init__(default_options) def _on_bind(self) -> None: - """Resolve robot-dependent resources from the owning engine.""" - 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 - - assert self.cfg.left_hand_open_qpos is not None - assert self.cfg.left_hand_close_qpos is not None - assert self.cfg.right_hand_open_qpos is not None - assert self.cfg.right_hand_close_qpos is not None - 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, - ) + """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." + ) + 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, @@ -572,10 +502,12 @@ def _resolve_target( held_state, ) - def _compute_segment_lengths(self, sample_count: int) -> dict[str, int]: + 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, self.cfg.hand_interp_steps) - n_hold = max(0, self.cfg.hold_steps) + 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 @@ -594,18 +526,28 @@ def _compute_segment_lengths(self, sample_count: int) -> dict[str, int]: "hold": n_hold, } - def get_segment_lengths(self, sample_count: int) -> dict[str, int]: + 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) + 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, @@ -696,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, @@ -719,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, ) @@ -763,32 +707,15 @@ def _plan_synchronized_object_motion( def plan( self, - invocation: ActionInvocation[CoordinatedPickGoal], + request: ResolvedActionRequest[CoordinatedPickGoal, CoordinatedPickmentOptions], context: PlanningContext, ) -> ActionPlan: """Plan a coordinated pick without committing the dual attachment.""" - target = self.require_goal(invocation) - bindings = ( - (invocation.binding.manipulator("left"), self.cfg.left_arm_control_part), - ( - invocation.binding.manipulator("right"), - self.cfg.right_arm_control_part, - ), - ( - invocation.binding.end_effector("left"), - self.cfg.left_hand_control_part, - ), - ( - invocation.binding.end_effector("right"), - self.cfg.right_hand_control_part, - ), - ) - if any(actual != expected for actual, expected in bindings): - raise ValueError( - "CoordinatedPickment bindings do not match its configured resources." - ) + target = self.require_goal(request) + options = request.skill_options + resources = self._resolve_resources(request) if ( - invocation.motion_policy.motion_source == "motion_gen" + request.motion_policy.motion_source == "motion_gen" and self.motion_generator.planner.cfg.planner_type == "curobo" ): raise ValueError( @@ -804,10 +731,14 @@ def plan( right_target_xpos, held_state, ) = self._resolve_target(target, context) - left_start_qpos, right_start_qpos = self._resolve_dual_arm_start(state) - segments = self._compute_segment_lengths(invocation.motion_policy.sample_count) - left_pre_grasp_xpos = self._compute_pre_grasp_xpos(left_grasp_xpos) - right_pre_grasp_xpos = self._compute_pre_grasp_xpos(right_grasp_xpos) + left_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 ) @@ -816,14 +747,14 @@ def plan( ) 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"], @@ -836,8 +767,9 @@ def plan( 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( @@ -845,20 +777,21 @@ def plan( 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, @@ -874,6 +807,8 @@ def plan( held_state.left_object_to_eef, held_state.right_object_to_eef, success_mask, + resources, + options, ) ) @@ -883,8 +818,9 @@ def plan( 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( @@ -901,6 +837,8 @@ def plan( held_state.left_object_to_eef, held_state.right_object_to_eef, success_mask, + resources, + options, ) ) @@ -910,8 +848,9 @@ def plan( 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( @@ -922,8 +861,9 @@ def plan( 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( @@ -949,19 +889,19 @@ def plan( right_grasp_xpos=right_target_xpos, ) return self.build_plan( - invocation, + request, context, success=success_mask, trajectory=full, expected_effects=StateDelta( held_object_updates={ - self.cfg.left_arm_control_part: None, - self.cfg.right_arm_control_part: None, + resources.left_arm.name: None, + resources.right_arm.name: None, }, coordinated_held_object_updates={ ( - self.cfg.left_arm_control_part, - self.cfg.right_arm_control_part, + resources.left_arm.name, + resources.right_arm.name, ): coordinated_held_object, }, ), @@ -971,5 +911,5 @@ def plan( __all__ = [ "CoordinatedPickGoal", "CoordinatedPickment", - "CoordinatedPickmentCfg", + "CoordinatedPickmentOptions", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py index 768c0bf69..1c0a0dcd2 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py @@ -24,16 +24,15 @@ 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 ( - ActionCfg, - AtomicAction, -) +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 ActionInvocation +from ..invocation import ActionOptions, ResolvedActionRequest from ..plans import ActionPlan from ..policies import MotionPolicy from ..state import HeldObjectState, PlanningContext @@ -58,7 +57,7 @@ class CoordinatedPlacementGoal: """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_goal( @@ -73,34 +72,9 @@ def __post_init__(self) -> None: ) -@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.""" @@ -123,100 +97,109 @@ 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[CoordinatedPlacementGoal]): + +class CoordinatedPlacement( + AtomicAction[CoordinatedPlacementGoal, CoordinatedPlacementOptions] +): """Coordinate two held objects: support object below, placing object above.""" 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, - cfg: CoordinatedPlacementCfg | None = None, + default_options: CoordinatedPlacementOptions | None = None, ) -> None: - super().__init__(cfg or CoordinatedPlacementCfg()) - self._validate_hand_qpos_cfg() + super().__init__(default_options) def _on_bind(self) -> None: - """Resolve robot-dependent resources from the owning engine.""" + """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) - - assert self.cfg.placing_hand_open_qpos is not None - assert self.cfg.placing_hand_close_qpos is not None - assert self.cfg.support_hand_close_qpos is not None - 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 plan( self, - invocation: ActionInvocation[CoordinatedPlacementGoal], + request: ResolvedActionRequest[ + CoordinatedPlacementGoal, CoordinatedPlacementOptions + ], context: PlanningContext, ) -> ActionPlan: """Plan coordinated placement without committing attachment changes.""" - target = self.require_goal(invocation) - bindings = ( - ( - invocation.binding.manipulator("placing"), - self.cfg.placing_arm_control_part, - ), - ( - invocation.binding.manipulator("support"), - self.cfg.support_arm_control_part, - ), - ( - invocation.binding.end_effector("placing"), - self.cfg.placing_hand_control_part, - ), - ( - invocation.binding.end_effector("support"), - self.cfg.support_hand_control_part, - ), - ) - if any(actual != expected for actual, expected in bindings): - raise ValueError( - "CoordinatedPlacement bindings do not match its configured resources." - ) + target = self.require_goal(request) + options = request.skill_options + resources = self._resolve_resources(request) if ( - invocation.motion_policy.motion_source == "motion_gen" + request.motion_policy.motion_source == "motion_gen" and self.motion_generator.planner.cfg.planner_type == "curobo" ): raise ValueError( @@ -229,45 +212,47 @@ def plan( release, placing_held_object, support_held_object, - ) = self._resolve_target(target, state) - placing_start_qpos, support_start_qpos = self._resolve_start_qpos(state) + ) = 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, invocation.motion_policy.sample_count + 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"], - invocation.motion_policy, + request.motion_policy, ) if not ok: logger.log_warning("CoordinatedPlacement failed to plan placing approach.") return self.failed_plan( - invocation, context, message="Placing approach failed." + 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"], - invocation.motion_policy, + request.motion_policy, ) if not ok: logger.log_warning("CoordinatedPlacement failed to plan support approach.") return self.failed_plan( - invocation, context, message="Support approach failed." + request, context, message="Support approach failed." ) placing_place_qpos = placing_approach_traj[:, -1] @@ -276,8 +261,9 @@ def plan( 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() @@ -286,8 +272,9 @@ def plan( 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() @@ -297,35 +284,39 @@ def plan( 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"], - invocation.motion_policy, + request.motion_policy, ) if not ok: logger.log_warning("CoordinatedPlacement failed to plan placing retreat.") - return self.failed_plan( - invocation, context, message="Placing retreat failed." - ) + 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( @@ -338,8 +329,8 @@ def plan( dim=1, ) involved_control_parts = { - self.cfg.placing_arm_control_part, - self.cfg.support_arm_control_part, + resources.placing_arm.name, + resources.support_arm.name, } coordinated_removals = { key: None @@ -347,34 +338,21 @@ def plan( if not involved_control_parts.isdisjoint(key) } return self.build_plan( - invocation, + request, context, success=True, trajectory=full, expected_effects=StateDelta( held_object_updates={ - self.cfg.placing_arm_control_part: ( + resources.placing_arm.name: ( None if release else placing_held_object ), - self.cfg.support_arm_control_part: support_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: PoseGoalValue, @@ -438,6 +416,8 @@ def _resolve_target( self, target: CoordinatedPlacementGoal, state: PlanningContext, + resources: _CoordinatedPlacementResources, + options: CoordinatedPlacementOptions, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -445,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 ) @@ -491,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, @@ -509,7 +491,9 @@ def _resolve_target( ) def _resolve_start_qpos( - self, state: PlanningContext + self, + state: PlanningContext, + resources: _CoordinatedPlacementResources, ) -> tuple[torch.Tensor, torch.Tensor]: if state.last_qpos.shape != (self.n_envs, self.robot_dof): logger.log_error( @@ -520,19 +504,20 @@ def _resolve_start_qpos( ) 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, sample_count: int, + options: CoordinatedPlacementOptions, ) -> dict[str, int]: """Split the invocation sample budget across placement phases.""" - 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_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( @@ -590,19 +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 __all__ = [ "CoordinatedPlacement", - "CoordinatedPlacementCfg", "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 aaca0495d..551f84732 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/hand_over.py +++ b/embodichain/lab/sim/atomic_actions/primitives/hand_over.py @@ -18,58 +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, - AtomicAction, - ObjectSemantics, -) +from ..bindings import ResolvedControlPart +from ..control import GRASP_COMMAND, OPEN_COMMAND +from ..core import AtomicAction, ObjectSemantics from ..effects import StateDelta -from ..invocation import ActionInvocation +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 @@ -107,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.""" -class HandOver(AtomicAction[GraspGoal]): + 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[GraspGoal, HandOverOptions]): """Hand an object from one arm to the other. The transferring arm (already holding the object) moves it to a middle @@ -119,80 +130,71 @@ class HandOver(AtomicAction[GraspGoal]): 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, - cfg: HandOverCfg | None = None, + default_options: HandOverOptions | None = None, ) -> None: - super().__init__(cfg or HandOverCfg()) - self._validate_pose_cfg() - self._validate_hand_qpos_cfg() + super().__init__(default_options) def _on_bind(self) -> None: - """Resolve robot-dependent resources from the owning engine.""" + """Resolve engine-wide resources from the owning engine.""" self.n_envs = self.robot.get_qpos().shape[0] self.robot_dof = self.robot.dof - assert self.cfg.middle_object_pose is not None - assert self.cfg.final_object_pose is not None - 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) - - assert self.cfg.transfer_hand_open_qpos is not None - assert self.cfg.transfer_hand_close_qpos is not None - assert self.cfg.receive_hand_open_qpos is not None - assert self.cfg.receive_hand_close_qpos is not None - 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 @@ -200,33 +202,16 @@ def _on_bind(self) -> None: def plan( self, - invocation: ActionInvocation[GraspGoal], + request: ResolvedActionRequest[GraspGoal, HandOverOptions], context: PlanningContext, ) -> ActionPlan: """Plan a handover without committing the attachment transfer.""" - target = self.require_goal(invocation) - expected_bindings = ( - ( - invocation.binding.manipulator("source"), - self.cfg.transfer_arm_control_part, - ), - ( - invocation.binding.manipulator("destination"), - self.cfg.receive_arm_control_part, - ), - ( - invocation.binding.end_effector("source"), - self.cfg.transfer_hand_control_part, - ), - ( - invocation.binding.end_effector("destination"), - self.cfg.receive_hand_control_part, - ), - ) - if any(actual != expected for actual, expected in expected_bindings): - raise ValueError("HandOver bindings do not match its configured resources.") + target = self.require_goal(request) + options = request.skill_options + self._validate_pose_options(options) + resources = self._resolve_resources(request) if ( - invocation.motion_policy.motion_source == "motion_gen" + request.motion_policy.motion_source == "motion_gen" and self.motion_generator.planner.cfg.planner_type == "curobo" ): raise ValueError( @@ -234,9 +219,24 @@ def plan( ) 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.clone() + 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] @@ -247,18 +247,21 @@ def plan( # 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.failed_plan(invocation, context, message="No receiving grasp.") + 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) @@ -266,68 +269,70 @@ def plan( 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(invocation.motion_policy.sample_count) + 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"], - invocation.motion_policy, + request.motion_policy, ) if not ok: logger.log_warning("HandOver failed to plan the transfer move.") - return self.failed_plan( - invocation, context, message="Transfer move failed." - ) + 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"], - invocation.motion_policy, + request.motion_policy, ) if not ok: logger.log_warning("HandOver failed to plan the receiving approach.") return self.failed_plan( - invocation, context, message="Receiving approach failed." + 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"], - invocation.motion_policy, + request.motion_policy, ) if not ok: logger.log_warning("HandOver failed to plan the transfer retreat.") return self.failed_plan( - invocation, context, message="Transfer retreat failed." + 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"], - invocation.motion_policy, + request.motion_policy, ) if not ok: logger.log_warning("HandOver failed to plan the receiving delivery.") return self.failed_plan( - invocation, context, message="Receiving delivery failed." + request, context, message="Receiving delivery failed." ) phases: list[torch.Tensor] = [] @@ -337,8 +342,13 @@ def plan( 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. @@ -347,8 +357,13 @@ def plan( 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. @@ -357,12 +372,15 @@ def plan( 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: @@ -371,8 +389,13 @@ def plan( 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. @@ -382,11 +405,14 @@ def plan( 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. @@ -395,8 +421,13 @@ def plan( 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) @@ -406,14 +437,14 @@ def plan( grasp_xpos=receive_grasp_xpos, ) return self.build_plan( - invocation, + request, context, success=True, trajectory=full, expected_effects=StateDelta( held_object_updates={ - self.cfg.transfer_arm_control_part: None, - self.cfg.receive_arm_control_part: held_object, + resources.transfer_arm.name: None, + resources.receive_arm.name: held_object, } ), ) @@ -422,21 +453,13 @@ def plan( # 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) @@ -450,12 +473,16 @@ def _resolve_matrix(self, matrix: torch.Tensor, name: str) -> torch.Tensor: ) return matrix - def _resolve_transfer_object_to_eef(self, state: PlanningContext) -> 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") @@ -465,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] @@ -493,7 +521,9 @@ def _resolve_receive_grasp( return grasp_xpos, is_success def _resolve_start_qpos( - self, state: PlanningContext + self, + state: PlanningContext, + resources: _HandOverResources, ) -> tuple[torch.Tensor, torch.Tensor]: if state.last_qpos.shape != (self.n_envs, self.robot_dof): logger.log_error( @@ -503,16 +533,18 @@ def _resolve_start_qpos( ) 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], + start_qpos[:, list(resources.transfer_arm.joint_ids)], + start_qpos[:, list(resources.receive_arm.joint_ids)], ) - def _compute_segment_lengths(self, sample_count: int) -> dict[str, int]: + 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, 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) + 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, (sample_count - reserved) // 2) n_approach = sample_count - reserved - n_transfer @@ -571,6 +603,8 @@ def _assemble_phase( 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 = ( @@ -579,11 +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 -__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 d35f1ee09..ec8a40f87 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py @@ -24,11 +24,9 @@ import torch from embodichain.lab.sim.planners import MoveType, PlanState -from embodichain.utils import configclass - -from ..core import ActionCfg, AtomicAction +from ..core import AtomicAction from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal -from ..invocation import ActionInvocation +from ..invocation import ActionOptions, ResolvedActionRequest from ..plans import ActionPlan, CompletionConditionKind from ..state import PlanningContext @@ -46,36 +44,36 @@ def __post_init__(self) -> None: validate_pose_goal(self.xpos, "xpos", allow_waypoints=True) -@configclass -class MoveEndEffectorCfg(ActionCfg): - """Skill-specific MoveEndEffector configuration.""" - - name: str = "move_end_effector" +@dataclass(frozen=True, slots=True, eq=False) +class MoveEndEffectorOptions(ActionOptions): + """Per-invocation behavior for :class:`MoveEndEffector`.""" -class MoveEndEffector(AtomicAction[EndEffectorPoseGoal]): +class MoveEndEffector(AtomicAction[EndEffectorPoseGoal, MoveEndEffectorOptions]): """Plan a free-space move for a bound manipulator.""" skill_id: ClassVar[str] = "move_end_effector" GoalType: ClassVar[type] = EndEffectorPoseGoal + OptionsType: ClassVar[type] = MoveEndEffectorOptions manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) def __init__( self, - cfg: MoveEndEffectorCfg | None = None, + default_options: MoveEndEffectorOptions | None = None, ) -> None: - super().__init__(cfg or MoveEndEffectorCfg()) + super().__init__(default_options) def plan( self, - invocation: ActionInvocation[EndEffectorPoseGoal], + request: ResolvedActionRequest[EndEffectorPoseGoal, MoveEndEffectorOptions], context: PlanningContext, ) -> ActionPlan: """Plan an end-effector pose goal from the observed joint state.""" - goal = self.require_goal(invocation) - control_part = invocation.binding.manipulator("primary") - joint_ids = self.robot.get_joint_ids(name=control_part) - arm_dof = len(joint_ids) + 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, @@ -90,20 +88,20 @@ def plan( result = self.builder.generate_arm_plan( target_states, start_qpos, - invocation.motion_policy.sample_count, + request.motion_policy.sample_count, control_part=control_part, arm_dof=arm_dof, - cfg=invocation.motion_policy, + 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=invocation.motion_policy.control_dt, + control_dt=request.motion_policy.control_dt, ) return self.build_plan( - invocation, + request, context, success=success, trajectory=trajectory, @@ -127,4 +125,8 @@ def _build_target_states( ] -__all__ = ["EndEffectorPoseGoal", "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 c0a94f2f6..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,16 +24,14 @@ 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 ( - ActionCfg, - AtomicAction, -) +from ..control import GRASP_COMMAND +from ..core import AtomicAction from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal -from ..invocation import ActionInvocation +from ..invocation import ActionOptions, ResolvedActionRequest from ..plans import ActionPlan from ..state import PlanningContext @@ -55,19 +53,9 @@ def __post_init__(self) -> None: ) -@configclass -class MoveHeldObjectCfg(ActionCfg): - name: str = "move_held_object" - """Name of the action, used for identification and logging.""" - - control_part: str = "arm" - """Manipulator resource used by this configured action instance.""" - - 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.""" @@ -75,56 +63,66 @@ 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[HeldObjectPoseGoal]): +class MoveHeldObject(AtomicAction[HeldObjectPoseGoal, MoveHeldObjectOptions]): """Move the held object to a target object pose; keep the gripper closed.""" 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, - cfg: MoveHeldObjectCfg | None = None, + default_options: MoveHeldObjectOptions | None = None, ) -> None: - super().__init__(cfg or MoveHeldObjectCfg()) - if self.cfg.hand_close_qpos is None: - logger.log_error( - "hand_close_qpos must be specified in MoveHeldObjectCfg", ValueError - ) + super().__init__(default_options) def _on_bind(self) -> None: - """Resolve robot-dependent resources from the owning engine.""" + """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 - assert self.cfg.hand_close_qpos is not None - self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) def plan( self, - invocation: ActionInvocation[HeldObjectPoseGoal], + request: ResolvedActionRequest[HeldObjectPoseGoal, MoveHeldObjectOptions], context: PlanningContext, ) -> ActionPlan: """Plan held-object transport without changing the attachment relation.""" - target = self.require_goal(invocation) - if invocation.binding.manipulator() != self.cfg.control_part: - raise ValueError( - "MoveHeldObject manipulator binding does not match its config." - ) - if invocation.binding.end_effector() != self.cfg.hand_control_part: - raise ValueError( - "MoveHeldObject end-effector binding does not match its config." - ) + 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(self.cfg.control_part) + 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( @@ -137,19 +135,20 @@ def plan( 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 @@ -158,7 +157,7 @@ def plan( 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 = [ @@ -168,10 +167,10 @@ def plan( success, arm_traj = self.builder.plan_arm_traj( target_states_list, start_arm_qpos, - invocation.motion_policy.sample_count, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=invocation.motion_policy, + request.motion_policy.sample_count, + control_part=control_part, + arm_dof=manipulator.dof, + cfg=request.motion_policy, ) full = torch.empty( @@ -180,11 +179,11 @@ def plan( 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 self.build_plan( - invocation, + request, context, success=success, trajectory=full, @@ -196,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] ) @@ -271,4 +271,8 @@ def _apply_automatic_transport_rotation( ) -__all__ = ["HeldObjectPoseGoal", "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 577888617..1397265a0 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py @@ -23,10 +23,8 @@ import torch -from embodichain.utils import configclass, logger - -from ..core import ActionCfg, AtomicAction -from ..invocation import ActionInvocation +from ..core import AtomicAction +from ..invocation import ActionOptions, ResolvedActionRequest from ..plans import ActionPlan, CompletionConditionKind from ..state import PlanningContext @@ -38,7 +36,7 @@ class JointPositionGoal: goal_kind: ClassVar[str] = "joint_position" target: torch.Tensor | str - """Joint qpos/waypoints or a name in ``MoveJointsCfg.named_joint_positions``.""" + """Joint qpos/waypoints or a named control-part profile command.""" def __post_init__(self) -> None: if isinstance(self.target, str): @@ -59,42 +57,43 @@ def __post_init__(self) -> None: ) -@configclass -class MoveJointsCfg(ActionCfg): - """Skill-specific MoveJoints configuration.""" - - name: str = "move_joints" - named_joint_positions: dict[str, torch.Tensor] | None = None - """Optional named joint-position targets. Motion settings belong to ``MotionPolicy``.""" +@dataclass(frozen=True, slots=True, eq=False) +class MoveJointsOptions(ActionOptions): + """Per-invocation behavior for :class:`MoveJoints`.""" -class MoveJoints(AtomicAction[JointPositionGoal]): +class MoveJoints(AtomicAction[JointPositionGoal, MoveJointsOptions]): """Plan joint motion from the observed state to one or more waypoints.""" 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, - cfg: MoveJointsCfg | None = None, + default_options: MoveJointsOptions | None = None, ) -> None: - super().__init__(cfg or MoveJointsCfg()) - self.named_joint_positions = self.cfg.named_joint_positions or {} + super().__init__(default_options) def plan( self, - invocation: ActionInvocation[JointPositionGoal], + request: ResolvedActionRequest[JointPositionGoal, MoveJointsOptions], context: PlanningContext, ) -> ActionPlan: """Plan a joint-space goal without mutating the robot or task state.""" - goal = self.require_goal(invocation) - control_part = invocation.binding.manipulator("primary") - joint_ids = self.robot.get_joint_ids(name=control_part) - joint_dof = len(joint_ids) + 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(goal), + self._resolve_target_qpos( + goal, + request=request, + context=context, + ), n_envs=context.batch_size, joint_dof=joint_dof, control_part=control_part, @@ -108,20 +107,20 @@ def plan( result = self.builder.generate_joint_plan( start_qpos, target_qpos, - invocation.motion_policy.sample_count, + request.motion_policy.sample_count, control_part=control_part, arm_dof=joint_dof, - cfg=invocation.motion_policy, + 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=invocation.motion_policy.control_dt, + control_dt=request.motion_policy.control_dt, ) return self.build_plan( - invocation, + request, context, success=success, trajectory=trajectory, @@ -131,21 +130,23 @@ def plan( def _resolve_target_qpos( self, goal: JointPositionGoal, + *, + request: ResolvedActionRequest[JointPositionGoal, MoveJointsOptions], + context: PlanningContext, ) -> torch.Tensor: """Resolve an explicit or named joint goal to a tensor.""" if isinstance(goal.target, torch.Tensor): return goal.target - if goal.target not in self.named_joint_positions: - logger.log_error( - f"Unknown named joint-position goal {goal.target!r}. Available " - f"goals: {sorted(self.named_joint_positions)}", - KeyError, - ) - return self.named_joint_positions[goal.target] + return request.binding.manipulator("primary").joint_positions( + goal.target, + n_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) __all__ = [ "JointPositionGoal", "MoveJoints", - "MoveJointsCfg", + "MoveJointsOptions", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index 8d145b51d..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,14 +35,12 @@ from ._helpers import arm_qpos_from_state from ..affordance import AntipodalAffordance -from ..core import ( - ActionCfg, - AtomicAction, - ObjectSemantics, -) +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 ActionInvocation +from ..invocation import ActionOptions, ResolvedActionRequest from ..plans import ActionPlan from ..policies import MotionPolicy from ..state import HeldObjectState, PlanningContext @@ -68,26 +66,13 @@ def __post_init__(self) -> None: 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.""" - - control_part: str = "arm" - """Manipulator resource used by this configured action instance.""" +@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].""" @@ -112,54 +97,61 @@ 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[GraspGoal]): + +class PickUp(AtomicAction[GraspGoal, PickUpOptions]): """Approach a grasp pose, close the gripper, lift.""" 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, - cfg: PickUpCfg | None = None, + default_options: PickUpOptions | None = None, ) -> None: - super().__init__(cfg or PickUpCfg()) - 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 - ) - 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, - ) + super().__init__(default_options) def _on_bind(self) -> None: - """Resolve robot-dependent resources from the owning engine.""" + """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 - assert self.cfg.hand_open_qpos is not None - assert self.cfg.hand_close_qpos is not None - 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 def _get_full_pickup_trajectory( self, @@ -167,14 +159,20 @@ def _get_full_pickup_trajectory( 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( motion_policy.sample_count, - self.cfg.hand_interp_steps, + options.hand_interp_steps, first_phase_name="approach", third_phase_name="lift", ) @@ -190,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, + 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)] @@ -208,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, + 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] @@ -225,31 +223,52 @@ 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 plan( self, - invocation: ActionInvocation[GraspGoal], + request: ResolvedActionRequest[GraspGoal, PickUpOptions], context: PlanningContext, ) -> ActionPlan: """Plan approach, close, and lift phases without committing attachment.""" - target = self.require_goal(invocation) - if invocation.binding.manipulator() != self.cfg.control_part: - raise ValueError("PickUp manipulator binding does not match its config.") - if invocation.binding.end_effector() != self.cfg.hand_control_part: - raise ValueError("PickUp end-effector binding does not match its config.") + 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( @@ -264,31 +283,41 @@ def plan( "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: - grasp_xpos = self._upright_adjusted_grasp_poses(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.failed_plan( - invocation, context, message="Failed to resolve a grasp pose." + 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, - invocation.motion_policy, + 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) @@ -297,30 +326,33 @@ def plan( semantics=sem, object_to_eef=object_to_eef, grasp_xpos=grasp_xpos ) coordinated_updates = { - key: None - for key in state.coordinated_held_objects - if self.cfg.control_part in key + key: None for key in state.coordinated_held_objects if control_part in key } return self.build_plan( - invocation, + request, context, success=is_success, trajectory=full, expected_effects=StateDelta( - held_object_updates={self.cfg.control_part: held}, + held_object_updates={control_part: held}, coordinated_held_object_updates=coordinated_updates, ), 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) @@ -346,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) @@ -362,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] @@ -370,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 ) @@ -403,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 ) @@ -422,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()) @@ -438,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]) @@ -463,49 +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 _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) @@ -518,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) @@ -528,4 +582,4 @@ def _upright_adjusted_grasp_poses( return adjusted_grasp_xpos -__all__ = ["GraspGoal", "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 12fdf6c57..7bfa24253 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/place.py +++ b/embodichain/lab/sim/atomic_actions/primitives/place.py @@ -24,18 +24,16 @@ 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 ( - ActionCfg, - AtomicAction, -) +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 ActionInvocation +from ..invocation import ActionOptions, ResolvedActionRequest from ..plans import ActionPlan from ..state import PlanningContext @@ -89,26 +87,13 @@ class AssembleGoal: """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.""" - - control_part: str = "arm" - """Manipulator resource used by this configured action instance.""" +@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.""" @@ -118,8 +103,16 @@ 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[PlaceGoal | AssembleGoal]): +class Place(AtomicAction[PlaceGoal | AssembleGoal, PlaceOptions]): """Lower the held object to a place pose, open the gripper, retract. The :class:`PlaceGoal` may carry either a single waypoint @@ -141,78 +134,79 @@ class Place(AtomicAction[PlaceGoal | AssembleGoal]): PlaceGoal, AssembleGoal, ) + OptionsType: ClassVar[type] = PlaceOptions manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) def __init__( self, - cfg: PlaceCfg | None = None, + default_options: PlaceOptions | None = None, ) -> None: - super().__init__(cfg or PlaceCfg()) - 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 - ) - if self.cfg.cartesian_waypoint_count < 1: - logger.log_error("cartesian_waypoint_count must be at least 1.", ValueError) + super().__init__(default_options) def _on_bind(self) -> None: - """Resolve robot-dependent resources from the owning engine.""" + """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 - assert self.cfg.hand_open_qpos is not None - assert self.cfg.hand_close_qpos is not None - self.hand_open_qpos = self.cfg.hand_open_qpos.to(self.device) - self.hand_close_qpos = self.cfg.hand_close_qpos.to(self.device) def plan( self, - invocation: ActionInvocation[PlaceGoal | AssembleGoal], + request: ResolvedActionRequest[PlaceGoal | AssembleGoal, PlaceOptions], context: PlanningContext, ) -> ActionPlan: """Plan approach, release, and retract without committing detachment.""" - target = self.require_goal(invocation) - if invocation.binding.manipulator() != self.cfg.control_part: - raise ValueError("Place manipulator binding does not match its config.") - if invocation.binding.end_effector() != self.cfg.hand_control_part: - raise ValueError("Place end-effector binding does not match its config.") + 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) + 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, 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( - invocation.motion_policy.sample_count, - 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 = [ [ @@ -225,14 +219,14 @@ def plan( target_states_list, start_arm_qpos, n_down, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=invocation.motion_policy, + 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 = [ [ @@ -245,14 +239,14 @@ def plan( target_states_list, reach_arm_qpos, n_back, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=invocation.motion_policy, + 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 @@ -265,36 +259,35 @@ def plan( 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 + 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 self.cfg.control_part in key + key: None for key in state.coordinated_held_objects if control_part in key } return self.build_plan( - invocation, + request, context, success=success, trajectory=full, expected_effects=StateDelta( - held_object_updates={self.cfg.control_part: None}, + held_object_updates={control_part: None}, coordinated_held_object_updates=coordinated_updates, ), phase_name="place", ) def _resolve_place_xpos( - self, target: PlaceGoal | AssembleGoal, state: PlanningContext + self, + target: PlaceGoal | AssembleGoal, + state: PlanningContext, + control_part: str, ) -> torch.Tensor: """Resolve the place EEF poses from a typed target. @@ -311,10 +304,13 @@ def _resolve_place_xpos( resolve_pose_goal(target.xpos, state, name="xpos"), n_envs=self.n_envs, ) - return self._resolve_assemble_place_xpos(target, state) + return self._resolve_assemble_place_xpos(target, state, control_part) def _resolve_assemble_place_xpos( - self, target: AssembleGoal, state: PlanningContext + self, + target: AssembleGoal, + state: PlanningContext, + control_part: str, ) -> torch.Tensor: """Derive the place EEF pose from an assembly affordance. @@ -332,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 AssembleGoal requires an object held by control " - f"part {self.cfg.control_part!r} (run PickUp first).", + f"part {control_part!r} (run PickUp first).", ValueError, ) affordance = target.affordance @@ -358,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, ) @@ -376,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 @@ -402,7 +403,10 @@ def _translation_keyframes( return keyframes.flatten(1, 2) 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() @@ -412,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]) @@ -433,4 +437,4 @@ def _select_tcp_symmetric_place_variant( ] -__all__ = ["AssembleGoal", "Place", "PlaceCfg", "PlaceGoal"] +__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 d9a6c73d2..0bb193fb2 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/press.py +++ b/embodichain/lab/sim/atomic_actions/primitives/press.py @@ -24,15 +24,13 @@ 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 ( - ActionCfg, - AtomicAction, -) +from ..control import GRASP_COMMAND +from ..core import AtomicAction from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal -from ..invocation import ActionInvocation +from ..invocation import ActionOptions, ResolvedActionRequest from ..plans import ActionPlan from ..state import PlanningContext @@ -50,89 +48,78 @@ def __post_init__(self) -> None: 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.""" - - control_part: str = "arm" - """Manipulator resource used by this configured action instance.""" +@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[PressGoal]): +class Press(AtomicAction[PressGoal, PressOptions]): """Close the gripper, press down to a target pose, then return.""" 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, - cfg: PressCfg | None = None, + default_options: PressOptions | None = None, ) -> None: - super().__init__(cfg or PressCfg()) - if self.cfg.hand_close_qpos is None: - logger.log_error( - "hand_close_qpos must be specified in PressCfg", ValueError - ) + super().__init__(default_options) def _on_bind(self) -> None: - """Resolve robot-dependent resources from the owning engine.""" + """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 - assert self.cfg.hand_close_qpos is not None - self.hand_close_qpos = self.builder.expand_hand_qpos( - self.cfg.hand_close_qpos, - n_envs=self.n_envs, - hand_dof=self.hand_dof, - ) - def plan( self, - invocation: ActionInvocation[PressGoal], + request: ResolvedActionRequest[PressGoal, PressOptions], context: PlanningContext, ) -> ActionPlan: """Plan a close, press, and retract sequence.""" - target = self.require_goal(invocation) - if invocation.binding.manipulator() != self.cfg.control_part: - raise ValueError("Press manipulator binding does not match its config.") - if invocation.binding.end_effector() != self.cfg.hand_control_part: - raise ValueError("Press end-effector binding does not match its config.") + 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, ) 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( - invocation.motion_policy.sample_count + 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, ) @@ -144,9 +131,9 @@ def plan( target_states_list, start_arm_qpos, n_down, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=invocation.motion_policy, + control_part=control_part, + arm_dof=manipulator.dof, + cfg=request.motion_policy, ) press_arm_qpos = down_arm[:, -1, :] @@ -154,9 +141,9 @@ def plan( press_arm_qpos, start_arm_qpos, n_back, - control_part=self.cfg.control_part, - arm_dof=self.arm_dof, - cfg=invocation.motion_policy, + control_part=control_part, + arm_dof=manipulator.dof, + cfg=request.motion_policy, ) success = down_success & back_success @@ -170,32 +157,30 @@ def plan( 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 self.build_plan( - invocation, + request, context, success=success, trajectory=full, phase_name="press", ) - def _compute_phase_waypoints(self, sample_count: int) -> tuple[int, int, int]: + def _compute_phase_waypoints( + self, sample_count: int, options: PressOptions + ) -> tuple[int, int, int]: """Split the invocation sample budget across press phases.""" - 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 - ) + n_close = options.hand_interp_steps motion_waypoints = sample_count - n_close n_down = motion_waypoints // 2 @@ -209,4 +194,4 @@ def _compute_phase_waypoints(self, sample_count: int) -> tuple[int, int, int]: return n_close, n_down, n_back -__all__ = ["Press", "PressCfg", "PressGoal"] +__all__ = ["Press", "PressGoal", "PressOptions"] diff --git a/embodichain/lab/sim/atomic_actions/runtime.py b/embodichain/lab/sim/atomic_actions/runtime.py index 426613ff2..608fc1416 100644 --- a/embodichain/lab/sim/atomic_actions/runtime.py +++ b/embodichain/lab/sim/atomic_actions/runtime.py @@ -18,10 +18,18 @@ 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 @@ -40,13 +48,26 @@ class ActionPlanningServices: 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) -> None: + 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: @@ -68,6 +89,16 @@ 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.""" @@ -77,5 +108,163 @@ def planner_name(self) -> str: 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/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 64f16e692..75e6a890b 100644 --- a/examples/sim/planners/curobo_planner.py +++ b/examples/sim/planners/curobo_planner.py @@ -64,7 +64,6 @@ AtomicActionEngine, EndEffectorPoseGoal, MoveEndEffector, - MoveEndEffectorCfg, MotionPolicy, ) from embodichain.data import get_data_path @@ -750,7 +749,7 @@ def main() -> None: ) ) engine = AtomicActionEngine(motion_generator) - engine.register(MoveEndEffector(MoveEndEffectorCfg())) + engine.register(MoveEndEffector()) binding = ActionBinding(manipulators={"primary": control_part}) motion_policy = MotionPolicy( motion_source="motion_gen", diff --git a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py index 6b079b572..67635eb54 100644 --- a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py +++ b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py @@ -234,7 +234,6 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: 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 @@ -264,7 +263,7 @@ 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(cfg=MoveEndEffectorCfg())) + atomic_engine.register(MoveEndEffector()) 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 797975e0b..ae722d520 100644 --- a/scripts/benchmark/atomic_action/move_held_object_benchmark.py +++ b/scripts/benchmark/atomic_action/move_held_object_benchmark.py @@ -178,13 +178,13 @@ def _prepare_held_state( ActionBinding, ActionInvocation, AtomicActionEngine, + ControlPartCommandProfile, EndEffectorPoseGoal, GraspGoal, MoveEndEffector, - MoveEndEffectorCfg, MotionPolicy, PickUp, - PickUpCfg, + PickUpOptions, ) from scripts.tutorials.atomic_action.move_held_object import ( build_grasp_generator_cfg, @@ -194,15 +194,19 @@ 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(cfg=MoveEndEffectorCfg())) + atomic_engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) + atomic_engine.register(MoveEndEffector()) atomic_engine.register( PickUp( - cfg=PickUpCfg( - control_part="arm", - hand_control_part="hand", - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, + default_options=PickUpOptions( approach_direction=resolve_pickup_approach_direction( pickup_approach, position_case, sim.device ), @@ -275,9 +279,9 @@ def _run_case( ActionBinding, ActionInvocation, AtomicActionEngine, + ControlPartCommandProfile, HeldObjectPoseGoal, MoveHeldObject, - MoveHeldObjectCfg, MotionPolicy, ) from scripts.tutorials.atomic_action.move_held_object import ( @@ -314,16 +318,15 @@ 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( - cfg=MoveHeldObjectCfg( - control_part="arm", - hand_control_part="hand", - hand_close_qpos=hand_close, - ), - ) + atomic_engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + grasp=hand_close, + ) + }, ) + atomic_engine.register(MoveHeldObject()) target_pose = _make_object_target_pose(sim.device, case.xyz) elapsed, mem_delta, peak_gpu, result = timed_call( lambda: atomic_engine.compile( diff --git a/scripts/benchmark/atomic_action/move_joints_benchmark.py b/scripts/benchmark/atomic_action/move_joints_benchmark.py index 8ee335184..ef9fbc5e3 100644 --- a/scripts/benchmark/atomic_action/move_joints_benchmark.py +++ b/scripts/benchmark/atomic_action/move_joints_benchmark.py @@ -240,8 +240,8 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: ensure_torch() from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, + ControlPartCommandProfile, MoveJoints, - MoveJointsCfg, ) from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg from embodichain.lab.sim.planners import ToppraPlannerCfg @@ -271,14 +271,13 @@ 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( - cfg=MoveJointsCfg( - named_joint_positions={"ready": ready_qpos}, - ), - ) + atomic_engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "arm": ControlPartCommandProfile.joint_positions(ready=ready_qpos), + }, ) + atomic_engine.register(MoveJoints()) results: list[dict[str, object]] = [] video_paths: list[str] = [] diff --git a/scripts/benchmark/atomic_action/pickup_benchmark.py b/scripts/benchmark/atomic_action/pickup_benchmark.py index 3bb3ec7cb..0c85f52b0 100644 --- a/scripts/benchmark/atomic_action/pickup_benchmark.py +++ b/scripts/benchmark/atomic_action/pickup_benchmark.py @@ -126,9 +126,10 @@ def _run_case( ActionBinding, ActionInvocation, AtomicActionEngine, + ControlPartCommandProfile, GraspGoal, PickUp, - PickUpCfg, + PickUpOptions, MotionPolicy, ) from scripts.tutorials.atomic_action.pickup import ( @@ -159,14 +160,18 @@ def _run_case( approach_direction_text = format_vector3( pickup_approach_direction_tuple(approach, position_case) ) - atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) atomic_engine.register( PickUp( - cfg=PickUpCfg( - control_part="arm", - hand_control_part="hand", - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, + default_options=PickUpOptions( approach_direction=approach_direction, pre_grasp_distance=0.15, lift_height=0.16, diff --git a/scripts/benchmark/atomic_action/place_benchmark.py b/scripts/benchmark/atomic_action/place_benchmark.py index 31f93c6c2..cdd348c69 100644 --- a/scripts/benchmark/atomic_action/place_benchmark.py +++ b/scripts/benchmark/atomic_action/place_benchmark.py @@ -177,9 +177,10 @@ def _prepare_held_state( ActionBinding, ActionInvocation, AtomicActionEngine, + ControlPartCommandProfile, GraspGoal, PickUp, - PickUpCfg, + PickUpOptions, MotionPolicy, ) from scripts.tutorials.atomic_action.place import ( @@ -191,14 +192,18 @@ 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 = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) atomic_engine.register( PickUp( - cfg=PickUpCfg( - control_part="arm", - hand_control_part="hand", - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, + default_options=PickUpOptions( approach_direction=resolve_pickup_approach_direction( pickup_approach, position_case, sim.device ), @@ -258,10 +263,11 @@ def _run_case( ActionBinding, ActionInvocation, AtomicActionEngine, + ControlPartCommandProfile, MotionPolicy, Place, - PlaceCfg, PlaceGoal, + PlaceOptions, ) from scripts.tutorials.atomic_action.place import ( compute_pick_close_end_step, @@ -298,14 +304,18 @@ 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 = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) atomic_engine.register( Place( - cfg=PlaceCfg( - control_part="arm", - hand_control_part="hand", - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, + default_options=PlaceOptions( lift_height=PLACE_LIFT_HEIGHT, hand_interp_steps=HAND_INTERP_STEPS, ), diff --git a/scripts/benchmark/atomic_action/press_benchmark.py b/scripts/benchmark/atomic_action/press_benchmark.py index 3cd5a2c90..cb6c9bda3 100644 --- a/scripts/benchmark/atomic_action/press_benchmark.py +++ b/scripts/benchmark/atomic_action/press_benchmark.py @@ -84,13 +84,13 @@ def _ensure_runtime_imports() -> None: ActionBinding as action_binding_cls, ActionInvocation as action_invocation_cls, AtomicActionEngine as atomic_action_engine_cls, + ControlPartCommandProfile as control_part_command_profile_cls, EndEffectorPoseGoal as end_effector_pose_target_cls, MoveEndEffector as move_end_effector_cls, - MoveEndEffectorCfg as move_end_effector_cfg_cls, MotionPolicy as motion_policy_cls, Press as press_cls, - PressCfg as press_cfg_cls, PressGoal as press_target_cls, + PressOptions as press_options_cls, ) from embodichain.lab.sim.cfg import ( RigidBodyAttributesCfg as rigid_body_attributes_cfg_cls, @@ -126,15 +126,15 @@ def _ensure_runtime_imports() -> None: "torch": torch_module, "SimulationManager": simulation_manager_cls, "AtomicActionEngine": atomic_action_engine_cls, + "ControlPartCommandProfile": control_part_command_profile_cls, "ActionBinding": action_binding_cls, "ActionInvocation": action_invocation_cls, "EndEffectorPoseGoal": end_effector_pose_target_cls, "MoveEndEffector": move_end_effector_cls, - "MoveEndEffectorCfg": move_end_effector_cfg_cls, "MotionPolicy": motion_policy_cls, "Press": press_cls, - "PressCfg": press_cfg_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, @@ -486,14 +486,18 @@ 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(cfg=MoveEndEffectorCfg())) + atomic_engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + grasp=hand_close, + ) + }, + ) + atomic_engine.register(MoveEndEffector()) atomic_engine.register( Press( - cfg=PressCfg( - control_part="arm", - hand_control_part="hand", - hand_close_qpos=hand_close, + default_options=PressOptions( hand_interp_steps=HAND_INTERP_STEPS, ), ) 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 c05bd2ded..41e8a85d2 100644 --- a/scripts/tutorials/atomic_action/assemble.py +++ b/scripts/tutorials/atomic_action/assemble.py @@ -45,11 +45,12 @@ AssembleAffordance, AssembleGoal, AtomicActionEngine, + ControlPartCommandProfile, GraspGoal, PickUp, - PickUpCfg, + PickUpOptions, Place, - PlaceCfg, + PlaceOptions, MotionPolicy, ) from embodichain.lab.sim.cfg import ( @@ -409,12 +410,7 @@ def run_assemble_demo( # Step 1 - the left arm picks the soda can up by its top part. pick_up_action = PickUp( - cfg=PickUpCfg( - name="pick_up", - control_part="left_arm", - hand_control_part="left_hand", - hand_open_qpos=left_open, - hand_close_qpos=left_close, + default_options=PickUpOptions( pick_object_part="top", pre_grasp_distance=PICKUP_PRE_GRASP_DISTANCE, lift_height=PICKUP_LIFT_HEIGHT, @@ -427,17 +423,20 @@ def run_assemble_demo( ) # Step 2 - the left arm places the can directly above the cube. place_action = Place( - cfg=PlaceCfg( - name="place", - control_part="left_arm", - hand_control_part="left_hand", - hand_open_qpos=left_open, - hand_close_qpos=left_close, + default_options=PlaceOptions( lift_height=PLACE_LIFT_HEIGHT, hand_interp_steps=PLACE_HAND_INTERP_STEPS, ), ) - engine = AtomicActionEngine(motion_generator=motion_gen) + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "left_hand": ControlPartCommandProfile.joint_positions( + open=left_open, + grasp=left_close, + ) + }, + ) engine.register(pick_up_action) engine.register(place_action) diff --git a/scripts/tutorials/atomic_action/coordinated_pickment.py b/scripts/tutorials/atomic_action/coordinated_pickment.py index 093128544..446a22ad1 100644 --- a/scripts/tutorials/atomic_action/coordinated_pickment.py +++ b/scripts/tutorials/atomic_action/coordinated_pickment.py @@ -44,9 +44,10 @@ ActionInvocation, Affordance, AtomicActionEngine, + ControlPartCommandProfile, CoordinatedPickGoal, CoordinatedPickment, - CoordinatedPickmentCfg, + CoordinatedPickmentOptions, ObjectSemantics, MotionPolicy, ) @@ -669,16 +670,7 @@ def run_coordinated_pickment_demo( robot, "right_hand", sim.device, preset.hand_close_qpos ) pickment_action = CoordinatedPickment( - 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, + default_options=CoordinatedPickmentOptions( pre_grasp_distance=PICKMENT_PRE_GRASP_DISTANCE, lift_height=PICKMENT_LIFT_HEIGHT, hand_interp_steps=PICKMENT_HAND_INTERP_STEPS, @@ -686,7 +678,19 @@ def run_coordinated_pickment_demo( object_motion_keyframes=PICKMENT_OBJECT_MOTION_KEYFRAMES, ), ) - engine = AtomicActionEngine(motion_generator=motion_gen) + engine = AtomicActionEngine( + motion_generator=motion_gen, + 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.register(pickment_action) left_grasp_pose, right_grasp_pose = build_object_grasp_poses( diff --git a/scripts/tutorials/atomic_action/coordinated_placement.py b/scripts/tutorials/atomic_action/coordinated_placement.py index 893d7b73d..d163e31e1 100644 --- a/scripts/tutorials/atomic_action/coordinated_placement.py +++ b/scripts/tutorials/atomic_action/coordinated_placement.py @@ -45,14 +45,15 @@ ActionInvocation, Affordance, AtomicActionEngine, + ControlPartCommandProfile, CoordinatedPlacement, - CoordinatedPlacementCfg, + CoordinatedPlacementOptions, CoordinatedPlacementGoal, GraspGoal, HeldObjectState, ObjectSemantics, PickUp, - PickUpCfg, + PickUpOptions, MotionPolicy, TaskState, ) @@ -796,37 +797,21 @@ def run_coordinated_placement_demo( ) left_open, left_close = get_hand_open_close_qpos(robot, "left_hand", sim.device) left_pick_action = PickUp( - cfg=PickUpCfg( - control_part="left_arm", - hand_control_part="left_hand", - hand_open_qpos=left_open, - hand_close_qpos=left_close, + default_options=PickUpOptions( pre_grasp_distance=PICK_APPROACH_DISTANCE, lift_height=0.12, hand_interp_steps=10, ), ) right_pick_action = PickUp( - cfg=PickUpCfg( - control_part="right_arm", - hand_control_part="right_hand", - hand_open_qpos=right_open, - hand_close_qpos=right_close, + default_options=PickUpOptions( pre_grasp_distance=PICK_APPROACH_DISTANCE, lift_height=0.10, hand_interp_steps=PAN_PICK_HAND_INTERP_STEPS, ), ) coordinated_action = CoordinatedPlacement( - 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, + default_options=CoordinatedPlacementOptions( release=True, placing_height_offset=BREAD_TARGET_HEIGHT_OFFSET, support_height_offset=SUPPORT_TARGET_HEIGHT_OFFSET, @@ -836,7 +821,19 @@ def run_coordinated_placement_demo( retreat_steps=18, ), ) - engine = AtomicActionEngine(motion_generator=motion_gen) + engine = AtomicActionEngine( + motion_generator=motion_gen, + 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.register(coordinated_action) full_joint_ids = list(range(robot.dof)) state = engine.initial_context() diff --git a/scripts/tutorials/atomic_action/hand_over.py b/scripts/tutorials/atomic_action/hand_over.py index a5366864a..c733661c7 100644 --- a/scripts/tutorials/atomic_action/hand_over.py +++ b/scripts/tutorials/atomic_action/hand_over.py @@ -42,10 +42,11 @@ ActionInvocation, GraspGoal, AtomicActionEngine, + ControlPartCommandProfile, HandOver, - HandOverCfg, + HandOverOptions, PickUp, - PickUpCfg, + PickUpOptions, MotionPolicy, ) from embodichain.lab.sim.cfg import ( @@ -346,12 +347,7 @@ def run_handover_demo( # Step 1 - the left arm picks the object up by its top part. pick_up_action = PickUp( - cfg=PickUpCfg( - name="pick_up", - control_part="left_arm", - hand_control_part="left_hand", - hand_open_qpos=left_open, - hand_close_qpos=left_close, + default_options=PickUpOptions( pick_object_part="top", pre_grasp_distance=PICKUP_PRE_GRASP_DISTANCE, lift_height=PICKUP_LIFT_HEIGHT, @@ -363,17 +359,7 @@ def run_handover_demo( ) # Step 2 - hand the object from the left arm to the right arm. handover_action = HandOver( - 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, + default_options=HandOverOptions( receive_pick_object_part="bottom", middle_object_pose=middle_pose, final_object_pose=final_pose, @@ -387,7 +373,19 @@ def run_handover_demo( ), ), ) - engine = AtomicActionEngine(motion_generator=motion_gen) + engine = AtomicActionEngine( + motion_generator=motion_gen, + 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.register(pick_up_action) engine.register(handover_action) diff --git a/scripts/tutorials/atomic_action/move_end_effector.py b/scripts/tutorials/atomic_action/move_end_effector.py index 8e2fa6446..e1081eea0 100644 --- a/scripts/tutorials/atomic_action/move_end_effector.py +++ b/scripts/tutorials/atomic_action/move_end_effector.py @@ -35,7 +35,6 @@ AtomicActionEngine, EndEffectorPoseGoal, MoveEndEffector, - MoveEndEffectorCfg, MotionPolicy, ) from embodichain.utils import logger @@ -75,7 +74,7 @@ def main() -> None: motion_gen = create_toppra_motion_generator(robot) engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register(MoveEndEffector(cfg=MoveEndEffectorCfg())) + engine.register(MoveEndEffector()) poses = torch.stack( [ diff --git a/scripts/tutorials/atomic_action/move_held_object.py b/scripts/tutorials/atomic_action/move_held_object.py index 81e6975da..fc641f6a7 100644 --- a/scripts/tutorials/atomic_action/move_held_object.py +++ b/scripts/tutorials/atomic_action/move_held_object.py @@ -34,15 +34,14 @@ ActionBinding, ActionInvocation, AtomicActionEngine, + ControlPartCommandProfile, EndEffectorPoseGoal, GraspGoal, HeldObjectPoseGoal, MoveEndEffector, - MoveEndEffectorCfg, MoveHeldObject, - MoveHeldObjectCfg, PickUp, - PickUpCfg, + PickUpOptions, MotionPolicy, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg @@ -129,26 +128,26 @@ 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(MoveEndEffectorCfg())) + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) + engine.register(MoveEndEffector()) engine.register( PickUp( - PickUpCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, + default_options=PickUpOptions( pre_grasp_distance=0.15, lift_height=0.16, hand_interp_steps=HAND_INTERP_STEPS, ), ) ) - engine.register( - MoveHeldObject( - MoveHeldObjectCfg( - hand_close_qpos=hand_close, - ), - ) - ) + engine.register(MoveHeldObject()) semantics = create_antipodal_semantics( obj, diff --git a/scripts/tutorials/atomic_action/move_joints.py b/scripts/tutorials/atomic_action/move_joints.py index 3ee97b9ad..1ce5482f1 100644 --- a/scripts/tutorials/atomic_action/move_joints.py +++ b/scripts/tutorials/atomic_action/move_joints.py @@ -33,9 +33,9 @@ ActionBinding, ActionInvocation, AtomicActionEngine, + ControlPartCommandProfile, JointPositionGoal, MoveJoints, - MoveJointsCfg, MotionPolicy, ) from embodichain.utils import logger @@ -79,14 +79,13 @@ def main() -> None: [0.0, -1.57, 1.57, -1.57, -1.57, 0.0], ) ) - engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register( - MoveJoints( - cfg=MoveJointsCfg( - named_joint_positions={"ready": ready}, - ), - ) + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "arm": ControlPartCommandProfile.joint_positions(ready=ready), + }, ) + engine.register(MoveJoints()) if not args.no_vis_eef_axis: draw_axis_marker( diff --git a/scripts/tutorials/atomic_action/pickup.py b/scripts/tutorials/atomic_action/pickup.py index 84b9fbb88..4eaaa4575 100644 --- a/scripts/tutorials/atomic_action/pickup.py +++ b/scripts/tutorials/atomic_action/pickup.py @@ -33,9 +33,10 @@ ActionBinding, ActionInvocation, AtomicActionEngine, + ControlPartCommandProfile, GraspGoal, PickUp, - PickUpCfg, + PickUpOptions, MotionPolicy, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg @@ -133,12 +134,18 @@ 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 = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) engine.register( PickUp( - cfg=PickUpCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, + default_options=PickUpOptions( approach_direction=resolve_approach_direction(args, sim.device), pre_grasp_distance=0.15, lift_height=0.16, diff --git a/scripts/tutorials/atomic_action/place.py b/scripts/tutorials/atomic_action/place.py index 9fe690367..b59a30840 100644 --- a/scripts/tutorials/atomic_action/place.py +++ b/scripts/tutorials/atomic_action/place.py @@ -33,12 +33,13 @@ ActionBinding, ActionInvocation, AtomicActionEngine, + ControlPartCommandProfile, GraspGoal, PickUp, - PickUpCfg, + PickUpOptions, Place, - PlaceCfg, PlaceGoal, + PlaceOptions, MotionPolicy, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg @@ -134,12 +135,18 @@ 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 = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) engine.register( PickUp( - cfg=PickUpCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, + default_options=PickUpOptions( pre_grasp_distance=0.15, lift_height=0.16, hand_interp_steps=HAND_INTERP_STEPS, @@ -148,9 +155,7 @@ def main() -> None: ) engine.register( Place( - cfg=PlaceCfg( - hand_open_qpos=hand_open, - hand_close_qpos=hand_close, + default_options=PlaceOptions( lift_height=PLACE_LIFT_HEIGHT, hand_interp_steps=HAND_INTERP_STEPS, ), diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py index 2df787461..756cbacea 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -33,11 +33,11 @@ ActionBinding, ActionInvocation, AtomicActionEngine, + ControlPartCommandProfile, EndEffectorPoseGoal, MoveEndEffector, - MoveEndEffectorCfg, Press, - PressCfg, + PressOptions, PressGoal, MotionPolicy, ) @@ -162,13 +162,20 @@ 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(MoveEndEffectorCfg())) + 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, + ) + }, + ) + engine.register(MoveEndEffector()) engine.register( Press( - PressCfg( - hand_close_qpos=hand_close, + default_options=PressOptions( hand_interp_steps=HAND_INTERP_STEPS, ), ) diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index 814e782b6..6a5a4bbb8 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -32,37 +32,38 @@ AssembleGoal, AtomicAction, AtomicActionEngine, + ControlPartCommandProfile, CoordinatedHeldObjectState, CoordinatedPickGoal, CoordinatedPickment, - CoordinatedPickmentCfg, + CoordinatedPickmentOptions, CoordinatedPlacement, - CoordinatedPlacementCfg, CoordinatedPlacementGoal, + CoordinatedPlacementOptions, EndEffectorPoseGoal, GraspGoal, HandOver, - HandOverCfg, + HandOverOptions, HeldObjectPoseGoal, HeldObjectState, JointPositionGoal, MotionPolicy, MoveEndEffector, - MoveEndEffectorCfg, + MoveEndEffectorOptions, MoveHeldObject, - MoveHeldObjectCfg, + MoveHeldObjectOptions, MoveJoints, - MoveJointsCfg, + MoveJointsOptions, ObjectSemantics, PickUp, - PickUpCfg, + PickUpOptions, Place, - PlaceCfg, PlaceGoal, + PlaceOptions, PlanningContext, Press, - PressCfg, PressGoal, + PressOptions, RobotObservation, SceneSnapshot, TaskState, @@ -109,18 +110,24 @@ def _robot() -> Mock: robot = Mock() robot.device = torch.device("cpu") 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 == "arm": + 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, ROBOT_DOF) def get_joint_ids(name: str | None = None) -> list[int]: - if name == "arm": + if name in {"arm", "alternate_arm"}: return list(range(ARM_DOF)) - if name == "hand": + if name in {"hand", "alternate_hand"}: return list(range(ARM_DOF, ROBOT_DOF)) return list(range(ROBOT_DOF)) @@ -158,13 +165,35 @@ def _motion_generator() -> Mock: return generator -def _bind_action(generator: Mock, action: ActionT) -> ActionT: +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.""" - engine = AtomicActionEngine(generator) + 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) engine.register(action) return action +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) + + def _context(task: TaskState | None = None) -> PlanningContext: qpos = torch.zeros(NUM_ENVS, ROBOT_DOF) return PlanningContext( @@ -220,6 +249,13 @@ def _dual_motion_generator() -> Mock: robot = Mock() robot.device = torch.device("cpu") 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"}: @@ -315,6 +351,26 @@ def test_builtin_descriptors_expose_goals_not_legacy_targets() -> None: 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__") + + 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) + + 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] @@ -342,10 +398,11 @@ def test_joint_position_goal_rejects_invalid_tensor_shape( def test_move_end_effector_returns_full_robot_timed_plan() -> None: generator = _motion_generator() - action = _bind_action(generator, MoveEndEffector(MoveEndEffectorCfg())) + action = _bind_action(generator, MoveEndEffector()) context = _context() - plan = action.plan( + plan = _plan_action( + action, _invocation( "move_end_effector", EndEffectorPoseGoal(torch.eye(4)), @@ -365,7 +422,10 @@ def test_move_joints_uses_binding_and_preserves_uncontrolled_joints() -> None: named = {"ready": torch.full((ARM_DOF,), 0.4)} action = _bind_action( generator, - MoveJoints(MoveJointsCfg(named_joint_positions=named)), + MoveJoints(), + control_profiles={ + "arm": ControlPartCommandProfile.joint_positions(**named), + }, ) qpos = torch.zeros(NUM_ENVS, ROBOT_DOF) qpos[:, ARM_DOF:] = 0.7 @@ -376,7 +436,8 @@ def test_move_joints_uses_binding_and_preserves_uncontrolled_joints() -> None: env_ids=torch.arange(NUM_ENVS), ) - plan = action.plan( + plan = _plan_action( + action, _invocation("move_joints", JointPositionGoal("ready"), sample_count=8), context, ) @@ -387,17 +448,13 @@ def test_move_joints_uses_binding_and_preserves_uncontrolled_joints() -> None: def test_pick_and_place_declare_effects_without_mutating_context() -> None: generator = _motion_generator() - hand_open = torch.zeros(HAND_DOF) - hand_close = torch.ones(HAND_DOF) - pick = _bind_action( - generator, - PickUp(PickUpCfg(hand_open_qpos=hand_open, hand_close_qpos=hand_close)), - ) + pick = _bind_action(generator, PickUp()) initial = _context() semantics = _semantics() grasp = torch.eye(4).repeat(NUM_ENVS, 1, 1) - pick_plan = pick.plan( + pick_plan = _plan_action( + pick, _invocation("pick_up", GraspGoal(semantics=semantics, grasp_xpos=grasp)), initial, ) @@ -406,17 +463,15 @@ def test_pick_and_place_declare_effects_without_mutating_context() -> None: assert initial.task.get_held_object("arm") is None assert picked_task.get_held_object("arm") is not None - place = _bind_action( - generator, - Place(PlaceCfg(hand_open_qpos=hand_open, hand_close_qpos=hand_close)), - ) + place = _bind_action(generator, Place()) picked_context = PlanningContext( robot=initial.robot, task=picked_task, scene=initial.scene, env_ids=initial.env_ids, ) - place_plan = place.plan( + place_plan = _plan_action( + place, _invocation("place", PlaceGoal(torch.eye(4))), picked_context, ) @@ -430,10 +485,7 @@ def test_pick_and_place_declare_effects_without_mutating_context() -> None: def test_move_held_object_requires_projected_attachment() -> None: generator = _motion_generator() - action = _bind_action( - generator, - MoveHeldObject(MoveHeldObjectCfg(hand_close_qpos=torch.ones(HAND_DOF))), - ) + action = _bind_action(generator, MoveHeldObject()) invocation = _invocation( "move_held_object", HeldObjectPoseGoal(torch.eye(4)), @@ -441,7 +493,7 @@ def test_move_held_object_requires_projected_attachment() -> None: ) with pytest.raises(ValueError, match="requires an object held"): - action.plan(invocation, _context()) + _plan_action(action, invocation, _context()) held = _held() task = TaskState( @@ -449,19 +501,17 @@ def test_move_held_object_requires_projected_attachment() -> None: device="cpu", held_objects={"arm": held}, ) - plan = action.plan(invocation, _context(task)) + plan = _plan_action(action, invocation, _context(task)) assert plan.plan_success.all() assert plan.expected_effects.is_empty def test_press_uses_invocation_sample_budget() -> None: generator = _motion_generator() - action = _bind_action( - generator, - Press(PressCfg(hand_close_qpos=torch.ones(HAND_DOF))), - ) + action = _bind_action(generator, Press()) - plan = action.plan( + plan = _plan_action( + action, _invocation("press", PressGoal(torch.eye(4)), sample_count=12), _context(), ) @@ -473,21 +523,21 @@ def test_press_uses_invocation_sample_budget() -> None: def test_motion_source_and_sample_count_are_not_action_config_fields() -> None: with pytest.raises(TypeError): - MoveEndEffectorCfg(motion_source="motion_gen") + MoveEndEffectorOptions(motion_source="motion_gen") # type: ignore[call-arg] with pytest.raises(TypeError): - MoveJointsCfg(sample_interval=10) + MoveJointsOptions(sample_interval=10) # type: ignore[call-arg] def test_move_joints_rejects_binding_with_wrong_goal_skill() -> None: generator = _motion_generator() - action = _bind_action(generator, MoveJoints(MoveJointsCfg())) + 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.plan(invocation, _context()) + action.resolve_request(invocation) def test_planner_timing_is_preserved_in_simple_action() -> None: @@ -504,7 +554,7 @@ def test_planner_timing_is_preserved_in_simple_action() -> None: NUM_ENVS, 1 ) generator.generate.return_value.duration = torch.full((NUM_ENVS,), 0.3) - action = _bind_action(generator, MoveJoints(MoveJointsCfg())) + action = _bind_action(generator, MoveJoints()) invocation = ActionInvocation( skill_id="move_joints", goal=JointPositionGoal(torch.ones(ARM_DOF)), @@ -512,7 +562,7 @@ def test_planner_timing_is_preserved_in_simple_action() -> None: motion_policy=MotionPolicy(motion_source="motion_gen", sample_count=3), ) - plan = action.plan(invocation, _context()) + plan = _plan_action(action, invocation, _context()) assert plan.trajectory.duration.tolist() == pytest.approx([0.3, 0.3]) assert plan.trajectory.velocities is not None @@ -539,7 +589,8 @@ def compute_ik( waypoints[:, 1, 0, 3] = 0.3 action = _bind_action(generator, MoveEndEffector()) - plan = action.plan( + plan = _plan_action( + action, _invocation( "move_end_effector", EndEffectorPoseGoal(waypoints), @@ -558,9 +609,10 @@ def test_move_joints_visits_waypoints_and_rejects_unknown_names() -> None: generator = _motion_generator() action = _bind_action( generator, - MoveJoints( - MoveJointsCfg(named_joint_positions={"ready": torch.zeros(ARM_DOF)}) - ), + MoveJoints(), + control_profiles={ + "arm": ControlPartCommandProfile.joint_positions(ready=torch.zeros(ARM_DOF)) + }, ) waypoints = torch.stack( [ @@ -570,7 +622,8 @@ def test_move_joints_visits_waypoints_and_rejects_unknown_names() -> None: dim=1, ) - plan = action.plan( + plan = _plan_action( + action, _invocation( "move_joints", JointPositionGoal(waypoints), @@ -581,8 +634,9 @@ def test_move_joints_visits_waypoints_and_rejects_unknown_names() -> None: 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="Unknown named joint-position goal"): - action.plan( + with pytest.raises(KeyError, match="has no command"): + _plan_action( + action, _invocation("move_joints", JointPositionGoal("missing")), _context(), ) @@ -602,17 +656,10 @@ def test_pick_explicit_grasp_bypasses_sampling_and_records_grasp() -> None: ) grasp = torch.eye(4).repeat(NUM_ENVS, 1, 1) grasp[:, 0, 3] = torch.tensor([0.1, 0.2]) - action = _bind_action( - generator, - PickUp( - PickUpCfg( - hand_open_qpos=torch.zeros(HAND_DOF), - hand_close_qpos=torch.ones(HAND_DOF), - ) - ), - ) + action = _bind_action(generator, PickUp()) - plan = action.plan( + plan = _plan_action( + action, _invocation( "pick_up", GraspGoal(semantics=semantics, grasp_xpos=grasp), @@ -628,6 +675,27 @@ def test_pick_explicit_grasp_bypasses_sampling_and_records_grasp() -> 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() + + 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 + + def test_press_closes_hand_without_changing_projected_attachment() -> None: held = _held() task = TaskState( @@ -638,10 +706,11 @@ def test_press_closes_hand_without_changing_projected_attachment() -> None: generator = _motion_generator() action = _bind_action( generator, - Press(PressCfg(hand_close_qpos=torch.ones(HAND_DOF), hand_interp_steps=4)), + Press(default_options=PressOptions(hand_interp_steps=4)), ) - plan = action.plan( + plan = _plan_action( + action, _invocation("press", PressGoal(torch.eye(4)), sample_count=12), _context(task), ) @@ -656,23 +725,19 @@ def test_press_closes_hand_without_changing_projected_attachment() -> None: 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( - HandOverCfg( - transfer_hand_open_qpos=torch.zeros(HAND_DOF), - transfer_hand_close_qpos=torch.ones(HAND_DOF), - receive_hand_open_qpos=torch.zeros(HAND_DOF), - receive_hand_close_qpos=torch.ones(HAND_DOF), - middle_object_pose=torch.eye(4), - final_object_pose=torch.eye(4), - hand_interp_steps=4, - hold_steps=2, - retreat_steps=5, - ) - ), + HandOver(default_options=handover_options), ) - original_final_pose = action.final_object_pose.clone() + 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() @@ -704,10 +769,10 @@ def plan_from_start( motion_policy=MotionPolicy(sample_count=30), ) - plan = action.plan(invocation, _dual_context(task)) + plan = _plan_action(action, invocation, _dual_context(task)) assert plan.plan_success.all() - assert torch.equal(action.final_object_pose, original_final_pose) + assert torch.equal(handover_options.final_object_pose, original_final_pose) def test_coordinated_pick_returns_full_dof_plan_and_projected_relation() -> None: @@ -715,15 +780,11 @@ def test_coordinated_pick_returns_full_dof_plan_and_projected_relation() -> None action = _bind_action( generator, CoordinatedPickment( - CoordinatedPickmentCfg( - left_hand_open_qpos=torch.zeros(HAND_DOF), - left_hand_close_qpos=torch.ones(HAND_DOF), - right_hand_open_qpos=torch.zeros(HAND_DOF), - right_hand_close_qpos=torch.ones(HAND_DOF), + default_options=CoordinatedPickmentOptions( hand_interp_steps=4, hold_steps=2, object_motion_keyframes=3, - ) + ), ), ) semantics = ObjectSemantics( @@ -743,7 +804,7 @@ def test_coordinated_pick_returns_full_dof_plan_and_projected_relation() -> None ) context = _dual_context() - plan = action.plan(invocation, context) + plan = _plan_action(action, invocation, context) projected = plan.expected_effects.apply(context.task, plan.plan_success) assert plan.plan_success.tolist() == [True, True] @@ -781,15 +842,11 @@ def fail_second_environment( action = _bind_action( generator, CoordinatedPickment( - CoordinatedPickmentCfg( - left_hand_open_qpos=torch.zeros(HAND_DOF), - left_hand_close_qpos=torch.ones(HAND_DOF), - right_hand_open_qpos=torch.zeros(HAND_DOF), - right_hand_close_qpos=torch.ones(HAND_DOF), + default_options=CoordinatedPickmentOptions( hand_interp_steps=4, hold_steps=2, object_motion_keyframes=3, - ) + ), ), ) target_pose = torch.eye(4) @@ -810,7 +867,7 @@ def fail_second_environment( ) context = _dual_context() - plan = action.plan(invocation, context) + plan = _plan_action(action, invocation, context) projected = plan.expected_effects.apply(context.task, plan.plan_success) assert plan.plan_success.tolist() == [True, False] @@ -829,14 +886,11 @@ def test_coordinated_placement_projects_release_and_support_attachment() -> None action = _bind_action( generator, CoordinatedPlacement( - CoordinatedPlacementCfg( - placing_hand_open_qpos=torch.zeros(HAND_DOF), - placing_hand_close_qpos=torch.ones(HAND_DOF), - support_hand_close_qpos=torch.ones(HAND_DOF), + default_options=CoordinatedPlacementOptions( hand_interp_steps=4, hold_steps=3, retreat_steps=5, - ) + ), ), ) placing = _held( @@ -861,7 +915,7 @@ def test_coordinated_placement_projects_release_and_support_attachment() -> None ) context = _dual_context(task) - plan = action.plan(invocation, context) + plan = _plan_action(action, invocation, context) projected = plan.expected_effects.apply(context.task, plan.plan_success) assert plan.plan_success.tolist() == [True, True] @@ -877,14 +931,7 @@ def test_coordinated_actions_reject_curobo_motion_generation() -> None: policy = MotionPolicy(motion_source="motion_gen", sample_count=30) pick = _bind_action( generator, - CoordinatedPickment( - CoordinatedPickmentCfg( - left_hand_open_qpos=torch.zeros(HAND_DOF), - left_hand_close_qpos=torch.ones(HAND_DOF), - right_hand_open_qpos=torch.zeros(HAND_DOF), - right_hand_close_qpos=torch.ones(HAND_DOF), - ) - ), + CoordinatedPickment(), ) pick_invocation = ActionInvocation( skill_id="coordinated_pickment", @@ -902,17 +949,11 @@ def test_coordinated_actions_reject_curobo_motion_generation() -> None: ) with pytest.raises(ValueError, match="not supported"): - pick.plan(pick_invocation, _dual_context()) + _plan_action(pick, pick_invocation, _dual_context()) placement = _bind_action( generator, - CoordinatedPlacement( - CoordinatedPlacementCfg( - placing_hand_open_qpos=torch.zeros(HAND_DOF), - placing_hand_close_qpos=torch.ones(HAND_DOF), - support_hand_close_qpos=torch.ones(HAND_DOF), - ) - ), + CoordinatedPlacement(), ) placement_invocation = ActionInvocation( skill_id="coordinated_placement", @@ -921,4 +962,4 @@ def test_coordinated_actions_reject_curobo_motion_generation() -> None: motion_policy=policy, ) with pytest.raises(ValueError, match="not supported"): - placement.plan(placement_invocation, _dual_context()) + _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_curobo_motion_source_e2e.py b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py index 1dde97ea8..0d254c870 100644 --- a/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py +++ b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py @@ -52,7 +52,6 @@ EndEffectorPoseGoal, MotionPolicy, MoveEndEffector, - MoveEndEffectorCfg, ) ROBOT_UID = "curobo_franka" @@ -89,7 +88,7 @@ def _make_franka_curobo_engine(): ) ) engine = AtomicActionEngine(mg) - engine.register(MoveEndEffector(MoveEndEffectorCfg())) + engine.register(MoveEndEffector()) return sim, robot, engine diff --git a/tests/sim/atomic_actions/test_engine.py b/tests/sim/atomic_actions/test_engine.py index 615146b82..f9338ec73 100644 --- a/tests/sim/atomic_actions/test_engine.py +++ b/tests/sim/atomic_actions/test_engine.py @@ -27,21 +27,25 @@ from embodichain.lab.sim.atomic_actions import ( ActionBinding, - ActionCfg, + ActionControlOverrides, ActionInvocation, + ActionOptions, ActionPlan, AtomicAction, AtomicActionEngine, + ControlPartCommandProfile, + JointPositionCommand, JointPositionGoal, MotionPolicy, PlanningContext, + ResolvedActionRequest, register_action, get_registered_actions, unregister_action, ) -class StubAction(AtomicAction[JointPositionGoal]): +class StubAction(AtomicAction[JointPositionGoal, ActionOptions]): """Deterministic test action that commands every robot joint.""" skill_id: ClassVar[str] = "stub" @@ -50,10 +54,10 @@ class StubAction(AtomicAction[JointPositionGoal]): def plan( self, - invocation: ActionInvocation[JointPositionGoal], + request: ResolvedActionRequest[JointPositionGoal, ActionOptions], context: PlanningContext, ) -> ActionPlan: - goal = self.require_goal(invocation) + goal = self.require_goal(request) assert isinstance(goal.target, torch.Tensor) target = goal.target.to(context.robot.qpos) if target.dim() == 1: @@ -66,7 +70,7 @@ def plan( target = torch.nan_to_num(target) trajectory = torch.stack([context.robot.qpos, target], dim=1) return self.build_plan( - invocation, + request, context, success=success, trajectory=trajectory, @@ -79,20 +83,28 @@ class OtherStubAction(StubAction): skill_id: ClassVar[str] = "other_stub" -def _engine(batch_size: int = 2, robot_dof: int = 3) -> AtomicActionEngine: +def _engine( + batch_size: int = 2, + robot_dof: int = 3, + control_profiles: dict[str, ControlPartCommandProfile] | None = None, +) -> AtomicActionEngine: 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 AtomicActionEngine(generator) + return AtomicActionEngine(generator, control_profiles=control_profiles) -def _invocation(qpos: torch.Tensor) -> ActionInvocation[JointPositionGoal]: +def _invocation( + qpos: torch.Tensor, +) -> ActionInvocation[JointPositionGoal, ActionOptions]: return ActionInvocation( skill_id="stub", goal=JointPositionGoal(qpos), @@ -113,7 +125,7 @@ def test_global_registry_uses_stable_skill_id() -> None: def test_engine_compile_projects_terminal_state_between_actions() -> None: engine = _engine() - engine.register(StubAction(ActionCfg(name="stub"))) + engine.register(StubAction()) first = torch.ones(2, 3) second = torch.full((2, 3), 2.0) @@ -128,7 +140,7 @@ def test_engine_compile_projects_terminal_state_between_actions() -> None: def test_engine_compile_holds_failed_rows_for_remaining_actions() -> None: engine = _engine() - engine.register(StubAction(ActionCfg(name="stub"))) + 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) @@ -159,8 +171,8 @@ def test_engine_rejects_unknown_skill() -> None: def test_engine_rejects_duplicate_instance_registration() -> None: engine = _engine() - first = StubAction(ActionCfg(name="first")) - second = StubAction(ActionCfg(name="second")) + first = StubAction() + second = StubAction() engine.register(first) with pytest.raises(ValueError, match="already registered"): engine.register(second) @@ -168,8 +180,8 @@ def test_engine_rejects_duplicate_instance_registration() -> None: def test_engine_binds_one_planning_service_to_every_action() -> None: engine = _engine() - first = StubAction(ActionCfg(name="first")) - second = OtherStubAction(ActionCfg(name="second")) + first = StubAction() + second = OtherStubAction() engine.register(first) engine.register(second) @@ -180,6 +192,59 @@ def test_engine_binds_one_planning_service_to_every_action() -> None: 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() @@ -189,7 +254,7 @@ def test_engine_motion_generator_is_read_only() -> None: def test_engine_plan_action_supports_unregistered_configured_instance() -> None: engine = _engine() - action = StubAction(ActionCfg(name="temporary")) + action = StubAction() plan = engine.plan_action( action, @@ -203,7 +268,7 @@ def test_engine_plan_action_supports_unregistered_configured_instance() -> None: def test_action_cannot_be_rebound_to_another_engine() -> None: - action = StubAction(ActionCfg(name="stub")) + action = StubAction() _engine().register(action) with pytest.raises(ValueError, match="another AtomicActionEngine"): @@ -211,28 +276,25 @@ def test_action_cannot_be_rebound_to_another_engine() -> None: def test_unbound_action_rejects_direct_planning() -> None: - action = StubAction(ActionCfg(name="stub")) + action = StubAction() with pytest.raises(RuntimeError, match="not bound"): - action.plan( - _invocation(torch.ones(2, 3)), - _engine().initial_context(), - ) + action.resolve_request(_invocation(torch.ones(2, 3))) def test_engine_rejects_plan_for_a_different_skill() -> None: engine = _engine() - action = StubAction(ActionCfg(name="stub")) + action = StubAction() original_plan = action.plan def wrong_skill_plan( - invocation: ActionInvocation[JointPositionGoal], + request: ResolvedActionRequest[JointPositionGoal, ActionOptions], context: PlanningContext, ) -> ActionPlan: - return replace(original_plan(invocation, context), skill_id="other") + 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 invocation"): + 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 5b9848248..bd2e0b4e3 100644 --- a/tests/sim/atomic_actions/test_engine_per_env.py +++ b/tests/sim/atomic_actions/test_engine_per_env.py @@ -26,8 +26,8 @@ from embodichain.lab.sim.atomic_actions import ( ActionBinding, - ActionCfg, ActionInvocation, + ActionOptions, ActionPlan, Affordance, AtomicAction, @@ -41,6 +41,7 @@ ObjectSemantics, PlanningContext, RecoveryPolicy, + ResolvedActionRequest, RobotObservation, SceneEntityPose, SceneSnapshot, @@ -50,7 +51,7 @@ from embodichain.lab.sim.atomic_actions.goals import resolve_pose_goal -class DynamicAction(AtomicAction[EndEffectorPoseGoal]): +class DynamicAction(AtomicAction[EndEffectorPoseGoal, ActionOptions]): """Test action whose terminal joint command follows a scene entity x pose.""" skill_id: ClassVar[str] = "dynamic" @@ -58,20 +59,22 @@ class DynamicAction(AtomicAction[EndEffectorPoseGoal]): manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) def __init__(self) -> None: - super().__init__(ActionCfg(name="dynamic")) + super().__init__() self.plan_count = 0 + self.requests: list[ResolvedActionRequest] = [] def plan( self, - invocation: ActionInvocation[EndEffectorPoseGoal], + request: ResolvedActionRequest[EndEffectorPoseGoal, ActionOptions], context: PlanningContext, ) -> ActionPlan: - goal = self.require_goal(invocation) + 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( - invocation, + request, context, success=True, trajectory=torch.stack([context.robot.qpos, target], dim=1), @@ -85,10 +88,10 @@ class EffectAction(DynamicAction): def plan( self, - invocation: ActionInvocation[EndEffectorPoseGoal], + request: ResolvedActionRequest[EndEffectorPoseGoal, ActionOptions], context: PlanningContext, ) -> ActionPlan: - goal = self.require_goal(invocation) + 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( @@ -100,7 +103,7 @@ def plan( grasp_xpos=torch.eye(4), ) return self.build_plan( - invocation, + request, context, success=True, trajectory=torch.stack([context.robot.qpos, target], dim=1), @@ -112,8 +115,10 @@ 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") @@ -194,9 +199,64 @@ def test_scene_motion_replans_late_bound_goal() -> None: 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( diff --git a/tests/sim/atomic_actions/test_motion_source_e2e.py b/tests/sim/atomic_actions/test_motion_source_e2e.py index 6b434a24a..5daaffbb6 100644 --- a/tests/sim/atomic_actions/test_motion_source_e2e.py +++ b/tests/sim/atomic_actions/test_motion_source_e2e.py @@ -31,7 +31,6 @@ EndEffectorPoseGoal, MotionPolicy, MoveEndEffector, - MoveEndEffectorCfg, ) @@ -60,7 +59,7 @@ def _setup(self, motion_source: str): MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=self.ROBOT_UID)) ) engine = AtomicActionEngine(mg) - engine.register(MoveEndEffector(MoveEndEffectorCfg())) + engine.register(MoveEndEffector()) return sim, robot, engine def _teardown(self, sim): diff --git a/tests/sim/planners/test_curobo_planner.py b/tests/sim/planners/test_curobo_planner.py index a4e01cb56..a08fc2d0b 100644 --- a/tests/sim/planners/test_curobo_planner.py +++ b/tests/sim/planners/test_curobo_planner.py @@ -631,7 +631,6 @@ def _make_curobo_engine( from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, MoveEndEffector, - MoveEndEffectorCfg, ) from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator @@ -645,7 +644,7 @@ def _make_curobo_engine( ) ) engine = AtomicActionEngine(motion_generator) - engine.register(MoveEndEffector(MoveEndEffectorCfg())) + engine.register(MoveEndEffector()) return engine From 54815b7b2652a790068aa9bb55ee202d49e665c1 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 3 Aug 2026 15:19:45 +0000 Subject: [PATCH 5/5] wip --- agent_context/MAP.yaml | 5 + .../topics/atomic-actions/atomic-actions.md | 44 +- .../sim/atomic_actions/builtin_actions.md | 16 +- .../overview/sim/atomic_actions/index.md | 205 ++++++-- docs/source/tutorial/atomic_actions.rst | 118 ++++- .../lab/sim/atomic_actions/__init__.py | 2 + embodichain/lab/sim/atomic_actions/engine.py | 54 +- .../sim/atomic_actions/primitives/__init__.py | 15 + examples/sim/planners/curobo_planner.py | 2 - .../move_end_effector_benchmark.py | 2 - .../move_held_object_benchmark.py | 25 +- .../atomic_action/move_joints_benchmark.py | 3 - .../atomic_action/pickup_benchmark.py | 17 +- .../atomic_action/place_benchmark.py | 34 +- .../atomic_action/press_benchmark.py | 15 +- scripts/tutorials/atomic_action/assemble.py | 224 ++------ .../atomic_action/coordinated_pickment.py | 366 +++----------- .../atomic_action/coordinated_placement.py | 477 +++++------------- scripts/tutorials/atomic_action/hand_over.py | 229 ++------- .../atomic_action/move_end_effector.py | 12 +- .../atomic_action/move_held_object.py | 43 +- .../tutorials/atomic_action/move_joints.py | 13 +- scripts/tutorials/atomic_action/pickup.py | 40 +- scripts/tutorials/atomic_action/place.py | 52 +- scripts/tutorials/atomic_action/press.py | 25 +- .../tutorials/atomic_action/scenario_utils.py | 351 +++++++++++++ .../tutorials/atomic_action/tutorial_utils.py | 91 +++- tests/sim/atomic_actions/test_actions.py | 6 +- .../test_curobo_motion_source_e2e.py | 2 - tests/sim/atomic_actions/test_engine.py | 73 ++- .../sim/atomic_actions/test_engine_per_env.py | 2 +- .../atomic_actions/test_motion_source_e2e.py | 2 - tests/sim/planners/test_curobo_planner.py | 2 - 33 files changed, 1230 insertions(+), 1337 deletions(-) create mode 100644 scripts/tutorials/atomic_action/scenario_utils.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index c597adca7..c9219adeb 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -451,6 +451,11 @@ topics: - 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: diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index c6f69d65a..18e57fff9 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -30,15 +30,33 @@ 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 after -`engine.register(action)` or -`engine.plan_action(action, invocation, context)`. A bound action cannot be -reused by another engine. +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. + +`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 -Register configured action instances by their class-level stable `skill_id`, -then call: +Built-ins are already registered by their class-level stable `skill_id`; call: ```python compiled = engine.compile(invocations, context=None) @@ -49,9 +67,8 @@ 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 `engine.plan_action(...)` for an unregistered configured instance when an -application needs multiple variants with the same stable `skill_id`; the action -still uses the engine's single motion generator. +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 @@ -101,6 +118,15 @@ 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 diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index 624c5c28c..87889a968 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -6,7 +6,9 @@ ``` EmbodiChain ships nine built-in action implementations with stable skill IDs; -applications register the configured instances they need with an engine. +`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. @@ -14,9 +16,14 @@ 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 during -`register()` or `plan_action()`. Generic motion and recovery choices belong to -the invocation. +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 current manipulation primitives consume semantic `open` and `grasp` @@ -280,7 +287,6 @@ engine = AtomicActionEngine( "left_arm": ControlPartCommandProfile.joint_positions(home=home_qpos), }, ) -engine.register(MoveJoints()) explicit_goal = JointPositionGoal(target=home_qpos) named_goal = JointPositionGoal(target="home") diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index b6f3bef89..e56bf960d 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -49,17 +49,17 @@ and whole-body control are not implemented by this module yet. | +-- MotionGenerator / planner backend | | +-- device and shared TrajectoryBuilder | | | -| registered AtomicAction.plan(...) -> ActionPlan | -+--------------------------+----------------------------------+ - | - +------------+-------------+ - | | - v v - compile(...) start(...) / tick(...) - fixed projection observed closed loop - | | - v v - CompiledTrajectory JointCommand + events +| 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 boundary is deliberate: @@ -90,10 +90,10 @@ manual_invocation = ActionInvocation( recovery_policy=RecoveryPolicy(max_replans=2), ) -# Inspect a single plan, compile a fixed sequence, or execute with recovery. -plan = engine.plan(manual_invocation, latest_context) -compiled = engine.compile((manual_invocation,), latest_context) -session = engine.start((manual_invocation,), latest_context) +# 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) ``` A manual caller may bypass the semantic-schema adapter only when its target and @@ -108,6 +108,37 @@ 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 +``` + +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. + +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`. + ## Core contracts The public contracts separate values with different owners and lifetimes. This @@ -238,52 +269,108 @@ calibrated commands. ### Engine-owned planning resources -One engine owns one motion generator. Actions borrow its planning services only -after `register()` or `plan_action()` binds them: +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) -engine.register(MoveEndEffector()) -engine.register(MoveJoints()) + +# 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: -- action constructors optionally contain only typed default options; +- 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. -Prefer invocation `skill_options` when behavior varies per call. If an -application still needs two instances with different default options and the -same stable skill ID, keep one or both outside the registry and call -`engine.plan_action(...)` explicitly: +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, +) + +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 -left_pick = PickUp(default_options=left_pick_options) -right_pick = PickUp(default_options=right_pick_options) +custom_engine = AtomicActionEngine(motion_generator, load_builtins=False) +custom_engine.register(MyAction()) -left_plan = engine.plan_action(left_pick, left_invocation, latest_context) -right_plan = engine.plan_action(right_pick, right_invocation, latest_context) +engine.register(CustomPickUp(), replace=True) ``` -Both instances still borrow the same engine-owned motion generator. +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. -## Which planning API to use +### Implementation and advanced APIs -| API | Use it for | Result / behavior | +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: + +| API | Intended caller | Behavior | |---|---|---| -| `AtomicAction.plan(request, context)` | Implementing a skill | Consumes an engine-resolved immutable request; application code normally calls it through the engine | -| `engine.plan(invocation, context)` | Planning one registered skill | Resolves the registered action, binds shared resources, and validates its plan | -| `engine.plan_action(action, invocation, context)` | Planning an unregistered configured instance | Supports multiple configurations with one `skill_id` and one engine backend | -| `engine.compile(invocations, context)` | Fixed-scene/offline sequence planning | Returns one concatenated `CompiledTrajectory` and a hypothetical projected context | -| `engine.start(invocations, context)` | Observed incremental execution | Returns an `ExecutionSession`; each `tick()` emits at most one command and recovery events | -| `session.revise_current(invocation)` | Explicit runtime parameter/goal update | Requires a newer revision of the active logical invocation and replans from the latest context | +| `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 +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), +) -`AtomicAction.plan()` is therefore not a second execution API. It is the -polymorphic implementation point used by the engine. Neither it nor the engine -mutates the simulator. +plan = engine.plan(invocation, latest_context) +if plan.plan_success.all(): + positions = plan.trajectory.positions +``` + +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 @@ -298,39 +385,55 @@ from embodichain.lab.sim.atomic_actions import ( ActionInvocation, AtomicActionEngine, EndEffectorPoseGoal, - ExecutionEventKind, - ExecutionStatus, MotionPolicy, - MoveEndEffector, - RecoveryPolicy, - SceneEntityPose, ) engine = AtomicActionEngine(motion_generator) -engine.register(MoveEndEffector()) +binding = ActionBinding(manipulators={"primary": "left_arm"}) +motion_policy = MotionPolicy(sample_count=80, control_dt=1.0 / 60.0) -invocation = ActionInvocation( +approach = 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), + 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, ) -compiled = engine.compile((invocation,)) +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 ``` 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. +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 +from embodichain.lab.sim.atomic_actions import ( + ExecutionStatus, + RecoveryPolicy, + SceneEntityPose, +) + moving_goal = ActionInvocation( skill_id="move_end_effector", goal=EndEffectorPoseGoal( diff --git a/docs/source/tutorial/atomic_actions.rst b/docs/source/tutorial/atomic_actions.rst index a5635cfc1..fc3b306f8 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -41,10 +41,48 @@ 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. Atomic action constructors accept only optional -typed default options; ``register()`` binds each action to the engine resources. Use -``engine.plan_action(action, invocation, context)`` for an unregistered, -default-option-specific action instance. +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 ----------------- @@ -109,11 +147,45 @@ engine is built: from its bound manipulator. Joint limits validate possible commands, but do not define their semantic meaning; supply calibrated robot commands in production. +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 ( + ActionBinding, + ActionInvocation, + EndEffectorPoseGoal, + MotionPolicy, + ) + + 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(): + 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 during planning: +the scene is treated as fixed and all goals in a sequence are known during +planning: .. code-block:: python @@ -123,25 +195,43 @@ the scene is treated as fixed during planning: AtomicActionEngine, EndEffectorPoseGoal, MotionPolicy, - MoveEndEffector, ) engine = AtomicActionEngine(motion_generator) - engine.register(MoveEndEffector()) + binding = ActionBinding(manipulators={"primary": "left_arm"}) + motion_policy = MotionPolicy(sample_count=80, control_dt=1.0 / 60.0) - invocation = ActionInvocation( + approach = 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), + goal=EndEffectorPoseGoal(xpos=approach_pose), + binding=binding, + motion_policy=motion_policy, ) - compiled = engine.compile((invocation,)) - trajectory = compiled.trajectory.positions + 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(): + 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. +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. Dynamic goals and closed-loop execution --------------------------------------- diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index f4ca7c71f..7564a52b4 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -74,6 +74,7 @@ from .runtime import ActionPlanningServices from .primitives import ( AssembleGoal, + BUILTIN_ACTION_TYPES, CoordinatedPickGoal, CoordinatedPickment, CoordinatedPickmentOptions, @@ -126,6 +127,7 @@ "AssembleGoal", "AtomicAction", "AtomicActionEngine", + "BUILTIN_ACTION_TYPES", "CompiledTrajectory", "CompletionCondition", "CompletionConditionKind", diff --git a/embodichain/lab/sim/atomic_actions/engine.py b/embodichain/lab/sim/atomic_actions/engine.py index f79ef2774..38a5190c0 100644 --- a/embodichain/lab/sim/atomic_actions/engine.py +++ b/embodichain/lab/sim/atomic_actions/engine.py @@ -36,11 +36,15 @@ from .execution import ExecutionSession -_global_action_registry: dict[str, type[AtomicAction]] = {} +_global_extension_registry: dict[str, type[AtomicAction]] = {} def register_action(action_class: type[AtomicAction]) -> None: - """Register an atomic action class under its stable skill identifier. + """Register an extension action type for process-wide discovery. + + 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. @@ -52,27 +56,27 @@ def register_action(action_class: type[AtomicAction]) -> None: 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_action_registry.get(descriptor.skill_id) + 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_action_registry[descriptor.skill_id] = action_class + _global_extension_registry[descriptor.skill_id] = action_class def unregister_action(skill_id: str) -> None: - """Remove a globally registered skill class if present. + """Remove a globally discoverable extension action type if present. Args: skill_id: Stable registered skill identifier. """ - _global_action_registry.pop(skill_id, None) + _global_extension_registry.pop(skill_id, None) def get_registered_actions() -> dict[str, type[AtomicAction]]: - """Return a copy of the global skill-class registry.""" - return dict(_global_action_registry) + """Return a copy of the process-wide extension action-type registry.""" + return dict(_global_extension_registry) class AtomicActionEngine: @@ -82,12 +86,24 @@ 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. + + 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: @@ -119,11 +135,14 @@ def actions(self) -> dict[str, AtomicAction]: """Registered action instances keyed by stable skill identifier.""" return dict(self._actions) - def register(self, action: AtomicAction) -> None: + 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. @@ -134,13 +153,22 @@ def register(self, action: AtomicAction) -> None: 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: + 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." ) 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 plan_action( self, action: AtomicAction, @@ -150,9 +178,9 @@ def plan_action( """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 supports multiple configured instances with the - same stable skill identifier while preserving one engine-owned motion - generator. + 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: action: Configured action implementation to invoke. diff --git a/embodichain/lab/sim/atomic_actions/primitives/__init__.py b/embodichain/lab/sim/atomic_actions/primitives/__init__.py index 9b85035e8..85de2c985 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/__init__.py +++ b/embodichain/lab/sim/atomic_actions/primitives/__init__.py @@ -18,6 +18,7 @@ from __future__ import annotations +from ..core import AtomicAction from .coordinated_pickment import ( CoordinatedPickGoal, CoordinatedPickment, @@ -44,8 +45,22 @@ from .place import AssembleGoal, Place, PlaceGoal, PlaceOptions from .press import Press, PressGoal, PressOptions +BUILTIN_ACTION_TYPES: tuple[type[AtomicAction], ...] = ( + MoveEndEffector, + MoveJoints, + PickUp, + MoveHeldObject, + Place, + Press, + CoordinatedPickment, + CoordinatedPlacement, + HandOver, +) +"""Built-in action implementations instantiated once per action engine.""" + __all__ = [ "AssembleGoal", + "BUILTIN_ACTION_TYPES", "CoordinatedPickGoal", "CoordinatedPickment", "CoordinatedPickmentOptions", diff --git a/examples/sim/planners/curobo_planner.py b/examples/sim/planners/curobo_planner.py index 75e6a890b..e12c0e717 100644 --- a/examples/sim/planners/curobo_planner.py +++ b/examples/sim/planners/curobo_planner.py @@ -63,7 +63,6 @@ ActionInvocation, AtomicActionEngine, EndEffectorPoseGoal, - MoveEndEffector, MotionPolicy, ) from embodichain.data import get_data_path @@ -749,7 +748,6 @@ def main() -> None: ) ) engine = AtomicActionEngine(motion_generator) - engine.register(MoveEndEffector()) binding = ActionBinding(manipulators={"primary": control_part}) motion_policy = MotionPolicy( motion_source="motion_gen", diff --git a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py index 67635eb54..8ad488594 100644 --- a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py +++ b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py @@ -233,7 +233,6 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: ensure_torch() from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, - MoveEndEffector, ) from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg from embodichain.lab.sim.planners import ToppraPlannerCfg @@ -263,7 +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()) 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 ae722d520..da623c0f1 100644 --- a/scripts/benchmark/atomic_action/move_held_object_benchmark.py +++ b/scripts/benchmark/atomic_action/move_held_object_benchmark.py @@ -181,9 +181,7 @@ def _prepare_held_state( ControlPartCommandProfile, EndEffectorPoseGoal, GraspGoal, - MoveEndEffector, MotionPolicy, - PickUp, PickUpOptions, ) from scripts.tutorials.atomic_action.move_held_object import ( @@ -203,19 +201,6 @@ def _prepare_held_state( ) }, ) - atomic_engine.register(MoveEndEffector()) - atomic_engine.register( - PickUp( - default_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, - ), - ) - ) semantics = create_antipodal_object_semantics( obj=obj, preset=object_preset, @@ -244,6 +229,14 @@ def _prepare_held_state( 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, + ), ), ) ) @@ -281,7 +274,6 @@ def _run_case( AtomicActionEngine, ControlPartCommandProfile, HeldObjectPoseGoal, - MoveHeldObject, MotionPolicy, ) from scripts.tutorials.atomic_action.move_held_object import ( @@ -326,7 +318,6 @@ def _run_case( ) }, ) - atomic_engine.register(MoveHeldObject()) target_pose = _make_object_target_pose(sim.device, case.xyz) elapsed, mem_delta, peak_gpu, result = timed_call( lambda: atomic_engine.compile( diff --git a/scripts/benchmark/atomic_action/move_joints_benchmark.py b/scripts/benchmark/atomic_action/move_joints_benchmark.py index ef9fbc5e3..0fc6e6272 100644 --- a/scripts/benchmark/atomic_action/move_joints_benchmark.py +++ b/scripts/benchmark/atomic_action/move_joints_benchmark.py @@ -241,7 +241,6 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, ControlPartCommandProfile, - MoveJoints, ) from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg from embodichain.lab.sim.planners import ToppraPlannerCfg @@ -277,8 +276,6 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: "arm": ControlPartCommandProfile.joint_positions(ready=ready_qpos), }, ) - atomic_engine.register(MoveJoints()) - 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 0c85f52b0..8e6686d88 100644 --- a/scripts/benchmark/atomic_action/pickup_benchmark.py +++ b/scripts/benchmark/atomic_action/pickup_benchmark.py @@ -128,7 +128,6 @@ def _run_case( AtomicActionEngine, ControlPartCommandProfile, GraspGoal, - PickUp, PickUpOptions, MotionPolicy, ) @@ -169,16 +168,6 @@ def _run_case( ) }, ) - atomic_engine.register( - PickUp( - default_options=PickUpOptions( - approach_direction=approach_direction, - pre_grasp_distance=0.15, - lift_height=0.16, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) - ) semantics = create_antipodal_object_semantics( obj=obj, preset=object_preset, @@ -197,6 +186,12 @@ def _run_case( 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, + ), ), ) ) diff --git a/scripts/benchmark/atomic_action/place_benchmark.py b/scripts/benchmark/atomic_action/place_benchmark.py index cdd348c69..071be0d51 100644 --- a/scripts/benchmark/atomic_action/place_benchmark.py +++ b/scripts/benchmark/atomic_action/place_benchmark.py @@ -179,7 +179,6 @@ def _prepare_held_state( AtomicActionEngine, ControlPartCommandProfile, GraspGoal, - PickUp, PickUpOptions, MotionPolicy, ) @@ -201,18 +200,6 @@ def _prepare_held_state( ) }, ) - atomic_engine.register( - PickUp( - default_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, - ), - ) - ) semantics = create_antipodal_object_semantics( obj=obj, preset=object_preset, @@ -230,6 +217,14 @@ def _prepare_held_state( 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, + ), ), ) ) @@ -265,7 +260,6 @@ def _run_case( AtomicActionEngine, ControlPartCommandProfile, MotionPolicy, - Place, PlaceGoal, PlaceOptions, ) @@ -313,14 +307,6 @@ def _run_case( ) }, ) - atomic_engine.register( - Place( - default_options=PlaceOptions( - lift_height=PLACE_LIFT_HEIGHT, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) - ) place_pose = _make_place_pose(sim.device, case.xyz) elapsed, mem_delta, peak_gpu, result = timed_call( lambda: atomic_engine.compile( @@ -333,6 +319,10 @@ def _run_case( 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, diff --git a/scripts/benchmark/atomic_action/press_benchmark.py b/scripts/benchmark/atomic_action/press_benchmark.py index cb6c9bda3..4d408ce41 100644 --- a/scripts/benchmark/atomic_action/press_benchmark.py +++ b/scripts/benchmark/atomic_action/press_benchmark.py @@ -86,9 +86,7 @@ def _ensure_runtime_imports() -> None: AtomicActionEngine as atomic_action_engine_cls, ControlPartCommandProfile as control_part_command_profile_cls, EndEffectorPoseGoal as end_effector_pose_target_cls, - MoveEndEffector as move_end_effector_cls, MotionPolicy as motion_policy_cls, - Press as press_cls, PressGoal as press_target_cls, PressOptions as press_options_cls, ) @@ -130,9 +128,7 @@ def _ensure_runtime_imports() -> None: "ActionBinding": action_binding_cls, "ActionInvocation": action_invocation_cls, "EndEffectorPoseGoal": end_effector_pose_target_cls, - "MoveEndEffector": move_end_effector_cls, "MotionPolicy": motion_policy_cls, - "Press": press_cls, "PressGoal": press_target_cls, "PressOptions": press_options_cls, "RigidBodyAttributesCfg": rigid_body_attributes_cfg_cls, @@ -494,14 +490,6 @@ def _build_atomic_engine( ) }, ) - atomic_engine.register(MoveEndEffector()) - atomic_engine.register( - Press( - default_options=PressOptions( - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) - ) return atomic_engine @@ -596,6 +584,9 @@ def _timed_atomic_run( PressGoal(xpos=press_target), binding, MotionPolicy(sample_count=PRESS_SAMPLE_INTERVAL), + skill_options=PressOptions( + hand_interp_steps=HAND_INTERP_STEPS, + ), ), ) ) diff --git a/scripts/tutorials/atomic_action/assemble.py b/scripts/tutorials/atomic_action/assemble.py index 41e8a85d2..22a57a7a5 100644 --- a/scripts/tutorials/atomic_action/assemble.py +++ b/scripts/tutorials/atomic_action/assemble.py @@ -34,10 +34,8 @@ 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, @@ -47,29 +45,27 @@ AtomicActionEngine, ControlPartCommandProfile, GraspGoal, - PickUp, PickUpOptions, - Place, PlaceOptions, MotionPolicy, ) -from embodichain.lab.sim.cfg import ( - JointDrivePropertiesCfg, - RigidBodyAttributesCfg, - RigidObjectCfg, - RobotCfg, - URDFCfg, -) +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, @@ -79,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 = ( @@ -134,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"}, + 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), ), - 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", - ), - 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, ) @@ -338,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) @@ -409,24 +278,20 @@ def run_assemble_demo( ) # Step 1 - the left arm picks the soda can up by its top part. - pick_up_action = PickUp( - default_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,), + 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( - default_options=PlaceOptions( - lift_height=PLACE_LIFT_HEIGHT, - hand_interp_steps=PLACE_HAND_INTERP_STEPS, - ), + place_options = PlaceOptions( + lift_height=PLACE_LIFT_HEIGHT, + hand_interp_steps=PLACE_HAND_INTERP_STEPS, ) engine = AtomicActionEngine( motion_generator=motion_gen, @@ -437,9 +302,6 @@ def run_assemble_demo( ) }, ) - 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..." ) @@ -465,12 +327,14 @@ def run_assemble_demo( 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, ), ) ) diff --git a/scripts/tutorials/atomic_action/coordinated_pickment.py b/scripts/tutorials/atomic_action/coordinated_pickment.py index 446a22ad1..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,54 +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 ( ActionBinding, ActionInvocation, - Affordance, AtomicActionEngine, ControlPartCommandProfile, CoordinatedPickGoal, CoordinatedPickment, CoordinatedPickmentOptions, - ObjectSemantics, 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 = ( @@ -155,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", @@ -184,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, ) @@ -344,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, @@ -370,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, @@ -434,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, @@ -556,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, @@ -660,23 +408,25 @@ 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_action = CoordinatedPickment( - default_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_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, ) engine = AtomicActionEngine( motion_generator=motion_gen, @@ -691,7 +441,9 @@ def run_coordinated_pickment_demo( ), }, ) - 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, @@ -752,6 +504,7 @@ def run_coordinated_pickment_demo( end_effectors={"left": "left_hand", "right": "right_hand"}, ), MotionPolicy(sample_count=PICKMENT_SAMPLE_INTERVAL), + skill_options=pickment_options, ), ) ) @@ -769,7 +522,10 @@ def run_coordinated_pickment_demo( "coordinated_pickment", traj, joint_ids, - pickment_action.get_segment_lengths(PICKMENT_SAMPLE_INTERVAL), + 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 d163e31e1..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,12 +36,10 @@ 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 ( ActionBinding, ActionInvocation, - Affordance, AtomicActionEngine, ControlPartCommandProfile, CoordinatedPlacement, @@ -52,29 +48,41 @@ GraspGoal, HeldObjectState, ObjectSemantics, - PickUp, 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, @@ -108,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" @@ -120,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 @@ -196,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: @@ -348,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, @@ -368,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, @@ -392,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, @@ -414,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] @@ -556,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, @@ -656,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, @@ -731,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, @@ -788,38 +554,38 @@ 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( - default_options=PickUpOptions( - pre_grasp_distance=PICK_APPROACH_DISTANCE, - lift_height=0.12, - 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( - default_options=PickUpOptions( - pre_grasp_distance=PICK_APPROACH_DISTANCE, - lift_height=0.10, - 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, ) - coordinated_action = CoordinatedPlacement( - default_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, - ), + 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, ) engine = AtomicActionEngine( motion_generator=motion_gen, @@ -834,12 +600,14 @@ def run_coordinated_placement_demo( ), }, ) - 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 = 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( @@ -850,9 +618,13 @@ def run_coordinated_placement_demo( sim.device, z_clearance=BREAD_GRASP_Z_CLEARANCE, ) - start_time = time.time() - left_pick_result = engine.plan_action( - left_pick_action, + pan_grasp_pose = build_pan_handle_grasp_pose( + pan_pose, + pan_vertices, + sim.device, + z_clearance=PAN_GRASP_Z_CLEARANCE, + ) + pick_invocations = ( ActionInvocation( skill_id="pick_up", goal=GraspGoal( @@ -864,36 +636,8 @@ def run_coordinated_placement_demo( end_effectors={"primary": "left_hand"}, ), motion_policy=MotionPolicy(sample_count=PICK_SAMPLE_INTERVAL), + skill_options=left_pick_options, ), - state, - ) - logger.log_info( - f"Plan left bread pick-up cost time: {time.time() - start_time:.2f} seconds" - ) - if not left_pick_result.plan_success.all(): - logger.log_warning("Failed to plan left bread pick-up trajectory.") - return - left_pick_traj = left_pick_result.trajectory.positions - state = state.project( - qpos=left_pick_traj[:, -1], - task=left_pick_result.expected_effects.apply( - state.task, left_pick_result.plan_success - ), - ) - 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 = engine.plan_action( - right_pick_action, ActionInvocation( skill_id="pick_up", goal=GraspGoal( @@ -905,25 +649,34 @@ def run_coordinated_placement_demo( 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.plan_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 + left_pick_traj = left_pick_result.trajectory.positions right_pick_traj = right_pick_result.trajectory.positions - state = state.project( - qpos=right_pick_traj[:, -1], - task=right_pick_result.expected_effects.apply( - state.task, right_pick_result.plan_success - ), - ) + 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: @@ -932,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 ): @@ -947,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() @@ -960,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 ) @@ -1034,7 +789,7 @@ def log_pick_execution(step_idx: int, total_steps: int) -> None: release=True, ) start_time = time.time() - compiled = engine.compile( + placement_compiled = engine.compile( ( ActionInvocation( skill_id="coordinated_placement", @@ -1050,13 +805,14 @@ def log_pick_execution(step_idx: int, total_steps: int) -> None: }, ), motion_policy=MotionPolicy(sample_count=COORDINATED_SAMPLE_INTERVAL), + skill_options=coordinated_options, ), ), state, ) - coordinated_success = compiled.plan_success - coordinated_traj = compiled.trajectory.positions - state = compiled.projected_context + 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" @@ -1070,8 +826,9 @@ def log_pick_execution(step_idx: int, total_steps: int) -> None: coordinated_traj, full_joint_ids, coordinated_action._compute_segment_lengths( - coordinated_action.cfg.release, + coordinated_options.release, COORDINATED_SAMPLE_INTERVAL, + coordinated_options, ), ) @@ -1088,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, @@ -1103,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 c733661c7..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,10 +31,8 @@ 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, @@ -43,27 +40,25 @@ GraspGoal, AtomicActionEngine, ControlPartCommandProfile, - HandOver, HandOverOptions, - PickUp, PickUpOptions, MotionPolicy, ) -from embodichain.lab.sim.cfg import ( - JointDrivePropertiesCfg, - RigidBodyAttributesCfg, - RigidObjectCfg, - RobotCfg, - URDFCfg, -) +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, @@ -73,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 = ( @@ -121,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"}, + 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), ), - 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", - ), - 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, ) @@ -291,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, @@ -346,31 +216,27 @@ def run_handover_demo( ) # Step 1 - the left arm picks the object up by its top part. - pick_up_action = PickUp( - default_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 - ), + 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( - default_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 - ), + 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( @@ -386,9 +252,6 @@ def run_handover_demo( ), }, ) - 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..." ) @@ -405,6 +268,7 @@ def run_handover_demo( end_effectors={"primary": "left_hand"}, ), MotionPolicy(sample_count=PICKUP_SAMPLE_INTERVAL), + skill_options=pick_up_options, ), ActionInvocation( "hand_over", @@ -420,6 +284,7 @@ def run_handover_demo( }, ), MotionPolicy(sample_count=HANDOVER_SAMPLE_INTERVAL), + skill_options=handover_options, ), ) ) diff --git a/scripts/tutorials/atomic_action/move_end_effector.py b/scripts/tutorials/atomic_action/move_end_effector.py index e1081eea0..27cef7267 100644 --- a/scripts/tutorials/atomic_action/move_end_effector.py +++ b/scripts/tutorials/atomic_action/move_end_effector.py @@ -28,13 +28,11 @@ 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, EndEffectorPoseGoal, - MoveEndEffector, MotionPolicy, ) from embodichain.utils import logger @@ -43,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, @@ -57,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() @@ -74,7 +71,6 @@ def main() -> None: motion_gen = create_toppra_motion_generator(robot) engine = AtomicActionEngine(motion_generator=motion_gen) - engine.register(MoveEndEffector()) poses = torch.stack( [ diff --git a/scripts/tutorials/atomic_action/move_held_object.py b/scripts/tutorials/atomic_action/move_held_object.py index fc641f6a7..22d4873e6 100644 --- a/scripts/tutorials/atomic_action/move_held_object.py +++ b/scripts/tutorials/atomic_action/move_held_object.py @@ -29,7 +29,6 @@ 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, @@ -38,9 +37,6 @@ EndEffectorPoseGoal, GraspGoal, HeldObjectPoseGoal, - MoveEndEffector, - MoveHeldObject, - PickUp, PickUpOptions, MotionPolicy, ) @@ -54,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, @@ -74,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() @@ -137,18 +131,6 @@ def main() -> None: ) }, ) - engine.register(MoveEndEffector()) - engine.register( - PickUp( - default_options=PickUpOptions( - pre_grasp_distance=0.15, - lift_height=0.16, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) - ) - engine.register(MoveHeldObject()) - semantics = create_antipodal_semantics( obj, label="paper_cup", @@ -183,6 +165,11 @@ def main() -> None: 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", @@ -203,14 +190,6 @@ 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, @@ -218,7 +197,7 @@ def clear_object_dynamics(step_idx: int, _: int) -> None: 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 1ce5482f1..0439bf876 100644 --- a/scripts/tutorials/atomic_action/move_joints.py +++ b/scripts/tutorials/atomic_action/move_joints.py @@ -28,20 +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, ControlPartCommandProfile, JointPositionGoal, - MoveJoints, 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, @@ -55,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() @@ -85,8 +82,6 @@ def main() -> None: "arm": ControlPartCommandProfile.joint_positions(ready=ready), }, ) - engine.register(MoveJoints()) - if not args.no_vis_eef_axis: draw_axis_marker( sim, diff --git a/scripts/tutorials/atomic_action/pickup.py b/scripts/tutorials/atomic_action/pickup.py index 4eaaa4575..403061292 100644 --- a/scripts/tutorials/atomic_action/pickup.py +++ b/scripts/tutorials/atomic_action/pickup.py @@ -28,14 +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, ControlPartCommandProfile, GraspGoal, - PickUp, PickUpOptions, MotionPolicy, ) @@ -48,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, @@ -71,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() @@ -143,16 +141,6 @@ def main() -> None: ) }, ) - engine.register( - PickUp( - default_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, - ), - ) - ) semantics = create_antipodal_semantics( obj, label="cube", @@ -175,6 +163,12 @@ def main() -> None: 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, + ), ), ) ) @@ -187,14 +181,6 @@ 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, @@ -202,7 +188,7 @@ def clear_object_dynamics(step_idx: int, _: int) -> None: 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 b59a30840..6531e59cc 100644 --- a/scripts/tutorials/atomic_action/place.py +++ b/scripts/tutorials/atomic_action/place.py @@ -28,16 +28,13 @@ 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, ControlPartCommandProfile, GraspGoal, - PickUp, PickUpOptions, - Place, PlaceGoal, PlaceOptions, MotionPolicy, @@ -53,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, @@ -73,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() @@ -144,24 +139,6 @@ def main() -> None: ) }, ) - engine.register( - PickUp( - default_options=PickUpOptions( - pre_grasp_distance=0.15, - lift_height=0.16, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) - ) - engine.register( - Place( - default_options=PlaceOptions( - lift_height=PLACE_LIFT_HEIGHT, - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) - ) - semantics = create_antipodal_semantics( obj, label="cube", @@ -190,6 +167,11 @@ def main() -> None: 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", @@ -200,6 +182,10 @@ def main() -> None: ), binding, MotionPolicy(sample_count=PLACE_SAMPLE_INTERVAL), + skill_options=PlaceOptions( + lift_height=PLACE_LIFT_HEIGHT, + hand_interp_steps=HAND_INTERP_STEPS, + ), ), ) ) @@ -212,14 +198,6 @@ 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, @@ -227,7 +205,7 @@ def clear_object_dynamics(step_idx: int, _: int) -> None: 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 756cbacea..83a8dc0ff 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -28,15 +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, ControlPartCommandProfile, EndEffectorPoseGoal, - MoveEndEffector, - Press, PressOptions, PressGoal, MotionPolicy, @@ -53,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, @@ -75,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() @@ -172,15 +169,6 @@ def main() -> None: ) }, ) - engine.register(MoveEndEffector()) - engine.register( - Press( - default_options=PressOptions( - hand_interp_steps=HAND_INTERP_STEPS, - ), - ) - ) - block_center = block.get_local_pose(to_matrix=True)[0, :3, 3] press_position = block_center.clone() press_position[2] += 0.5 * BLOCK_SIZE[2] + PRESS_SURFACE_OFFSET @@ -212,6 +200,9 @@ def main() -> None: PressGoal(press_target), binding, MotionPolicy(sample_count=PRESS_SAMPLE_INTERVAL), + skill_options=PressOptions( + hand_interp_steps=HAND_INTERP_STEPS, + ), ), ) ) 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_actions.py b/tests/sim/atomic_actions/test_actions.py index 6a5a4bbb8..d0e8a2d6d 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -180,7 +180,11 @@ def _bind_action( if "hand" in name } profiles.update({} if control_profiles is None else control_profiles) - engine = AtomicActionEngine(generator, control_profiles=profiles) + engine = AtomicActionEngine( + generator, + control_profiles=profiles, + load_builtins=False, + ) engine.register(action) return action 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 0d254c870..67c0a15dc 100644 --- a/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py +++ b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py @@ -51,7 +51,6 @@ AtomicActionEngine, EndEffectorPoseGoal, MotionPolicy, - MoveEndEffector, ) ROBOT_UID = "curobo_franka" @@ -88,7 +87,6 @@ def _make_franka_curobo_engine(): ) ) engine = AtomicActionEngine(mg) - engine.register(MoveEndEffector()) return sim, robot, engine diff --git a/tests/sim/atomic_actions/test_engine.py b/tests/sim/atomic_actions/test_engine.py index f9338ec73..c30ca1d1a 100644 --- a/tests/sim/atomic_actions/test_engine.py +++ b/tests/sim/atomic_actions/test_engine.py @@ -33,11 +33,14 @@ ActionPlan, AtomicAction, AtomicActionEngine, + BUILTIN_ACTION_TYPES, ControlPartCommandProfile, JointPositionCommand, JointPositionGoal, MotionPolicy, PlanningContext, + PressGoal, + PressOptions, ResolvedActionRequest, register_action, get_registered_actions, @@ -83,11 +86,10 @@ class OtherStubAction(StubAction): skill_id: ClassVar[str] = "other_stub" -def _engine( +def _motion_generator( batch_size: int = 2, robot_dof: int = 3, - control_profiles: dict[str, ControlPartCommandProfile] | None = None, -) -> AtomicActionEngine: +) -> Mock: robot = Mock() robot.device = torch.device("cpu") robot.dof = robot_dof @@ -99,7 +101,21 @@ def _engine( generator.robot = robot generator.device = torch.device("cpu") generator.planner.cfg.planner_type = "stub_planner" - return AtomicActionEngine(generator, control_profiles=control_profiles) + 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( @@ -123,6 +139,43 @@ def test_global_registry_uses_stable_skill_id() -> None: 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 + ) + + +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 + + def test_engine_compile_projects_terminal_state_between_actions() -> None: engine = _engine() engine.register(StubAction()) @@ -178,6 +231,18 @@ def test_engine_rejects_duplicate_instance_registration() -> None: 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() diff --git a/tests/sim/atomic_actions/test_engine_per_env.py b/tests/sim/atomic_actions/test_engine_per_env.py index bd2e0b4e3..38179d517 100644 --- a/tests/sim/atomic_actions/test_engine_per_env.py +++ b/tests/sim/atomic_actions/test_engine_per_env.py @@ -123,7 +123,7 @@ def _engine() -> tuple[AtomicActionEngine, DynamicAction]: generator.robot = robot generator.device = torch.device("cpu") generator.planner.cfg.planner_type = "stub" - engine = AtomicActionEngine(generator) + engine = AtomicActionEngine(generator, load_builtins=False) action = DynamicAction() engine.register(action) return engine, action diff --git a/tests/sim/atomic_actions/test_motion_source_e2e.py b/tests/sim/atomic_actions/test_motion_source_e2e.py index 5daaffbb6..66778af01 100644 --- a/tests/sim/atomic_actions/test_motion_source_e2e.py +++ b/tests/sim/atomic_actions/test_motion_source_e2e.py @@ -30,7 +30,6 @@ AtomicActionEngine, EndEffectorPoseGoal, MotionPolicy, - MoveEndEffector, ) @@ -59,7 +58,6 @@ def _setup(self, motion_source: str): MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=self.ROBOT_UID)) ) engine = AtomicActionEngine(mg) - engine.register(MoveEndEffector()) return sim, robot, engine def _teardown(self, sim): diff --git a/tests/sim/planners/test_curobo_planner.py b/tests/sim/planners/test_curobo_planner.py index a08fc2d0b..3d952c443 100644 --- a/tests/sim/planners/test_curobo_planner.py +++ b/tests/sim/planners/test_curobo_planner.py @@ -630,7 +630,6 @@ def _make_curobo_engine( ) -> object: from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, - MoveEndEffector, ) from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator @@ -644,7 +643,6 @@ def _make_curobo_engine( ) ) engine = AtomicActionEngine(motion_generator) - engine.register(MoveEndEffector()) return engine