diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index df7ac14aa..dc289e4cb 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -545,6 +545,11 @@ topics: - whole body resource - leaf resource claim - SceneEntityPose + - PressGoal.target_pose + - SlideGoal.target_pose + - TwistGoal.target_pose + - open_loop + - axis_translation_keyframes - dynamic goal - error recovery - ActionOptions diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index 06bcf00c1..921ffd612 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -309,7 +309,9 @@ Scene dependencies must match the poses each primitive actually consumes: | `CoordinatedPickment` | Goal-owned target/initial `SceneEntityPose` values; the semantic `entity_id` only when `object_initial_pose` is omitted and semantic grounding supplies that pose. | | `Place` | A `SceneEntityPose` in ordinary `xpos`; for `AssembleGoal`, `base_pose` when supplied. Omitting `base_pose` uses the deprecated live `AssembleAffordance.base_object_entity` fallback with no dependency. | | `MoveHeldObject` | A `SceneEntityPose` in `object_target_pose`; current object orientation is derived from observed EEF pose plus verified `object_to_eef`, not a scene-object read. | -| `Press` | A `SceneEntityPose` in `xpos`. | +| `Press` | `PressGoal.target_pose` when it is a `SceneEntityPose`; affordance data is entity-free. | +| `Slide` | `SlideGoal.target_pose` when it is a `SceneEntityPose`; the local grasp mesh does not own the link. | +| `Twist` | `TwistGoal.target_pose` when it is a `SceneEntityPose`; affordance data is entity-free. | | `CoordinatedPlacement` | `SceneEntityPose` values in the placing or support object target pose. | | `HandOver` | No semantic-object scene dependency. It verifies stable attachment identity and derives current pose from held state; its middle/final option poses are tensors, and the reused `GraspGoal.grasp_xpos` field is ignored. | @@ -554,10 +556,29 @@ the legacy core mapping. | `move_held_object` | `HeldObjectPoseGoal` | manipulator/end effector `primary` | | `place` | `PlaceGoal`, `AssembleGoal` | manipulator/end effector `primary` | | `press` | `PressGoal` | manipulator/end effector `primary` | +| `slide` | `SlideGoal` | manipulator/end effector `primary` | +| `twist` | `TwistGoal` | manipulator/end effector `primary` | | `coordinated_pickment` | `CoordinatedPickGoal` | `left`, `right` | | `coordinated_placement` | `CoordinatedPlacementGoal` | `placing`, `support` | | `hand_over` | `GraspGoal` | `source`, `destination` | +`PressAffordance`, `SlideAffordance`, and `TwistAffordance` contain only +target-local geometry and interaction semantics. Their goals own an explicit +`target_pose`, which may be a deterministic tensor snapshot or a late-bound +`SceneEntityPose`. Never put an `Articulation`, `RigidObject`, or live link pose +reader in these affordances. + +`Press` and `Slide` use dense axis-aligned Cartesian targets for their contact +motion. The linear motion-generator path solves every output sample with IK; +it does not resample sparse IK endpoints in joint space. `Press` has a distinct +contact segment before penetration. `TwistAffordance.axis_origin` and +`twist_axis` together define the full 3D rotation axis. + +These three motion-centric primitives declare `SkillDescriptor.open_loop=True` +and an empty `StateDelta`. Their completion means motion execution only, not +verified button actuation, grasp retention, or articulation travel. Applications +that need semantic completion must observe and verify those physical outcomes. + `GraspGoal.grasp_xpos` accepts an explicit pose tensor, a late-bound `SceneEntityPose`, or `None` for affordance sampling. A `SceneEntityPose` registers the referenced entity as a recovery dependency, allowing an executing diff --git a/docs/source/_static/atomic_actions/press.gif b/docs/source/_static/atomic_actions/press.gif index 1eafdc319..8013cd5db 100644 Binary files a/docs/source/_static/atomic_actions/press.gif and b/docs/source/_static/atomic_actions/press.gif differ diff --git a/docs/source/_static/atomic_actions/slide_pull.gif b/docs/source/_static/atomic_actions/slide_pull.gif new file mode 100644 index 000000000..415b5b5e5 Binary files /dev/null and b/docs/source/_static/atomic_actions/slide_pull.gif differ diff --git a/docs/source/_static/atomic_actions/slide_push.gif b/docs/source/_static/atomic_actions/slide_push.gif new file mode 100644 index 000000000..4da01f04f Binary files /dev/null and b/docs/source/_static/atomic_actions/slide_push.gif differ diff --git a/docs/source/_static/atomic_actions/twist.gif b/docs/source/_static/atomic_actions/twist.gif new file mode 100644 index 000000000..2c4b720fe Binary files /dev/null and b/docs/source/_static/atomic_actions/twist.gif differ diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.primitives.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.primitives.rst index 32e46c1e5..5a8c12268 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 @@ -29,6 +29,10 @@ full-robot timed trajectory and uncommitted expected effects. PlaceOptions Press PressOptions + Slide + SlideOptions + Twist + TwistOptions CoordinatedPickment CoordinatedPickmentOptions CoordinatedPlacement @@ -47,6 +51,8 @@ full-robot timed trajectory and uncommitted expected effects. PlaceGoal AssembleGoal PressGoal + SlideGoal + TwistGoal CoordinatedPickGoal CoordinatedPlacementGoal @@ -100,6 +106,22 @@ Press :show-inheritance: :exclude-members: __init__, copy, replace, to_dict +Slide +----- + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives.slide + :members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict + +Twist +----- + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives.twist + :members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict + CoordinatedPickment ------------------- 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 3c9703039..8cf64587a 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 @@ -80,6 +80,14 @@ embodichain.lab.sim.atomic_actions PlaceGoal AssembleGoal PressGoal + PressOptions + PressAffordance + SlideGoal + SlideOptions + SlideAffordance + TwistGoal + TwistOptions + TwistAffordance CoordinatedPickGoal CoordinatedPlacementGoal MoveEndEffector @@ -88,6 +96,8 @@ embodichain.lab.sim.atomic_actions MoveHeldObject Place Press + Slide + Twist CoordinatedPickment CoordinatedPlacement HandOver diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index 0205dcc8e..046e92eb4 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -5,12 +5,12 @@ ```{currentmodule} embodichain.lab.sim.atomic_actions ``` -EmbodiChain ships nine built-in action implementations with stable skill IDs; +EmbodiChain ships eleven built-in action implementations with stable skill IDs; `AtomicActionEngine` creates and registers a fresh instance of every built-in by default. Applications select them by stable skill ID rather than registering routine instances themselves. `Place` additionally accepts an `AssembleGoal`, so assembly reuses the same -release primitive instead of introducing a tenth skill ID. +release primitive instead of introducing another skill ID. All built-ins implement `plan(request, context) -> ActionPlan`, where `request` is the engine-resolved @@ -99,11 +99,30 @@ The animations below are the focused simulator demos under :link: builtin-press :link-type: ref -`press` · close, contact, and return +`press` · close, approach, press, and retract Press demo ::: +:::{grid-item-card} `Slide` +:link: builtin-slide +:link-type: ref + +`slide` · grasped translation along a constrained axis + +Slide pull demo +Slide push demo +::: + +:::{grid-item-card} `Twist` +:link: builtin-twist +:link-type: ref + +`twist` · grasped rotation about a configured axis + +Twist demo +::: + :::{grid-item-card} `CoordinatedPickment` :link: builtin-coordinated-pickment :link-type: ref @@ -140,12 +159,14 @@ The animations below are the focused simulator demos under | `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 exclusively by `primary` | preserve attachment | -| `place` | `PlaceGoal`, `AssembleGoal` | manipulator + end effector `primary` | primary: `open`, `grasp` | an active attachment must be exclusive to `primary`; `AssembleGoal` requires one | 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 | attach the shared object to both manipulators | -| `coordinated_placement` | `CoordinatedPlacementGoal` | manipulator + end effector `placing`, `support` | placing: `open`, `grasp`; support: `grasp` | two distinct objects, each held exclusively by its arm | optionally detach placing object; preserve support attachment | -| `hand_over` | `GraspGoal` | manipulator + end effector `source`, `destination` | both: `open`, `grasp` | object held exclusively by source arm | transfer attachment to destination arm | +| `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` | `PressAffordance` + target pose | open-loop motion; application verifies contact/actuation | +| `slide` | `SlideGoal` | manipulator + end effector `primary` | primary: `open`, `grasp` | `SlideAffordance` + link pose | open-loop motion; application verifies joint travel/grasp | +| `twist` | `TwistGoal` | manipulator + end effector `primary` | primary: `open`, `grasp` | `TwistAffordance` + target pose | open-loop motion; application verifies joint travel/grasp | +| `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 @@ -199,7 +220,6 @@ entity as a recovery dependency. | `MoveJoints.target` | no | no | | `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.grasp_xpos` | yes | yes | @@ -471,27 +491,107 @@ migration. ## `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. +Plans **close hand -> approach target -> contact -> press along axis -> return +to the approach pose**. `PressAffordance` is entity-free and stores an explicit +target-local surface `press_position` and `press_axis`. `PressGoal.target_pose` +is either a pose snapshot or `SceneEntityPose`, which resolves through the +current `PlanningContext.scene` and participates in dynamic-goal recovery. + +The contact, press, and retract segments use axis-aligned Cartesian keyframes; +each output sample is grounded with IK instead of being interpolated only in +joint space. The generated tool frame uses an adaptive reference axis and is a +right-handed orthonormal rotation even for vertical or oblique press axes. | Contract | Value | |---|---| | Skill ID | `press` | -| Goal | `PressGoal(xpos=...)` | +| Goal | `PressGoal(semantics=..., target_pose=...)` | | 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` | +| Motion | close, approach, contact, axis-constrained press, axis-constrained retract | +| Effect | explicitly open-loop; no physical button/contact effect is claimed | -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. +`PressOptions` controls hand-close interpolation, approach distance, +press distance, and an optional target-local `press_position`. An options-level +position overrides the affordance's explicit surface point. The bound +end-effector profile must provide `grasp`; the action keeps the gripper closed +for all arm-motion segments. **Example:** `scripts/tutorials/atomic_action/press.py` +(builtin-slide)= + +## `Slide` + +Plans a grasped linear interaction for one articulation link. The entity-free +`SlideAffordance` stores the link-local grasp mesh, `translation_axis`, and +optional joint name/limits. `SlideGoal.target_pose` supplies the link pose as a +snapshot or `SceneEntityPose`. The positive axis direction means approach and +push/close; pull/open uses its negative direction. The affordance inherits +`AntipodalAffordance` and selects a grasp with `get_best_grasp_poses()`. The grasp +approach direction is the link-frame translation axis transformed by the +current link rotation. + +With `direction="pull"`, the sequence is **approach -> reach -> close -> pull -> +open**. With `direction="push"`, it is **approach -> reach -> close -> push -> open +-> return**, where `return` moves the open gripper back to the original approach +pose. + +| Contract | Value | +|---|---| +| Skill ID | `slide` | +| Goal | `SlideGoal(semantics=..., target_pose=...)` | +| Binding | manipulator + end effector role `primary` | +| Motion | pull: approach, reach, close, pull, open; push adds return to approach | +| Effect | explicitly open-loop; no articulation travel or grasp success is claimed | + +`SlideOptions` controls `direction`, hand close/open +interpolation, approach distance, and translation distance. The link-frame +translation axis belongs to `SlideAffordance`; the bound end-effector profile +must provide `open` and `grasp`. Reach, pull/push, and push-return use +axis-aligned Cartesian samples rather than sparse joint-space endpoints. + +**Example:** `scripts/tutorials/atomic_action/slide.py` +plans and replays a pull first, then replans a push from the drawer's measured +post-pull link pose. + +(builtin-twist)= + +## `Twist` + +Plans **approach -> reach -> close -> twist -> open -> retract** for an +articulation link or a rigid object. The entity-free `TwistAffordance` stores an +explicit local `grasp_position`, `twist_axis`, and `axis_origin`, plus optional +joint name/limits. `TwistGoal.target_pose` supplies the grounded target pose. + +The grasp frame's z-axis follows the world-transformed twist axis; an adaptive +reference completes a right-handed orthonormal frame. Twist keyframes rotate +around the full 3D axis defined by `axis_origin + twist_axis`, not implicitly +around the target link origin. + +| Contract | Value | +|---|---| +| Skill ID | `twist` | +| Goal | `TwistGoal(semantics=..., target_pose=...)` | +| Binding | manipulator + end effector role `primary` | +| Motion | approach, reach, close, rotate about the target-local axis, open, retract | +| Effect | explicitly open-loop; no articulation travel or grasp success is claimed | + +`TwistOptions` controls the pre-grasp distance, close/open interpolation, +Cartesian twist keyframes, and twist angle. The pre-grasp pose is offset along +the grasp pose's negative z-axis; the target-local twist axis belongs to +`TwistAffordance`. + +`Twist` is intentionally a pure-rotation primitive. Thread pitch, coupled axial +translation, and regrasping are outside its contract; an `Unscrew` action should +model those behaviors separately. + +For all three primitives, `SkillDescriptor.open_loop` is `True`. Trajectory +completion therefore means commanded motion completion only. Applications that +need semantic success must observe button/contact or articulation state and +verify it outside the side-effect-free planner. + +**Example:** `scripts/tutorials/atomic_action/twist.py` + (builtin-coordinated-pickment)= ## `CoordinatedPickment` diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index 20c451cf2..35fe80966 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -316,7 +316,7 @@ instances to the engine's planning services: ```python engine = AtomicActionEngine(motion_generator, control_profiles=profiles) -# All nine built-ins are immediately usable by stable skill ID. +# All eleven built-ins are immediately usable by stable skill ID. assert "move_end_effector" in engine.actions assert "pick_up" in engine.actions ``` diff --git a/docs/source/overview/sim/planners/curobo_planner.md b/docs/source/overview/sim/planners/curobo_planner.md index 12ec63b27..82eaa16db 100644 --- a/docs/source/overview/sim/planners/curobo_planner.md +++ b/docs/source/overview/sim/planners/curobo_planner.md @@ -341,7 +341,7 @@ assert result.success.all() Single-arm MoveEndEffector is supported through the normal `strategy="motion_gen"` route. MoveJoints can opt in to collision-aware joint-space planning with `strategy="motion_gen"`; the action uses the planner -already owned by its MotionGenerator. Movement phases of PickUp, Place, Press, +already owned by its MotionGenerator. Movement phases of PickUp, Place, and MoveHeldObject can use the same single-arm static-world route. This first release intentionally has the following limits: diff --git a/docs/source/tutorial/atomic_actions.rst b/docs/source/tutorial/atomic_actions.rst index e26d7d5e5..b56488960 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -102,6 +102,8 @@ Focused examples live under ``scripts/tutorials/atomic_action``: * ``place.py`` * ``assemble.py`` * ``press.py`` +* ``slide.py`` +* ``twist.py`` * ``coordinated_pickment.py`` * ``coordinated_placement.py`` * ``hand_over.py`` diff --git a/embodichain/data/assets/obj_assets.py b/embodichain/data/assets/obj_assets.py index 4901bb45e..a1b77d023 100644 --- a/embodichain/data/assets/obj_assets.py +++ b/embodichain/data/assets/obj_assets.py @@ -242,6 +242,20 @@ def __init__(self, data_root: str = None): super().__init__(prefix, data_descriptor, path) +class MicrowaveOven(EmbodiChainDataset): + """get_data_path("MicrowaveOven/microwave_oven.urdf")""" + + def __init__(self, data_root: str = None): + data_descriptor = o3d.data.DataDescriptor( + os.path.join(EMBODICHAIN_DOWNLOAD_PREFIX, obj_assets, "MicrowaveOven.zip"), + "5672da2d5a888a12469d6277636646b0", + ) + prefix = type(self).__name__ + path = EMBODICHAIN_DEFAULT_DATA_ROOT if data_root is None else data_root + + super().__init__(prefix, data_descriptor, path) + + class PlasticTray(EmbodiChainDataset): """get_data_path("PlasticTray/plastic_tray.glb")""" @@ -268,3 +282,17 @@ def __init__(self, data_root: str = None): path = EMBODICHAIN_DEFAULT_DATA_ROOT if data_root is None else data_root super().__init__(prefix, data_descriptor, path) + + +class Drawer(EmbodiChainDataset): + """get_data_path("Drawer/model_split_links_with_inertials.urdf")""" + + def __init__(self, data_root: str = None): + data_descriptor = o3d.data.DataDescriptor( + os.path.join(EMBODICHAIN_DOWNLOAD_PREFIX, obj_assets, "Drawer.zip"), + "eba30c852074388c2e5b634b1ae37572", + ) + prefix = type(self).__name__ + path = EMBODICHAIN_DEFAULT_DATA_ROOT if data_root is None else data_root + + super().__init__(prefix, data_descriptor, path) diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index 99cf83e04..58cf9fe13 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -33,6 +33,9 @@ AntipodalAffordance, AssembleAffordance, InteractionPoints, + PressAffordance, + SlideAffordance, + TwistAffordance, ) from .bindings import ActionBinding, ResolvedActionBinding, ResolvedControlPart from .control import ( @@ -109,6 +112,12 @@ Press, PressGoal, PressOptions, + Slide, + SlideGoal, + SlideOptions, + Twist, + TwistGoal, + TwistOptions, ) from .runner import ( CommandAcknowledgement, @@ -223,8 +232,16 @@ "PlanningContext", "PoseGoalValue", "Press", + "PressAffordance", "PressGoal", "PressOptions", + "SlideAffordance", + "Slide", + "SlideGoal", + "SlideOptions", + "Twist", + "TwistGoal", + "TwistOptions", "RecoveryPolicy", "RigidObjectSceneProvider", "RigidObjectSceneProviderCfg", @@ -247,5 +264,6 @@ "SimulationExecutionAdapter", "TaskState", "TimedTrajectory", + "TwistAffordance", "TrajectorySegment", ] diff --git a/embodichain/lab/sim/atomic_actions/affordance.py b/embodichain/lab/sim/atomic_actions/affordance.py index 206aa33ac..ad6b7c8af 100644 --- a/embodichain/lab/sim/atomic_actions/affordance.py +++ b/embodichain/lab/sim/atomic_actions/affordance.py @@ -63,7 +63,11 @@ def get_batch_size(self) -> int: @dataclass class AntipodalAffordance(Affordance): - """Antipodal grasp affordance for parallel-jaw grippers.""" + """Antipodal grasp affordance for parallel-jaw grippers. + + The affordance owns only target-local triangle-mesh data. Simulator entity + handles and live poses belong to scene grounding, not semantic geometry. + """ mesh_vertices: torch.Tensor | None = None """Object mesh vertices, shape [N, 3].""" @@ -172,15 +176,41 @@ def get_best_grasp_poses( [0, 0, -1], dtype=torch.float32 ), ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return the best antipodal grasp for each object pose. + + Args: + obj_poses: Batched object poses with shape ``(B, 4, 4)``. + approach_direction: One shared ``(3,)`` world-frame direction or + per-object directions with shape ``(B, 3)``. + + Returns: + A success mask, best grasp poses, and gripper opening lengths with + batch dimension ``B``. + + Raises: + ValueError: If ``approach_direction`` has an incompatible shape. + """ if self._generator is None: self._init_generator() approach_direction = self._resolve_approach_direction(approach_direction) + if approach_direction.shape == (3,): + approach_directions = approach_direction.unsqueeze(0).expand( + obj_poses.shape[0], -1 + ) + elif approach_direction.shape == (obj_poses.shape[0], 3): + approach_directions = approach_direction + else: + raise ValueError( + "approach_direction must have shape (3,) or " + f"({obj_poses.shape[0]}, 3), got " + f"{tuple(approach_direction.shape)}." + ) grasp_xpos_list: list[torch.Tensor] = [] is_success_list: list[bool] = [] open_length_list: list[float] = [] for i, obj_pose in enumerate(obj_poses): is_success, grasp_xpos, open_length = self._generator.get_grasp_poses( - obj_pose, approach_direction + obj_pose, approach_directions[i] ) if is_success: grasp_xpos_list.append(grasp_xpos.unsqueeze(0)) @@ -203,6 +233,277 @@ def get_best_grasp_poses( return is_success_t, grasp_xpos, open_length_t +@dataclass +class TwistAffordance(Affordance): + """Target-local grasp point and rotation-axis geometry for twisting.""" + + grasp_position: tuple[float, float, float] = field(kw_only=True) + """Explicit target-local center of the gripper contact region.""" + + axis_origin: tuple[float, float, float] = field(kw_only=True) + """Explicit point on the rotation axis in the target-local frame.""" + + twist_axis: torch.Tensor = field( + default_factory=lambda: torch.tensor([0.0, 1.0, 0.0]) + ) + """Twist axis expressed in the target object's local frame.""" + + joint_name: str | None = None + """Optional stable articulation-joint name associated with the axis.""" + + joint_limits: tuple[float, float] | None = None + """Optional lower and upper angular limits in radians.""" + + def __post_init__(self) -> None: + if ( + not isinstance(self.twist_axis, torch.Tensor) + or self.twist_axis.shape != (3,) + or not torch.isfinite(self.twist_axis).all() + ): + raise ValueError("TwistAffordance.twist_axis must be a finite (3,) tensor.") + if torch.linalg.vector_norm(self.twist_axis) <= 1.0e-6: + raise ValueError("TwistAffordance.twist_axis must be non-zero.") + self.twist_axis = self.twist_axis.clone() + self.grasp_position = _validate_local_point( + self.grasp_position, "TwistAffordance.grasp_position" + ) + self.axis_origin = _validate_local_point( + self.axis_origin, "TwistAffordance.axis_origin" + ) + _validate_joint_metadata(self.joint_name, self.joint_limits) + + def get_grasp_pose(self, target_pose: torch.Tensor) -> torch.Tensor: + """Construct a deterministic world grasp pose from local geometry. + + The pose z-axis follows :attr:`twist_axis`. The remaining axes are + formed with an adaptive reference so the result is always in SO(3). + + Returns: + Batched world-frame grasp poses with shape ``(B, 4, 4)``. + + Raises: + ValueError: If ``target_pose`` is not a batched pose tensor. + """ + if target_pose.dim() != 3 or target_pose.shape[1:] != (4, 4): + raise ValueError("Target pose must have shape (B, 4, 4).") + target_pose = target_pose.to(dtype=torch.float32) + device = target_pose.device + twist_axis = self.twist_axis.to(device=device, dtype=torch.float32) + twist_axis = twist_axis / torch.linalg.vector_norm(twist_axis) + + z_axis = torch.matmul(target_pose[:, :3, :3], twist_axis) + z_axis = torch.nn.functional.normalize(z_axis, dim=1) + x_axis, y_axis = _orthogonal_xy_from_z(z_axis) + + grasp_pose = torch.eye(4, dtype=torch.float32, device=device).repeat( + target_pose.shape[0], 1, 1 + ) + grasp_pose[:, :3, 0] = x_axis + grasp_pose[:, :3, 1] = y_axis + grasp_pose[:, :3, 2] = z_axis + local_grasp = torch.tensor( + self.grasp_position, dtype=torch.float32, device=device + ) + grasp_pose[:, :3, 3] = ( + torch.matmul(target_pose[:, :3, :3], local_grasp) + target_pose[:, :3, 3] + ) + return grasp_pose + + +@dataclass +class SlideAffordance(AntipodalAffordance): + """Target-local antipodal grasp and translation-axis geometry. + + The positive translation-axis direction denotes approaching and pushing + the articulated part closed. Pulling moves in the opposite direction. + The mesh describes the actual graspable contact surface. The target pose is + supplied separately by :class:`~.goals.SceneEntityPose` or a pose snapshot. + """ + + mesh_vertices: torch.Tensor = field(kw_only=True) + """Target-local vertices for the graspable contact surface.""" + + mesh_triangles: torch.Tensor = field(kw_only=True) + """Triangle indices for the graspable contact surface.""" + + translation_axis: torch.Tensor = field( + default_factory=lambda: torch.tensor([0.0, 1.0, 0.0]) + ) + """Approach and push/close direction in the articulation-link frame.""" + + joint_name: str | None = None + """Optional stable prismatic-joint name associated with the link.""" + + joint_limits: tuple[float, float] | None = None + """Optional lower and upper translation limits in metres.""" + + def __post_init__(self) -> None: + if self.mesh_vertices.dim() != 2 or self.mesh_vertices.shape[1] != 3: + raise ValueError("SlideAffordance.mesh_vertices must have shape (N, 3).") + if ( + self.mesh_vertices.shape[0] == 0 + or not torch.isfinite(self.mesh_vertices).all() + ): + raise ValueError( + "SlideAffordance.mesh_vertices must be finite and non-empty." + ) + if self.mesh_triangles.dim() != 2 or self.mesh_triangles.shape[1] != 3: + raise ValueError("SlideAffordance.mesh_triangles must have shape (M, 3).") + if ( + not isinstance(self.translation_axis, torch.Tensor) + or self.translation_axis.shape != (3,) + or not torch.isfinite(self.translation_axis).all() + ): + raise ValueError( + "SlideAffordance.translation_axis must be a finite (3,) tensor." + ) + if torch.linalg.vector_norm(self.translation_axis) <= 1.0e-6: + raise ValueError("SlideAffordance.translation_axis must be non-zero.") + self.translation_axis = self.translation_axis.clone() + _validate_joint_metadata(self.joint_name, self.joint_limits) + + +@dataclass +class PressAffordance(Affordance): + """Explicit target-local contact point and pressing direction.""" + + press_axis: torch.Tensor = field( + default_factory=lambda: torch.tensor([0.0, 0.0, 1.0]) + ) + """Press direction expressed in the target object's local frame.""" + + press_position: tuple[float, float, float] = field(kw_only=True) + """Explicit local-frame point on the pressable contact surface.""" + + def __post_init__(self) -> None: + if ( + not isinstance(self.press_axis, torch.Tensor) + or self.press_axis.shape != (3,) + or not torch.isfinite(self.press_axis).all() + ): + raise ValueError("PressAffordance.press_axis must be a finite (3,) tensor.") + if torch.linalg.vector_norm(self.press_axis) <= 1.0e-6: + raise ValueError("PressAffordance.press_axis must be non-zero.") + self.press_axis = self.press_axis.clone() + self.press_position = _validate_local_point( + self.press_position, "PressAffordance.press_position" + ) + + def get_press_pose( + self, + target_pose: torch.Tensor, + press_position: tuple[float, float, float] | None = None, + ) -> torch.Tensor: + """Construct a press pose at the configured surface point. + + The end-effector z-axis follows :attr:`press_axis` in world space. An + adaptive reference produces an orthonormal, right-handed frame. + + Args: + target_pose: Current target world pose with shape ``(B, 4, 4)``. + press_position: Optional per-call exact local-frame press position. + It overrides :attr:`press_position`. + + Returns: + Batched world-frame press poses with shape ``(B, 4, 4)``. + + Raises: + ValueError: If an input has an invalid shape or value. + """ + if target_pose.dim() != 3 or target_pose.shape[1:] != (4, 4): + raise ValueError("Target pose must have shape (B, 4, 4).") + target_pose = target_pose.to(dtype=torch.float32) + device = target_pose.device + press_axis = self.press_axis.to(device=device, dtype=torch.float32) + press_axis = press_axis / torch.linalg.vector_norm(press_axis) + configured_position = self._validate_press_position( + press_position, + field_name="press_position", + ) + configured_position = ( + self.press_position if configured_position is None else configured_position + ) + local_press_position = torch.tensor( + configured_position, + dtype=torch.float32, + device=device, + ) + + z_axis = torch.matmul(target_pose[:, :3, :3], press_axis) + z_axis = torch.nn.functional.normalize(z_axis, dim=1) + x_axis, y_axis = _orthogonal_xy_from_z(z_axis) + + press_pose = torch.eye(4, dtype=torch.float32, device=device).repeat( + target_pose.shape[0], 1, 1 + ) + press_pose[:, :3, 0] = x_axis + press_pose[:, :3, 1] = y_axis + press_pose[:, :3, 2] = z_axis + press_pose[:, :3, 3] = ( + torch.matmul(target_pose[:, :3, :3], local_press_position) + + target_pose[:, :3, 3] + ) + return press_pose + + @staticmethod + def _validate_press_position( + value: tuple[float, float, float] | None, + *, + field_name: str, + ) -> tuple[float, float, float] | None: + """Validate and normalize an optional local-frame press position.""" + if value is None: + return None + position = torch.as_tensor(value, dtype=torch.float32) + if position.shape != (3,) or not torch.isfinite(position).all(): + raise ValueError(f"{field_name} must be a finite (x, y, z) tuple.") + return tuple(float(component) for component in position) + + +def _orthogonal_xy_from_z(z_axis: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Complete normalized z axes into right-handed orthonormal frames.""" + basis = torch.eye(3, dtype=z_axis.dtype, device=z_axis.device) + reference_indices = torch.argmin(torch.abs(z_axis), dim=1) + reference = basis[reference_indices] + y_axis = torch.nn.functional.normalize( + torch.linalg.cross(reference, z_axis, dim=1), dim=1 + ) + x_axis = torch.nn.functional.normalize( + torch.linalg.cross(y_axis, z_axis, dim=1), dim=1 + ) + return x_axis, y_axis + + +def _validate_local_point( + value: tuple[float, float, float], field_name: str +) -> tuple[float, float, float]: + """Validate and normalize one target-local 3D point.""" + point = torch.as_tensor(value, dtype=torch.float32) + if point.shape != (3,) or not torch.isfinite(point).all(): + raise ValueError(f"{field_name} must be a finite (x, y, z) tuple.") + return tuple(float(component) for component in point) + + +def _validate_joint_metadata( + joint_name: str | None, + joint_limits: tuple[float, float] | None, +) -> None: + """Validate optional articulation joint metadata.""" + if joint_name is not None and ( + not isinstance(joint_name, str) or not joint_name.strip() + ): + raise ValueError("joint_name must be a non-empty string when provided.") + if joint_limits is None: + return + limits = torch.as_tensor(joint_limits, dtype=torch.float32) + if ( + limits.shape != (2,) + or not torch.isfinite(limits).all() + or limits[0] > limits[1] + ): + raise ValueError("joint_limits must be finite and ordered (lower, upper).") + + @dataclass class InteractionPoints(Affordance): """Batch of 3D interaction points on an object surface.""" @@ -316,6 +617,9 @@ def get_assemble_object_pose(self, base_pose: torch.Tensor) -> torch.Tensor: __all__ = [ "Affordance", "AntipodalAffordance", + "SlideAffordance", + "PressAffordance", + "TwistAffordance", "InteractionPoints", "AssembleAffordance", ] diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index d4b7906bb..e33eafc77 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -152,6 +152,8 @@ class SkillDescriptor: manipulator_roles: tuple[str, ...] = () end_effector_roles: tuple[str, ...] = () agent_visible: bool = True + open_loop: bool = False + """Whether completion reports motion execution without physical-effect proof.""" binding_contract: SkillBindingContract | None = None """Explicit generic resource contract used by the semantic skill layer.""" @@ -169,6 +171,8 @@ def __post_init__(self) -> None: raise TypeError( "SkillDescriptor.options_type must be an ActionOptions subclass." ) + if not isinstance(self.open_loop, bool): + raise TypeError("SkillDescriptor.open_loop must be a bool.") for field_name in ("manipulator_roles", "end_effector_roles"): roles = tuple(getattr(self, field_name)) if len(set(roles)) != len(roles) or not all( @@ -214,6 +218,9 @@ class AtomicAction(Generic[GoalT, OptionsT], ABC): agent_visible: ClassVar[bool] = True """Whether an Action Agent should expose this skill by default.""" + open_loop: ClassVar[bool] = False + """Whether the skill intentionally declares no verified physical effect.""" + binding_contract: ClassVar[SkillBindingContract | None] = None """Explicit robot-independent requirements for semantic discovery. @@ -319,6 +326,7 @@ def descriptor(cls) -> SkillDescriptor: manipulator_roles=cls.manipulator_roles, end_effector_roles=cls.end_effector_roles, agent_visible=cls.agent_visible, + open_loop=cls.open_loop, binding_contract=cls.__dict__.get("binding_contract"), ) diff --git a/embodichain/lab/sim/atomic_actions/policies.py b/embodichain/lab/sim/atomic_actions/policies.py index 8f82b49a6..d326eb304 100644 --- a/embodichain/lab/sim/atomic_actions/policies.py +++ b/embodichain/lab/sim/atomic_actions/policies.py @@ -117,6 +117,7 @@ def to_motion_gen_options( start_qpos: "torch.Tensor", control_part: str, sample_count: int | None = None, + cartesian_linear: bool = False, ) -> "MotionGenOptions": """Translate this atomic policy into motion-generator options. @@ -124,6 +125,8 @@ def to_motion_gen_options( start_qpos: Observed controlled-joint start positions. control_part: Bound robot control-part name. sample_count: Optional segment-local sample-count override. + cartesian_linear: Whether every supplied Cartesian keyframe is a + required linear-path sample rather than a sparse endpoint. Returns: Independently owned options for :class:`MotionGenerator`. @@ -139,6 +142,8 @@ def to_motion_gen_options( control_part=control_part, plan_opts=self.plan_opts, is_interpolate=True, + is_linear=cartesian_linear, + preserve_cartesian_samples=cartesian_linear, ) diff --git a/embodichain/lab/sim/atomic_actions/primitives/__init__.py b/embodichain/lab/sim/atomic_actions/primitives/__init__.py index 85de2c985..718e6780b 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/__init__.py +++ b/embodichain/lab/sim/atomic_actions/primitives/__init__.py @@ -44,6 +44,12 @@ from .pick_up import GraspGoal, PickUp, PickUpOptions from .place import AssembleGoal, Place, PlaceGoal, PlaceOptions from .press import Press, PressGoal, PressOptions +from .slide import ( + Slide, + SlideGoal, + SlideOptions, +) +from .twist import Twist, TwistGoal, TwistOptions BUILTIN_ACTION_TYPES: tuple[type[AtomicAction], ...] = ( MoveEndEffector, @@ -52,6 +58,8 @@ MoveHeldObject, Place, Press, + Slide, + Twist, CoordinatedPickment, CoordinatedPlacement, HandOver, @@ -87,4 +95,10 @@ "Press", "PressGoal", "PressOptions", + "Slide", + "SlideGoal", + "SlideOptions", + "Twist", + "TwistGoal", + "TwistOptions", ] diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py index 7829f72ec..26d148ab1 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py @@ -26,33 +26,45 @@ from embodichain.utils import logger from embodichain.utils.math import matrix_from_quat, pose_inv, quat_from_matrix -from ..affordance import AntipodalAffordance -from ..bindings import ResolvedControlPart -from ..control import GRASP_COMMAND, OPEN_COMMAND, JointPositionCommand -from ..core import AtomicAction, ObjectSemantics -from ..effects import StateDelta -from ..goals import ( +from embodichain.lab.sim.atomic_actions.affordance import AntipodalAffordance +from embodichain.lab.sim.atomic_actions.bindings import ResolvedControlPart +from embodichain.lab.sim.atomic_actions.control import ( + GRASP_COMMAND, + OPEN_COMMAND, + JointPositionCommand, +) +from embodichain.lab.sim.atomic_actions.core import AtomicAction, ObjectSemantics +from embodichain.lab.sim.atomic_actions.effects import StateDelta +from embodichain.lab.sim.atomic_actions.goals import ( ObjectActionGoal, PoseGoalValue, _resolve_object_pose, resolve_pose_goal, validate_pose_goal, ) -from ..invocation import ActionOptions, ResolvedActionRequest -from ..plans import ActionPlan, normalize_success_mask -from ..requirements import ( +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan, normalize_success_mask +from embodichain.lab.sim.atomic_actions.requirements import ( DisjointResourceSlots, INVERSE_KINEMATICS_CAPABILITY, SkillBindingContract, ) -from ..state import HeldObjectState, PlanningContext -from ..trajectory_ops import interpolate_joint_trajectory, translate_pose_world -from ._helpers import ( +from embodichain.lab.sim.atomic_actions.primitives._binding_contracts import ( + make_manipulation_slot, +) +from embodichain.lab.sim.atomic_actions.primitives._helpers import ( assemble_full_robot_trajectory, repeat_qpos, resolve_batched_pose, ) -from ._binding_contracts import make_manipulation_slot +from embodichain.lab.sim.atomic_actions.state import HeldObjectState, PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( + interpolate_joint_trajectory, + translate_pose_world, +) @dataclass(frozen=True, slots=True, eq=False) diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py index 9db4563f7..a6f366d05 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py @@ -25,31 +25,44 @@ from embodichain.utils import logger -from ..bindings import ResolvedControlPart -from ..control import GRASP_COMMAND, OPEN_COMMAND, JointPositionCommand -from ..core import AtomicAction -from ..effects import StateDelta -from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal -from ..invocation import ActionOptions, ResolvedActionRequest -from ..plans import ActionPlan, normalize_success_mask -from ..requirements import ( +from embodichain.lab.sim.atomic_actions.bindings import ResolvedControlPart +from embodichain.lab.sim.atomic_actions.control import ( + GRASP_COMMAND, + OPEN_COMMAND, + JointPositionCommand, +) +from embodichain.lab.sim.atomic_actions.core import AtomicAction +from embodichain.lab.sim.atomic_actions.effects import StateDelta +from embodichain.lab.sim.atomic_actions.goals import ( + PoseGoalValue, + resolve_pose_goal, + validate_pose_goal, +) +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan, normalize_success_mask +from embodichain.lab.sim.atomic_actions.requirements import ( CARTESIAN_POSE_CAPABILITY, DisjointResourceSlots, SkillBindingContract, ) -from ..state import HeldObjectState, PlanningContext -from ..trajectory_ops import ( - interpolate_hand_qpos, - translate_pose_world, +from embodichain.lab.sim.atomic_actions.primitives._binding_contracts import ( + make_manipulation_slot, ) -from ._helpers import ( +from embodichain.lab.sim.atomic_actions.primitives._helpers import ( assemble_full_robot_trajectory, plan_named_arm_trajectory, repeat_qpos, resolve_batched_pose, resolve_object_target, ) -from ._binding_contracts import make_manipulation_slot +from embodichain.lab.sim.atomic_actions.state import HeldObjectState, PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( + interpolate_hand_qpos, + translate_pose_world, +) @dataclass(frozen=True, slots=True, eq=False) diff --git a/embodichain/lab/sim/atomic_actions/primitives/hand_over.py b/embodichain/lab/sim/atomic_actions/primitives/hand_over.py index 04b635c04..ec325aa99 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/hand_over.py +++ b/embodichain/lab/sim/atomic_actions/primitives/hand_over.py @@ -26,31 +26,44 @@ from embodichain.utils import logger from embodichain.utils.math import pose_inv -from ..bindings import ResolvedControlPart -from ..control import GRASP_COMMAND, OPEN_COMMAND, JointPositionCommand -from ..core import AtomicAction, ObjectSemantics, _same_object_identity -from ..effects import StateDelta -from ..invocation import ActionOptions, ResolvedActionRequest -from ..plans import ActionPlan, normalize_success_mask -from ..requirements import ( +from embodichain.lab.sim.atomic_actions.bindings import ResolvedControlPart +from embodichain.lab.sim.atomic_actions.control import ( + GRASP_COMMAND, + OPEN_COMMAND, + JointPositionCommand, +) +from embodichain.lab.sim.atomic_actions.core import ( + AtomicAction, + ObjectSemantics, + _same_object_identity, +) +from embodichain.lab.sim.atomic_actions.effects import StateDelta +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan, normalize_success_mask +from embodichain.lab.sim.atomic_actions.requirements import ( CARTESIAN_POSE_CAPABILITY, DisjointResourceSlots, FORWARD_KINEMATICS_CAPABILITY, SkillBindingContract, ) -from ..state import HeldObjectState, PlanningContext -from ..trajectory_ops import ( - interpolate_hand_qpos, - translate_pose_world, +from embodichain.lab.sim.atomic_actions.primitives._binding_contracts import ( + make_manipulation_slot, ) -from ._helpers import ( +from embodichain.lab.sim.atomic_actions.primitives._helpers import ( assemble_full_robot_trajectory, plan_named_arm_trajectory, repeat_qpos, resolve_batched_pose, ) -from ._binding_contracts import make_manipulation_slot -from .pick_up import GraspGoal +from embodichain.lab.sim.atomic_actions.primitives.pick_up import GraspGoal +from embodichain.lab.sim.atomic_actions.state import HeldObjectState, PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( + interpolate_hand_qpos, + translate_pose_world, +) @dataclass(frozen=True, slots=True, eq=False) 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 78a35fcde..4ece1aed4 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_end_effector.py @@ -23,18 +23,30 @@ import torch -from ..core import AtomicAction -from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal -from ..invocation import ActionOptions, ResolvedActionRequest -from ..plans import ActionPlan -from ..requirements import CARTESIAN_POSE_CAPABILITY, SkillBindingContract -from ..state import PlanningContext -from ..trajectory_ops import ( +from embodichain.lab.sim.atomic_actions.core import AtomicAction +from embodichain.lab.sim.atomic_actions.goals import ( + PoseGoalValue, + resolve_pose_goal, + validate_pose_goal, +) +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan +from embodichain.lab.sim.atomic_actions.requirements import ( + CARTESIAN_POSE_CAPABILITY, + SkillBindingContract, +) +from embodichain.lab.sim.atomic_actions.primitives._binding_contracts import ( + make_motion_slot, +) +from embodichain.lab.sim.atomic_actions.state import PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( build_pose_plan_states, resolve_pose_target, to_full_robot_trajectory, ) -from ._binding_contracts import make_motion_slot @dataclass(frozen=True, slots=True, eq=False) 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 0f8d014e0..ce26168f0 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_held_object.py @@ -29,20 +29,35 @@ pose_inv, ) -from ._helpers import arm_qpos_from_state, resolve_object_target -from ..control import GRASP_COMMAND, JointPositionCommand -from ..core import AtomicAction -from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal -from ..invocation import ActionOptions, ResolvedActionRequest -from ..plans import ActionPlan -from ..requirements import ( +from embodichain.lab.sim.atomic_actions.control import ( + GRASP_COMMAND, + JointPositionCommand, +) +from embodichain.lab.sim.atomic_actions.core import AtomicAction +from embodichain.lab.sim.atomic_actions.goals import ( + PoseGoalValue, + resolve_pose_goal, + validate_pose_goal, +) +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan +from embodichain.lab.sim.atomic_actions.requirements import ( CARTESIAN_POSE_CAPABILITY, FORWARD_KINEMATICS_CAPABILITY, SkillBindingContract, ) -from ..state import PlanningContext -from ..trajectory_ops import build_pose_plan_states -from ._binding_contracts import make_manipulation_slot +from embodichain.lab.sim.atomic_actions.primitives._binding_contracts import ( + make_manipulation_slot, +) +from embodichain.lab.sim.atomic_actions.primitives._helpers import ( + arm_qpos_from_state, + resolve_object_target, +) +from embodichain.lab.sim.atomic_actions.state import PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import build_pose_plan_states @dataclass(frozen=True, slots=True, eq=False) diff --git a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py index 60d9383da..278745c64 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py @@ -23,17 +23,25 @@ import torch -from ..core import AtomicAction -from ..invocation import ActionOptions, ResolvedActionRequest -from ..plans import ActionPlan -from ..requirements import JOINT_POSITION_CAPABILITY, SkillBindingContract -from ..state import PlanningContext -from ..trajectory_ops import ( +from embodichain.lab.sim.atomic_actions.core import AtomicAction +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan +from embodichain.lab.sim.atomic_actions.requirements import ( + JOINT_POSITION_CAPABILITY, + SkillBindingContract, +) +from embodichain.lab.sim.atomic_actions.primitives._binding_contracts import ( + make_motion_slot, +) +from embodichain.lab.sim.atomic_actions.state import PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( build_joint_plan_states, resolve_joint_target, to_full_robot_trajectory, ) -from ._binding_contracts import make_motion_slot @dataclass(frozen=True, slots=True, eq=False) diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index dea984d1d..ffbd145bc 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -32,37 +32,46 @@ quat_from_matrix, ) -from ._helpers import arm_qpos_from_state -from ..affordance import AntipodalAffordance -from ..bindings import ResolvedControlPart -from ..control import GRASP_COMMAND, OPEN_COMMAND, JointPositionCommand -from ..core import AtomicAction, ObjectSemantics -from ..effects import StateDelta -from ..goals import ( +from embodichain.lab.sim.atomic_actions.affordance import AntipodalAffordance +from embodichain.lab.sim.atomic_actions.bindings import ResolvedControlPart +from embodichain.lab.sim.atomic_actions.control import ( + GRASP_COMMAND, + OPEN_COMMAND, + JointPositionCommand, +) +from embodichain.lab.sim.atomic_actions.core import AtomicAction, ObjectSemantics +from embodichain.lab.sim.atomic_actions.effects import StateDelta +from embodichain.lab.sim.atomic_actions.goals import ( ObjectActionGoal, PoseGoalValue, _resolve_object_pose, resolve_pose_goal, validate_pose_goal, ) -from ..invocation import ActionOptions, ResolvedActionRequest -from ..plans import ActionPlan, normalize_success_mask -from ..policies import MotionPolicy -from ..requirements import ( +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan, normalize_success_mask +from embodichain.lab.sim.atomic_actions.policies import MotionPolicy +from embodichain.lab.sim.atomic_actions.primitives._binding_contracts import ( + make_manipulation_slot, +) +from embodichain.lab.sim.atomic_actions.primitives._helpers import arm_qpos_from_state +from embodichain.lab.sim.atomic_actions.requirements import ( BATCH_INVERSE_KINEMATICS_CAPABILITY, CARTESIAN_POSE_CAPABILITY, FORWARD_KINEMATICS_CAPABILITY, SkillBindingContract, ) -from ..state import HeldObjectState, PlanningContext -from ..trajectory_ops import ( +from embodichain.lab.sim.atomic_actions.state import HeldObjectState, PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( build_pose_plan_states, interpolate_hand_qpos, resolve_pose_target, split_three_segments, translate_pose_world, ) -from ._binding_contracts import make_manipulation_slot @dataclass(frozen=True, slots=True, eq=False) diff --git a/embodichain/lab/sim/atomic_actions/primitives/place.py b/embodichain/lab/sim/atomic_actions/primitives/place.py index 6bee6fd38..fe7a9c60e 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/place.py +++ b/embodichain/lab/sim/atomic_actions/primitives/place.py @@ -26,32 +26,44 @@ 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 ..control import GRASP_COMMAND, OPEN_COMMAND, JointPositionCommand -from ..core import AtomicAction -from ..effects import StateDelta -from ..goals import ( +from embodichain.lab.sim.atomic_actions.affordance import AssembleAffordance +from embodichain.lab.sim.atomic_actions.control import ( + GRASP_COMMAND, + OPEN_COMMAND, + JointPositionCommand, +) +from embodichain.lab.sim.atomic_actions.core import AtomicAction +from embodichain.lab.sim.atomic_actions.effects import StateDelta +from embodichain.lab.sim.atomic_actions.goals import ( PoseGoalValue, SceneEntityPose, resolve_pose_goal, validate_pose_goal, ) -from ..invocation import ActionOptions, ResolvedActionRequest -from ..plans import ActionPlan -from ..requirements import ( +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan +from embodichain.lab.sim.atomic_actions.primitives._helpers import ( + arm_qpos_from_state, + resolve_object_target, +) +from embodichain.lab.sim.atomic_actions.requirements import ( CARTESIAN_POSE_CAPABILITY, FORWARD_KINEMATICS_CAPABILITY, SkillBindingContract, ) -from ..state import PlanningContext -from ..trajectory_ops import ( +from embodichain.lab.sim.atomic_actions.primitives._binding_contracts import ( + make_manipulation_slot, +) +from embodichain.lab.sim.atomic_actions.state import PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( build_pose_plan_states, interpolate_hand_qpos, resolve_pose_target, split_three_segments, ) -from ._binding_contracts import make_manipulation_slot TcpSymmetry = Literal["none", "z_roll_180"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/press.py b/embodichain/lab/sim/atomic_actions/primitives/press.py index 4310451f7..8c4f860b8 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/press.py +++ b/embodichain/lab/sim/atomic_actions/primitives/press.py @@ -18,63 +18,111 @@ from __future__ import annotations +import math from dataclasses import dataclass from typing import ClassVar import torch -from ._helpers import arm_qpos_from_state -from ..control import GRASP_COMMAND, JointPositionCommand -from ..core import AtomicAction -from ..goals import PoseGoalValue, resolve_pose_goal, validate_pose_goal -from ..invocation import ActionOptions, ResolvedActionRequest -from ..plans import ActionPlan -from ..requirements import ( +from embodichain.lab.sim.atomic_actions.affordance import PressAffordance +from embodichain.lab.sim.atomic_actions.control import ( + GRASP_COMMAND, + JointPositionCommand, +) +from embodichain.lab.sim.atomic_actions.core import AtomicAction, ObjectSemantics +from embodichain.lab.sim.atomic_actions.effects import StateDelta +from embodichain.lab.sim.atomic_actions.goals import ( + ObjectActionGoal, + PoseGoalValue, + resolve_pose_goal, + validate_pose_goal, +) +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan +from embodichain.lab.sim.atomic_actions.primitives._binding_contracts import ( + make_manipulation_slot, +) +from embodichain.lab.sim.atomic_actions.primitives._helpers import arm_qpos_from_state +from embodichain.lab.sim.atomic_actions.requirements import ( CARTESIAN_POSE_CAPABILITY, - JOINT_POSITION_CAPABILITY, + FORWARD_KINEMATICS_CAPABILITY, SkillBindingContract, ) -from ..state import PlanningContext -from ..trajectory_ops import ( - build_joint_plan_states, +from embodichain.utils.math import get_relative_rotation +from embodichain.lab.sim.atomic_actions.state import PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( + axis_translation_keyframes, build_pose_plan_states, interpolate_hand_qpos, resolve_pose_target, + translate_pose_world, ) -from ._binding_contracts import make_manipulation_slot @dataclass(frozen=True, slots=True, eq=False) -class PressGoal: - """Single end-effector contact pose used by :class:`Press`.""" +class PressGoal(ObjectActionGoal): + """Target object described by a press affordance.""" - xpos: PoseGoalValue - """Contact pose, shape ``(4, 4)`` or ``(num_envs, 4, 4)``.""" + goal_kind: ClassVar[str] = "press" + + target_pose: PoseGoalValue + """Target pose snapshot or late-bound stable scene-entity reference.""" def __post_init__(self) -> None: - validate_pose_goal(self.xpos, "xpos", allow_waypoints=False) + ObjectActionGoal.__post_init__(self) + validate_pose_goal(self.target_pose, "target_pose", allow_waypoints=False) @dataclass(frozen=True, slots=True, eq=False) class PressOptions(ActionOptions): - """Per-invocation press behavior.""" + """Per-invocation pressing behavior.""" hand_interp_steps: int = 5 - """Number of waypoints for closing the gripper before pressing.""" + """Number of waypoints used to close the hand.""" + + approach_distance: float = 0.1 + """Distance from the press position opposite the press direction.""" + + press_distance: float = 0.05 + """Distance traveled into the target along its press axis.""" + + press_position: tuple[float, float, float] | None = None + """Optional local-frame position overriding the affordance press position.""" def __post_init__(self) -> None: if self.hand_interp_steps < 1: raise ValueError("hand_interp_steps must be at least 1.") + if not math.isfinite(self.approach_distance): + raise ValueError("approach_distance must be finite.") + if self.approach_distance < 0.0: + raise ValueError("approach_distance must be non-negative.") + if not math.isfinite(self.press_distance): + raise ValueError("press_distance must be finite.") + if self.press_distance <= 0.0: + raise ValueError("press_distance must be positive.") + if self.press_position is not None: + position = torch.as_tensor(self.press_position, dtype=torch.float32) + if position.shape != (3,) or not torch.isfinite(position).all(): + raise ValueError("press_position must be a finite (x, y, z) tuple.") + object.__setattr__( + self, + "press_position", + tuple(float(component) for component in position), + ) class Press(AtomicAction[PressGoal, PressOptions]): - """Close the gripper, press down to a target pose, then return.""" + """Open-loop motion primitive that approaches, presses, and retracts.""" 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",) + open_loop: ClassVar[bool] = True binding_contract: ClassVar[SkillBindingContract] = SkillBindingContract( slots=( make_manipulation_slot( @@ -82,7 +130,7 @@ class Press(AtomicAction[PressGoal, PressOptions]): motion_capabilities=frozenset( { CARTESIAN_POSE_CAPABILITY, - JOINT_POSITION_CAPABILITY, + FORWARD_KINEMATICS_CAPABILITY, } ), grasp_commands={GRASP_COMMAND: JointPositionCommand}, @@ -90,121 +138,218 @@ class Press(AtomicAction[PressGoal, PressOptions]): ), ) + def __init__(self, default_options: PressOptions | None = None) -> None: + super().__init__(default_options) + + def _on_bind(self) -> None: + """Resolve dimensions owned by the engine's robot.""" + self.num_envs = self.robot.get_qpos().shape[0] + self.robot_dof = self.robot.dof + + def _find_symmetric_nearest_xpos( + self, target_xpos: torch.Tensor, reference_xpos: torch.Tensor + ) -> torch.Tensor: + """Find the nearest symmetric pose to the reference pose.""" + symmetric_xpos = target_xpos.clone() + symmetric_xpos[:, :3, 0] = -symmetric_xpos[:, :3, 0] + symmetric_xpos[:, :3, 1] = -symmetric_xpos[:, :3, 1] + angle_a = get_relative_rotation( + reference_xpos[:, :3, :3], target_xpos[:, :3, :3] + ) + angle_b = get_relative_rotation( + reference_xpos[:, :3, :3], symmetric_xpos[:, :3, :3] + ) + choose_target = (angle_a < angle_b)[..., None, None] + target_xpos = torch.where(choose_target, target_xpos, symmetric_xpos) + return target_xpos + def _plan( self, request: ResolvedActionRequest[PressGoal, PressOptions], context: PlanningContext, ) -> ActionPlan: - """Plan a close, press, and retract sequence.""" - target = request.goal + """Plan close, approach, press, and retract without stepping simulation.""" + target = self.require_goal(request) + affordance = self._require_press_affordance(target.semantics) options = request.skill_options - binding = request.binding - manipulator = binding.manipulator() - end_effector = binding.end_effector() - control_part = manipulator.name + manipulator = request.binding.manipulator() + end_effector = request.binding.end_effector() arm_joint_ids = list(manipulator.joint_ids) hand_joint_ids = list(end_effector.joint_ids) - hand_close_qpos = end_effector.joint_positions( + start_arm_qpos = arm_qpos_from_state(context, arm_joint_ids) + start_hand_qpos = context.last_qpos[:, hand_joint_ids] + hand_grasp_qpos = end_effector.joint_positions( GRASP_COMMAND, - num_envs=self.num_envs, + num_envs=context.batch_size, device=self.device, dtype=context.robot.qpos.dtype, ) - state = context - press_xpos = resolve_pose_target( - resolve_pose_goal(target.xpos, context, name="xpos"), + + target_pose = resolve_pose_target( + resolve_pose_goal(target.target_pose, context, name="target_pose"), num_envs=self.num_envs, device=self.device, ) - start_arm_qpos = arm_qpos_from_state(state, arm_joint_ids) - start_hand_qpos = state.last_qpos[:, hand_joint_ids] - - n_close, n_down, n_back = self._compute_segment_waypoints( - request.motion_policy.sample_count, options + contact_xpos = affordance.get_press_pose( + target_pose, + press_position=options.press_position, + ).to(device=self.device, dtype=torch.float32) + contact_xpos = self._find_symmetric_nearest_xpos( + contact_xpos, + reference_xpos=self.robot.compute_fk( + qpos=start_arm_qpos, name=manipulator.name, to_matrix=True + ), ) - - hand_close_path = interpolate_hand_qpos( + approach_xpos = translate_pose_world( + contact_xpos, + -contact_xpos[:, :3, 2] * options.approach_distance, + ) + pressed_xpos = translate_pose_world( + contact_xpos, + contact_xpos[:, :3, 2] * options.press_distance, + ) + n_approach, n_contact, n_press, n_retract = self._motion_segment_lengths( + request.motion_policy.sample_count, + options.hand_interp_steps, + ) + hand_close = interpolate_hand_qpos( start_hand_qpos, - hand_close_qpos, - n_waypoints=n_close, + hand_grasp_qpos, + n_waypoints=options.hand_interp_steps, ) - - down_result = self.motion_generator.generate( - build_pose_plan_states(press_xpos), - options=request.motion_policy.to_motion_gen_options( - start_qpos=start_arm_qpos, - control_part=control_part, - sample_count=n_down, - ), + approach_success, approach_arm = self._plan_pose_segment( + approach_xpos, + start_arm_qpos, + manipulator.name, + request, + n_approach, ) - assert isinstance(down_result.success, torch.Tensor) - assert down_result.positions is not None - down_success = down_result.success - down_arm = down_result.positions - - press_arm_qpos = down_arm[:, -1, :] - back_result = self.motion_generator.generate( - build_joint_plan_states(start_arm_qpos), - options=request.motion_policy.to_motion_gen_options( - start_qpos=press_arm_qpos, - control_part=control_part, - sample_count=n_back, - ), + contact_keyframes = axis_translation_keyframes( + approach_xpos, + contact_xpos, + contact_xpos[:, :3, 2], + n_waypoints=n_contact - 1, + ) + contact_success, contact_arm = self._plan_pose_segment( + contact_keyframes, + approach_arm[:, -1], + manipulator.name, + request, + n_contact, + cartesian_linear=True, + ) + press_keyframes = axis_translation_keyframes( + contact_xpos, + pressed_xpos, + contact_xpos[:, :3, 2], + n_waypoints=n_press - 1, + ) + press_success, press_arm = self._plan_pose_segment( + press_keyframes, + contact_arm[:, -1], + manipulator.name, + request, + n_press, + cartesian_linear=True, + ) + retract_keyframes = axis_translation_keyframes( + pressed_xpos, + approach_xpos, + contact_xpos[:, :3, 2], + n_waypoints=n_retract - 1, + ) + retract_success, retract_arm = self._plan_pose_segment( + retract_keyframes, + press_arm[:, -1], + manipulator.name, + request, + n_retract, + cartesian_linear=True, ) - assert isinstance(back_result.success, torch.Tensor) - assert back_result.positions is not None - back_success = back_result.success - back_arm = back_result.positions - success = down_success & back_success + success = approach_success & contact_success & press_success & retract_success - # Allocate from the actually returned segment lengths so collision-aware - # planners (which preserve their own sample count) are accommodated. - n_down_actual = down_arm.shape[1] - n_back_actual = back_arm.shape[1] + parts = (hand_close, approach_arm, contact_arm, press_arm, retract_arm) + lengths = tuple(part.shape[1] for part in parts) full = torch.empty( - (self.num_envs, n_close + n_down_actual + n_back_actual, self.robot_dof), - dtype=torch.float32, + (self.num_envs, sum(lengths), self.robot_dof), + dtype=context.robot.qpos.dtype, device=self.device, ) - full[:, :, :] = state.last_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 :, arm_joint_ids] = back_arm - full[:, n_close + n_down_actual :, hand_joint_ids] = hand_close_qpos.unsqueeze( - 1 - ) + full[:] = context.last_qpos.unsqueeze(1) + offset = 0 + + stop = offset + hand_close.shape[1] + full[:, offset:stop, arm_joint_ids] = start_arm_qpos.unsqueeze(1) + full[:, offset:stop, hand_joint_ids] = hand_close + offset = stop + + for arm in (approach_arm, contact_arm, press_arm, retract_arm): + stop = offset + arm.shape[1] + full[:, offset:stop, arm_joint_ids] = arm + full[:, offset:stop, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1) + offset = stop return self.build_plan( request, context, success=success, trajectory=full, + expected_effects=StateDelta(), segment_lengths={ - "close": n_close, - "press": n_down_actual, - "retract": n_back_actual, + "close": lengths[0], + "approach": lengths[1], + "contact": lengths[2], + "press": lengths[3], + "retract": lengths[4], }, ) - def _compute_segment_waypoints( - self, sample_count: int, options: PressOptions - ) -> tuple[int, int, int]: - """Split the invocation sample budget across press segments.""" - n_close = options.hand_interp_steps + @staticmethod + def _require_press_affordance( + semantics: ObjectSemantics, + ) -> PressAffordance: + affordance = semantics.affordance + if not isinstance(affordance, PressAffordance): + raise ValueError("Press requires a PressAffordance.") + return affordance - 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: + @staticmethod + def _motion_segment_lengths( + sample_count: int, + hand_interp_steps: int, + ) -> tuple[int, int, int, int]: + motion_count = sample_count - hand_interp_steps + if motion_count < 8: raise ValueError( - "Not enough waypoints for press trajectory. Increase " - "MotionPolicy.sample_count or decrease hand_interp_steps." + "Not enough waypoints for Press. Increase sample_count or " + "decrease hand_interp_steps." ) - return n_close, n_down, n_back + base, remainder = divmod(motion_count, 4) + values = [base + (index < remainder) for index in range(4)] + return values[0], values[1], values[2], values[3] + + def _plan_pose_segment( + self, + target_pose: torch.Tensor, + start_qpos: torch.Tensor, + control_part: str, + request: ResolvedActionRequest[PressGoal, PressOptions], + sample_count: int, + *, + cartesian_linear: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + result = self.motion_generator.generate( + build_pose_plan_states(target_pose), + options=request.motion_policy.to_motion_gen_options( + start_qpos=start_qpos, + control_part=control_part, + sample_count=sample_count, + cartesian_linear=cartesian_linear, + ), + ) + assert isinstance(result.success, torch.Tensor) + assert result.positions is not None + return result.success, result.positions __all__ = ["Press", "PressGoal", "PressOptions"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/slide.py b/embodichain/lab/sim/atomic_actions/primitives/slide.py new file mode 100644 index 000000000..c305a737e --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/slide.py @@ -0,0 +1,405 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Slide atomic action implementation.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import ClassVar, Literal + +import torch + +from embodichain.lab.sim.atomic_actions.affordance import SlideAffordance +from embodichain.lab.sim.atomic_actions.control import ( + GRASP_COMMAND, + OPEN_COMMAND, + JointPositionCommand, +) +from embodichain.lab.sim.atomic_actions.core import AtomicAction, ObjectSemantics +from embodichain.lab.sim.atomic_actions.effects import StateDelta +from embodichain.lab.sim.atomic_actions.goals import ( + ObjectActionGoal, + PoseGoalValue, + resolve_pose_goal, + validate_pose_goal, +) +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan, normalize_success_mask +from embodichain.lab.sim.atomic_actions.primitives._helpers import arm_qpos_from_state +from embodichain.lab.sim.atomic_actions.requirements import ( + ActionBindingRoute, + CARTESIAN_POSE_CAPABILITY, + DisjointSlotEndpoints, + GRASP_CAPABILITY, + SkillBindingContract, + SkillEndpointRequirement, + SkillResourceSlot, +) +from embodichain.lab.sim.atomic_actions.state import PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( + axis_translation_keyframes, + build_pose_plan_states, + interpolate_hand_qpos, + resolve_pose_target, + translate_pose_world, +) + + +@dataclass(frozen=True, slots=True, eq=False) +class SlideGoal(ObjectActionGoal): + """Translating articulation link described by a slide affordance.""" + + goal_kind: ClassVar[str] = "slide" + + target_pose: PoseGoalValue + """Link pose snapshot or late-bound stable scene-entity reference.""" + + def __post_init__(self) -> None: + ObjectActionGoal.__post_init__(self) + validate_pose_goal(self.target_pose, "target_pose", allow_waypoints=False) + + +@dataclass(frozen=True, slots=True, eq=False) +class SlideOptions(ActionOptions): + """Per-invocation sliding behavior for a translating articulation link.""" + + direction: Literal["pull", "push"] = "pull" + """Whether to pull the part open or push it closed.""" + + hand_interp_steps: int = 5 + """Number of waypoints used for each close/open hand segment.""" + + approach_distance: float = 0.1 + """Pre-grasp distance opposite the approach/push axis.""" + + translation_distance: float = 0.15 + """Distance traveled along the pull or push direction.""" + + def __post_init__(self) -> None: + if self.direction not in ("pull", "push"): + raise ValueError("direction must be either 'pull' or 'push'.") + if self.hand_interp_steps < 1: + raise ValueError("hand_interp_steps must be at least 1.") + if not math.isfinite(self.approach_distance): + raise ValueError("approach_distance must be finite.") + if self.approach_distance < 0.0: + raise ValueError("approach_distance must be non-negative.") + if not math.isfinite(self.translation_distance): + raise ValueError("translation_distance must be finite.") + if self.translation_distance <= 0.0: + raise ValueError("translation_distance must be positive.") + + +class Slide(AtomicAction[SlideGoal, SlideOptions]): + """Open-loop approach, grasp, and axis-constrained sliding motion.""" + + skill_id: ClassVar[str] = "slide" + GoalType: ClassVar[type] = SlideGoal + OptionsType: ClassVar[type] = SlideOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) + open_loop: ClassVar[bool] = True + binding_contract: ClassVar[SkillBindingContract] = SkillBindingContract( + slots=( + SkillResourceSlot( + slot_id="primary", + endpoints=( + SkillEndpointRequirement( + endpoint_id="motion", + capabilities=frozenset({CARTESIAN_POSE_CAPABILITY}), + route=ActionBindingRoute("manipulator", "primary"), + ), + SkillEndpointRequirement( + endpoint_id="grasp", + capabilities=frozenset({GRASP_CAPABILITY}), + required_commands={ + OPEN_COMMAND: JointPositionCommand, + GRASP_COMMAND: JointPositionCommand, + }, + route=ActionBindingRoute("end_effector", "primary"), + ), + ), + constraints=(DisjointSlotEndpoints(("motion", "grasp")),), + ), + ), + ) + + def __init__( + self, + default_options: SlideOptions | None = None, + ) -> None: + super().__init__(default_options) + + def _on_bind(self) -> None: + """Resolve dimensions owned by the engine's robot.""" + self.num_envs = self.robot.get_qpos().shape[0] + self.robot_dof = self.robot.dof + + def _plan( + self, + request: ResolvedActionRequest[ + SlideGoal, + SlideOptions, + ], + context: PlanningContext, + ) -> ActionPlan: + """Plan the complete pull/push sequence without stepping simulation.""" + target = self.require_goal(request) + affordance = self._require_slide_affordance(target.semantics) + options = request.skill_options + manipulator = request.binding.manipulator() + end_effector = request.binding.end_effector() + arm_joint_ids = list(manipulator.joint_ids) + hand_joint_ids = list(end_effector.joint_ids) + start_arm_qpos = arm_qpos_from_state(context, arm_joint_ids) + hand_open_qpos = end_effector.joint_positions( + OPEN_COMMAND, + num_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + hand_grasp_qpos = end_effector.joint_positions( + GRASP_COMMAND, + num_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + + link_pose = resolve_pose_target( + resolve_pose_goal(target.target_pose, context, name="target_pose"), + num_envs=self.num_envs, + device=self.device, + ) + translation_axis = affordance.translation_axis.to( + device=self.device, dtype=torch.float32 + ) + translation_axis = translation_axis / torch.linalg.vector_norm(translation_axis) + translation_axis_world = torch.matmul(link_pose[:, :3, :3], translation_axis) + grasp_success, grasp_xpos, _ = affordance.get_best_grasp_poses( + obj_poses=link_pose, + approach_direction=translation_axis_world, + ) + grasp_xpos = grasp_xpos.to(device=self.device, dtype=torch.float32) + grasp_success = normalize_success_mask( + grasp_success, + num_envs=self.num_envs, + device=self.device, + name="Slide grasp-pose success", + ) + if not grasp_success.any(): + return self.failed_plan( + request, + context, + message="Failed to resolve an articulated-part grasp pose.", + ) + approach_xpos = translate_pose_world( + grasp_xpos, + -translation_axis_world * options.approach_distance, + ) + translation_sign = -1.0 if options.direction == "pull" else 1.0 + translated_xpos = translate_pose_world( + grasp_xpos, + translation_axis_world * (translation_sign * options.translation_distance), + ) + + motion_lengths = self._motion_segment_lengths( + request.motion_policy.sample_count, + options.hand_interp_steps, + direction=options.direction, + ) + approach_success, approach_arm = self._plan_pose_segment( + approach_xpos, + start_arm_qpos, + manipulator.name, + request, + motion_lengths[0], + ) + reach_keyframes = axis_translation_keyframes( + approach_xpos, + grasp_xpos, + translation_axis_world, + n_waypoints=motion_lengths[1] - 1, + ) + reach_success, reach_arm = self._plan_pose_segment( + reach_keyframes, + approach_arm[:, -1], + manipulator.name, + request, + motion_lengths[1], + cartesian_linear=True, + ) + translate_keyframes = axis_translation_keyframes( + grasp_xpos, + translated_xpos, + translation_axis_world, + n_waypoints=motion_lengths[2] - 1, + ) + translate_success, translate_arm = self._plan_pose_segment( + translate_keyframes, + reach_arm[:, -1], + manipulator.name, + request, + motion_lengths[2], + cartesian_linear=True, + ) + success = grasp_success & approach_success & reach_success & translate_success + + return_arm: torch.Tensor | None = None + if options.direction == "push": + return_keyframes = axis_translation_keyframes( + translated_xpos, + approach_xpos, + translation_axis_world, + n_waypoints=motion_lengths[3] - 1, + ) + return_success, return_arm = self._plan_pose_segment( + return_keyframes, + translate_arm[:, -1], + manipulator.name, + request, + motion_lengths[3], + cartesian_linear=True, + ) + success = success & return_success + + hand_close = interpolate_hand_qpos( + hand_open_qpos, + hand_grasp_qpos, + n_waypoints=options.hand_interp_steps, + ) + hand_open = interpolate_hand_qpos( + hand_grasp_qpos, + hand_open_qpos, + n_waypoints=options.hand_interp_steps, + ) + named_parts: list[tuple[str, torch.Tensor]] = [ + ("approach", approach_arm), + ("reach", reach_arm), + ("close", hand_close), + (options.direction, translate_arm), + ("open", hand_open), + ] + if return_arm is not None: + named_parts.append(("return", return_arm)) + + segment_lengths = {name: part.shape[1] for name, part in named_parts} + full = torch.empty( + (self.num_envs, sum(segment_lengths.values()), self.robot_dof), + dtype=context.robot.qpos.dtype, + device=self.device, + ) + full[:] = context.last_qpos.unsqueeze(1) + offset = 0 + + for arm in (approach_arm, reach_arm): + stop = offset + arm.shape[1] + full[:, offset:stop, arm_joint_ids] = arm + full[:, offset:stop, hand_joint_ids] = hand_open_qpos.unsqueeze(1) + offset = stop + + stop = offset + hand_close.shape[1] + full[:, offset:stop, arm_joint_ids] = reach_arm[:, -1].unsqueeze(1) + full[:, offset:stop, hand_joint_ids] = hand_close + offset = stop + + stop = offset + translate_arm.shape[1] + full[:, offset:stop, arm_joint_ids] = translate_arm + full[:, offset:stop, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1) + offset = stop + + stop = offset + hand_open.shape[1] + full[:, offset:stop, arm_joint_ids] = translate_arm[:, -1].unsqueeze(1) + full[:, offset:stop, hand_joint_ids] = hand_open + offset = stop + + if return_arm is not None: + full[:, offset:, arm_joint_ids] = return_arm + full[:, offset:, hand_joint_ids] = hand_open_qpos.unsqueeze(1) + + return self.build_plan( + request, + context, + success=success, + trajectory=full, + expected_effects=StateDelta(), + segment_lengths=segment_lengths, + ) + + @staticmethod + def _require_slide_affordance( + semantics: ObjectSemantics, + ) -> SlideAffordance: + affordance = semantics.affordance + if not isinstance(affordance, SlideAffordance): + raise ValueError("Slide requires a SlideAffordance.") + return affordance + + @staticmethod + def _motion_segment_lengths( + sample_count: int, + hand_interp_steps: int, + *, + direction: Literal["pull", "push"], + ) -> tuple[int, ...]: + motion_segment_count = 3 if direction == "pull" else 4 + motion_count = sample_count - 2 * hand_interp_steps + if motion_count < 2 * motion_segment_count: + raise ValueError( + "Not enough waypoints for Slide. Increase " + "sample_count or decrease hand_interp_steps." + ) + base, remainder = divmod(motion_count, motion_segment_count) + return tuple( + base + (index < remainder) for index in range(motion_segment_count) + ) + + def _plan_pose_segment( + self, + target_pose: torch.Tensor, + start_qpos: torch.Tensor, + control_part: str, + request: ResolvedActionRequest[ + SlideGoal, + SlideOptions, + ], + sample_count: int, + *, + cartesian_linear: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + result = self.motion_generator.generate( + build_pose_plan_states(target_pose), + options=request.motion_policy.to_motion_gen_options( + start_qpos=start_qpos, + control_part=control_part, + sample_count=sample_count, + cartesian_linear=cartesian_linear, + ), + ) + assert isinstance(result.success, torch.Tensor) + assert result.positions is not None + return result.success, result.positions + + +__all__ = [ + "Slide", + "SlideGoal", + "SlideOptions", +] diff --git a/embodichain/lab/sim/atomic_actions/primitives/twist.py b/embodichain/lab/sim/atomic_actions/primitives/twist.py new file mode 100644 index 000000000..92b9bbb72 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/twist.py @@ -0,0 +1,421 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Twist atomic action implementation.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import ClassVar + +import torch + +from embodichain.utils.math import ( + axis_angle_to_rotation_matrix, + pose_inv, + get_relative_rotation, +) + +from embodichain.lab.sim.atomic_actions.affordance import TwistAffordance +from embodichain.lab.sim.atomic_actions.control import ( + GRASP_COMMAND, + OPEN_COMMAND, + JointPositionCommand, +) +from embodichain.lab.sim.atomic_actions.core import AtomicAction, ObjectSemantics +from embodichain.lab.sim.atomic_actions.effects import StateDelta +from embodichain.lab.sim.atomic_actions.goals import ( + ObjectActionGoal, + PoseGoalValue, + resolve_pose_goal, + validate_pose_goal, +) +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan +from embodichain.lab.sim.atomic_actions.primitives._helpers import arm_qpos_from_state +from embodichain.lab.sim.atomic_actions.requirements import ( + ActionBindingRoute, + CARTESIAN_POSE_CAPABILITY, + DisjointSlotEndpoints, + FORWARD_KINEMATICS_CAPABILITY, + GRASP_CAPABILITY, + SkillBindingContract, + SkillEndpointRequirement, + SkillResourceSlot, +) +from embodichain.lab.sim.atomic_actions.state import PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( + build_pose_plan_states, + interpolate_hand_qpos, + resolve_pose_target, + translate_pose_world, +) + + +@dataclass(frozen=True, slots=True, eq=False) +class TwistGoal(ObjectActionGoal): + """Target object described by a twist affordance.""" + + goal_kind: ClassVar[str] = "twist" + + target_pose: PoseGoalValue + """Target pose snapshot or late-bound stable scene-entity reference.""" + + def __post_init__(self) -> None: + ObjectActionGoal.__post_init__(self) + validate_pose_goal(self.target_pose, "target_pose", allow_waypoints=False) + + +@dataclass(frozen=True, slots=True, eq=False) +class TwistOptions(ActionOptions): + """Per-invocation twisting behavior.""" + + hand_interp_steps: int = 5 + """Number of waypoints used for each close/open hand segment.""" + + twist_waypoint_count: int = 8 + """Number of Cartesian keyframes along the target's circular twist arc.""" + + pre_grasp_distance: float = 0.1 + """Distance from the grasp pose along its negative z-axis.""" + + twist_angle: float = math.pi / 4 + """Requested twist rotation in radians.""" + + def __post_init__(self) -> None: + if self.hand_interp_steps < 1: + raise ValueError("hand_interp_steps must be at least 1.") + if self.twist_waypoint_count < 1: + raise ValueError("twist_waypoint_count must be at least 1.") + if not math.isfinite(self.pre_grasp_distance): + raise ValueError("pre_grasp_distance must be finite.") + if self.pre_grasp_distance < 0.0: + raise ValueError("pre_grasp_distance must be non-negative.") + if not math.isfinite(self.twist_angle): + raise ValueError("twist_angle must be finite.") + + +class Twist(AtomicAction[TwistGoal, TwistOptions]): + """Open-loop approach, grasp, twist, release, and retract motion.""" + + skill_id: ClassVar[str] = "twist" + GoalType: ClassVar[type] = TwistGoal + OptionsType: ClassVar[type] = TwistOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) + open_loop: ClassVar[bool] = True + binding_contract: ClassVar[SkillBindingContract] = SkillBindingContract( + slots=( + SkillResourceSlot( + slot_id="primary", + endpoints=( + SkillEndpointRequirement( + endpoint_id="motion", + capabilities=frozenset( + { + CARTESIAN_POSE_CAPABILITY, + FORWARD_KINEMATICS_CAPABILITY, + } + ), + route=ActionBindingRoute("manipulator", "primary"), + ), + SkillEndpointRequirement( + endpoint_id="grasp", + capabilities=frozenset({GRASP_CAPABILITY}), + required_commands={ + OPEN_COMMAND: JointPositionCommand, + GRASP_COMMAND: JointPositionCommand, + }, + route=ActionBindingRoute("end_effector", "primary"), + ), + ), + constraints=(DisjointSlotEndpoints(("motion", "grasp")),), + ), + ), + ) + + def __init__(self, default_options: TwistOptions | None = None) -> None: + super().__init__(default_options) + + def _on_bind(self) -> None: + """Resolve dimensions owned by the engine's robot.""" + self.num_envs = self.robot.get_qpos().shape[0] + self.robot_dof = self.robot.dof + + def _find_symmetric_nearest_xpos( + self, target_xpos: torch.Tensor, reference_xpos: torch.Tensor + ) -> torch.Tensor: + """Find the nearest symmetric pose to the reference pose.""" + symmetric_xpos = target_xpos.clone() + symmetric_xpos[:, :3, 0] = -symmetric_xpos[:, :3, 0] + symmetric_xpos[:, :3, 1] = -symmetric_xpos[:, :3, 1] + angle_a = get_relative_rotation( + reference_xpos[:, :3, :3], target_xpos[:, :3, :3] + ) + angle_b = get_relative_rotation( + reference_xpos[:, :3, :3], symmetric_xpos[:, :3, :3] + ) + choose_target = (angle_a < angle_b)[..., None, None] + target_xpos = torch.where(choose_target, target_xpos, symmetric_xpos) + return target_xpos + + def _plan( + self, + request: ResolvedActionRequest[TwistGoal, TwistOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan all six twisting segments without stepping simulation.""" + target = self.require_goal(request) + affordance = self._require_twist_affordance(target.semantics) + options = request.skill_options + manipulator = request.binding.manipulator() + end_effector = request.binding.end_effector() + arm_joint_ids = list(manipulator.joint_ids) + hand_joint_ids = list(end_effector.joint_ids) + start_arm_qpos = arm_qpos_from_state(context, arm_joint_ids) + hand_open_qpos = end_effector.joint_positions( + OPEN_COMMAND, + num_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + hand_grasp_qpos = end_effector.joint_positions( + GRASP_COMMAND, + num_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + + link_pose = resolve_pose_target( + resolve_pose_goal(target.target_pose, context, name="target_pose"), + num_envs=self.num_envs, + device=self.device, + ) + grasp_xpos = affordance.get_grasp_pose(link_pose).to( + device=self.device, dtype=torch.float32 + ) + grasp_xpos = self._find_symmetric_nearest_xpos( + grasp_xpos, + reference_xpos=self.robot.compute_fk( + qpos=start_arm_qpos, name=manipulator.name, to_matrix=True + ), + ) + pre_grasp_xpos = translate_pose_world( + grasp_xpos, + -grasp_xpos[:, :3, 2] * options.pre_grasp_distance, + ) + twist_xpos = self._twisted_grasp_poses( + link_pose, + grasp_xpos, + affordance.twist_axis, + affordance.axis_origin, + options.twist_angle, + options.twist_waypoint_count, + ) + + n_approach, n_reach, n_twist, n_retract = self._motion_segment_lengths( + request.motion_policy.sample_count, + options.hand_interp_steps, + ) + + approach_success, approach_arm = self._plan_pose_segment( + pre_grasp_xpos, + start_arm_qpos, + manipulator.name, + request, + n_approach, + ) + reach_success, reach_arm = self._plan_pose_segment( + grasp_xpos, + approach_arm[:, -1], + manipulator.name, + request, + n_reach, + ) + twist_success, twist_arm = self._plan_pose_segment( + twist_xpos, + reach_arm[:, -1], + manipulator.name, + request, + n_twist, + ) + retract_success, retract_arm = self._plan_pose_segment( + pre_grasp_xpos, + twist_arm[:, -1], + manipulator.name, + request, + n_retract, + ) + success = approach_success & reach_success & twist_success & retract_success + + hand_close = interpolate_hand_qpos( + hand_open_qpos, + hand_grasp_qpos, + n_waypoints=options.hand_interp_steps, + ) + hand_open = interpolate_hand_qpos( + hand_grasp_qpos, + hand_open_qpos, + n_waypoints=options.hand_interp_steps, + ) + parts = ( + approach_arm, + reach_arm, + hand_close, + twist_arm, + hand_open, + retract_arm, + ) + lengths = tuple(part.shape[1] for part in parts) + full = torch.empty( + (self.num_envs, sum(lengths), self.robot_dof), + dtype=context.robot.qpos.dtype, + device=self.device, + ) + full[:] = context.last_qpos.unsqueeze(1) + offset = 0 + arm_parts = (approach_arm, reach_arm, twist_arm, retract_arm) + arm_hands = ( + hand_open_qpos, + hand_open_qpos, + hand_grasp_qpos, + hand_open_qpos, + ) + for arm, hand in zip(arm_parts[:2], arm_hands[:2]): + stop = offset + arm.shape[1] + full[:, offset:stop, arm_joint_ids] = arm + full[:, offset:stop, hand_joint_ids] = hand.unsqueeze(1) + offset = stop + stop = offset + hand_close.shape[1] + full[:, offset:stop, arm_joint_ids] = reach_arm[:, -1].unsqueeze(1) + full[:, offset:stop, hand_joint_ids] = hand_close + offset = stop + stop = offset + twist_arm.shape[1] + full[:, offset:stop, arm_joint_ids] = twist_arm + full[:, offset:stop, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1) + offset = stop + stop = offset + hand_open.shape[1] + full[:, offset:stop, arm_joint_ids] = twist_arm[:, -1].unsqueeze(1) + full[:, offset:stop, hand_joint_ids] = hand_open + offset = stop + full[:, offset:, arm_joint_ids] = retract_arm + full[:, offset:, hand_joint_ids] = hand_open_qpos.unsqueeze(1) + + return self.build_plan( + request, + context, + success=success, + trajectory=full, + expected_effects=StateDelta(), + segment_lengths={ + "approach": lengths[0], + "reach": lengths[1], + "close": lengths[2], + "twist": lengths[3], + "open": lengths[4], + "retract": lengths[5], + }, + ) + + @staticmethod + def _require_twist_affordance( + semantics: ObjectSemantics, + ) -> TwistAffordance: + affordance = semantics.affordance + if not isinstance(affordance, TwistAffordance): + raise ValueError("Twist requires a TwistAffordance.") + return affordance + + @staticmethod + def _motion_segment_lengths( + sample_count: int, + hand_interp_steps: int, + ) -> tuple[int, int, int, int]: + motion_count = sample_count - 2 * hand_interp_steps + if motion_count < 8: + raise ValueError( + "Not enough waypoints for Twist. Increase sample_count or " + "decrease hand_interp_steps." + ) + base, remainder = divmod(motion_count, 4) + values = [base + (index < remainder) for index in range(4)] + return values[0], values[1], values[2], values[3] + + def _plan_pose_segment( + self, + target_pose: torch.Tensor, + start_qpos: torch.Tensor, + control_part: str, + request: ResolvedActionRequest[TwistGoal, TwistOptions], + sample_count: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + result = self.motion_generator.generate( + build_pose_plan_states(target_pose), + options=request.motion_policy.to_motion_gen_options( + start_qpos=start_qpos, + control_part=control_part, + sample_count=sample_count, + ), + ) + assert isinstance(result.success, torch.Tensor) + assert result.positions is not None + return result.success, result.positions + + def _twisted_grasp_poses( + self, + link_pose: torch.Tensor, + grasp_xpos: torch.Tensor, + twist_axis: torch.Tensor, + axis_origin: tuple[float, float, float], + twist_angle: float, + waypoint_count: int, + ) -> torch.Tensor: + """Build Cartesian EEF keyframes that follow the target's twist arc.""" + axis = twist_axis.to(device=self.device, dtype=torch.float32) + axis = axis / torch.linalg.vector_norm(axis) + angles = torch.linspace( + twist_angle / waypoint_count, + twist_angle, + waypoint_count, + dtype=torch.float32, + device=self.device, + ) + rotations = ( + torch.eye(4, dtype=torch.float32, device=self.device) + .reshape(1, 4, 4) + .repeat(waypoint_count, 1, 1) + ) + rotations[:, :3, :3] = axis_angle_to_rotation_matrix(angles[:, None] * axis) + link_to_eef = torch.bmm(pose_inv(link_pose), grasp_xpos) + origin = torch.tensor(axis_origin, dtype=torch.float32, device=self.device) + to_origin = torch.eye(4, dtype=torch.float32, device=self.device) + from_origin = torch.eye(4, dtype=torch.float32, device=self.device) + to_origin[:3, 3] = origin + from_origin[:3, 3] = -origin + local_rotations = torch.matmul( + torch.matmul(to_origin[None], rotations), from_origin[None] + ) + return torch.matmul( + torch.matmul(link_pose[:, None], local_rotations[None]), + link_to_eef[:, None], + ) + + +__all__ = ["Twist", "TwistGoal", "TwistOptions"] diff --git a/embodichain/lab/sim/atomic_actions/trajectory_ops.py b/embodichain/lab/sim/atomic_actions/trajectory_ops.py index e64af32b7..cc544ef34 100644 --- a/embodichain/lab/sim/atomic_actions/trajectory_ops.py +++ b/embodichain/lab/sim/atomic_actions/trajectory_ops.py @@ -172,6 +172,76 @@ def translate_pose_world(pose: torch.Tensor, offset: torch.Tensor) -> torch.Tens return result +def axis_translation_keyframes( + start_pose: torch.Tensor, + end_pose: torch.Tensor, + axis: torch.Tensor, + *, + n_waypoints: int, +) -> torch.Tensor: + """Build exact Cartesian translation targets along one world-space axis. + + The returned targets exclude ``start_pose`` and include ``end_pose``. This + matches motion generation, where the observed start configuration is added + separately. Rotation remains fixed for the entire constrained segment. + + Args: + start_pose: Batched segment-start poses, shape ``(B, 4, 4)``. + end_pose: Batched segment-end poses, shape ``(B, 4, 4)``. + axis: Shared ``(3,)`` or batched ``(B, 3)`` world-space axis. + n_waypoints: Number of target poses, excluding the segment start. + + Returns: + Batched keyframes with shape ``(B, n_waypoints, 4, 4)``. + + Raises: + ValueError: If poses, axis, count, rotation, or displacement are invalid. + """ + if ( + start_pose.dim() != 3 + or start_pose.shape[1:] != (4, 4) + or end_pose.shape != start_pose.shape + ): + raise ValueError("start_pose and end_pose must have shape (B, 4, 4).") + if n_waypoints < 1: + raise ValueError("n_waypoints must be at least 1.") + axis = axis.to(device=start_pose.device, dtype=start_pose.dtype) + if axis.shape == (3,): + axis = axis.unsqueeze(0).expand(start_pose.shape[0], -1) + if axis.shape != (start_pose.shape[0], 3) or not torch.isfinite(axis).all(): + raise ValueError("axis must be finite with shape (3,) or (B, 3).") + axis_norm = torch.linalg.vector_norm(axis, dim=1, keepdim=True) + if torch.any(axis_norm <= 1.0e-6): + raise ValueError("axis must be non-zero.") + axis = axis / axis_norm + if not torch.allclose( + start_pose[:, :3, :3], + end_pose[:, :3, :3], + rtol=1.0e-5, + atol=1.0e-6, + ): + raise ValueError("Axis translation requires a fixed segment rotation.") + displacement = end_pose[:, :3, 3] - start_pose[:, :3, 3] + orthogonal = displacement - (displacement * axis).sum(dim=1, keepdim=True) * axis + if torch.any(torch.linalg.vector_norm(orthogonal, dim=1) > 1.0e-5): + raise ValueError("Segment displacement must be parallel to axis.") + + weights = torch.linspace( + 0.0, + 1.0, + n_waypoints + 1, + dtype=start_pose.dtype, + device=start_pose.device, + )[1:] + result = start_pose[:, None].expand(-1, n_waypoints, -1, -1).clone() + result[:, :, :3, 3] = torch.lerp( + start_pose[:, None, :3, 3], + end_pose[:, None, :3, 3], + weights[None, :, None], + ) + return result + + def split_three_segments( sample_count: int, hand_interp_steps: int, @@ -298,6 +368,7 @@ def embed_derivative(value: torch.Tensor | None) -> torch.Tensor | None: __all__ = [ + "axis_translation_keyframes", "build_joint_plan_states", "build_pose_plan_states", "interpolate_hand_qpos", diff --git a/embodichain/lab/sim/planners/motion_generator.py b/embodichain/lab/sim/planners/motion_generator.py index f26d411da..61093c26a 100644 --- a/embodichain/lab/sim/planners/motion_generator.py +++ b/embodichain/lab/sim/planners/motion_generator.py @@ -113,6 +113,13 @@ class MotionGenOptions: is_linear: bool = False """If True, use cartesian linear interpolation, else joint space""" + preserve_cartesian_samples: bool = False + """Treat Cartesian targets as exact output samples and solve each with IK. + + This constrained mode requires exactly ``sample_count - 1`` target states; + the observed start configuration supplies the first output sample. + """ + interpolate_position_step: float = 0.002 """Step size for interpolation. If is_linear is True, this is the step size in Cartesian space (meters). If is_linear is False, this is the step size in joint space (radians).""" @@ -426,9 +433,13 @@ def generate( names = sorted(move_type.name for move_type in move_types) raise ValueError(f"All target states must share move_type; got {names}.") move_type = target_states[0].move_type - use_interpolation = options.strategy == "ik_interp" or ( - move_type is MoveType.JOINT_MOVE - and not self.planner.supports_move_type(MoveType.JOINT_MOVE) + use_interpolation = ( + options.preserve_cartesian_samples + or options.strategy == "ik_interp" + or ( + move_type is MoveType.JOINT_MOVE + and not self.planner.supports_move_type(MoveType.JOINT_MOVE) + ) ) if use_interpolation: raw_result = self._generate_ik_interpolation(target_states, options) @@ -633,11 +644,21 @@ def _generate_ik_interpolation( [start_qpos.unsqueeze(1), torch.stack(solved_waypoints, dim=1)], dim=1, ) - positions = interpolate_with_distance( - trajectory=keyframes, - interp_num=options.sample_count, - device=device, - ) + if options.preserve_cartesian_samples: + if keyframes.shape[1] != options.sample_count: + raise ValueError( + "Linear Cartesian targets must provide sample_count - 1 " + "keyframes so every output sample is IK-grounded; got " + f"{len(target_states)} targets for sample_count " + f"{options.sample_count}." + ) + positions = keyframes + else: + positions = interpolate_with_distance( + trajectory=keyframes, + interp_num=options.sample_count, + device=device, + ) held = start_qpos.unsqueeze(1).expand_as(positions) positions = torch.where(success[:, None, None], positions, held) return PlanResult(success=success, positions=positions) diff --git a/scripts/benchmark/atomic_action/press_benchmark.py b/scripts/benchmark/atomic_action/press_benchmark.py deleted file mode 100644 index a5687af7f..000000000 --- a/scripts/benchmark/atomic_action/press_benchmark.py +++ /dev/null @@ -1,994 +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. -# ---------------------------------------------------------------------------- - -"""Benchmark Press atomic action across object presets and start positions. - -The benchmark sweeps object presets such as bottle and mug against multiple -initial XY positions that cover all four workspace quadrants. It reports -planning latency, memory usage, planning success, and whether the generated -trajectory reaches the object's top center. -Run: embodichain benchmark atomic-action --action press -""" - -from __future__ import annotations - -import argparse -import math -import os -import resource -import sys -import time -from dataclasses import dataclass -from datetime import datetime -from pathlib import Path - -from scripts.benchmark.atomic_action.common import ( - add_profile_benchmark_args, - add_video_benchmark_args, - build_video_output_path, - COVERAGE_POSITION_CASE_NAMES, - park_rigid_object, - replay_trajectory_with_recording, - reset_rigid_object, - reset_rigid_object_xy, - resolve_profile, - should_record_case, - SMOKE_POSITION_CASE_NAMES, -) - -try: - import psutil -except ModuleNotFoundError: - psutil = None - -CPU_MEMORY_BACKEND = "psutil" if psutil is not None else "resource" -_RUNTIME_IMPORTS_READY = False - -# Keep these constants aligned with scripts/tutorials/atomic_action/press.py. -DEFAULT_PRESS_TOLERANCE = 0.01 -MOVE_SAMPLE_INTERVAL = 60 -PRESS_SAMPLE_INTERVAL = 90 -HAND_INTERP_STEPS = 12 -TABLE_TOP_Z = -0.045 -PRESS_CLEARANCE = 0.13 -PRESS_SURFACE_OFFSET = 0.003 - - -def _ensure_runtime_imports() -> None: - """Import simulation dependencies only when the benchmark is executed.""" - global _RUNTIME_IMPORTS_READY - if _RUNTIME_IMPORTS_READY: - return - - repo_root = Path(__file__).resolve().parents[3] - if str(repo_root) not in sys.path: - sys.path.insert(0, str(repo_root)) - - try: - import torch as torch_module - from embodichain.lab.sim import SimulationManager as simulation_manager_cls - from embodichain.lab.sim.atomic_actions import ( - 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, - MotionPolicy as motion_policy_cls, - PressGoal as press_target_cls, - PressOptions as press_options_cls, - ) - from embodichain.lab.sim.cfg import ( - RigidBodyAttributesCfg as rigid_body_attributes_cfg_cls, - RigidObjectCfg as rigid_object_cfg_cls, - ) - from embodichain.lab.sim.material import ( - VisualMaterialCfg as visual_material_cfg_cls, - ) - from embodichain.lab.sim.objects import RigidObject as rigid_object_cls - from embodichain.lab.sim.objects import Robot as robot_cls - from embodichain.lab.sim.planners import ( - MotionGenerator as motion_generator_cls, - MotionGenCfg as motion_gen_cfg_cls, - ToppraPlannerCfg as toppra_planner_cfg_cls, - ) - from embodichain.lab.sim.shapes import CubeCfg as cube_cfg_cls - from scripts.tutorials.atomic_action.press import ( - create_robot as create_robot_fn, - create_table as create_table_fn, - get_hand_close_qpos as get_hand_close_qpos_fn, - initialize_simulation as initialize_simulation_fn, - make_top_down_eef_pose as make_top_down_eef_pose_fn, - settle_object as settle_object_fn, - ) - except ModuleNotFoundError as exc: - raise RuntimeError( - "Atomic action benchmark requires the EmbodiChain simulation runtime " - f"and PyTorch. Missing module: {exc.name}." - ) from exc - - globals().update( - { - "torch": torch_module, - "SimulationManager": simulation_manager_cls, - "AtomicActionEngine": atomic_action_engine_cls, - "ControlPartCommandProfile": control_part_command_profile_cls, - "ActionBinding": action_binding_cls, - "ActionInvocation": action_invocation_cls, - "EndEffectorPoseGoal": end_effector_pose_target_cls, - "MotionPolicy": motion_policy_cls, - "PressGoal": press_target_cls, - "PressOptions": press_options_cls, - "RigidBodyAttributesCfg": rigid_body_attributes_cfg_cls, - "RigidObjectCfg": rigid_object_cfg_cls, - "VisualMaterialCfg": visual_material_cfg_cls, - "RigidObject": rigid_object_cls, - "Robot": robot_cls, - "MotionGenerator": motion_generator_cls, - "MotionGenCfg": motion_gen_cfg_cls, - "ToppraPlannerCfg": toppra_planner_cfg_cls, - "CubeCfg": cube_cfg_cls, - "create_robot": create_robot_fn, - "create_table": create_table_fn, - "get_hand_close_qpos": get_hand_close_qpos_fn, - "initialize_simulation": initialize_simulation_fn, - "make_top_down_eef_pose": make_top_down_eef_pose_fn, - "settle_object": settle_object_fn, - } - ) - _RUNTIME_IMPORTS_READY = True - - -@dataclass(frozen=True) -class ObjectPreset: - """Primitive object preset used by the atomic-action benchmark.""" - - object_type: str - material_name: str - size: tuple[float, float, float] - base_color: tuple[float, float, float, float] - roughness: float - dynamic_friction: float = 0.8 - static_friction: float = 0.9 - - -@dataclass(frozen=True) -class PositionCase: - """Initial object position case with a quadrant label.""" - - name: str - quadrant: str - xy: tuple[float, float] - - -@dataclass(frozen=True) -class PressCaseResult: - """Result for one Press benchmark case.""" - - case_id: str - object_type: str - material_name: str - quadrant: str - position_case: str - init_xy: tuple[float, float] - repeat_index: int - planning_success: bool - center_hit: bool - cost_time_ms: float - cpu_delta_mb: float - gpu_delta_mb: float - peak_gpu_mb: float - xy_error_m: float | None - hit_step: int | None - trajectory_waypoints: int - failure_reason: str - video_path: str = "" - - -OBJECT_PRESETS: dict[str, ObjectPreset] = { - "bottle": ObjectPreset( - object_type="bottle", - material_name="green_plastic", - size=(0.06, 0.06, 0.16), - base_color=(0.10, 0.45, 0.32, 1.0), - roughness=0.55, - ), - "mug": ObjectPreset( - object_type="mug", - material_name="ceramic", - size=(0.10, 0.08, 0.10), - base_color=(0.88, 0.85, 0.78, 1.0), - roughness=0.35, - ), - "wooden_block": ObjectPreset( - object_type="wooden_block", - material_name="wood", - size=(0.12, 0.12, 0.06), - base_color=(0.58, 0.32, 0.14, 1.0), - roughness=0.85, - ), -} - -POSITION_CASES: dict[str, PositionCase] = { - "q1_near": PositionCase(name="q1_near", quadrant="q1", xy=(0.02, 0.18)), - "q1_far": PositionCase(name="q1_far", quadrant="q1", xy=(0.12, 0.36)), - "q2_near": PositionCase(name="q2_near", quadrant="q2", xy=(-0.42, 0.18)), - "q2_far": PositionCase(name="q2_far", quadrant="q2", xy=(-0.62, 0.36)), - "q3_near": PositionCase(name="q3_near", quadrant="q3", xy=(-0.42, -0.18)), - "q3_far": PositionCase(name="q3_far", quadrant="q3", xy=(-0.62, -0.36)), - "q4_near": PositionCase(name="q4_near", quadrant="q4", xy=(0.02, -0.18)), - "q4_far": PositionCase(name="q4_far", quadrant="q4", xy=(0.12, -0.36)), -} - -DEFAULT_OBJECT_TYPES = ("bottle", "mug") -FULL_OBJECT_TYPES = tuple(OBJECT_PRESETS.keys()) -SMOKE_OBJECT_TYPES = ("bottle",) - - -def add_benchmark_args(parser: argparse.ArgumentParser) -> None: - """Add atomic-action benchmark arguments to an argument parser.""" - add_profile_benchmark_args(parser) - parser.add_argument( - "--object_types", - nargs="+", - choices=(*OBJECT_PRESETS.keys(), "all"), - default=None, - help=( - "Object presets to benchmark. Defaults are selected by --profile; " - "use 'all' to include every preset." - ), - ) - parser.add_argument( - "--position_cases", - nargs="+", - choices=(*POSITION_CASES.keys(), "all"), - default=None, - help=( - "Initial position cases to benchmark. Defaults are selected by " - "--profile; use 'all' for all near/far cases." - ), - ) - parser.add_argument( - "--repeat", - type=int, - default=1, - help="Number of repeats for every object-position case.", - ) - parser.add_argument( - "--smoke", - action="store_true", - help="Alias for --profile smoke.", - ) - parser.add_argument( - "--device", - type=str, - default="cpu", - help="Simulation device, e.g. 'cpu' or 'cuda'.", - ) - parser.add_argument( - "--renderer", - type=str, - choices=("auto", "hybrid", "fast-rt", "rt"), - default="auto", - help="Renderer backend used by SimulationManager.", - ) - add_video_benchmark_args(parser) - parser.add_argument( - "--press_tolerance", - type=float, - default=DEFAULT_PRESS_TOLERANCE, - help="XY tolerance in meters for the press-center check.", - ) - - -def _parse_args() -> argparse.Namespace: - """Parse command line arguments for the atomic-action benchmark.""" - parser = argparse.ArgumentParser( - description=( - "Benchmark Press atomic action over object presets and initial " - "workspace quadrants." - ) - ) - add_benchmark_args(parser) - return parser.parse_args() - - -def _sync_cuda() -> None: - """Synchronize CUDA stream when available.""" - if torch.cuda.is_available(): - torch.cuda.synchronize() - - -def _reset_peak_gpu_memory() -> None: - """Reset PyTorch peak GPU memory stats when CUDA is available.""" - if torch.cuda.is_available(): - torch.cuda.reset_peak_memory_stats() - - -def _peak_gpu_memory_mb() -> float: - """Return peak GPU memory allocated by PyTorch in MB.""" - if not torch.cuda.is_available(): - return 0.0 - return torch.cuda.max_memory_allocated() / 1024**2 - - -def _memory_snapshot() -> dict[str, float]: - """Return current process memory usage snapshot in MB.""" - if psutil is not None: - process = psutil.Process(os.getpid()) - cpu_mb = process.memory_info().rss / 1024**2 - else: - cpu_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0 - gpu_mb = ( - torch.cuda.memory_allocated() / 1024**2 if torch.cuda.is_available() else 0.0 - ) - return {"cpu_mb": cpu_mb, "gpu_mb": gpu_mb} - - -def _format_float(value: float | None, precision: int = 6) -> str: - """Format finite floats for tables and use N/A for missing values.""" - if value is None or not math.isfinite(value): - return "N/A" - return f"{value:.{precision}f}" - - -def _format_markdown_table(rows: list[dict[str, object]]) -> list[str]: - """Format rows into a markdown table.""" - if not rows: - return ["No data."] - - headers = list(rows[0].keys()) - lines = [ - "| " + " | ".join(headers) + " |", - "| " + " | ".join(["---"] * len(headers)) + " |", - ] - for row in rows: - lines.append("| " + " | ".join(str(row[h]) for h in headers) + " |") - return lines - - -def _write_markdown_report( - benchmark_name: str, - perf_rows: list[dict[str, object]], - metric_rows: list[dict[str, object]], - leaderboard_rows: list[dict[str, object]], - notes: list[str] | None = None, -) -> Path: - """Write benchmark results into one markdown report file.""" - output_dir = Path("outputs/benchmarks") - output_dir.mkdir(parents=True, exist_ok=True) - - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - report_path = output_dir / f"{benchmark_name}_{timestamp}.md" - - lines: list[str] = [ - f"# {benchmark_name} Benchmark Report", - "", - f"Generated at: {datetime.now().isoformat(timespec='seconds')}", - "", - "## Time & Memory", - "", - ] - lines.extend(_format_markdown_table(perf_rows)) - lines.extend(["", "## Success & Other Metrics", ""]) - lines.extend(_format_markdown_table(metric_rows)) - lines.extend(["", "## Leaderboard", ""]) - lines.extend(_format_markdown_table(leaderboard_rows)) - - if notes: - lines.extend(["", "## Notes", ""]) - lines.extend([f"- {note}" for note in notes]) - - report_path.write_text("\n".join(lines) + "\n", encoding="utf-8") - return report_path - - -def _default_object_types_for_profile(profile: str) -> tuple[str, ...]: - """Return default Press primitive object names for a profile.""" - if profile in ("smoke", "coverage", "full"): - return SMOKE_OBJECT_TYPES - raise ValueError(f"Unsupported benchmark profile: {profile}") - - -def _select_object_presets( - object_types: list[str] | None, - profile: str, -) -> list[ObjectPreset]: - """Resolve selected object preset names.""" - if not object_types: - object_types = list(_default_object_types_for_profile(profile)) - if "all" in object_types: - return list(OBJECT_PRESETS.values()) - return [OBJECT_PRESETS[name] for name in object_types] - - -def _default_position_cases_for_profile(profile: str) -> tuple[str, ...]: - """Return default Press position case names for a profile.""" - if profile == "smoke": - return SMOKE_POSITION_CASE_NAMES - if profile in ("coverage", "full"): - return COVERAGE_POSITION_CASE_NAMES - raise ValueError(f"Unsupported benchmark profile: {profile}") - - -def _select_position_cases( - position_cases: list[str] | None, - profile: str, -) -> list[PositionCase]: - """Resolve selected position case names.""" - if not position_cases: - position_cases = list(_default_position_cases_for_profile(profile)) - if "all" in position_cases: - return list(POSITION_CASES.values()) - return [POSITION_CASES[name] for name in position_cases] - - -def _create_benchmark_object( - sim: SimulationManager, - preset: ObjectPreset, - position_case: PositionCase, - repeat_index: int, -) -> RigidObject: - """Create one static benchmark object at the requested initial position.""" - init_pos = ( - position_case.xy[0], - position_case.xy[1], - TABLE_TOP_Z + 0.5 * preset.size[2], - ) - uid = f"atomic_benchmark_{preset.object_type}_{position_case.name}_{repeat_index}" - cfg = RigidObjectCfg( - uid=uid, - shape=CubeCfg( - size=list(preset.size), - visual_material=VisualMaterialCfg( - uid=f"{preset.object_type}_{preset.material_name}_mat", - base_color=list(preset.base_color), - roughness=preset.roughness, - ), - ), - body_type="static", - attrs=RigidBodyAttributesCfg( - dynamic_friction=preset.dynamic_friction, - static_friction=preset.static_friction, - ), - init_pos=init_pos, - ) - return sim.add_rigid_object(cfg=cfg) - - -def _reset_robot(robot: Robot, initial_qpos: torch.Tensor) -> None: - """Reset current and target robot qpos to the benchmark initial posture.""" - for target in (False, True): - robot.set_qpos(initial_qpos, target=target) - robot.clear_dynamics() - - -def _build_atomic_engine( - motion_gen: MotionGenerator, - robot: Robot, - device: torch.device, -) -> AtomicActionEngine: - """Build a Press benchmark engine with MoveEndEffector pre-positioning.""" - hand_close = get_hand_close_qpos(robot, device) - atomic_engine = AtomicActionEngine( - motion_generator=motion_gen, - control_profiles={ - "hand": ControlPartCommandProfile.joint_positions( - grasp=hand_close, - ) - }, - ) - return atomic_engine - - -def _make_press_targets( - obj: RigidObject, - preset: ObjectPreset, -) -> tuple[torch.Tensor, torch.Tensor]: - """Create pre-press and press poses for the object's top center.""" - obj_pose = obj.get_local_pose(to_matrix=True) - object_center = obj_pose[0, :3, 3].clone() - object_top_z = object_center[2] + 0.5 * preset.size[2] - - press_position = object_center.clone() - press_position[2] = object_top_z + PRESS_SURFACE_OFFSET - move_position = press_position.clone() - move_position[2] = object_top_z + PRESS_CLEARANCE - - return make_top_down_eef_pose(move_position), make_top_down_eef_pose(press_position) - - -def _compute_press_center_check( - robot: Robot, - traj: torch.Tensor, - obj: RigidObject, - object_height: float, - tolerance: float, -) -> tuple[bool, float, int]: - """Check whether the planned Press trajectory reaches the object top center.""" - if traj.numel() == 0: - return False, float("inf"), -1 - - arm_joint_ids = robot.get_joint_ids(name="arm") - n_down = (PRESS_SAMPLE_INTERVAL - HAND_INTERP_STEPS) // 2 - press_segment_start = MOVE_SAMPLE_INTERVAL + HAND_INTERP_STEPS - press_segment_end = min(press_segment_start + n_down, traj.shape[1]) - arm_traj = traj[:, press_segment_start:press_segment_end, arm_joint_ids] - if arm_traj.shape[1] == 0: - return False, float("inf"), -1 - - fk_pose = torch.stack( - [ - robot.compute_fk( - qpos=waypoint.unsqueeze(0), - name="arm", - to_matrix=True, - )[0] - for waypoint in arm_traj[0] - ], - dim=0, - ) - - obj_pose = obj.get_local_pose(to_matrix=True) - object_center = obj_pose[0, :3, 3] - object_top_z = object_center[2] + 0.5 * object_height - target_xy = object_center[:2] - target_z = object_top_z + PRESS_SURFACE_OFFSET - - xy_error = torch.linalg.norm(fk_pose[:, :2, 3] - target_xy, dim=1) - z_error = torch.abs(fk_pose[:, 2, 3] - target_z) - combined_error = xy_error + z_error - best_idx = int(torch.argmin(combined_error).item()) - best_pos = fk_pose[best_idx, :3, 3] - center_error = float(torch.linalg.norm(best_pos[:2] - target_xy).item()) - return center_error <= tolerance, center_error, press_segment_start + best_idx - - -def _timed_atomic_run( - atomic_engine: AtomicActionEngine, - move_target: torch.Tensor, - press_target: torch.Tensor, -) -> tuple[float, dict[str, float], float, bool, torch.Tensor]: - """Run a timed atomic-action sequence and return timing/memory/results.""" - _reset_peak_gpu_memory() - mem_before = _memory_snapshot() - _sync_cuda() - - start = time.perf_counter() - binding = ActionBinding( - manipulators={"primary": "arm"}, - end_effectors={"primary": "hand"}, - ) - result = atomic_engine.compile( - ( - ActionInvocation( - "move_end_effector", - EndEffectorPoseGoal(xpos=move_target), - binding, - MotionPolicy(sample_count=MOVE_SAMPLE_INTERVAL), - ), - ActionInvocation( - "press", - PressGoal(xpos=press_target), - binding, - MotionPolicy(sample_count=PRESS_SAMPLE_INTERVAL), - skill_options=PressOptions( - hand_interp_steps=HAND_INTERP_STEPS, - ), - ), - ) - ) - is_success = bool(result.plan_success.all().item()) - traj = result.trajectory.positions - _sync_cuda() - elapsed = time.perf_counter() - start - - mem_after = _memory_snapshot() - deltas = { - "cpu_mb": mem_after["cpu_mb"] - mem_before["cpu_mb"], - "gpu_mb": mem_after["gpu_mb"] - mem_before["gpu_mb"], - } - return elapsed, deltas, _peak_gpu_memory_mb(), is_success, traj - - -def _run_press_case( - sim: SimulationManager, - robot: Robot, - atomic_engine: AtomicActionEngine, - initial_qpos: torch.Tensor, - obj: RigidObject, - base_obj_pose: torch.Tensor, - preset: ObjectPreset, - position_case: PositionCase, - repeat_index: int, - press_tolerance: float, - args: argparse.Namespace, - recorded_count: int, -) -> PressCaseResult: - """Run one object-position Press benchmark case.""" - case_id = f"{preset.object_type}:{position_case.name}:r{repeat_index}" - try: - _reset_robot(robot, initial_qpos) - initial_obj_pose = reset_rigid_object_xy( - obj=obj, - base_pose=base_obj_pose, - xy=position_case.xy, - sim=sim, - settle_steps=2, - ) - move_target, press_target = _make_press_targets(obj, preset) - - elapsed, mem_delta, peak_gpu, planning_success, traj = _timed_atomic_run( - atomic_engine=atomic_engine, - move_target=move_target, - press_target=press_target, - ) - video_path = None - if should_record_case(args, recorded_count, bool(planning_success)): - _reset_robot(robot, initial_qpos) - reset_rigid_object(obj, initial_obj_pose) - video_path = replay_trajectory_with_recording( - sim=sim, - robot=robot, - traj=traj, - args=args, - video_path=build_video_output_path( - args, - "atomic_action_press", - (f"{preset.object_type}_{position_case.name}" f"_r{repeat_index}"), - ), - ) - _reset_robot(robot, initial_qpos) - reset_rigid_object(obj, initial_obj_pose) - - center_hit = False - xy_error_m: float | None = None - hit_step: int | None = None - failure_reason = "" - if planning_success: - center_hit, xy_error_m, raw_hit_step = _compute_press_center_check( - robot=robot, - traj=traj, - obj=obj, - object_height=preset.size[2], - tolerance=press_tolerance, - ) - hit_step = raw_hit_step if raw_hit_step >= 0 else None - if not center_hit: - failure_reason = "center_miss" - else: - failure_reason = "planning_failed" - - return PressCaseResult( - case_id=case_id, - object_type=preset.object_type, - material_name=preset.material_name, - quadrant=position_case.quadrant, - position_case=position_case.name, - init_xy=position_case.xy, - repeat_index=repeat_index, - planning_success=planning_success, - center_hit=center_hit, - cost_time_ms=elapsed * 1000.0, - cpu_delta_mb=mem_delta["cpu_mb"], - gpu_delta_mb=mem_delta["gpu_mb"], - peak_gpu_mb=peak_gpu, - xy_error_m=xy_error_m, - hit_step=hit_step, - trajectory_waypoints=int(traj.shape[1]) if traj.ndim >= 2 else 0, - failure_reason=failure_reason, - video_path=str(video_path) if video_path is not None else "", - ) - except Exception as exc: - return PressCaseResult( - case_id=case_id, - object_type=preset.object_type, - material_name=preset.material_name, - quadrant=position_case.quadrant, - position_case=position_case.name, - init_xy=position_case.xy, - repeat_index=repeat_index, - planning_success=False, - center_hit=False, - cost_time_ms=0.0, - cpu_delta_mb=0.0, - gpu_delta_mb=0.0, - peak_gpu_mb=0.0, - xy_error_m=None, - hit_step=None, - trajectory_waypoints=0, - failure_reason=f"exception:{type(exc).__name__}:{exc}", - ) - - -def _build_perf_rows(results: list[PressCaseResult]) -> list[dict[str, object]]: - """Build Time & Memory table rows.""" - rows: list[dict[str, object]] = [] - for result in results: - rows.append( - { - "sample_size": 1, - "impl": "press", - "case_id": result.case_id, - "object_type": result.object_type, - "material": result.material_name, - "quadrant": result.quadrant, - "position_case": result.position_case, - "init_xy": f"({result.init_xy[0]:.3f},{result.init_xy[1]:.3f})", - "repeat": result.repeat_index, - "cost_time_ms": _format_float(result.cost_time_ms), - "cpu_delta_mb": _format_float(result.cpu_delta_mb), - "gpu_delta_mb": _format_float(result.gpu_delta_mb), - "peak_gpu_mb": _format_float(result.peak_gpu_mb), - } - ) - return rows - - -def _build_metric_rows(results: list[PressCaseResult]) -> list[dict[str, object]]: - """Build Success & Other Metrics table rows.""" - rows: list[dict[str, object]] = [] - for result in results: - overall_success = result.planning_success and result.center_hit - rows.append( - { - "sample_size": 1, - "impl": "press", - "case_id": result.case_id, - "object_type": result.object_type, - "material": result.material_name, - "quadrant": result.quadrant, - "position_case": result.position_case, - "success_rate": f"{float(overall_success):.6f}", - "planning_success_rate": f"{float(result.planning_success):.6f}", - "center_hit_rate": f"{float(result.center_hit):.6f}", - "xy_error_m": _format_float(result.xy_error_m), - "hit_step": result.hit_step if result.hit_step is not None else "N/A", - "trajectory_waypoints": result.trajectory_waypoints, - "failure_reason": result.failure_reason or "N/A", - } - ) - return rows - - -def _build_leaderboard_rows(results: list[PressCaseResult]) -> list[dict[str, object]]: - """Aggregate and rank object-conditioned Press variants by success rate.""" - aggregate: dict[str, dict[str, float | set[str]]] = {} - for result in results: - algorithm = f"press:{result.object_type}" - if algorithm not in aggregate: - aggregate[algorithm] = { - "overall_success_sum": 0.0, - "planning_success_sum": 0.0, - "xy_error_sum": 0.0, - "xy_error_count": 0.0, - "cost_time_sum": 0.0, - "case_count": 0.0, - "quadrants": set(), - } - - stats = aggregate[algorithm] - stats["overall_success_sum"] = float(stats["overall_success_sum"]) + float( - result.planning_success and result.center_hit - ) - stats["planning_success_sum"] = float(stats["planning_success_sum"]) + float( - result.planning_success - ) - if result.xy_error_m is not None and math.isfinite(result.xy_error_m): - stats["xy_error_sum"] = float(stats["xy_error_sum"]) + result.xy_error_m - stats["xy_error_count"] = float(stats["xy_error_count"]) + 1.0 - stats["cost_time_sum"] = float(stats["cost_time_sum"]) + result.cost_time_ms - stats["case_count"] = float(stats["case_count"]) + 1.0 - quadrants = stats["quadrants"] - if isinstance(quadrants, set): - quadrants.add(result.quadrant) - - ranked = sorted( - aggregate.items(), - key=lambda item: ( - float(item[1]["overall_success_sum"]) - / max(float(item[1]["case_count"]), 1.0), - -float(item[1]["cost_time_sum"]) / max(float(item[1]["case_count"]), 1.0), - ), - reverse=True, - ) - - rows: list[dict[str, object]] = [] - for rank, (algorithm, stats) in enumerate(ranked, start=1): - case_count = max(float(stats["case_count"]), 1.0) - xy_error_count = float(stats["xy_error_count"]) - avg_xy_error = ( - float(stats["xy_error_sum"]) / xy_error_count - if xy_error_count > 0.0 - else None - ) - quadrants = stats["quadrants"] - quadrant_coverage = ( - ",".join(sorted(quadrants)) if isinstance(quadrants, set) else "" - ) - rows.append( - { - "rank": rank, - "algorithm": algorithm, - "overall_success_rate": ( - f"{float(stats['overall_success_sum']) / case_count:.2%}" - ), - "planning_success_rate": ( - f"{float(stats['planning_success_sum']) / case_count:.2%}" - ), - "avg_xy_error_m": _format_float(avg_xy_error), - "avg_cost_time_ms": _format_float( - float(stats["cost_time_sum"]) / case_count - ), - "evaluated_cases": int(case_count), - "quadrant_coverage": quadrant_coverage, - } - ) - return rows - - -def _print_case_result(result: PressCaseResult) -> None: - """Print one aligned case result line.""" - overall_success = result.planning_success and result.center_hit - print( - f" {result.case_id:<28} " - f"time={result.cost_time_ms:>10.2f} ms | " - f"CPU delta={result.cpu_delta_mb:+.1f} MB " - f"GPU delta={result.gpu_delta_mb:+.1f} MB " - f"peak GPU={result.peak_gpu_mb:.1f} MB | " - f"success={overall_success} " - f"xy_error={_format_float(result.xy_error_m, precision=4)}" - ) - if result.failure_reason: - print(f" reason={result.failure_reason}") - - -def _build_notes( - object_presets: list[ObjectPreset], - position_cases: list[PositionCase], - repeat: int, - video_paths: list[str], - profile: str, -) -> list[str]: - """Build report notes with benchmark coverage metadata.""" - quadrant_counts: dict[str, int] = {} - for position_case in position_cases: - quadrant_counts[position_case.quadrant] = ( - quadrant_counts.get(position_case.quadrant, 0) + 1 - ) - return [ - f"Profile: {profile}", - "Object presets: " - + ", ".join( - f"{preset.object_type}/{preset.material_name}/size={preset.size}" - for preset in object_presets - ), - "Position cases per quadrant: " - + ", ".join( - f"{quadrant}={count}" for quadrant, count in sorted(quadrant_counts.items()) - ), - f"CPU memory backend: {CPU_MEMORY_BACKEND}", - f"Repeat per object-position case: {repeat}", - "Replay videos: " + (", ".join(video_paths) if video_paths else "disabled"), - "success_rate is 1 only when planning succeeds and the Press trajectory " - "reaches the object top center.", - ] - - -def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: - """Run all atomic-action benchmarks and write the markdown report.""" - args = _parse_args() if args is None else args - if args.repeat < 1: - raise ValueError("--repeat must be at least 1.") - profile = resolve_profile(args) - _ensure_runtime_imports() - - object_presets = _select_object_presets(args.object_types, profile) - position_cases = _select_position_cases(args.position_cases, profile) - repeat = 1 if profile == "smoke" else args.repeat - - print("=" * 60) - print("Atomic Action Press Benchmark") - print("=" * 60) - print( - "Coverage: " - f"profile={profile}, {len(object_presets)} object presets x " - f"{len(position_cases)} position cases x {repeat} repeat(s)" - ) - - sim = initialize_simulation(args) - robot = create_robot(sim) - create_table(sim) - initial_qpos = robot.get_qpos().clone() - motion_gen = MotionGenerator( - cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) - ) - atomic_engine = _build_atomic_engine(motion_gen, robot, sim.device) - object_pool = {} - for object_index, preset in enumerate(object_presets): - obj = _create_benchmark_object(sim, preset, position_cases[0], object_index) - settle_object(sim, obj, step=2) - base_pose = obj.get_local_pose(to_matrix=True).clone() - park_rigid_object(obj, base_pose, index=object_index, sim=sim) - object_pool[preset.object_type] = (obj, base_pose) - - results: list[PressCaseResult] = [] - video_paths: list[str] = [] - print("\n=== Press Object/Position Sweep ===") - for preset in object_presets: - obj, base_pose = object_pool[preset.object_type] - for parked_index, parked_preset in enumerate(object_presets): - if parked_preset.object_type == preset.object_type: - continue - parked_obj, parked_base_pose = object_pool[parked_preset.object_type] - park_rigid_object(parked_obj, parked_base_pose, index=parked_index, sim=sim) - for position_case in position_cases: - for repeat_index in range(repeat): - result = _run_press_case( - sim=sim, - robot=robot, - atomic_engine=atomic_engine, - initial_qpos=initial_qpos, - obj=obj, - base_obj_pose=base_pose, - preset=preset, - position_case=position_case, - repeat_index=repeat_index, - press_tolerance=args.press_tolerance, - args=args, - recorded_count=len(video_paths), - ) - results.append(result) - if result.video_path: - video_paths.append(result.video_path) - _print_case_result(result) - - perf_rows = _build_perf_rows(results) - metric_rows = _build_metric_rows(results) - leaderboard_rows = _build_leaderboard_rows(results) - report_path = _write_markdown_report( - benchmark_name="atomic_action_press", - perf_rows=perf_rows, - metric_rows=metric_rows, - leaderboard_rows=leaderboard_rows, - notes=_build_notes( - object_presets, - position_cases, - repeat, - video_paths, - profile, - ), - ) - - print("\n" + "=" * 60) - print("Benchmarks complete.") - print(f"Markdown report saved: {report_path}") - print("=" * 60) - return report_path - - -def main() -> None: - """Run the CLI entry point.""" - try: - run_all_benchmarks() - except RuntimeError as exc: - raise SystemExit(str(exc)) from exc - - -if __name__ == "__main__": - main() - - -__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/run_benchmark.py b/scripts/benchmark/atomic_action/run_benchmark.py index ab20494e6..d50107756 100644 --- a/scripts/benchmark/atomic_action/run_benchmark.py +++ b/scripts/benchmark/atomic_action/run_benchmark.py @@ -17,7 +17,7 @@ """Dispatch benchmarks for all atomic actions. Run a single action benchmark or all action benchmarks in sequence. -Run: embodichain benchmark atomic-action --action press +Run: embodichain benchmark atomic-action --action move_end_effector """ from __future__ import annotations @@ -43,11 +43,9 @@ "pick_up": "scripts.benchmark.atomic_action.pickup_benchmark", "move_held_object": "scripts.benchmark.atomic_action.move_held_object_benchmark", "place": "scripts.benchmark.atomic_action.place_benchmark", - "press": "scripts.benchmark.atomic_action.press_benchmark", } DEFAULT_ACTIONS = tuple(ACTION_MODULES.keys()) MESH_OBJECT_ACTIONS = {"pick_up", "move_held_object", "place"} -PRESS_OBJECT_TYPES = {"bottle", "mug", "wooden_block", "all"} MESH_OBJECT_TYPES = {*MESH_OBJECT_PRESETS.keys(), "all"} @@ -57,7 +55,7 @@ def add_benchmark_args(parser: argparse.ArgumentParser) -> None: "--action", nargs="+", choices=(*ACTION_MODULES.keys(), "all"), - default=["press"], + default=["move_end_effector"], help="Atomic action benchmark(s) to run. Use 'all' for every action.", ) parser.add_argument( @@ -148,8 +146,6 @@ def _validate_object_types_for_actions( for action_name in selected_actions: if action_name in MESH_OBJECT_ACTIONS: validators[action_name] = MESH_OBJECT_TYPES - elif action_name == "press": - validators[action_name] = PRESS_OBJECT_TYPES invalid_parts = [] for action_name, valid_types in validators.items(): @@ -177,7 +173,6 @@ def _make_child_args(args: argparse.Namespace) -> argparse.Namespace: renderer=args.renderer, object_types=args.object_types, position_cases=args.position_cases, - press_tolerance=0.01, pose_cases=["all"], sequence_cases=["all"], approach_cases=args.approach_cases, @@ -224,7 +219,7 @@ def _make_child_cli_args(args: argparse.Namespace, action_name: str) -> list[str "--video_hold_steps", str(args.video_hold_steps), ] - if action_name in {"pick_up", "move_held_object", "place", "press"}: + if action_name in {"pick_up", "move_held_object", "place"}: if args.object_types: child_args.append("--object_types") child_args.extend(args.object_types) diff --git a/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md b/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md index 66463bb2c..8e561933c 100644 --- a/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md +++ b/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md @@ -561,7 +561,6 @@ Suggested coverage: | PickUp | Approach/lift plan succeeds, `held_object` is created, minimum object lift is reached, no drop | | MoveHeldObject | Object reaches target pose, grasp remains stable, object drift/tilt stays within threshold | | Place | Place pose reached, release succeeds, final object pose is correct and stable | -| Press | Press depth and valid contact/force reached, retract succeeds, no abnormal object motion | | Pick-Move-Place | Every stage succeeds in sequence; final object pose and release state are correct | Record: @@ -675,7 +674,6 @@ tracks: - pick_up - move_held_object - place - - press - pick_move_place scenario_overrides: @@ -960,7 +958,7 @@ Minimum tests: - Parameterize planner construction in Atomic Action benchmarks. - Explicitly separate `ik_interp` and `motion_gen`. - Reuse current object/position/approach profiles and physical-success rules. -- Add MoveEndEffector, PickUp, MoveHeldObject, Place, Press, and +- Add MoveEndEffector, PickUp, MoveHeldObject, Place, and Pick-Move-Place. - Add controller tracking, collision/contact, and stable-hold metrics. diff --git a/scripts/tutorials/atomic_action/hand_over.py b/scripts/tutorials/atomic_action/hand_over.py index 5c75bbb73..e03175909 100644 --- a/scripts/tutorials/atomic_action/hand_over.py +++ b/scripts/tutorials/atomic_action/hand_over.py @@ -224,9 +224,6 @@ def run_handover_demo( 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_options = HandOverOptions( diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py index ca30a1927..dff06c935 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Demonstrate Press on the center of a regular wooden block.""" +"""Demonstrate Press on an articulation link or rigid object.""" from __future__ import annotations @@ -28,138 +28,171 @@ import torch +from embodichain.data import get_data_path from embodichain.lab.sim.atomic_actions import ( ActionBinding, ActionInvocation, AtomicActionEngine, ControlPartCommandProfile, - EndEffectorPoseGoal, - PressOptions, - PressGoal, + EntityState, MotionPolicy, + ObjectSemantics, + PressAffordance, + PressGoal, + PressOptions, + SceneEntityPose, + SceneSnapshot, ) from embodichain.lab.sim.cfg import ( - RigidBodyAttributesCfg, + ArticulationCfg, + JointDrivePropertiesCfg, RigidObjectCfg, ) -from embodichain.lab.sim.material import VisualMaterialCfg -from embodichain.lab.sim.objects import RigidObject +from embodichain.lab.sim.objects import Articulation, RigidObject from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_tutorial_robot, - broadcast_pose_batch, - create_curobo_motion_generator, + add_ur5_gripper_robot, + create_toppra_motion_generator, create_tutorial_argument_parser, create_tutorial_simulation, - draw_axis_marker, - format_tensor, get_hand_open_close_qpos, - make_top_down_eef_pose, prepare_tutorial_scene, replay_trajectory, run_tutorial, ) -MOVE_SAMPLE_INTERVAL = 60 -PRESS_SAMPLE_INTERVAL = 90 +MICROWAVE_ASSET = "MicrowaveOven/microwave_oven_with_inertials.urdf" +BUTTON_LINK_NAME = "button_cap" +MICROWAVE_POSITION = (-1.0, -0.30, 0.4) +MICROWAVE_ORIENTATION = (0.0, 0.0, 90) # degrees +PRESS_SAMPLE_INTERVAL = 140 HAND_INTERP_STEPS = 12 -POST_TRAJECTORY_STEPS = 180 -BLOCK_SIZE = (0.12, 0.12, 0.06) -PRESS_CLEARANCE = 0.13 -PRESS_SURFACE_OFFSET = 0.003 -DEFAULT_PRESS_TOLERANCE = 0.01 +POST_TRAJECTORY_STEPS = 240 +RIGID_BUTTON_POSITION = (-0.7, -0.00, 0.70) +RIGID_BUTTON_SIZE = (0.04, 0.02, 0.04) +BUTTON_SCENE_ENTITY_ID = "press-target" def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the Press tutorial.""" parser = create_tutorial_argument_parser( - "Demonstrate Press on a wooden block.", - features=("debug_state", "visualize_axes"), + "Demonstrate Press on an articulation-link or rigid button.", + features=("visualize_axes",), + ) + parser.add_argument("--press_distance", type=float, default=0.03) + parser.add_argument( + "--press_position", + type=float, + nargs=3, + default=None, + metavar=("X", "Y", "Z"), + help="Optional target-local press position overriding the affordance.", ) parser.add_argument( - "--press_tolerance", type=float, default=DEFAULT_PRESS_TOLERANCE + "--rigid_object", + action="store_true", + help="Use a standalone rigid button instead of the microwave link.", ) - parser.add_argument("--block_pos", type=float, nargs=2, default=(-0.30, -0.12)) return parser.parse_args() -def create_wooden_block(sim, center: list[float]) -> RigidObject: - """Create the static block used as a press target.""" - return sim.add_rigid_object( - cfg=RigidObjectCfg( - uid="wooden_block", - shape=CubeCfg( - size=list(BLOCK_SIZE), - visual_material=VisualMaterialCfg( - uid="wooden_block_mat", - base_color=[0.58, 0.32, 0.14, 1.0], - roughness=0.85, - ), +def create_microwave(sim) -> Articulation: + """Create the fixed-base microwave articulation used by the demo.""" + microwave = sim.add_articulation( + cfg=ArticulationCfg( + uid="microwave", + fpath=get_data_path(MICROWAVE_ASSET), + init_pos=MICROWAVE_POSITION, + init_qpos=(0, 0, 0, 0), + init_rot=MICROWAVE_ORIENTATION, + drive_pros=JointDrivePropertiesCfg( + stiffness=1e-3, damping=1e2, max_effort=1e-2 ), - body_type="static", - attrs=RigidBodyAttributesCfg(dynamic_friction=0.8, static_friction=0.9), - init_pos=center, + fix_base=True, ) ) + sim.update(step=10) + return microwave -def compute_press_center_check( - robot, - trajectory: torch.Tensor, - block: RigidObject, - tolerance: float, -) -> tuple[bool, float, int, torch.Tensor, torch.Tensor]: - """Return whether the press trajectory reaches the block center tolerance.""" - arm_joint_ids = robot.get_joint_ids(name="arm") - start = MOVE_SAMPLE_INTERVAL + HAND_INTERP_STEPS - arm_traj = trajectory[ - :, start : MOVE_SAMPLE_INTERVAL + PRESS_SAMPLE_INTERVAL, arm_joint_ids - ] - fk_pose = torch.stack( - [ - robot.compute_fk(qpos=qpos, name="arm", to_matrix=True) - for qpos in arm_traj.unbind(dim=1) - ], - dim=1, - ) - block_center = block.get_local_pose(to_matrix=True)[:, :3, 3] - target_z = block_center[:, 2] + 0.5 * BLOCK_SIZE[2] + PRESS_SURFACE_OFFSET - xy_error = torch.linalg.norm( - fk_pose[:, :, :2, 3] - block_center[:, None, :2], dim=2 - ) - z_error = torch.abs(fk_pose[:, :, 2, 3] - target_z[:, None]) - best_idx = (xy_error + z_error).argmin(dim=1) - env_idx = torch.arange(trajectory.shape[0], device=trajectory.device) - best_pos = fk_pose[env_idx, best_idx, :3, 3] - center_error = torch.linalg.norm(best_pos[:, :2] - block_center[:, :2], dim=1) - worst_env = int(center_error.argmax().item()) - expected = torch.stack( - [block_center[worst_env, 0], block_center[worst_env, 1], target_z[worst_env]] +def create_rigid_button(sim) -> RigidObject: + """Create the standalone static rigid button used by the optional demo.""" + button = sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="rigid_button", + shape=CubeCfg(size=list(RIGID_BUTTON_SIZE)), + body_type="static", + init_pos=RIGID_BUTTON_POSITION, + ) ) + sim.update(step=10) + return button + + +def create_button_semantics( + target: Articulation | RigidObject, +) -> tuple[ObjectSemantics, torch.Tensor]: + """Create press semantics for an articulation-link or rigid button.""" + if isinstance(target, Articulation): + vertices, _ = target.get_link_vert_face(BUTTON_LINK_NAME) + target_pose = target.get_link_pose(BUTTON_LINK_NAME, to_matrix=True) + press_axis = torch.tensor([0.0, 0.0, -1.0], device=target.device) + affordance = PressAffordance( + # button_cap's local -z direction matches the prismatic joint's + # inward press direction in this asset. + press_axis=press_axis, + press_position=_surface_center(vertices, press_axis), + ) + label = "microwave_start_button" + else: + vertices = target.get_vertices(env_ids=[0], scale=True)[0] + target_pose = target.get_local_pose(to_matrix=True) + press_axis = torch.tensor([-1.0, 0.0, 0.0], device=target.device) + affordance = PressAffordance( + press_axis=press_axis, + press_position=_surface_center(vertices, press_axis), + ) + label = "rigid_button" return ( - bool(torch.all(center_error <= tolerance)), - float(center_error[worst_env].item()), - start + int(best_idx[worst_env].item()), - best_pos[worst_env], - expected, + ObjectSemantics( + label=label, + geometry={}, + entity_id=BUTTON_SCENE_ENTITY_ID, + affordance=affordance, + ), + target_pose, ) +def _surface_center( + vertices: torch.Tensor, + inward_axis: torch.Tensor, +) -> tuple[float, float, float]: + """Return the center of the outermost mesh face opposite inward travel.""" + vertices = torch.as_tensor(vertices, dtype=torch.float32, device=inward_axis.device) + axis = inward_axis.to(dtype=torch.float32) + axis = axis / torch.linalg.vector_norm(axis) + projection = torch.matmul(vertices, axis) + surface = vertices[torch.isclose(projection, projection.min(), atol=1.0e-5)] + point = surface.mean(dim=0) + return tuple(float(value) for value in point) + + def main() -> None: - """Plan, verify, and replay MoveEndEffector followed by Press.""" + """Plan and replay Press for the selected target object type.""" args = parse_arguments() sim = create_tutorial_simulation(args) - robot = add_tutorial_robot(sim, args.robot) - block = create_wooden_block(sim, [*args.block_pos, 0.5 * BLOCK_SIZE[2]]) - if sim.device.type == "cuda": - sim.init_gpu_physics() - block.reset() - sim.update(step=5) - block.clear_dynamics() + robot = add_ur5_gripper_robot( + sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0] + ) + target = create_rigid_button(sim) if args.rigid_object else create_microwave(sim) + hand_open, hand_close = get_hand_open_close_qpos(robot, close_qpos=0.040) + motion_gen = create_toppra_motion_generator(robot) + semantics, target_pose = create_button_semantics(target) + affordance = semantics.affordance + assert isinstance(affordance, PressAffordance) - motion_gen = create_curobo_motion_generator(robot) - hand_open, hand_close = get_hand_open_close_qpos(robot) engine = AtomicActionEngine( motion_generator=motion_gen, control_profiles={ @@ -169,89 +202,76 @@ def main() -> None: ) }, ) - block_center = block.get_local_pose(to_matrix=True)[0, :3, 3] - press_position = block_center.clone() - press_position[2] += 0.5 * BLOCK_SIZE[2] + PRESS_SURFACE_OFFSET - move_position = press_position.clone() - move_position[2] += PRESS_CLEARANCE - PRESS_SURFACE_OFFSET - num_envs = robot.get_qpos().shape[0] - move_target = broadcast_pose_batch(make_top_down_eef_pose(move_position), num_envs) - press_target = broadcast_pose_batch( - make_top_down_eef_pose(press_position), num_envs - ) - if not args.no_vis_eef_axis: - draw_axis_marker(sim, "press_target_axis", press_target) wait_for_user = prepare_tutorial_scene( - sim, args, "Inspect the wooden block, then press Enter to plan..." + sim, + args, + "Inspect the button target, then press Enter to plan Press...", ) - binding = ActionBinding( - manipulators={"primary": "arm"}, - end_effectors={"primary": "hand"}, - ) compiled = engine.compile( ( ActionInvocation( - "move_end_effector", - EndEffectorPoseGoal(move_target), - binding, - MotionPolicy( - strategy="motion_gen", - sample_count=MOVE_SAMPLE_INTERVAL, + skill_id="press", + goal=PressGoal( + semantics, + SceneEntityPose(BUTTON_SCENE_ENTITY_ID), ), - ), - ActionInvocation( - "press", - PressGoal(press_target), - binding, - MotionPolicy( - strategy="motion_gen", - sample_count=PRESS_SAMPLE_INTERVAL, + binding=ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, ), + motion_policy=MotionPolicy(sample_count=PRESS_SAMPLE_INTERVAL), skill_options=PressOptions( hand_interp_steps=HAND_INTERP_STEPS, + approach_distance=0.12, + press_distance=args.press_distance, + press_position=( + None + if args.press_position is None + else tuple(args.press_position) + ), ), ), - ) + ), + context=engine.initial_context( + scene=SceneSnapshot( + timestamp=0.0, + version=0, + entities={BUTTON_SCENE_ENTITY_ID: EntityState(target_pose)}, + ) + ), ) 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) - ) - logger.log_info( - "Press center check: " - f"success={is_center_hit}, xy_error={center_error:.4f} m, hit_step={hit_step}, " - f"hit_pos={format_tensor(hit_pos)}, expected={format_tensor(expected_pos)}" - ) - if not is_center_hit: - logger.log_warning( - "Press trajectory did not reach the block center within tolerance." - ) + logger.log_warning("Failed to plan the Press demo trajectory.") return + if isinstance(target, RigidObject): + focus_pose = target.get_local_pose(to_matrix=True) + elif isinstance(target, Articulation): + focus_pose = target.get_link_pose(BUTTON_LINK_NAME, to_matrix=True) + else: + raise ValueError("Unsupported target type for Press demo.") + focus_position = [focus_pose[0, 0, 3], focus_pose[0, 1, 3], focus_pose[0, 2, 3]] + camera_position = [ + focus_position[0] + 0.0, + focus_position[1] + 0.3, + focus_position[2] + 0.2, + ] + look_at = [camera_position, focus_position, [0, 0, 1]] if wait_for_user: input("Press Enter to replay the Press demo...") - - def log_state(step_idx: int, total_steps: int) -> None: - if args.debug_state and ( - step_idx % max(1, total_steps // 10) == 0 or step_idx == total_steps - 1 - ): - logger.log_info( - f"replay step {step_idx}/{total_steps - 1}: " - f"pos={format_tensor(block.get_local_pose(to_matrix=True)[0, :3, 3])}" - ) - replay_trajectory( sim, robot, compiled.trajectory, args, - video_prefix="press_auto_play", + video_prefix=( + "press_rigid_button_auto_play" + if args.rigid_object + else "press_microwave_button_auto_play" + ), hold_steps=POST_TRAJECTORY_STEPS, - on_trajectory_step=log_state, + look_at=look_at, ) if wait_for_user: input("Press Enter to exit the simulation...") diff --git a/scripts/tutorials/atomic_action/slide.py b/scripts/tutorials/atomic_action/slide.py new file mode 100644 index 000000000..599a99861 --- /dev/null +++ b/scripts/tutorials/atomic_action/slide.py @@ -0,0 +1,310 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Demonstrate Slide on a translating drawer.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from typing import Literal + +_REPO_ROOT = Path(__file__).resolve().parents[3] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +import torch + +from embodichain.data import get_data_path +from embodichain.lab.sim import SimulationManager +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + AtomicActionEngine, + ControlPartCommandProfile, + EntityState, + MotionPolicy, + ObjectSemantics, + SlideAffordance, + SlideGoal, + SlideOptions, + SceneEntityPose, + SceneSnapshot, +) +from embodichain.lab.sim.cfg import ( + ArticulationCfg, + JointDrivePropertiesCfg, + RigidBodyAttributesCfg, +) +from embodichain.lab.sim.objects import Articulation +from embodichain.toolkits.graspkit.pg_grasp.antipodal_generator import ( + AntipodalSamplerCfg, + GraspGeneratorCfg, +) +from embodichain.toolkits.graspkit.pg_grasp.gripper_collision_checker import ( + GripperCollisionCfg, +) +from embodichain.utils import logger +from scripts.tutorials.atomic_action.tutorial_utils import ( + add_ur5_gripper_robot, + create_toppra_motion_generator, + create_tutorial_argument_parser, + create_tutorial_simulation, + draw_axis_marker, + get_hand_open_close_qpos, + prepare_tutorial_scene, + replay_trajectory, + run_tutorial, +) + +DRAWER_ASSET = "Drawer/model_split_links_with_inertials.urdf" +HANDLE_LINK_NAME = "large_handle_bar" +DRAWER_POSITION = (-1.1, 0.0, 0.0) +DRAWER_ORIENTATION = (0.0, 0.0, 90.0) # degrees +TRANSLATION_AXIS = (0.0, 1.0, 0.0) # handle-link frame, approach/push direction +TRAJECTORY_SAMPLE_COUNT = 140 +HAND_INTERP_STEPS = 12 +POST_TRAJECTORY_STEPS = 240 +HANDLE_SCENE_ENTITY_ID = "drawer-large-handle" + + +def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the drawer pull/push tutorial.""" + parser = create_tutorial_argument_parser( + "Pull a drawer open, then push it closed with Slide.", + features=("grasp_sampling", "visualize_axes"), + ) + parser.add_argument("--translation_distance", type=float, default=0.18) + parser.add_argument("--approach_distance", type=float, default=0.10) + return parser.parse_args() + + +def create_drawer( + sim: SimulationManager, +) -> Articulation: + """Create the fixed-base drawer in its closed initial state.""" + drawer = sim.add_articulation( + cfg=ArticulationCfg( + uid="drawer", + fpath=get_data_path(DRAWER_ASSET), + init_pos=DRAWER_POSITION, + init_rot=DRAWER_ORIENTATION, + init_qpos=(0.0,), + drive_pros=JointDrivePropertiesCfg(drive_type="none"), + attrs=RigidBodyAttributesCfg( + static_friction=1.0, + dynamic_friction=1.0, + ), + fix_base=True, + ) + ) + sim.update(step=10) + return drawer + + +def create_drawer_semantics( + drawer: Articulation, + *, + n_sample: int, + force_reannotate: bool, +) -> ObjectSemantics: + """Create sampled-grasp translation semantics for the drawer handle. + + Args: + drawer: Drawer articulation that owns the target handle link. + n_sample: Number of antipodal surface samples. + force_reannotate: Whether to ignore a cached grasp annotation. + + Returns: + Pure target-local semantics for the handle's pull/push affordance. + """ + vertices, triangles = drawer.get_link_vert_face(HANDLE_LINK_NAME) + return ObjectSemantics( + label="drawer_large_handle", + geometry={}, + entity_id=HANDLE_SCENE_ENTITY_ID, + affordance=SlideAffordance( + mesh_vertices=torch.as_tensor(vertices), + mesh_triangles=torch.as_tensor(triangles), + translation_axis=torch.tensor( + TRANSLATION_AXIS, + dtype=torch.float32, + device=drawer.device, + ), + generator_cfg=GraspGeneratorCfg( + antipodal_sampler_cfg=AntipodalSamplerCfg( + n_sample=n_sample, + max_length=0.1, + min_length=0.003, + ), + is_partial_annotate=False, + is_filter_ground_collision=False, + ), + gripper_collision_cfg=GripperCollisionCfg( + max_open_length=0.1, + finger_length=0.1, + y_thickness=0.04, + root_z_width=0.096, + open_check_margin=0.03, + point_sample_dense=0.012, + ), + force_reannotate=force_reannotate, + ), + ) + + +def create_invocation( + semantics: ObjectSemantics, + *, + direction: Literal["pull", "push"], + approach_distance: float, + translation_distance: float, +) -> ActionInvocation: + """Create one pull or push invocation for the shared drawer target. + + Args: + semantics: Drawer-handle semantics shared by both operations. + direction: Whether this invocation pulls open or pushes closed. + approach_distance: Pre-grasp offset opposite the approach axis. + translation_distance: Drawer travel distance for this operation. + + Returns: + A grounded pull/push invocation for the tutorial UR5. + """ + return ActionInvocation( + skill_id="slide", + goal=SlideGoal( + semantics, + SceneEntityPose(HANDLE_SCENE_ENTITY_ID), + ), + binding=ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ), + motion_policy=MotionPolicy(sample_count=TRAJECTORY_SAMPLE_COUNT), + skill_options=SlideOptions( + direction=direction, + hand_interp_steps=HAND_INTERP_STEPS, + approach_distance=approach_distance, + translation_distance=translation_distance, + ), + ) + + +def main() -> None: + """Plan and replay a drawer pull followed by a push.""" + args = parse_arguments() + if args.translation_distance <= 0.0: + raise ValueError("--translation_distance must be positive.") + if args.translation_distance > 0.285: + raise ValueError( + "--translation_distance must not exceed the drawer limit 0.285." + ) + + sim = create_tutorial_simulation(args) + robot = add_ur5_gripper_robot( + sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0], tcp_z=0.15 + ) + drawer = create_drawer(sim) + hand_open, hand_close = get_hand_open_close_qpos(robot) + motion_gen = create_toppra_motion_generator(robot) + semantics = create_drawer_semantics( + drawer, + n_sample=args.n_sample, + force_reannotate=args.force_reannotate, + ) + affordance = semantics.affordance + assert isinstance(affordance, SlideAffordance) + if not args.no_vis_eef_axis: + draw_axis_marker( + sim, + "drawer_handle_link_pose", + drawer.get_link_pose(HANDLE_LINK_NAME, to_matrix=True), + ) + + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) + wait_for_user = prepare_tutorial_scene( + sim, + args, + "Inspect the closed drawer, then press Enter to plan the pull...", + ) + + for scene_version, direction in enumerate(("pull", "push")): + if direction == "push" and wait_for_user: + input( + "Pull replay finished. Press Enter to read the moved handle " + "pose and plan the push..." + ) + + handle_pose = drawer.get_link_pose(HANDLE_LINK_NAME, to_matrix=True) + compiled = engine.compile( + ( + create_invocation( + semantics, + direction=direction, + approach_distance=args.approach_distance, + translation_distance=args.translation_distance, + ), + ), + context=engine.initial_context( + scene=SceneSnapshot( + timestamp=float(scene_version), + version=scene_version, + entities={ + HANDLE_SCENE_ENTITY_ID: EntityState(handle_pose), + }, + ) + ), + ) + if not compiled.plan_success.all(): + logger.log_warning(f"Failed to plan the Slide {direction} trajectory.") + return + + if wait_for_user: + input(f"Press Enter to replay the drawer {direction}...") + focus_pose = drawer.get_link_pose(HANDLE_LINK_NAME, to_matrix=True) + focus_position = [focus_pose[0, 0, 3], focus_pose[0, 1, 3], focus_pose[0, 2, 3]] + camera_position = [ + focus_position[0] + 0.5, + focus_position[1] + 0.5, + focus_position[2] + 0.5, + ] + look_at = [camera_position, focus_position, [0, 0, 1]] + replay_trajectory( + sim, + robot, + compiled.trajectory, + args, + video_prefix=f"{direction}_drawer_auto_play", + hold_steps=POST_TRAJECTORY_STEPS, + look_at=look_at, + ) + + if wait_for_user: + input("Press Enter to exit the simulation...") + + +if __name__ == "__main__": + run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/tutorial_utils.py b/scripts/tutorials/atomic_action/tutorial_utils.py index b1ddd5a2b..f90c1b127 100644 --- a/scripts/tutorials/atomic_action/tutorial_utils.py +++ b/scripts/tutorials/atomic_action/tutorial_utils.py @@ -69,7 +69,6 @@ GRIPPER_URDF_PATH = "DH_PGI_140_80/DH_PGI_140_80.urdf" GRIPPER_HAND_JOINT_PATTERN = "gripper_finger1_joint_1" -GRIPPER_TCP_Z = 0.15 GRIPPER_MAX_OPEN_WIDTH = 0.100 GRIPPER_MIN_OPEN_WIDTH = 0.003 GRIPPER_FINGER_LENGTH = 0.10 @@ -78,10 +77,11 @@ DEFAULT_GRIPPER_CLOSE_QPOS = 0.024 DEFAULT_TUTORIAL_LIGHT_POS = (1.0, 0.0, 3.0) _FRANKA_TUTORIAL_BASE_ROTATION = (0.0, 0.0, 180.0) +_DEFAULT_GRIPPER_TCP_Z = 0.17 _GRIPPER_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, 1.0, _DEFAULT_GRIPPER_TCP_Z), (0.0, 0.0, 0.0, 1.0), ) TOP_DOWN_EEF_ROTATION = ( @@ -246,6 +246,7 @@ def add_ur5_gripper_robot( sim: SimulationManager, init_pos: Sequence[float] = (0.0, 0.0, 0.0), init_qpos: Sequence[float] | None = None, + tcp_z: float = _DEFAULT_GRIPPER_TCP_Z, ) -> Robot: """Add the standard UR5 plus PGI gripper tutorial robot. @@ -257,7 +258,11 @@ def add_ur5_gripper_robot( The added robot instance. """ return sim.add_robot( - cfg=create_ur5_gripper_robot_cfg(init_pos=init_pos, init_qpos=init_qpos) + cfg=create_ur5_gripper_robot_cfg( + init_pos=init_pos, + init_qpos=init_qpos, + tcp_z=tcp_z, + ) ) @@ -840,6 +845,7 @@ def clone_local_pose_from_first_env(entity) -> torch.Tensor: def create_ur5_gripper_robot_cfg( init_pos: Sequence[float] = (0.0, 0.0, 0.0), init_qpos: Sequence[float] | None = None, + tcp_z: float = _DEFAULT_GRIPPER_TCP_Z, ) -> RobotCfg: """Build a UR5 arm + DH_PGI_140_80 gripper robot configuration. @@ -895,7 +901,16 @@ def create_ur5_gripper_robot_cfg( GRIPPER_HAND_JOINT_PATTERN: 1e4, }, }, - "solver_cfg": {"arm": {"tcp": _GRIPPER_TCP}}, + "solver_cfg": { + "arm": { + "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], + ] + } + }, "init_qpos": qpos, "init_pos": init_pos, } @@ -1004,7 +1019,6 @@ def create_tutorial_robot_cfg( "DEFAULT_GRIPPER_CLOSE_QPOS", "DEFAULT_TUTORIAL_LIGHT_POS", "GRIPPER_HAND_JOINT_PATTERN", - "GRIPPER_TCP_Z", "GRIPPER_URDF_PATH", "TOP_DOWN_EEF_ROTATION", "TutorialCliFeature", diff --git a/scripts/tutorials/atomic_action/twist.py b/scripts/tutorials/atomic_action/twist.py new file mode 100644 index 000000000..9d85388be --- /dev/null +++ b/scripts/tutorials/atomic_action/twist.py @@ -0,0 +1,257 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Demonstrate Twist on an articulation link or rigid object.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[3] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +import torch + +from embodichain.data import get_data_path +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + AtomicActionEngine, + ControlPartCommandProfile, + EntityState, + MotionPolicy, + ObjectSemantics, + TwistAffordance, + TwistGoal, + TwistOptions, + SceneEntityPose, + SceneSnapshot, +) +from embodichain.lab.sim.cfg import ( + ArticulationCfg, + JointDrivePropertiesCfg, + RigidObjectCfg, +) +from embodichain.lab.sim.objects import Articulation, RigidObject +from embodichain.lab.sim.shapes import CubeCfg +from embodichain.utils import logger +from scripts.tutorials.atomic_action.tutorial_utils import ( + add_ur5_gripper_robot, + create_toppra_motion_generator, + create_tutorial_argument_parser, + create_tutorial_simulation, + get_hand_open_close_qpos, + prepare_tutorial_scene, + replay_trajectory, + run_tutorial, +) + +MICROWAVE_ASSET = "MicrowaveOven/microwave_oven_with_inertials.urdf" +KNOB_LINK_NAME = "cap_1" +MICROWAVE_POSITION = (-1.0, -0.30, 0.4) +MICROWAVE_ORIENTATION = (0.0, 0.0, 90) # degrees +TWIST_SAMPLE_INTERVAL = 140 +HAND_INTERP_STEPS = 12 +POST_TRAJECTORY_STEPS = 240 +RIGID_KNOB_POSITION = (-0.7, -0.00, 0.70) +RIGID_KNOB_SIZE = (0.05, 0.05, 0.05) +KNOB_SCENE_ENTITY_ID = "twist-target" +KNOB_AXIS_ORIGIN = (0.0, 0.0, 0.0) + + +def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the Twist tutorial.""" + parser = create_tutorial_argument_parser( + "Demonstrate Twist on an articulation-link or rigid knob.", + features=("visualize_axes",), + ) + parser.add_argument("--twist_angle", type=float, default=-0.7853981634) + parser.add_argument( + "--rigid_object", + action="store_true", + help="Use a standalone rigid knob instead of the microwave link.", + ) + return parser.parse_args() + + +def create_microwave(sim) -> Articulation: + """Create the fixed-base microwave articulation used by the demo.""" + microwave = sim.add_articulation( + cfg=ArticulationCfg( + uid="microwave", + fpath=get_data_path(MICROWAVE_ASSET), + init_pos=MICROWAVE_POSITION, + init_rot=MICROWAVE_ORIENTATION, + drive_pros=JointDrivePropertiesCfg( + stiffness=1e-3, damping=1e2, max_effort=1e-2 + ), + fix_base=True, + ) + ) + sim.update(step=10) + return microwave + + +def create_rigid_knob(sim) -> RigidObject: + """Create the standalone static rigid knob used by the optional demo.""" + knob = sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="rigid_knob", + shape=CubeCfg(size=list(RIGID_KNOB_SIZE)), + body_type="static", + init_pos=RIGID_KNOB_POSITION, + ) + ) + sim.update(step=10) + return knob + + +def create_knob_semantics( + target: Articulation | RigidObject, +) -> tuple[ObjectSemantics, torch.Tensor]: + """Create twist semantics for an articulation-link or rigid knob.""" + if isinstance(target, Articulation): + vertices, _ = target.get_link_vert_face(KNOB_LINK_NAME) + target_pose = target.get_link_pose(KNOB_LINK_NAME, to_matrix=True) + affordance = TwistAffordance( + grasp_position=_mesh_center(vertices), + # The cap_1 revolute axis passes through its link-frame origin. + axis_origin=KNOB_AXIS_ORIGIN, + twist_axis=torch.tensor([0.0, 0.0, -1.0], device=target.device), + ) + label = "microwave_power_knob" + else: + vertices = target.get_vertices(env_ids=[0], scale=True)[0] + target_pose = target.get_local_pose(to_matrix=True) + affordance = TwistAffordance( + grasp_position=_mesh_center(vertices), + axis_origin=KNOB_AXIS_ORIGIN, + twist_axis=torch.tensor([-1.0, 0.0, 0.0], device=target.device), + ) + label = "rigid_knob" + return ( + ObjectSemantics( + label=label, + geometry={}, + entity_id=KNOB_SCENE_ENTITY_ID, + affordance=affordance, + ), + target_pose, + ) + + +def _mesh_center(vertices: torch.Tensor) -> tuple[float, float, float]: + """Return an explicit local gripper-center point for a knob mesh.""" + center = torch.as_tensor(vertices, dtype=torch.float32).mean(dim=0) + return tuple(float(value) for value in center) + + +def main() -> None: + """Plan and replay Twist for the selected target object type.""" + args = parse_arguments() + sim = create_tutorial_simulation(args) + robot = add_ur5_gripper_robot( + sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0] + ) + target = create_rigid_knob(sim) if args.rigid_object else create_microwave(sim) + hand_open, hand_close = get_hand_open_close_qpos(robot) + motion_gen = create_toppra_motion_generator(robot) + semantics, target_pose = create_knob_semantics(target) + + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) + wait_for_user = prepare_tutorial_scene( + sim, + args, + "Inspect the knob target, then press Enter to plan Twist...", + ) + + compiled = engine.compile( + ( + ActionInvocation( + skill_id="twist", + goal=TwistGoal( + semantics, + SceneEntityPose(KNOB_SCENE_ENTITY_ID), + ), + binding=ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ), + motion_policy=MotionPolicy(sample_count=TWIST_SAMPLE_INTERVAL), + skill_options=TwistOptions( + hand_interp_steps=HAND_INTERP_STEPS, + pre_grasp_distance=0.12, + twist_angle=args.twist_angle, + ), + ), + ), + context=engine.initial_context( + scene=SceneSnapshot( + timestamp=0.0, + version=0, + entities={KNOB_SCENE_ENTITY_ID: EntityState(target_pose)}, + ) + ), + ) + if not compiled.plan_success.all(): + logger.log_warning("Failed to plan the Twist demo trajectory.") + return + + if isinstance(target, RigidObject): + focus_pose = target.get_local_pose(to_matrix=True) + elif isinstance(target, Articulation): + focus_pose = target.get_link_pose(KNOB_LINK_NAME, to_matrix=True) + else: + raise ValueError("Unsupported target type for Press demo.") + focus_position = [focus_pose[0, 0, 3], focus_pose[0, 1, 3], focus_pose[0, 2, 3]] + camera_position = [ + focus_position[0] + 0.3, + focus_position[1] + 0.3, + focus_position[2] + 0.3, + ] + look_at = [camera_position, focus_position, [0, 0, 1]] + if wait_for_user: + input("Press Enter to replay the Twist demo...") + replay_trajectory( + sim, + robot, + compiled.trajectory, + args, + video_prefix=( + "twist_rigid_knob_auto_play" + if args.rigid_object + else "twist_microwave_knob_auto_play" + ), + hold_steps=POST_TRAJECTORY_STEPS, + look_at=look_at, + ) + if wait_for_user: + input("Press Enter to exit the simulation...") + + +if __name__ == "__main__": + run_tutorial(main) diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index 06bafe4da..5bfcebc92 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -18,7 +18,8 @@ from __future__ import annotations -from typing import TypeVar +import math +from typing import Literal, TypeVar from unittest.mock import Mock import pytest @@ -64,13 +65,23 @@ PlaceOptions, PlanningContext, Press, + PressAffordance, PressGoal, PressOptions, + SlideAffordance, + Slide, + SlideGoal, + SlideOptions, RobotObservation, SceneEntityPose, SceneSnapshot, TaskState, + TwistAffordance, + Twist, + TwistGoal, + TwistOptions, ) +from embodichain.lab.sim.atomic_actions.goals import collect_scene_dependencies from embodichain.lab.sim.common import BatchEntity from embodichain.lab.sim.planners import ( MotionGenerator, @@ -422,11 +433,21 @@ def test_builtin_descriptors_expose_goals_not_legacy_targets() -> None: assert MoveHeldObject.GoalType is HeldObjectPoseGoal assert Place.GoalType == (PlaceGoal, AssembleGoal) assert Press.GoalType is PressGoal + assert Slide.GoalType is SlideGoal + assert Twist.GoalType is TwistGoal assert CoordinatedPickment.GoalType is CoordinatedPickGoal assert CoordinatedPlacement.GoalType is CoordinatedPlacementGoal assert HandOver.GoalType is GraspGoal +def test_interaction_primitives_use_motion_centric_skill_ids() -> None: + assert (Press.skill_id, Slide.skill_id, Twist.skill_id) == ( + "press", + "slide", + "twist", + ) + + @pytest.mark.parametrize( "options", ( @@ -434,6 +455,8 @@ def test_builtin_descriptors_expose_goals_not_legacy_targets() -> None: MoveHeldObjectOptions(), PlaceOptions(), PressOptions(), + SlideOptions(), + TwistOptions(), CoordinatedPickmentOptions(), CoordinatedPlacementOptions(), HandOverOptions(), @@ -656,62 +679,6 @@ def test_move_held_object_requires_projected_attachment() -> None: semantics.entity.get_local_pose.assert_not_called() -def test_move_held_object_moves_only_exclusively_held_rows() -> None: - generator = _motion_generator() - - def move_ik( - pose: torch.Tensor, - name: str, - joint_seed: torch.Tensor, - **_: object, - ) -> tuple[torch.Tensor, torch.Tensor]: - return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed + 0.1 - - generator.robot.compute_ik.side_effect = move_ik - action = _bind_action(generator, MoveHeldObject()) - semantics = _semantics() - task = TaskState( - batch_size=NUM_ENVS, - device="cpu", - held_objects={ - "arm": _held(semantics), - "alternate_arm": _held( - semantics, - env_mask=torch.tensor([True, False]), - ), - }, - ) - context = _context(task) - - plan = _plan_action( - action, - _invocation("move_held_object", HeldObjectPoseGoal(torch.eye(4))), - context, - ) - - assert plan.plan_success.tolist() == [False, True] - assert torch.allclose( - plan.trajectory.positions[0], - context.robot.qpos[0].unsqueeze(0).expand(plan.trajectory.waypoint_count, -1), - ) - assert not torch.allclose(plan.trajectory.positions[1], context.robot.qpos[1]) - - -def test_press_uses_invocation_sample_budget() -> None: - generator = _motion_generator() - action = _bind_action(generator, Press()) - - plan = _plan_action( - action, - _invocation("press", PressGoal(torch.eye(4)), sample_count=12), - _context(), - ) - - assert plan.plan_success.tolist() == [True, True] - assert plan.trajectory.waypoint_count == 12 - assert plan.expected_effects.is_empty - - def test_strategy_and_sample_count_are_not_action_config_fields() -> None: with pytest.raises(TypeError): MoveEndEffectorOptions(strategy="motion_gen") # type: ignore[call-arg] @@ -1083,31 +1050,777 @@ def test_pick_uses_binding_control_part_as_effect_resource() -> None: assert projected.get_held_object("arm") is None -def test_press_closes_hand_without_changing_projected_attachment() -> None: - held = _held() - task = TaskState( - batch_size=NUM_ENVS, - device="cpu", - held_objects={"arm": held}, +def test_twist_plans_six_segments_from_articulation_link() -> None: + affordance = TwistAffordance( + grasp_position=(0.0, 0.0, 0.0), + axis_origin=(0.0, 0.0, 0.0), + twist_axis=torch.tensor([0.0, 1.0, 0.0]), + ) + semantics = ObjectSemantics( + affordance=affordance, + geometry={}, + label="knob", ) generator = _motion_generator() - action = _bind_action( + action = _bind_action(generator, Twist()) + + plan = _plan_action( + action, + ActionInvocation( + skill_id="twist", + goal=TwistGoal(semantics, torch.eye(4)), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=24), + skill_options=TwistOptions(hand_interp_steps=3), + ), + _context(), + ) + + assert plan.plan_success.tolist() == [True, True] + assert plan.trajectory.positions.shape == (NUM_ENVS, 24, ROBOT_DOF) + assert [segment.name for segment in plan.segments] == [ + "approach", + "reach", + "close", + "twist", + "open", + "retract", + ] + assert torch.all( + plan.trajectory.positions[:, plan.segment("close").stop - 1, ARM_DOF:] == 1.0 + ) + assert torch.all( + plan.trajectory.positions[:, plan.segment("open").stop - 1, ARM_DOF:] == 0.0 + ) + first_target = generator.robot.compute_ik.call_args_list[0].kwargs["pose"] + grasp_pose = affordance.get_grasp_pose(torch.eye(4).repeat(NUM_ENVS, 1, 1)) + expected_pre_grasp_position = ( + grasp_pose[:, :3, 3] - grasp_pose[:, :3, 2] * TwistOptions().pre_grasp_distance + ) + assert torch.allclose(first_target[:, :3, 3], expected_pre_grasp_position) + + +def test_twist_plans_from_explicit_rigid_object_pose_snapshot() -> None: + semantics = ObjectSemantics( + affordance=TwistAffordance( + grasp_position=(0.0, 0.0, 0.0), + axis_origin=(0.0, 0.0, 0.0), + twist_axis=torch.tensor([0.0, 1.0, 0.0]), + ), + geometry={}, + label="rigid-knob", + ) + + plan = _plan_action( + _bind_action(_motion_generator(), Twist()), + ActionInvocation( + skill_id="twist", + goal=TwistGoal(semantics, torch.eye(4)), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=24), + skill_options=TwistOptions(hand_interp_steps=3), + ), + _context(), + ) + + assert plan.plan_success.tolist() == [True, True] + + +def test_twist_rotates_grasp_about_explicit_axis_origin() -> None: + action = _bind_action(_motion_generator(), Twist()) + target_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + grasp_pose = target_pose.clone() + grasp_pose[:, 0, 3] = 2.0 + + twisted = action._twisted_grasp_poses( + target_pose, + grasp_pose, + torch.tensor([0.0, 0.0, 1.0]), + (1.0, 0.0, 0.0), + math.pi / 2, + 4, + ) + + assert torch.allclose( + twisted[:, -1, :3, 3], + torch.tensor([1.0, 1.0, 0.0]).expand(NUM_ENVS, -1), + atol=1.0e-6, + ) + + +@pytest.mark.parametrize( + ("goal_factory", "affordance"), + ( + ( + PressGoal, + PressAffordance( + press_axis=torch.tensor([1.0, 0.0, 0.0]), + press_position=(0.0, 0.0, 0.0), + ), + ), + ( + SlideGoal, + SlideAffordance( + mesh_vertices=torch.zeros(3, 3), + mesh_triangles=torch.tensor([[0, 1, 2]]), + ), + ), + ( + TwistGoal, + TwistAffordance( + grasp_position=(0.0, 0.0, 0.0), + axis_origin=(0.0, 0.0, 0.0), + ), + ), + ), +) +def test_interaction_goal_collects_target_scene_dependency( + goal_factory, + affordance, +) -> None: + semantics = ObjectSemantics(affordance=affordance, geometry={}, label="target") + goal = goal_factory(semantics, SceneEntityPose("target-link")) + + assert collect_scene_dependencies(goal) == ("target-link",) + + +def test_open_loop_interaction_primitives_are_explicitly_described() -> None: + assert Press.descriptor().open_loop is True + assert Slide.descriptor().open_loop is True + assert Twist.descriptor().open_loop is True + + +def test_twist_session_replans_when_scene_target_moves() -> None: + generator = _motion_generator() + engine = AtomicActionEngine( generator, - Press(default_options=PressOptions(hand_interp_steps=4)), + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=torch.zeros(HAND_DOF), + grasp=torch.ones(HAND_DOF), + ) + }, + load_builtins=False, + ) + engine.register(Twist()) + semantics = ObjectSemantics( + affordance=TwistAffordance( + grasp_position=(0.0, 0.0, 0.0), + axis_origin=(0.0, 0.0, 0.0), + ), + geometry={}, + label="moving-knob", + ) + invocation = ActionInvocation( + skill_id="twist", + goal=TwistGoal(semantics, SceneEntityPose("target")), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=24), + skill_options=TwistOptions(hand_interp_steps=3), + ) + initial_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + initial_context = _context( + scene=_target_scene(initial_pose, timestamp=0.0, version=0) + ) + session = engine.start((invocation,), initial_context) + session.tick(initial_context) + moved_pose = initial_pose.clone() + moved_pose[:, 1, 3] = 0.3 + + recovered = session.tick( + _context( + scene=_target_scene(moved_pose, timestamp=0.1, version=1), + timestamp=0.1, + ) + ) + + event_kinds = {event.kind for event in recovered.events} + assert ExecutionEventKind.DYNAMIC_GOAL_CHANGED in event_kinds + assert ExecutionEventKind.REPLANNED in event_kinds + + +@pytest.mark.parametrize( + ("direction", "expected_segments", "translation_sign"), + ( + ("pull", ["approach", "reach", "close", "pull", "open"], -1.0), + ( + "push", + ["approach", "reach", "close", "push", "open", "return"], + 1.0, + ), + ), +) +def test_slide_plans_expected_segments( + direction: Literal["pull", "push"], + expected_segments: list[str], + translation_sign: float, + monkeypatch: pytest.MonkeyPatch, +) -> None: + vertices = torch.tensor( + [ + [-0.1, 0.0, 0.0], + [0.1, 0.0, 0.0], + [0.0, 0.0, 0.0], + ] + ) + link_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + affordance = SlideAffordance( + mesh_vertices=vertices, + mesh_triangles=torch.tensor([[0, 1, 2]]), + translation_axis=torch.tensor([0.0, -1.0, 0.0]), + ) + grasp_calls: list[tuple[torch.Tensor, torch.Tensor]] = [] + + def sample_grasp( + self: SlideAffordance, + obj_poses: torch.Tensor, + approach_direction: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + grasp_calls.append((obj_poses, approach_direction)) + return ( + torch.ones(NUM_ENVS, dtype=torch.bool), + torch.eye(4).repeat(NUM_ENVS, 1, 1), + torch.full((NUM_ENVS,), 0.03), + ) + + monkeypatch.setattr( + SlideAffordance, + "get_best_grasp_poses", + sample_grasp, + ) + semantics = ObjectSemantics( + affordance=affordance, + geometry={}, + label="drawer_handle", + ) + generator = _motion_generator() + action = _bind_action(generator, Slide()) + options = SlideOptions( + direction=direction, + hand_interp_steps=3, + approach_distance=0.1, + translation_distance=0.15, ) plan = _plan_action( action, - _invocation("press", PressGoal(torch.eye(4)), sample_count=12), - _context(task), + ActionInvocation( + skill_id="slide", + goal=SlideGoal(semantics, link_pose), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=24), + skill_options=options, + ), + _context(), ) - 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) + assert plan.plan_success.tolist() == [True, True] + assert plan.trajectory.positions.shape == (NUM_ENVS, 24, ROBOT_DOF) + assert [segment.name for segment in plan.segments] == expected_segments + assert torch.all( + plan.trajectory.positions[:, plan.segment("close").stop - 1, ARM_DOF:] == 1.0 + ) + assert torch.all( + plan.trajectory.positions[:, plan.segment("open").stop - 1, ARM_DOF:] == 0.0 + ) + assert len(grasp_calls) == 1 + assert torch.equal(grasp_calls[0][0], link_pose) + assert torch.allclose( + grasp_calls[0][1], + torch.tensor([0.0, -1.0, 0.0]).expand(NUM_ENVS, -1), + ) + planned_targets = [ + call.kwargs["pose"] for call in generator.robot.compute_ik.call_args_list + ] + expected_axis = torch.tensor([0.0, -1.0, 0.0]) + motion_lengths = Slide._motion_segment_lengths( + 24, + options.hand_interp_steps, + direction=direction, + ) + assert torch.allclose( + planned_targets[0][:, :3, 3], + -expected_axis.expand(NUM_ENVS, -1) * options.approach_distance, + ) + reach_stop = 1 + motion_lengths[1] - 1 + assert torch.allclose( + planned_targets[reach_stop - 1][:, :3, 3], + torch.zeros(NUM_ENVS, 3), + ) + translate_stop = reach_stop + motion_lengths[2] - 1 + translated_targets = torch.stack( + [pose[:, :3, 3] for pose in planned_targets[reach_stop:translate_stop]], + dim=1, + ) + assert torch.allclose( + translated_targets[:, -1], + expected_axis.expand(NUM_ENVS, -1) + * (translation_sign * options.translation_distance), + ) + orthogonal = ( + translated_targets + - (translated_targets * expected_axis).sum(dim=-1, keepdim=True) * expected_axis + ) + assert torch.allclose(orthogonal, torch.zeros_like(orthogonal), atol=1.0e-6) + if direction == "push": + assert torch.allclose( + planned_targets[-1][:, :3, 3], + -expected_axis.expand(NUM_ENVS, -1) * options.approach_distance, + ) + + +def test_slide_holds_failed_environment() -> None: + affordance = SlideAffordance( + mesh_vertices=torch.zeros(3, 3), + mesh_triangles=torch.tensor([[0, 1, 2]]), + translation_axis=torch.tensor([0.0, -1.0, 0.0]), + ) + affordance.get_best_grasp_poses = Mock( + return_value=( + torch.tensor([True, False]), + torch.eye(4).repeat(NUM_ENVS, 1, 1), + torch.full((NUM_ENVS,), 0.03), + ) + ) + semantics = ObjectSemantics( + affordance=affordance, + geometry={}, + label="drawer_handle", + ) + generator = _motion_generator() + + def successful_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), torch.ones_like(joint_seed) + + generator.robot.compute_ik.side_effect = successful_ik + action = _bind_action(generator, Slide()) + context = _context() + + plan = _plan_action( + action, + ActionInvocation( + skill_id="slide", + goal=SlideGoal(semantics, torch.eye(4)), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=18), + skill_options=SlideOptions(hand_interp_steps=3), + ), + context, + ) + + 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).expand(18, -1), + ) + + +def test_slide_fk_path_remains_on_translation_axis( + monkeypatch: pytest.MonkeyPatch, +) -> None: + generator = _motion_generator() + + def position_ik( + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + qpos = joint_seed.clone() + qpos[:, :3] = pose[:, :3, 3] + return torch.ones(NUM_ENVS, dtype=torch.bool), qpos + + def position_fk( + qpos: torch.Tensor, + name: str, + to_matrix: bool, + ) -> torch.Tensor: + pose = torch.eye(4).repeat(qpos.shape[0], 1, 1) + pose[:, :3, 3] = qpos[:, :3] + return pose + + generator.robot.compute_ik.side_effect = position_ik + generator.robot.compute_fk.side_effect = position_fk + affordance = SlideAffordance( + mesh_vertices=torch.zeros(3, 3), + mesh_triangles=torch.tensor([[0, 1, 2]]), + translation_axis=torch.tensor([0.0, -1.0, 0.0]), + ) + monkeypatch.setattr( + affordance, + "get_best_grasp_poses", + Mock( + return_value=( + torch.ones(NUM_ENVS, dtype=torch.bool), + torch.eye(4).repeat(NUM_ENVS, 1, 1), + torch.full((NUM_ENVS,), 0.03), + ) + ), + ) + semantics = ObjectSemantics(affordance=affordance, geometry={}, label="handle") + action = _bind_action(generator, Slide()) + plan = _plan_action( + action, + ActionInvocation( + skill_id="slide", + goal=SlideGoal(semantics, torch.eye(4)), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=24), + skill_options=SlideOptions(direction="pull", hand_interp_steps=3), + ), + _context(), + ) + + pull_segment = plan.segment("pull") + arm_path = plan.trajectory.positions[ + :, pull_segment.start : pull_segment.stop, :ARM_DOF + ] + fk_path = position_fk(arm_path.reshape(-1, ARM_DOF), "arm", True).reshape( + NUM_ENVS, -1, 4, 4 + ) + positions = fk_path[:, :, :3, 3] + axis = torch.tensor([0.0, -1.0, 0.0]) + orthogonal = positions - (positions * axis).sum(dim=-1, keepdim=True) * axis + assert torch.allclose(orthogonal, torch.zeros_like(orthogonal), atol=1.0e-6) + + +def test_press_plans_close_approach_press_and_retract() -> None: + affordance = PressAffordance( + press_axis=torch.tensor([1.0, 0.0, 0.0]), + press_position=(0.0, 0.0, 0.0), + ) + semantics = ObjectSemantics( + affordance=affordance, + geometry={}, + label="button", + ) + generator = _motion_generator() + action = _bind_action(generator, Press()) + options = PressOptions( + hand_interp_steps=3, + approach_distance=0.1, + press_distance=0.02, + ) + + plan = _plan_action( + action, + ActionInvocation( + skill_id="press", + goal=PressGoal(semantics, torch.eye(4)), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=24), + skill_options=options, + ), + _context(), + ) + + assert plan.plan_success.tolist() == [True, True] + assert plan.trajectory.positions.shape == (NUM_ENVS, 24, ROBOT_DOF) + assert [segment.name for segment in plan.segments] == [ + "close", + "approach", + "contact", + "press", + "retract", + ] + assert torch.all( + plan.trajectory.positions[:, plan.segment("close").stop - 1, ARM_DOF:] == 1.0 + ) + contact_pose = affordance.get_press_pose(torch.eye(4).repeat(NUM_ENVS, 1, 1)) + expected_approach = ( + contact_pose[:, :3, 3] - contact_pose[:, :3, 2] * options.approach_distance + ) + expected_pressed = ( + contact_pose[:, :3, 3] + contact_pose[:, :3, 2] * options.press_distance + ) + planned_targets = [ + call.kwargs["pose"] for call in generator.robot.compute_ik.call_args_list + ] + motion_lengths = Press._motion_segment_lengths(24, options.hand_interp_steps) + contact_stop = 1 + motion_lengths[1] - 1 + press_stop = contact_stop + motion_lengths[2] - 1 + assert torch.allclose(planned_targets[0][:, :3, 3], expected_approach) + assert torch.allclose( + planned_targets[contact_stop - 1][:, :3, 3], contact_pose[:, :3, 3] + ) + assert torch.allclose(planned_targets[press_stop - 1][:, :3, 3], expected_pressed) + assert torch.allclose(planned_targets[-1][:, :3, 3], expected_approach) + + +def test_press_plans_from_rigid_object_pose_snapshot_with_option_position() -> None: + affordance = PressAffordance( + press_axis=torch.tensor([1.0, 0.0, 0.0]), + press_position=(0.5, 0.5, 0.5), + ) + semantics = ObjectSemantics( + affordance=affordance, + geometry={}, + label="rigid-button", + ) + generator = _motion_generator() + + plan = _plan_action( + _bind_action(generator, Press()), + ActionInvocation( + skill_id="press", + goal=PressGoal(semantics, torch.eye(4)), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=24), + skill_options=PressOptions( + hand_interp_steps=3, + press_position=(0.1, 0.2, 0.3), + ), + ), + _context(), + ) + + assert plan.plan_success.tolist() == [True, True] + planned_approach = generator.robot.compute_ik.call_args_list[0].kwargs["pose"] + assert torch.allclose( + planned_approach[:, :3, 3], + torch.tensor([0.0, 0.2, 0.3]).expand(NUM_ENVS, -1), + ) + + +def test_press_fk_path_passes_contact_and_remains_on_press_axis() -> None: + generator = _motion_generator() + + def position_ik( + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + qpos = joint_seed.clone() + qpos[:, :3] = pose[:, :3, 3] + return torch.ones(NUM_ENVS, dtype=torch.bool), qpos + + def position_fk( + qpos: torch.Tensor, + name: str, + to_matrix: bool, + ) -> torch.Tensor: + pose = torch.eye(4).repeat(qpos.shape[0], 1, 1) + pose[:, :3, 3] = qpos[:, :3] + return pose + + generator.robot.compute_ik.side_effect = position_ik + generator.robot.compute_fk.side_effect = position_fk + semantics = ObjectSemantics( + affordance=PressAffordance( + press_axis=torch.tensor([1.0, 0.0, 0.0]), + press_position=(0.0, 0.0, 0.0), + ), + geometry={}, + label="button", + ) + action = _bind_action(generator, Press()) + plan = _plan_action( + action, + ActionInvocation( + skill_id="press", + goal=PressGoal(semantics, torch.eye(4)), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=24), + skill_options=PressOptions(hand_interp_steps=3, press_distance=0.04), + ), + _context(), + ) + + contact_arm = plan.trajectory.positions[ + :, plan.segment("contact").stop - 1, :ARM_DOF + ] + contact_fk = position_fk(contact_arm, "arm", True) + assert torch.allclose(contact_fk[:, :3, 3], torch.zeros(NUM_ENVS, 3)) + press_segment = plan.segment("press") + press_arm = plan.trajectory.positions[ + :, press_segment.start : press_segment.stop, :ARM_DOF + ] + press_fk = position_fk(press_arm.reshape(-1, ARM_DOF), "arm", True).reshape( + NUM_ENVS, -1, 4, 4 + ) + positions = press_fk[:, :, :3, 3] + axis = torch.tensor([1.0, 0.0, 0.0]) + orthogonal = positions - (positions * axis).sum(dim=-1, keepdim=True) * axis + assert torch.allclose(orthogonal, torch.zeros_like(orthogonal), atol=1.0e-6) + assert torch.allclose(positions[:, -1], torch.tensor([0.04, 0.0, 0.0])) + + +def test_press_preserves_failed_environment_at_observed_qpos() -> None: + semantics = ObjectSemantics( + affordance=PressAffordance( + press_axis=torch.tensor([1.0, 0.0, 0.0]), + press_position=(0.0, 0.0, 0.0), + ), + geometry={}, + label="button", + ) + generator = _motion_generator() + + def partial_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.tensor([True, False]), torch.ones_like(joint_seed) + + generator.robot.compute_ik.side_effect = partial_ik + action = _bind_action(generator, Press()) + context = _context() + + plan = _plan_action( + action, + ActionInvocation( + skill_id="press", + goal=PressGoal(semantics, torch.eye(4)), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=18), + skill_options=PressOptions(hand_interp_steps=3), + ), + context, + ) + + 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).expand(18, -1), + ) + + +def test_press_rejects_non_press_affordance() -> None: + semantics = ObjectSemantics( + affordance=AntipodalAffordance( + mesh_vertices=torch.zeros(8, 3), + mesh_triangles=torch.zeros(4, 3, dtype=torch.long), + ), + geometry={}, + label="mesh-button", + ) + action = _bind_action(_motion_generator(), Press()) + + with pytest.raises(ValueError, match="PressAffordance"): + _plan_action( + action, + _invocation("press", PressGoal(semantics, torch.eye(4))), + _context(), + ) + + +def test_press_requires_primary_arm_and_end_effector_bindings() -> None: + semantics = ObjectSemantics( + affordance=AntipodalAffordance(), + geometry={}, + label="button", + ) + action = _bind_action(_motion_generator(), Press()) + invocation = ActionInvocation( + skill_id="press", + goal=PressGoal(semantics, torch.eye(4)), + binding=ActionBinding(manipulators={"primary": "arm"}), + ) + + with pytest.raises(KeyError, match="No end effector is bound to role 'primary'"): + action.resolve_request(invocation) + + +def test_press_axis_belongs_to_affordance_not_action_options() -> None: + assert "press_axis" not in PressOptions.__dataclass_fields__ + + +@pytest.mark.parametrize( + "press_position", + ((0.0, 1.0), (0.0, 1.0, float("nan"))), +) +def test_press_options_reject_invalid_press_position( + press_position: tuple[float, ...], +) -> None: + with pytest.raises(ValueError, match="press_position"): + PressOptions(press_position=press_position) # type: ignore[arg-type] + + +def test_twist_rejects_non_twist_affordance() -> None: + semantics = ObjectSemantics( + affordance=AntipodalAffordance( + mesh_vertices=torch.zeros(8, 3), + mesh_triangles=torch.zeros(4, 3, dtype=torch.long), + ), + geometry={}, + label="mesh-knob", + ) + action = _bind_action(_motion_generator(), Twist()) + + with pytest.raises(ValueError, match="TwistAffordance"): + _plan_action( + action, + _invocation("twist", TwistGoal(semantics, torch.eye(4))), + _context(), + ) + + +def test_twist_axis_belongs_to_affordance_not_action_options() -> None: + assert "twist_axis" not in TwistOptions.__dataclass_fields__ + assert "approach_direction" not in TwistOptions.__dataclass_fields__ + + +def test_twist_options_reject_non_finite_pre_grasp_distance() -> None: + with pytest.raises(ValueError, match="pre_grasp_distance must be finite"): + TwistOptions(pre_grasp_distance=float("nan")) + + +def test_slide_rejects_non_slide_affordance() -> None: + semantics = ObjectSemantics( + affordance=AntipodalAffordance( + mesh_vertices=torch.zeros(8, 3), + mesh_triangles=torch.zeros(4, 3, dtype=torch.long), + ), + geometry={}, + label="mesh-handle", + ) + action = _bind_action(_motion_generator(), Slide()) + + with pytest.raises(ValueError, match="SlideAffordance"): + _plan_action( + action, + _invocation( + "slide", + SlideGoal(semantics, torch.eye(4)), + ), + _context(), + ) + + +def test_slide_requires_primary_end_effector() -> None: + semantics = ObjectSemantics( + affordance=AntipodalAffordance(), + geometry={}, + label="drawer_handle", + ) + action = _bind_action(_motion_generator(), Slide()) + invocation = ActionInvocation( + skill_id="slide", + goal=SlideGoal(semantics, torch.eye(4)), + binding=ActionBinding(manipulators={"primary": "arm"}), + ) + + with pytest.raises(KeyError, match="No end effector is bound to role 'primary'"): + action.resolve_request(invocation) + + +def test_slide_axis_belongs_to_affordance_not_action_options() -> None: + assert "translation_axis" not in SlideOptions.__dataclass_fields__ + + +def test_slide_options_reject_invalid_direction() -> None: + with pytest.raises(ValueError, match="direction"): + SlideOptions(direction="open") # type: ignore[arg-type] def test_handover_does_not_mutate_cached_final_pose( diff --git a/tests/sim/atomic_actions/test_affordance.py b/tests/sim/atomic_actions/test_affordance.py index 097a1ffe8..0e123eefe 100644 --- a/tests/sim/atomic_actions/test_affordance.py +++ b/tests/sim/atomic_actions/test_affordance.py @@ -28,6 +28,9 @@ AntipodalAffordance, AssembleAffordance, InteractionPoints, + PressAffordance, + SlideAffordance, + TwistAffordance, ) @@ -118,6 +121,186 @@ def test_best_grasp_poses_casts_approach_direction_to_generator_device(self): assert approach_direction.device == generator.device +class TestTwistAffordance: + def test_requires_explicit_grasp_position_and_axis_origin(self): + with pytest.raises(TypeError, match="grasp_position"): + TwistAffordance() # type: ignore[call-arg] + + @pytest.mark.parametrize( + "twist_axis", + ( + torch.tensor([1.0, 0.0, 1.0]), + torch.tensor([0.0, 0.0, 1.0]), + torch.tensor([0.0, 0.0, -1.0]), + ), + ) + def test_builds_right_handed_orthonormal_grasp_frame(self, twist_axis): + link_pose = torch.eye(4).repeat(2, 1, 1) + link_pose[:, :3, 3] = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + affordance = TwistAffordance( + grasp_position=(0.25, -0.5, 0.75), + axis_origin=(0.1, 0.2, 0.3), + twist_axis=twist_axis, + ) + + grasp_pose = affordance.get_grasp_pose(link_pose) + rotation = grasp_pose[:, :3, :3] + + assert torch.allclose( + grasp_pose[:, :3, 3], + link_pose[:, :3, 3] + torch.tensor([0.25, -0.5, 0.75]).expand(2, -1), + ) + assert torch.allclose( + torch.matmul(rotation.transpose(1, 2), rotation), + torch.eye(3).expand(2, -1, -1), + atol=1.0e-6, + ) + assert torch.allclose(torch.linalg.det(rotation), torch.ones(2), atol=1.0e-6) + + +class TestSlideAffordance: + def test_uses_local_antipodal_mesh_with_batched_directions(self): + vertices = torch.tensor( + [ + [-1.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 0.0, 0.0], + ] + ) + triangles = torch.tensor([[0, 1, 2]]) + link_pose = torch.eye(4).repeat(2, 1, 1) + link_pose[:, :3, 3] = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + affordance = SlideAffordance( + mesh_vertices=vertices, + mesh_triangles=triangles, + translation_axis=torch.tensor([0.0, -1.0, 0.0]), + ) + generator = Mock() + generator.device = torch.device("cpu") + first_grasp = torch.eye(4) + first_grasp[:3, 3] = torch.tensor([1.0, 2.0, 3.0]) + second_grasp = torch.eye(4) + second_grasp[:3, 3] = torch.tensor([4.0, 5.0, 6.0]) + generator.get_grasp_poses.side_effect = ( + (True, first_grasp, 0.03), + (True, second_grasp, 0.04), + ) + affordance._generator = generator + approach_directions = torch.tensor([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0]]) + + success, grasp_poses, open_lengths = affordance.get_best_grasp_poses( + link_pose, + approach_direction=approach_directions, + ) + + assert isinstance(affordance, AntipodalAffordance) + assert success.tolist() == [True, True] + assert torch.allclose(grasp_poses, torch.stack([first_grasp, second_grasp])) + assert torch.allclose(open_lengths, torch.tensor([0.03, 0.04])) + assert torch.equal( + generator.get_grasp_poses.call_args_list[0].args[1], + approach_directions[0], + ) + assert torch.equal( + generator.get_grasp_poses.call_args_list[1].args[1], + approach_directions[1], + ) + + def test_requires_local_antipodal_geometry(self): + with pytest.raises(TypeError, match="mesh_vertices"): + SlideAffordance() + + @pytest.mark.parametrize( + "translation_axis", + ( + torch.zeros(3), + torch.tensor([float("nan"), 0.0, 0.0]), + torch.zeros(2), + ), + ) + def test_rejects_invalid_translation_axis(self, translation_axis): + with pytest.raises(ValueError, match="translation_axis"): + SlideAffordance( + mesh_vertices=torch.ones(3, 3), + mesh_triangles=torch.tensor([[0, 1, 2]]), + translation_axis=translation_axis, + ) + + +class TestPressAffordance: + def test_requires_explicit_surface_press_position(self): + with pytest.raises(TypeError, match="press_position"): + PressAffordance() # type: ignore[call-arg] + + @pytest.mark.parametrize( + "press_axis", + ( + torch.tensor([1.0, 0.0, 1.0]), + torch.tensor([0.0, 0.0, 1.0]), + torch.tensor([0.0, 0.0, -1.0]), + ), + ) + def test_builds_right_handed_orthonormal_press_frame(self, press_axis): + link_pose = torch.eye(4).repeat(2, 1, 1) + link_pose[:, :3, 3] = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + affordance = PressAffordance( + press_axis=press_axis, + press_position=(0.25, -0.5, 0.75), + ) + + press_pose = affordance.get_press_pose(link_pose) + rotation = press_pose[:, :3, :3] + + assert torch.allclose( + press_pose[:, :3, 3], + link_pose[:, :3, 3] + torch.tensor([0.25, -0.5, 0.75]).expand(2, -1), + ) + assert torch.allclose( + torch.matmul(rotation.transpose(1, 2), rotation), + torch.eye(3).expand(2, -1, -1), + atol=1.0e-6, + ) + assert torch.allclose(torch.linalg.det(rotation), torch.ones(2), atol=1.0e-6) + + def test_rejects_zero_press_axis(self): + with pytest.raises(ValueError, match="press_axis must be non-zero"): + PressAffordance( + press_axis=torch.zeros(3), + press_position=(0.0, 0.0, 0.0), + ) + + def test_uses_configured_press_position(self): + object_pose = torch.eye(4).repeat(2, 1, 1) + object_pose[:, :3, 3] = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + affordance = PressAffordance( + press_axis=torch.tensor([1.0, 0.0, 0.0]), + press_position=(0.25, -0.5, 0.75), + ) + + press_pose = affordance.get_press_pose(object_pose) + + assert torch.allclose( + press_pose[:, :3, 3], + object_pose[:, :3, 3] + torch.tensor([0.25, -0.5, 0.75]).expand(2, -1), + ) + + def test_per_call_press_position_overrides_affordance_position(self): + affordance = PressAffordance( + press_axis=torch.tensor([1.0, 0.0, 0.0]), + press_position=(1.0, 1.0, 1.0), + ) + + press_pose = affordance.get_press_pose( + torch.eye(4).unsqueeze(0), + press_position=(0.1, 0.2, 0.3), + ) + + assert torch.allclose( + press_pose[0, :3, 3], + torch.tensor([0.1, 0.2, 0.3]), + ) + + class TestInteractionPoints: def test_default_points_shape(self): assert InteractionPoints().points.shape == (1, 3) diff --git a/tests/sim/atomic_actions/test_engine.py b/tests/sim/atomic_actions/test_engine.py index 8704c357a..b2b286630 100644 --- a/tests/sim/atomic_actions/test_engine.py +++ b/tests/sim/atomic_actions/test_engine.py @@ -39,8 +39,6 @@ JointPositionGoal, MotionPolicy, PlanningContext, - PressGoal, - PressOptions, ResolvedActionRequest, ) @@ -165,26 +163,6 @@ 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.actions["press"].resolve_request(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()) diff --git a/tests/sim/atomic_actions/test_module_imports.py b/tests/sim/atomic_actions/test_module_imports.py new file mode 100644 index 000000000..23db8966c --- /dev/null +++ b/tests/sim/atomic_actions/test_module_imports.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. +# ---------------------------------------------------------------------------- + +"""Regression tests for atomic-action module imports and file entry points.""" + +from __future__ import annotations + +from pathlib import Path +import subprocess +import sys + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +PRIMITIVES_DIRECTORY = ( + REPOSITORY_ROOT / "embodichain" / "lab" / "sim" / "atomic_actions" / "primitives" +) +TUTORIAL_DIRECTORY = REPOSITORY_ROOT / "scripts" / "tutorials" / "atomic_action" + +PUBLIC_PRIMITIVE_SCRIPTS = tuple( + path + for path in sorted(PRIMITIVES_DIRECTORY.glob("*.py")) + if not path.name.startswith("_") +) +TUTORIAL_SCRIPTS = tuple(sorted(TUTORIAL_DIRECTORY.glob("*.py"))) + +RUN_PUBLIC_PRIMITIVES_CODE = """ +import runpy +import sys + +for module_path in sys.argv[1:]: + runpy.run_path(module_path, run_name="__main__") +""" + +IMPORT_TUTORIALS_CODE = """ +import runpy +import sys + +for module_path in sys.argv[1:]: + runpy.run_path(module_path, run_name="__atomic_action_import_check__") +""" + + +def _run_import_check( + code: str, paths: tuple[Path, ...] +) -> subprocess.CompletedProcess[str]: + """Run multiple module files in one isolated Python process.""" + return subprocess.run( + [sys.executable, "-c", code, *(str(path) for path in paths)], + cwd=REPOSITORY_ROOT, + capture_output=True, + text=True, + check=False, + ) + + +def test_all_public_primitive_modules_can_run_as_files() -> None: + """Public primitive files should resolve package imports when run directly.""" + result = _run_import_check(RUN_PUBLIC_PRIMITIVES_CODE, PUBLIC_PRIMITIVE_SCRIPTS) + + assert result.returncode == 0, result.stderr + + +def test_all_atomic_action_tutorials_import_without_running_main() -> None: + """Tutorial modules should import without starting their simulations.""" + result = _run_import_check(IMPORT_TUTORIALS_CODE, TUTORIAL_SCRIPTS) + + assert result.returncode == 0, result.stderr diff --git a/tests/sim/atomic_actions/test_trajectory_ops.py b/tests/sim/atomic_actions/test_trajectory_ops.py index a23384834..24362b8b9 100644 --- a/tests/sim/atomic_actions/test_trajectory_ops.py +++ b/tests/sim/atomic_actions/test_trajectory_ops.py @@ -28,6 +28,7 @@ resolve_object_target, ) from embodichain.lab.sim.atomic_actions.trajectory_ops import ( + axis_translation_keyframes, build_joint_plan_states, build_pose_plan_states, interpolate_hand_qpos, @@ -322,6 +323,45 @@ def test_incompatible_offset_batch_raises(self): translate_pose_world(pose, offset) +class TestAxisTranslationKeyframes: + def test_excludes_start_includes_end_and_stays_on_axis(self): + start = torch.eye(4).repeat(2, 1, 1) + start[:, :3, 3] = torch.tensor([[-0.1, 0.2, 0.3], [0.4, -0.2, 0.1]]) + axis = torch.tensor([[1.0, 0.0, 1.0], [0.0, -1.0, 0.0]]) + axis = torch.nn.functional.normalize(axis, dim=1) + end = start.clone() + end[:, :3, 3] += axis * torch.tensor([[0.5], [-0.3]]) + + keyframes = axis_translation_keyframes( + start, + end, + axis, + n_waypoints=5, + ) + + displacement = keyframes[:, :, :3, 3] - start[:, None, :3, 3] + orthogonal = ( + displacement + - (displacement * axis[:, None]).sum(dim=-1, keepdim=True) * axis[:, None] + ) + assert keyframes.shape == (2, 5, 4, 4) + assert torch.allclose(keyframes[:, -1], end) + assert torch.allclose(orthogonal, torch.zeros_like(orthogonal), atol=1.0e-6) + + def test_rejects_off_axis_displacement(self): + start = torch.eye(4).unsqueeze(0) + end = start.clone() + end[:, 1, 3] = 0.1 + + with pytest.raises(ValueError, match="parallel to axis"): + axis_translation_keyframes( + start, + end, + torch.tensor([1.0, 0.0, 0.0]), + n_waypoints=2, + ) + + def test_interpolate_hand_qpos_preserves_endpoints(): start = torch.tensor([[0.0, 0.0]]) end = torch.tensor([[1.0, 1.0]]) diff --git a/tests/sim/atomic_actions/test_tutorial_utils.py b/tests/sim/atomic_actions/test_tutorial_utils.py index def9e17a3..1c837184a 100644 --- a/tests/sim/atomic_actions/test_tutorial_utils.py +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -67,7 +67,7 @@ [ [1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.15], + [0.0, 0.0, 1.0, 0.17], [0.0, 0.0, 0.0, 1.0], ] ) diff --git a/tests/sim/planners/test_motion_generator_batched.py b/tests/sim/planners/test_motion_generator_batched.py index 32d18fd9e..03ab3e660 100644 --- a/tests/sim/planners/test_motion_generator_batched.py +++ b/tests/sim/planners/test_motion_generator_batched.py @@ -485,6 +485,51 @@ def test_ik_interp_solves_batched_poses_without_calling_backend(self): ) generator.planner.plan.assert_not_called() + def test_linear_cartesian_motion_grounds_every_output_sample_with_ik(self): + generator = _mock_generator() + + def encode_position( + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + qpos = joint_seed.clone() + qpos[:, :3] = pose[:, :3, 3] + return torch.ones(BATCH_SIZE, dtype=torch.bool), qpos + + generator.robot.compute_ik.side_effect = encode_position + weights = torch.linspace(1.0 / (SAMPLE_COUNT - 1), 1.0, SAMPLE_COUNT - 1) + targets = [] + for weight in weights: + pose = torch.eye(4).repeat(BATCH_SIZE, 1, 1) + pose[:, 0, 3] = weight + targets.append(PlanState.from_xpos(pose)) + + result = generator.generate( + targets, + MotionGenOptions( + strategy="motion_gen", + sample_count=SAMPLE_COUNT, + start_qpos=torch.zeros(BATCH_SIZE, CONTROLLED_DOF), + control_part="arm", + is_linear=True, + preserve_cartesian_samples=True, + ), + ) + + assert result.positions is not None + assert result.positions.shape == ( + BATCH_SIZE, + SAMPLE_COUNT, + CONTROLLED_DOF, + ) + expected_x = torch.linspace(0.0, 1.0, SAMPLE_COUNT) + assert torch.allclose( + result.positions[:, :, 0], expected_x.expand(BATCH_SIZE, -1) + ) + assert generator.robot.compute_ik.call_count == SAMPLE_COUNT - 1 + generator.planner.plan.assert_not_called() + def test_motion_gen_delegates_and_resamples_backend_result(self): raw_sample_count = 5 generator = _mock_generator(