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 5a8c12268..0ae43fd87 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 @@ -23,8 +23,12 @@ full-robot timed trajectory and uncommitted expected effects. MoveJointsOptions PickUp PickUpOptions + AxisAlign + AxisAlignOptions MoveHeldObject MoveHeldObjectOptions + Pour + PourOptions Place PlaceOptions Press @@ -38,6 +42,7 @@ full-robot timed trajectory and uncommitted expected effects. CoordinatedPlacement CoordinatedPlacementOptions HandOver + HandOverGoal HandOverOptions .. rubric:: Built-in Goal Contracts @@ -47,7 +52,9 @@ full-robot timed trajectory and uncommitted expected effects. EndEffectorPoseGoal JointPositionGoal GraspGoal + AxisAlignGoal HeldObjectPoseGoal + PourGoal PlaceGoal AssembleGoal PressGoal @@ -55,6 +62,7 @@ full-robot timed trajectory and uncommitted expected effects. TwistGoal CoordinatedPickGoal CoordinatedPlacementGoal + HandOverGoal .. currentmodule:: embodichain.lab.sim.atomic_actions.primitives @@ -82,6 +90,14 @@ PickUp :show-inheritance: :exclude-members: __init__, copy, replace, to_dict +AxisAlign +--------- + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives.axis_align + :members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict + MoveHeldObject -------------- @@ -90,6 +106,14 @@ MoveHeldObject :show-inheritance: :exclude-members: __init__, copy, replace, to_dict +Pour +---- + +.. automodule:: embodichain.lab.sim.atomic_actions.primitives.pour + :members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict + Place ----- 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 ae595999d..0063609f0 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 @@ -83,7 +83,13 @@ embodichain.lab.sim.atomic_actions EndEffectorPoseGoal JointPositionGoal GraspGoal + HandOverGoal + AxisAlignGoal + AxisAlignOptions + AxisAlignAffordance HeldObjectPoseGoal + PourGoal + PourOptions PlaceGoal AssembleGoal PressGoal @@ -100,7 +106,9 @@ embodichain.lab.sim.atomic_actions MoveEndEffector MoveJoints PickUp + AxisAlign MoveHeldObject + Pour Place Press Slide diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index 2a26f15c7..067ea2ae5 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -5,7 +5,7 @@ ```{currentmodule} embodichain.lab.sim.atomic_actions ``` -EmbodiChain ships eleven built-in action implementations with stable skill IDs; +EmbodiChain ships thirteen 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. @@ -69,6 +69,13 @@ The animations below are the focused simulator demos under PickUp demo ::: +:::{grid-item-card} `AxisAlign` +:link: builtin-axis-align +:link-type: ref + +`axis_align` · grasp, lift, align an object-local axis, and release +::: + :::{grid-item-card} `MoveHeldObject` :link: builtin-move-held-object :link-type: ref @@ -78,6 +85,13 @@ The animations below are the focused simulator demos under MoveHeldObject demo ::: +:::{grid-item-card} `Pour` +:link: builtin-pour +:link-type: ref + +`pour` · rotate an already-held object about its local internal axis +::: + :::{grid-item-card} `Place` :link: builtin-place :link-type: ref @@ -146,7 +160,7 @@ The animations below are the focused simulator demos under :link: builtin-hand-over :link-type: ref -`hand_over` · transfer an attachment between arms +`hand_over` · pick, transfer, place, and release with two arms HandOver demo ::: @@ -160,14 +174,16 @@ The animations below are the focused simulator demos under | `move_end_effector` | `EndEffectorPoseGoal` | `primary.motion` | none | none | none | | `move_joints` | `JointPositionGoal` | `primary.motion` | named target only: command matching `target` on `primary.motion` | none | none | | `pick_up` | `GraspGoal` | `primary.motion`, `primary.grasp` | `primary.grasp`: `open`, `grasp` | semantic object/entity | attach object to the `primary.motion` target | +| `axis_align` | `AxisAlignGoal` | `primary.motion`, `primary.grasp` | `primary.grasp`: `open`, `grasp` | unheld object with `AxisAlignAffordance` | open-loop pick, align, lower, and release | | `move_held_object` | `HeldObjectPoseGoal` | `primary.motion`, `primary.grasp` | `primary.grasp`: `grasp` | object held exclusively by the `primary.motion` target | preserve attachment | +| `pour` | `PourGoal` | `primary.motion`, `primary.grasp` | `primary.grasp`: `grasp` | exclusively held object with `AxisAlignAffordance` | preserve attachment; open-loop rotate and return | | `place` | `PlaceGoal`, `AssembleGoal` | `primary.motion`, `primary.grasp` | `primary.grasp`: `open`, `grasp` | any active attachment must be exclusive to `primary.motion`; `AssembleGoal` requires one | detach object | | `press` | `PressGoal` | `primary.motion`, `primary.grasp` | `primary.grasp`: `grasp` | `PressAffordance` + target pose | open-loop motion; application verifies contact/actuation | | `slide` | `SlideGoal` | `primary.motion`, `primary.grasp` | `primary.grasp`: `open`, `grasp` | `SlideAffordance` + link pose | open-loop motion; application verifies joint travel/grasp | | `twist` | `TwistGoal` | `primary.motion`, `primary.grasp` | `primary.grasp`: `open`, `grasp` | `TwistAffordance` + target pose | open-loop motion; application verifies joint travel/grasp | | `coordinated_pickment` | `CoordinatedPickGoal` | `left.motion`, `left.grasp`, `right.motion`, `right.grasp` | both grasp endpoints: `open`, `grasp` | semantic object/entity | attach the shared object to both motion targets | | `coordinated_placement` | `CoordinatedPlacementGoal` | `placing.motion`, `placing.grasp`, `support.motion`, `support.grasp` | `placing.grasp`: `open`, `grasp`; `support.grasp`: `grasp` | two distinct objects, each held exclusively by its motion target | optionally detach placing object; preserve support attachment | -| `hand_over` | `GraspGoal` | `source.motion`, `source.grasp`, `destination.motion`, `destination.grasp` | both grasp endpoints: `open`, `grasp` | object held exclusively by the source motion target | transfer attachment to the destination motion target | +| `hand_over` | `HandOverGoal` | `source.motion`, `source.grasp`, `destination.motion`, `destination.grasp` | both grasp endpoints: `open`, `grasp` | both candidate motion targets unoccupied; unheld object with `AntipodalAffordance` | open-loop pick, transfer, place, and release | ### Participant slot meanings @@ -180,8 +196,7 @@ commands to a robot resource, then adapters produce the generic | Slot | Used by | Meaning | |---|---|---| | `primary` | Single-participant skills | Principal participant for this invocation; it has no inherent left/right or default-robot meaning | -| `source` | `hand_over` | Participant that initially holds and transfers the object | -| `destination` | `hand_over` | Participant that receives the object | +| `source`, `destination` | `hand_over` | Two candidate participants; the action assigns the nearer one to pickup and the other one to receive | | `left`, `right` | `coordinated_pickment` | Participants on whose sides the affordance samples left/right grasps | | `placing` | `coordinated_placement` | Participant that aligns and optionally releases the placing object | | `support` | `coordinated_placement` | Participant that keeps holding and positioning the support object | @@ -227,10 +242,13 @@ entity as a recovery dependency. | `CoordinatedPlacementGoal` placing/support poses | yes | yes | | `PickUp.grasp_xpos` | yes | yes | | `PickUp` `ObjectSemantics.entity_id` grounding | implicit snapshot reference | yes; always consumed for the object pose | +| `AxisAlign.grasp_xpos` | yes | yes | +| `AxisAlign` `ObjectSemantics.entity_id` grounding | implicit snapshot reference | yes; always consumed for the object pose | +| `HandOverGoal.target_pose` | yes | yes | +| `HandOver` `ObjectSemantics.entity_id` grounding | implicit snapshot reference | yes; always consumed for the initial object pose | | Coordinated pickup implicit initial pose via `ObjectSemantics.entity_id` | implicit snapshot reference | yes; only when `object_initial_pose` is omitted | | `AssembleGoal.base_pose` | yes | yes | | Deprecated `ObjectSemantics.entity` / `AssembleAffordance.base_object_entity` fallback | no | no | -| `HandOver` current held-object pose | no scene lookup | no; derived from observed EEF pose and verified attachment state | ### Object identity and grounding @@ -245,7 +263,7 @@ create a scene-motion dependency. `collect_scene_dependencies()` intentionally does not recurse into `ObjectSemantics`; each primitive declares a semantic ID only when its planner actually consumes that object's snapshot pose. -Attachment and handover identity are not based on `label`. The core resolves an +Attachment identity is not based on `label`. The core resolves an explicit `entity_id` only against another explicit ID. If either compared side has one, both sides must have the same explicit value; an equal legacy `entity.uid` does not match it. When both explicit IDs are absent, two non-empty @@ -406,6 +424,49 @@ deprecated entity-only fallback. For canonical snapshot grounding and moving target recovery, see `scripts/tutorials/atomic_action/moving_target_recovery.py`. +(builtin-axis-align)= + +## `AxisAlign` + +Executes **approach -> reach -> close -> lift -> align -> lower -> open** while +grouping arm motion into two planner calls: the open-gripper `approach` phase +contains the pre-grasp and grasp targets, and the closed-gripper `manipulate` +phase contains lift, alignment, and lowering targets. The `close` and `open` +segments are local hand interpolation and do not call the motion generator. +Only the final aligned pose is sent to the planner; the alignment sample budget +controls trajectory resolution without expanding the rotation into one CuRobo +`plan_pose` call per intermediate orientation. +The object's `AxisAlignAffordance.internal_axis` is expressed in the +object-local frame, while `AxisAlignOptions.target_axis` is expressed in the +world frame. The final alignment target applies the shortest rotation about +the lifted object's origin so that +`aligned_object_rotation @ internal_axis == normalized_target_axis`, then +derives every end-effector keyframe through the fixed grasp transform. + +| Contract | Value | +|---|---| +| Skill ID | `axis_align` | +| Goal | `AxisAlignGoal(semantics=..., grasp_xpos=None)` | +| Binding | manipulator + end effector role `primary` | +| Precondition | an `AxisAlignAffordance`; the object pose resolves from `ObjectSemantics.entity_id` or the deprecated live entity fallback | +| Motion | approach, grasp, lift, rotate in place, lower, release | +| Effect | explicitly open-loop; no final object-pose success is claimed | + +An explicit `grasp_xpos` accepts the same pose forms as `GraspGoal`; omitting it +prefers valid antipodal grasps whose TCP y-axis is perpendicular to the object +rotation axis, using grasp cost as the tie-breaker. When a currently horizontal +object axis is aligned to world-up, the initial grasp orientation is pre-rotated +45 degrees opposite the alignment rotation. This reduces the arm's table-side +sweep during upright manipulation. `AxisAlignOptions` extends `PickUpOptions` +with `target_axis` and `lower_distance`. Shared and per-environment target axes +use shapes `(3,)` and `(B, 3)` respectively. Zero or non-finite axes are +rejected, and exactly opposite axes use a deterministic 180-degree rotation +rather than an unstable cross-product direction. + +**Example:** `scripts/tutorials/atomic_action/axis_align.py` provides +`--alignment upright` (align object-local X to world Z) and +`--alignment horizontal_align` (align object-local X to world Y). + (builtin-move-held-object)= ## `MoveHeldObject` @@ -437,6 +498,38 @@ marked unsuccessful and held in place. **Example:** `scripts/tutorials/atomic_action/move_held_object.py` +(builtin-pour)= + +## `Pour` + +Rotates the object already held by the bound `primary` manipulator. `PourGoal` +contains no object pose because the action consumes the verified +`HeldObjectState` created by `PickUp`. `PourOptions` contains only the signed +`rotate_angle` in radians. + +The held object's semantics must use `AxisAlignAffordance`. The action obtains +the current EEF pose from FK at the observed starting arm qpos, reconstructs +the current object pose using `eef_pose @ inverse(object_to_eef)`, transforms +the affordance's object-local `internal_axis` into world space, and applies the +requested rotation while keeping the object origin fixed. The resulting EEF +target is `target_object_pose @ object_to_eef`. A second target then returns to +the EEF pose observed by FK, thereby reversing the rotation by the same angle. +Both targets are submitted in one motion-generation call so a collision-aware +planner can chain the outbound and return legs. + +| Contract | Value | +|---|---| +| Skill ID | `pour` | +| Goal | `PourGoal()` | +| Options | `PourOptions(rotate_angle=...)` | +| Binding | manipulator + end effector role `primary` | +| Precondition | an exclusive `HeldObjectState` whose semantics use `AxisAlignAffordance` | +| Motion | rotate by `rotate_angle`, then reverse by the same angle, with the hand held at `grasp` | +| Effect | none; the existing attachment is preserved | + +**Example:** `scripts/tutorials/atomic_action/pour.py` compiles a horizontal +`PickUp` followed by `Pour`. + (builtin-place)= ## `Place` @@ -633,9 +726,10 @@ options to partition the object into left/right grasp regions and select the lowest-cost grasp on each side. Each derived `object_to_eef` transform is stored in the corresponding projected `HeldObjectState`. Later object-centric skills can inspect those per-manipulator entries directly; sharing the same -`ObjectSemantics` instance identifies the common object. Single-arm transport, -release, and handover skills reject those shared rows rather than moving or -detaching just one participant. +`ObjectSemantics` instance identifies the common object. Single-arm transport +and release skills reject those shared rows rather than moving or detaching +just one participant. The unified `HandOver` action starts before pickup and +therefore requires both candidate arms to be unoccupied. The object target and optional initial pose may use `SceneEntityPose`. Those references declare their own scene dependencies. When `object_initial_pose` is @@ -699,35 +793,56 @@ pickment applies to dual-arm `strategy="motion_gen"` planning. ## `HandOver` -Transfers an already held object from one arm to another: **move source to the -handover pose -> destination approaches and grasps -> source releases and -retreats -> destination delivers**. +Runs the full two-arm manipulation as one action: **choose the nearer arm -> +pick the object's top half -> lift and move to the computed middle point -> the +other arm grasps the bottom half -> transfer the grasp -> place and release**. | Contract | Value | |---|---| | Skill ID | `hand_over` | -| Goal | `GraspGoal(semantics=...)` | +| Goal | `HandOverGoal(semantics=..., target_pose=...)` | | Binding contract | disjoint `source` and `destination` slots, each with disjoint `motion` and `grasp` endpoints | -| Precondition | source motion target exclusively has a verified `HeldObjectState`; goal semantics identify that object and support destination grasp selection | -| Effect | remove source attachment and create destination `HeldObjectState` | -| Verification | attachment transfer must be externally verified | - -Both source and destination grasp endpoints must provide `open` and `grasp`. -`HandOverOptions` owns the destination grasp region and approach -direction, middle/final object poses, and segment distances/counts. The -source/destination motion and grasp endpoints come exclusively from the -corresponding generic `ActionBinding` slots. The destination attachment reuses -the source relation's canonical `ObjectSemantics` instance. - -The middle and final poses are currently option tensors rather than -`SceneEntityPose` goal fields. Consequently, handover supports tracking-error -and timeout recovery, but does not automatically invalidate a moving handover -point. An application can submit a newer invocation revision with updated -`HandOverOptions`. The action verifies that the goal and source attachment have -the same stable object identity, then derives the current object orientation -from the observed source EEF pose and verified `object_to_eef` relation. -The reused `GraspGoal.grasp_xpos` field is not consumed by `HandOver` and does -not create a scene dependency. +| Precondition | both candidate arms start unoccupied; object semantics use `AntipodalAffordance` | +| Effect | none; both grippers are open after placing the object | +| Verification | open-loop physical pickup, transfer, placement, and release | + +Both grasp endpoints must provide `open` and `grasp`. The two binding slots +are candidate motion/grasp pairs rather than a caller-selected transfer +direction. For every environment, the action compares the observed object +position with both configured solver root-link positions. The nearer arm picks +up the object; the other arm receives it. + +To keep both grasps spatially separated, HandOver deterministically samples at +most 1000 points from the object's triangle-mesh surface, transforms them by +the current object pose, and uses SVD to find the widest distribution direction +(`obj_longest_axis`). The handover arm selects the projected end nearest its +current TCP, and the receiving arm selects the opposite end. Grasp generation +receives this world-frame axis plus `is_positive_part`; a `None` axis retains +the ordinary unpartitioned center behavior for other callers. + +When the longest axis is within 45 degrees of world Z, HandOver uses vertical +mode: each approach's horizontal projection points from the acting arm's +current TCP toward the corresponding object position and is tilted downward +by 45 degrees. Otherwise it uses horizontal mode and both approaches are +world-Z downward. Pickup uses the observed object position, while receiving +uses the predicted middle object position because that pose is not observed +again during open-loop planning. + +After pickup, the action lifts in world Z and computes the middle object pose +by finding the root-link separation's largest-magnitude coordinate and setting +only that object coordinate to the two roots' midpoint. From the first grasp +through middle transfer, and from the receiving grasp through final lowering, +EEF waypoint rotations remain fixed; only translations change. The final +object translation comes from `HandOverGoal.target_pose`, while its execution +orientation stays consistent with the handover grasp. `HandOverOptions` owns +only the approach/lift distances and gripper interpolation count. The first +placement waypoint changes only horizontal coordinates and preserves the +handover height exactly; the second waypoint lowers to the final target. + +Planning failures are reported with semantic waypoint names and affected +environment IDs. For a failed motion phase, HandOver diagnoses each target +with IK; if all target waypoints are reachable, the report identifies the +interval between them as a likely path- or collision-planning failure. As with the other coordinated primitive, cuRobo does not currently support its dual-arm `strategy="motion_gen"` path. diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index ebe26490e..882a06f51 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -753,10 +753,10 @@ should pass their fresh context explicitly. ```{attention} Automatic dynamic-goal invalidation is dependency-driven. A goal must contain a `SceneEntityPose`, or an object-centric primitive must explicitly declare the -`ObjectSemantics.entity_id` whose snapshot pose it consumes. `PickUp` and the -implicit-initial-pose path of coordinated pickup declare that dependency -automatically. The deprecated live-entity fallback does not trigger -scene-motion replanning. +`ObjectSemantics.entity_id` whose snapshot pose it consumes. `PickUp`, +`HandOver`, and the implicit-initial-pose path of coordinated pickup declare +that dependency automatically. The deprecated live-entity fallback does not +trigger scene-motion replanning. Dynamic collision invalidation is provider-driven. Only registered, pose-updatable collision entities are supported; adding/removing obstacles or @@ -766,7 +766,7 @@ changing their geometry requires rebuilding the planner world. ## Planning success versus physical success `ActionPlan.plan_success` only means a valid command plan was produced for an -environment row. Pick, place, handover, and coordinated skills also return an +environment row. Pick, place, and coordinated skills also return an uncommitted `StateDelta` describing the attachment state expected after execution. @@ -774,8 +774,9 @@ execution. Multi-arm grasps use multiple entries that share the same `ObjectSemantics`; there is no parallel coordinated-attachment representation to synchronize. Consumers query per-environment active and exclusive-hold masks from that one -map. A single-arm transport, release, or handover row fails safely while a -second manipulator still holds the same semantic object or live entity. +map. A single-arm transport or release row fails safely while a second +manipulator still holds the same semantic object or live entity. The unified +`HandOver` action instead requires both candidate arms to start unoccupied. At the terminal waypoint, an `ExecutionSession` requests an external per-environment verification mask before committing a non-empty effect: @@ -788,7 +789,7 @@ if tick.pending_effect is not None: ``` This prevents a collision-free or well-tracked command plan from being -misreported as a successful grasp, release, or handover. The typed +misreported as a successful grasp or release. The typed `EffectVerificationRequest` persists on subsequent ticks while waiting; `EFFECT_VERIFICATION_REQUIRED` remains a one-time observability event. diff --git a/docs/source/tutorial/atomic_actions.rst b/docs/source/tutorial/atomic_actions.rst index f8c756cec..093cf7d58 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -123,6 +123,7 @@ Focused examples live under ``scripts/tutorials/atomic_action``: * ``control_dt.py`` * ``pickup.py`` * ``move_held_object.py`` +* ``pour.py`` * ``place.py`` * ``assemble.py`` * ``press.py`` @@ -143,6 +144,7 @@ video under ``outputs/videos``: python scripts/tutorials/atomic_action/move_end_effector.py --headless --auto_play --device cpu python scripts/tutorials/atomic_action/control_dt.py --headless --auto_play --device cpu python scripts/tutorials/atomic_action/pickup.py --headless --auto_play --device cpu + python scripts/tutorials/atomic_action/pour.py --headless --auto_play --device cpu python scripts/tutorials/atomic_action/assemble.py --headless --auto_play --device cpu python scripts/tutorials/atomic_action/hand_over.py --headless --auto_play --device cpu @@ -465,7 +467,7 @@ create a scene dependency. Task-state effects ------------------ -Pick, place, handover, and coordinated skills declare attachment changes as a +Pick, place, and coordinated skills declare attachment changes as a :class:`~embodichain.lab.sim.atomic_actions.StateDelta`. Planning does not commit those changes. During closed-loop execution, a non-empty effect requires an external per-environment verification mask: diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index ab6a5f760..1fa805803 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -32,6 +32,7 @@ Affordance, AntipodalAffordance, AssembleAffordance, + AxisAlignAffordance, InteractionPoints, PressAffordance, SlideAffordance, @@ -97,6 +98,9 @@ from .transports import EndpointCommandRouter, EndpointCommandTransport from .primitives import ( AssembleGoal, + AxisAlign, + AxisAlignGoal, + AxisAlignOptions, BUILTIN_ACTION_TYPES, CoordinatedPickGoal, CoordinatedPickment, @@ -107,6 +111,7 @@ EndEffectorPoseGoal, GraspGoal, HandOver, + HandOverGoal, HandOverOptions, HeldObjectPoseGoal, JointPositionGoal, @@ -121,6 +126,9 @@ Place, PlaceGoal, PlaceOptions, + Pour, + PourGoal, + PourOptions, Press, PressGoal, PressOptions, @@ -174,6 +182,10 @@ "AntipodalAffordance", "AssembleAffordance", "AssembleGoal", + "AxisAlign", + "AxisAlignAffordance", + "AxisAlignGoal", + "AxisAlignOptions", "AtomicAction", "AtomicActionEngine", "BUILTIN_ACTION_TYPES", @@ -217,6 +229,7 @@ "GRASP_CAPABILITY", "GraspGoal", "HandOver", + "HandOverGoal", "HandOverOptions", "HeldObjectPoseGoal", "HeldObjectState", @@ -245,6 +258,9 @@ "Place", "PlaceGoal", "PlaceOptions", + "Pour", + "PourGoal", + "PourOptions", "PlannerDiagnostics", "PlanningContext", "PoseGoalValue", diff --git a/embodichain/lab/sim/atomic_actions/affordance.py b/embodichain/lab/sim/atomic_actions/affordance.py index ad6b7c8af..33172d16e 100644 --- a/embodichain/lab/sim/atomic_actions/affordance.py +++ b/embodichain/lab/sim/atomic_actions/affordance.py @@ -16,9 +16,10 @@ from __future__ import annotations -import torch from dataclasses import dataclass, field -from typing import Any, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, ClassVar + +import torch from embodichain.toolkits.graspkit.pg_grasp import ( GraspGenerator, @@ -86,6 +87,9 @@ class AntipodalAffordance(Affordance): _generator: GraspGenerator | None = field(default=None, init=False, repr=False) + MAX_SURFACE_POINT_COUNT: ClassVar[int] = 1000 + """Maximum point-cloud size used for geometry-distribution analysis.""" + def _init_generator(self) -> None: if self.mesh_vertices is None or self.mesh_triangles is None: logger.log_error( @@ -117,17 +121,35 @@ def get_valid_grasp_poses( approach_direction: torch.Tensor = torch.tensor( [0, 0, -1], dtype=torch.float32 ), - object_part: str = "center", + obj_longest_axis: torch.Tensor | None = None, + is_positive_part: bool | torch.Tensor = True, ) -> list[tuple[torch.Tensor, torch.Tensor]]: + """Return valid grasps, optionally restricted to one end of an axis. + + Args: + obj_poses: Batched object poses with shape ``(B, 4, 4)``. + approach_direction: Shared world-frame approach direction. + obj_longest_axis: Optional shared ``(3,)`` or batched ``(B, 3)`` + world-frame object axis. ``None`` keeps the default center mode + and does not partition grasp candidates. + is_positive_part: Select the positive or negative projected half of + ``obj_longest_axis``. May be one bool or a ``(B,)`` bool tensor. + + Returns: + Per-object ``(grasp_poses, costs)`` tuples. + """ if self._generator is None: self._init_generator() approach_direction = self._resolve_approach_direction(approach_direction) + axes = self._resolve_grasp_region_axes(obj_poses, obj_longest_axis) + positive_parts = self._resolve_positive_parts(obj_poses, is_positive_part) results = [] for i, obj_pose in enumerate(obj_poses): is_success, grasp_poses, _, costs = self._generator.get_valid_grasp_poses( object_pose=obj_pose, approach_direction=approach_direction, - object_part=object_part, + obj_longest_axis=None if axes is None else axes[i], + is_positive_part=bool(positive_parts[i].item()), ) if grasp_poses.shape == (4, 4): grasp_poses = grasp_poses.unsqueeze(0) @@ -146,6 +168,190 @@ def get_valid_grasp_poses( results.append((grasp_poses, costs)) return results + def sample_surface_points(self, max_points: int = 1000) -> torch.Tensor: + """Deterministically sample at most 1000 target-local surface points. + + Triangle faces are selected proportionally to area and barycentric + coordinates use deterministic low-discrepancy sequences. Degenerate or + triangle-free meshes fall back to an evenly subsampled vertex cloud. + + Args: + max_points: Requested point cap in ``[1, 1000]``. + + Returns: + Target-local surface points with shape ``(N, 3)``. + + Raises: + ValueError: If mesh geometry or ``max_points`` is invalid. + """ + if not isinstance(max_points, int) or isinstance(max_points, bool): + raise TypeError("max_points must be an integer.") + if not 1 <= max_points <= self.MAX_SURFACE_POINT_COUNT: + raise ValueError( + f"max_points must be between 1 and {self.MAX_SURFACE_POINT_COUNT}." + ) + if self.mesh_vertices is None: + raise ValueError("AntipodalAffordance requires mesh_vertices.") + vertices = self.mesh_vertices.to(dtype=torch.float32) + if vertices.ndim != 2 or vertices.shape[1] != 3 or vertices.shape[0] == 0: + raise ValueError("mesh_vertices must have shape (N, 3) with N > 0.") + + triangles = self.mesh_triangles + if triangles is None or triangles.numel() == 0: + return self._evenly_subsample_points(vertices, max_points) + triangles = triangles.to(device=vertices.device, dtype=torch.long) + if triangles.ndim != 2 or triangles.shape[1] != 3: + raise ValueError("mesh_triangles must have shape (M, 3).") + if triangles.min() < 0 or triangles.max() >= vertices.shape[0]: + raise ValueError("mesh_triangles contains an invalid vertex index.") + + face_vertices = vertices[triangles] + face_areas = 0.5 * torch.linalg.vector_norm( + torch.cross( + face_vertices[:, 1] - face_vertices[:, 0], + face_vertices[:, 2] - face_vertices[:, 0], + dim=1, + ), + dim=1, + ) + valid_faces = face_areas > torch.finfo(vertices.dtype).eps + if not valid_faces.any(): + return self._evenly_subsample_points(vertices, max_points) + face_vertices = face_vertices[valid_faces] + face_areas = face_areas[valid_faces] + + sample_index = torch.arange( + max_points, device=vertices.device, dtype=vertices.dtype + ) + area_quantiles = (sample_index + 0.5) / max_points + cumulative_area = torch.cumsum(face_areas / face_areas.sum(), dim=0) + face_indices = torch.searchsorted(cumulative_area, area_quantiles).clamp_max( + face_vertices.shape[0] - 1 + ) + sampled_faces = face_vertices[face_indices] + + barycentric_u = torch.frac((sample_index + 0.5) * 0.7548776662466927) + barycentric_v = torch.frac((sample_index + 0.5) * 0.5698402909980532) + sqrt_u = torch.sqrt(barycentric_u) + weights = torch.stack( + ( + 1.0 - sqrt_u, + sqrt_u * (1.0 - barycentric_v), + sqrt_u * barycentric_v, + ), + dim=1, + ) + return torch.sum(sampled_faces * weights.unsqueeze(2), dim=1) + + def get_object_longest_axis( + self, + obj_poses: torch.Tensor, + *, + max_points: int = 1000, + ) -> torch.Tensor: + """Find the widest surface-point distribution axis in world space. + + Args: + obj_poses: Current object poses with shape ``(B, 4, 4)``. + max_points: Surface point cap, never greater than 1000. + + Returns: + Normalized first right-singular vectors with shape ``(B, 3)``. + + Raises: + ValueError: If poses or sampled geometry are invalid or degenerate. + """ + if obj_poses.ndim != 3 or obj_poses.shape[1:] != (4, 4): + raise ValueError("obj_poses must have shape (B, 4, 4).") + points = self.sample_surface_points(max_points=max_points).to( + device=obj_poses.device, + dtype=torch.float32, + ) + poses = obj_poses.to(dtype=torch.float32) + world_points = ( + torch.matmul(points.unsqueeze(0), poses[:, :3, :3].transpose(1, 2)) + + poses[:, None, :3, 3] + ) + centered = world_points - world_points.mean(dim=1, keepdim=True) + if torch.any(torch.linalg.vector_norm(centered, dim=2).amax(dim=1) <= 1.0e-8): + raise ValueError("Object surface point distribution is degenerate.") + _, singular_values, vh = torch.linalg.svd(centered, full_matrices=False) + if torch.any(singular_values[:, 0] <= 1.0e-8): + raise ValueError("Object surface point distribution has no principal axis.") + return torch.nn.functional.normalize(vh[:, 0, :], dim=1) + + @staticmethod + def _evenly_subsample_points( + points: torch.Tensor, + max_points: int, + ) -> torch.Tensor: + """Return an evenly spaced deterministic subset of ``points``.""" + if points.shape[0] <= max_points: + return points.clone() + indices = ( + torch.linspace( + 0, + points.shape[0] - 1, + max_points, + device=points.device, + ) + .round() + .to(torch.long) + ) + return points[indices] + + def _resolve_grasp_region_axes( + self, + obj_poses: torch.Tensor, + obj_longest_axis: torch.Tensor | None, + ) -> torch.Tensor | None: + """Validate and broadcast optional world-frame grasp-region axes.""" + if obj_longest_axis is None: + return None + axes = torch.as_tensor( + obj_longest_axis, + dtype=torch.float32, + device=self._generator.device, + ) + if axes.shape == (3,): + axes = axes.unsqueeze(0).expand(obj_poses.shape[0], -1) + if axes.shape != (obj_poses.shape[0], 3): + raise ValueError( + "obj_longest_axis must have shape (3,) or " + f"({obj_poses.shape[0]}, 3)." + ) + if not torch.isfinite(axes).all() or torch.any( + torch.linalg.vector_norm(axes, dim=1) <= 1.0e-8 + ): + raise ValueError("obj_longest_axis must be finite and non-zero.") + return torch.nn.functional.normalize(axes, dim=1) + + def _resolve_positive_parts( + self, + obj_poses: torch.Tensor, + is_positive_part: bool | torch.Tensor, + ) -> torch.Tensor: + """Validate and broadcast positive/negative grasp-region selections.""" + if isinstance(is_positive_part, bool): + return torch.full( + (obj_poses.shape[0],), + is_positive_part, + dtype=torch.bool, + device=self._generator.device, + ) + positive_parts = torch.as_tensor( + is_positive_part, + device=self._generator.device, + ) + if positive_parts.dtype != torch.bool or positive_parts.shape != ( + obj_poses.shape[0], + ): + raise ValueError( + "is_positive_part must be a bool or a bool tensor with shape " + f"({obj_poses.shape[0]},)." + ) + return positive_parts + def get_dual_arm_valid_grasp_poses( self, obj_poses: torch.Tensor, @@ -233,6 +439,29 @@ def get_best_grasp_poses( return is_success_t, grasp_xpos, open_length_t +@dataclass +class AxisAlignAffordance(AntipodalAffordance): + """Antipodal grasp affordance with an object-local alignment axis.""" + + internal_axis: torch.Tensor = field( + default_factory=lambda: torch.tensor([0.0, 0.0, 1.0]) + ) + """Axis expressed in the target object's local frame.""" + + def __post_init__(self) -> None: + if ( + not isinstance(self.internal_axis, torch.Tensor) + or self.internal_axis.shape != (3,) + or not torch.isfinite(self.internal_axis).all() + ): + raise ValueError( + "AxisAlignAffordance.internal_axis must be a finite (3,) tensor." + ) + if torch.linalg.vector_norm(self.internal_axis) <= 1.0e-6: + raise ValueError("AxisAlignAffordance.internal_axis must be non-zero.") + self.internal_axis = self.internal_axis.clone() + + @dataclass class TwistAffordance(Affordance): """Target-local grasp point and rotation-axis geometry for twisting.""" @@ -617,6 +846,7 @@ def get_assemble_object_pose(self, base_pose: torch.Tensor) -> torch.Tensor: __all__ = [ "Affordance", "AntipodalAffordance", + "AxisAlignAffordance", "SlideAffordance", "PressAffordance", "TwistAffordance", diff --git a/embodichain/lab/sim/atomic_actions/primitives/__init__.py b/embodichain/lab/sim/atomic_actions/primitives/__init__.py index 718e6780b..311b4166e 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/__init__.py +++ b/embodichain/lab/sim/atomic_actions/primitives/__init__.py @@ -19,6 +19,7 @@ from __future__ import annotations from ..core import AtomicAction +from .axis_align import AxisAlign, AxisAlignGoal, AxisAlignOptions from .coordinated_pickment import ( CoordinatedPickGoal, CoordinatedPickment, @@ -29,7 +30,7 @@ CoordinatedPlacementGoal, CoordinatedPlacementOptions, ) -from .hand_over import HandOver, HandOverOptions +from .hand_over import HandOver, HandOverGoal, HandOverOptions from .move_end_effector import ( EndEffectorPoseGoal, MoveEndEffector, @@ -44,6 +45,7 @@ from .pick_up import GraspGoal, PickUp, PickUpOptions from .place import AssembleGoal, Place, PlaceGoal, PlaceOptions from .press import Press, PressGoal, PressOptions +from .pour import Pour, PourGoal, PourOptions from .slide import ( Slide, SlideGoal, @@ -55,8 +57,10 @@ MoveEndEffector, MoveJoints, PickUp, + AxisAlign, MoveHeldObject, Place, + Pour, Press, Slide, Twist, @@ -68,6 +72,9 @@ __all__ = [ "AssembleGoal", + "AxisAlign", + "AxisAlignGoal", + "AxisAlignOptions", "BUILTIN_ACTION_TYPES", "CoordinatedPickGoal", "CoordinatedPickment", @@ -78,6 +85,7 @@ "EndEffectorPoseGoal", "GraspGoal", "HandOver", + "HandOverGoal", "HandOverOptions", "HeldObjectPoseGoal", "JointPositionGoal", @@ -92,6 +100,9 @@ "Place", "PlaceGoal", "PlaceOptions", + "Pour", + "PourGoal", + "PourOptions", "Press", "PressGoal", "PressOptions", diff --git a/embodichain/lab/sim/atomic_actions/primitives/axis_align.py b/embodichain/lab/sim/atomic_actions/primitives/axis_align.py new file mode 100644 index 000000000..0a06ff78e --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/axis_align.py @@ -0,0 +1,624 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""AxisAlign atomic action implementation.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import ClassVar + +import torch + +from embodichain.utils import logger +from embodichain.utils.math import ( + axis_angle_to_rotation_matrix, + get_relative_rotation, + pose_inv, +) + +from embodichain.lab.sim.atomic_actions.affordance import AxisAlignAffordance +from embodichain.lab.sim.atomic_actions.bindings import JointPositionTarget +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 embodichain.lab.sim.atomic_actions.invocation import ResolvedActionRequest +from embodichain.lab.sim.atomic_actions.plans import ( + ActionPlan, + TimedTrajectory, + normalize_success_mask, +) +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.primitives.pick_up import PickUpOptions +from embodichain.lab.sim.atomic_actions.requirements import ( + CARTESIAN_POSE_CAPABILITY, + FORWARD_KINEMATICS_CAPABILITY, + SkillBindingContract, +) +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 AxisAlignGoal(ObjectActionGoal): + """Object whose local axis should be aligned after an antipodal grasp.""" + + goal_kind: ClassVar[str] = "axis_align" + + grasp_xpos: PoseGoalValue | None = None + """Optional explicit end-effector grasp pose; omitted poses are sampled.""" + + def __post_init__(self) -> None: + ObjectActionGoal.__post_init__(self) + if self.grasp_xpos is not None: + validate_pose_goal(self.grasp_xpos, "grasp_xpos", allow_waypoints=False) + + +@dataclass(frozen=True, slots=True, eq=False) +class AxisAlignOptions(PickUpOptions): + """Per-invocation grasp-and-axis-alignment behavior.""" + + target_axis: torch.Tensor = torch.tensor([0.0, 0.0, 1.0]) + """Desired world-frame axis, shape ``(3,)`` or ``(B, 3)``.""" + + lower_distance: float = 0.03 + """World-Z distance (m) to lower the aligned object before release.""" + + def __post_init__(self) -> None: + PickUpOptions.__post_init__(self) + if ( + not isinstance(self.target_axis, torch.Tensor) + or self.target_axis.dim() not in (1, 2) + or self.target_axis.shape[-1] != 3 + or not torch.isfinite(self.target_axis).all() + ): + raise ValueError("target_axis must be a finite (3,) or (B, 3) tensor.") + if torch.any(torch.linalg.vector_norm(self.target_axis, dim=-1) <= 1.0e-6): + raise ValueError("target_axis must be non-zero.") + if not math.isfinite(self.lower_distance): + raise ValueError("lower_distance must be finite.") + if self.lower_distance < 0.0: + raise ValueError("lower_distance must be non-negative.") + object.__setattr__(self, "target_axis", self.target_axis.clone()) + + +class AxisAlign(AtomicAction[AxisAlignGoal, AxisAlignOptions]): + """Grasp an object, align its local axis to a world axis, and release it.""" + + skill_id: ClassVar[str] = "axis_align" + GoalType: ClassVar[type] = AxisAlignGoal + OptionsType: ClassVar[type] = AxisAlignOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) + open_loop: ClassVar[bool] = True + _UPRIGHT_HORIZONTAL_MAX_ABS_Z: ClassVar[float] = 0.5 + _UPRIGHT_TARGET_MIN_Z: ClassVar[float] = math.cos(math.pi / 6.0) + _UPRIGHT_GRASP_PRE_ROTATION: ClassVar[float] = math.pi / 4.0 + binding_contract: ClassVar[SkillBindingContract] = SkillBindingContract( + slots=( + make_manipulation_slot( + "primary", + motion_capabilities=frozenset( + { + CARTESIAN_POSE_CAPABILITY, + FORWARD_KINEMATICS_CAPABILITY, + } + ), + grasp_commands={ + OPEN_COMMAND: JointPositionCommand, + GRASP_COMMAND: JointPositionCommand, + }, + ), + ), + ) + + def __init__(self, default_options: AxisAlignOptions | None = None) -> None: + super().__init__(default_options) + + def _scene_dependencies( + self, + request: ResolvedActionRequest[AxisAlignGoal, AxisAlignOptions], + ) -> tuple[str, ...]: + """Include the semantic object when it has a stable scene identity.""" + dependencies = set(super()._scene_dependencies(request)) + entity_id = request.goal.semantics.entity_id + if entity_id is not None: + dependencies.add(entity_id) + return tuple(sorted(dependencies)) + + 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[AxisAlignGoal, AxisAlignOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan all seven physical actions in two arm-planning phases.""" + target = request.goal + options = request.skill_options + affordance = self._require_axis_align_affordance(target.semantics) + motion_endpoint = request.binding.endpoint("primary", "motion") + grasp_endpoint = request.binding.endpoint("primary", "grasp") + manipulator = motion_endpoint.require_target(JointPositionTarget) + end_effector = grasp_endpoint.require_target(JointPositionTarget) + 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 = grasp_endpoint.joint_positions( + OPEN_COMMAND, + num_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + hand_grasp_qpos = grasp_endpoint.joint_positions( + GRASP_COMMAND, + num_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + approach_direction = options.approach_direction.to( + device=self.device, dtype=torch.float32 + ) + approach_direction = approach_direction / torch.linalg.vector_norm( + approach_direction + ) + object_pose = _resolve_object_pose( + target.semantics, + context, + name="axis_align_object_pose", + ) + + # Resolve the shortest object rotation before selecting a grasp. The + # source axis is ``object_rotation @ internal_axis`` in world space; + # ``rotation_axis`` is the normalized cross product from that source to + # the requested world-space target. For opposite axes the helper picks + # a deterministic perpendicular axis instead of dividing by a near-zero + # cross product. + source_axis, target_axis, rotation_axis, rotation_angle = ( + self._axis_alignment_parameters( + object_pose, + affordance.internal_axis, + options.target_axis, + ) + ) + + # If no explicit grasp was supplied, _resolve_grasp_pose first filters + # invalid affordance samples, then gives priority to the candidates whose + # TCP y-axis is most perpendicular to ``rotation_axis``. Grasp-generator + # cost only breaks ties between equally perpendicular candidates. This + # orientation keeps the gripper's roll axis away from the object's + # rotation axis and generally leaves the arm more room for the alignment + # motion. The antipodal pose has a 180-degree symmetric alternative; + # after choosing the sample, select whichever symmetric orientation is + # closer to the arm's currently observed FK pose. + grasp_success, grasp_xpos = self._resolve_grasp_pose( + target, + affordance, + object_pose, + context, + approach_direction, + rotation_axis, + rotation_angle, + object_part=options.pick_object_part, + ) + grasp_xpos = self._find_symmetric_nearest_xpos( + grasp_xpos, + reference_xpos=self.robot.compute_fk( + qpos=start_arm_qpos, + name=manipulator.control_part, + to_matrix=True, + ), + ) + grasp_success = normalize_success_mask( + grasp_success, + num_envs=self.num_envs, + device=self.device, + name="Axis-align grasp success", + ) + if not grasp_success.any(): + logger.log_warning("AxisAlign failed to resolve a grasp pose.") + return self.failed_plan( + request, context, message="Failed to resolve a grasp pose." + ) + + # Upright handling is enabled independently for each environment when + # the current object axis is mostly horizontal (|world z| <= 0.5) and + # the requested target points mostly upward (within 30 degrees of +Z). + # Before deriving the fixed object-to-EEF grasp transform, rotate only + # the grasp orientation by 45 degrees *opposite* ``rotation_axis``; its + # position is unchanged. The subsequent alignment still rotates the + # object through the full shortest arc, but the arm starts that arc with + # a 45-degree bias, reducing the link sweep near the table. + upright_mask = ( + source_axis[:, 2].abs() <= self._UPRIGHT_HORIZONTAL_MAX_ABS_Z + ) & (target_axis[:, 2] >= self._UPRIGHT_TARGET_MIN_Z) + grasp_xpos = self._apply_upright_grasp_pre_rotation( + grasp_xpos, + rotation_axis, + upright_mask, + ) + + pre_grasp_xpos = translate_pose_world( + grasp_xpos, -approach_direction * options.pre_grasp_distance + ) + lift_xpos = translate_pose_world( + grasp_xpos, + torch.tensor( + [0.0, 0.0, options.lift_height], + device=self.device, + dtype=torch.float32, + ), + ) + object_to_eef = torch.bmm(pose_inv(object_pose), grasp_xpos) + lifted_object_pose = torch.bmm(lift_xpos, pose_inv(object_to_eef)) + + n_approach, n_reach, n_lift, n_align, n_lower = self._motion_segment_lengths( + request.motion_policy.sample_count, + options.hand_interp_steps, + ) + interpolation_dt = context.require_control_dt() + # Only the final aligned pose is a planner target. Supplying n_align + # intermediate Cartesian keyframes would make CuRobo call plan_pose once + # per keyframe; n_align is instead retained as the output sample budget + # for the continuous post-close phase. + align_xpos = self._axis_alignment_eef_keyframes( + lifted_object_pose, + object_to_eef, + affordance.internal_axis, + options.target_axis, + waypoint_count=1, + ) + lower_xpos = translate_pose_world( + align_xpos[:, -1], + torch.tensor( + [0.0, 0.0, -options.lower_distance], + device=self.device, + dtype=torch.float32, + ), + ) + + # CuRobo planning is grouped by gripper state. The open-gripper phase + # contains both the pre-grasp and grasp waypoints, so one generate call + # replaces the former independent approach and reach calls. + pre_close_xpos = torch.stack([pre_grasp_xpos, grasp_xpos], dim=1) + pre_close_success, pre_close_arm = self._plan_pose_phase( + pre_close_xpos, + start_arm_qpos, + manipulator, + request, + n_approach + n_reach, + interpolation_dt, + ) + + # Once the gripper is closed, lifting, alignment, and lowering form one + # continuous held-object phase. Passing only those three semantic + # endpoints retains the required ordering without expanding the rotation + # into many CuRobo plan_pose calls. Together with the open-gripper phase, + # the action now uses two MotionGenerator.generate calls and five backend + # target plans instead of n_align + 4 backend target plans. + post_close_xpos = torch.cat( + [lift_xpos[:, None], align_xpos, lower_xpos[:, None]], dim=1 + ) + post_close_success, post_close_arm = self._plan_pose_phase( + post_close_xpos, + pre_close_arm[:, -1], + manipulator, + request, + n_lift + n_align + n_lower, + interpolation_dt, + ) + success = grasp_success & normalize_success_mask( + pre_close_success & post_close_success, + num_envs=self.num_envs, + device=self.device, + name="Axis-align trajectory 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, + ) + segment_lengths = { + "approach": pre_close_arm.shape[1], + "close": hand_close.shape[1], + "manipulate": post_close_arm.shape[1], + "open": hand_open.shape[1], + } + 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 = pre_close_arm.shape[1] + full[:, :offset, arm_joint_ids] = pre_close_arm + full[:, :offset, hand_joint_ids] = hand_open_qpos.unsqueeze(1) + stop = offset + hand_close.shape[1] + full[:, offset:stop, arm_joint_ids] = pre_close_arm[:, -1].unsqueeze(1) + full[:, offset:stop, hand_joint_ids] = hand_close + offset = stop + stop = offset + post_close_arm.shape[1] + full[:, offset:stop, arm_joint_ids] = post_close_arm + full[:, offset:stop, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1) + offset = stop + full[:, offset:, arm_joint_ids] = post_close_arm[:, -1].unsqueeze(1) + full[:, offset:, hand_joint_ids] = hand_open + + return self.build_plan( + request, + context, + success=success, + trajectory=TimedTrajectory.from_uniform_step( + full, + env_ids=context.env_ids, + step_dt=interpolation_dt, + ), + expected_effects=StateDelta(), + segment_lengths=segment_lengths, + ) + + def _resolve_grasp_pose( + self, + goal: AxisAlignGoal, + affordance: AxisAlignAffordance, + object_pose: torch.Tensor, + context: PlanningContext, + approach_direction: torch.Tensor, + rotation_axis: torch.Tensor, + rotation_angle: torch.Tensor, + *, + object_part: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Resolve an explicit grasp or select the lowest-cost sampled grasp.""" + if goal.grasp_xpos is not None: + grasp_xpos = resolve_pose_target( + resolve_pose_goal(goal.grasp_xpos, context, name="grasp_xpos"), + num_envs=self.num_envs, + device=self.device, + ) + return ( + torch.ones(self.num_envs, dtype=torch.bool, device=self.device), + grasp_xpos, + ) + + sampled = affordance.get_valid_grasp_poses( + obj_poses=object_pose, + approach_direction=approach_direction, + obj_longest_axis=( + None + if object_part == "center" + else torch.tensor( + [0.0, 0.0, 1.0], + dtype=torch.float32, + device=self.device, + ) + ), + is_positive_part=object_part == "top", + ) + poses: list[torch.Tensor] = [] + success: list[bool] = [] + for env_index, (candidates, costs) in enumerate(sampled): + candidates = candidates.to(device=self.device, dtype=torch.float32) + costs = costs.to(device=self.device, dtype=torch.float32) + valid = candidates.shape[0] > 0 and bool(torch.isfinite(costs).any()) + if valid: + finite_cost = torch.isfinite(costs) + if rotation_angle[env_index] > 1.0e-6: + grasp_y_axis = torch.nn.functional.normalize( + candidates[:, :3, 1], dim=1 + ) + perpendicularity_error = torch.abs( + torch.matmul(grasp_y_axis, rotation_axis[env_index]) + ) + best_error = perpendicularity_error[finite_cost].min() + preferred = finite_cost & torch.isclose( + perpendicularity_error, + best_error, + atol=1.0e-6, + rtol=1.0e-5, + ) + ranked_costs = torch.where( + preferred, + costs, + torch.full_like(costs, torch.inf), + ) + best_index = int(torch.argmin(ranked_costs).item()) + else: + best_index = int(torch.argmin(costs).item()) + poses.append(candidates[best_index]) + else: + poses.append(torch.eye(4, device=self.device, dtype=torch.float32)) + success.append(valid) + return ( + torch.tensor(success, dtype=torch.bool, device=self.device), + torch.stack(poses), + ) + + def _plan_pose_phase( + self, + target_pose: torch.Tensor, + start_qpos: torch.Tensor, + manipulator: JointPositionTarget, + request: ResolvedActionRequest[AxisAlignGoal, AxisAlignOptions], + sample_count: int, + interpolation_dt: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Plan one continuous arm phase with a fixed gripper command.""" + 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=manipulator.control_part, + sample_count=sample_count, + interpolation_dt=interpolation_dt, + ), + ) + assert isinstance(result.success, torch.Tensor) + assert result.positions is not None + return result.success, result.positions + + def _axis_alignment_eef_keyframes( + self, + object_pose: torch.Tensor, + object_to_eef: torch.Tensor, + internal_axis: torch.Tensor, + target_axis: torch.Tensor, + *, + waypoint_count: int, + ) -> torch.Tensor: + """Rotate the object in place along the shortest axis-alignment arc.""" + _, _, axis, angle = self._axis_alignment_parameters( + object_pose, + internal_axis, + target_axis, + ) + + fractions = torch.linspace( + 1.0 / waypoint_count, + 1.0, + waypoint_count, + dtype=torch.float32, + device=self.device, + ) + rotation_vectors = ( + axis[:, None, :] * angle[:, None, None] * fractions[None, :, None] + ) + delta_rotation = axis_angle_to_rotation_matrix(rotation_vectors) + object_keyframes = object_pose[:, None].repeat(1, waypoint_count, 1, 1) + object_keyframes[:, :, :3, :3] = torch.matmul( + delta_rotation, object_pose[:, None, :3, :3] + ) + return torch.matmul(object_keyframes, object_to_eef[:, None]) + + def _axis_alignment_parameters( + self, + object_pose: torch.Tensor, + internal_axis: torch.Tensor, + target_axis: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Return normalized source/target axes and the shortest rotation.""" + internal = internal_axis.to(device=self.device, dtype=torch.float32) + internal = internal / torch.linalg.vector_norm(internal) + source = torch.matmul(object_pose[:, :3, :3], internal) + source = torch.nn.functional.normalize(source, dim=1) + target = target_axis.to(device=self.device, dtype=torch.float32) + if target.shape == (3,): + target = target.unsqueeze(0).expand(self.num_envs, -1) + elif target.shape != (self.num_envs, 3): + raise ValueError( + f"target_axis must have shape (3,) or ({self.num_envs}, 3)." + ) + target = torch.nn.functional.normalize(target, dim=1) + + cross = torch.linalg.cross(source, target, dim=1) + sin_angle = torch.linalg.vector_norm(cross, dim=1) + cos_angle = torch.sum(source * target, dim=1).clamp(-1.0, 1.0) + axis = cross / sin_angle.clamp_min(1.0e-8).unsqueeze(1) + basis = torch.eye(3, dtype=torch.float32, device=self.device) + reference = basis[torch.argmin(torch.abs(source), dim=1)] + fallback_axis = torch.nn.functional.normalize( + torch.linalg.cross(source, reference, dim=1), dim=1 + ) + degenerate = sin_angle <= 1.0e-6 + axis = torch.where(degenerate.unsqueeze(1), fallback_axis, axis) + angle = torch.atan2(sin_angle, cos_angle) + opposite = degenerate & (cos_angle < 0.0) + angle = torch.where(opposite, torch.full_like(angle, torch.pi), angle) + return source, target, axis, angle + + def _apply_upright_grasp_pre_rotation( + self, + grasp_xpos: torch.Tensor, + rotation_axis: torch.Tensor, + upright_mask: torch.Tensor, + ) -> torch.Tensor: + """Pre-rotate upright grasps to reduce the arm's table-side sweep.""" + if not upright_mask.any(): + return grasp_xpos + delta = axis_angle_to_rotation_matrix( + -rotation_axis * self._UPRIGHT_GRASP_PRE_ROTATION + ) + rotated = grasp_xpos.clone() + rotated[:, :3, :3] = torch.matmul(delta, grasp_xpos[:, :3, :3]) + return torch.where(upright_mask[:, None, None], rotated, grasp_xpos) + + @staticmethod + def _motion_segment_lengths( + sample_count: int, + hand_interp_steps: int, + ) -> tuple[int, int, int, int, int]: + motion_count = sample_count - 2 * hand_interp_steps + if motion_count < 5: + raise ValueError( + "Not enough waypoints for AxisAlign. Increase sample_count or " + "decrease hand_interp_steps." + ) + base, remainder = divmod(motion_count, 5) + values = [base + (index < remainder) for index in range(5)] + return values[0], values[1], values[2], values[3], values[4] + + @staticmethod + def _require_axis_align_affordance( + semantics: ObjectSemantics, + ) -> AxisAlignAffordance: + affordance = semantics.affordance + if not isinstance(affordance, AxisAlignAffordance): + raise ValueError("AxisAlign requires an AxisAlignAffordance.") + return affordance + + +__all__ = ["AxisAlign", "AxisAlignGoal", "AxisAlignOptions"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/hand_over.py b/embodichain/lab/sim/atomic_actions/primitives/hand_over.py index 6e7fa9269..535c5d850 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/hand_over.py +++ b/embodichain/lab/sim/atomic_actions/primitives/hand_over.py @@ -14,30 +14,39 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""HandOver atomic action implementation.""" +"""Pick-up-and-handover atomic action implementation.""" from __future__ import annotations +import math +from collections.abc import Mapping from dataclasses import dataclass from typing import ClassVar import torch from embodichain.utils import logger -from embodichain.utils.math import pose_inv +from embodichain.utils.math import get_relative_rotation, pose_inv -from embodichain.lab.sim.atomic_actions.bindings import JointPositionTarget +from embodichain.lab.sim.atomic_actions.affordance import AntipodalAffordance +from embodichain.lab.sim.atomic_actions.bindings import ( + EndpointBinding, + JointPositionTarget, +) 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.core import AtomicAction 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 embodichain.lab.sim.atomic_actions.invocation import ( ActionOptions, ResolvedActionRequest, @@ -47,122 +56,120 @@ TimedTrajectory, normalize_success_mask, ) +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, + plan_named_arm_trajectory, + repeat_qpos, + resolve_batched_pose, +) from embodichain.lab.sim.atomic_actions.requirements import ( CARTESIAN_POSE_CAPABILITY, DisjointResourceSlots, FORWARD_KINEMATICS_CAPABILITY, SkillBindingContract, ) -from embodichain.lab.sim.atomic_actions.state import HeldObjectState, PlanningContext +from embodichain.lab.sim.atomic_actions.state import PlanningContext from embodichain.lab.sim.atomic_actions.trajectory_ops import ( interpolate_hand_qpos, translate_pose_world, ) -from embodichain.lab.sim.atomic_actions.primitives._helpers import ( - assemble_full_robot_trajectory, - plan_named_arm_trajectory, - repeat_qpos, - resolve_batched_pose, -) -from embodichain.lab.sim.atomic_actions.primitives._binding_contracts import ( - make_manipulation_slot, -) -from embodichain.lab.sim.atomic_actions.primitives.pick_up import GraspGoal @dataclass(frozen=True, slots=True, eq=False) -class HandOverOptions(ActionOptions): - """Per-invocation handover behavior and object-pose targets.""" +class HandOverGoal(ObjectActionGoal): + """Object to pick and hand over, plus its final object pose.""" - receive_pick_object_part: str = "bottom" - """Object part the receiving arm grasps during the handover - (see :meth:`AntipodalAffordance.get_valid_grasp_poses`).""" + goal_kind: ClassVar[str] = "hand_over" - middle_object_pose: torch.Tensor | None = None - """Object pose at the handover point where the receiving arm grasps it, - shape ``(4, 4)`` or ``(num_envs, 4, 4)``. Must be set by the caller.""" + target_pose: PoseGoalValue + """Final object pose after the receiving arm lowers and releases it.""" - final_object_pose: torch.Tensor | None = None - """Object pose the receiving arm delivers the object to, shape ``(4, 4)`` - or ``(num_envs, 4, 4)``. Must be set by the caller.""" + def __post_init__(self) -> None: + ObjectActionGoal.__post_init__(self) + validate_pose_goal(self.target_pose, "target_pose", allow_waypoints=False) - receive_approach_direction: torch.Tensor = torch.tensor( - [0.0, 0.0, -1.0], dtype=torch.float32 - ) - """World-frame approach direction used to sample and approach the receiving - grasp. Tune this (e.g. to ``[0, 0, 1]`` for a from-below receive grasp) so - the receiving arm does not collide with the transferring arm.""" + +@dataclass(frozen=True, slots=True, eq=False) +class HandOverOptions(ActionOptions): + """Per-invocation pick-up, handover, and placement behavior.""" pre_grasp_distance: float = 0.10 - """World distance to offset back from the receiving grasp pose along the - negative approach direction.""" + """Distance from each grasp pose to its approach pose, in metres.""" - lift_height: float = 0.08 - """World-Z lift distance for the transferring arm after it releases.""" + lift_height: float = 0.10 + """World-Z distance used to lift the object after the first grasp.""" hand_interp_steps: int = 10 - """Number of waypoints used for the receiving-hand close and the - transferring-hand release interpolations.""" - - hold_steps: int = 4 - """Number of waypoints to hold the handoff pose before releasing.""" - - retreat_steps: int = 24 - """Number of waypoints used for the final deliver/retreat segment.""" + """Waypoints used by every gripper open/close interpolation.""" def __post_init__(self) -> None: - if not isinstance(self.receive_pick_object_part, str) or not ( - self.receive_pick_object_part - ): - raise ValueError("receive_pick_object_part must be non-empty.") - if self.receive_approach_direction.shape != (3,): - raise ValueError("receive_approach_direction must have shape (3,).") - if not torch.isfinite(self.receive_approach_direction).all() or ( - torch.linalg.vector_norm(self.receive_approach_direction) <= 1.0e-6 - ): - raise ValueError("receive_approach_direction must be finite and non-zero.") - if self.pre_grasp_distance < 0.0 or self.lift_height < 0.0: - raise ValueError("pre_grasp_distance and lift_height must be non-negative.") - for name in ("hand_interp_steps", "hold_steps", "retreat_steps"): - if getattr(self, name) < 0: - raise ValueError(f"{name} must be non-negative.") - object.__setattr__( - self, - "receive_approach_direction", - self.receive_approach_direction.clone(), - ) - for name in ("middle_object_pose", "final_object_pose"): + for name in ("pre_grasp_distance", "lift_height"): value = getattr(self, name) - if value is not None: - object.__setattr__(self, name, value.clone()) + if not math.isfinite(value) or value < 0.0: + raise ValueError(f"{name} must be finite and non-negative.") + if self.hand_interp_steps < 1: + raise ValueError("hand_interp_steps must be at least 1.") + + +@dataclass(frozen=True, slots=True, eq=False) +class _Participant: + """One candidate arm, its hand, and resolved semantic hand commands.""" + + arm: JointPositionTarget + hand: JointPositionTarget + hand_open_qpos: torch.Tensor + hand_grasp_qpos: torch.Tensor @dataclass(frozen=True, slots=True, eq=False) class _HandOverResources: - """Invocation-bound control parts and compatible hand commands.""" + """The two invocation-bound candidate participants.""" + + first: _Participant + second: _Participant - transfer_arm: JointPositionTarget - receive_arm: JointPositionTarget - transfer_hand: JointPositionTarget - receive_hand: JointPositionTarget - transfer_hand_open_qpos: torch.Tensor - transfer_hand_close_qpos: torch.Tensor - receive_hand_open_qpos: torch.Tensor - receive_hand_close_qpos: torch.Tensor +@dataclass(frozen=True, slots=True, eq=False) +class _DirectionalPlan: + """One fixed handover-arm/receiving-arm assignment.""" + + success: torch.Tensor + trajectory: torch.Tensor + segment_lengths: dict[str, int] -class HandOver(AtomicAction[GraspGoal, HandOverOptions]): - """Hand an object from one arm to the other. - The transferring arm (already holding the object) moves it to a middle - handover pose, the receiving arm approaches and grasps a different part of - the object, the transferring arm releases and retreats, and the receiving - arm carries the object to a final pose. +class HandOver(AtomicAction[HandOverGoal, HandOverOptions]): + """Pick an object with the nearer arm, hand it over, and place it. + + For each environment, the action chooses the arm whose root link is closer + to the observed object pose. It samples at most 1000 mesh-surface points and + applies SVD in the current object pose to find ``obj_longest_axis``. When + that axis is closer to world Z than to the horizontal plane, both grasp + approaches point toward the object horizontally and tilt downward by 45 + degrees. Otherwise both approaches are world-Z downward. + + The first arm grasps the projected end of ``obj_longest_axis`` nearest its + current TCP; the receiving arm grasps the opposite end at the predicted + middle object pose. This keeps the two hands from selecting the same object + region regardless of whether a long object is standing or lying down. + + After each grasp waypoint, subsequent EEF waypoints preserve that grasp + rotation and change translation only. In particular, placement first moves + strictly horizontally at the handover height and then lowers to the final + target pose before releasing the object. """ skill_id: ClassVar[str] = "hand_over" - GoalType: ClassVar[type] = GraspGoal + GoalType: ClassVar[type] = HandOverGoal OptionsType: ClassVar[type] = HandOverOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("source", "destination") + end_effector_roles: ClassVar[tuple[str, ...]] = ("source", "destination") + open_loop: ClassVar[bool] = True + _SURFACE_POINT_COUNT: ClassVar[int] = 1000 + _VERTICAL_MODE_MIN_ABS_Z: ClassVar[float] = math.sqrt(0.5) binding_contract: ClassVar[SkillBindingContract] = SkillBindingContract( slots=( make_manipulation_slot( @@ -180,7 +187,12 @@ class HandOver(AtomicAction[GraspGoal, HandOverOptions]): ), make_manipulation_slot( "destination", - motion_capabilities=frozenset({CARTESIAN_POSE_CAPABILITY}), + motion_capabilities=frozenset( + { + CARTESIAN_POSE_CAPABILITY, + FORWARD_KINEMATICS_CAPABILITY, + } + ), grasp_commands={ OPEN_COMMAND: JointPositionCommand, GRASP_COMMAND: JointPositionCommand, @@ -189,84 +201,72 @@ class HandOver(AtomicAction[GraspGoal, HandOverOptions]): ), constraints=(DisjointResourceSlots(("source", "destination")),), ) - _repeat_qpos = staticmethod(repeat_qpos) def _scene_dependencies( self, - request: ResolvedActionRequest[GraspGoal, HandOverOptions], + request: ResolvedActionRequest[HandOverGoal, HandOverOptions], ) -> tuple[str, ...]: - """Return no goal-pose dependency because handover ignores grasp_xpos.""" - del request - return () + """Track both the initial object and any late-bound final target.""" + dependencies = set(super()._scene_dependencies(request)) + entity_id = request.goal.semantics.entity_id + if entity_id is not None: + dependencies.add(entity_id) + return tuple(sorted(dependencies)) def _resolve_resources( self, - request: ResolvedActionRequest[GraspGoal, HandOverOptions], + request: ResolvedActionRequest[HandOverGoal, HandOverOptions], ) -> _HandOverResources: - """Resolve source/destination roles from robot control parts.""" + """Resolve the two bound candidate arm/hand pairs.""" binding = request.binding - transfer_motion = binding.endpoint("source", "motion") - receive_motion = binding.endpoint("destination", "motion") - transfer_grasp = binding.endpoint("source", "grasp") - receive_grasp = binding.endpoint("destination", "grasp") - transfer_arm = transfer_motion.require_target(JointPositionTarget) - receive_arm = receive_motion.require_target(JointPositionTarget) - transfer_hand = transfer_grasp.require_target(JointPositionTarget) - receive_hand = receive_grasp.require_target(JointPositionTarget) - if transfer_arm.control_part == receive_arm.control_part: - raise ValueError( - "HandOver source and destination must use different manipulator " - "control parts." - ) - if transfer_hand.control_part == receive_hand.control_part: - raise ValueError( - "HandOver source and destination must use different end-effector " - "control parts." + first_motion = binding.endpoint("source", "motion") + second_motion = binding.endpoint("destination", "motion") + first_grasp = binding.endpoint("source", "grasp") + second_grasp = binding.endpoint("destination", "grasp") + first_arm = first_motion.require_target(JointPositionTarget) + second_arm = second_motion.require_target(JointPositionTarget) + first_hand = first_grasp.require_target(JointPositionTarget) + second_hand = second_grasp.require_target(JointPositionTarget) + if first_arm.control_part == second_arm.control_part: + raise ValueError("HandOver requires two different manipulator parts.") + if first_hand.control_part == second_hand.control_part: + raise ValueError("HandOver requires two different end-effector parts.") + + def participant( + arm: JointPositionTarget, + hand: JointPositionTarget, + grasp_endpoint: EndpointBinding, + ) -> _Participant: + return _Participant( + arm=arm, + hand=hand, + hand_open_qpos=grasp_endpoint.joint_positions( + OPEN_COMMAND, + num_envs=self.num_envs, + device=self.device, + dtype=torch.float32, + ), + hand_grasp_qpos=grasp_endpoint.joint_positions( + GRASP_COMMAND, + num_envs=self.num_envs, + device=self.device, + dtype=torch.float32, + ), ) + return _HandOverResources( - transfer_arm=transfer_arm, - receive_arm=receive_arm, - transfer_hand=transfer_hand, - receive_hand=receive_hand, - transfer_hand_open_qpos=transfer_grasp.joint_positions( - OPEN_COMMAND, - num_envs=self.num_envs, - device=self.device, - dtype=torch.float32, - ), - transfer_hand_close_qpos=transfer_grasp.joint_positions( - GRASP_COMMAND, - num_envs=self.num_envs, - device=self.device, - dtype=torch.float32, - ), - receive_hand_open_qpos=receive_grasp.joint_positions( - OPEN_COMMAND, - num_envs=self.num_envs, - device=self.device, - dtype=torch.float32, - ), - receive_hand_close_qpos=receive_grasp.joint_positions( - GRASP_COMMAND, - num_envs=self.num_envs, - device=self.device, - dtype=torch.float32, - ), + first=participant(first_arm, first_hand, first_grasp), + second=participant(second_arm, second_hand, second_grasp), ) - # ------------------------------------------------------------------ - # Public contract - # ------------------------------------------------------------------ - def _plan( self, - request: ResolvedActionRequest[GraspGoal, HandOverOptions], + request: ResolvedActionRequest[HandOverGoal, HandOverOptions], context: PlanningContext, ) -> ActionPlan: - """Plan a handover without committing the attachment transfer.""" - target = request.goal + """Plan the complete pick-up, handover, placement, and release.""" + goal = self.require_goal(request) options = request.skill_options - self._validate_pose_options(options) resources = self._resolve_resources(request) if ( request.motion_policy.strategy == "motion_gen" @@ -275,459 +275,891 @@ def _plan( raise ValueError( "Coordinated dual-arm planning is not supported by the cuRobo backend." ) - state = context - transfer_control_part = resources.transfer_arm.control_part - transfer_held_object = self._resolve_transfer_held_object( - state, transfer_control_part + if not isinstance(goal.semantics.affordance, AntipodalAffordance): + raise ValueError("HandOver requires an AntipodalAffordance.") + + object_pose = _resolve_object_pose( + goal.semantics, + context, + name="handover_object_pose", + ) + obj_longest_axis = goal.semantics.affordance.get_object_longest_axis( + object_pose, + max_points=self._SURFACE_POINT_COUNT, + ) + final_object_pose = resolve_batched_pose( + resolve_pose_goal( + goal.target_pose, + context, + name="handover_target_pose", + ), + num_envs=self.num_envs, + device=self.device, + name="handover_target_pose", + ) + first_root_pose = self._root_link_pose(resources.first.arm, context.env_ids) + second_root_pose = self._root_link_pose(resources.second.arm, context.env_ids) + first_distance = torch.linalg.vector_norm( + object_pose[:, :3, 3] - first_root_pose[:, :3, 3], dim=1 + ) + second_distance = torch.linalg.vector_norm( + object_pose[:, :3, 3] - second_root_pose[:, :3, 3], dim=1 ) - self._validate_requested_object( - target.semantics, transfer_held_object.semantics + first_is_handover = first_distance <= second_distance + + # This unified action starts before pickup. Rows where either bound arm + # already holds an object are therefore ineligible and remain at the + # observed robot state. + eligible = ~context.task.held_object_mask( + resources.first.arm.control_part + ) & ~context.task.held_object_mask(resources.second.arm.control_part) + self._report_waypoint_failure( + context, + "candidate_arms_unoccupied", + ~eligible, + "one or both candidate arms already hold an object", ) - semantics = transfer_held_object.semantics - eligible = context.task.exclusive_held_object_mask(transfer_control_part) if not eligible.any(): - logger.log_warning("HandOver requires an exclusively held source object.") return self.failed_plan( request, context, - message="Source object must be held exclusively.", + message="HandOver requires both candidate arms to start unoccupied.", + ) + + segment_lengths = self._compute_segment_lengths( + request.motion_policy.sample_count, + options, + ) + if first_is_handover.all(): + selected = self._plan_direction( + context, + request, + goal.semantics.affordance, + object_pose, + obj_longest_axis, + final_object_pose, + first_root_pose, + second_root_pose, + resources.first, + resources.second, + segment_lengths, + eligible, + ) + success = selected.success & eligible + full = selected.trajectory + elif (~first_is_handover).all(): + selected = self._plan_direction( + context, + request, + goal.semantics.affordance, + object_pose, + obj_longest_axis, + final_object_pose, + second_root_pose, + first_root_pose, + resources.second, + resources.first, + segment_lengths, + eligible, ) - transfer_object_to_eef = self._resolve_matrix( - transfer_held_object.object_to_eef, - "held_object.object_to_eef", - ) - transfer_start_qpos, receive_start_qpos = self._resolve_start_qpos( - state, - resources, - ) - assert options.middle_object_pose is not None - assert options.final_object_pose is not None - middle_object_pose = self._resolve_matrix( - options.middle_object_pose, "middle_object_pose" - ) - final_object_pose = self._resolve_matrix( - options.final_object_pose, "final_object_pose" - ) - receive_approach_direction = options.receive_approach_direction.to( - device=self.device, dtype=torch.float32 - ) - receive_approach_direction = ( - receive_approach_direction - / torch.linalg.vector_norm(receive_approach_direction) - ) - # Keep the requested object orientation consistent with the verified - # attachment and the transferring arm's current measured pose. - transfer_current_eef = self.robot.compute_fk( - qpos=transfer_start_qpos, - name=resources.transfer_arm.control_part, + success = selected.success & eligible + full = selected.trajectory + else: + first_to_second = self._plan_direction( + context, + request, + goal.semantics.affordance, + object_pose, + obj_longest_axis, + final_object_pose, + first_root_pose, + second_root_pose, + resources.first, + resources.second, + segment_lengths, + first_is_handover & eligible, + ) + second_to_first = self._plan_direction( + context, + request, + goal.semantics.affordance, + object_pose, + obj_longest_axis, + final_object_pose, + second_root_pose, + first_root_pose, + resources.second, + resources.first, + segment_lengths, + ~first_is_handover & eligible, + ) + if first_to_second.trajectory.shape != second_to_first.trajectory.shape: + raise ValueError( + "Both HandOver arm assignments must produce matching trajectory shapes." + ) + success = ( + torch.where( + first_is_handover, + first_to_second.success, + second_to_first.success, + ) + & eligible + ) + full = torch.where( + first_is_handover[:, None, None], + first_to_second.trajectory, + second_to_first.trajectory, + ) + + return self.build_plan( + request, + context, + success=success, + trajectory=TimedTrajectory.from_uniform_step( + full, + env_ids=context.env_ids, + step_dt=context.require_control_dt(), + ), + expected_effects=StateDelta(), + segment_lengths=segment_lengths, + ) + + 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_direction( + self, + context: PlanningContext, + request: ResolvedActionRequest[HandOverGoal, HandOverOptions], + affordance: AntipodalAffordance, + object_pose: torch.Tensor, + obj_longest_axis: torch.Tensor, + final_object_pose: torch.Tensor, + handover_root_pose: torch.Tensor, + receive_root_pose: torch.Tensor, + handover: _Participant, + receive: _Participant, + segment_lengths: dict[str, int], + active_mask: torch.Tensor, + ) -> _DirectionalPlan: + """Plan one concrete handover-arm to receiving-arm assignment. + + The pickup grasp is sampled on the observed object pose, whereas the + receiving grasp is sampled on the predicted pose after lift and middle + transfer. Vertical-mode approaches tilt down by 45 degrees; horizontal + mode approaches vertically downward. The two grasps use opposite ends + of the SVD-derived longest object axis. + """ + options = request.skill_options + state = context + start_qpos = state.last_qpos.to(device=self.device, dtype=torch.float32) + handover_start_qpos = start_qpos[:, list(handover.arm.joint_ids)] + handover_start_eef = self.robot.compute_fk( + qpos=handover_start_qpos, + name=handover.arm.control_part, + to_matrix=True, + ) + receive_start_qpos = start_qpos[:, list(receive.arm.joint_ids)] + receive_start_eef = self.robot.compute_fk( + qpos=receive_start_qpos, + name=receive.arm.control_part, to_matrix=True, ) - current_object_pose = torch.bmm( - transfer_current_eef, - pose_inv(transfer_object_to_eef), + + vertical_mode = ( + torch.abs(obj_longest_axis[:, 2]) >= self._VERTICAL_MODE_MIN_ABS_Z + ) + vertical_down = torch.tensor( + [0.0, 0.0, -1.0], + dtype=torch.float32, + device=self.device, + ).expand(self.num_envs, -1) + + # In vertical mode, point from the pickup TCP toward the observed + # object horizontally and tilt down by 45 degrees. In horizontal mode, + # approach vertically downward. Only vertical rows require nonzero + # horizontal TCP-to-object separation. + handover_diagonal, handover_diagonal_valid = ( + self._downward_diagonal_approach_direction( + handover_start_eef[:, :3, 3], object_pose[:, :3, 3] + ) + ) + handover_direction = torch.where( + vertical_mode[:, None], handover_diagonal, vertical_down ) - middle_object_pose[:, :3, :3] = current_object_pose[:, :3, :3] - final_object_pose[:, :3, :3] = current_object_pose[:, :3, :3] + handover_direction_valid = ~vertical_mode | handover_diagonal_valid - # 2.1 - EEF target that keeps the object at the handover pose. - transfer_middle_eef = torch.bmm(middle_object_pose, transfer_object_to_eef) + # SVD axes have arbitrary sign. Selecting the sign whose projected end + # points toward the pickup TCP makes the physical choice sign-invariant; + # the receiving arm always takes the opposite projected end. + handover_is_positive_part = ( + torch.sum( + (handover_start_eef[:, :3, 3] - object_pose[:, :3, 3]) + * obj_longest_axis, + dim=1, + ) + >= 0.0 + ) + handover_grasp, handover_grasp_success = self._resolve_grasp( + affordance, + object_pose, + handover_direction, + obj_longest_axis=obj_longest_axis, + is_positive_part=handover_is_positive_part, + ) + handover_grasp = self._find_symmetric_nearest_xpos( + handover_grasp, handover_start_eef + ) + handover_pre_grasp = translate_pose_world( + handover_grasp, + -handover_direction * options.pre_grasp_distance, + ) + handover_object_to_eef = torch.bmm(pose_inv(object_pose), handover_grasp) - # 2.2 - receiving grasp on the requested object part at the handover pose. - receive_grasp_xpos, grasp_success = self._resolve_receive_grasp( - semantics, + lifted_object_pose = object_pose.clone() + lifted_object_pose[:, 2, 3] += options.lift_height + middle_object_pose = self._middle_object_pose( + lifted_object_pose, + handover_root_pose, + receive_root_pose, + ) + handover_lift_eef = torch.bmm( + lifted_object_pose, + handover_object_to_eef, + ) + handover_middle_eef = torch.bmm( middle_object_pose, - options.receive_pick_object_part, - receive_approach_direction, + handover_object_to_eef, ) - success_mask = normalize_success_mask( - grasp_success, - num_envs=self.num_envs, - device=self.device, - name="Receiving-grasp success", + # From the grasp waypoint through lift and transfer, only translation + # may change. Pinning the rotations explicitly also avoids numerical + # drift from the object/EFF transform multiplications. + handover_lift_eef[:, :3, :3] = handover_grasp[:, :3, :3] + handover_middle_eef[:, :3, :3] = handover_grasp[:, :3, :3] + + # Apply the same mode to receiving. The object rotation is unchanged by + # lift and middle transfer, so its world-space longest axis is unchanged. + receive_diagonal, receive_diagonal_valid = ( + self._downward_diagonal_approach_direction( + receive_start_eef[:, :3, 3], middle_object_pose[:, :3, 3] + ) ) - success_mask &= eligible - if not success_mask.any(): - logger.log_warning("HandOver failed to resolve a receiving grasp pose.") - return self.failed_plan(request, context, message="No receiving grasp.") - receive_object_to_eef = torch.bmm( - pose_inv(middle_object_pose), receive_grasp_xpos - ) - receive_grasp_z = receive_grasp_xpos[..., :3, 2] - receive_pre_grasp_eef = translate_pose_world( - receive_grasp_xpos, - -receive_grasp_z * options.pre_grasp_distance, - ) - # 2.4 - receiving arm delivers the object to the final pose. - receive_final_eef = torch.bmm(final_object_pose, receive_object_to_eef) - # 2.3 - transferring arm retreats upward after releasing. - transfer_retreat_eef = translate_pose_world( - transfer_middle_eef, - torch.tensor( - [0.0, 0.0, options.lift_height], - dtype=torch.float32, - device=self.device, - ), + receive_direction = torch.where( + vertical_mode[:, None], receive_diagonal, vertical_down + ) + receive_direction_valid = ~vertical_mode | receive_diagonal_valid + receive_grasp, receive_grasp_success = self._resolve_grasp( + affordance, + middle_object_pose, + receive_direction, + obj_longest_axis=obj_longest_axis, + is_positive_part=~handover_is_positive_part, + ) + receive_grasp = self._find_symmetric_nearest_xpos( + receive_grasp, receive_start_eef ) - segments = self._compute_segment_lengths( - request.motion_policy.sample_count, options + receive_pre_grasp = translate_pose_world( + receive_grasp, + -receive_direction * options.pre_grasp_distance, + ) + receive_object_to_eef = torch.bmm( + pose_inv(middle_object_pose), + receive_grasp, + ) + placed_object_pose = final_object_pose.clone() + placed_object_pose[:, :3, :3] = middle_object_pose[:, :3, :3] + above_object_pose = placed_object_pose.clone() + # Move to the target's horizontal coordinates while preserving the + # middle handover height exactly. The following target performs the + # only vertical motion and reaches the requested final object pose. + above_object_pose[:, 2, 3] = middle_object_pose[:, 2, 3] + lowering_direction_valid = ( + above_object_pose[:, 2, 3] - placed_object_pose[:, 2, 3] > 1.0e-6 + ) + receive_above_eef = torch.bmm(above_object_pose, receive_object_to_eef) + receive_final_eef = torch.bmm(placed_object_pose, receive_object_to_eef) + # Likewise, receiving-grasp through final lowering reuses the same EEF + # rotation and changes translation only. + receive_above_eef[:, :3, :3] = receive_grasp[:, :3, :3] + receive_final_eef[:, :3, :3] = receive_grasp[:, :3, :3] + self._report_waypoint_failure( + context, + "pickup_grasp", + active_mask & ~handover_grasp_success, + "no finite grasp candidate on the pickup-side object end for arm " + f"{handover.arm.control_part!r}", + ) + self._report_waypoint_failure( + context, + "pickup_approach_direction", + active_mask & ~handover_direction_valid, + "handover TCP and observed object position have no horizontal " + f"separation for arm {handover.arm.control_part!r}", + ) + self._report_waypoint_failure( + context, + "receive_approach_direction", + active_mask & ~receive_direction_valid, + "receive TCP and predicted object position have no horizontal " + f"separation for arm {receive.arm.control_part!r}", + ) + self._report_waypoint_failure( + context, + "receive_grasp", + active_mask & ~receive_grasp_success, + "no finite grasp candidate on the opposite object end for arm " + f"{receive.arm.control_part!r}", + ) + self._report_waypoint_failure( + context, + "target_final", + active_mask & ~lowering_direction_valid, + "final target is not below the horizontal-transfer height", ) - segment_success, transfer_move_traj = plan_named_arm_trajectory( + success = ( + handover_direction_valid + & handover_grasp_success + & receive_direction_valid + & receive_grasp_success + & lowering_direction_valid + ) + pickup_approach_targets = torch.stack( + [handover_pre_grasp, handover_grasp], dim=1 + ) + phase_success, pickup_approach = plan_named_arm_trajectory( self.motion_generator, - resources.transfer_arm.control_part, - transfer_start_qpos, - transfer_middle_eef.unsqueeze(1), - segments["transfer"], + handover.arm.control_part, + handover_start_qpos, + pickup_approach_targets, + segment_lengths["pickup_approach"], request.motion_policy, context.control_dt, ) - success_mask &= normalize_success_mask( - segment_success, + pickup_approach_success = normalize_success_mask( + phase_success, num_envs=self.num_envs, device=self.device, - name="Transfer-move success", + name="HandOver pickup-approach success", ) - if not success_mask.any(): - logger.log_warning("HandOver failed to plan the transfer move.") - return self.failed_plan(request, context, message="Transfer move failed.") + self._report_phase_failure( + context, + phase_name="pickup_approach", + waypoint_names=("pickup_pre_grasp", "pickup_grasp"), + target_poses=pickup_approach_targets, + start_qpos=handover_start_qpos, + arm=handover.arm, + failed_mask=active_mask & ~pickup_approach_success, + ) + success &= pickup_approach_success + handover_grasp_qpos = pickup_approach[:, -1] - segment_success, receive_approach_traj = plan_named_arm_trajectory( + pickup_transport_targets = torch.stack( + [handover_lift_eef, handover_middle_eef], dim=1 + ) + phase_success, pickup_transport = plan_named_arm_trajectory( self.motion_generator, - resources.receive_arm.control_part, - receive_start_qpos, - torch.stack([receive_pre_grasp_eef, receive_grasp_xpos], dim=1), - segments["approach"], + handover.arm.control_part, + handover_grasp_qpos, + pickup_transport_targets, + segment_lengths["pickup_transport"], request.motion_policy, context.control_dt, ) - success_mask &= normalize_success_mask( - segment_success, + pickup_transport_success = normalize_success_mask( + phase_success, num_envs=self.num_envs, device=self.device, - name="Receiving-approach success", + name="HandOver pickup-transport success", ) - if not success_mask.any(): - logger.log_warning("HandOver failed to plan the receiving approach.") - return self.failed_plan( - request, context, message="Receiving approach failed." - ) - - transfer_hold_qpos = transfer_move_traj[:, -1] - receive_grasp_qpos = receive_approach_traj[:, -1] + self._report_phase_failure( + context, + phase_name="pickup_transport", + waypoint_names=("pickup_lift", "handover_middle"), + target_poses=pickup_transport_targets, + start_qpos=handover_grasp_qpos, + arm=handover.arm, + failed_mask=active_mask & ~pickup_transport_success, + ) + success &= pickup_transport_success + handover_middle_qpos = pickup_transport[:, -1] - segment_success, transfer_retreat_traj = plan_named_arm_trajectory( + receive_approach_targets = torch.stack( + [receive_pre_grasp, receive_grasp], dim=1 + ) + phase_success, receive_approach = plan_named_arm_trajectory( self.motion_generator, - resources.transfer_arm.control_part, - transfer_hold_qpos, - transfer_retreat_eef.unsqueeze(1), - segments["deliver"], + receive.arm.control_part, + receive_start_qpos, + receive_approach_targets, + segment_lengths["receive_approach"], request.motion_policy, context.control_dt, ) - success_mask &= normalize_success_mask( - segment_success, + receive_approach_success = normalize_success_mask( + phase_success, num_envs=self.num_envs, device=self.device, - name="Transfer-retreat success", + name="HandOver receive-approach success", ) - if not success_mask.any(): - logger.log_warning("HandOver failed to plan the transfer retreat.") - return self.failed_plan( - request, context, message="Transfer retreat failed." - ) + self._report_phase_failure( + context, + phase_name="receive_approach", + waypoint_names=("receive_pre_grasp", "receive_grasp"), + target_poses=receive_approach_targets, + start_qpos=receive_start_qpos, + arm=receive.arm, + failed_mask=active_mask & ~receive_approach_success, + ) + success &= receive_approach_success + receive_grasp_qpos = receive_approach[:, -1] - segment_success, receive_deliver_traj = plan_named_arm_trajectory( + placement_targets = torch.stack([receive_above_eef, receive_final_eef], dim=1) + phase_success, receive_place = plan_named_arm_trajectory( self.motion_generator, - resources.receive_arm.control_part, + receive.arm.control_part, receive_grasp_qpos, - receive_final_eef.unsqueeze(1), - segments["deliver"], + placement_targets, + segment_lengths["place"], request.motion_policy, context.control_dt, ) - success_mask &= normalize_success_mask( - segment_success, + placement_success = normalize_success_mask( + phase_success, num_envs=self.num_envs, device=self.device, - name="Receiving-delivery success", + name="HandOver placement success", ) - if not success_mask.any(): - logger.log_warning("HandOver failed to plan the receiving delivery.") - return self.failed_plan( - request, context, message="Receiving delivery failed." - ) + self._report_phase_failure( + context, + phase_name="place", + waypoint_names=("target_above", "target_final"), + target_poses=placement_targets, + start_qpos=receive_grasp_qpos, + arm=receive.arm, + failed_mask=active_mask & ~placement_success, + ) + success &= placement_success + receive_final_qpos = receive_place[:, -1] - segment_trajectories: list[torch.Tensor] = [] - # 2.1 transfer: transferring arm carries the object to the middle pose. - segment_trajectories.append( + segments = [ self._assemble_segment( state, - transfer_move_traj, - self._repeat_qpos(receive_start_qpos, segments["transfer"]), - self._repeat_qpos( - resources.transfer_hand_close_qpos, segments["transfer"] + pickup_approach, + repeat_qpos(receive_start_qpos, segment_lengths["pickup_approach"]), + repeat_qpos( + handover.hand_open_qpos, segment_lengths["pickup_approach"] ), - self._repeat_qpos( - resources.receive_hand_open_qpos, segments["transfer"] + repeat_qpos(receive.hand_open_qpos, segment_lengths["pickup_approach"]), + handover, + receive, + ), + self._assemble_segment( + state, + repeat_qpos(handover_grasp_qpos, segment_lengths["pickup_close"]), + repeat_qpos(receive_start_qpos, segment_lengths["pickup_close"]), + interpolate_hand_qpos( + handover.hand_open_qpos, + handover.hand_grasp_qpos, + n_waypoints=segment_lengths["pickup_close"], ), - resources=resources, - ) - ) - # 2.2 approach: receiving arm moves to the grasp pose; transferring arm holds. - segment_trajectories.append( + repeat_qpos(receive.hand_open_qpos, segment_lengths["pickup_close"]), + handover, + receive, + ), self._assemble_segment( state, - self._repeat_qpos(transfer_hold_qpos, segments["approach"]), - receive_approach_traj, - self._repeat_qpos( - resources.transfer_hand_close_qpos, segments["approach"] + pickup_transport, + repeat_qpos(receive_start_qpos, segment_lengths["pickup_transport"]), + repeat_qpos( + handover.hand_grasp_qpos, + segment_lengths["pickup_transport"], ), - self._repeat_qpos( - resources.receive_hand_open_qpos, segments["approach"] + repeat_qpos( + receive.hand_open_qpos, segment_lengths["pickup_transport"] ), - resources=resources, - ) - ) - # 2.2 close: receiving hand closes; transferring arm keeps holding. - segment_trajectories.append( + handover, + receive, + ), self._assemble_segment( state, - self._repeat_qpos(transfer_hold_qpos, segments["close"]), - self._repeat_qpos(receive_grasp_qpos, segments["close"]), - self._repeat_qpos( - resources.transfer_hand_close_qpos, segments["close"] + repeat_qpos(handover_middle_qpos, segment_lengths["receive_approach"]), + receive_approach, + repeat_qpos( + handover.hand_grasp_qpos, + segment_lengths["receive_approach"], ), + repeat_qpos( + receive.hand_open_qpos, segment_lengths["receive_approach"] + ), + handover, + receive, + ), + self._assemble_segment( + state, + repeat_qpos(handover_middle_qpos, segment_lengths["receive_close"]), + repeat_qpos(receive_grasp_qpos, segment_lengths["receive_close"]), + repeat_qpos(handover.hand_grasp_qpos, segment_lengths["receive_close"]), interpolate_hand_qpos( - resources.receive_hand_open_qpos, - resources.receive_hand_close_qpos, - n_waypoints=segments["close"], + receive.hand_open_qpos, + receive.hand_grasp_qpos, + n_waypoints=segment_lengths["receive_close"], ), - resources=resources, - ) - ) - if segments["hold"] > 0: - segment_trajectories.append( - self._assemble_segment( - state, - self._repeat_qpos(transfer_hold_qpos, segments["hold"]), - self._repeat_qpos(receive_grasp_qpos, segments["hold"]), - self._repeat_qpos( - resources.transfer_hand_close_qpos, segments["hold"] - ), - self._repeat_qpos( - resources.receive_hand_close_qpos, segments["hold"] - ), - resources=resources, - ) - ) - # 2.3 release: transferring hand opens; receiving arm keeps holding. - segment_trajectories.append( + handover, + receive, + ), self._assemble_segment( state, - self._repeat_qpos(transfer_hold_qpos, segments["release"]), - self._repeat_qpos(receive_grasp_qpos, segments["release"]), + repeat_qpos(handover_middle_qpos, segment_lengths["handover_release"]), + repeat_qpos(receive_grasp_qpos, segment_lengths["handover_release"]), interpolate_hand_qpos( - resources.transfer_hand_close_qpos, - resources.transfer_hand_open_qpos, - n_waypoints=segments["release"], + handover.hand_grasp_qpos, + handover.hand_open_qpos, + n_waypoints=segment_lengths["handover_release"], ), - self._repeat_qpos( - resources.receive_hand_close_qpos, segments["release"] + repeat_qpos( + receive.hand_grasp_qpos, + segment_lengths["handover_release"], ), - resources=resources, - ) - ) - # 2.4 deliver: receiving arm carries the object away; transferring arm retreats. - segment_trajectories.append( + handover, + receive, + ), + self._assemble_segment( + state, + repeat_qpos(handover_middle_qpos, segment_lengths["place"]), + receive_place, + repeat_qpos(handover.hand_open_qpos, segment_lengths["place"]), + repeat_qpos(receive.hand_grasp_qpos, segment_lengths["place"]), + handover, + receive, + ), self._assemble_segment( state, - transfer_retreat_traj, - receive_deliver_traj, - self._repeat_qpos( - resources.transfer_hand_open_qpos, segments["deliver"] + repeat_qpos(handover_middle_qpos, segment_lengths["receive_release"]), + repeat_qpos(receive_final_qpos, segment_lengths["receive_release"]), + repeat_qpos( + handover.hand_open_qpos, segment_lengths["receive_release"] ), - self._repeat_qpos( - resources.receive_hand_close_qpos, segments["deliver"] + interpolate_hand_qpos( + receive.hand_grasp_qpos, + receive.hand_open_qpos, + n_waypoints=segment_lengths["receive_release"], ), - resources=resources, - ) - ) - full = torch.cat(segment_trajectories, dim=1) - segment_names = ["transfer", "approach", "close"] - if segments["hold"] > 0: - segment_names.append("hold") - segment_names.extend(("release", "deliver")) - segment_lengths = { - name: trajectory.shape[1] - for name, trajectory in zip( - segment_names, segment_trajectories, strict=True - ) + handover, + receive, + ), + ] + trajectory = torch.cat(segments, dim=1) + actual_lengths = { + name: segment.shape[1] + for name, segment in zip(segment_lengths, segments, strict=True) } - held_object = HeldObjectState( - semantics=semantics, - object_to_eef=receive_object_to_eef, - grasp_xpos=receive_grasp_xpos, + if actual_lengths != segment_lengths: + logger.log_warning( + "HandOver planner returned segment lengths that differ from the request." + ) + return _DirectionalPlan( + success=success, + trajectory=trajectory, + segment_lengths=actual_lengths, ) - return self.build_plan( - request, - context, - success=success_mask, - trajectory=TimedTrajectory.from_uniform_step( - full, - env_ids=context.env_ids, - step_dt=context.require_control_dt(), - ), - expected_effects=StateDelta( - held_object_updates={ - resources.transfer_arm.control_part: None, - resources.receive_arm.control_part: held_object, - } - ), - segment_lengths=segment_lengths, + + @staticmethod + def _report_waypoint_failure( + context: PlanningContext, + waypoint_name: str, + failed_mask: torch.Tensor, + reason: str, + ) -> None: + """Log one semantic waypoint failure with affected environment IDs.""" + if not failed_mask.any(): + return + env_ids = context.env_ids.to(failed_mask.device)[failed_mask] + logger.log_warning( + f"HandOver waypoint '{waypoint_name}' failed for " + f"env_ids={env_ids.detach().cpu().tolist()}: {reason}." ) - # ------------------------------------------------------------------ - # Resolution helpers - # ------------------------------------------------------------------ + def _report_phase_failure( + self, + context: PlanningContext, + *, + phase_name: str, + waypoint_names: tuple[str, ...], + target_poses: torch.Tensor, + start_qpos: torch.Tensor, + arm: JointPositionTarget, + failed_mask: torch.Tensor, + ) -> None: + """Identify failed waypoint IK, or report a path/collision failure.""" + if not failed_mask.any(): + return + identified = torch.zeros_like(failed_mask) + joint_seed = start_qpos + try: + for waypoint_index, waypoint_name in enumerate(waypoint_names): + ik_success, waypoint_qpos = self.robot.compute_ik( + pose=target_poses[:, waypoint_index], + name=arm.control_part, + joint_seed=joint_seed, + ) + ik_success = normalize_success_mask( + ik_success, + num_envs=self.num_envs, + device=self.device, + name=f"HandOver diagnostic IK for {waypoint_name}", + ) + waypoint_qpos = torch.as_tensor( + waypoint_qpos, + dtype=joint_seed.dtype, + device=self.device, + ) + if waypoint_qpos.shape != joint_seed.shape: + raise ValueError( + "diagnostic IK returned qpos with shape " + f"{tuple(waypoint_qpos.shape)}, expected " + f"{tuple(joint_seed.shape)}" + ) + waypoint_failed = failed_mask & ~ik_success + self._report_waypoint_failure( + context, + waypoint_name, + waypoint_failed, + f"IK failed for arm {arm.control_part!r}", + ) + identified |= waypoint_failed + joint_seed = torch.where(ik_success[:, None], waypoint_qpos, joint_seed) + except Exception as exc: + env_ids = context.env_ids.to(failed_mask.device)[failed_mask] + logger.log_warning( + f"HandOver phase '{phase_name}' failed for " + f"arm {arm.control_part!r}, " + f"env_ids={env_ids.detach().cpu().tolist()}, " + "but waypoint IK " + f"diagnostics could not run: {exc}." + ) + return - @staticmethod - def _validate_pose_options(options: HandOverOptions) -> None: - for name in ("middle_object_pose", "final_object_pose"): - if getattr(options, name) is None: - raise ValueError(f"{name} must be specified in HandOverOptions") + unresolved = failed_mask & ~identified + if unresolved.any(): + env_ids = context.env_ids.to(unresolved.device)[unresolved] + logger.log_warning( + f"HandOver phase '{phase_name}' failed between waypoints " + f"{list(waypoint_names)} for " + f"arm {arm.control_part!r}, " + f"env_ids={env_ids.detach().cpu().tolist()}; " + "individual " + "waypoint IK succeeded, so the likely cause is path or " + "collision planning." + ) - def _resolve_matrix(self, matrix: torch.Tensor, name: str) -> torch.Tensor: + def _root_link_pose( + self, + arm: JointPositionTarget, + env_ids: torch.Tensor, + ) -> torch.Tensor: + """Read the root-link pose configured for ``arm``.""" + robot_cfg = getattr(self.robot, "cfg", None) + solver_cfg = getattr(robot_cfg, "solver_cfg", None) + if not isinstance(solver_cfg, Mapping) or arm.control_part not in solver_cfg: + raise ValueError( + "HandOver requires " f"solver_cfg[{arm.control_part!r}].root_link_name." + ) + root_link_name = getattr(solver_cfg[arm.control_part], "root_link_name", None) + if not isinstance(root_link_name, str) or not root_link_name: + raise ValueError( + "HandOver requires a root_link_name for arm " f"{arm.control_part!r}." + ) + pose = self.robot.get_link_pose( + link_name=root_link_name, + env_ids=env_ids.tolist(), + to_matrix=True, + ) return resolve_batched_pose( - matrix, + pose, num_envs=self.num_envs, device=self.device, - name=name, + name=f"{arm.control_part} root-link pose", ) - def _resolve_transfer_held_object( + def _resolve_grasp( self, - state: PlanningContext, - transfer_control_part: str, - ) -> HeldObjectState: - held = state.get_held_object(transfer_control_part) - if held is None: + affordance: AntipodalAffordance, + object_pose: torch.Tensor, + approach_direction: torch.Tensor, + *, + obj_longest_axis: torch.Tensor, + is_positive_part: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Select the lowest-cost grasp on one projected end of the object.""" + if object_pose.shape != (self.num_envs, 4, 4): raise ValueError( - "HandOver requires an object held by transfer control part " - f"{transfer_control_part!r} (run PickUp first)." + "HandOver grasp object_pose must have shape " + f"({self.num_envs}, 4, 4)." ) - return held - - @staticmethod - def _validate_requested_object( - requested: ObjectSemantics, - held: ObjectSemantics, - ) -> None: - """Reject a request that names a different grounded object.""" - if not _same_object_identity(requested, held): + if approach_direction.shape != (self.num_envs, 3): raise ValueError( - "HandOver goal semantics must identify the object held by the " - "source control part." + "HandOver grasp approach_direction must have shape " + f"({self.num_envs}, 3)." + ) + if obj_longest_axis.shape != (self.num_envs, 3): + raise ValueError( + f"HandOver obj_longest_axis must have shape ({self.num_envs}, 3)." + ) + if is_positive_part.dtype != torch.bool or is_positive_part.shape != ( + self.num_envs, + ): + raise ValueError( + "HandOver is_positive_part must be a bool tensor with shape " + f"({self.num_envs},)." ) - def _resolve_receive_grasp( - self, - semantics: ObjectSemantics, - object_pose: torch.Tensor, - object_part: str, - approach_direction: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Select the lowest-cost receiving grasp on ``object_part`` at ``object_pose``.""" - grasp_poses_result = semantics.affordance.get_valid_grasp_poses( - obj_poses=object_pose, - approach_direction=approach_direction, - object_part=object_part, - ) - num_envs = object_pose.shape[0] - grasp_xpos = ( - torch.eye(4, device=self.device, dtype=torch.float32) - .unsqueeze(0) - .repeat(num_envs, 1, 1) - ) - is_success = torch.ones(num_envs, dtype=torch.bool, device=self.device) - for i in range(num_envs): - poses, costs = grasp_poses_result[i] - poses = poses.to(device=self.device, dtype=torch.float32) + # AntipodalAffordance iterates over object poses but its underlying + # grasp generator accepts one direction vector with shape (3,) per + # object. Sample each environment separately so the automatically + # derived per-environment directions are not forwarded as (B, 3). + sampled: list[tuple[torch.Tensor, torch.Tensor]] = [] + for env_index in range(self.num_envs): + env_sampled = affordance.get_valid_grasp_poses( + obj_poses=object_pose[env_index : env_index + 1], + approach_direction=approach_direction[env_index], + obj_longest_axis=obj_longest_axis[env_index], + is_positive_part=bool(is_positive_part[env_index].item()), + ) + if len(env_sampled) != 1: + raise ValueError( + "HandOver expected exactly one grasp-sampling result per " + "environment." + ) + sampled.append(env_sampled[0]) + poses = torch.eye( + 4, + dtype=torch.float32, + device=self.device, + ).repeat(self.num_envs, 1, 1) + success = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + for env_index, (candidates, costs) in enumerate(sampled): + candidates = candidates.to(device=self.device, dtype=torch.float32) costs = costs.to(device=self.device, dtype=torch.float32) - if poses.shape[0] == 0: - is_success[i] = False + if candidates.shape[0] == 0 or not torch.isfinite(costs).any(): continue - best_idx = torch.argmin(costs) - if not torch.isfinite(costs[best_idx]): - is_success[i] = False - grasp_xpos[i] = poses[best_idx] - return grasp_xpos, is_success + finite_costs = torch.where( + torch.isfinite(costs), + costs, + torch.full_like(costs, torch.inf), + ) + poses[env_index] = candidates[torch.argmin(finite_costs)] + success[env_index] = True + return poses, success - def _resolve_start_qpos( - self, - state: PlanningContext, - resources: _HandOverResources, + @staticmethod + def _downward_diagonal_approach_direction( + start_position: torch.Tensor, + target_position: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - if state.last_qpos.shape != (self.num_envs, self.robot_dof): - raise ValueError( - f"PlanningContext.last_qpos must have shape " - f"({self.num_envs}, {self.robot_dof}), but got " - f"{state.last_qpos.shape}" - ) - start_qpos = state.last_qpos.to(device=self.device, dtype=torch.float32) - return ( - start_qpos[:, list(resources.transfer_arm.joint_ids)], - start_qpos[:, list(resources.receive_arm.joint_ids)], + """Return TCP-to-target horizontal directions tilted down by 45 degrees. + + The direction is valid only when the TCP and target have nonzero + horizontal separation; callers report the corresponding semantic + approach waypoint when that construction is undefined. + """ + horizontal_delta = target_position[:, :2] - start_position[:, :2] + horizontal_norm = torch.linalg.vector_norm(horizontal_delta, dim=1) + valid = horizontal_norm > 1.0e-6 + horizontal_unit = horizontal_delta / horizontal_norm.clamp_min( + 1.0e-6 + ).unsqueeze(1) + component = math.sqrt(0.5) + direction = torch.zeros( + (start_position.shape[0], 3), + dtype=start_position.dtype, + device=start_position.device, ) + direction[:, :2] = horizontal_unit * component + direction[:, 2] = -component + return direction, valid + @staticmethod + def _middle_object_pose( + lifted_object_pose: torch.Tensor, + handover_root_pose: torch.Tensor, + receive_root_pose: torch.Tensor, + ) -> torch.Tensor: + """Move only the dominant root-separation coordinate to its midpoint.""" + handover_root_position = handover_root_pose[:, :3, 3] + receive_root_position = receive_root_pose[:, :3, 3] + dominant_axis = torch.argmax( + torch.abs(handover_root_position - receive_root_position), dim=1 + ) + root_midpoint = 0.5 * (handover_root_position + receive_root_position) + middle = lifted_object_pose.clone() + middle_position = middle[:, :3, 3] + selected_midpoint = root_midpoint.gather(1, dominant_axis[:, None]) + middle_position.scatter_(1, dominant_axis[:, None], selected_midpoint) + return middle + + @staticmethod def _compute_segment_lengths( - self, sample_count: int, options: HandOverOptions + sample_count: int, + options: HandOverOptions, ) -> dict[str, int]: - """Split the invocation sample budget across handover segments.""" - n_close = max(2, options.hand_interp_steps) - n_release = max(2, options.hand_interp_steps) - n_deliver = max(2, options.retreat_steps) - n_hold = max(0, options.hold_steps) - reserved = n_close + n_release + n_deliver + n_hold - n_transfer = max(2, (sample_count - reserved) // 2) - n_approach = sample_count - reserved - n_transfer - if n_approach < 2: + """Split the sample budget across four arm and four hand phases.""" + hand_count = options.hand_interp_steps + motion_budget = sample_count - 4 * hand_count + if motion_budget < 8: raise ValueError( - "Not enough waypoints for handover. Increase sample_count or " - "decrease hand_interp_steps/hold_steps/retreat_steps." + "Not enough HandOver waypoints. Increase sample_count or decrease " + "hand_interp_steps." ) + motion_counts = [motion_budget // 4] * 4 + for index in range(motion_budget % 4): + motion_counts[index] += 1 return { - "transfer": n_transfer, - "approach": n_approach, - "close": n_close, - "hold": n_hold, - "release": n_release, - "deliver": n_deliver, + "pickup_approach": motion_counts[0], + "pickup_close": hand_count, + "pickup_transport": motion_counts[1], + "receive_approach": motion_counts[2], + "receive_close": hand_count, + "handover_release": hand_count, + "place": motion_counts[3], + "receive_release": hand_count, } - # ------------------------------------------------------------------ - # Planning / assembly helpers - # ------------------------------------------------------------------ - + @staticmethod def _assemble_segment( - self, state: PlanningContext, - transfer_arm_traj: torch.Tensor, - receive_arm_traj: torch.Tensor, - transfer_hand_traj: torch.Tensor, - receive_hand_traj: torch.Tensor, - *, - resources: _HandOverResources, + handover_arm_trajectory: torch.Tensor, + receive_arm_trajectory: torch.Tensor, + handover_hand_trajectory: torch.Tensor, + receive_hand_trajectory: torch.Tensor, + handover: _Participant, + receive: _Participant, ) -> torch.Tensor: + """Embed both arm and hand paths in the full robot joint order.""" return assemble_full_robot_trajectory( state.last_qpos, ( - (resources.transfer_arm.joint_ids, transfer_arm_traj), - (resources.receive_arm.joint_ids, receive_arm_traj), - (resources.transfer_hand.joint_ids, transfer_hand_traj), - (resources.receive_hand.joint_ids, receive_hand_traj), + (handover.arm.joint_ids, handover_arm_trajectory), + (receive.arm.joint_ids, receive_arm_trajectory), + (handover.hand.joint_ids, handover_hand_trajectory), + (receive.hand.joint_ids, receive_hand_trajectory), ), ) -__all__ = ["HandOver", "HandOverOptions"] +__all__ = ["HandOver", "HandOverGoal", "HandOverOptions"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index 13437db44..8ad98e7ab 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -132,8 +132,10 @@ class PickUpOptions(ActionOptions): def __post_init__(self) -> None: if self.hand_interp_steps < 1: raise ValueError("hand_interp_steps must be at least 1.") - if not isinstance(self.pick_object_part, str) or not self.pick_object_part: - raise ValueError("pick_object_part must be a non-empty string.") + if self.pick_object_part not in {"center", "top", "bottom"}: + raise ValueError( + "pick_object_part must be one of 'center', 'top', or 'bottom'." + ) if self.lift_height < 0.0: raise ValueError("lift_height must be non-negative.") if self.pre_grasp_distance < 0.0: @@ -428,10 +430,20 @@ def _resolve_grasp_pose( options: PickUpOptions, approach_direction: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: + obj_longest_axis = None + is_positive_part = True + if options.pick_object_part != "center": + obj_longest_axis = torch.tensor( + [0.0, 0.0, 1.0], + dtype=torch.float32, + device=self.device, + ) + is_positive_part = options.pick_object_part == "top" grasp_poses_result = semantics.affordance.get_valid_grasp_poses( obj_poses=object_pose, approach_direction=approach_direction, - object_part=options.pick_object_part, + obj_longest_axis=obj_longest_axis, + is_positive_part=is_positive_part, ) num_envs = object_pose.shape[0] n_max_pose = max(r[0].shape[0] for r in grasp_poses_result) diff --git a/embodichain/lab/sim/atomic_actions/primitives/pour.py b/embodichain/lab/sim/atomic_actions/primitives/pour.py new file mode 100644 index 000000000..9d20ab1c0 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/pour.py @@ -0,0 +1,213 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Pour 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 + +from embodichain.lab.sim.atomic_actions.affordance import AxisAlignAffordance +from embodichain.lab.sim.atomic_actions.bindings import JointPositionTarget +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.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan, TimedTrajectory +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, + FORWARD_KINEMATICS_CAPABILITY, + SkillBindingContract, +) +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) +class PourGoal: + """Rotate the object currently held by the bound manipulator.""" + + goal_kind: ClassVar[str] = "pour" + + +@dataclass(frozen=True, slots=True, eq=False) +class PourOptions(ActionOptions): + """Per-invocation pouring behavior.""" + + rotate_angle: float = math.pi / 4.0 + """Signed rotation about the held object's local internal axis, in radians.""" + + def __post_init__(self) -> None: + if not math.isfinite(self.rotate_angle): + raise ValueError("rotate_angle must be finite.") + + +class Pour(AtomicAction[PourGoal, PourOptions]): + """Rotate and return an exclusively held object about its internal axis.""" + + skill_id: ClassVar[str] = "pour" + GoalType: ClassVar[type] = PourGoal + OptionsType: ClassVar[type] = PourOptions + 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( + "primary", + motion_capabilities=frozenset( + { + CARTESIAN_POSE_CAPABILITY, + FORWARD_KINEMATICS_CAPABILITY, + } + ), + grasp_commands={GRASP_COMMAND: JointPositionCommand}, + ), + ), + ) + + def __init__(self, default_options: PourOptions | None = None) -> None: + super().__init__(default_options) + + def _plan( + self, + request: ResolvedActionRequest[PourGoal, PourOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan a held-object rotation followed by the inverse rotation.""" + self.require_goal(request) + options = request.skill_options + motion_endpoint = request.binding.endpoint("primary", "motion") + grasp_endpoint = request.binding.endpoint("primary", "grasp") + manipulator = motion_endpoint.require_target(JointPositionTarget) + end_effector = grasp_endpoint.require_target(JointPositionTarget) + control_part = manipulator.control_part + arm_joint_ids = list(manipulator.joint_ids) + hand_joint_ids = list(end_effector.joint_ids) + + held_object = context.get_held_object(control_part) + if held_object is None: + raise ValueError( + "Pour requires an object held by control part " + f"{control_part!r} - run PickUp first." + ) + affordance = held_object.semantics.affordance + if not isinstance(affordance, AxisAlignAffordance): + raise ValueError( + "Pour requires the held object to use an AxisAlignAffordance." + ) + eligible = context.task.exclusive_held_object_mask(control_part) + if not eligible.any(): + return self.failed_plan( + request, + context, + message="Held object is not exclusive to the control part.", + ) + + start_arm_qpos = arm_qpos_from_state(context, arm_joint_ids) + current_eef_pose = self.robot.compute_fk( + qpos=start_arm_qpos, + name=control_part, + to_matrix=True, + ) + object_to_eef = held_object.object_to_eef.to( + device=self.device, + dtype=torch.float32, + ) + current_object_pose = torch.bmm(current_eef_pose, pose_inv(object_to_eef)) + + internal_axis = affordance.internal_axis.to( + device=self.device, + dtype=torch.float32, + ) + internal_axis = internal_axis / torch.linalg.vector_norm(internal_axis) + world_axis = torch.matmul(current_object_pose[:, :3, :3], internal_axis) + rotation_delta = axis_angle_to_rotation_matrix( + world_axis * options.rotate_angle + ) + target_object_pose = current_object_pose.clone() + target_object_pose[:, :3, :3] = torch.bmm( + rotation_delta, + current_object_pose[:, :3, :3], + ) + poured_eef_pose = torch.bmm(target_object_pose, object_to_eef) + + # The object remains rigidly attached to the EEF throughout Pour. The + # first target performs the requested rotation about the object's + # internal axis; the second target returns to the EEF pose observed by + # FK at start_qpos, which is exactly the inverse rotation by the same + # signed angle. Supplying both targets in one call lets planners such + # as cuRobo chain the two collision-aware legs without another public + # MotionGenerator.generate invocation. + eef_targets = torch.stack([poured_eef_pose, current_eef_pose], dim=1) + + result = self.motion_generator.generate( + build_pose_plan_states(eef_targets), + options=request.motion_policy.to_motion_gen_options( + start_qpos=start_arm_qpos, + control_part=control_part, + interpolation_dt=context.require_control_dt(), + ), + ) + assert isinstance(result.success, torch.Tensor) + assert result.positions is not None + assert result.dt is not None + success = result.success & eligible + + hand_grasp_qpos = grasp_endpoint.joint_positions( + GRASP_COMMAND, + num_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + full = torch.empty( + (self.num_envs, result.positions.shape[1], self.robot_dof), + dtype=context.robot.qpos.dtype, + device=self.device, + ) + full[:] = context.last_qpos.unsqueeze(1) + full[:, :, arm_joint_ids] = result.positions + full[:, :, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1) + + return self.build_plan( + request, + context, + success=success, + trajectory=TimedTrajectory.from_positions( + full, + env_ids=context.env_ids, + dt=result.dt, + ), + segment_lengths={"pour": full.shape[1]}, + ) + + +__all__ = ["Pour", "PourGoal", "PourOptions"] diff --git a/embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py b/embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py index abc0466ef..96cc1b0c2 100644 --- a/embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py +++ b/embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py @@ -611,9 +611,24 @@ def get_valid_grasp_poses( self, object_pose: torch.Tensor, approach_direction: torch.Tensor, - object_part: str = "center", + obj_longest_axis: torch.Tensor | None = None, + is_positive_part: bool = True, visualize_collision: bool = False, ): + """Filter valid grasps, optionally to one projected half of the object. + + Args: + object_pose: Current object pose with shape ``(4, 4)``. + approach_direction: World-frame gripper approach direction. + obj_longest_axis: Optional world-frame object axis. When ``None``, + all annotated antipodal pairs remain eligible (center mode). + is_positive_part: When an axis is supplied, select the positive + projected half if true and the negative half otherwise. + visualize_collision: Whether to visualize collision checks. + + Returns: + Success, grasp poses, opening lengths, and grasp costs. + """ if self._hit_point_pairs is None: logger.log_warning( "No antipodal point pairs available. " @@ -631,27 +646,39 @@ def get_valid_grasp_poses( hit_points_ = self._apply_transform(hit_points, object_pose) mesh_vert_transformed = self._apply_transform(self.vertices, object_pose) - if object_part == "bottom": - z_max = mesh_vert_transformed[:, 2].max() - z_min = mesh_vert_transformed[:, 2].min() - z_threshold = z_min + (z_max - z_min) * 0.45 - z_mask = (origin_points_[:, 2] < z_threshold) | ( - hit_points_[:, 2] < z_threshold - ) - origin_points_masked = origin_points_[z_mask] - hit_points_masked = hit_points_[z_mask] - elif object_part == "top": - z_max = mesh_vert_transformed[:, 2].max() - z_min = mesh_vert_transformed[:, 2].min() - z_threshold = z_min + (z_max - z_min) * 0.6 - z_mask = (origin_points_[:, 2] > z_threshold) | ( - hit_points_[:, 2] > z_threshold - ) - origin_points_masked = origin_points_[z_mask] - hit_points_masked = hit_points_[z_mask] - else: + if obj_longest_axis is None: origin_points_masked = origin_points_ hit_points_masked = hit_points_ + else: + axis = torch.as_tensor( + obj_longest_axis, + dtype=torch.float32, + device=self.device, + ) + if axis.shape != (3,) or not torch.isfinite(axis).all(): + raise ValueError("obj_longest_axis must be a finite (3,) tensor.") + axis_norm = torch.linalg.vector_norm(axis) + if axis_norm <= 1.0e-8: + raise ValueError("obj_longest_axis must be non-zero.") + if not isinstance(is_positive_part, bool): + raise TypeError("is_positive_part must be a bool.") + axis = axis / axis_norm + mesh_projection = torch.matmul(mesh_vert_transformed, axis) + mesh_projection_range = mesh_projection.max() - mesh_projection.min() + projection_posi_threshold = ( + mesh_projection.min() + 0.6 * mesh_projection_range + ) + projection_nega_threshold = ( + mesh_projection.min() + 0.4 * mesh_projection_range + ) + pair_centers = 0.5 * (origin_points_ + hit_points_) + pair_projection = torch.matmul(pair_centers, axis) + if is_positive_part: + part_mask = pair_projection > projection_posi_threshold + else: + part_mask = pair_projection < projection_nega_threshold + origin_points_masked = origin_points_[part_mask] + hit_points_masked = hit_points_[part_mask] return self._filter_valid_grasp_poses( origin_points_=origin_points_masked, hit_points_=hit_points_masked, @@ -910,7 +937,9 @@ def get_grasp_poses( """ is_success, valid_grasp_poses, valid_open_lengths, total_cost = ( self.get_valid_grasp_poses( - object_pose, approach_direction, visualize_collision + object_pose, + approach_direction, + visualize_collision=visualize_collision, ) ) if not is_success: diff --git a/scripts/benchmark/atomic_action/common.py b/scripts/benchmark/atomic_action/common.py index 7220ab625..81e846da2 100644 --- a/scripts/benchmark/atomic_action/common.py +++ b/scripts/benchmark/atomic_action/common.py @@ -576,10 +576,14 @@ def get_valid_grasp_poses( self, obj_poses, approach_direction, + obj_longest_axis=None, + is_positive_part=True, ): results = super().get_valid_grasp_poses( obj_poses=obj_poses, approach_direction=approach_direction, + obj_longest_axis=obj_longest_axis, + is_positive_part=is_positive_part, ) if not _is_horizontal_approach_direction(approach_direction): return results diff --git a/scripts/tutorials/atomic_action/axis_align.py b/scripts/tutorials/atomic_action/axis_align.py new file mode 100644 index 000000000..fbe513591 --- /dev/null +++ b/scripts/tutorials/atomic_action/axis_align.py @@ -0,0 +1,233 @@ +# ---------------------------------------------------------------------------- +# 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 AxisAlign on the same cube and robot used by PickUp.""" + +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 typing import Sequence +from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + AxisAlignAffordance, + AxisAlignGoal, + AxisAlignOptions, + ControlPartCommandProfile, + MotionPolicy, + ObjectSemantics, +) +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.objects import RigidObject +from embodichain.lab.sim.shapes import CubeCfg +from embodichain.utils import logger +from scripts.tutorials.atomic_action.tutorial_utils import ( + add_ur5_gripper_robot, + clone_local_pose_from_first_env, + create_antipodal_semantics, + create_curobo_motion_generator, + create_tutorial_argument_parser, + create_tutorial_simulation, + draw_axis_marker, + get_hand_open_close_qpos, + initialize_pre_pick_robot_pose, + make_clear_dynamics_callback, + prepare_tutorial_scene, + replay_trajectory, + run_tutorial, +) + +OBJECT_SIZE = (0.05, 0.05, 0.05) +OBJECT_XY = (-0.42, -0.08) +DEFAULT_INTERNAL_AXIS = (1.0, 0.0, 0.0) +DEFAULT_TARGET_AXIS = (0.0, 0.0, 1.0) +HORIZONTAL_TARGET_AXIS = (0.0, 1.0, 0.0) +ALIGNMENT_AXES = { + "upright": (DEFAULT_INTERNAL_AXIS, DEFAULT_TARGET_AXIS), + "horizontal_align": (DEFAULT_INTERNAL_AXIS, HORIZONTAL_TARGET_AXIS), +} +ALIGN_SAMPLE_INTERVAL = 180 +HAND_INTERP_STEPS = 12 +POST_TRAJECTORY_STEPS = 240 + + +def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the AxisAlign tutorial.""" + parser = create_tutorial_argument_parser( + "Demonstrate upright or horizontal AxisAlign on a cube.", + features=("grasp_sampling", "visualize_axes"), + ) + parser.add_argument( + "--alignment", + choices=tuple(ALIGNMENT_AXES), + default="upright", + help="Choose the object-axis alignment example.", + ) + return parser.parse_args() + + +def create_align_object( + sim, + obj_position=None, +) -> RigidObject: + """Create the same settled cube used by the PickUp tutorial.""" + init_pos = ( + [*OBJECT_XY, OBJECT_SIZE[2]] if obj_position is None else list(obj_position) + ) + obj = sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=list(OBJECT_SIZE)), + attrs=RigidBodyAttributesCfg( + mass=0.05, + dynamic_friction=0.97, + static_friction=0.99, + ), + max_convex_hull_num=16, + init_pos=init_pos, + ) + ) + sim.update(step=10) + clone_local_pose_from_first_env(obj) + obj.clear_dynamics() + return obj + + +def create_axis_align_semantics( + obj: RigidObject, args: argparse.Namespace, obj_internal_axis: Sequence[float] +) -> ObjectSemantics: + """Extend the tutorial antipodal affordance with a local alignment axis.""" + semantics = create_antipodal_semantics( + obj, + label="cube", + n_sample=args.n_sample, + force_reannotate=args.force_reannotate, + ) + antipodal = semantics.affordance + return ObjectSemantics( + label=semantics.label, + geometry=semantics.geometry, + properties=semantics.properties, + entity=semantics.entity, + entity_id=semantics.entity_id, + affordance=AxisAlignAffordance( + mesh_vertices=antipodal.mesh_vertices, + mesh_triangles=antipodal.mesh_triangles, + generator_cfg=antipodal.generator_cfg, + gripper_collision_cfg=antipodal.gripper_collision_cfg, + force_reannotate=antipodal.force_reannotate, + internal_axis=torch.tensor(obj_internal_axis, dtype=torch.float32), + ), + ) + + +def main() -> None: + """Plan and replay a grasp, axis alignment, lowering, and release.""" + args = parse_arguments() + sim = create_tutorial_simulation(args) + robot = add_ur5_gripper_robot(sim, tcp_z=0.15) + obj = create_align_object(sim) + hand_open, hand_close = get_hand_open_close_qpos(robot) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + motion_gen = create_curobo_motion_generator(robot) + + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) + # apply object internal axis to the semantics creation + obj_internal_axis = ALIGNMENT_AXES[args.alignment][0] + # apply target axis for the alignment skill + align_target_axis = ALIGNMENT_AXES[args.alignment][1] + semantics = create_axis_align_semantics(obj, args, obj_internal_axis) + if not args.no_vis_eef_axis: + draw_axis_marker( + sim, + "axis_align_object_axis", + obj.get_local_pose(to_matrix=True), + ) + wait_for_user = prepare_tutorial_scene( + sim, + args, + f"Inspect the cube, then press Enter to plan {args.alignment} AxisAlign...", + ) + + compiled = engine.compile( + ( + engine.make_invocation( + "axis_align", + AxisAlignGoal(semantics), + control_parts={ + "primary": {"motion": "arm", "grasp": "hand"}, + }, + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=ALIGN_SAMPLE_INTERVAL, + ), + skill_options=AxisAlignOptions( + target_axis=torch.tensor( + align_target_axis, + dtype=torch.float32, + device=sim.device, + ), + approach_direction=torch.tensor( + [0.0, 0.0, -1.0], + dtype=torch.float32, + device=sim.device, + ), + pre_grasp_distance=0.15, + lift_height=0.16, + lower_distance=0.03, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ), + ), + engine.initial_context(control_dt=sim.sim_config.physics_dt), + ) + if not compiled.plan_success.all(): + logger.log_warning("Failed to plan AxisAlign demo trajectory.") + return + + if wait_for_user: + input("Press Enter to replay the AxisAlign demo...") + clear_after_step = compiled.segment(0, "manipulate").start + replay_trajectory( + sim, + robot, + compiled.trajectory, + args, + video_prefix=f"axis_align_{args.alignment}_cube_auto_play", + hold_steps=POST_TRAJECTORY_STEPS, + on_trajectory_step=make_clear_dynamics_callback(obj, clear_after_step), + ) + if wait_for_user: + input("Press Enter to exit the simulation...") + + +if __name__ == "__main__": + run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/hand_over.py b/scripts/tutorials/atomic_action/hand_over.py index 446596c86..47d85397b 100644 --- a/scripts/tutorials/atomic_action/hand_over.py +++ b/scripts/tutorials/atomic_action/hand_over.py @@ -14,12 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Demonstrate a dual-arm handover with a single textured mesh object. - -The left arm picks the object up by its top part, hands it to the right arm at a -middle handover pose, the right arm grasps the bottom part, the left arm -releases, and the right arm carries the object to the other side. -""" +"""Demonstrate the unified dual-arm PickUp-to-HandOver atomic action.""" from __future__ import annotations @@ -35,11 +30,10 @@ from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.atomic_actions import ( - GraspGoal, AtomicActionEngine, ControlPartCommandProfile, + HandOverGoal, HandOverOptions, - PickUpOptions, MotionPolicy, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg @@ -67,7 +61,12 @@ serve_tutorial_scene, ) -OBJECT_MESH_PATH = get_data_path("SodaCan/simple_cola_can.obj") +VERTICAL_OBJECT_MESH_PATH = get_data_path("SodaCan/simple_cola_can.obj") +HORIZONTAL_OBJECT_MESH_PATH = get_data_path( + "CoordinatedPlacementAndPickment/pencil.glb" +) +VERTICAL_OBJECT_SCALE = (0.56, 0.56, 0.56) +HORIZONTAL_OBJECT_SCALE = (2.0, 2.0, 2.0) GRIPPER_TCP_Z = 0.155 SUPPORT_SURFACE_Z = 0.50 SUPPORT_SURFACE_SIZE = (0.8, 1.2, 0.02) @@ -78,31 +77,23 @@ ) # --- Adjustable scene placeholders ----------------------------------------- -# The object starts on the left side, is handed over at a lifted middle pose, -# and is delivered to the right side. Tweak these to match the mesh geometry -# and the selected dual-arm robot's reach. +# The object starts on one side and is delivered to the other. HandOver chooses +# the nearer arm for pickup and computes the middle handover position itself. OBJECT_INIT_XY = (0.0, 0.02) -MIDDLE_OBJECT_XYZ = (0.0, 0.02, 0.82) -MIDDLE_OBJECT_YAW_DEG = 0.0 -FINAL_OBJECT_XYZ = (0.22, 0.02, 0.72) -FINAL_OBJECT_YAW_DEG = 0.0 +OBJECT_ROT_VERTICAL = (90.0, 0.0, 0.0) +OBJECT_ROT_HORIZONTAL = (90.0, 0.0, 0.0) +FINAL_OBJECT_XYZ = (0.0, -0.2, 0.6) # --------------------------------------------------------------------------- -HAND_CLOSE_QPOS = 0.026 -PICKUP_SAMPLE_INTERVAL = 80 -PICKUP_HAND_INTERP_STEPS = 5 -PICKUP_PRE_GRASP_DISTANCE = 0.08 -PICKUP_LIFT_HEIGHT = 0.1 -HANDOVER_SAMPLE_INTERVAL = 140 +HAND_CLOSE_QPOS = 0.04 +HANDOVER_SAMPLE_INTERVAL = 220 HANDOVER_HAND_INTERP_STEPS = 10 -HANDOVER_HOLD_STEPS = 4 -HANDOVER_RETREAT_STEPS = 28 HANDOVER_PRE_GRASP_DISTANCE = 0.08 -HANDOVER_LIFT_HEIGHT = 0.08 +HANDOVER_LIFT_HEIGHT = 0.15 TRAJECTORY_SIM_STEPS = 4 HANDOVER_RECORD_LOOK_AT = ( - (-0.25, 0.02, 2.5), - (0.0, 0.02, 0.75), + (-1.0, 0.2, 1.8), + (-0.4, 0.0, 0.7), (0.0, 0.0, 1.0), ) @@ -115,6 +106,11 @@ def parse_arguments() -> argparse.Namespace: default_device="cpu", default_renderer="hybrid", ) + parser.add_argument( + "--is_horizontal", + action="store_true", + help="Use the horizontal WaterBasin object instead of the vertical soda can.", + ) return parser.parse_args() @@ -130,9 +126,9 @@ def create_dual_robot( urdf_name=f"dual_{robot_type}_hand_over", tcp_z=GRIPPER_TCP_Z, ur_ik_nearest_weight=(1.0, 4.0, 1.0, 1.0, 1.0, 1.0), - hand_stiffness=1e2, - hand_damping=1e1, - hand_max_effort=1e3, + hand_stiffness=1e3, + hand_damping=1e2, + hand_max_effort=1e4, ) @@ -145,12 +141,18 @@ def create_support_surface(sim: SimulationManager) -> RigidObject: ) -def create_handover_object(sim: SimulationManager) -> RigidObject: - """Create the textured mesh object on the support surface.""" +def create_handover_object(sim: SimulationManager, args) -> RigidObject: + """Create the mode-specific mesh object on the support surface.""" + mesh_path = ( + HORIZONTAL_OBJECT_MESH_PATH if args.is_horizontal else VERTICAL_OBJECT_MESH_PATH + ) + body_scale = ( + HORIZONTAL_OBJECT_SCALE if args.is_horizontal else VERTICAL_OBJECT_SCALE + ) return sim.add_rigid_object( cfg=RigidObjectCfg( uid="handover_object", - shape=MeshCfg(fpath=OBJECT_MESH_PATH, compute_uv=False), + shape=MeshCfg(fpath=mesh_path, compute_uv=False), attrs=RigidBodyAttributesCfg( mass=0.01, dynamic_friction=0.97, @@ -164,10 +166,12 @@ def create_handover_object(sim: SimulationManager) -> RigidObject: min_velocity_iters=8, max_depenetration_velocity=2.0, ), - max_convex_hull_num=1, + max_convex_hull_num=16, init_pos=[OBJECT_INIT_XY[0], OBJECT_INIT_XY[1], SUPPORT_SURFACE_Z + 0.12], - init_rot=[90.0, 0.0, 0.0], - body_scale=(0.56, 0.56, 0.56), + init_rot=( + OBJECT_ROT_VERTICAL if not args.is_horizontal else OBJECT_ROT_HORIZONTAL + ), + body_scale=body_scale, ) ) @@ -177,9 +181,9 @@ def run_handover_demo( sim: SimulationManager, robot: Robot, ) -> None: - """Plan and optionally execute a pick-up followed by a handover.""" + """Plan and optionally execute one unified pick-up and handover.""" create_support_surface(sim) - obj = create_handover_object(sim) + obj = create_handover_object(sim, args) settle_object(sim, obj, step=0) clone_local_pose_from_first_env(obj) obj.clear_dynamics() @@ -196,46 +200,13 @@ def run_handover_demo( robot, hand_control_part="right_hand", close_qpos=HAND_CLOSE_QPOS ) - middle_pose = torch.as_tensor( - [ - [1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, -1.0, 0.0], - [0.0, 1.0, 0.0, 0.7], - [0.0, 0.0, 0.0, 1.0], - ], - dtype=torch.float32, - ) - - final_pose = torch.as_tensor( - [ - [1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, -1.0, -0.2], - [0.0, 1.0, 0.0, 0.7], - [0.0, 0.0, 0.0, 1.0], - ], - dtype=torch.float32, - ) + final_pose = torch.eye(4, dtype=torch.float32) + final_pose[:3, 3] = torch.as_tensor(FINAL_OBJECT_XYZ) - # Step 1 - the left arm picks the object up by its top part. - pick_up_options = PickUpOptions( - pick_object_part="top", - pre_grasp_distance=PICKUP_PRE_GRASP_DISTANCE, - lift_height=PICKUP_LIFT_HEIGHT, - hand_interp_steps=PICKUP_HAND_INTERP_STEPS, - ) - # Step 2 - hand the object from the left arm to the right arm. handover_options = HandOverOptions( - receive_pick_object_part="bottom", - middle_object_pose=middle_pose, - final_object_pose=final_pose, pre_grasp_distance=HANDOVER_PRE_GRASP_DISTANCE, lift_height=HANDOVER_LIFT_HEIGHT, hand_interp_steps=HANDOVER_HAND_INTERP_STEPS, - hold_steps=HANDOVER_HOLD_STEPS, - retreat_steps=HANDOVER_RETREAT_STEPS, - receive_approach_direction=torch.as_tensor( - [0.0, 707106781, -707106781], dtype=torch.float32 - ), ) engine = AtomicActionEngine( motion_generator=motion_gen, @@ -258,21 +229,14 @@ def run_handover_demo( sim.update(step=10) compiled = engine.compile( ( - engine.make_invocation( - "pick_up", - GraspGoal(object_semantics), - control_parts={"primary": {"motion": "left_arm", "grasp": "left_hand"}}, - motion_policy=MotionPolicy( - strategy="motion_gen", - sample_count=PICKUP_SAMPLE_INTERVAL, - ), - skill_options=pick_up_options, - ), engine.make_invocation( "hand_over", - GraspGoal(object_semantics), + HandOverGoal(object_semantics, target_pose=final_pose), control_parts={ - "source": {"motion": "left_arm", "grasp": "left_hand"}, + "source": { + "motion": "left_arm", + "grasp": "left_hand", + }, "destination": { "motion": "right_arm", "grasp": "right_hand", @@ -291,7 +255,7 @@ def run_handover_demo( traj = compiled.trajectory.positions if not success.all(): - logger.log_warning("Failed to plan the full pick-up + handover trajectory.") + logger.log_warning("Failed to plan the unified HandOver trajectory.") return if args.diagnose_plan: diff --git a/scripts/tutorials/atomic_action/pour.py b/scripts/tutorials/atomic_action/pour.py new file mode 100644 index 000000000..30346c19a --- /dev/null +++ b/scripts/tutorials/atomic_action/pour.py @@ -0,0 +1,171 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Pick up a cube from the side, then rotate it with the Pour action.""" + +from __future__ import annotations + +import argparse +import math +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[3] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +import torch + +from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + ControlPartCommandProfile, + GraspGoal, + MotionPolicy, + PickUpOptions, + PourGoal, + PourOptions, +) +from embodichain.utils import logger +from scripts.tutorials.atomic_action.axis_align import ( + create_align_object, + create_axis_align_semantics, +) +from scripts.tutorials.atomic_action.tutorial_utils import ( + add_ur5_gripper_robot, + create_curobo_motion_generator, + create_tutorial_argument_parser, + create_tutorial_simulation, + draw_axis_marker, + get_hand_open_close_qpos, + initialize_pre_pick_robot_pose, + make_clear_dynamics_callback, + prepare_tutorial_scene, + replay_trajectory, + run_tutorial, +) + +POUR_INTERNAL_AXIS = (1.0, 0.0, 0.0) +APPROACH_DIRECTION = (-0.707, 0, -0.707) +PICK_SAMPLE_INTERVAL = 120 +POUR_SAMPLE_INTERVAL = 80 +HAND_INTERP_STEPS = 12 +POST_TRAJECTORY_STEPS = 240 +OBJ_POSITION = (-0.5, 0.0, 0.0) + + +def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the Pour tutorial.""" + parser = create_tutorial_argument_parser( + "Pick up a cube horizontally, then pour it about a local axis.", + features=("grasp_sampling", "visualize_axes"), + ) + parser.add_argument( + "--rotate_angle", + type=float, + default=math.pi / 4.0, + help="Signed pouring rotation in radians.", + ) + return parser.parse_args() + + +def main() -> None: + """Plan and replay PickUp followed by Pour.""" + args = parse_arguments() + sim = create_tutorial_simulation(args) + robot = add_ur5_gripper_robot(sim, tcp_z=0.15) + obj = create_align_object( + sim, + obj_position=OBJ_POSITION, + ) + hand_open, hand_close = get_hand_open_close_qpos(robot) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + motion_gen = create_curobo_motion_generator(robot) + + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) + semantics = create_axis_align_semantics(obj, args, POUR_INTERNAL_AXIS) + if not args.no_vis_eef_axis: + draw_axis_marker(sim, "pour_object_axis", obj.get_local_pose(to_matrix=True)) + wait_for_user = prepare_tutorial_scene( + sim, + args, + "Inspect the cube, then press Enter to plan PickUp followed by Pour...", + ) + + control_parts = {"primary": {"motion": "arm", "grasp": "hand"}} + compiled = engine.compile( + ( + engine.make_invocation( + "pick_up", + GraspGoal(semantics), + control_parts=control_parts, + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=PICK_SAMPLE_INTERVAL, + ), + skill_options=PickUpOptions( + approach_direction=torch.tensor( + APPROACH_DIRECTION, + dtype=torch.float32, + device=sim.device, + ), + pre_grasp_distance=0.15, + lift_height=0.16, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ), + engine.make_invocation( + "pour", + PourGoal(), + control_parts=control_parts, + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=POUR_SAMPLE_INTERVAL, + ), + skill_options=PourOptions(rotate_angle=args.rotate_angle), + ), + ), + engine.initial_context(control_dt=sim.sim_config.physics_dt), + ) + if not compiled.plan_success.all(): + logger.log_warning("Failed to plan PickUp followed by Pour.") + return + + if wait_for_user: + input("Press Enter to replay the PickUp + Pour trajectory...") + clear_after_step = compiled.segment(0, "lift").start + replay_trajectory( + sim, + robot, + compiled.trajectory, + args, + video_prefix="pour_cube_auto_play", + hold_steps=POST_TRAJECTORY_STEPS, + on_trajectory_step=make_clear_dynamics_callback(obj, clear_after_step), + ) + 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 89118dda4..24971990d 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -35,6 +35,10 @@ AssembleGoal, AtomicAction, AtomicActionEngine, + AxisAlign, + AxisAlignAffordance, + AxisAlignGoal, + AxisAlignOptions, ControlPartCommandProfile, CoordinatedPickGoal, CoordinatedPickment, @@ -47,6 +51,7 @@ ExecutionEventKind, GraspGoal, HandOver, + HandOverGoal, HandOverOptions, HeldObjectPoseGoal, HeldObjectState, @@ -66,6 +71,9 @@ Place, PlaceGoal, PlaceOptions, + Pour, + PourGoal, + PourOptions, PlanningContext, Press, PressAffordance, @@ -93,7 +101,7 @@ PlanOptions, PlanResult, ) -from embodichain.utils.math import pose_inv +from embodichain.utils.math import axis_angle_to_rotation_matrix, pose_inv NUM_ENVS = 2 ARM_DOF = 6 @@ -380,6 +388,21 @@ def _dual_motion_generator() -> MotionGenerator: "left_hand": object(), "right_hand": object(), } + robot.cfg = Mock() + robot.cfg.solver_cfg = { + "left_arm": Mock(root_link_name="left_root"), + "right_arm": Mock(root_link_name="right_root"), + } + + def get_link_pose( + link_name: str, + env_ids: list[int] | None = None, + to_matrix: bool = True, + ) -> torch.Tensor: + del env_ids, to_matrix + pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + pose[:, 0, 3] = -1.0 if link_name == "left_root" else 1.0 + return pose def get_qpos(name: str | None = None) -> torch.Tensor: if name in {"left_arm", "right_arm"}: @@ -427,6 +450,7 @@ def compute_fk( robot.get_joint_ids.side_effect = get_joint_ids robot.compute_ik.side_effect = compute_ik robot.compute_fk.side_effect = compute_fk + robot.get_link_pose.side_effect = get_link_pose generator = object.__new__(MotionGenerator) generator.robot = robot @@ -497,14 +521,16 @@ def test_builtin_descriptors_expose_goals_not_legacy_targets() -> None: assert MoveEndEffector.GoalType is EndEffectorPoseGoal assert MoveJoints.GoalType is JointPositionGoal assert PickUp.GoalType is GraspGoal + assert AxisAlign.GoalType is AxisAlignGoal assert MoveHeldObject.GoalType is HeldObjectPoseGoal assert Place.GoalType == (PlaceGoal, AssembleGoal) + assert Pour.GoalType is PourGoal 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 + assert HandOver.GoalType is HandOverGoal def test_interaction_primitives_use_motion_centric_skill_ids() -> None: @@ -519,8 +545,10 @@ def test_interaction_primitives_use_motion_centric_skill_ids() -> None: "options", ( PickUpOptions(), + AxisAlignOptions(), MoveHeldObjectOptions(), PlaceOptions(), + PourOptions(), PressOptions(), SlideOptions(), TwistOptions(), @@ -751,6 +779,180 @@ def test_move_held_object_requires_projected_attachment() -> None: semantics.entity.get_local_pose.assert_not_called() +def test_pour_rotates_held_object_about_internal_axis_and_returns() -> None: + generator = _motion_generator() + solved_poses: list[torch.Tensor] = [] + + def compute_ik( + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + **_: object, + ) -> tuple[torch.Tensor, torch.Tensor]: + solved_poses.append(pose.clone()) + return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed + 0.1 + + generator.robot.compute_ik.side_effect = compute_ik + current_eef_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + current_eef_pose[:, 0, 3] = torch.tensor([0.4, 0.7]) + generator.robot.compute_fk.return_value = current_eef_pose + generator.robot.compute_fk.side_effect = None + action = _bind_action(generator, Pour()) + semantics = ObjectSemantics( + affordance=AxisAlignAffordance(internal_axis=torch.tensor([1.0, 0.0, 0.0])), + geometry={}, + label="pourable-object", + ) + task = TaskState( + batch_size=NUM_ENVS, + device="cpu", + held_objects={"arm": _held(semantics)}, + ) + context = _context(task) + + plan = _plan_action( + action, + ActionInvocation( + skill_id="pour", + goal=PourGoal(), + binding=_binding(action), + motion_policy=MotionPolicy(sample_count=10), + skill_options=PourOptions(rotate_angle=math.pi / 2.0), + ), + context, + ) + + expected_rotation = axis_angle_to_rotation_matrix( + torch.tensor([math.pi / 2.0, 0.0, 0.0]) + ) + trajectory = _joint_trajectory(plan) + assert plan.plan_success.tolist() == [True, True] + assert trajectory.positions.shape == (NUM_ENVS, 10, ROBOT_DOF) + assert trajectory.duration.tolist() == pytest.approx([9.0 / 60.0] * NUM_ENVS) + assert [segment.name for segment in plan.segments] == ["pour"] + assert torch.allclose( + solved_poses[0][:, :3, :3], + expected_rotation.expand(NUM_ENVS, -1, -1), + atol=1.0e-6, + ) + assert torch.allclose( + solved_poses[0][:, :3, 3], + current_eef_pose[:, :3, 3], + ) + assert len(solved_poses) == 2 + assert torch.allclose(solved_poses[1], current_eef_pose, atol=1.0e-6) + assert torch.all(trajectory.positions[:, :, ARM_DOF:] == 1.0) + assert plan.expected_effects.is_empty + assert context.task is task + + +def test_engine_compiles_pickup_followed_by_pour() -> None: + generator = _motion_generator() + engine = AtomicActionEngine( + generator, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=torch.zeros(HAND_DOF), + grasp=torch.ones(HAND_DOF), + ) + }, + ) + semantics = ObjectSemantics( + affordance=AxisAlignAffordance(internal_axis=torch.tensor([1.0, 0.0, 0.0])), + geometry={}, + label="pourable-object", + entity_id="target", + ) + context = _context( + scene=_target_scene( + torch.eye(4).repeat(NUM_ENVS, 1, 1), + timestamp=0.0, + version=0, + ) + ) + + compiled = engine.compile( + ( + engine.make_invocation( + "pick_up", + GraspGoal(semantics, grasp_xpos=torch.eye(4)), + control_parts={ + "primary": {"motion": "arm", "grasp": "hand"}, + }, + motion_policy=MotionPolicy(sample_count=20), + ), + engine.make_invocation( + "pour", + PourGoal(), + control_parts={ + "primary": {"motion": "arm", "grasp": "hand"}, + }, + motion_policy=MotionPolicy(sample_count=10), + skill_options=PourOptions(rotate_angle=math.pi / 2.0), + ), + ), + context, + ) + + assert compiled.plan_success.tolist() == [True, True] + assert [plan.skill_id for plan in compiled.action_plans] == ["pick_up", "pour"] + assert compiled.projected_context.get_held_object("arm") is not None + + +def test_pour_requires_exclusively_held_axis_align_affordance() -> None: + generator = _motion_generator() + action = _bind_action(generator, Pour()) + invocation = ActionInvocation( + skill_id="pour", + goal=PourGoal(), + binding=_binding(action), + motion_policy=MotionPolicy(sample_count=10), + ) + + with pytest.raises(ValueError, match="run PickUp first"): + _plan_action(action, invocation, _context()) + + invalid_task = TaskState( + batch_size=NUM_ENVS, + device="cpu", + held_objects={"arm": _held(_semantics())}, + ) + with pytest.raises(ValueError, match="AxisAlignAffordance"): + _plan_action(action, invocation, _context(invalid_task)) + + semantics = ObjectSemantics( + affordance=AxisAlignAffordance(), + geometry={}, + label="shared-pourable-object", + ) + task = TaskState( + batch_size=NUM_ENVS, + device="cpu", + held_objects={ + "arm": _held(semantics), + "alternate_arm": _held( + semantics, + env_mask=torch.tensor([True, False]), + ), + }, + ) + plan = _plan_action(action, invocation, _context(task)) + + assert plan.plan_success.tolist() == [False, True] + trajectory = _joint_trajectory(plan) + assert torch.allclose( + trajectory.positions[0], + _context(task).robot.qpos[0].unsqueeze(0).expand(trajectory.waypoint_count, -1), + ) + + +def test_pour_options_only_contain_rotate_angle_and_require_finite_value() -> None: + assert set(PourOptions.__dataclass_fields__) == {"rotate_angle"} + assert PourOptions().rotate_angle == pytest.approx(math.pi / 4.0) + with pytest.raises(ValueError, match="rotate_angle must be finite"): + PourOptions(rotate_angle=float("nan")) + + def test_move_held_object_moves_only_exclusively_held_rows() -> None: generator = _motion_generator() @@ -760,6 +962,7 @@ def move_ik( joint_seed: torch.Tensor, **_: object, ) -> tuple[torch.Tensor, torch.Tensor]: + del pose, name return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed + 0.1 generator.robot.compute_ik.side_effect = move_ik @@ -995,6 +1198,230 @@ def test_pick_explicit_grasp_bypasses_sampling_and_records_grasp() -> None: assert plan.segment("close").stop == plan.segment("lift").start +def test_axis_align_plans_two_arm_phases_and_aligns_the_object_axis() -> None: + generator = _motion_generator() + generator.generate = Mock(wraps=generator.generate) + solved_poses: list[torch.Tensor] = [] + + def compute_ik( + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + **_: object, + ) -> tuple[torch.Tensor, torch.Tensor]: + solved_poses.append(pose.clone()) + return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed + 0.1 + + generator.robot.compute_ik.side_effect = compute_ik + action = _bind_action(generator, AxisAlign()) + object_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + semantics = ObjectSemantics( + affordance=AxisAlignAffordance(internal_axis=torch.tensor([0.0, 0.0, 1.0])), + geometry={}, + label="axis-object", + entity_id="target", + ) + context = _context(scene=_target_scene(object_pose, timestamp=0.0, version=0)) + original_task = context.task + + plan = _plan_action( + action, + ActionInvocation( + skill_id="axis_align", + goal=AxisAlignGoal(semantics=semantics, grasp_xpos=torch.eye(4)), + binding=_binding(action), + motion_policy=MotionPolicy(sample_count=20), + skill_options=AxisAlignOptions( + target_axis=torch.tensor([1.0, 0.0, 0.0]), + lift_height=0.1, + lower_distance=0.03, + ), + ), + context, + ) + + trajectory = _joint_trajectory(plan) + assert plan.plan_success.tolist() == [True, True] + assert trajectory.positions.shape == (NUM_ENVS, 20, ROBOT_DOF) + assert torch.equal(trajectory.env_ids, context.env_ids) + assert trajectory.duration.tolist() == pytest.approx([19.0 / 60.0] * NUM_ENVS) + assert [segment.name for segment in plan.segments] == [ + "approach", + "close", + "manipulate", + "open", + ] + assert generator.generate.call_count == 2 + assert len(solved_poses) == 5 + assert plan.expected_effects.is_empty + assert context.task is original_task + assert plan.scene_dependencies == ("target",) + final_object_rotation = solved_poses[-1][:, :3, :3] + final_world_axis = torch.matmul( + final_object_rotation, + torch.tensor([0.0, 0.0, 1.0]), + ) + assert torch.allclose( + final_world_axis, + torch.tensor([1.0, 0.0, 0.0]).expand(NUM_ENVS, -1), + atol=1.0e-6, + ) + assert solved_poses[-1][:, 2, 3].tolist() == pytest.approx([0.07, 0.07]) + + +def test_axis_align_upright_prefers_perpendicular_grasp_and_pre_rotates() -> None: + generator = _motion_generator() + solved_poses: list[torch.Tensor] = [] + + def compute_ik( + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + **_: object, + ) -> tuple[torch.Tensor, torch.Tensor]: + solved_poses.append(pose.clone()) + return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed + 0.1 + + generator.robot.compute_ik.side_effect = compute_ik + action = _bind_action(generator, AxisAlign()) + parallel_grasp = torch.eye(4) + perpendicular_grasp = torch.eye(4) + perpendicular_grasp[:3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + candidates = torch.stack([parallel_grasp, perpendicular_grasp]) + affordance = AxisAlignAffordance(internal_axis=torch.tensor([1.0, 0.0, 0.0])) + affordance.get_valid_grasp_poses = Mock( + return_value=[ + (candidates.clone(), torch.tensor([0.0, 10.0])) for _ in range(NUM_ENVS) + ] + ) + semantics = ObjectSemantics( + affordance=affordance, + geometry={}, + label="upright-axis-object", + entity_id="target", + ) + object_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + + plan = _plan_action( + action, + ActionInvocation( + skill_id="axis_align", + goal=AxisAlignGoal(semantics=semantics), + binding=_binding(action), + motion_policy=MotionPolicy(sample_count=20), + skill_options=AxisAlignOptions( + target_axis=torch.tensor([0.0, 0.0, 1.0]), + ), + ), + _context(scene=_target_scene(object_pose, timestamp=0.0, version=0)), + ) + + assert plan.plan_success.all() + expected_rotation = ( + axis_angle_to_rotation_matrix(torch.tensor([0.0, math.pi / 4.0, 0.0])) + @ perpendicular_grasp[:3, :3] + ) + assert torch.allclose( + solved_poses[1][:, :3, :3], + expected_rotation.expand(NUM_ENVS, -1, -1), + atol=1.0e-6, + ) + + +def test_axis_align_holds_only_failed_environment_rows() -> None: + generator = _motion_generator() + + def compute_ik( + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + **_: object, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.tensor([True, False]), joint_seed + 0.1 + + generator.robot.compute_ik.side_effect = compute_ik + action = _bind_action(generator, AxisAlign()) + object_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + semantics = ObjectSemantics( + affordance=AxisAlignAffordance(), + geometry={}, + label="partially-alignable-object", + entity_id="target", + ) + context = _context(scene=_target_scene(object_pose, timestamp=0.0, version=0)) + + plan = _plan_action( + action, + _invocation( + action, + AxisAlignGoal(semantics=semantics, grasp_xpos=torch.eye(4)), + sample_count=20, + ), + context, + ) + + trajectory = _joint_trajectory(plan) + assert plan.plan_success.tolist() == [True, False] + assert not torch.allclose(trajectory.positions[0], context.robot.qpos[0]) + assert torch.allclose( + trajectory.positions[1], + context.robot.qpos[1].unsqueeze(0).expand(trajectory.waypoint_count, -1), + ) + + +def test_axis_align_validates_goal_and_binding_contract() -> None: + action = _bind_action(_motion_generator(), AxisAlign()) + semantics = ObjectSemantics( + affordance=AxisAlignAffordance(), + geometry={}, + label="axis-object", + ) + + with pytest.raises(TypeError, match="expects goal AxisAlignGoal"): + action.resolve_request( + ActionInvocation( + skill_id="axis_align", + goal=object(), + binding=_binding(action), + ) + ) + with pytest.raises(ValueError, match="missing=.*grasp"): + action.resolve_request( + ActionInvocation( + skill_id="axis_align", + goal=AxisAlignGoal(semantics), + binding=ActionBinding( + owner_id=_ACTION_ENGINES[id(action)].binding_owner_id, + ), + ) + ) + + +def test_axis_align_handles_opposite_axes_without_nan() -> None: + action = _bind_action(_motion_generator(), AxisAlign()) + identity = torch.eye(4).repeat(NUM_ENVS, 1, 1) + + eef_keyframes = action._axis_alignment_eef_keyframes( + identity, + identity, + torch.tensor([1.0, 0.0, 0.0]), + torch.tensor([-1.0, 0.0, 0.0]), + waypoint_count=3, + ) + + final_axis = torch.matmul( + eef_keyframes[:, -1, :3, :3], torch.tensor([1.0, 0.0, 0.0]) + ) + assert torch.isfinite(eef_keyframes).all() + assert torch.allclose( + final_axis, + torch.tensor([-1.0, 0.0, 0.0]).expand(NUM_ENVS, -1), + atol=1.0e-6, + ) + + def test_pick_holds_only_environment_without_a_feasible_grasp() -> None: generator = _motion_generator() entity = Mock() @@ -1997,45 +2424,104 @@ def test_slide_options_reject_invalid_direction() -> None: SlideOptions(direction="open") # type: ignore[arg-type] -def test_handover_does_not_mutate_cached_final_pose( +def _handover_semantics( + longest_axis: torch.Tensor = torch.tensor([0.0, 0.0, 1.0]), +) -> tuple[ObjectSemantics, AntipodalAffordance]: + affordance = AntipodalAffordance() + + def get_object_longest_axis( + obj_poses: torch.Tensor, + *, + max_points: int, + ) -> torch.Tensor: + assert max_points <= 1000 + return longest_axis.to(dtype=torch.float32).expand(obj_poses.shape[0], -1) + + def sample_grasps( + obj_poses: torch.Tensor, + approach_direction: torch.Tensor, + obj_longest_axis: torch.Tensor, + is_positive_part: bool, + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + assert obj_poses.shape == (1, 4, 4) + assert approach_direction.shape == (3,) + assert obj_longest_axis.shape == (3,) + grasp_poses = obj_poses.clone() + if is_positive_part: + grasp_poses[:, :3, :3] = torch.tensor( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]] + ) + else: + grasp_poses[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + return [ + (grasp_poses[index].unsqueeze(0), torch.zeros(1)) + for index in range(obj_poses.shape[0]) + ] + + affordance.get_object_longest_axis = Mock(side_effect=get_object_longest_axis) + affordance.get_valid_grasp_poses = Mock(side_effect=sample_grasps) + semantics = ObjectSemantics( + affordance=affordance, + geometry={}, + label="handover-object", + entity_id="handover_object", + ) + return semantics, affordance + + +def _handover_context( + object_pose: torch.Tensor, + task: TaskState | None = None, +) -> PlanningContext: + return _dual_context( + task, + scene=SceneSnapshot( + timestamp=0.0, + version=0, + entities={"handover_object": EntityState(object_pose)}, + ), + ) + + +def test_handover_picks_with_nearer_arm_and_preserves_waypoint_rotations( monkeypatch: pytest.MonkeyPatch, ) -> None: generator = _dual_motion_generator() - handover_options = HandOverOptions( - middle_object_pose=torch.eye(4), - final_object_pose=torch.eye(4), - hand_interp_steps=4, - hold_steps=2, - retreat_steps=5, - ) - action = _bind_action( - generator, - HandOver(default_options=handover_options), - ) - assert handover_options.final_object_pose is not None - original_final_pose = handover_options.final_object_pose.clone() - semantics = _semantics(entity_id="handover_object") - held = _held(semantics) - held.object_to_eef[:, :3, :3] = torch.tensor( - [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] - ) - held.object_to_eef[:, 0, 3] = torch.tensor([0.1, 0.2]) - task = TaskState( - batch_size=NUM_ENVS, - device="cpu", - held_objects={"left_arm": held}, - ) - current_eef = torch.eye(4).repeat(NUM_ENVS, 1, 1) - current_eef[:, :3, :3] = torch.tensor( + action = _bind_action(generator, HandOver()) + semantics, _ = _handover_semantics() + object_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + object_pose[:, :3, 3] = torch.tensor([-0.8, 0.1, 0.5]) + final_pose = torch.eye(4) + final_pose[:3, :3] = torch.tensor( [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]] ) - current_eef[:, 1, 3] = torch.tensor([0.3, 0.5]) - generator.robot.compute_fk.return_value = current_eef - generator.robot.compute_fk.side_effect = None - receive_grasp = torch.eye(4).repeat(NUM_ENVS, 1, 1) - action._resolve_receive_grasp = Mock( - return_value=(receive_grasp, torch.ones(NUM_ENVS, dtype=torch.bool)) - ) + final_pose[:3, 3] = torch.tensor([0.4, -0.2, 0.65]) + planned_parts: list[str] = [] + planned_targets: list[torch.Tensor] = [] + + def resolve_grasp( + sampled_affordance: AntipodalAffordance, + sampled_object_pose: torch.Tensor, + approach_direction: torch.Tensor, + *, + obj_longest_axis: torch.Tensor, + is_positive_part: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + del sampled_affordance, approach_direction, obj_longest_axis + grasp_pose = sampled_object_pose.clone() + if bool(is_positive_part[0].item()): + grasp_pose[:, :3, :3] = torch.tensor( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]] + ) + else: + grasp_pose[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + return grasp_pose, torch.ones(NUM_ENVS, dtype=torch.bool) + + action._resolve_grasp = Mock(side_effect=resolve_grasp) def plan_from_start( motion_generator: MotionGenerator, @@ -2045,9 +2531,12 @@ def plan_from_start( n_waypoints: int, motion_policy: MotionPolicy, interpolation_dt: float | None, - ) -> tuple[bool, torch.Tensor]: - del interpolation_dt - return True, start_qpos.unsqueeze(1).repeat(1, n_waypoints, 1) + ) -> tuple[torch.Tensor, torch.Tensor]: + del motion_generator, motion_policy, interpolation_dt + planned_parts.append(control_part) + planned_targets.append(target_poses.clone()) + trajectory = (start_qpos + 0.1).unsqueeze(1).repeat(1, n_waypoints, 1) + return torch.ones(NUM_ENVS, dtype=torch.bool), trajectory monkeypatch.setattr( "embodichain.lab.sim.atomic_actions.primitives.hand_over." @@ -2056,196 +2545,348 @@ def plan_from_start( ) invocation = ActionInvocation( skill_id="hand_over", - goal=GraspGoal( - semantics=semantics, - grasp_xpos=SceneEntityPose("unused_grasp_pose"), - ), + goal=HandOverGoal(semantics, target_pose=final_pose), binding=_dual_binding(action, "source", "destination"), - motion_policy=MotionPolicy(sample_count=30), + motion_policy=MotionPolicy(sample_count=24), + skill_options=HandOverOptions( + pre_grasp_distance=0.1, + lift_height=0.2, + hand_interp_steps=2, + ), ) + context = _handover_context(object_pose) + original_task = context.task - plan = _plan_action(action, invocation, _dual_context(task)) + plan = _plan_action(action, invocation, context) - assert plan.plan_success.all() - assert plan.scene_dependencies == () - handover_object_pose = action._resolve_receive_grasp.call_args.args[1] - expected_current_object_pose = torch.bmm( - current_eef, - pose_inv(held.object_to_eef), + trajectory = _joint_trajectory(plan) + assert plan.plan_success.tolist() == [True, True] + assert trajectory.positions.shape == (NUM_ENVS, 24, DUAL_ROBOT_DOF) + assert planned_parts == ["left_arm", "left_arm", "right_arm", "right_arm"] + assert [segment.name for segment in plan.segments] == [ + "pickup_approach", + "pickup_close", + "pickup_transport", + "receive_approach", + "receive_close", + "handover_release", + "place", + "receive_release", + ] + + pickup_call, receive_call = action._resolve_grasp.call_args_list + expected_axis = torch.tensor([[0.0, 0.0, 1.0]]).expand(NUM_ENVS, -1) + assert torch.equal(pickup_call.kwargs["obj_longest_axis"], expected_axis) + assert pickup_call.kwargs["is_positive_part"].tolist() == [False, False] + diagonal_component = math.sqrt(0.5) + pickup_horizontal = object_pose[:, :2, 3] + pickup_horizontal = pickup_horizontal / torch.linalg.vector_norm( + pickup_horizontal, dim=1, keepdim=True + ) + expected_pickup_direction = torch.zeros(NUM_ENVS, 3) + expected_pickup_direction[:, :2] = pickup_horizontal * diagonal_component + expected_pickup_direction[:, 2] = -diagonal_component + assert torch.allclose(pickup_call.args[2], expected_pickup_direction) + assert torch.equal(receive_call.kwargs["obj_longest_axis"], expected_axis) + assert receive_call.kwargs["is_positive_part"].tolist() == [True, True] + predicted_middle_pose = receive_call.args[1] + assert torch.allclose( + predicted_middle_pose[:, :3, 3], + torch.tensor([[0.0, 0.1, 0.7], [0.0, 0.1, 0.7]]), + ) + expected_receive_direction = torch.tensor( + [ + [0.0, diagonal_component, -diagonal_component], + [0.0, diagonal_component, -diagonal_component], + ] ) + assert torch.allclose(receive_call.args[2], expected_receive_direction) + + pickup_grasp_rotation = planned_targets[0][:, 1, :3, :3] assert torch.allclose( - handover_object_pose[:, :3, :3], - expected_current_object_pose[:, :3, :3], + planned_targets[1][:, :, :3, :3], + pickup_grasp_rotation[:, None].expand(-1, 2, -1, -1), ) - assert torch.equal(handover_options.final_object_pose, original_final_pose) - semantics.entity.get_local_pose.assert_not_called() - assert [segment.name for segment in plan.segments] == [ - "transfer", - "approach", - "close", - "hold", - "release", - "deliver", - ] + receive_grasp_rotation = planned_targets[2][:, 1, :3, :3] + assert torch.allclose( + planned_targets[3][:, :, :3, :3], + receive_grasp_rotation[:, None].expand(-1, 2, -1, -1), + ) + assert torch.allclose( + planned_targets[3][:, 0, :3, 3], + torch.tensor([[0.4, -0.2, 0.7], [0.4, -0.2, 0.7]]), + ) + assert torch.allclose( + planned_targets[3][:, 1, :3, 3], + torch.tensor([[0.4, -0.2, 0.65], [0.4, -0.2, 0.65]]), + ) + assert torch.allclose( + planned_targets[3][:, 0, 2, 3], + planned_targets[2][:, 1, 2, 3], + ) + + pickup_close_end = plan.segment("pickup_close").stop - 1 + receive_close_end = plan.segment("receive_close").stop - 1 + handover_release_end = plan.segment("handover_release").stop - 1 + receive_release_end = plan.segment("receive_release").stop - 1 + positions = trajectory.positions + assert torch.all( + positions[:, pickup_close_end, DUAL_ARM_DOF : DUAL_ARM_DOF + 2] == 1 + ) + assert torch.all(positions[:, pickup_close_end, DUAL_ARM_DOF + 2 :] == 0) + assert torch.all(positions[:, receive_close_end, DUAL_ARM_DOF:] == 1) + assert torch.all( + positions[:, handover_release_end, DUAL_ARM_DOF : DUAL_ARM_DOF + 2] == 0 + ) + assert torch.all(positions[:, handover_release_end, DUAL_ARM_DOF + 2 :] == 1) + assert torch.all(positions[:, receive_release_end, DUAL_ARM_DOF:] == 0) + assert plan.expected_effects.is_empty + assert context.task is original_task + assert plan.scene_dependencies == ("handover_object",) -def test_handover_holds_only_environment_with_ik_failure() -> None: +def test_handover_horizontal_mode_uses_downward_opposite_end_grasps() -> None: generator = _dual_motion_generator() - original_compute_ik = generator.robot.compute_ik.side_effect + action = _bind_action(generator, HandOver()) + semantics, _ = _handover_semantics(torch.tensor([1.0, 0.0, 0.0])) + object_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + object_pose[:, :3, 3] = torch.tensor([-0.8, 0.2, 0.5]) - def fail_second_receiving_arm( - pose: torch.Tensor, - name: str, - joint_seed: torch.Tensor, - **kwargs: object, + def resolve_grasp( + affordance: AntipodalAffordance, + sampled_object_pose: torch.Tensor, + approach_direction: torch.Tensor, + *, + obj_longest_axis: torch.Tensor, + is_positive_part: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - success, qpos = original_compute_ik( - pose=pose, - name=name, - joint_seed=joint_seed, - **kwargs, - ) - if name == "right_arm": - success = success.clone() - success[1] = False - return success, qpos + del affordance, approach_direction, obj_longest_axis, is_positive_part + return sampled_object_pose.clone(), torch.ones(NUM_ENVS, dtype=torch.bool) - generator.robot.compute_ik.side_effect = fail_second_receiving_arm - semantics = _semantics(entity_id="handover_object") - task = TaskState( - batch_size=NUM_ENVS, - device="cpu", - held_objects={"left_arm": _held(semantics)}, - ) - action = _bind_action( - generator, - HandOver( - default_options=HandOverOptions( - middle_object_pose=torch.eye(4), - final_object_pose=torch.eye(4), - hand_interp_steps=4, - hold_steps=2, - retreat_steps=5, - ) - ), - ) - action._resolve_receive_grasp = Mock( - return_value=( - torch.eye(4).repeat(NUM_ENVS, 1, 1), - torch.ones(NUM_ENVS, dtype=torch.bool), - ) + action._resolve_grasp = Mock(side_effect=resolve_grasp) + invocation = ActionInvocation( + skill_id="hand_over", + goal=HandOverGoal(semantics, target_pose=torch.eye(4)), + binding=_dual_binding(action, "source", "destination"), + motion_policy=MotionPolicy(sample_count=24), + skill_options=HandOverOptions(hand_interp_steps=2), ) - context = _dual_context(task) + + plan = _plan_action(action, invocation, _handover_context(object_pose)) + + assert plan.plan_success.tolist() == [True, True] + pickup_call, receive_call = action._resolve_grasp.call_args_list + downward = torch.tensor([0.0, 0.0, -1.0]) + expected_axis = torch.tensor([[1.0, 0.0, 0.0]]).expand(NUM_ENVS, -1) + assert torch.equal(pickup_call.args[2], downward.expand(NUM_ENVS, -1)) + assert torch.equal(pickup_call.kwargs["obj_longest_axis"], expected_axis) + assert pickup_call.kwargs["is_positive_part"].tolist() == [True, True] + assert torch.equal(receive_call.args[2], downward.expand(NUM_ENVS, -1)) + assert torch.equal(receive_call.kwargs["obj_longest_axis"], expected_axis) + assert receive_call.kwargs["is_positive_part"].tolist() == [False, False] + + +def test_handover_selects_arm_per_environment() -> None: + generator = _dual_motion_generator() + action = _bind_action(generator, HandOver()) + semantics, _ = _handover_semantics() + object_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + object_pose[:, 0, 3] = torch.tensor([-0.8, 0.8]) + object_pose[:, 1, 3] = 0.2 invocation = ActionInvocation( skill_id="hand_over", - goal=GraspGoal(semantics=semantics), + goal=HandOverGoal(semantics, target_pose=torch.eye(4)), binding=_dual_binding(action, "source", "destination"), - motion_policy=MotionPolicy(sample_count=30), + motion_policy=MotionPolicy(sample_count=24), + skill_options=HandOverOptions(hand_interp_steps=2), ) - request = action.resolve_request(invocation) - plan = action.plan(request, context) - projected = plan.expected_effects.apply(context.task, plan.plan_success) + plan = _plan_action(action, invocation, _handover_context(object_pose)) - assert plan.plan_success.tolist() == [True, False] trajectory = _joint_trajectory(plan) - assert not torch.allclose(trajectory.positions[0], context.robot.qpos[0]) - assert torch.allclose( - trajectory.positions[1], - context.robot.qpos[1].unsqueeze(0).expand(30, -1), - ) - assert all(not frame.active_mask[1].item() for frame in plan.commands.frames) - received = projected.get_held_object("right_arm") - assert received is not None - assert received.env_mask.tolist() == [True, False] - semantics.entity.get_local_pose.assert_not_called() + assert plan.plan_success.tolist() == [True, True] + pickup_close_end = plan.segment("pickup_close").stop - 1 + positions = trajectory.positions[:, pickup_close_end] + assert torch.all(positions[0, DUAL_ARM_DOF : DUAL_ARM_DOF + 2] == 1) + assert torch.all(positions[0, DUAL_ARM_DOF + 2 :] == 0) + assert torch.all(positions[1, DUAL_ARM_DOF : DUAL_ARM_DOF + 2] == 0) + assert torch.all(positions[1, DUAL_ARM_DOF + 2 :] == 1) -def test_handover_rejects_goal_for_a_different_held_object() -> None: +def test_handover_holds_rows_whose_candidate_arm_is_occupied() -> None: generator = _dual_motion_generator() - action = _bind_action( - generator, - HandOver( - default_options=HandOverOptions( - middle_object_pose=torch.eye(4), - final_object_pose=torch.eye(4), - ) - ), - ) - held_semantics = _semantics(entity_id="held_object") - goal_semantics = _semantics(entity_id="other_object") + action = _bind_action(generator, HandOver()) + semantics, _ = _handover_semantics() + object_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + object_pose[:, 0, 3] = -0.8 + object_pose[:, 1, 3] = 0.2 + occupied = _held(env_mask=torch.tensor([True, False])) task = TaskState( batch_size=NUM_ENVS, device="cpu", - held_objects={"left_arm": _held(held_semantics)}, + held_objects={"left_arm": occupied}, ) invocation = ActionInvocation( skill_id="hand_over", - goal=GraspGoal(semantics=goal_semantics), + goal=HandOverGoal(semantics, target_pose=torch.eye(4)), binding=_dual_binding(action, "source", "destination"), + motion_policy=MotionPolicy(sample_count=24), + skill_options=HandOverOptions(hand_interp_steps=2), ) + context = _handover_context(object_pose, task) - with pytest.raises(ValueError, match="must identify the object held"): - _plan_action(action, invocation, _dual_context(task)) + plan = _plan_action(action, invocation, context) - held_semantics.entity.get_local_pose.assert_not_called() - goal_semantics.entity.get_local_pose.assert_not_called() + assert plan.plan_success.tolist() == [False, True] + trajectory = _joint_trajectory(plan) + assert torch.allclose( + trajectory.positions[0], + context.robot.qpos[0].unsqueeze(0).expand(24, -1), + ) -def test_handover_transfers_only_exclusively_held_rows() -> None: +def test_handover_reports_failed_semantic_waypoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: generator = _dual_motion_generator() - semantics = _semantics() - task = TaskState( - batch_size=NUM_ENVS, - device="cpu", - held_objects={ - "left_arm": _held(semantics), - "right_arm": _held( - semantics, - env_mask=torch.tensor([True, False]), - ), - }, + action = _bind_action(generator, HandOver()) + semantics, _ = _handover_semantics() + object_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + object_pose[:, 0, 3] = -0.8 + object_pose[:, 1, 3] = 0.2 + grasp_call_count = 0 + + def resolve_grasp( + affordance: AntipodalAffordance, + sampled_object_pose: torch.Tensor, + approach_direction: torch.Tensor, + *, + obj_longest_axis: torch.Tensor, + is_positive_part: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + nonlocal grasp_call_count + del affordance, approach_direction, obj_longest_axis, is_positive_part + grasp_call_count += 1 + success = torch.ones(NUM_ENVS, dtype=torch.bool) + if grasp_call_count == 2: + success[0] = False + return sampled_object_pose.clone(), success + + warnings: list[str] = [] + action._resolve_grasp = Mock(side_effect=resolve_grasp) + monkeypatch.setattr( + "embodichain.lab.sim.atomic_actions.primitives.hand_over.logger.log_warning", + warnings.append, ) - action = _bind_action( - generator, - HandOver( - default_options=HandOverOptions( - middle_object_pose=torch.eye(4), - final_object_pose=torch.eye(4), - hand_interp_steps=4, - hold_steps=2, - retreat_steps=5, - ) - ), + invocation = ActionInvocation( + skill_id="hand_over", + goal=HandOverGoal(semantics, target_pose=torch.eye(4)), + binding=_dual_binding(action, "source", "destination"), + motion_policy=MotionPolicy(sample_count=24), + skill_options=HandOverOptions(hand_interp_steps=2), ) - action._resolve_receive_grasp = Mock( - return_value=( - torch.eye(4).repeat(NUM_ENVS, 1, 1), - torch.ones(NUM_ENVS, dtype=torch.bool), + + plan = _plan_action(action, invocation, _handover_context(object_pose)) + + assert plan.plan_success.tolist() == [False, True] + assert any( + "waypoint 'receive_grasp'" in warning and "env_ids=[0]" in warning + for warning in warnings + ) + + +def test_handover_reports_path_failure_between_reachable_waypoints( + monkeypatch: pytest.MonkeyPatch, +) -> None: + generator = _dual_motion_generator() + action = _bind_action(generator, HandOver()) + semantics, _ = _handover_semantics() + object_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + object_pose[:, 0, 3] = -0.8 + object_pose[:, 1, 3] = 0.2 + phase_call_count = 0 + + def plan_phase( + motion_generator: MotionGenerator, + control_part: str, + start_qpos: torch.Tensor, + target_poses: torch.Tensor, + n_waypoints: int, + motion_policy: MotionPolicy, + interpolation_dt: float | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + nonlocal phase_call_count + del ( + motion_generator, + control_part, + target_poses, + motion_policy, + interpolation_dt, ) + phase_call_count += 1 + success = torch.ones(NUM_ENVS, dtype=torch.bool) + if phase_call_count == 1: + success[0] = False + trajectory = start_qpos.unsqueeze(1).repeat(1, n_waypoints, 1) + return success, trajectory + + warnings: list[str] = [] + monkeypatch.setattr( + "embodichain.lab.sim.atomic_actions.primitives.hand_over." + "plan_named_arm_trajectory", + plan_phase, + ) + monkeypatch.setattr( + "embodichain.lab.sim.atomic_actions.primitives.hand_over.logger.log_warning", + warnings.append, ) - context = _dual_context(task) invocation = ActionInvocation( skill_id="hand_over", - goal=GraspGoal(semantics=semantics), + goal=HandOverGoal(semantics, target_pose=torch.eye(4)), binding=_dual_binding(action, "source", "destination"), - motion_policy=MotionPolicy(sample_count=30), + motion_policy=MotionPolicy(sample_count=24), + skill_options=HandOverOptions(hand_interp_steps=2), ) - plan = _plan_action(action, invocation, context) - projected = plan.expected_effects.apply(context.task, plan.plan_success) + plan = _plan_action(action, invocation, _handover_context(object_pose)) assert plan.plan_success.tolist() == [False, True] - trajectory = _joint_trajectory(plan) - assert torch.allclose( - trajectory.positions[0], - context.robot.qpos[0].unsqueeze(0).expand(30, -1), + assert any( + "phase 'pickup_approach' failed between waypoints" in warning + and "pickup_pre_grasp" in warning + and "pickup_grasp" in warning + and "env_ids=[0]" in warning + for warning in warnings ) - transferred = projected.get_held_object("left_arm") - received = projected.get_held_object("right_arm") - assert transferred is not None and transferred.env_mask.tolist() == [True, False] - assert received is not None and received.env_mask.tolist() == [True, True] - assert received.semantics is semantics + + +def test_handover_requires_antipodal_affordance_and_valid_options() -> None: + generator = _dual_motion_generator() + action = _bind_action(generator, HandOver()) + semantics = _semantics(entity_id="target") + invocation = ActionInvocation( + skill_id="hand_over", + goal=HandOverGoal(semantics, target_pose=torch.eye(4)), + binding=_dual_binding(action, "source", "destination"), + ) + + with pytest.raises(ValueError, match="AntipodalAffordance"): + _plan_action( + action, + invocation, + _dual_context( + scene=_target_scene( + torch.eye(4).repeat(NUM_ENVS, 1, 1), + timestamp=0.0, + version=0, + ) + ), + ) + + with pytest.raises(ValueError, match="lift_height"): + HandOverOptions(lift_height=float("nan")) + with pytest.raises(ValueError, match="hand_interp_steps"): + HandOverOptions(hand_interp_steps=0) def test_coordinated_pick_returns_full_dof_plan_and_projected_relation() -> None: diff --git a/tests/sim/atomic_actions/test_affordance.py b/tests/sim/atomic_actions/test_affordance.py index 0e123eefe..683f50e8c 100644 --- a/tests/sim/atomic_actions/test_affordance.py +++ b/tests/sim/atomic_actions/test_affordance.py @@ -23,10 +23,13 @@ import pytest import torch +from embodichain.toolkits.graspkit.pg_grasp import GraspGenerator + from embodichain.lab.sim.atomic_actions.affordance import ( Affordance, AntipodalAffordance, AssembleAffordance, + AxisAlignAffordance, InteractionPoints, PressAffordance, SlideAffordance, @@ -50,6 +53,40 @@ def test_base_get_batch_size_is_one(self): class TestAntipodalAffordance: + @staticmethod + def _long_box_mesh() -> tuple[torch.Tensor, torch.Tensor]: + vertices = torch.tensor( + [ + [-0.1, -0.1, -1.0], + [0.1, -0.1, -1.0], + [0.1, 0.1, -1.0], + [-0.1, 0.1, -1.0], + [-0.1, -0.1, 1.0], + [0.1, -0.1, 1.0], + [0.1, 0.1, 1.0], + [-0.1, 0.1, 1.0], + ], + dtype=torch.float32, + ) + triangles = torch.tensor( + [ + [0, 1, 2], + [0, 2, 3], + [4, 6, 5], + [4, 7, 6], + [0, 4, 5], + [0, 5, 1], + [1, 5, 6], + [1, 6, 2], + [2, 6, 7], + [2, 7, 3], + [3, 7, 4], + [3, 4, 0], + ], + dtype=torch.long, + ) + return vertices, triangles + def test_stores_mesh_fields_directly(self): v = torch.randn(8, 3) t = torch.randint(0, 8, (5, 3)) @@ -104,6 +141,114 @@ def test_valid_grasp_poses_casts_approach_direction_to_generator_device(self): assert approach_direction.dtype == torch.float32 assert approach_direction.device == generator.device + def test_valid_grasp_poses_forwards_batched_axis_end_selection(self): + aff = AntipodalAffordance() + generator = Mock() + generator.device = torch.device("cpu") + generator.get_valid_grasp_poses.return_value = ( + True, + torch.eye(4).unsqueeze(0), + torch.ones(1), + torch.zeros(1), + ) + aff._generator = generator + poses = torch.eye(4).repeat(2, 1, 1) + axes = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, 2.0]]) + + aff.get_valid_grasp_poses( + poses, + obj_longest_axis=axes, + is_positive_part=torch.tensor([True, False]), + ) + + first, second = generator.get_valid_grasp_poses.call_args_list + assert torch.equal( + first.kwargs["obj_longest_axis"], torch.tensor([1.0, 0.0, 0.0]) + ) + assert first.kwargs["is_positive_part"] is True + assert torch.equal( + second.kwargs["obj_longest_axis"], torch.tensor([0.0, 0.0, 1.0]) + ) + assert second.kwargs["is_positive_part"] is False + + def test_surface_svd_uses_at_most_1000_points_in_current_pose(self): + vertices, triangles = self._long_box_mesh() + aff = AntipodalAffordance( + mesh_vertices=vertices, + mesh_triangles=triangles, + ) + first_points = aff.sample_surface_points(max_points=1000) + second_points = aff.sample_surface_points(max_points=1000) + poses = torch.eye(4).repeat(2, 1, 1) + poses[1, :3, :3] = torch.tensor( + [[0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]] + ) + + axes = aff.get_object_longest_axis(poses, max_points=1000) + + assert first_points.shape == (1000, 3) + assert torch.equal(first_points, second_points) + assert torch.abs(axes[0, 2]) > 0.99 + assert torch.abs(axes[1, 0]) > 0.99 + + def test_generator_partitions_pairs_by_axis_projection(self): + generator = object.__new__(GraspGenerator) + generator.device = torch.device("cpu") + centers = torch.tensor( + [ + [-0.75, 0.0, 0.0], + [-0.25, 0.0, 0.0], + [0.25, 0.0, 0.0], + [0.75, 0.0, 0.0], + ] + ) + jaw_offset = torch.tensor([0.0, 0.05, 0.0]) + generator._hit_point_pairs = torch.stack( + [centers - jaw_offset, centers + jaw_offset], dim=1 + ) + generator.vertices = torch.tensor( + [ + [-1.0, -0.1, -0.1], + [-1.0, 0.1, 0.1], + [1.0, -0.1, -0.1], + [1.0, 0.1, 0.1], + ] + ) + expected_result = (True, torch.eye(4), torch.ones(1), torch.zeros(1)) + generator._filter_valid_grasp_poses = Mock(return_value=expected_result) + + positive_result = generator.get_valid_grasp_poses( + torch.eye(4), + torch.tensor([0.0, 0.0, -1.0]), + obj_longest_axis=torch.tensor([1.0, 0.0, 0.0]), + is_positive_part=True, + ) + positive_kwargs = generator._filter_valid_grasp_poses.call_args.kwargs + generator._filter_valid_grasp_poses.reset_mock() + negative_result = generator.get_valid_grasp_poses( + torch.eye(4), + torch.tensor([0.0, 0.0, -1.0]), + obj_longest_axis=torch.tensor([1.0, 0.0, 0.0]), + is_positive_part=False, + ) + negative_kwargs = generator._filter_valid_grasp_poses.call_args.kwargs + generator._filter_valid_grasp_poses.reset_mock() + center_result = generator.get_valid_grasp_poses( + torch.eye(4), + torch.tensor([0.0, 0.0, -1.0]), + obj_longest_axis=None, + ) + center_kwargs = generator._filter_valid_grasp_poses.call_args.kwargs + + assert positive_result is expected_result + assert negative_result is expected_result + assert center_result is expected_result + assert torch.all(positive_kwargs["origin_points_"][:, 0] > 0.0) + assert torch.all(negative_kwargs["origin_points_"][:, 0] < 0.0) + assert positive_kwargs["origin_points_"].shape[0] == 2 + assert negative_kwargs["origin_points_"].shape[0] == 2 + assert center_kwargs["origin_points_"].shape[0] == 4 + def test_best_grasp_poses_casts_approach_direction_to_generator_device(self): aff = AntipodalAffordance() generator = Mock() @@ -121,6 +266,29 @@ def test_best_grasp_poses_casts_approach_direction_to_generator_device(self): assert approach_direction.device == generator.device +class TestAxisAlignAffordance: + def test_extends_antipodal_affordance_with_owned_internal_axis(self): + internal_axis = torch.tensor([1.0, 0.0, 0.0]) + + affordance = AxisAlignAffordance(internal_axis=internal_axis) + internal_axis[0] = 0.0 + + assert isinstance(affordance, AntipodalAffordance) + assert torch.equal(affordance.internal_axis, torch.tensor([1.0, 0.0, 0.0])) + + @pytest.mark.parametrize( + "internal_axis", + ( + torch.zeros(3), + torch.tensor([float("nan"), 0.0, 0.0]), + torch.zeros(2), + ), + ) + def test_rejects_invalid_internal_axis(self, internal_axis): + with pytest.raises(ValueError, match="internal_axis"): + AxisAlignAffordance(internal_axis=internal_axis) + + class TestTwistAffordance: def test_requires_explicit_grasp_position_and_axis_origin(self): with pytest.raises(TypeError, match="grasp_position"): diff --git a/tests/sim/atomic_actions/test_tutorial_utils.py b/tests/sim/atomic_actions/test_tutorial_utils.py index c40b330a8..fe8962ac5 100644 --- a/tests/sim/atomic_actions/test_tutorial_utils.py +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -19,6 +19,7 @@ from __future__ import annotations import importlib +import math from argparse import Namespace from types import SimpleNamespace from unittest.mock import MagicMock, call, patch @@ -73,6 +74,7 @@ ) ATOMIC_ACTION_TUTORIAL_MODULES = ( "assemble", + "axis_align", "coordinated_pickment", "coordinated_placement", "dynamic_obstacle_recovery", @@ -405,6 +407,43 @@ def test_all_atomic_action_tutorials_accept_both_robot_choices( assert franka_args.robot == "franka" +def test_axis_align_tutorial_exposes_upright_and_horizontal_modes() -> None: + module = importlib.import_module("scripts.tutorials.atomic_action.axis_align") + + with patch("sys.argv", ["axis_align.py"]): + upright_args = module.parse_arguments() + with patch( + "sys.argv", + ["axis_align.py", "--alignment", "horizontal_align"], + ): + horizontal_args = module.parse_arguments() + + assert upright_args.alignment == "upright" + assert horizontal_args.alignment == "horizontal_align" + assert module.ALIGNMENT_AXES["upright"] == ( + (1.0, 0.0, 0.0), + (0.0, 0.0, 1.0), + ) + assert module.ALIGNMENT_AXES["horizontal_align"] == ( + (1.0, 0.0, 0.0), + (0.0, 1.0, 0.0), + ) + + +def test_pour_tutorial_uses_configured_pickup_and_local_rotation_axis() -> None: + module = importlib.import_module("scripts.tutorials.atomic_action.pour") + + with patch("sys.argv", ["pour.py"]): + default_args = module.parse_arguments() + with patch("sys.argv", ["pour.py", "--rotate_angle", "-1.25"]): + configured_args = module.parse_arguments() + + assert default_args.rotate_angle == pytest.approx(math.pi / 4.0) + assert configured_args.rotate_angle == pytest.approx(-1.25) + assert module.APPROACH_DIRECTION == pytest.approx((-0.707, 0.0, -0.707)) + assert module.POUR_INTERNAL_AXIS == (1.0, 0.0, 0.0) + + def test_replay_timed_trajectory_uses_arrival_intervals() -> None: sim = MagicMock() sim.sim_config.physics_dt = 0.1