diff --git a/docs/source/_static/atomic_actions/coordinated_placement.gif b/docs/source/_static/atomic_actions/coordinated_placement.gif new file mode 100644 index 000000000..a10785652 Binary files /dev/null and b/docs/source/_static/atomic_actions/coordinated_placement.gif differ 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 783f76170..9b72735b0 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 @@ -17,6 +17,7 @@ embodichain.lab.sim.atomic_actions GraspTarget HeldObjectPoseTarget CoordinatedPickmentTarget + CoordinatedPlacementTarget Target HeldObjectState CoordinatedHeldObjectState @@ -39,6 +40,8 @@ embodichain.lab.sim.atomic_actions Press CoordinatedPickmentCfg CoordinatedPickment + CoordinatedPlacementCfg + CoordinatedPlacement AtomicActionEngine .. currentmodule:: embodichain.lab.sim.atomic_actions @@ -86,6 +89,10 @@ Core :members: :show-inheritance: +.. autoclass:: CoordinatedPlacementTarget + :members: + :show-inheritance: + .. autodata:: Target .. autoclass:: HeldObjectState @@ -185,6 +192,15 @@ Actions :members: :show-inheritance: +.. autoclass:: CoordinatedPlacementCfg + :members: + :exclude-members: __init__, copy, replace, to_dict + :show-inheritance: + +.. autoclass:: CoordinatedPlacement + :members: + :show-inheritance: + Engine & Registry ----------------- diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index 52859c988..40a4983a7 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -16,6 +16,7 @@ The following actions are available out of the box: | `Place` | Single | `EndEffectorPoseTarget` — EEF release pose | Lower → open gripper → retract | Place | | `Press` | Single | `EndEffectorPoseTarget` — EEF press pose | Close gripper → press down → return | Press | | `CoordinatedPickment` | Dual | `CoordinatedPickmentTarget` — shared-object pose | Approach both ends → close both grippers → lift → move object | CoordinatedPickment | +| `CoordinatedPlacement` | Dual | `CoordinatedPlacementTarget` — two held-object poses | Move support object → align placing object → release placing hand → retreat | CoordinatedPlacement | --- @@ -171,3 +172,41 @@ while keeping both grippers closed. On success, the returned `WorldState` carrie semantics, and left/right object-to-EEF transforms. ![CoordinatedPickment demo](../../../_static/atomic_actions/coordinated_pickment.gif) + +--- + +## `CoordinatedPlacement` + +Dual-arm object-centric placement. The support arm moves its held object to a lower +target pose and keeps its gripper closed. The placing arm moves its held object to +the aligned upper target pose, optionally opens the placing hand, then lifts away. + +`CoordinatedPlacement` is intentionally explicit about dual-arm state: the target +contains both `placing_held_object` and `support_held_object`. This avoids relying +on the engine's single `WorldState.held_object` slot to infer two simultaneously +held objects. + +| Config field | Default | Description | +|---|---|---| +| `control_part` | `"dual_arm"` | Robot control part containing both arms | +| `placing_arm_control_part` | `"left_arm"` | Arm that releases the placed object | +| `support_arm_control_part` | `"right_arm"` | Arm that keeps holding the support object | +| `placing_hand_control_part` | `"left_hand"` | Placing gripper control part | +| `support_hand_control_part` | `"right_hand"` | Support gripper control part | +| `placing_hand_open_qpos` | `None` | **Required.** Placing gripper open joint positions | +| `placing_hand_close_qpos` | `None` | **Required.** Placing gripper closed joint positions | +| `support_hand_close_qpos` | `None` | **Required.** Support gripper closed joint positions | +| `release` | `True` | Whether to open the placing gripper | +| `placing_height_offset` | `0.0` | World-Z offset applied to the placing object target pose | +| `support_height_offset` | `0.0` | World-Z offset applied to the support object target pose | +| `lift_height` | `0.08` | Placing-arm lift distance after release (m) | +| `hand_interp_steps` | `10` | Waypoints for placing-hand release | +| `hold_steps` | `4` | Alignment hold waypoints before release | +| `retreat_steps` | `16` | Placing-arm retreat waypoints | +| `sample_interval` | `100` | Total waypoints across all phases | + +**Target:** `CoordinatedPlacementTarget(...)` with placing/support object target +poses plus the corresponding `HeldObjectState` values. On success, the returned +`WorldState.held_object` is the support object's held state. + +![CoordinatedPlacement demo](../../../_static/atomic_actions/coordinated_placement.gif) diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index d57f71c80..78b2438f3 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -30,6 +30,7 @@ EndEffectorPoseTarget(xpos) │ JointPositionTarget(qpos) ├─ IK solve when pose-based NamedJointPositionTarget(name) ├─ Motion plan / interpolation HeldObjectPoseTarget(pose) └─ Gripper interpolation when needed +CoordinatedPlacementTarget(...) │ │ ActionResult (success, full-DoF traj, next_state) @@ -62,6 +63,7 @@ and every action declares the target type, or tuple of target types, it accepts | `NamedJointPositionTarget` | `NamedJointPositionTarget(name)` | `MoveJoints` | | `GraspTarget` | `GraspTarget(semantics)` | `PickUp` | | `HeldObjectPoseTarget` | `HeldObjectPoseTarget(object_target_pose)` | `MoveHeldObject` | +| `CoordinatedPlacementTarget` | `CoordinatedPlacementTarget(...)` | `CoordinatedPlacement` | `Target` is the union of these typed target dataclasses. @@ -99,6 +101,7 @@ action's `TargetType` before calling `execute`: | `NamedJointPositionTarget(name)` | Name resolved from `MoveJointsCfg.named_joint_positions` | `MoveJoints` | | `GraspTarget(semantics)` | `ObjectSemantics` (affordance + entity) | `PickUp` | | `HeldObjectPoseTarget(object_target_pose)` | Desired held-object pose tensor | `MoveHeldObject` | +| `CoordinatedPlacementTarget(...)` | Two held-object states plus object-centric placing/support target poses | `CoordinatedPlacement` | `WorldState` is threaded between actions and carries the robot's `last_qpos` plus an optional `held_object: HeldObjectState`. The built-in actions update it as follows: @@ -108,6 +111,7 @@ action's `TargetType` before calling `execute`: | `PickUp` | Populates it (computed object-to-EEF transform) | | `MoveHeldObject` | Requires it; preserves it unchanged | | `Place` | Clears it to `None` | +| `CoordinatedPlacement` | Returns the support arm's `HeldObjectState`; the placing object is released | | `MoveEndEffector` | Leaves it unchanged | | `MoveJoints` | Leaves it unchanged | | `Press` | Leaves it unchanged | @@ -307,6 +311,7 @@ builtin_actions - `scripts/tutorials/atomic_action/move_held_object.py` - `scripts/tutorials/atomic_action/place.py` - `scripts/tutorials/atomic_action/press.py` + - `scripts/tutorials/atomic_action/coordinated_placement.py` Run a demo in headless CPU mode with `--auto_play --headless --device cpu` to record an MP4 under `outputs/videos`. For example: diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index 279e3230f..1fa8a7bcd 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -33,6 +33,7 @@ AtomicAction, CoordinatedHeldObjectState, CoordinatedPickmentTarget, + CoordinatedPlacementTarget, GraspTarget, HeldObjectState, HeldObjectPoseTarget, @@ -46,6 +47,8 @@ from .actions import ( CoordinatedPickment, CoordinatedPickmentCfg, + CoordinatedPlacement, + CoordinatedPlacementCfg, MoveEndEffector, MoveJoints, MoveHeldObject, @@ -75,6 +78,7 @@ "ObjectSemantics", "CoordinatedHeldObjectState", "CoordinatedPickmentTarget", + "CoordinatedPlacementTarget", "HeldObjectState", "HeldObjectPoseTarget", "JointPositionTarget", @@ -88,6 +92,7 @@ "AtomicAction", # Action implementations "CoordinatedPickment", + "CoordinatedPlacement", "MoveEndEffector", "MoveJoints", "MoveHeldObject", @@ -95,6 +100,7 @@ "Place", "Press", "CoordinatedPickmentCfg", + "CoordinatedPlacementCfg", "MoveEndEffectorCfg", "MoveJointsCfg", "MoveHeldObjectCfg", diff --git a/embodichain/lab/sim/atomic_actions/actions.py b/embodichain/lab/sim/atomic_actions/actions.py index 9436b2706..c4ff096b2 100644 --- a/embodichain/lab/sim/atomic_actions/actions.py +++ b/embodichain/lab/sim/atomic_actions/actions.py @@ -16,8 +16,9 @@ """Concrete atomic actions built on :class:`AtomicAction` and :class:`TrajectoryBuilder`. -Six sibling actions live here: :class:`MoveEndEffector`, :class:`MoveJoints`, -:class:`PickUp`, :class:`MoveHeldObject`, :class:`Place`, and :class:`Press`. +Built-in sibling actions live here: :class:`MoveEndEffector`, :class:`MoveJoints`, +:class:`PickUp`, :class:`MoveHeldObject`, :class:`Place`, :class:`Press`, +:class:`CoordinatedPickment`, and :class:`CoordinatedPlacement`. Each inherits :class:`AtomicAction` directly and composes a :class:`TrajectoryBuilder` for shared trajectory math. ``execute`` takes a typed target plus a @@ -41,6 +42,7 @@ AtomicAction, CoordinatedHeldObjectState, CoordinatedPickmentTarget, + CoordinatedPlacementTarget, GraspTarget, HeldObjectState, HeldObjectPoseTarget, @@ -198,6 +200,60 @@ class CoordinatedPickmentCfg(ActionCfg): """Number of waypoints to hold the final object target pose.""" +@configclass +class CoordinatedPlacementCfg(ActionCfg): + name: str = "coordinated_placement" + """Name of the action, used for identification and logging.""" + + control_part: str = "dual_arm" + """Robot control part containing both placing and support arms.""" + + placing_arm_control_part: str = "left_arm" + """Arm that places and releases its held object.""" + + support_arm_control_part: str = "right_arm" + """Arm that moves the support object and keeps holding it.""" + + placing_hand_control_part: str = "left_hand" + """Hand attached to the placing arm.""" + + support_hand_control_part: str = "right_hand" + """Hand attached to the support arm.""" + + placing_hand_open_qpos: torch.Tensor | None = None + """Placing-hand qpos for the open state, shape ``[hand_dof,]``.""" + + placing_hand_close_qpos: torch.Tensor | None = None + """Placing-hand qpos for the closed state, shape ``[hand_dof,]``.""" + + support_hand_close_qpos: torch.Tensor | None = None + """Support-hand qpos for the closed state, shape ``[hand_dof,]``.""" + + release: bool = True + """Whether to open the placing hand at the aligned placement pose.""" + + placing_height_offset: float = 0.0 + """Default World-Z offset above the placing object target pose.""" + + support_height_offset: float = 0.0 + """Default World-Z offset above the support object target pose.""" + + lift_height: float = 0.08 + """World-Z lift distance for the placing arm after release.""" + + sample_interval: int = 100 + """Number of waypoints for the full coordinated placement trajectory.""" + + hand_interp_steps: int = 10 + """Number of waypoints for the placing-hand release interpolation.""" + + hold_steps: int = 4 + """Number of waypoints to hold alignment before releasing.""" + + retreat_steps: int = 16 + """Number of waypoints used for the placing-arm lift retreat.""" + + @configclass class PressCfg(ActionCfg): name: str = "press" @@ -1597,6 +1653,400 @@ def execute( ) +# ============================================================================= +# CoordinatedPlacement +# ============================================================================= + + +class CoordinatedPlacement(AtomicAction): + """Coordinate two held objects: support object below, placing object above.""" + + TargetType: ClassVar[type] = CoordinatedPlacementTarget + + def __init__( + self, + motion_generator, + cfg: CoordinatedPlacementCfg | None = None, + ) -> None: + super().__init__(motion_generator, cfg or CoordinatedPlacementCfg()) + self.builder = TrajectoryBuilder(motion_generator) + self.n_envs = self.robot.get_qpos().shape[0] + self.robot_dof = self.robot.dof + + self.dual_arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part) + self.placing_arm_joint_ids = self.robot.get_joint_ids( + name=self.cfg.placing_arm_control_part + ) + self.support_arm_joint_ids = self.robot.get_joint_ids( + name=self.cfg.support_arm_control_part + ) + self.placing_hand_joint_ids = self.robot.get_joint_ids( + name=self.cfg.placing_hand_control_part + ) + self.support_hand_joint_ids = self.robot.get_joint_ids( + name=self.cfg.support_hand_control_part + ) + self.joint_ids = ( + self.dual_arm_joint_ids + + self.placing_hand_joint_ids + + self.support_hand_joint_ids + ) + self.placing_arm_dof = len(self.placing_arm_joint_ids) + self.support_arm_dof = len(self.support_arm_joint_ids) + self.placing_hand_dof = len(self.placing_hand_joint_ids) + self.support_hand_dof = len(self.support_hand_joint_ids) + + self._validate_hand_qpos_cfg() + self.placing_hand_open_qpos = self.builder.expand_hand_qpos( + self.cfg.placing_hand_open_qpos, + n_envs=self.n_envs, + hand_dof=self.placing_hand_dof, + ) + self.placing_hand_close_qpos = self.builder.expand_hand_qpos( + self.cfg.placing_hand_close_qpos, + n_envs=self.n_envs, + hand_dof=self.placing_hand_dof, + ) + self.support_hand_close_qpos = self.builder.expand_hand_qpos( + self.cfg.support_hand_close_qpos, + n_envs=self.n_envs, + hand_dof=self.support_hand_dof, + ) + + def execute( + self, target: CoordinatedPlacementTarget, state: WorldState + ) -> ActionResult: + placing_xpos, support_xpos, release, support_held_object = self._resolve_target( + target + ) + placing_start_qpos, support_start_qpos = self._resolve_start_qpos(state) + segments = self._compute_segment_lengths(release) + + placing_lift_xpos = self.builder.apply_local_offset( + placing_xpos, + torch.tensor( + [0.0, 0.0, self.cfg.lift_height], + dtype=torch.float32, + device=self.device, + ), + ) + + ok, placing_approach_traj = self._plan_named_arm_trajectory( + self.cfg.placing_arm_control_part, + placing_start_qpos, + torch.stack([placing_lift_xpos, placing_xpos], dim=1), + segments["approach"], + ) + if not ok: + logger.log_warning("CoordinatedPlacement failed to plan placing approach.") + return self._fail(state) + + ok, support_approach_traj = self._plan_named_arm_trajectory( + self.cfg.support_arm_control_part, + support_start_qpos, + support_xpos.unsqueeze(1), + segments["approach"], + ) + if not ok: + logger.log_warning("CoordinatedPlacement failed to plan support approach.") + return self._fail(state) + + placing_place_qpos = placing_approach_traj[:, -1] + support_place_qpos = support_approach_traj[:, -1] + approach_trajectory = self._assemble_phase( + state.last_qpos, + placing_approach_traj, + support_approach_traj, + self._repeat_qpos(self.placing_hand_close_qpos, segments["approach"]), + self._repeat_qpos(self.support_hand_close_qpos, segments["approach"]), + ) + + hold_trajectory = self._empty_phase(state) + if segments["hold"] > 0: + hold_trajectory = self._assemble_phase( + state.last_qpos, + self._repeat_qpos(placing_place_qpos, segments["hold"]), + self._repeat_qpos(support_place_qpos, segments["hold"]), + self._repeat_qpos(self.placing_hand_close_qpos, segments["hold"]), + self._repeat_qpos(self.support_hand_close_qpos, segments["hold"]), + ) + + release_trajectory = self._empty_phase(state) + if release: + release_trajectory = self._assemble_phase( + state.last_qpos, + self._repeat_qpos(placing_place_qpos, segments["release"]), + self._repeat_qpos(support_place_qpos, segments["release"]), + self.builder.interpolate_hand_qpos( + self.placing_hand_close_qpos, + self.placing_hand_open_qpos, + n_waypoints=segments["release"], + ), + self._repeat_qpos(self.support_hand_close_qpos, segments["release"]), + ) + + ok, placing_retreat_traj = self._plan_named_arm_trajectory( + self.cfg.placing_arm_control_part, + placing_place_qpos, + placing_lift_xpos.unsqueeze(1), + segments["retreat"], + ) + if not ok: + logger.log_warning("CoordinatedPlacement failed to plan placing retreat.") + return self._fail(state) + + placing_hand_retreat_qpos = ( + self.placing_hand_open_qpos if release else self.placing_hand_close_qpos + ) + retreat_trajectory = self._assemble_phase( + state.last_qpos, + placing_retreat_traj, + self._repeat_qpos(support_place_qpos, segments["retreat"]), + self._repeat_qpos(placing_hand_retreat_qpos, segments["retreat"]), + self._repeat_qpos(self.support_hand_close_qpos, segments["retreat"]), + ) + + full = torch.cat( + [ + approach_trajectory, + hold_trajectory, + release_trajectory, + retreat_trajectory, + ], + dim=1, + ) + return ActionResult( + success=True, + trajectory=full, + next_state=WorldState( + last_qpos=full[:, -1, :].clone(), + held_object=support_held_object, + ), + ) + + def _validate_hand_qpos_cfg(self) -> None: + """Ensure all hand state tensors are provided.""" + required_names = ( + "placing_hand_open_qpos", + "placing_hand_close_qpos", + "support_hand_close_qpos", + ) + for name in required_names: + if getattr(self.cfg, name) is None: + logger.log_error( + f"{name} must be specified in CoordinatedPlacementCfg", + ValueError, + ) + + def _resolve_object_pose( + self, + pose: torch.Tensor, + height_offset: float, + name: str, + ) -> torch.Tensor: + """Resolve an object target pose into a batched pose with height offset.""" + object_pose = _resolve_object_target( + pose, + n_envs=self.n_envs, + device=self.device, + ) + return self.builder.apply_local_offset( + object_pose, + torch.tensor( + [0.0, 0.0, height_offset], + dtype=torch.float32, + device=self.device, + ), + ) + + def _resolve_object_to_eef( + self, + held_state: HeldObjectState, + name: str, + ) -> torch.Tensor: + """Resolve a held-object transform into batched shape.""" + return self._resolve_held_matrix( + held_state.object_to_eef, + f"{name}.object_to_eef", + ) + + def _resolve_held_matrix(self, matrix: torch.Tensor, name: str) -> torch.Tensor: + """Resolve a held-object pose-like matrix into batched shape.""" + matrix = matrix.to(device=self.device, dtype=torch.float32) + if matrix.shape == (4, 4): + matrix = matrix.unsqueeze(0).repeat(self.n_envs, 1, 1) + if matrix.shape != (self.n_envs, 4, 4): + logger.log_error( + f"{name} must have shape (4, 4) or ({self.n_envs}, 4, 4), " + f"but got {matrix.shape}", + ValueError, + ) + return matrix + + def _resolve_held_state( + self, + held_state: HeldObjectState, + name: str, + object_to_eef: torch.Tensor, + ) -> HeldObjectState: + """Return a held-object state normalized to batched tensor shapes.""" + return HeldObjectState( + semantics=held_state.semantics, + object_to_eef=object_to_eef, + grasp_xpos=self._resolve_held_matrix( + held_state.grasp_xpos, + f"{name}.grasp_xpos", + ), + ) + + def _resolve_target( + self, + target: CoordinatedPlacementTarget, + ) -> tuple[torch.Tensor, torch.Tensor, bool, HeldObjectState]: + """Resolve object-centric target into placing and support TCP poses.""" + placing_height_offset = ( + self.cfg.placing_height_offset + if target.placing_height_offset is None + else target.placing_height_offset + ) + support_height_offset = ( + self.cfg.support_height_offset + if target.support_height_offset is None + else target.support_height_offset + ) + placing_object_pose = self._resolve_object_pose( + target.placing_object_target_pose, + placing_height_offset, + "placing_object_target_pose", + ) + support_object_pose = self._resolve_object_pose( + target.support_object_target_pose, + support_height_offset, + "support_object_target_pose", + ) + placing_object_to_eef = self._resolve_object_to_eef( + target.placing_held_object, + "placing_held_object", + ) + support_object_to_eef = self._resolve_object_to_eef( + target.support_held_object, + "support_held_object", + ) + placing_xpos = torch.bmm(placing_object_pose, placing_object_to_eef) + support_xpos = torch.bmm(support_object_pose, support_object_to_eef) + release = self.cfg.release if target.release is None else target.release + return ( + placing_xpos, + support_xpos, + release, + self._resolve_held_state( + target.support_held_object, + "support_held_object", + support_object_to_eef, + ), + ) + + def _resolve_start_qpos( + self, state: WorldState + ) -> tuple[torch.Tensor, torch.Tensor]: + """Extract per-arm start qpos from full-robot WorldState qpos.""" + if state.last_qpos.shape != (self.n_envs, self.robot_dof): + logger.log_error( + f"WorldState.last_qpos must have shape ({self.n_envs}, {self.robot_dof}), " + f"but got {state.last_qpos.shape}", + ValueError, + ) + start_qpos = state.last_qpos.to(device=self.device, dtype=torch.float32) + return ( + start_qpos[:, self.placing_arm_joint_ids], + start_qpos[:, self.support_arm_joint_ids], + ) + + def _compute_segment_lengths(self, release: bool) -> dict[str, int]: + """Compute waypoint counts for coordinated placement phases.""" + n_release = max(2, self.cfg.hand_interp_steps) if release else 0 + n_hold = max(0, self.cfg.hold_steps) + n_retreat = max(2, self.cfg.retreat_steps) + n_approach = self.cfg.sample_interval - n_hold - n_release - n_retreat + if n_approach < 2: + logger.log_error( + "Not enough waypoints for coordinated placement. Increase " + "sample_interval or decrease hold/release/retreat steps.", + ValueError, + ) + return { + "approach": n_approach, + "hold": n_hold, + "release": n_release, + "retreat": n_retreat, + } + + def _plan_named_arm_trajectory( + self, + control_part: str, + start_qpos: torch.Tensor, + target_poses: torch.Tensor, + n_waypoints: int, + ) -> tuple[bool, torch.Tensor]: + """Plan a batched arm trajectory for a named control part.""" + target_states_list = [ + [ + PlanState(xpos=target_poses[i, j], move_type=MoveType.EEF_MOVE) + for j in range(target_poses.shape[1]) + ] + for i in range(self.n_envs) + ] + return self.builder.plan_arm_traj( + target_states_list, + start_qpos, + n_waypoints, + control_part=control_part, + arm_dof=start_qpos.shape[-1], + ) + + @staticmethod + def _repeat_qpos(qpos: torch.Tensor, n_waypoints: int) -> torch.Tensor: + """Repeat batched qpos across waypoints.""" + return qpos.unsqueeze(1).repeat(1, n_waypoints, 1) + + def _empty_phase(self, state: WorldState) -> torch.Tensor: + """Return an empty full-DoF phase for optional segments.""" + return torch.empty( + (self.n_envs, 0, self.robot_dof), + dtype=torch.float32, + device=self.device, + ) + + def _assemble_phase( + self, + base_full_qpos: torch.Tensor, + placing_arm_traj: torch.Tensor, + support_arm_traj: torch.Tensor, + placing_hand_traj: torch.Tensor, + support_hand_traj: torch.Tensor, + ) -> torch.Tensor: + """Assemble arm and hand trajectories into full-robot DoF order.""" + n_waypoints = placing_arm_traj.shape[1] + full = base_full_qpos.to(device=self.device, dtype=torch.float32) + full = full.unsqueeze(1).repeat(1, n_waypoints, 1).clone() + full[:, :, self.placing_arm_joint_ids] = placing_arm_traj + full[:, :, self.support_arm_joint_ids] = support_arm_traj + full[:, :, self.placing_hand_joint_ids] = placing_hand_traj + full[:, :, self.support_hand_joint_ids] = support_hand_traj + return full + + def _fail(self, state: WorldState) -> ActionResult: + return ActionResult( + success=False, + trajectory=torch.empty( + (self.n_envs, 0, self.robot_dof), + dtype=torch.float32, + device=self.device, + ), + next_state=state, + ) + + # ============================================================================= # Press # ============================================================================= @@ -1730,6 +2180,8 @@ def _fail(self, state: WorldState) -> ActionResult: __all__ = [ "CoordinatedPickment", "CoordinatedPickmentCfg", + "CoordinatedPlacement", + "CoordinatedPlacementCfg", "MoveEndEffector", "MoveEndEffectorCfg", "MoveJoints", diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index 4c937a0e8..bf6e5a6a1 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -139,16 +139,6 @@ class CoordinatedPickmentTarget: """Optional initial object pose. Defaults to ``object_semantics.entity`` pose.""" -Target = ( - EndEffectorPoseTarget - | JointPositionTarget - | NamedJointPositionTarget - | GraspTarget - | HeldObjectPoseTarget - | CoordinatedPickmentTarget -) - - # ============================================================================= # World state threaded between actions # ============================================================================= @@ -188,6 +178,47 @@ class CoordinatedHeldObjectState: """Right end-effector grasp pose for the shared object, shape [n_envs, 4, 4].""" +@dataclass(frozen=True) +class CoordinatedPlacementTarget: + """Object-centric target for dual-arm coordinated placement. + + The placing arm moves its held object to the upper target and releases it. + The support arm moves its held object to the lower target and keeps holding. + """ + + placing_object_target_pose: torch.Tensor + """Target pose for the object released by the placing arm.""" + + support_object_target_pose: torch.Tensor + """Target pose for the object held by the support arm.""" + + placing_held_object: HeldObjectState + """Held-object state for the placing arm.""" + + support_held_object: HeldObjectState + """Held-object state for the support arm.""" + + placing_height_offset: float | None = None + """World-Z offset above the placing object target pose.""" + + support_height_offset: float | None = None + """World-Z offset above the support object target pose.""" + + release: bool | None = None + """Whether the placing hand releases. ``None`` uses action config.""" + + +Target = ( + EndEffectorPoseTarget + | JointPositionTarget + | NamedJointPositionTarget + | GraspTarget + | HeldObjectPoseTarget + | CoordinatedPickmentTarget + | CoordinatedPlacementTarget +) + + @dataclass class WorldState: """State the engine threads through a sequence of actions.""" @@ -278,6 +309,7 @@ def execute(self, target: Target, state: WorldState) -> ActionResult: "AtomicAction", "CoordinatedHeldObjectState", "CoordinatedPickmentTarget", + "CoordinatedPlacementTarget", "GraspTarget", "HeldObjectState", "HeldObjectPoseTarget", diff --git a/scripts/benchmark/__main__.py b/scripts/benchmark/__main__.py index b8b24150e..7b197ead8 100644 --- a/scripts/benchmark/__main__.py +++ b/scripts/benchmark/__main__.py @@ -22,6 +22,7 @@ python -m scripts.benchmark rl --rebuild-report-only python -m scripts.benchmark robotics-kinematic-solver -s pytorch python -m scripts.benchmark planners-neural-planner --num-waypoints 1 3 5 + python -m scripts.benchmark atomic-action --smoke """ from __future__ import annotations @@ -66,6 +67,13 @@ def _run_neural_planner_cli(args: argparse.Namespace) -> None: ) +def _run_atomic_action_cli(_: argparse.Namespace) -> None: + """Run atomic action benchmark CLI entrypoint.""" + from scripts.benchmark.atomic_action.run_benchmark import main as atomic_main + + atomic_main() + + def main() -> None: """Dispatch to the appropriate benchmark sub-command CLI.""" parser = argparse.ArgumentParser( @@ -167,6 +175,16 @@ def main() -> None: ) neural_planner_parser.set_defaults(func=_run_neural_planner_cli) + # -- atomic-action ------------------------------------------------------- + atomic_action_parser = subparsers.add_parser( + "atomic-action", + help="Benchmark atomic actions over object presets and positions.", + ) + from scripts.benchmark.atomic_action.run_benchmark import add_benchmark_args + + add_benchmark_args(atomic_action_parser) + atomic_action_parser.set_defaults(func=_run_atomic_action_cli) + # -- Parse --------------------------------------------------------------- # If no sub-command is given, print help and exit. if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help"): @@ -193,3 +211,6 @@ def main() -> None: if __name__ == "__main__": main() + + +__all__ = ["main"] diff --git a/scripts/benchmark/atomic_action/__init__.py b/scripts/benchmark/atomic_action/__init__.py new file mode 100644 index 000000000..4709da3d9 --- /dev/null +++ b/scripts/benchmark/atomic_action/__init__.py @@ -0,0 +1,21 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Benchmarks for atomic actions.""" + +from __future__ import annotations + +__all__ = [] diff --git a/scripts/benchmark/atomic_action/common.py b/scripts/benchmark/atomic_action/common.py new file mode 100644 index 000000000..7220ab625 --- /dev/null +++ b/scripts/benchmark/atomic_action/common.py @@ -0,0 +1,1194 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Shared helpers for atomic-action benchmark scripts.""" + +from __future__ import annotations + +import argparse +import math +import os +import re +import resource +import sys +import time +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Callable + +try: + import psutil +except ModuleNotFoundError: + psutil = None + +CPU_MEMORY_BACKEND = "psutil" if psutil is not None else "resource" +DEFAULT_VIDEO_DIR = Path("outputs/benchmark_videos") +DEFAULT_VIDEO_FPS = 20 +DEFAULT_VIDEO_MAX_MEMORY_MB = 2048 +DEFAULT_VIDEO_WIDTH = 640 +DEFAULT_VIDEO_HEIGHT = 480 +DEFAULT_VIDEO_HOLD_STEPS = 120 +DEFAULT_VIDEO_CASE_LIMIT = 0 +PHYSICAL_PICK_MIN_LIFT_M = 0.04 +PHYSICAL_PLACE_XY_TOLERANCE_M = 0.10 +PHYSICAL_MOVE_HELD_OBJECT_XYZ_TOLERANCE_M = 0.12 +PHYSICAL_VALIDATION_HOLD_STEPS = 80 +SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z = 0.35 +SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT = 0.5 +BENCHMARK_PROFILES = ("smoke", "coverage", "full") +DEFAULT_BENCHMARK_PROFILE = "coverage" +DEFAULT_VIDEO_LOOK_AT = ( + (-1.25, -1.15, 0.95), + (-0.25, -0.02, 0.25), + (0.0, 0.0, 1.0), +) + + +@dataclass(frozen=True) +class PositionCase: + """Initial object position case with a quadrant label.""" + + name: str + quadrant: str + xy: tuple[float, float] + + +@dataclass(frozen=True) +class MeshObjectPreset: + """Real mesh object preset used by object-conditioned benchmarks.""" + + object_type: str + material_name: str + label: str + init_rot: tuple[float, float, float] + body_scale: tuple[float, float, float] + mass: float + initial_z: float + mesh_path: str = "" + shape_type: str = "mesh" + cube_size: tuple[float, float, float] | None = None + use_usd_properties: bool = False + dynamic_friction: float = 0.97 + static_friction: float = 0.99 + restitution: float = 0.0 + contact_offset: float = 0.002 + rest_offset: float = 0.0 + linear_damping: float = 0.7 + angular_damping: float = 0.7 + max_depenetration_velocity: float = 10.0 + min_position_iters: int = 4 + min_velocity_iters: int = 1 + max_linear_velocity: float = 100.0 + max_angular_velocity: float = 100.0 + max_convex_hull_num: int = 16 + enable_ccd: bool = False + + +POSITION_CASES: dict[str, PositionCase] = { + "q1_near": PositionCase(name="q1_near", quadrant="q1", xy=(0.02, 0.18)), + "q1_far": PositionCase(name="q1_far", quadrant="q1", xy=(0.12, 0.36)), + "q2_near": PositionCase(name="q2_near", quadrant="q2", xy=(-0.42, 0.18)), + "q2_far": PositionCase(name="q2_far", quadrant="q2", xy=(-0.62, 0.36)), + "q3_near": PositionCase(name="q3_near", quadrant="q3", xy=(-0.42, -0.18)), + "q3_far": PositionCase(name="q3_far", quadrant="q3", xy=(-0.62, -0.36)), + "q4_near": PositionCase(name="q4_near", quadrant="q4", xy=(0.02, -0.18)), + "q4_far": PositionCase(name="q4_far", quadrant="q4", xy=(0.12, -0.36)), +} +FULL_POSITION_CASE_NAMES = tuple(POSITION_CASES.keys()) +COVERAGE_POSITION_CASE_NAMES = FULL_POSITION_CASE_NAMES +SMOKE_POSITION_CASE_NAMES = ("q3_near",) + +MESH_OBJECT_PRESETS: dict[str, MeshObjectPreset] = { + "sugar_box": MeshObjectPreset( + object_type="sugar_box", + material_name="cardboard", + label="sugar_box", + mesh_path="SugarBox/sugar_box_usd/sugar_box.usda", + init_rot=(0.0, 0.0, 0.0), + body_scale=(0.8, 0.8, 0.8), + mass=0.05, + initial_z=0.05, + use_usd_properties=False, + ), + "coffee_cup": MeshObjectPreset( + object_type="coffee_cup", + material_name="ceramic", + label="coffee_cup", + mesh_path="CoffeeCup/cup.ply", + init_rot=(0.0, 0.0, -90.0), + body_scale=(4.0, 4.0, 4.0), + mass=0.01, + initial_z=0.01, + use_usd_properties=False, + ), + "cube": MeshObjectPreset( + object_type="cube", + material_name="plastic", + label="cube", + shape_type="cube", + cube_size=(0.05, 0.05, 0.05), + init_rot=(0.0, 0.0, 0.0), + body_scale=(1.0, 1.0, 1.0), + mass=0.05, + initial_z=0.05, + use_usd_properties=False, + dynamic_friction=0.5, + static_friction=0.5, + contact_offset=0.003, + rest_offset=0.001, + max_depenetration_velocity=10.0, + min_position_iters=32, + min_velocity_iters=8, + max_convex_hull_num=1, + ), + "paper_cup": MeshObjectPreset( + object_type="paper_cup", + material_name="paper", + label="paper_cup", + mesh_path="PaperCup/paper_cup.ply", + init_rot=(0.0, 0.0, 0.0), + body_scale=(0.75, 0.75, 1.0), + mass=0.01, + initial_z=0.05, + use_usd_properties=False, + dynamic_friction=1.0, + static_friction=1.0, + contact_offset=0.003, + rest_offset=0.001, + linear_damping=2.0, + angular_damping=2.0, + max_depenetration_velocity=2.0, + min_position_iters=32, + min_velocity_iters=8, + max_linear_velocity=5.0, + max_angular_velocity=10.0, + max_convex_hull_num=8, + ), + "scanned_bottle": MeshObjectPreset( + object_type="scanned_bottle", + material_name="plastic", + label="scanned_bottle", + mesh_path="ScannedBottle/yibao_processed.ply", + init_rot=(0.0, 0.0, 0.0), + body_scale=(1.0, 1.0, 1.0), + mass=0.05, + initial_z=0.05, + use_usd_properties=False, + ), +} +COVERAGE_MESH_OBJECT_TYPES = ("sugar_box", "cube", "paper_cup") +FULL_MESH_OBJECT_TYPES = COVERAGE_MESH_OBJECT_TYPES +SMOKE_MESH_OBJECT_TYPES = ("sugar_box",) +PICKUP_APPROACH_CASES = ("top", "side") +SMOKE_PICKUP_APPROACH_CASES = ("top",) + + +def ensure_repo_root() -> None: + """Add the repository root to sys.path for module execution.""" + repo_root = Path(__file__).resolve().parents[3] + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + + +def ensure_torch(): + """Import torch or raise a clear benchmark runtime error.""" + try: + import torch + except ModuleNotFoundError as exc: + raise RuntimeError( + "Atomic action benchmark requires the EmbodiChain simulation runtime " + f"and PyTorch. Missing module: {exc.name}." + ) from exc + return torch + + +def add_common_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add common atomic-action benchmark CLI arguments.""" + add_profile_benchmark_args(parser) + parser.add_argument( + "--repeat", + type=int, + default=1, + help="Number of repeats for every benchmark case.", + ) + parser.add_argument( + "--smoke", + action="store_true", + help="Alias for --profile smoke.", + ) + parser.add_argument( + "--device", + type=str, + default="cpu", + help="Simulation device, e.g. 'cpu' or 'cuda'.", + ) + parser.add_argument( + "--renderer", + type=str, + choices=("auto", "hybrid", "fast-rt", "rt"), + default="auto", + help="Renderer backend used by SimulationManager.", + ) + add_video_benchmark_args(parser) + + +def add_video_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add optional benchmark video recording CLI arguments.""" + parser.add_argument( + "--record_video", + action="store_true", + help="Record trajectory replay videos for selected successful cases.", + ) + parser.add_argument( + "--record_failed_video", + action="store_true", + help=( + "With --record_video, also record failed cases when a partial " + "trajectory or static debug scene is available." + ), + ) + parser.add_argument( + "--video_case_limit", + type=int, + default=DEFAULT_VIDEO_CASE_LIMIT, + help="Maximum cases to record. Use 0 to record all selected cases.", + ) + parser.add_argument( + "--video_dir", + type=Path, + default=DEFAULT_VIDEO_DIR, + help="Directory for benchmark replay videos.", + ) + parser.add_argument( + "--video_fps", + type=int, + default=DEFAULT_VIDEO_FPS, + help="Recorded video frames per second.", + ) + parser.add_argument( + "--video_max_memory", + type=int, + default=DEFAULT_VIDEO_MAX_MEMORY_MB, + help="Maximum recorder frame-buffer memory in MB.", + ) + parser.add_argument( + "--video_width", + type=int, + default=DEFAULT_VIDEO_WIDTH, + help="Recorded video width in pixels.", + ) + parser.add_argument( + "--video_height", + type=int, + default=DEFAULT_VIDEO_HEIGHT, + help="Recorded video height in pixels.", + ) + parser.add_argument( + "--video_hold_steps", + type=int, + default=DEFAULT_VIDEO_HOLD_STEPS, + help="Extra simulation steps to hold the final replay pose.", + ) + + +def add_grasp_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add common grasp-affordance setup arguments.""" + parser.add_argument( + "--n_sample", + type=int, + default=10000, + help="Number of samples for antipodal grasp generation.", + ) + parser.add_argument( + "--force_reannotate", + action="store_true", + help="Force grasp region re-annotation instead of using cached data.", + ) + + +def add_profile_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add unified benchmark profile CLI arguments.""" + parser.add_argument( + "--profile", + choices=BENCHMARK_PROFILES, + default=DEFAULT_BENCHMARK_PROFILE, + help=( + "Benchmark profile: smoke is one fast case, coverage is the default " + "core object/position matrix, and full sweeps all configured cases." + ), + ) + + +def add_object_position_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add object and initial-position selection arguments.""" + parser.add_argument( + "--object_types", + nargs="+", + choices=(*MESH_OBJECT_PRESETS.keys(), "all"), + default=None, + help=( + "Real mesh object presets to benchmark. Defaults are selected by " + "--profile; use 'all' to include every default full preset." + ), + ) + parser.add_argument( + "--position_cases", + nargs="+", + choices=(*POSITION_CASES.keys(), "all"), + default=None, + help=( + "Initial object position cases to benchmark. Defaults are selected by " + "--profile; use 'all' for all near/far cases." + ), + ) + + +def add_pickup_approach_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add common PickUp approach selection arguments.""" + parser.add_argument( + "--approach_cases", + nargs="+", + choices=(*PICKUP_APPROACH_CASES, "all"), + default=None, + help=( + "PickUp approach cases to benchmark. Defaults are selected by " + "--profile; side uses the current object's initial XY direction." + ), + ) + + +def resolve_profile(args: argparse.Namespace) -> str: + """Resolve the effective benchmark profile from CLI arguments.""" + profile = getattr(args, "profile", DEFAULT_BENCHMARK_PROFILE) + if getattr(args, "smoke", False): + profile = "smoke" + if profile not in BENCHMARK_PROFILES: + raise ValueError( + f"Unsupported benchmark profile {profile!r}. " + f"Expected one of {BENCHMARK_PROFILES}." + ) + return profile + + +def default_position_case_names_for_profile(profile: str) -> tuple[str, ...]: + """Return default position case names for a profile.""" + if profile == "smoke": + return SMOKE_POSITION_CASE_NAMES + if profile == "coverage": + return COVERAGE_POSITION_CASE_NAMES + if profile == "full": + return FULL_POSITION_CASE_NAMES + raise ValueError(f"Unsupported benchmark profile: {profile}") + + +def select_position_cases( + case_names: Sequence[str] | None, + profile: str, +) -> list[PositionCase]: + """Resolve requested position cases, falling back to profile defaults.""" + names = ( + default_position_case_names_for_profile(profile) + if not case_names + else tuple(case_names) + ) + if "all" in names: + names = FULL_POSITION_CASE_NAMES + return [POSITION_CASES[name] for name in names] + + +def default_mesh_object_types_for_profile(profile: str) -> tuple[str, ...]: + """Return default real mesh object preset names for a profile.""" + if profile == "smoke": + return SMOKE_MESH_OBJECT_TYPES + if profile == "coverage": + return COVERAGE_MESH_OBJECT_TYPES + if profile == "full": + return FULL_MESH_OBJECT_TYPES + raise ValueError(f"Unsupported benchmark profile: {profile}") + + +def select_mesh_object_presets( + object_types: Sequence[str] | None, + profile: str, +) -> list[MeshObjectPreset]: + """Resolve requested mesh object presets, falling back to profile defaults.""" + names = ( + default_mesh_object_types_for_profile(profile) + if not object_types + else tuple(object_types) + ) + if "all" in names: + names = FULL_MESH_OBJECT_TYPES + return [MESH_OBJECT_PRESETS[name] for name in names] + + +def default_pickup_approach_cases_for_profile(profile: str) -> tuple[str, ...]: + """Return default PickUp approach case names for a profile.""" + if profile == "smoke": + return SMOKE_PICKUP_APPROACH_CASES + if profile in ("coverage", "full"): + return PICKUP_APPROACH_CASES + raise ValueError(f"Unsupported benchmark profile: {profile}") + + +def select_pickup_approaches( + approach_cases: Sequence[str] | None, + profile: str, +) -> list[str]: + """Resolve PickUp approach cases, falling back to profile defaults.""" + names = ( + default_pickup_approach_cases_for_profile(profile) + if not approach_cases + else tuple(approach_cases) + ) + if "all" in names: + names = PICKUP_APPROACH_CASES + return list(names) + + +def pickup_approach_direction_tuple( + approach: str, + position_case: PositionCase, +) -> tuple[float, float, float]: + """Resolve a PickUp approach name into a normalized world-frame tuple.""" + if approach == "top": + direction = (0.0, 0.0, -1.0) + elif approach == "side": + direction = (-position_case.xy[0], -position_case.xy[1], 0.0) + else: + raise ValueError(f"Unsupported PickUp approach case: {approach}") + + norm = math.sqrt(sum(value * value for value in direction)) + if norm < 1e-6: + raise ValueError(f"PickUp approach direction is zero for {position_case.name}.") + return tuple(value / norm for value in direction) + + +def resolve_pickup_approach_direction( + approach: str, + position_case: PositionCase, + device, +): + """Resolve a PickUp approach name into a normalized world-frame vector.""" + torch = ensure_torch() + return torch.tensor( + pickup_approach_direction_tuple(approach, position_case), + dtype=torch.float32, + device=device, + ) + + +def format_vector3(vector) -> str: + """Format a 3D vector-like value for benchmark reports.""" + if hasattr(vector, "detach"): + vector = vector.detach().to("cpu").tolist() + return f"({float(vector[0]):.3f},{float(vector[1]):.3f},{float(vector[2]):.3f})" + + +def _is_horizontal_approach_direction(approach_direction) -> bool: + """Return true when the requested approach is a side/horizontal approach.""" + if not hasattr(approach_direction, "detach"): + return abs(float(approach_direction[2])) < 1e-4 + direction = approach_direction.detach() + if direction.ndim > 1: + direction = direction.reshape(-1, direction.shape[-1])[0] + return abs(float(direction[2].to("cpu"))) < 1e-4 + + +def create_benchmark_object( + sim, + preset: MeshObjectPreset, + position_case: PositionCase, + uid_suffix: str, +): + """Create one benchmark object at a selected initial position.""" + from embodichain.data import get_data_path + from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg + from embodichain.lab.sim.shapes import CubeCfg, MeshCfg + + if preset.shape_type == "mesh": + shape = MeshCfg(fpath=get_data_path(preset.mesh_path)) + elif preset.shape_type == "cube": + if preset.cube_size is None: + raise ValueError(f"Cube preset {preset.object_type!r} misses cube_size.") + shape = CubeCfg(size=list(preset.cube_size)) + else: + raise ValueError( + f"Unsupported benchmark object shape_type {preset.shape_type!r}." + ) + + cfg = RigidObjectCfg( + uid=f"benchmark_{preset.label}_{position_case.name}_{uid_suffix}", + shape=shape, + attrs=RigidBodyAttributesCfg( + mass=preset.mass, + dynamic_friction=preset.dynamic_friction, + static_friction=preset.static_friction, + restitution=preset.restitution, + contact_offset=preset.contact_offset, + rest_offset=preset.rest_offset, + linear_damping=preset.linear_damping, + angular_damping=preset.angular_damping, + max_depenetration_velocity=preset.max_depenetration_velocity, + min_position_iters=preset.min_position_iters, + min_velocity_iters=preset.min_velocity_iters, + max_linear_velocity=preset.max_linear_velocity, + max_angular_velocity=preset.max_angular_velocity, + enable_ccd=preset.enable_ccd, + ), + max_convex_hull_num=preset.max_convex_hull_num, + init_pos=[position_case.xy[0], position_case.xy[1], preset.initial_z], + init_rot=preset.init_rot, + body_scale=preset.body_scale, + use_usd_properties=preset.use_usd_properties, + ) + obj = sim.add_rigid_object(cfg=cfg) + sim.update(step=10) + return obj + + +create_mesh_benchmark_object = create_benchmark_object + + +def _make_benchmark_antipodal_affordance_class(): + """Create an AntipodalAffordance subclass after project imports are available.""" + from embodichain.lab.sim.atomic_actions import AntipodalAffordance + + class BenchmarkAntipodalAffordance(AntipodalAffordance): + """Benchmark affordance that biases side grasps to horizontal closing.""" + + def get_valid_grasp_poses( + self, + obj_poses, + approach_direction, + ): + results = super().get_valid_grasp_poses( + obj_poses=obj_poses, + approach_direction=approach_direction, + ) + if not _is_horizontal_approach_direction(approach_direction): + return results + + adjusted_results = [] + for grasp_poses, costs in results: + if grasp_poses.ndim < 3 or grasp_poses.shape[0] == 0: + adjusted_results.append((grasp_poses, costs)) + continue + + opening_axis_abs_z = grasp_poses[:, :3, 0].abs()[:, 2] + keep = opening_axis_abs_z <= SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z + if bool(keep.any()): + adjusted_results.append((grasp_poses[keep], costs[keep])) + continue + + adjusted_costs = costs + ( + opening_axis_abs_z * SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT + ) + adjusted_results.append((grasp_poses, adjusted_costs)) + return adjusted_results + + return BenchmarkAntipodalAffordance + + +def create_antipodal_object_semantics( + obj, + preset: MeshObjectPreset, + args: argparse.Namespace, + build_gripper_collision_cfg: Callable[[], object], + build_grasp_generator_cfg: Callable[[argparse.Namespace], object], +): + """Create object semantics with an antipodal grasp affordance.""" + from embodichain.lab.sim.atomic_actions import ObjectSemantics + + mesh_vertices = obj.get_vertices(env_ids=[0], scale=True)[0] + mesh_triangles = obj.get_triangles(env_ids=[0])[0] + affordance_cls = _make_benchmark_antipodal_affordance_class() + return ObjectSemantics( + label=preset.label, + geometry={ + "mesh_vertices": mesh_vertices, + "mesh_triangles": mesh_triangles, + }, + affordance=affordance_cls( + mesh_vertices=mesh_vertices, + mesh_triangles=mesh_triangles, + gripper_collision_cfg=build_gripper_collision_cfg(), + generator_cfg=build_grasp_generator_cfg(args), + force_reannotate=args.force_reannotate, + ), + entity=obj, + ) + + +def describe_object_preset(preset: MeshObjectPreset) -> str: + """Describe an object preset for benchmark report notes.""" + if preset.shape_type == "cube": + return ( + f"{preset.object_type}/{preset.material_name}/" + f"CubeCfg(size={preset.cube_size})" + ) + return f"{preset.object_type}/{preset.material_name}/{preset.mesh_path}" + + +def sync_cuda() -> None: + """Synchronize CUDA stream when available.""" + torch = ensure_torch() + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def reset_peak_gpu_memory() -> None: + """Reset PyTorch peak GPU memory stats when CUDA is available.""" + torch = ensure_torch() + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + + +def peak_gpu_memory_mb() -> float: + """Return peak GPU memory allocated by PyTorch in MB.""" + torch = ensure_torch() + if not torch.cuda.is_available(): + return 0.0 + return torch.cuda.max_memory_allocated() / 1024**2 + + +def memory_snapshot() -> dict[str, float]: + """Return current process memory usage snapshot in MB.""" + torch = ensure_torch() + if psutil is not None: + process = psutil.Process(os.getpid()) + cpu_mb = process.memory_info().rss / 1024**2 + else: + cpu_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0 + gpu_mb = ( + torch.cuda.memory_allocated() / 1024**2 if torch.cuda.is_available() else 0.0 + ) + return {"cpu_mb": cpu_mb, "gpu_mb": gpu_mb} + + +def timed_call( + callable_fn: Callable[[], object], +) -> tuple[float, dict[str, float], float, object]: + """Time a callable and return elapsed seconds, memory deltas, peak GPU, result.""" + reset_peak_gpu_memory() + before = memory_snapshot() + sync_cuda() + + start = time.perf_counter() + result = callable_fn() + sync_cuda() + elapsed = time.perf_counter() - start + + after = memory_snapshot() + deltas = { + "cpu_mb": after["cpu_mb"] - before["cpu_mb"], + "gpu_mb": after["gpu_mb"] - before["gpu_mb"], + } + return elapsed, deltas, peak_gpu_memory_mb(), result + + +def reset_robot(robot, initial_qpos) -> None: + """Reset current and target robot qpos to the benchmark initial posture.""" + for target in (False, True): + robot.set_qpos(initial_qpos, target=target) + robot.clear_dynamics() + + +def reset_rigid_object(obj, initial_pose) -> None: + """Reset a rigid object pose and clear residual dynamics.""" + obj.set_local_pose(initial_pose) + obj.clear_dynamics() + + +def reset_rigid_object_xy( + obj, + base_pose, + xy: tuple[float, float], + sim=None, + settle_steps: int = 0, +): + """Reset a rigid object to a new XY position while preserving base orientation.""" + pose = base_pose.clone() + pose[:, 0, 3] = xy[0] + pose[:, 1, 3] = xy[1] + reset_rigid_object(obj, pose) + if sim is not None and settle_steps > 0: + sim.update(step=settle_steps) + return pose + + +def park_rigid_object( + obj, + base_pose, + index: int = 0, + sim=None, +) -> None: + """Move an inactive benchmark object outside the robot workspace.""" + pose = base_pose.clone() + pose[:, 0, 3] = 8.0 + float(index) + pose[:, 1, 3] = 8.0 + pose[:, 2, 3] = 1.0 + reset_rigid_object(obj, pose) + if sim is not None: + sim.update(step=1) + + +def object_position_tuple(obj) -> tuple[float, float, float]: + """Return the current object-frame origin position as a CPU tuple.""" + pose = obj.get_local_pose(to_matrix=True) + xyz = pose[0, :3, 3] + if hasattr(xyz, "detach"): + xyz = xyz.detach().to("cpu").tolist() + return (float(xyz[0]), float(xyz[1]), float(xyz[2])) + + +def xy_distance_m( + position: Sequence[float], + target: Sequence[float], +) -> float: + """Return XY Euclidean distance in meters.""" + return math.sqrt( + (float(position[0]) - float(target[0])) ** 2 + + (float(position[1]) - float(target[1])) ** 2 + ) + + +def xyz_distance_m( + position: Sequence[float], + target: Sequence[float], +) -> float: + """Return XYZ Euclidean distance in meters.""" + return math.sqrt( + (float(position[0]) - float(target[0])) ** 2 + + (float(position[1]) - float(target[1])) ** 2 + + (float(position[2]) - float(target[2])) ** 2 + ) + + +def replay_trajectory_for_physical_validation( + sim, + robot, + obj, + traj, + on_step: Callable[[int], None] | None = None, + hold_steps: int = PHYSICAL_VALIDATION_HOLD_STEPS, +) -> tuple[float, float, float] | None: + """Replay a trajectory in physics and return the final object position.""" + if traj is None or getattr(traj, "ndim", 0) < 3 or traj.shape[1] == 0: + return None + + for waypoint_index in range(traj.shape[1]): + robot.set_qpos(traj[:, waypoint_index, :]) + sim.update(step=4) + if on_step is not None: + on_step(waypoint_index) + + final_qpos = traj[:, -1, :] + for _ in range(hold_steps): + robot.set_qpos(final_qpos) + sim.update(step=2) + return object_position_tuple(obj) + + +def should_record_case( + args: argparse.Namespace, + recorded_count: int, + success: bool, +) -> bool: + """Return whether a benchmark case should emit a replay video.""" + if not getattr(args, "record_video", False): + return False + if not success and not getattr(args, "record_failed_video", False): + return False + + case_limit = getattr(args, "video_case_limit", DEFAULT_VIDEO_CASE_LIMIT) + if case_limit < 0: + raise ValueError("--video_case_limit must be non-negative.") + return case_limit == 0 or recorded_count < case_limit + + +def build_video_output_path( + args: argparse.Namespace, + benchmark_name: str, + case_id: str, +) -> Path: + """Build a deterministic, timestamped output path for one replay video.""" + output_dir = Path(getattr(args, "video_dir", DEFAULT_VIDEO_DIR)) + output_dir.mkdir(parents=True, exist_ok=True) + safe_case_id = re.sub(r"[^A-Za-z0-9_.-]+", "_", case_id).strip("_") + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + return output_dir / f"{benchmark_name}_{safe_case_id}_{timestamp}.mp4" + + +def summarize_video_recording( + args: argparse.Namespace, + results: Sequence[dict[str, object]], + video_paths: Sequence[str], +) -> list[str]: + """Build report notes that make video coverage explicit.""" + evaluated_count = len(results) + success_count = sum(1 for result in results if bool(result.get("success"))) + failure_count = evaluated_count - success_count + notes = [ + ( + "Case/video summary: " + f"evaluated={evaluated_count}, success={success_count}, " + f"failure={failure_count}, videos={len(video_paths)}" + ) + ] + if getattr(args, "record_video", False): + if getattr(args, "record_failed_video", False): + notes.append( + "Video policy: records report-success replays and failed-case " + "debug videos when trajectory/static scene capture is available." + ) + else: + notes.append( + "Video policy: records report-success replays only; failed " + "cases are reported in the tables but do not emit videos." + ) + else: + notes.append("Video policy: disabled.") + notes.append( + "Replay videos: " + (", ".join(video_paths) if video_paths else "disabled") + ) + return notes + + +def _replay_trajectory_with_recording( + sim, + robot, + traj, + args: argparse.Namespace, + video_path: Path, + on_step: Callable[[int], None] | None = None, + look_at: tuple[ + Sequence[float], + Sequence[float], + Sequence[float], + ] = DEFAULT_VIDEO_LOOK_AT, +) -> Path | None: + """Replay a planned trajectory and record it with the simulation recorder.""" + if traj is None or getattr(traj, "ndim", 0) < 3 or traj.shape[1] == 0: + return None + + video_fps = getattr(args, "video_fps", DEFAULT_VIDEO_FPS) + video_max_memory = getattr(args, "video_max_memory", DEFAULT_VIDEO_MAX_MEMORY_MB) + video_width = getattr(args, "video_width", DEFAULT_VIDEO_WIDTH) + video_height = getattr(args, "video_height", DEFAULT_VIDEO_HEIGHT) + video_hold_steps = getattr(args, "video_hold_steps", DEFAULT_VIDEO_HOLD_STEPS) + + original_width = sim.sim_config.width + original_height = sim.sim_config.height + recording_started = False + try: + sim.sim_config.width = video_width + sim.sim_config.height = video_height + recording_started = sim.start_window_record( + save_path=str(video_path), + fps=video_fps, + max_memory=video_max_memory, + look_at=look_at, + use_sim_time=True, + ) + finally: + sim.sim_config.width = original_width + sim.sim_config.height = original_height + + if not recording_started: + return None + + stop_success = False + try: + for waypoint_index in range(traj.shape[1]): + robot.set_qpos(traj[:, waypoint_index, :]) + sim.update(step=4) + if on_step is not None: + on_step(waypoint_index) + + final_qpos = traj[:, -1, :] + for _ in range(video_hold_steps): + robot.set_qpos(final_qpos) + sim.update(step=2) + finally: + if sim.is_window_recording(): + stop_success = sim.stop_window_record() + sim.wait_window_record_saves() + + return video_path if stop_success else None + + +def _record_static_scene_video( + sim, + args: argparse.Namespace, + video_path: Path, + look_at: tuple[ + Sequence[float], + Sequence[float], + Sequence[float], + ] = DEFAULT_VIDEO_LOOK_AT, +) -> Path | None: + """Record the current scene without replaying a planned trajectory.""" + video_fps = getattr(args, "video_fps", DEFAULT_VIDEO_FPS) + video_max_memory = getattr(args, "video_max_memory", DEFAULT_VIDEO_MAX_MEMORY_MB) + video_width = getattr(args, "video_width", DEFAULT_VIDEO_WIDTH) + video_height = getattr(args, "video_height", DEFAULT_VIDEO_HEIGHT) + video_hold_steps = getattr(args, "video_hold_steps", DEFAULT_VIDEO_HOLD_STEPS) + + original_width = sim.sim_config.width + original_height = sim.sim_config.height + recording_started = False + try: + sim.sim_config.width = video_width + sim.sim_config.height = video_height + recording_started = sim.start_window_record( + save_path=str(video_path), + fps=video_fps, + max_memory=video_max_memory, + look_at=look_at, + use_sim_time=True, + ) + finally: + sim.sim_config.width = original_width + sim.sim_config.height = original_height + + if not recording_started: + return None + + stop_success = False + try: + for _ in range(video_hold_steps): + sim.update(step=2) + finally: + if sim.is_window_recording(): + stop_success = sim.stop_window_record() + sim.wait_window_record_saves() + + return video_path if stop_success else None + + +def replay_trajectory_with_recording( + sim, + robot, + traj, + args: argparse.Namespace, + video_path: Path, + on_step: Callable[[int], None] | None = None, + look_at: tuple[ + Sequence[float], + Sequence[float], + Sequence[float], + ] = DEFAULT_VIDEO_LOOK_AT, +) -> Path | None: + """Best-effort replay recording that never changes benchmark success.""" + try: + return _replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=traj, + args=args, + video_path=video_path, + on_step=on_step, + look_at=look_at, + ) + except Exception as exc: + try: + if sim.is_window_recording(): + sim.stop_window_record() + sim.wait_window_record_saves() + except Exception: + pass + print( + "Warning: failed to record benchmark replay video " + f"{video_path}: {type(exc).__name__}: {exc}" + ) + return None + + +def record_static_scene_video( + sim, + args: argparse.Namespace, + video_path: Path, + look_at: tuple[ + Sequence[float], + Sequence[float], + Sequence[float], + ] = DEFAULT_VIDEO_LOOK_AT, +) -> Path | None: + """Best-effort static scene recording that never changes benchmark success.""" + try: + return _record_static_scene_video( + sim=sim, + args=args, + video_path=video_path, + look_at=look_at, + ) + except Exception as exc: + try: + if sim.is_window_recording(): + sim.stop_window_record() + sim.wait_window_record_saves() + except Exception: + pass + print( + "Warning: failed to record benchmark static debug video " + f"{video_path}: {type(exc).__name__}: {exc}" + ) + return None + + +def format_float(value: float | None, precision: int = 6) -> str: + """Format finite floats for tables and use N/A for missing values.""" + if value is None or not math.isfinite(value): + return "N/A" + return f"{value:.{precision}f}" + + +def format_markdown_table(rows: list[dict[str, object]]) -> list[str]: + """Format rows into a markdown table.""" + if not rows: + return ["No data."] + + headers = list(rows[0].keys()) + lines = [ + "| " + " | ".join(headers) + " |", + "| " + " | ".join(["---"] * len(headers)) + " |", + ] + for row in rows: + lines.append("| " + " | ".join(str(row[h]) for h in headers) + " |") + return lines + + +def write_markdown_report( + benchmark_name: str, + perf_rows: list[dict[str, object]], + metric_rows: list[dict[str, object]], + leaderboard_rows: list[dict[str, object]], + notes: list[str] | None = None, +) -> Path: + """Write benchmark results into one markdown report file.""" + output_dir = Path("outputs/benchmarks") + output_dir.mkdir(parents=True, exist_ok=True) + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + report_path = output_dir / f"{benchmark_name}_{timestamp}.md" + + lines: list[str] = [ + f"# {benchmark_name} Benchmark Report", + "", + f"Generated at: {datetime.now().isoformat(timespec='seconds')}", + "", + "## Time & Memory", + "", + ] + lines.extend(format_markdown_table(perf_rows)) + lines.extend(["", "## Success & Other Metrics", ""]) + lines.extend(format_markdown_table(metric_rows)) + lines.extend(["", "## Leaderboard", ""]) + lines.extend(format_markdown_table(leaderboard_rows)) + + if notes: + lines.extend(["", "## Notes", ""]) + lines.extend([f"- {note}" for note in notes]) + + report_path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return report_path + + +def build_single_action_leaderboard( + action_name: str, + metric_rows: list[dict[str, object]], +) -> list[dict[str, object]]: + """Aggregate rows for one action by benchmark case.""" + if not metric_rows: + return [] + + success_sum = sum(float(row["success_rate"]) for row in metric_rows) + count = len(metric_rows) + return [ + { + "rank": 1, + "algorithm": action_name, + "overall_success_rate": f"{success_sum / max(count, 1):.2%}", + "evaluated_cases": count, + } + ] + + +__all__ = [ + "BENCHMARK_PROFILES", + "CPU_MEMORY_BACKEND", + "COVERAGE_MESH_OBJECT_TYPES", + "COVERAGE_POSITION_CASE_NAMES", + "DEFAULT_BENCHMARK_PROFILE", + "add_common_benchmark_args", + "add_grasp_benchmark_args", + "add_object_position_benchmark_args", + "add_profile_benchmark_args", + "add_video_benchmark_args", + "build_video_output_path", + "build_single_action_leaderboard", + "create_antipodal_object_semantics", + "create_benchmark_object", + "create_mesh_benchmark_object", + "DEFAULT_VIDEO_DIR", + "default_pickup_approach_cases_for_profile", + "default_mesh_object_types_for_profile", + "default_position_case_names_for_profile", + "describe_object_preset", + "ensure_repo_root", + "ensure_torch", + "format_float", + "format_vector3", + "FULL_MESH_OBJECT_TYPES", + "FULL_POSITION_CASE_NAMES", + "MESH_OBJECT_PRESETS", + "MeshObjectPreset", + "PICKUP_APPROACH_CASES", + "PHYSICAL_MOVE_HELD_OBJECT_XYZ_TOLERANCE_M", + "PHYSICAL_PICK_MIN_LIFT_M", + "PHYSICAL_PLACE_XY_TOLERANCE_M", + "PHYSICAL_VALIDATION_HOLD_STEPS", + "POSITION_CASES", + "PositionCase", + "object_position_tuple", + "park_rigid_object", + "pickup_approach_direction_tuple", + "record_static_scene_video", + "replay_trajectory_for_physical_validation", + "replay_trajectory_with_recording", + "reset_rigid_object", + "reset_rigid_object_xy", + "reset_robot", + "resolve_pickup_approach_direction", + "resolve_profile", + "select_mesh_object_presets", + "select_pickup_approaches", + "select_position_cases", + "should_record_case", + "SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z", + "SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT", + "SMOKE_PICKUP_APPROACH_CASES", + "SMOKE_MESH_OBJECT_TYPES", + "SMOKE_POSITION_CASE_NAMES", + "timed_call", + "summarize_video_recording", + "write_markdown_report", + "xy_distance_m", + "xyz_distance_m", +] diff --git a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py new file mode 100644 index 000000000..6267e5418 --- /dev/null +++ b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py @@ -0,0 +1,318 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Benchmark MoveEndEffector atomic action across target poses. + +Measures planning latency, memory usage, trajectory success, and final TCP +translation error for several reachable pose targets. +Run: python -m scripts.benchmark.atomic_action.move_end_effector_benchmark +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + CPU_MEMORY_BACKEND, + add_common_benchmark_args, + build_single_action_leaderboard, + build_video_output_path, + ensure_repo_root, + ensure_torch, + format_float, + replay_trajectory_with_recording, + reset_robot, + resolve_profile, + should_record_case, + timed_call, + write_markdown_report, +) + + +@dataclass(frozen=True) +class PoseCase: + """End-effector target pose benchmark case.""" + + name: str + xyz: tuple[float, float, float] + + +POSE_CASES = { + "front_left": PoseCase("front_left", (0.30, -0.20, 0.36)), + "front_right": PoseCase("front_right", (0.30, 0.20, 0.36)), + "near_center": PoseCase("near_center", (0.18, 0.00, 0.42)), + "far_center": PoseCase("far_center", (0.42, 0.00, 0.34)), +} +DEFAULT_POSE_CASES = tuple(POSE_CASES.keys()) +MOVE_SAMPLE_INTERVAL = 80 +SUCCESS_TOLERANCE_M = 0.01 + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add MoveEndEffector benchmark CLI arguments.""" + parser.add_argument( + "--pose_cases", + nargs="+", + choices=(*POSE_CASES.keys(), "all"), + default=list(DEFAULT_POSE_CASES), + help="Target pose cases to benchmark. Use 'all' for every case.", + ) + add_common_benchmark_args(parser) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Benchmark MoveEndEffector over reachable target poses." + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _select_pose_cases(case_names: list[str]) -> list[PoseCase]: + """Resolve selected pose case names.""" + if "all" in case_names: + return list(POSE_CASES.values()) + return [POSE_CASES[name] for name in case_names] + + +def _make_pose(device, xyz: tuple[float, float, float]): + """Create a top-down target pose.""" + torch = ensure_torch() + pose = torch.eye(4, dtype=torch.float32, device=device) + pose[:3, :3] = torch.tensor( + [ + [-0.0539, -0.9985, -0.0022], + [-0.9977, 0.0540, -0.0401], + [0.0401, 0.0000, -0.9992], + ], + dtype=torch.float32, + device=device, + ) + pose[:3, 3] = torch.tensor(xyz, dtype=torch.float32, device=device) + return pose + + +def _run_case( + sim, + robot, + atomic_engine, + initial_qpos, + pose_case: PoseCase, + repeat: int, + args: argparse.Namespace, + recorded_count: int, +): + """Run one MoveEndEffector case.""" + torch = ensure_torch() + from embodichain.lab.sim.atomic_actions import EndEffectorPoseTarget + + reset_robot(robot, initial_qpos) + target_pose = _make_pose(sim.device, pose_case.xyz) + + elapsed, mem_delta, peak_gpu, result = timed_call( + lambda: atomic_engine.run( + steps=[("move_end_effector", EndEffectorPoseTarget(xpos=target_pose))] + ) + ) + is_success, traj, _ = result + video_path = None + if should_record_case(args, recorded_count, bool(is_success)): + reset_robot(robot, initial_qpos) + video_path = replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=traj, + args=args, + video_path=build_video_output_path( + args, "atomic_action_move_end_effector", f"{pose_case.name}_r{repeat}" + ), + ) + reset_robot(robot, initial_qpos) + + final_error_m = None + if is_success and traj.shape[1] > 0: + arm_joint_ids = robot.get_joint_ids(name="arm") + final_arm_qpos = traj[:, -1, arm_joint_ids] + final_pose = robot.compute_fk(qpos=final_arm_qpos, name="arm", to_matrix=True)[ + 0 + ] + final_error_m = float(torch.linalg.norm(final_pose[:3, 3] - target_pose[:3, 3])) + target_reached = bool( + is_success + and final_error_m is not None + and final_error_m <= SUCCESS_TOLERANCE_M + ) + return { + "case_id": f"{pose_case.name}:r{repeat}", + "target_case": pose_case.name, + "repeat": repeat, + "planning_success": bool(is_success), + "target_reached": target_reached, + "cost_time_ms": elapsed * 1000.0, + "cpu_delta_mb": mem_delta["cpu_mb"], + "gpu_delta_mb": mem_delta["gpu_mb"], + "peak_gpu_mb": peak_gpu, + "final_error_m": final_error_m, + "trajectory_waypoints": int(traj.shape[1]) if traj.ndim >= 2 else 0, + "failure_reason": "" if target_reached else "target_not_reached", + "video_path": str(video_path) if video_path is not None else "", + } + + +def _build_rows(results: list[dict[str, object]]): + """Build report rows for MoveEndEffector.""" + perf_rows = [] + metric_rows = [] + for result in results: + perf_rows.append( + { + "sample_size": 1, + "impl": "move_end_effector", + "case_id": result["case_id"], + "target_case": result["target_case"], + "repeat": result["repeat"], + "cost_time_ms": format_float(result["cost_time_ms"]), + "cpu_delta_mb": format_float(result["cpu_delta_mb"]), + "gpu_delta_mb": format_float(result["gpu_delta_mb"]), + "peak_gpu_mb": format_float(result["peak_gpu_mb"]), + } + ) + metric_rows.append( + { + "sample_size": 1, + "impl": "move_end_effector", + "case_id": result["case_id"], + "target_case": result["target_case"], + "success_rate": f"{float(result['target_reached']):.6f}", + "planning_success_rate": f"{float(result['planning_success']):.6f}", + "translation_err_m": format_float(result["final_error_m"]), + "trajectory_waypoints": result["trajectory_waypoints"], + "failure_reason": result["failure_reason"] or "N/A", + } + ) + return perf_rows, metric_rows + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: + """Run MoveEndEffector benchmark and write a markdown report.""" + args = _parse_args() if args is None else args + if args.repeat < 1: + raise ValueError("--repeat must be at least 1.") + profile = resolve_profile(args) + + ensure_repo_root() + ensure_torch() + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + MoveEndEffector, + MoveEndEffectorCfg, + ) + from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.planners import ToppraPlannerCfg + from scripts.tutorials.atomic_action.move_end_effector import ( + create_robot, + initialize_simulation, + ) + + pose_cases = _select_pose_cases(args.pose_cases) + repeat = 1 if profile == "smoke" else args.repeat + if profile == "smoke": + pose_cases = [POSE_CASES["front_left"]] + + print("=" * 60) + print("MoveEndEffector Atomic Action Benchmark") + print("=" * 60) + print( + "Coverage: " + f"profile={profile}, {len(pose_cases)} pose case(s) x " + f"{repeat} repeat(s)" + ) + + sim = initialize_simulation(args) + robot = create_robot(sim) + initial_qpos = robot.get_qpos().clone() + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + MoveEndEffector( + motion_gen, + cfg=MoveEndEffectorCfg( + control_part="arm", sample_interval=MOVE_SAMPLE_INTERVAL + ), + ) + ) + + results: list[dict[str, object]] = [] + video_paths: list[str] = [] + print("\n=== MoveEndEffector Target Sweep ===") + for pose_case in pose_cases: + for repeat_index in range(repeat): + result = _run_case( + sim, + robot, + atomic_engine, + initial_qpos, + pose_case, + repeat_index, + args, + len(video_paths), + ) + results.append(result) + if result["video_path"]: + video_paths.append(str(result["video_path"])) + print( + f" {result['case_id']:<24} " + f"time={result['cost_time_ms']:>10.2f} ms | " + f"success={result['target_reached']} " + f"err={format_float(result['final_error_m'], precision=4)}" + ) + + perf_rows, metric_rows = _build_rows(results) + leaderboard_rows = build_single_action_leaderboard("move_end_effector", metric_rows) + report_path = write_markdown_report( + benchmark_name="atomic_action_move_end_effector", + perf_rows=perf_rows, + metric_rows=metric_rows, + leaderboard_rows=leaderboard_rows, + notes=[ + f"Profile: {profile}", + f"CPU memory backend: {CPU_MEMORY_BACKEND}", + f"Success tolerance: {SUCCESS_TOLERANCE_M} m translation error.", + "Replay videos: " + (", ".join(video_paths) if video_paths else "disabled"), + ], + ) + print(f"Markdown report saved: {report_path}") + return report_path + + +def main() -> None: + """Run the CLI entry point.""" + try: + run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/move_held_object_benchmark.py b/scripts/benchmark/atomic_action/move_held_object_benchmark.py new file mode 100644 index 000000000..cabfb1582 --- /dev/null +++ b/scripts/benchmark/atomic_action/move_held_object_benchmark.py @@ -0,0 +1,766 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Benchmark MoveHeldObject atomic action after a PickUp precondition. + +Measures MoveHeldObject-only planning latency and memory usage once a held +object state has been produced by PickUp. +Run: python -m scripts.benchmark.atomic_action.move_held_object_benchmark +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + CPU_MEMORY_BACKEND, + add_common_benchmark_args, + add_grasp_benchmark_args, + add_object_position_benchmark_args, + add_pickup_approach_benchmark_args, + build_single_action_leaderboard, + build_video_output_path, + create_antipodal_object_semantics, + create_benchmark_object, + describe_object_preset, + ensure_repo_root, + ensure_torch, + format_float, + format_vector3, + MeshObjectPreset, + object_position_tuple, + park_rigid_object, + PHYSICAL_MOVE_HELD_OBJECT_XYZ_TOLERANCE_M, + PHYSICAL_PICK_MIN_LIFT_M, + pickup_approach_direction_tuple, + PositionCase, + record_static_scene_video, + replay_trajectory_for_physical_validation, + replay_trajectory_with_recording, + reset_rigid_object, + reset_rigid_object_xy, + reset_robot, + resolve_pickup_approach_direction, + resolve_profile, + select_mesh_object_presets, + select_pickup_approaches, + select_position_cases, + should_record_case, + SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z, + SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT, + summarize_video_recording, + timed_call, + write_markdown_report, + xy_distance_m, + xyz_distance_m, +) + + +@dataclass(frozen=True) +class HeldObjectCase: + """Held-object target pose benchmark case.""" + + name: str + xyz: tuple[float, float, float] + + +HELD_OBJECT_CASES = { + "front_left": HeldObjectCase("front_left", (-0.30, -0.30, 0.50)), + "front_right": HeldObjectCase("front_right", (-0.30, 0.30, 0.50)), + "center_high": HeldObjectCase("center_high", (-0.42, 0.00, 0.58)), +} +PICK_SAMPLE_INTERVAL = 120 +MOVE_SAMPLE_INTERVAL = 60 +MOVE_HELD_OBJECT_SAMPLE_INTERVAL = 120 +HAND_INTERP_STEPS = 12 + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add MoveHeldObject benchmark CLI arguments.""" + add_pickup_approach_benchmark_args(parser) + parser.add_argument( + "--held_object_cases", + nargs="+", + choices=(*HELD_OBJECT_CASES.keys(), "all"), + default=None, + help=( + "Held-object target cases to benchmark. Defaults are selected by " + "--profile; use 'all' for every case." + ), + ) + add_object_position_benchmark_args(parser) + add_grasp_benchmark_args(parser) + add_common_benchmark_args(parser) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Benchmark MoveHeldObject after a PickUp precondition." + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _selected_cases( + case_names: list[str] | None, + profile: str, +) -> list[HeldObjectCase]: + """Resolve selected held-object case names.""" + if not case_names: + return [HELD_OBJECT_CASES["front_left"]] + if "all" in case_names: + return list(HELD_OBJECT_CASES.values()) + return [HELD_OBJECT_CASES[name] for name in case_names] + + +def _make_object_target_pose(device, xyz: tuple[float, float, float]): + """Create a held-object target pose.""" + torch = ensure_torch() + pose = torch.eye(4, dtype=torch.float32, device=device) + pose[:3, :3] = torch.tensor( + [ + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0], + [1.0, 0.0, 0.0], + ], + dtype=torch.float32, + device=device, + ) + pose[:3, 3] = torch.tensor(xyz, dtype=torch.float32, device=device) + return pose + + +def _make_pickup_args( + args: argparse.Namespace, + object_preset: MeshObjectPreset, + profile: str, +) -> argparse.Namespace: + """Build a tutorial-compatible argparse namespace.""" + return argparse.Namespace( + object=object_preset.label, + n_sample=1000 if profile == "smoke" else args.n_sample, + force_reannotate=args.force_reannotate, + device=args.device, + renderer=args.renderer, + headless=True, + ) + + +def _prepare_held_state( + sim, + robot, + obj, + motion_gen, + args, + object_preset: MeshObjectPreset, + position_case: PositionCase, + pickup_approach: str, + profile: str, +): + """Run PickUp precondition outside the timed MoveHeldObject block.""" + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + EndEffectorPoseTarget, + GraspTarget, + MoveEndEffector, + MoveEndEffectorCfg, + PickUp, + PickUpCfg, + ) + from scripts.tutorials.atomic_action.move_held_object import ( + build_grasp_generator_cfg, + build_gripper_collision_cfg, + get_hand_open_close_qpos, + make_pre_pick_eef_pose, + ) + + hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + MoveEndEffector( + motion_gen, + cfg=MoveEndEffectorCfg( + control_part="arm", + sample_interval=MOVE_SAMPLE_INTERVAL, + ), + ) + ) + atomic_engine.register( + PickUp( + motion_gen, + cfg=PickUpCfg( + control_part="arm", + hand_control_part="hand", + hand_open_qpos=hand_open, + hand_close_qpos=hand_close, + approach_direction=resolve_pickup_approach_direction( + pickup_approach, position_case, sim.device + ), + pre_grasp_distance=0.15, + lift_height=0.16, + sample_interval=PICK_SAMPLE_INTERVAL, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ) + ) + semantics = create_antipodal_object_semantics( + obj=obj, + preset=object_preset, + args=_make_pickup_args(args, object_preset, profile), + build_gripper_collision_cfg=build_gripper_collision_cfg, + build_grasp_generator_cfg=build_grasp_generator_cfg, + ) + obj_pose = obj.get_local_pose(to_matrix=True) + move_position = obj_pose[0, :3, 3].clone() + move_position[2] = 0.36 + move_target = make_pre_pick_eef_pose(robot, move_position) + is_success, traj, state = atomic_engine.run( + steps=[ + ("move_end_effector", EndEffectorPoseTarget(xpos=move_target)), + ("pick_up", GraspTarget(semantics=semantics)), + ] + ) + if not is_success or state.held_object is None: + raise RuntimeError( + "Failed to prepare held-object state for MoveHeldObject benchmark." + ) + robot.set_qpos(state.last_qpos) + return state, hand_close, traj + + +def _run_case( + sim, + robot, + motion_gen, + initial_qpos, + args, + obj, + base_obj_pose, + object_preset: MeshObjectPreset, + position_case: PositionCase, + pickup_approach: str, + case: HeldObjectCase, + repeat: int, + recorded_count: int, + profile: str, +): + """Run one MoveHeldObject benchmark case.""" + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + HeldObjectPoseTarget, + MoveHeldObject, + MoveHeldObjectCfg, + ) + from scripts.tutorials.atomic_action.move_held_object import ( + compute_pick_close_end_step, + ) + + case_id = ( + f"{object_preset.object_type}:{position_case.name}:" + f"{pickup_approach}:{case.name}:r{repeat}" + ) + try: + reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + reset_rigid_object(obj, initial_obj_pose) + initial_obj_position = object_position_tuple(obj) + state, hand_close, precondition_traj = _prepare_held_state( + sim, + robot, + obj, + motion_gen, + args, + object_preset, + position_case, + pickup_approach, + profile, + ) + approach_direction_text = format_vector3( + pickup_approach_direction_tuple(pickup_approach, position_case) + ) + precondition_waypoints = int(precondition_traj.shape[1]) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + MoveHeldObject( + motion_gen, + cfg=MoveHeldObjectCfg( + control_part="arm", + hand_control_part="hand", + hand_close_qpos=hand_close, + sample_interval=MOVE_HELD_OBJECT_SAMPLE_INTERVAL, + ), + ) + ) + target_pose = _make_object_target_pose(sim.device, case.xyz) + elapsed, mem_delta, peak_gpu, result = timed_call( + lambda: atomic_engine.run( + steps=[ + ( + "move_held_object", + HeldObjectPoseTarget(object_target_pose=target_pose), + ) + ], + state=state, + ) + ) + is_success, traj, final_state = result + torch = ensure_torch() + precondition_obj_position = None + final_obj_position = None + object_lift_delta_m = None + held_object_xy_error_m = None + held_object_z_error_m = None + held_object_xyz_error_m = None + + if ( + getattr(precondition_traj, "ndim", 0) >= 3 + and precondition_traj.shape[1] > 0 + ): + if bool(is_success) and getattr(traj, "ndim", 0) >= 3 and traj.shape[1] > 0: + validation_traj = torch.cat((precondition_traj, traj), dim=1) + else: + validation_traj = precondition_traj + + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + post_grasp_clear_step = compute_pick_close_end_step() + should_clear_object_dynamics = True + + def _on_validation_step(waypoint_index: int) -> None: + nonlocal should_clear_object_dynamics, precondition_obj_position + if ( + should_clear_object_dynamics + and waypoint_index + 1 >= post_grasp_clear_step + ): + obj.clear_dynamics() + should_clear_object_dynamics = False + if ( + precondition_obj_position is None + and waypoint_index + 1 >= precondition_waypoints + ): + precondition_obj_position = object_position_tuple(obj) + + final_obj_position = replay_trajectory_for_physical_validation( + sim=sim, + robot=robot, + obj=obj, + traj=validation_traj, + on_step=_on_validation_step, + ) + if precondition_obj_position is not None: + object_lift_delta_m = ( + precondition_obj_position[2] - initial_obj_position[2] + ) + if bool(is_success) and final_obj_position is not None: + held_object_xy_error_m = xy_distance_m(final_obj_position, case.xyz) + held_object_z_error_m = abs(final_obj_position[2] - case.xyz[2]) + held_object_xyz_error_m = xyz_distance_m(final_obj_position, case.xyz) + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + still_held = bool(is_success and final_state.held_object is not None) + physical_pick_success = bool( + object_lift_delta_m is not None + and object_lift_delta_m >= PHYSICAL_PICK_MIN_LIFT_M + ) + physical_move_success = bool( + still_held + and physical_pick_success + and held_object_xyz_error_m is not None + and held_object_xyz_error_m <= PHYSICAL_MOVE_HELD_OBJECT_XYZ_TOLERANCE_M + ) + success = physical_move_success + if success: + failure_reason = "" + elif not is_success: + failure_reason = "planning_failed" + elif not still_held: + failure_reason = "held_object_lost" + elif not physical_pick_success: + failure_reason = "physical_pick_failed" + elif held_object_xyz_error_m is None: + failure_reason = "physical_replay_missing" + else: + failure_reason = "held_object_target_miss" + + video_path = None + if should_record_case(args, recorded_count, success): + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + video_case_suffix = ( + f"{object_preset.object_type}_{position_case.name}_" + f"{pickup_approach}_{case.name}_r{repeat}" + ) + if getattr(traj, "ndim", 0) >= 3 and traj.shape[1] > 0: + replay_traj = torch.cat((precondition_traj, traj), dim=1) + else: + replay_traj = precondition_traj + + if getattr(replay_traj, "ndim", 0) >= 3 and replay_traj.shape[1] > 0: + post_grasp_clear_step = compute_pick_close_end_step() + should_clear_object_dynamics = True + + def _on_step(waypoint_index: int) -> None: + nonlocal should_clear_object_dynamics + if ( + should_clear_object_dynamics + and waypoint_index + 1 >= post_grasp_clear_step + ): + obj.clear_dynamics() + should_clear_object_dynamics = False + + video_path = replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=replay_traj, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_move_held_object", + video_case_suffix, + ), + on_step=_on_step, + ) + else: + video_path = record_static_scene_video( + sim=sim, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_move_held_object", + f"{video_case_suffix}_failed_static", + ), + ) + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + return { + "case_id": case_id, + "object_type": object_preset.object_type, + "material": object_preset.material_name, + "quadrant": position_case.quadrant, + "position_case": position_case.name, + "init_xy": position_case.xy, + "pickup_approach": pickup_approach, + "approach_direction": approach_direction_text, + "held_object_case": case.name, + "repeat": repeat, + "planning_success": bool(is_success), + "still_held": still_held, + "physical_pick_success": physical_pick_success, + "physical_move_success": physical_move_success, + "success": success, + "cost_time_ms": elapsed * 1000.0, + "cpu_delta_mb": mem_delta["cpu_mb"], + "gpu_delta_mb": mem_delta["gpu_mb"], + "peak_gpu_mb": peak_gpu, + "object_lift_delta_m": object_lift_delta_m, + "object_final_x_m": ( + final_obj_position[0] if final_obj_position is not None else None + ), + "object_final_y_m": ( + final_obj_position[1] if final_obj_position is not None else None + ), + "object_final_z_m": ( + final_obj_position[2] if final_obj_position is not None else None + ), + "held_object_xy_error_m": held_object_xy_error_m, + "held_object_z_error_m": held_object_z_error_m, + "held_object_xyz_error_m": held_object_xyz_error_m, + "precondition_waypoints": precondition_waypoints, + "trajectory_waypoints": int(traj.shape[1]) if traj.ndim >= 2 else 0, + "failure_reason": failure_reason, + "video_path": str(video_path) if video_path is not None else "", + } + except Exception as exc: + video_path = None + if should_record_case(args, recorded_count, False): + try: + reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + video_path = record_static_scene_video( + sim=sim, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_move_held_object", + ( + f"{object_preset.object_type}_{position_case.name}_" + f"{pickup_approach}_{case.name}_r{repeat}" + "_exception_static" + ), + ), + ) + reset_rigid_object(obj, initial_obj_pose) + except Exception: + video_path = None + return { + "case_id": case_id, + "object_type": object_preset.object_type, + "material": object_preset.material_name, + "quadrant": position_case.quadrant, + "position_case": position_case.name, + "init_xy": position_case.xy, + "pickup_approach": pickup_approach, + "approach_direction": "N/A", + "held_object_case": case.name, + "repeat": repeat, + "planning_success": False, + "still_held": False, + "physical_pick_success": False, + "physical_move_success": False, + "success": False, + "cost_time_ms": 0.0, + "cpu_delta_mb": 0.0, + "gpu_delta_mb": 0.0, + "peak_gpu_mb": 0.0, + "object_lift_delta_m": None, + "object_final_x_m": None, + "object_final_y_m": None, + "object_final_z_m": None, + "held_object_xy_error_m": None, + "held_object_z_error_m": None, + "held_object_xyz_error_m": None, + "precondition_waypoints": 0, + "trajectory_waypoints": 0, + "failure_reason": f"exception:{type(exc).__name__}:{exc}", + "video_path": str(video_path) if video_path is not None else "", + } + + +def _build_rows(results: list[dict[str, object]]): + """Build report rows for MoveHeldObject.""" + perf_rows = [] + metric_rows = [] + for result in results: + perf_rows.append( + { + "sample_size": 1, + "impl": "move_held_object", + "case_id": result["case_id"], + "object_type": result["object_type"], + "material": result["material"], + "quadrant": result["quadrant"], + "position_case": result["position_case"], + "init_xy": f"({result['init_xy'][0]:.3f},{result['init_xy'][1]:.3f})", + "pickup_approach": result["pickup_approach"], + "approach_direction": result["approach_direction"], + "held_object_case": result["held_object_case"], + "repeat": result["repeat"], + "cost_time_ms": format_float(result["cost_time_ms"]), + "cpu_delta_mb": format_float(result["cpu_delta_mb"]), + "gpu_delta_mb": format_float(result["gpu_delta_mb"]), + "peak_gpu_mb": format_float(result["peak_gpu_mb"]), + } + ) + metric_rows.append( + { + "sample_size": 1, + "impl": "move_held_object", + "case_id": result["case_id"], + "object_type": result["object_type"], + "material": result["material"], + "quadrant": result["quadrant"], + "position_case": result["position_case"], + "pickup_approach": result["pickup_approach"], + "approach_direction": result["approach_direction"], + "held_object_case": result["held_object_case"], + "success_rate": f"{float(result['success']):.6f}", + "planning_success_rate": f"{float(result['planning_success']):.6f}", + "held_object_rate": f"{float(result['still_held']):.6f}", + "physical_pick_success_rate": ( + f"{float(result['physical_pick_success']):.6f}" + ), + "physical_move_success_rate": ( + f"{float(result['physical_move_success']):.6f}" + ), + "object_lift_delta_m": format_float(result["object_lift_delta_m"]), + "object_final_x_m": format_float(result["object_final_x_m"]), + "object_final_y_m": format_float(result["object_final_y_m"]), + "object_final_z_m": format_float(result["object_final_z_m"]), + "held_object_xy_error_m": format_float( + result["held_object_xy_error_m"] + ), + "held_object_z_error_m": format_float(result["held_object_z_error_m"]), + "held_object_xyz_error_m": format_float( + result["held_object_xyz_error_m"] + ), + "precondition_waypoints": result["precondition_waypoints"], + "trajectory_waypoints": result["trajectory_waypoints"], + "failure_reason": result["failure_reason"] or "N/A", + } + ) + return perf_rows, metric_rows + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: + """Run MoveHeldObject benchmark and write a markdown report.""" + args = _parse_args() if args is None else args + if args.repeat < 1: + raise ValueError("--repeat must be at least 1.") + profile = resolve_profile(args) + + ensure_repo_root() + ensure_torch() + from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.planners import ToppraPlannerCfg + from scripts.tutorials.atomic_action.move_held_object import ( + create_robot, + initialize_simulation, + ) + + object_presets = select_mesh_object_presets(args.object_types, profile) + position_cases = select_position_cases(args.position_cases, profile) + approaches = select_pickup_approaches(args.approach_cases, profile) + cases = _selected_cases(args.held_object_cases, profile) + repeat = 1 if profile == "smoke" else args.repeat + + print("=" * 60) + print("MoveHeldObject Atomic Action Benchmark") + print("=" * 60) + print( + "Coverage: " + f"profile={profile}, {len(object_presets)} object(s) x " + f"{len(position_cases)} position(s) x {len(approaches)} pickup approach(es) " + f"x {len(cases)} held target(s) " + f"x {repeat} repeat(s)" + ) + + sim = initialize_simulation(args) + robot = create_robot(sim) + initial_qpos = robot.get_qpos().clone() + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + object_pool = {} + for object_index, object_preset in enumerate(object_presets): + obj = create_benchmark_object( + sim=sim, + preset=object_preset, + position_case=position_cases[0], + uid_suffix="pool", + ) + base_pose = obj.get_local_pose(to_matrix=True).clone() + park_rigid_object(obj, base_pose, index=object_index, sim=sim) + object_pool[object_preset.object_type] = (obj, base_pose) + + results: list[dict[str, object]] = [] + video_paths: list[str] = [] + print("\n=== MoveHeldObject Object/Position/Target Sweep ===") + for object_preset in object_presets: + obj, base_pose = object_pool[object_preset.object_type] + for parked_index, parked_preset in enumerate(object_presets): + if parked_preset.object_type == object_preset.object_type: + continue + parked_obj, parked_base_pose = object_pool[parked_preset.object_type] + park_rigid_object(parked_obj, parked_base_pose, index=parked_index, sim=sim) + for position_case in position_cases: + for pickup_approach in approaches: + for case in cases: + for repeat_index in range(repeat): + result = _run_case( + sim, + robot, + motion_gen, + initial_qpos, + args, + obj, + base_pose, + object_preset, + position_case, + pickup_approach, + case, + repeat_index, + len(video_paths), + profile, + ) + results.append(result) + if result["video_path"]: + video_paths.append(str(result["video_path"])) + print( + f" {result['case_id']:<54} " + f"time={result['cost_time_ms']:>10.2f} ms | " + f"success={result['success']}" + ) + + perf_rows, metric_rows = _build_rows(results) + leaderboard_rows = build_single_action_leaderboard("move_held_object", metric_rows) + report_path = write_markdown_report( + benchmark_name="atomic_action_move_held_object", + perf_rows=perf_rows, + metric_rows=metric_rows, + leaderboard_rows=leaderboard_rows, + notes=[ + "Timed block includes MoveHeldObject only; PickUp precondition is " + "prepared outside timing.", + f"Profile: {profile}", + "Object presets: " + + ", ".join(describe_object_preset(preset) for preset in object_presets), + "Position cases: " + + ", ".join( + f"{case.name}/{case.quadrant}/xy={case.xy}" for case in position_cases + ), + "PickUp approach cases: " + ", ".join(approaches), + "Held-object target cases: " + ", ".join(case.name for case in cases), + f"CPU memory backend: {CPU_MEMORY_BACKEND}", + f"n_sample: {1000 if profile == 'smoke' else args.n_sample}", + ( + "Side grasp opening-axis rule: prefer candidates with " + f"abs(opening_axis_z) <= {SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z:.2f}; " + "fallback cost += " + f"abs(opening_axis_z) * {SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT:.2f}." + ), + ( + "Physical success rule: PickUp precondition lift delta >= " + f"{PHYSICAL_PICK_MIN_LIFT_M:.3f} m, held-object state is kept, " + "and final object XYZ error <= " + f"{PHYSICAL_MOVE_HELD_OBJECT_XYZ_TOLERANCE_M:.3f} m." + ), + *summarize_video_recording(args, results, video_paths), + ], + ) + print(f"Markdown report saved: {report_path}") + return report_path + + +def main() -> None: + """Run the CLI entry point.""" + try: + run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/move_joints_benchmark.py b/scripts/benchmark/atomic_action/move_joints_benchmark.py new file mode 100644 index 000000000..25a1f3a05 --- /dev/null +++ b/scripts/benchmark/atomic_action/move_joints_benchmark.py @@ -0,0 +1,331 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Benchmark MoveJoints atomic action across joint-space target sequences. + +Measures planning latency, memory usage, trajectory success, and final arm joint +error for named and explicit joint-position targets. +Run: python -m scripts.benchmark.atomic_action.move_joints_benchmark +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + CPU_MEMORY_BACKEND, + add_common_benchmark_args, + build_single_action_leaderboard, + build_video_output_path, + ensure_repo_root, + ensure_torch, + format_float, + replay_trajectory_with_recording, + reset_robot, + resolve_profile, + should_record_case, + timed_call, + write_markdown_report, +) + + +@dataclass(frozen=True) +class JointSequenceCase: + """MoveJoints target sequence benchmark case.""" + + name: str + sequence: tuple[str, ...] + + +READY_QPOS = (0.35, -1.20, 1.30, -1.65, -1.57, 0.20) +HOME_QPOS = (0.0, -1.57, 1.57, -1.57, -1.57, 0.0) +EXTENDED_QPOS = (0.55, -1.05, 1.10, -1.45, -1.35, 0.35) +JOINT_TARGETS = { + "ready": READY_QPOS, + "home": HOME_QPOS, + "extended": EXTENDED_QPOS, +} +SEQUENCE_CASES = { + "ready_home": JointSequenceCase("ready_home", ("ready", "home")), + "extended_home": JointSequenceCase("extended_home", ("extended", "home")), + "ready_extended_home": JointSequenceCase( + "ready_extended_home", ("ready", "extended", "home") + ), +} +DEFAULT_SEQUENCE_CASES = tuple(SEQUENCE_CASES.keys()) +MOVE_JOINTS_SAMPLE_INTERVAL = 80 +SUCCESS_TOLERANCE_RAD = 1e-4 + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add MoveJoints benchmark CLI arguments.""" + parser.add_argument( + "--sequence_cases", + nargs="+", + choices=(*SEQUENCE_CASES.keys(), "all"), + default=list(DEFAULT_SEQUENCE_CASES), + help="Joint sequence cases to benchmark. Use 'all' for every case.", + ) + add_common_benchmark_args(parser) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Benchmark MoveJoints over named and explicit qpos targets." + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _select_cases(case_names: list[str]) -> list[JointSequenceCase]: + """Resolve selected sequence case names.""" + if "all" in case_names: + return list(SEQUENCE_CASES.values()) + return [SEQUENCE_CASES[name] for name in case_names] + + +def _qpos(values, device): + """Create a torch qpos tensor.""" + torch = ensure_torch() + return torch.tensor(values, dtype=torch.float32, device=device) + + +def _targets_for_sequence(sequence_case: JointSequenceCase, device): + """Build typed MoveJoints targets for a sequence case.""" + from embodichain.lab.sim.atomic_actions import ( + JointPositionTarget, + NamedJointPositionTarget, + ) + + targets = [] + for index, name in enumerate(sequence_case.sequence): + if index == 0 and name == "ready": + targets.append(("move_joints", NamedJointPositionTarget(name="ready"))) + else: + targets.append( + ( + "move_joints", + JointPositionTarget(qpos=_qpos(JOINT_TARGETS[name], device)), + ) + ) + return targets + + +def _run_case( + sim, + robot, + atomic_engine, + initial_qpos, + case: JointSequenceCase, + repeat: int, + args: argparse.Namespace, + recorded_count: int, +): + """Run one MoveJoints case.""" + torch = ensure_torch() + reset_robot(robot, initial_qpos) + steps = _targets_for_sequence(case, sim.device) + elapsed, mem_delta, peak_gpu, result = timed_call(lambda: atomic_engine.run(steps)) + is_success, traj, _ = result + video_path = None + if should_record_case(args, recorded_count, bool(is_success)): + reset_robot(robot, initial_qpos) + video_path = replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=traj, + args=args, + video_path=build_video_output_path( + args, "atomic_action_move_joints", f"{case.name}_r{repeat}" + ), + ) + reset_robot(robot, initial_qpos) + + final_error_rad = None + if is_success and traj.shape[1] > 0: + arm_joint_ids = robot.get_joint_ids(name="arm") + final_qpos = traj[:, -1, arm_joint_ids][0] + expected = _qpos(JOINT_TARGETS[case.sequence[-1]], sim.device) + final_error_rad = float(torch.linalg.norm(final_qpos - expected)) + target_reached = bool( + is_success + and final_error_rad is not None + and final_error_rad <= SUCCESS_TOLERANCE_RAD + ) + return { + "case_id": f"{case.name}:r{repeat}", + "sequence_case": case.name, + "repeat": repeat, + "planning_success": bool(is_success), + "target_reached": target_reached, + "cost_time_ms": elapsed * 1000.0, + "cpu_delta_mb": mem_delta["cpu_mb"], + "gpu_delta_mb": mem_delta["gpu_mb"], + "peak_gpu_mb": peak_gpu, + "final_error_rad": final_error_rad, + "trajectory_waypoints": int(traj.shape[1]) if traj.ndim >= 2 else 0, + "failure_reason": "" if target_reached else "target_not_reached", + "video_path": str(video_path) if video_path is not None else "", + } + + +def _build_rows(results: list[dict[str, object]]): + """Build report rows for MoveJoints.""" + perf_rows = [] + metric_rows = [] + for result in results: + perf_rows.append( + { + "sample_size": 1, + "impl": "move_joints", + "case_id": result["case_id"], + "sequence_case": result["sequence_case"], + "repeat": result["repeat"], + "cost_time_ms": format_float(result["cost_time_ms"]), + "cpu_delta_mb": format_float(result["cpu_delta_mb"]), + "gpu_delta_mb": format_float(result["gpu_delta_mb"]), + "peak_gpu_mb": format_float(result["peak_gpu_mb"]), + } + ) + metric_rows.append( + { + "sample_size": 1, + "impl": "move_joints", + "case_id": result["case_id"], + "sequence_case": result["sequence_case"], + "success_rate": f"{float(result['target_reached']):.6f}", + "planning_success_rate": f"{float(result['planning_success']):.6f}", + "joint_err_rad": format_float(result["final_error_rad"]), + "trajectory_waypoints": result["trajectory_waypoints"], + "failure_reason": result["failure_reason"] or "N/A", + } + ) + return perf_rows, metric_rows + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: + """Run MoveJoints benchmark and write a markdown report.""" + args = _parse_args() if args is None else args + if args.repeat < 1: + raise ValueError("--repeat must be at least 1.") + profile = resolve_profile(args) + + ensure_repo_root() + ensure_torch() + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + MoveJoints, + MoveJointsCfg, + ) + from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.planners import ToppraPlannerCfg + from scripts.tutorials.atomic_action.move_joints import ( + create_robot, + initialize_simulation, + ) + + cases = _select_cases(args.sequence_cases) + repeat = 1 if profile == "smoke" else args.repeat + if profile == "smoke": + cases = [SEQUENCE_CASES["ready_home"]] + + print("=" * 60) + print("MoveJoints Atomic Action Benchmark") + print("=" * 60) + print( + "Coverage: " + f"profile={profile}, {len(cases)} sequence case(s) x " + f"{repeat} repeat(s)" + ) + + sim = initialize_simulation(args) + robot = create_robot(sim) + initial_qpos = robot.get_qpos().clone() + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + ready_qpos = _qpos(READY_QPOS, sim.device) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + MoveJoints( + motion_gen, + cfg=MoveJointsCfg( + control_part="arm", + sample_interval=MOVE_JOINTS_SAMPLE_INTERVAL, + named_joint_positions={"ready": ready_qpos}, + ), + ) + ) + + results: list[dict[str, object]] = [] + video_paths: list[str] = [] + print("\n=== MoveJoints Sequence Sweep ===") + for case in cases: + for repeat_index in range(repeat): + result = _run_case( + sim, + robot, + atomic_engine, + initial_qpos, + case, + repeat_index, + args, + len(video_paths), + ) + results.append(result) + if result["video_path"]: + video_paths.append(str(result["video_path"])) + print( + f" {result['case_id']:<24} " + f"time={result['cost_time_ms']:>10.2f} ms | " + f"success={result['target_reached']} " + f"err={format_float(result['final_error_rad'], precision=6)}" + ) + + perf_rows, metric_rows = _build_rows(results) + leaderboard_rows = build_single_action_leaderboard("move_joints", metric_rows) + report_path = write_markdown_report( + benchmark_name="atomic_action_move_joints", + perf_rows=perf_rows, + metric_rows=metric_rows, + leaderboard_rows=leaderboard_rows, + notes=[ + f"Profile: {profile}", + f"CPU memory backend: {CPU_MEMORY_BACKEND}", + f"Success tolerance: {SUCCESS_TOLERANCE_RAD} rad joint error.", + "Replay videos: " + (", ".join(video_paths) if video_paths else "disabled"), + ], + ) + print(f"Markdown report saved: {report_path}") + return report_path + + +def main() -> None: + """Run the CLI entry point.""" + try: + run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/pickup_benchmark.py b/scripts/benchmark/atomic_action/pickup_benchmark.py new file mode 100644 index 000000000..37395ca83 --- /dev/null +++ b/scripts/benchmark/atomic_action/pickup_benchmark.py @@ -0,0 +1,572 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Benchmark PickUp atomic action across grasp approach directions. + +Measures planning latency, memory usage, grasp planning success, held-object +state creation, and trajectory length for the PickUp action. +Run: python -m scripts.benchmark.atomic_action.pickup_benchmark +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + CPU_MEMORY_BACKEND, + add_common_benchmark_args, + add_grasp_benchmark_args, + add_object_position_benchmark_args, + add_pickup_approach_benchmark_args, + build_single_action_leaderboard, + build_video_output_path, + create_antipodal_object_semantics, + create_benchmark_object, + describe_object_preset, + ensure_repo_root, + ensure_torch, + format_float, + format_vector3, + MeshObjectPreset, + object_position_tuple, + park_rigid_object, + PHYSICAL_PICK_MIN_LIFT_M, + pickup_approach_direction_tuple, + PositionCase, + record_static_scene_video, + replay_trajectory_for_physical_validation, + replay_trajectory_with_recording, + reset_rigid_object, + reset_rigid_object_xy, + reset_robot, + resolve_pickup_approach_direction, + resolve_profile, + select_mesh_object_presets, + select_pickup_approaches, + select_position_cases, + should_record_case, + SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z, + SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT, + summarize_video_recording, + timed_call, + write_markdown_report, +) + +PICK_SAMPLE_INTERVAL = 120 +HAND_INTERP_STEPS = 12 + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add PickUp benchmark CLI arguments.""" + add_pickup_approach_benchmark_args(parser) + add_object_position_benchmark_args(parser) + add_grasp_benchmark_args(parser) + add_common_benchmark_args(parser) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Benchmark PickUp over grasp approach directions." + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _make_pickup_args( + args: argparse.Namespace, + approach: str, + object_preset: MeshObjectPreset, + profile: str, +) -> argparse.Namespace: + """Build a tutorial-compatible argparse namespace.""" + return argparse.Namespace( + object=object_preset.label, + n_sample=1000 if profile == "smoke" else args.n_sample, + force_reannotate=args.force_reannotate, + approach="custom" if approach == "side" else approach, + custom_approach_direction=None, + device=args.device, + renderer=args.renderer, + headless=True, + ) + + +def _run_case( + sim, + robot, + motion_gen, + initial_qpos, + args, + obj, + base_obj_pose, + object_preset: MeshObjectPreset, + position_case: PositionCase, + approach: str, + repeat: int, + recorded_count: int, + profile: str, +): + """Run one PickUp benchmark case.""" + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + GraspTarget, + PickUp, + PickUpCfg, + ) + from scripts.tutorials.atomic_action.pickup import ( + build_grasp_generator_cfg, + build_gripper_collision_cfg, + get_hand_open_close_qpos, + initialize_pre_pick_robot_pose, + compute_pick_close_end_step, + ) + + case_id = f"{object_preset.object_type}:{position_case.name}:{approach}:r{repeat}" + try: + reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + initial_obj_position = object_position_tuple(obj) + hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + case_args = _make_pickup_args(args, approach, object_preset, profile) + approach_direction = resolve_pickup_approach_direction( + approach, position_case, sim.device + ) + approach_direction_text = format_vector3( + pickup_approach_direction_tuple(approach, position_case) + ) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + PickUp( + motion_gen, + cfg=PickUpCfg( + control_part="arm", + hand_control_part="hand", + hand_open_qpos=hand_open, + hand_close_qpos=hand_close, + approach_direction=approach_direction, + pre_grasp_distance=0.15, + lift_height=0.16, + sample_interval=PICK_SAMPLE_INTERVAL, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ) + ) + semantics = create_antipodal_object_semantics( + obj=obj, + preset=object_preset, + args=case_args, + build_gripper_collision_cfg=build_gripper_collision_cfg, + build_grasp_generator_cfg=build_grasp_generator_cfg, + ) + elapsed, mem_delta, peak_gpu, result = timed_call( + lambda: atomic_engine.run( + steps=[("pick_up", GraspTarget(semantics=semantics))] + ) + ) + is_success, traj, final_state = result + final_obj_position = None + object_lift_delta_m = None + object_xy_drift_m = None + + if bool(is_success) and getattr(traj, "ndim", 0) >= 3 and traj.shape[1] > 0: + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + post_grasp_clear_step = compute_pick_close_end_step() + should_clear_object_dynamics = True + + def _on_validation_step(waypoint_index: int) -> None: + nonlocal should_clear_object_dynamics + if ( + should_clear_object_dynamics + and waypoint_index + 1 >= post_grasp_clear_step + ): + obj.clear_dynamics() + should_clear_object_dynamics = False + + final_obj_position = replay_trajectory_for_physical_validation( + sim=sim, + robot=robot, + obj=obj, + traj=traj, + on_step=_on_validation_step, + ) + if final_obj_position is not None: + object_lift_delta_m = final_obj_position[2] - initial_obj_position[2] + object_xy_drift_m = ( + (final_obj_position[0] - initial_obj_position[0]) ** 2 + + (final_obj_position[1] - initial_obj_position[1]) ** 2 + ) ** 0.5 + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + held_created = bool(is_success and final_state.held_object is not None) + physical_pick_success = bool( + held_created + and object_lift_delta_m is not None + and object_lift_delta_m >= PHYSICAL_PICK_MIN_LIFT_M + ) + + video_path = None + if should_record_case(args, recorded_count, physical_pick_success): + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + video_case_suffix = ( + f"{object_preset.object_type}_{position_case.name}_" + f"{approach}_r{repeat}" + ) + if getattr(traj, "ndim", 0) >= 3 and traj.shape[1] > 0: + post_grasp_clear_step = compute_pick_close_end_step() + should_clear_object_dynamics = True + + def _on_step(waypoint_index: int) -> None: + nonlocal should_clear_object_dynamics + if ( + should_clear_object_dynamics + and waypoint_index + 1 >= post_grasp_clear_step + ): + obj.clear_dynamics() + should_clear_object_dynamics = False + + video_path = replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=traj, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_pick_up", + video_case_suffix, + ), + on_step=_on_step, + ) + else: + video_path = record_static_scene_video( + sim=sim, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_pick_up", + f"{video_case_suffix}_failed_static", + ), + ) + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + lift_height_m = None + if is_success and traj.shape[1] > 0: + arm_joint_ids = robot.get_joint_ids(name="arm") + start_pose = robot.compute_fk( + qpos=traj[:, 0, arm_joint_ids], name="arm", to_matrix=True + )[0] + final_pose = robot.compute_fk( + qpos=traj[:, -1, arm_joint_ids], name="arm", to_matrix=True + )[0] + lift_height_m = float(final_pose[2, 3] - start_pose[2, 3]) + success = physical_pick_success + if success: + failure_reason = "" + elif not is_success: + failure_reason = "planning_failed" + elif not held_created: + failure_reason = "held_object_missing" + elif object_lift_delta_m is None: + failure_reason = "physical_replay_missing" + else: + failure_reason = "physical_pick_lift_too_low" + return { + "case_id": case_id, + "object_type": object_preset.object_type, + "material": object_preset.material_name, + "quadrant": position_case.quadrant, + "position_case": position_case.name, + "init_xy": position_case.xy, + "approach": approach, + "approach_direction": approach_direction_text, + "repeat": repeat, + "planning_success": bool(is_success), + "held_created": held_created, + "physical_pick_success": physical_pick_success, + "success": success, + "cost_time_ms": elapsed * 1000.0, + "cpu_delta_mb": mem_delta["cpu_mb"], + "gpu_delta_mb": mem_delta["gpu_mb"], + "peak_gpu_mb": peak_gpu, + "lift_height_m": lift_height_m, + "object_initial_z_m": initial_obj_position[2], + "object_final_z_m": ( + final_obj_position[2] if final_obj_position is not None else None + ), + "object_lift_delta_m": object_lift_delta_m, + "object_xy_drift_m": object_xy_drift_m, + "trajectory_waypoints": int(traj.shape[1]) if traj.ndim >= 2 else 0, + "failure_reason": failure_reason, + "video_path": str(video_path) if video_path is not None else "", + } + except Exception as exc: + video_path = None + if should_record_case(args, recorded_count, False): + try: + reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + video_path = record_static_scene_video( + sim=sim, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_pick_up", + ( + f"{object_preset.object_type}_{position_case.name}_" + f"{approach}_r{repeat}_exception_static" + ), + ), + ) + reset_rigid_object(obj, initial_obj_pose) + except Exception: + video_path = None + return { + "case_id": case_id, + "object_type": object_preset.object_type, + "material": object_preset.material_name, + "quadrant": position_case.quadrant, + "position_case": position_case.name, + "init_xy": position_case.xy, + "approach": approach, + "approach_direction": "N/A", + "repeat": repeat, + "planning_success": False, + "held_created": False, + "physical_pick_success": False, + "success": False, + "cost_time_ms": 0.0, + "cpu_delta_mb": 0.0, + "gpu_delta_mb": 0.0, + "peak_gpu_mb": 0.0, + "lift_height_m": None, + "object_initial_z_m": None, + "object_final_z_m": None, + "object_lift_delta_m": None, + "object_xy_drift_m": None, + "trajectory_waypoints": 0, + "failure_reason": f"exception:{type(exc).__name__}:{exc}", + "video_path": str(video_path) if video_path is not None else "", + } + + +def _build_rows(results: list[dict[str, object]]): + """Build report rows for PickUp.""" + perf_rows = [] + metric_rows = [] + for result in results: + perf_rows.append( + { + "sample_size": 1, + "impl": "pick_up", + "case_id": result["case_id"], + "object_type": result["object_type"], + "material": result["material"], + "quadrant": result["quadrant"], + "position_case": result["position_case"], + "init_xy": f"({result['init_xy'][0]:.3f},{result['init_xy'][1]:.3f})", + "approach": result["approach"], + "approach_direction": result["approach_direction"], + "repeat": result["repeat"], + "cost_time_ms": format_float(result["cost_time_ms"]), + "cpu_delta_mb": format_float(result["cpu_delta_mb"]), + "gpu_delta_mb": format_float(result["gpu_delta_mb"]), + "peak_gpu_mb": format_float(result["peak_gpu_mb"]), + } + ) + metric_rows.append( + { + "sample_size": 1, + "impl": "pick_up", + "case_id": result["case_id"], + "object_type": result["object_type"], + "material": result["material"], + "quadrant": result["quadrant"], + "position_case": result["position_case"], + "approach": result["approach"], + "approach_direction": result["approach_direction"], + "success_rate": f"{float(result['success']):.6f}", + "planning_success_rate": f"{float(result['planning_success']):.6f}", + "held_object_rate": f"{float(result['held_created']):.6f}", + "physical_pick_success_rate": ( + f"{float(result['physical_pick_success']):.6f}" + ), + "lift_height_m": format_float(result["lift_height_m"]), + "object_initial_z_m": format_float(result["object_initial_z_m"]), + "object_final_z_m": format_float(result["object_final_z_m"]), + "object_lift_delta_m": format_float(result["object_lift_delta_m"]), + "object_xy_drift_m": format_float(result["object_xy_drift_m"]), + "trajectory_waypoints": result["trajectory_waypoints"], + "failure_reason": result["failure_reason"] or "N/A", + } + ) + return perf_rows, metric_rows + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: + """Run PickUp benchmark and write a markdown report.""" + args = _parse_args() if args is None else args + if args.repeat < 1: + raise ValueError("--repeat must be at least 1.") + profile = resolve_profile(args) + + ensure_repo_root() + ensure_torch() + from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.planners import ToppraPlannerCfg + from scripts.tutorials.atomic_action.pickup import ( + create_robot, + initialize_simulation, + ) + + object_presets = select_mesh_object_presets(args.object_types, profile) + position_cases = select_position_cases(args.position_cases, profile) + approaches = select_pickup_approaches(args.approach_cases, profile) + repeat = 1 if profile == "smoke" else args.repeat + + print("=" * 60) + print("PickUp Atomic Action Benchmark") + print("=" * 60) + print( + "Coverage: " + f"profile={profile}, {len(object_presets)} object(s) x " + f"{len(position_cases)} position(s) x {len(approaches)} approach(es) " + f"x {repeat} repeat(s)" + ) + + sim = initialize_simulation(args) + robot = create_robot(sim) + initial_qpos = robot.get_qpos().clone() + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + object_pool = {} + for object_index, object_preset in enumerate(object_presets): + obj = create_benchmark_object( + sim=sim, + preset=object_preset, + position_case=position_cases[0], + uid_suffix="pool", + ) + base_pose = obj.get_local_pose(to_matrix=True).clone() + park_rigid_object(obj, base_pose, index=object_index, sim=sim) + object_pool[object_preset.object_type] = (obj, base_pose) + + results: list[dict[str, object]] = [] + video_paths: list[str] = [] + print("\n=== PickUp Object/Position/Approach Sweep ===") + for object_preset in object_presets: + obj, base_pose = object_pool[object_preset.object_type] + for parked_index, parked_preset in enumerate(object_presets): + if parked_preset.object_type == object_preset.object_type: + continue + parked_obj, parked_base_pose = object_pool[parked_preset.object_type] + park_rigid_object(parked_obj, parked_base_pose, index=parked_index, sim=sim) + for position_case in position_cases: + for approach in approaches: + for repeat_index in range(repeat): + result = _run_case( + sim, + robot, + motion_gen, + initial_qpos, + args, + obj, + base_pose, + object_preset, + position_case, + approach, + repeat_index, + len(video_paths), + profile, + ) + results.append(result) + if result["video_path"]: + video_paths.append(str(result["video_path"])) + print( + f" {result['case_id']:<42} " + f"time={result['cost_time_ms']:>10.2f} ms | " + f"success={result['success']} " + f"lift={format_float(result['lift_height_m'], precision=4)}" + ) + + perf_rows, metric_rows = _build_rows(results) + leaderboard_rows = build_single_action_leaderboard("pick_up", metric_rows) + report_path = write_markdown_report( + benchmark_name="atomic_action_pick_up", + perf_rows=perf_rows, + metric_rows=metric_rows, + leaderboard_rows=leaderboard_rows, + notes=[ + f"Profile: {profile}", + "Object presets: " + + ", ".join(describe_object_preset(preset) for preset in object_presets), + "Position cases: " + + ", ".join( + f"{case.name}/{case.quadrant}/xy={case.xy}" for case in position_cases + ), + "PickUp approach cases: " + ", ".join(approaches), + f"CPU memory backend: {CPU_MEMORY_BACKEND}", + f"n_sample: {1000 if profile == 'smoke' else args.n_sample}", + ( + "Side grasp opening-axis rule: prefer candidates with " + f"abs(opening_axis_z) <= {SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z:.2f}; " + "fallback cost += " + f"abs(opening_axis_z) * {SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT:.2f}." + ), + ( + "Physical success rule: planning success, held-object state " + f"created, and object lift delta >= {PHYSICAL_PICK_MIN_LIFT_M:.3f} m." + ), + *summarize_video_recording(args, results, video_paths), + ], + ) + print(f"Markdown report saved: {report_path}") + return report_path + + +def main() -> None: + """Run the CLI entry point.""" + try: + run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/place_benchmark.py b/scripts/benchmark/atomic_action/place_benchmark.py new file mode 100644 index 000000000..6ea98d68d --- /dev/null +++ b/scripts/benchmark/atomic_action/place_benchmark.py @@ -0,0 +1,736 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Benchmark Place atomic action after a PickUp precondition. + +Measures Place-only planning latency and memory usage once a held-object state +has been produced by the PickUp action. +Run: python -m scripts.benchmark.atomic_action.place_benchmark +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + CPU_MEMORY_BACKEND, + add_common_benchmark_args, + add_grasp_benchmark_args, + add_object_position_benchmark_args, + add_pickup_approach_benchmark_args, + build_single_action_leaderboard, + build_video_output_path, + create_antipodal_object_semantics, + create_benchmark_object, + describe_object_preset, + ensure_repo_root, + ensure_torch, + format_float, + format_vector3, + MeshObjectPreset, + object_position_tuple, + park_rigid_object, + PHYSICAL_PICK_MIN_LIFT_M, + PHYSICAL_PLACE_XY_TOLERANCE_M, + pickup_approach_direction_tuple, + PositionCase, + record_static_scene_video, + replay_trajectory_for_physical_validation, + replay_trajectory_with_recording, + reset_rigid_object, + reset_rigid_object_xy, + reset_robot, + resolve_pickup_approach_direction, + resolve_profile, + select_mesh_object_presets, + select_pickup_approaches, + select_position_cases, + should_record_case, + SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z, + SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT, + summarize_video_recording, + timed_call, + write_markdown_report, + xy_distance_m, +) + + +@dataclass(frozen=True) +class PlaceCase: + """Place target benchmark case.""" + + name: str + xyz: tuple[float, float, float] + + +PLACE_CASES = { + "left_bin": PlaceCase("left_bin", (-0.20, 0.28, 0.10)), + "right_bin": PlaceCase("right_bin", (-0.20, -0.28, 0.10)), + "center": PlaceCase("center", (-0.35, 0.00, 0.12)), +} +PICK_SAMPLE_INTERVAL = 120 +PLACE_SAMPLE_INTERVAL = 120 +HAND_INTERP_STEPS = 12 +PLACE_LIFT_HEIGHT = 0.14 + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add Place benchmark CLI arguments.""" + add_pickup_approach_benchmark_args(parser) + parser.add_argument( + "--place_cases", + nargs="+", + choices=(*PLACE_CASES.keys(), "all"), + default=None, + help=( + "Place target cases to benchmark. Defaults are selected by " + "--profile; use 'all' for every case." + ), + ) + add_object_position_benchmark_args(parser) + add_grasp_benchmark_args(parser) + add_common_benchmark_args(parser) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Benchmark Place after a PickUp precondition." + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _selected_cases( + case_names: list[str] | None, + profile: str, +) -> list[PlaceCase]: + """Resolve selected place case names.""" + if not case_names: + return [PLACE_CASES["left_bin"]] + if "all" in case_names: + return list(PLACE_CASES.values()) + return [PLACE_CASES[name] for name in case_names] + + +def _make_place_pose(device, xyz: tuple[float, float, float]): + """Create a top-down place target pose.""" + torch = ensure_torch() + pose = torch.eye(4, dtype=torch.float32, device=device) + pose[:3, :3] = torch.tensor( + [ + [-0.0539, -0.9985, -0.0022], + [-0.9977, 0.0540, -0.0401], + [0.0401, 0.0000, -0.9992], + ], + dtype=torch.float32, + device=device, + ) + pose[:3, 3] = torch.tensor(xyz, dtype=torch.float32, device=device) + return pose + + +def _make_pickup_args( + args: argparse.Namespace, + object_preset: MeshObjectPreset, + profile: str, +) -> argparse.Namespace: + """Build a tutorial-compatible argparse namespace.""" + return argparse.Namespace( + object=object_preset.label, + n_sample=1000 if profile == "smoke" else args.n_sample, + force_reannotate=args.force_reannotate, + device=args.device, + renderer=args.renderer, + headless=True, + ) + + +def _prepare_held_state( + sim, + robot, + obj, + motion_gen, + args, + object_preset: MeshObjectPreset, + position_case: PositionCase, + pickup_approach: str, + profile: str, +): + """Run PickUp precondition outside the timed Place block.""" + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + GraspTarget, + PickUp, + PickUpCfg, + ) + from scripts.tutorials.atomic_action.place import ( + build_grasp_generator_cfg, + build_gripper_collision_cfg, + get_hand_open_close_qpos, + initialize_pre_pick_robot_pose, + ) + + hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + PickUp( + motion_gen, + cfg=PickUpCfg( + control_part="arm", + hand_control_part="hand", + hand_open_qpos=hand_open, + hand_close_qpos=hand_close, + approach_direction=resolve_pickup_approach_direction( + pickup_approach, position_case, sim.device + ), + pre_grasp_distance=0.15, + lift_height=0.16, + sample_interval=PICK_SAMPLE_INTERVAL, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ) + ) + semantics = create_antipodal_object_semantics( + obj=obj, + preset=object_preset, + args=_make_pickup_args(args, object_preset, profile), + build_gripper_collision_cfg=build_gripper_collision_cfg, + build_grasp_generator_cfg=build_grasp_generator_cfg, + ) + is_success, traj, state = atomic_engine.run( + steps=[("pick_up", GraspTarget(semantics=semantics))] + ) + if not is_success or state.held_object is None: + raise RuntimeError("Failed to prepare held-object state for Place benchmark.") + robot.set_qpos(state.last_qpos) + return state, hand_open, hand_close, traj + + +def _run_case( + sim, + robot, + motion_gen, + initial_qpos, + args, + obj, + base_obj_pose, + object_preset: MeshObjectPreset, + position_case: PositionCase, + pickup_approach: str, + case: PlaceCase, + repeat: int, + recorded_count: int, + profile: str, +): + """Run one Place benchmark case.""" + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + EndEffectorPoseTarget, + Place, + PlaceCfg, + ) + from scripts.tutorials.atomic_action.place import ( + compute_pick_close_end_step, + initialize_pre_pick_robot_pose, + ) + + case_id = ( + f"{object_preset.object_type}:{position_case.name}:" + f"{pickup_approach}:{case.name}:r{repeat}" + ) + try: + reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + reset_rigid_object(obj, initial_obj_pose) + initial_obj_position = object_position_tuple(obj) + state, hand_open, hand_close, precondition_traj = _prepare_held_state( + sim, + robot, + obj, + motion_gen, + args, + object_preset, + position_case, + pickup_approach, + profile, + ) + approach_direction_text = format_vector3( + pickup_approach_direction_tuple(pickup_approach, position_case) + ) + precondition_waypoints = int(precondition_traj.shape[1]) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + Place( + motion_gen, + cfg=PlaceCfg( + control_part="arm", + hand_control_part="hand", + hand_open_qpos=hand_open, + hand_close_qpos=hand_close, + lift_height=PLACE_LIFT_HEIGHT, + sample_interval=PLACE_SAMPLE_INTERVAL, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ) + ) + place_pose = _make_place_pose(sim.device, case.xyz) + elapsed, mem_delta, peak_gpu, result = timed_call( + lambda: atomic_engine.run( + steps=[("place", EndEffectorPoseTarget(xpos=place_pose))], + state=state, + ) + ) + is_success, traj, final_state = result + torch = ensure_torch() + precondition_obj_position = None + final_obj_position = None + object_lift_delta_m = None + place_xy_error_m = None + place_z_error_m = None + + if ( + getattr(precondition_traj, "ndim", 0) >= 3 + and precondition_traj.shape[1] > 0 + ): + if bool(is_success) and getattr(traj, "ndim", 0) >= 3 and traj.shape[1] > 0: + validation_traj = torch.cat((precondition_traj, traj), dim=1) + else: + validation_traj = precondition_traj + + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + post_grasp_clear_step = compute_pick_close_end_step() + should_clear_object_dynamics = True + + def _on_validation_step(waypoint_index: int) -> None: + nonlocal should_clear_object_dynamics, precondition_obj_position + if ( + should_clear_object_dynamics + and waypoint_index + 1 >= post_grasp_clear_step + ): + obj.clear_dynamics() + should_clear_object_dynamics = False + if ( + precondition_obj_position is None + and waypoint_index + 1 >= precondition_waypoints + ): + precondition_obj_position = object_position_tuple(obj) + + final_obj_position = replay_trajectory_for_physical_validation( + sim=sim, + robot=robot, + obj=obj, + traj=validation_traj, + on_step=_on_validation_step, + ) + if precondition_obj_position is not None: + object_lift_delta_m = ( + precondition_obj_position[2] - initial_obj_position[2] + ) + if bool(is_success) and final_obj_position is not None: + place_xy_error_m = xy_distance_m(final_obj_position, case.xyz) + place_z_error_m = abs(final_obj_position[2] - case.xyz[2]) + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + released = bool(is_success and final_state.held_object is None) + physical_pick_success = bool( + object_lift_delta_m is not None + and object_lift_delta_m >= PHYSICAL_PICK_MIN_LIFT_M + ) + physical_place_success = bool( + released + and physical_pick_success + and place_xy_error_m is not None + and place_xy_error_m <= PHYSICAL_PLACE_XY_TOLERANCE_M + ) + success = physical_place_success + if success: + failure_reason = "" + elif not is_success: + failure_reason = "planning_failed" + elif not released: + failure_reason = "held_object_not_released" + elif not physical_pick_success: + failure_reason = "physical_pick_failed" + elif place_xy_error_m is None: + failure_reason = "physical_replay_missing" + else: + failure_reason = "place_target_miss" + + video_path = None + if should_record_case(args, recorded_count, success): + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + initialize_pre_pick_robot_pose(robot, obj, hand_open) + video_case_suffix = ( + f"{object_preset.object_type}_{position_case.name}_" + f"{pickup_approach}_{case.name}_r{repeat}" + ) + if getattr(traj, "ndim", 0) >= 3 and traj.shape[1] > 0: + replay_traj = torch.cat((precondition_traj, traj), dim=1) + else: + replay_traj = precondition_traj + + if getattr(replay_traj, "ndim", 0) >= 3 and replay_traj.shape[1] > 0: + post_grasp_clear_step = compute_pick_close_end_step() + should_clear_object_dynamics = True + + def _on_step(waypoint_index: int) -> None: + nonlocal should_clear_object_dynamics + if ( + should_clear_object_dynamics + and waypoint_index + 1 >= post_grasp_clear_step + ): + obj.clear_dynamics() + should_clear_object_dynamics = False + + video_path = replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=replay_traj, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_place", + video_case_suffix, + ), + on_step=_on_step, + ) + else: + video_path = record_static_scene_video( + sim=sim, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_place", + f"{video_case_suffix}_failed_static", + ), + ) + reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + return { + "case_id": case_id, + "object_type": object_preset.object_type, + "material": object_preset.material_name, + "quadrant": position_case.quadrant, + "position_case": position_case.name, + "init_xy": position_case.xy, + "pickup_approach": pickup_approach, + "approach_direction": approach_direction_text, + "place_case": case.name, + "repeat": repeat, + "planning_success": bool(is_success), + "released": released, + "physical_pick_success": physical_pick_success, + "physical_place_success": physical_place_success, + "success": success, + "cost_time_ms": elapsed * 1000.0, + "cpu_delta_mb": mem_delta["cpu_mb"], + "gpu_delta_mb": mem_delta["gpu_mb"], + "peak_gpu_mb": peak_gpu, + "object_lift_delta_m": object_lift_delta_m, + "object_final_x_m": ( + final_obj_position[0] if final_obj_position is not None else None + ), + "object_final_y_m": ( + final_obj_position[1] if final_obj_position is not None else None + ), + "object_final_z_m": ( + final_obj_position[2] if final_obj_position is not None else None + ), + "place_xy_error_m": place_xy_error_m, + "place_z_error_m": place_z_error_m, + "precondition_waypoints": precondition_waypoints, + "trajectory_waypoints": int(traj.shape[1]) if traj.ndim >= 2 else 0, + "failure_reason": failure_reason, + "video_path": str(video_path) if video_path is not None else "", + } + except Exception as exc: + video_path = None + if should_record_case(args, recorded_count, False): + try: + reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + video_path = record_static_scene_video( + sim=sim, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_place", + ( + f"{object_preset.object_type}_{position_case.name}_" + f"{pickup_approach}_{case.name}_r{repeat}" + "_exception_static" + ), + ), + ) + reset_rigid_object(obj, initial_obj_pose) + except Exception: + video_path = None + return { + "case_id": case_id, + "object_type": object_preset.object_type, + "material": object_preset.material_name, + "quadrant": position_case.quadrant, + "position_case": position_case.name, + "init_xy": position_case.xy, + "pickup_approach": pickup_approach, + "approach_direction": "N/A", + "place_case": case.name, + "repeat": repeat, + "planning_success": False, + "released": False, + "physical_pick_success": False, + "physical_place_success": False, + "success": False, + "cost_time_ms": 0.0, + "cpu_delta_mb": 0.0, + "gpu_delta_mb": 0.0, + "peak_gpu_mb": 0.0, + "object_lift_delta_m": None, + "object_final_x_m": None, + "object_final_y_m": None, + "object_final_z_m": None, + "place_xy_error_m": None, + "place_z_error_m": None, + "precondition_waypoints": 0, + "trajectory_waypoints": 0, + "failure_reason": f"exception:{type(exc).__name__}:{exc}", + "video_path": str(video_path) if video_path is not None else "", + } + + +def _build_rows(results: list[dict[str, object]]): + """Build report rows for Place.""" + perf_rows = [] + metric_rows = [] + for result in results: + perf_rows.append( + { + "sample_size": 1, + "impl": "place", + "case_id": result["case_id"], + "object_type": result["object_type"], + "material": result["material"], + "quadrant": result["quadrant"], + "position_case": result["position_case"], + "init_xy": f"({result['init_xy'][0]:.3f},{result['init_xy'][1]:.3f})", + "pickup_approach": result["pickup_approach"], + "approach_direction": result["approach_direction"], + "place_case": result["place_case"], + "repeat": result["repeat"], + "cost_time_ms": format_float(result["cost_time_ms"]), + "cpu_delta_mb": format_float(result["cpu_delta_mb"]), + "gpu_delta_mb": format_float(result["gpu_delta_mb"]), + "peak_gpu_mb": format_float(result["peak_gpu_mb"]), + } + ) + metric_rows.append( + { + "sample_size": 1, + "impl": "place", + "case_id": result["case_id"], + "object_type": result["object_type"], + "material": result["material"], + "quadrant": result["quadrant"], + "position_case": result["position_case"], + "pickup_approach": result["pickup_approach"], + "approach_direction": result["approach_direction"], + "place_case": result["place_case"], + "success_rate": f"{float(result['success']):.6f}", + "planning_success_rate": f"{float(result['planning_success']):.6f}", + "release_success_rate": f"{float(result['released']):.6f}", + "physical_pick_success_rate": ( + f"{float(result['physical_pick_success']):.6f}" + ), + "physical_place_success_rate": ( + f"{float(result['physical_place_success']):.6f}" + ), + "object_lift_delta_m": format_float(result["object_lift_delta_m"]), + "object_final_x_m": format_float(result["object_final_x_m"]), + "object_final_y_m": format_float(result["object_final_y_m"]), + "object_final_z_m": format_float(result["object_final_z_m"]), + "place_xy_error_m": format_float(result["place_xy_error_m"]), + "place_z_error_m": format_float(result["place_z_error_m"]), + "precondition_waypoints": result["precondition_waypoints"], + "trajectory_waypoints": result["trajectory_waypoints"], + "failure_reason": result["failure_reason"] or "N/A", + } + ) + return perf_rows, metric_rows + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: + """Run Place benchmark and write a markdown report.""" + args = _parse_args() if args is None else args + if args.repeat < 1: + raise ValueError("--repeat must be at least 1.") + profile = resolve_profile(args) + + ensure_repo_root() + ensure_torch() + from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.planners import ToppraPlannerCfg + from scripts.tutorials.atomic_action.place import ( + create_robot, + initialize_simulation, + ) + + object_presets = select_mesh_object_presets(args.object_types, profile) + position_cases = select_position_cases(args.position_cases, profile) + approaches = select_pickup_approaches(args.approach_cases, profile) + cases = _selected_cases(args.place_cases, profile) + repeat = 1 if profile == "smoke" else args.repeat + + print("=" * 60) + print("Place Atomic Action Benchmark") + print("=" * 60) + print( + "Coverage: " + f"profile={profile}, {len(object_presets)} object(s) x " + f"{len(position_cases)} position(s) x {len(approaches)} pickup approach(es) " + f"x {len(cases)} place target(s) " + f"x {repeat} repeat(s)" + ) + + sim = initialize_simulation(args) + robot = create_robot(sim) + initial_qpos = robot.get_qpos().clone() + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + object_pool = {} + for object_index, object_preset in enumerate(object_presets): + obj = create_benchmark_object( + sim=sim, + preset=object_preset, + position_case=position_cases[0], + uid_suffix="pool", + ) + base_pose = obj.get_local_pose(to_matrix=True).clone() + park_rigid_object(obj, base_pose, index=object_index, sim=sim) + object_pool[object_preset.object_type] = (obj, base_pose) + + results: list[dict[str, object]] = [] + video_paths: list[str] = [] + print("\n=== Place Object/Position/Target Sweep ===") + for object_preset in object_presets: + obj, base_pose = object_pool[object_preset.object_type] + for parked_index, parked_preset in enumerate(object_presets): + if parked_preset.object_type == object_preset.object_type: + continue + parked_obj, parked_base_pose = object_pool[parked_preset.object_type] + park_rigid_object(parked_obj, parked_base_pose, index=parked_index, sim=sim) + for position_case in position_cases: + for pickup_approach in approaches: + for case in cases: + for repeat_index in range(repeat): + result = _run_case( + sim, + robot, + motion_gen, + initial_qpos, + args, + obj, + base_pose, + object_preset, + position_case, + pickup_approach, + case, + repeat_index, + len(video_paths), + profile, + ) + results.append(result) + if result["video_path"]: + video_paths.append(str(result["video_path"])) + print( + f" {result['case_id']:<48} " + f"time={result['cost_time_ms']:>10.2f} ms | " + f"success={result['success']}" + ) + + perf_rows, metric_rows = _build_rows(results) + leaderboard_rows = build_single_action_leaderboard("place", metric_rows) + report_path = write_markdown_report( + benchmark_name="atomic_action_place", + perf_rows=perf_rows, + metric_rows=metric_rows, + leaderboard_rows=leaderboard_rows, + notes=[ + "Timed block includes Place only; PickUp precondition is prepared " + "outside timing.", + f"Profile: {profile}", + "Object presets: " + + ", ".join(describe_object_preset(preset) for preset in object_presets), + "Position cases: " + + ", ".join( + f"{case.name}/{case.quadrant}/xy={case.xy}" for case in position_cases + ), + "PickUp approach cases: " + ", ".join(approaches), + "Place target cases: " + ", ".join(case.name for case in cases), + f"CPU memory backend: {CPU_MEMORY_BACKEND}", + f"n_sample: {1000 if profile == 'smoke' else args.n_sample}", + ( + "Side grasp opening-axis rule: prefer candidates with " + f"abs(opening_axis_z) <= {SIDE_GRASP_MAX_OPEN_AXIS_ABS_Z:.2f}; " + "fallback cost += " + f"abs(opening_axis_z) * {SIDE_GRASP_OPEN_AXIS_Z_COST_WEIGHT:.2f}." + ), + ( + "Physical success rule: PickUp precondition lift delta >= " + f"{PHYSICAL_PICK_MIN_LIFT_M:.3f} m, Place released the object, " + f"and final object XY error <= {PHYSICAL_PLACE_XY_TOLERANCE_M:.3f} m." + ), + *summarize_video_recording(args, results, video_paths), + ], + ) + print(f"Markdown report saved: {report_path}") + return report_path + + +def main() -> None: + """Run the CLI entry point.""" + try: + run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/press_benchmark.py b/scripts/benchmark/atomic_action/press_benchmark.py new file mode 100644 index 000000000..a2c874a73 --- /dev/null +++ b/scripts/benchmark/atomic_action/press_benchmark.py @@ -0,0 +1,985 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Benchmark Press atomic action across object presets and start positions. + +The benchmark sweeps object presets such as bottle and mug against multiple +initial XY positions that cover all four workspace quadrants. It reports +planning latency, memory usage, planning success, and whether the generated +trajectory reaches the object's top center. +Run: python -m scripts.benchmark.atomic_action.press_benchmark +""" + +from __future__ import annotations + +import argparse +import math +import os +import resource +import sys +import time +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + add_profile_benchmark_args, + add_video_benchmark_args, + build_video_output_path, + COVERAGE_POSITION_CASE_NAMES, + park_rigid_object, + replay_trajectory_with_recording, + reset_rigid_object, + reset_rigid_object_xy, + resolve_profile, + should_record_case, + SMOKE_POSITION_CASE_NAMES, +) + +try: + import psutil +except ModuleNotFoundError: + psutil = None + +CPU_MEMORY_BACKEND = "psutil" if psutil is not None else "resource" +_RUNTIME_IMPORTS_READY = False + +# Keep these constants aligned with scripts/tutorials/atomic_action/press.py. +DEFAULT_PRESS_TOLERANCE = 0.01 +MOVE_SAMPLE_INTERVAL = 60 +PRESS_SAMPLE_INTERVAL = 90 +HAND_INTERP_STEPS = 12 +TABLE_TOP_Z = -0.045 +PRESS_CLEARANCE = 0.13 +PRESS_SURFACE_OFFSET = 0.003 + + +def _ensure_runtime_imports() -> None: + """Import simulation dependencies only when the benchmark is executed.""" + global _RUNTIME_IMPORTS_READY + if _RUNTIME_IMPORTS_READY: + return + + repo_root = Path(__file__).resolve().parents[3] + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + + try: + import torch as torch_module + from embodichain.lab.sim import SimulationManager as simulation_manager_cls + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine as atomic_action_engine_cls, + EndEffectorPoseTarget as end_effector_pose_target_cls, + MoveEndEffector as move_end_effector_cls, + MoveEndEffectorCfg as move_end_effector_cfg_cls, + Press as press_cls, + PressCfg as press_cfg_cls, + ) + from embodichain.lab.sim.cfg import ( + RigidBodyAttributesCfg as rigid_body_attributes_cfg_cls, + RigidObjectCfg as rigid_object_cfg_cls, + ) + from embodichain.lab.sim.material import ( + VisualMaterialCfg as visual_material_cfg_cls, + ) + from embodichain.lab.sim.objects import RigidObject as rigid_object_cls + from embodichain.lab.sim.objects import Robot as robot_cls + from embodichain.lab.sim.planners import ( + MotionGenerator as motion_generator_cls, + MotionGenCfg as motion_gen_cfg_cls, + ToppraPlannerCfg as toppra_planner_cfg_cls, + ) + from embodichain.lab.sim.shapes import CubeCfg as cube_cfg_cls + from scripts.tutorials.atomic_action.press import ( + create_robot as create_robot_fn, + create_table as create_table_fn, + get_hand_close_qpos as get_hand_close_qpos_fn, + initialize_simulation as initialize_simulation_fn, + make_top_down_eef_pose as make_top_down_eef_pose_fn, + settle_object as settle_object_fn, + ) + except ModuleNotFoundError as exc: + raise RuntimeError( + "Atomic action benchmark requires the EmbodiChain simulation runtime " + f"and PyTorch. Missing module: {exc.name}." + ) from exc + + globals().update( + { + "torch": torch_module, + "SimulationManager": simulation_manager_cls, + "AtomicActionEngine": atomic_action_engine_cls, + "EndEffectorPoseTarget": end_effector_pose_target_cls, + "MoveEndEffector": move_end_effector_cls, + "MoveEndEffectorCfg": move_end_effector_cfg_cls, + "Press": press_cls, + "PressCfg": press_cfg_cls, + "RigidBodyAttributesCfg": rigid_body_attributes_cfg_cls, + "RigidObjectCfg": rigid_object_cfg_cls, + "VisualMaterialCfg": visual_material_cfg_cls, + "RigidObject": rigid_object_cls, + "Robot": robot_cls, + "MotionGenerator": motion_generator_cls, + "MotionGenCfg": motion_gen_cfg_cls, + "ToppraPlannerCfg": toppra_planner_cfg_cls, + "CubeCfg": cube_cfg_cls, + "create_robot": create_robot_fn, + "create_table": create_table_fn, + "get_hand_close_qpos": get_hand_close_qpos_fn, + "initialize_simulation": initialize_simulation_fn, + "make_top_down_eef_pose": make_top_down_eef_pose_fn, + "settle_object": settle_object_fn, + } + ) + _RUNTIME_IMPORTS_READY = True + + +@dataclass(frozen=True) +class ObjectPreset: + """Primitive object preset used by the atomic-action benchmark.""" + + object_type: str + material_name: str + size: tuple[float, float, float] + base_color: tuple[float, float, float, float] + roughness: float + dynamic_friction: float = 0.8 + static_friction: float = 0.9 + + +@dataclass(frozen=True) +class PositionCase: + """Initial object position case with a quadrant label.""" + + name: str + quadrant: str + xy: tuple[float, float] + + +@dataclass(frozen=True) +class PressCaseResult: + """Result for one Press benchmark case.""" + + case_id: str + object_type: str + material_name: str + quadrant: str + position_case: str + init_xy: tuple[float, float] + repeat_index: int + planning_success: bool + center_hit: bool + cost_time_ms: float + cpu_delta_mb: float + gpu_delta_mb: float + peak_gpu_mb: float + xy_error_m: float | None + hit_step: int | None + trajectory_waypoints: int + failure_reason: str + video_path: str = "" + + +OBJECT_PRESETS: dict[str, ObjectPreset] = { + "bottle": ObjectPreset( + object_type="bottle", + material_name="green_plastic", + size=(0.06, 0.06, 0.16), + base_color=(0.10, 0.45, 0.32, 1.0), + roughness=0.55, + ), + "mug": ObjectPreset( + object_type="mug", + material_name="ceramic", + size=(0.10, 0.08, 0.10), + base_color=(0.88, 0.85, 0.78, 1.0), + roughness=0.35, + ), + "wooden_block": ObjectPreset( + object_type="wooden_block", + material_name="wood", + size=(0.12, 0.12, 0.06), + base_color=(0.58, 0.32, 0.14, 1.0), + roughness=0.85, + ), +} + +POSITION_CASES: dict[str, PositionCase] = { + "q1_near": PositionCase(name="q1_near", quadrant="q1", xy=(0.02, 0.18)), + "q1_far": PositionCase(name="q1_far", quadrant="q1", xy=(0.12, 0.36)), + "q2_near": PositionCase(name="q2_near", quadrant="q2", xy=(-0.42, 0.18)), + "q2_far": PositionCase(name="q2_far", quadrant="q2", xy=(-0.62, 0.36)), + "q3_near": PositionCase(name="q3_near", quadrant="q3", xy=(-0.42, -0.18)), + "q3_far": PositionCase(name="q3_far", quadrant="q3", xy=(-0.62, -0.36)), + "q4_near": PositionCase(name="q4_near", quadrant="q4", xy=(0.02, -0.18)), + "q4_far": PositionCase(name="q4_far", quadrant="q4", xy=(0.12, -0.36)), +} + +DEFAULT_OBJECT_TYPES = ("bottle", "mug") +FULL_OBJECT_TYPES = tuple(OBJECT_PRESETS.keys()) +SMOKE_OBJECT_TYPES = ("bottle",) + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add atomic-action benchmark arguments to an argument parser.""" + add_profile_benchmark_args(parser) + parser.add_argument( + "--object_types", + nargs="+", + choices=(*OBJECT_PRESETS.keys(), "all"), + default=None, + help=( + "Object presets to benchmark. Defaults are selected by --profile; " + "use 'all' to include every preset." + ), + ) + parser.add_argument( + "--position_cases", + nargs="+", + choices=(*POSITION_CASES.keys(), "all"), + default=None, + help=( + "Initial position cases to benchmark. Defaults are selected by " + "--profile; use 'all' for all near/far cases." + ), + ) + parser.add_argument( + "--repeat", + type=int, + default=1, + help="Number of repeats for every object-position case.", + ) + parser.add_argument( + "--smoke", + action="store_true", + help="Alias for --profile smoke.", + ) + parser.add_argument( + "--device", + type=str, + default="cpu", + help="Simulation device, e.g. 'cpu' or 'cuda'.", + ) + parser.add_argument( + "--renderer", + type=str, + choices=("auto", "hybrid", "fast-rt", "rt"), + default="auto", + help="Renderer backend used by SimulationManager.", + ) + add_video_benchmark_args(parser) + parser.add_argument( + "--press_tolerance", + type=float, + default=DEFAULT_PRESS_TOLERANCE, + help="XY tolerance in meters for the press-center check.", + ) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments for the atomic-action benchmark.""" + parser = argparse.ArgumentParser( + description=( + "Benchmark Press atomic action over object presets and initial " + "workspace quadrants." + ) + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _sync_cuda() -> None: + """Synchronize CUDA stream when available.""" + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def _reset_peak_gpu_memory() -> None: + """Reset PyTorch peak GPU memory stats when CUDA is available.""" + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + + +def _peak_gpu_memory_mb() -> float: + """Return peak GPU memory allocated by PyTorch in MB.""" + if not torch.cuda.is_available(): + return 0.0 + return torch.cuda.max_memory_allocated() / 1024**2 + + +def _memory_snapshot() -> dict[str, float]: + """Return current process memory usage snapshot in MB.""" + if psutil is not None: + process = psutil.Process(os.getpid()) + cpu_mb = process.memory_info().rss / 1024**2 + else: + cpu_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0 + gpu_mb = ( + torch.cuda.memory_allocated() / 1024**2 if torch.cuda.is_available() else 0.0 + ) + return {"cpu_mb": cpu_mb, "gpu_mb": gpu_mb} + + +def _format_float(value: float | None, precision: int = 6) -> str: + """Format finite floats for tables and use N/A for missing values.""" + if value is None or not math.isfinite(value): + return "N/A" + return f"{value:.{precision}f}" + + +def _format_markdown_table(rows: list[dict[str, object]]) -> list[str]: + """Format rows into a markdown table.""" + if not rows: + return ["No data."] + + headers = list(rows[0].keys()) + lines = [ + "| " + " | ".join(headers) + " |", + "| " + " | ".join(["---"] * len(headers)) + " |", + ] + for row in rows: + lines.append("| " + " | ".join(str(row[h]) for h in headers) + " |") + return lines + + +def _write_markdown_report( + benchmark_name: str, + perf_rows: list[dict[str, object]], + metric_rows: list[dict[str, object]], + leaderboard_rows: list[dict[str, object]], + notes: list[str] | None = None, +) -> Path: + """Write benchmark results into one markdown report file.""" + output_dir = Path("outputs/benchmarks") + output_dir.mkdir(parents=True, exist_ok=True) + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + report_path = output_dir / f"{benchmark_name}_{timestamp}.md" + + lines: list[str] = [ + f"# {benchmark_name} Benchmark Report", + "", + f"Generated at: {datetime.now().isoformat(timespec='seconds')}", + "", + "## Time & Memory", + "", + ] + lines.extend(_format_markdown_table(perf_rows)) + lines.extend(["", "## Success & Other Metrics", ""]) + lines.extend(_format_markdown_table(metric_rows)) + lines.extend(["", "## Leaderboard", ""]) + lines.extend(_format_markdown_table(leaderboard_rows)) + + if notes: + lines.extend(["", "## Notes", ""]) + lines.extend([f"- {note}" for note in notes]) + + report_path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return report_path + + +def _default_object_types_for_profile(profile: str) -> tuple[str, ...]: + """Return default Press primitive object names for a profile.""" + if profile in ("smoke", "coverage", "full"): + return SMOKE_OBJECT_TYPES + raise ValueError(f"Unsupported benchmark profile: {profile}") + + +def _select_object_presets( + object_types: list[str] | None, + profile: str, +) -> list[ObjectPreset]: + """Resolve selected object preset names.""" + if not object_types: + object_types = list(_default_object_types_for_profile(profile)) + if "all" in object_types: + return list(OBJECT_PRESETS.values()) + return [OBJECT_PRESETS[name] for name in object_types] + + +def _default_position_cases_for_profile(profile: str) -> tuple[str, ...]: + """Return default Press position case names for a profile.""" + if profile == "smoke": + return SMOKE_POSITION_CASE_NAMES + if profile in ("coverage", "full"): + return COVERAGE_POSITION_CASE_NAMES + raise ValueError(f"Unsupported benchmark profile: {profile}") + + +def _select_position_cases( + position_cases: list[str] | None, + profile: str, +) -> list[PositionCase]: + """Resolve selected position case names.""" + if not position_cases: + position_cases = list(_default_position_cases_for_profile(profile)) + if "all" in position_cases: + return list(POSITION_CASES.values()) + return [POSITION_CASES[name] for name in position_cases] + + +def _create_benchmark_object( + sim: SimulationManager, + preset: ObjectPreset, + position_case: PositionCase, + repeat_index: int, +) -> RigidObject: + """Create one static benchmark object at the requested initial position.""" + init_pos = ( + position_case.xy[0], + position_case.xy[1], + TABLE_TOP_Z + 0.5 * preset.size[2], + ) + uid = f"atomic_benchmark_{preset.object_type}_{position_case.name}_{repeat_index}" + cfg = RigidObjectCfg( + uid=uid, + shape=CubeCfg( + size=list(preset.size), + visual_material=VisualMaterialCfg( + uid=f"{preset.object_type}_{preset.material_name}_mat", + base_color=list(preset.base_color), + roughness=preset.roughness, + ), + ), + body_type="static", + attrs=RigidBodyAttributesCfg( + dynamic_friction=preset.dynamic_friction, + static_friction=preset.static_friction, + ), + init_pos=init_pos, + ) + return sim.add_rigid_object(cfg=cfg) + + +def _reset_robot(robot: Robot, initial_qpos: torch.Tensor) -> None: + """Reset current and target robot qpos to the benchmark initial posture.""" + for target in (False, True): + robot.set_qpos(initial_qpos, target=target) + robot.clear_dynamics() + + +def _build_atomic_engine( + motion_gen: MotionGenerator, + robot: Robot, + device: torch.device, +) -> AtomicActionEngine: + """Build a Press benchmark engine with MoveEndEffector pre-positioning.""" + hand_close = get_hand_close_qpos(robot, device) + atomic_engine = AtomicActionEngine(motion_generator=motion_gen) + atomic_engine.register( + MoveEndEffector( + motion_gen, + cfg=MoveEndEffectorCfg( + control_part="arm", + sample_interval=MOVE_SAMPLE_INTERVAL, + ), + ) + ) + atomic_engine.register( + Press( + motion_gen, + cfg=PressCfg( + control_part="arm", + hand_control_part="hand", + hand_close_qpos=hand_close, + sample_interval=PRESS_SAMPLE_INTERVAL, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ) + ) + return atomic_engine + + +def _make_press_targets( + obj: RigidObject, + preset: ObjectPreset, +) -> tuple[torch.Tensor, torch.Tensor]: + """Create pre-press and press poses for the object's top center.""" + obj_pose = obj.get_local_pose(to_matrix=True) + object_center = obj_pose[0, :3, 3].clone() + object_top_z = object_center[2] + 0.5 * preset.size[2] + + press_position = object_center.clone() + press_position[2] = object_top_z + PRESS_SURFACE_OFFSET + move_position = press_position.clone() + move_position[2] = object_top_z + PRESS_CLEARANCE + + return make_top_down_eef_pose(move_position), make_top_down_eef_pose(press_position) + + +def _compute_press_center_check( + robot: Robot, + traj: torch.Tensor, + obj: RigidObject, + object_height: float, + tolerance: float, +) -> tuple[bool, float, int]: + """Check whether the planned Press trajectory reaches the object top center.""" + if traj.numel() == 0: + return False, float("inf"), -1 + + arm_joint_ids = robot.get_joint_ids(name="arm") + n_down = (PRESS_SAMPLE_INTERVAL - HAND_INTERP_STEPS) // 2 + press_segment_start = MOVE_SAMPLE_INTERVAL + HAND_INTERP_STEPS + press_segment_end = min(press_segment_start + n_down, traj.shape[1]) + arm_traj = traj[:, press_segment_start:press_segment_end, arm_joint_ids] + if arm_traj.shape[1] == 0: + return False, float("inf"), -1 + + fk_pose = torch.stack( + [ + robot.compute_fk( + qpos=waypoint.unsqueeze(0), + name="arm", + to_matrix=True, + )[0] + for waypoint in arm_traj[0] + ], + dim=0, + ) + + obj_pose = obj.get_local_pose(to_matrix=True) + object_center = obj_pose[0, :3, 3] + object_top_z = object_center[2] + 0.5 * object_height + target_xy = object_center[:2] + target_z = object_top_z + PRESS_SURFACE_OFFSET + + xy_error = torch.linalg.norm(fk_pose[:, :2, 3] - target_xy, dim=1) + z_error = torch.abs(fk_pose[:, 2, 3] - target_z) + combined_error = xy_error + z_error + best_idx = int(torch.argmin(combined_error).item()) + best_pos = fk_pose[best_idx, :3, 3] + center_error = float(torch.linalg.norm(best_pos[:2] - target_xy).item()) + return center_error <= tolerance, center_error, press_segment_start + best_idx + + +def _timed_atomic_run( + atomic_engine: AtomicActionEngine, + move_target: torch.Tensor, + press_target: torch.Tensor, +) -> tuple[float, dict[str, float], float, bool, torch.Tensor]: + """Run a timed atomic-action sequence and return timing/memory/results.""" + _reset_peak_gpu_memory() + mem_before = _memory_snapshot() + _sync_cuda() + + start = time.perf_counter() + is_success, traj, _ = atomic_engine.run( + steps=[ + ("move_end_effector", EndEffectorPoseTarget(xpos=move_target)), + ("press", EndEffectorPoseTarget(xpos=press_target)), + ] + ) + _sync_cuda() + elapsed = time.perf_counter() - start + + mem_after = _memory_snapshot() + deltas = { + "cpu_mb": mem_after["cpu_mb"] - mem_before["cpu_mb"], + "gpu_mb": mem_after["gpu_mb"] - mem_before["gpu_mb"], + } + return elapsed, deltas, _peak_gpu_memory_mb(), is_success, traj + + +def _run_press_case( + sim: SimulationManager, + robot: Robot, + atomic_engine: AtomicActionEngine, + initial_qpos: torch.Tensor, + obj: RigidObject, + base_obj_pose: torch.Tensor, + preset: ObjectPreset, + position_case: PositionCase, + repeat_index: int, + press_tolerance: float, + args: argparse.Namespace, + recorded_count: int, +) -> PressCaseResult: + """Run one object-position Press benchmark case.""" + case_id = f"{preset.object_type}:{position_case.name}:r{repeat_index}" + try: + _reset_robot(robot, initial_qpos) + initial_obj_pose = reset_rigid_object_xy( + obj=obj, + base_pose=base_obj_pose, + xy=position_case.xy, + sim=sim, + settle_steps=2, + ) + move_target, press_target = _make_press_targets(obj, preset) + + elapsed, mem_delta, peak_gpu, planning_success, traj = _timed_atomic_run( + atomic_engine=atomic_engine, + move_target=move_target, + press_target=press_target, + ) + video_path = None + if should_record_case(args, recorded_count, bool(planning_success)): + _reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + video_path = replay_trajectory_with_recording( + sim=sim, + robot=robot, + traj=traj, + args=args, + video_path=build_video_output_path( + args, + "atomic_action_press", + (f"{preset.object_type}_{position_case.name}" f"_r{repeat_index}"), + ), + ) + _reset_robot(robot, initial_qpos) + reset_rigid_object(obj, initial_obj_pose) + + center_hit = False + xy_error_m: float | None = None + hit_step: int | None = None + failure_reason = "" + if planning_success: + center_hit, xy_error_m, raw_hit_step = _compute_press_center_check( + robot=robot, + traj=traj, + obj=obj, + object_height=preset.size[2], + tolerance=press_tolerance, + ) + hit_step = raw_hit_step if raw_hit_step >= 0 else None + if not center_hit: + failure_reason = "center_miss" + else: + failure_reason = "planning_failed" + + return PressCaseResult( + case_id=case_id, + object_type=preset.object_type, + material_name=preset.material_name, + quadrant=position_case.quadrant, + position_case=position_case.name, + init_xy=position_case.xy, + repeat_index=repeat_index, + planning_success=planning_success, + center_hit=center_hit, + cost_time_ms=elapsed * 1000.0, + cpu_delta_mb=mem_delta["cpu_mb"], + gpu_delta_mb=mem_delta["gpu_mb"], + peak_gpu_mb=peak_gpu, + xy_error_m=xy_error_m, + hit_step=hit_step, + trajectory_waypoints=int(traj.shape[1]) if traj.ndim >= 2 else 0, + failure_reason=failure_reason, + video_path=str(video_path) if video_path is not None else "", + ) + except Exception as exc: + return PressCaseResult( + case_id=case_id, + object_type=preset.object_type, + material_name=preset.material_name, + quadrant=position_case.quadrant, + position_case=position_case.name, + init_xy=position_case.xy, + repeat_index=repeat_index, + planning_success=False, + center_hit=False, + cost_time_ms=0.0, + cpu_delta_mb=0.0, + gpu_delta_mb=0.0, + peak_gpu_mb=0.0, + xy_error_m=None, + hit_step=None, + trajectory_waypoints=0, + failure_reason=f"exception:{type(exc).__name__}:{exc}", + ) + + +def _build_perf_rows(results: list[PressCaseResult]) -> list[dict[str, object]]: + """Build Time & Memory table rows.""" + rows: list[dict[str, object]] = [] + for result in results: + rows.append( + { + "sample_size": 1, + "impl": "press", + "case_id": result.case_id, + "object_type": result.object_type, + "material": result.material_name, + "quadrant": result.quadrant, + "position_case": result.position_case, + "init_xy": f"({result.init_xy[0]:.3f},{result.init_xy[1]:.3f})", + "repeat": result.repeat_index, + "cost_time_ms": _format_float(result.cost_time_ms), + "cpu_delta_mb": _format_float(result.cpu_delta_mb), + "gpu_delta_mb": _format_float(result.gpu_delta_mb), + "peak_gpu_mb": _format_float(result.peak_gpu_mb), + } + ) + return rows + + +def _build_metric_rows(results: list[PressCaseResult]) -> list[dict[str, object]]: + """Build Success & Other Metrics table rows.""" + rows: list[dict[str, object]] = [] + for result in results: + overall_success = result.planning_success and result.center_hit + rows.append( + { + "sample_size": 1, + "impl": "press", + "case_id": result.case_id, + "object_type": result.object_type, + "material": result.material_name, + "quadrant": result.quadrant, + "position_case": result.position_case, + "success_rate": f"{float(overall_success):.6f}", + "planning_success_rate": f"{float(result.planning_success):.6f}", + "center_hit_rate": f"{float(result.center_hit):.6f}", + "xy_error_m": _format_float(result.xy_error_m), + "hit_step": result.hit_step if result.hit_step is not None else "N/A", + "trajectory_waypoints": result.trajectory_waypoints, + "failure_reason": result.failure_reason or "N/A", + } + ) + return rows + + +def _build_leaderboard_rows(results: list[PressCaseResult]) -> list[dict[str, object]]: + """Aggregate and rank object-conditioned Press variants by success rate.""" + aggregate: dict[str, dict[str, float | set[str]]] = {} + for result in results: + algorithm = f"press:{result.object_type}" + if algorithm not in aggregate: + aggregate[algorithm] = { + "overall_success_sum": 0.0, + "planning_success_sum": 0.0, + "xy_error_sum": 0.0, + "xy_error_count": 0.0, + "cost_time_sum": 0.0, + "case_count": 0.0, + "quadrants": set(), + } + + stats = aggregate[algorithm] + stats["overall_success_sum"] = float(stats["overall_success_sum"]) + float( + result.planning_success and result.center_hit + ) + stats["planning_success_sum"] = float(stats["planning_success_sum"]) + float( + result.planning_success + ) + if result.xy_error_m is not None and math.isfinite(result.xy_error_m): + stats["xy_error_sum"] = float(stats["xy_error_sum"]) + result.xy_error_m + stats["xy_error_count"] = float(stats["xy_error_count"]) + 1.0 + stats["cost_time_sum"] = float(stats["cost_time_sum"]) + result.cost_time_ms + stats["case_count"] = float(stats["case_count"]) + 1.0 + quadrants = stats["quadrants"] + if isinstance(quadrants, set): + quadrants.add(result.quadrant) + + ranked = sorted( + aggregate.items(), + key=lambda item: ( + float(item[1]["overall_success_sum"]) + / max(float(item[1]["case_count"]), 1.0), + -float(item[1]["cost_time_sum"]) / max(float(item[1]["case_count"]), 1.0), + ), + reverse=True, + ) + + rows: list[dict[str, object]] = [] + for rank, (algorithm, stats) in enumerate(ranked, start=1): + case_count = max(float(stats["case_count"]), 1.0) + xy_error_count = float(stats["xy_error_count"]) + avg_xy_error = ( + float(stats["xy_error_sum"]) / xy_error_count + if xy_error_count > 0.0 + else None + ) + quadrants = stats["quadrants"] + quadrant_coverage = ( + ",".join(sorted(quadrants)) if isinstance(quadrants, set) else "" + ) + rows.append( + { + "rank": rank, + "algorithm": algorithm, + "overall_success_rate": ( + f"{float(stats['overall_success_sum']) / case_count:.2%}" + ), + "planning_success_rate": ( + f"{float(stats['planning_success_sum']) / case_count:.2%}" + ), + "avg_xy_error_m": _format_float(avg_xy_error), + "avg_cost_time_ms": _format_float( + float(stats["cost_time_sum"]) / case_count + ), + "evaluated_cases": int(case_count), + "quadrant_coverage": quadrant_coverage, + } + ) + return rows + + +def _print_case_result(result: PressCaseResult) -> None: + """Print one aligned case result line.""" + overall_success = result.planning_success and result.center_hit + print( + f" {result.case_id:<28} " + f"time={result.cost_time_ms:>10.2f} ms | " + f"CPU delta={result.cpu_delta_mb:+.1f} MB " + f"GPU delta={result.gpu_delta_mb:+.1f} MB " + f"peak GPU={result.peak_gpu_mb:.1f} MB | " + f"success={overall_success} " + f"xy_error={_format_float(result.xy_error_m, precision=4)}" + ) + if result.failure_reason: + print(f" reason={result.failure_reason}") + + +def _build_notes( + object_presets: list[ObjectPreset], + position_cases: list[PositionCase], + repeat: int, + video_paths: list[str], + profile: str, +) -> list[str]: + """Build report notes with benchmark coverage metadata.""" + quadrant_counts: dict[str, int] = {} + for position_case in position_cases: + quadrant_counts[position_case.quadrant] = ( + quadrant_counts.get(position_case.quadrant, 0) + 1 + ) + return [ + f"Profile: {profile}", + "Object presets: " + + ", ".join( + f"{preset.object_type}/{preset.material_name}/size={preset.size}" + for preset in object_presets + ), + "Position cases per quadrant: " + + ", ".join( + f"{quadrant}={count}" for quadrant, count in sorted(quadrant_counts.items()) + ), + f"CPU memory backend: {CPU_MEMORY_BACKEND}", + f"Repeat per object-position case: {repeat}", + "Replay videos: " + (", ".join(video_paths) if video_paths else "disabled"), + "success_rate is 1 only when planning succeeds and the Press trajectory " + "reaches the object top center.", + ] + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: + """Run all atomic-action benchmarks and write the markdown report.""" + args = _parse_args() if args is None else args + if args.repeat < 1: + raise ValueError("--repeat must be at least 1.") + profile = resolve_profile(args) + _ensure_runtime_imports() + + object_presets = _select_object_presets(args.object_types, profile) + position_cases = _select_position_cases(args.position_cases, profile) + repeat = 1 if profile == "smoke" else args.repeat + + print("=" * 60) + print("Atomic Action Press Benchmark") + print("=" * 60) + print( + "Coverage: " + f"profile={profile}, {len(object_presets)} object presets x " + f"{len(position_cases)} position cases x {repeat} repeat(s)" + ) + + sim = initialize_simulation(args) + robot = create_robot(sim) + create_table(sim) + initial_qpos = robot.get_qpos().clone() + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + atomic_engine = _build_atomic_engine(motion_gen, robot, sim.device) + object_pool = {} + for object_index, preset in enumerate(object_presets): + obj = _create_benchmark_object(sim, preset, position_cases[0], object_index) + settle_object(sim, obj, step=2) + base_pose = obj.get_local_pose(to_matrix=True).clone() + park_rigid_object(obj, base_pose, index=object_index, sim=sim) + object_pool[preset.object_type] = (obj, base_pose) + + results: list[PressCaseResult] = [] + video_paths: list[str] = [] + print("\n=== Press Object/Position Sweep ===") + for preset in object_presets: + obj, base_pose = object_pool[preset.object_type] + for parked_index, parked_preset in enumerate(object_presets): + if parked_preset.object_type == preset.object_type: + continue + parked_obj, parked_base_pose = object_pool[parked_preset.object_type] + park_rigid_object(parked_obj, parked_base_pose, index=parked_index, sim=sim) + for position_case in position_cases: + for repeat_index in range(repeat): + result = _run_press_case( + sim=sim, + robot=robot, + atomic_engine=atomic_engine, + initial_qpos=initial_qpos, + obj=obj, + base_obj_pose=base_pose, + preset=preset, + position_case=position_case, + repeat_index=repeat_index, + press_tolerance=args.press_tolerance, + args=args, + recorded_count=len(video_paths), + ) + results.append(result) + if result.video_path: + video_paths.append(result.video_path) + _print_case_result(result) + + perf_rows = _build_perf_rows(results) + metric_rows = _build_metric_rows(results) + leaderboard_rows = _build_leaderboard_rows(results) + report_path = _write_markdown_report( + benchmark_name="atomic_action_press", + perf_rows=perf_rows, + metric_rows=metric_rows, + leaderboard_rows=leaderboard_rows, + notes=_build_notes( + object_presets, + position_cases, + repeat, + video_paths, + profile, + ), + ) + + print("\n" + "=" * 60) + print("Benchmarks complete.") + print(f"Markdown report saved: {report_path}") + print("=" * 60) + return report_path + + +def main() -> None: + """Run the CLI entry point.""" + try: + run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/run_benchmark.py b/scripts/benchmark/atomic_action/run_benchmark.py new file mode 100644 index 000000000..26c9749b3 --- /dev/null +++ b/scripts/benchmark/atomic_action/run_benchmark.py @@ -0,0 +1,340 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Dispatch benchmarks for all atomic actions. + +Run a single action benchmark or all action benchmarks in sequence. +Run: python -m scripts.benchmark.atomic_action.run_benchmark --action press +""" + +from __future__ import annotations + +import argparse +import shlex +import subprocess +import sys +from pathlib import Path + +from scripts.benchmark.atomic_action.common import ( + MESH_OBJECT_PRESETS, + PICKUP_APPROACH_CASES, + POSITION_CASES, + add_profile_benchmark_args, + add_video_benchmark_args, + resolve_profile, +) + +ACTION_MODULES = { + "move_end_effector": "scripts.benchmark.atomic_action.move_end_effector_benchmark", + "move_joints": "scripts.benchmark.atomic_action.move_joints_benchmark", + "pick_up": "scripts.benchmark.atomic_action.pickup_benchmark", + "move_held_object": "scripts.benchmark.atomic_action.move_held_object_benchmark", + "place": "scripts.benchmark.atomic_action.place_benchmark", + "press": "scripts.benchmark.atomic_action.press_benchmark", +} +DEFAULT_ACTIONS = tuple(ACTION_MODULES.keys()) +MESH_OBJECT_ACTIONS = {"pick_up", "move_held_object", "place"} +PRESS_OBJECT_TYPES = {"bottle", "mug", "wooden_block", "all"} +MESH_OBJECT_TYPES = {*MESH_OBJECT_PRESETS.keys(), "all"} + + +def add_benchmark_args(parser: argparse.ArgumentParser) -> None: + """Add atomic-action aggregate benchmark CLI arguments.""" + parser.add_argument( + "--action", + nargs="+", + choices=(*ACTION_MODULES.keys(), "all"), + default=["press"], + help="Atomic action benchmark(s) to run. Use 'all' for every action.", + ) + parser.add_argument( + "--smoke", + action="store_true", + help="Alias for --profile smoke.", + ) + add_profile_benchmark_args(parser) + parser.add_argument( + "--object_types", + nargs="+", + default=None, + help=( + "Optional object presets forwarded to selected object-conditioned " + "benchmarks. Values must be valid for each selected action." + ), + ) + parser.add_argument( + "--position_cases", + nargs="+", + choices=(*POSITION_CASES.keys(), "all"), + default=None, + help=( + "Optional initial position cases forwarded to selected " + "object-conditioned benchmarks." + ), + ) + parser.add_argument( + "--approach_cases", + nargs="+", + choices=(*PICKUP_APPROACH_CASES, "all"), + default=None, + help=( + "Optional PickUp approach cases forwarded to pick/place/held-object " + "benchmarks." + ), + ) + parser.add_argument( + "--device", + type=str, + default="cpu", + help="Simulation device forwarded to each selected benchmark.", + ) + parser.add_argument( + "--renderer", + type=str, + choices=("auto", "hybrid", "fast-rt", "rt"), + default="auto", + help="Renderer backend forwarded to each selected benchmark.", + ) + add_video_benchmark_args(parser) + parser.add_argument( + "--in_process", + action="store_true", + help=( + "Run selected benchmarks in the current Python process. This is useful " + "for debugging but can leave DexSim resources alive across actions." + ), + ) + + +def _parse_args() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Run one or more atomic-action benchmarks." + ) + add_benchmark_args(parser) + return parser.parse_args() + + +def _selected_actions(actions: list[str]) -> list[str]: + """Resolve selected action names.""" + if "all" in actions: + return list(DEFAULT_ACTIONS) + return actions + + +def _validate_object_types_for_actions( + selected_actions: list[str], + object_types: list[str] | None, +) -> None: + """Validate optional object filters against selected action namespaces.""" + if not object_types: + return + + requested = set(object_types) + validators: dict[str, set[str]] = {} + for action_name in selected_actions: + if action_name in MESH_OBJECT_ACTIONS: + validators[action_name] = MESH_OBJECT_TYPES + elif action_name == "press": + validators[action_name] = PRESS_OBJECT_TYPES + + invalid_parts = [] + for action_name, valid_types in validators.items(): + invalid = sorted(requested - valid_types) + if invalid: + valid = ", ".join(sorted(valid_types)) + invalid_parts.append( + f"{action_name}: invalid {invalid}; valid values are {valid}" + ) + if invalid_parts: + raise RuntimeError( + "--object_types must be valid for every selected object-conditioned " + "benchmark. " + " | ".join(invalid_parts) + ) + + +def _make_child_args(args: argparse.Namespace) -> argparse.Namespace: + """Build minimal child benchmark arguments for aggregate dispatch.""" + profile = resolve_profile(args) + return argparse.Namespace( + smoke=profile == "smoke", + profile=profile, + repeat=1, + device=args.device, + renderer=args.renderer, + object_types=args.object_types, + position_cases=args.position_cases, + press_tolerance=0.01, + pose_cases=["all"], + sequence_cases=["all"], + approach_cases=args.approach_cases, + place_cases=None, + held_object_cases=None, + n_sample=1000 if profile == "smoke" else 10000, + force_reannotate=False, + record_video=args.record_video, + record_failed_video=args.record_failed_video, + video_case_limit=args.video_case_limit, + video_dir=args.video_dir, + video_fps=args.video_fps, + video_max_memory=args.video_max_memory, + video_width=args.video_width, + video_height=args.video_height, + video_hold_steps=args.video_hold_steps, + ) + + +def _make_child_cli_args(args: argparse.Namespace, action_name: str) -> list[str]: + """Build CLI arguments forwarded to one isolated benchmark subprocess.""" + profile = resolve_profile(args) + child_args = [ + "--profile", + profile, + "--repeat", + "1", + "--device", + args.device, + "--renderer", + args.renderer, + "--video_case_limit", + str(args.video_case_limit), + "--video_dir", + str(args.video_dir), + "--video_fps", + str(args.video_fps), + "--video_max_memory", + str(args.video_max_memory), + "--video_width", + str(args.video_width), + "--video_height", + str(args.video_height), + "--video_hold_steps", + str(args.video_hold_steps), + ] + if action_name in {"pick_up", "move_held_object", "place", "press"}: + if args.object_types: + child_args.append("--object_types") + child_args.extend(args.object_types) + if args.position_cases: + child_args.append("--position_cases") + child_args.extend(args.position_cases) + if action_name in {"pick_up", "move_held_object", "place"}: + if args.approach_cases: + child_args.append("--approach_cases") + child_args.extend(args.approach_cases) + if args.record_video: + child_args.append("--record_video") + if args.record_failed_video: + child_args.append("--record_failed_video") + return child_args + + +def _run_action_subprocess(action_name: str, args: argparse.Namespace) -> Path: + """Run one action benchmark in an isolated Python subprocess.""" + module_name = ACTION_MODULES[action_name] + command = [ + sys.executable, + "-m", + module_name, + *_make_child_cli_args(args, action_name), + ] + print("\n" + "=" * 60, flush=True) + print( + f"Running {action_name} benchmark in an isolated subprocess", + flush=True, + ) + print( + "Command: " + " ".join(shlex.quote(part) for part in command), + flush=True, + ) + print("=" * 60, flush=True) + + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + encoding="utf-8", + errors="replace", + bufsize=1, + ) + report_path: Path | None = None + if process.stdout is not None: + for line in process.stdout: + print(line, end="", flush=True) + if line.startswith("Markdown report saved:"): + report_path = Path(line.split(":", maxsplit=1)[1].strip()) + + return_code = process.wait() + if return_code != 0: + raise RuntimeError( + f"{action_name} benchmark subprocess failed with exit code " + f"{return_code}." + ) + if report_path is None: + raise RuntimeError( + f"{action_name} benchmark subprocess finished without reporting " + "a markdown report path." + ) + return report_path + + +def _run_in_process_benchmarks( + args: argparse.Namespace, + selected_actions: list[str], +) -> list[Path]: + """Run selected benchmarks in the current Python process.""" + reports: list[Path] = [] + for action_name in selected_actions: + module_name = ACTION_MODULES[action_name] + module = __import__(module_name, fromlist=["run_all_benchmarks"]) + child_args = _make_child_args(args) + report_path = module.run_all_benchmarks(child_args) + reports.append(report_path) + return reports + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> list[Path]: + """Run selected atomic-action benchmarks.""" + args = _parse_args() if args is None else args + selected_actions = _selected_actions(args.action) + _validate_object_types_for_actions(selected_actions, args.object_types) + + if args.in_process: + return _run_in_process_benchmarks(args, selected_actions) + + return [ + _run_action_subprocess(action_name, args) for action_name in selected_actions + ] + + +def main() -> None: + """Run the CLI entry point.""" + try: + reports = run_all_benchmarks() + except RuntimeError as exc: + raise SystemExit(str(exc)) from exc + if reports: + print("Generated reports:") + for report in reports: + print(f" {report}") + + +if __name__ == "__main__": + main() + + +__all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/tutorials/atomic_action/coordinated_placement.py b/scripts/tutorials/atomic_action/coordinated_placement.py new file mode 100644 index 000000000..5e9275e38 --- /dev/null +++ b/scripts/tutorials/atomic_action/coordinated_placement.py @@ -0,0 +1,1202 @@ +# ---------------------------------------------------------------------------- +# 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 dual-arm coordinated placement with bread and pan meshes. + +The left UR5 picks up bread. The right UR5 picks up a pan and moves it to the +lower alignment pose. The left UR5 places the bread above the pan and releases +it while the right hand keeps holding the pan. +""" + +from __future__ import annotations + +import argparse +import math +import os +import sys +import time +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 numpy as np +import torch +from scipy.spatial.transform import Rotation as SciRotation + +from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.atomic_actions import ( + Affordance, + CoordinatedPlacement, + CoordinatedPlacementCfg, + CoordinatedPlacementTarget, + HeldObjectState, + ObjectSemantics, + TrajectoryBuilder, + WorldState, +) +from embodichain.lab.sim.cfg import ( + JointDrivePropertiesCfg, + LightCfg, + RenderCfg, + RigidBodyAttributesCfg, + RigidObjectCfg, + RobotCfg, + URDFCfg, +) +from embodichain.lab.sim.objects import RigidObject, Robot +from embodichain.lab.sim.planners import ( + MotionGenCfg, + MotionGenerator, + MoveType, + PlanState, + ToppraPlannerCfg, +) +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.solvers import URSolverCfg +from embodichain.utils import logger +from scripts.tutorials.atomic_action.tutorial_utils import ( + draw_axis_marker, + get_tutorial_window_size, + make_ur5_solver_cfg, + start_auto_play_recording, + stop_auto_play_recording, +) + +DEFAULT_MESH_FRAME_CORRECTION_EULER_DEG = (-90.0, 0.0, 0.0) +# DexSim imports this pan GLB with gym raw +Z mapped to local -Y. The +90deg X +# correction therefore makes the pan opening point upward in world Z. +PAN_MESH_FRAME_CORRECTION_EULER_DEG = (90.0, 0.0, 0.0) +PAN_WORLD_YAW_CORRECTION_DEG = 270.0 + + +def transform_baseline_pose( + init_pos: tuple[float, float, float], + init_rot: tuple[float, float, float], + *, + z_offset: float = 0.0, + mesh_frame_correction_euler_deg: tuple[float, float, float] = ( + DEFAULT_MESH_FRAME_CORRECTION_EULER_DEG + ), + world_yaw_correction_deg: float = 0.0, +) -> tuple[tuple[float, float, float], tuple[float, float, float]]: + """Apply mesh-frame correction while preserving baseline world placement.""" + pos = np.asarray(init_pos, dtype=np.float64) + pos[2] += z_offset + rot = ( + SciRotation.from_euler("Z", world_yaw_correction_deg, degrees=True) + * SciRotation.from_euler("XYZ", init_rot, degrees=True) + * SciRotation.from_euler("XYZ", mesh_frame_correction_euler_deg, degrees=True) + ).as_euler("XYZ", degrees=True) + return tuple(float(value) for value in pos), tuple(float(value) for value in rot) + + +ARM_URDF_PATH = "UniversalRobots/UR5/UR5.urdf" +GRIPPER_URDF_PATH = "DH_PGI_140_80/DH_PGI_140_80.urdf" +SCRIPT_DIR = Path(__file__).resolve().parent +TABLE_MESH_PATH = SCRIPT_DIR / "table.glb" +BREAD_MESH_PATH = SCRIPT_DIR / "bread.glb" +PAN_MESH_PATH = SCRIPT_DIR / "pan.glb" +BREAD_LABEL = "bread" +PAN_LABEL = "pan" +GRIPPER_TCP_Z = 0.121 +PICK_SAMPLE_INTERVAL = 100 +COORDINATED_SAMPLE_INTERVAL = 120 +ROBOT_INIT_POS = (1.85, 0.0, 0.1) +ROBOT_INIT_ROT = (0.0, 0.0, -90.0) +LEFT_ARM_HOME = (0.0, 0.0, -1.57, -1.57, 1.57, 1.57) +RIGHT_ARM_HOME = (-1.57, -1.57, -1.57, -1.57, 0.0, 0.0) +TABLE_TOP_Z = 0.65 +BASELINE_TABLE_TOP_Z = 0.3621708124799265 +SCENE_Z_OFFSET = TABLE_TOP_Z - BASELINE_TABLE_TOP_Z +BASELINE_TABLE_INIT_POS = ( + 0.00014585733079742588, + 0.00023304896730074557, + -0.019599792839044783, +) +BASELINE_TABLE_INIT_ROT = ( + 0.0001074673904926984, + 0.00865572768366991, + -90.6562109309317, +) +BASELINE_BREAD_INIT_POS = ( + 0.007266042530919159, + 0.17218712515099063, + 0.38805152145807564, +) +BASELINE_BREAD_INIT_ROT = ( + 179.93952112929065, + -0.12776179446053365, + 85.59207565132371, +) +BASELINE_PAN_INIT_POS = ( + 0.0009683294205463406, + -0.14189524793277888, + 0.38900474548025743, +) +BASELINE_PAN_INIT_ROT = ( + -179.23950670370294, + -0.4795764805552328, + 98.19364391929443, +) +TABLE_INIT_POS, TABLE_INIT_ROT = transform_baseline_pose( + BASELINE_TABLE_INIT_POS, + BASELINE_TABLE_INIT_ROT, + z_offset=SCENE_Z_OFFSET, +) +BREAD_INIT_POS, BREAD_INIT_ROT = transform_baseline_pose( + BASELINE_BREAD_INIT_POS, + BASELINE_BREAD_INIT_ROT, + z_offset=SCENE_Z_OFFSET, +) +PAN_INIT_POS, PAN_INIT_ROT = transform_baseline_pose( + BASELINE_PAN_INIT_POS, + BASELINE_PAN_INIT_ROT, + z_offset=SCENE_Z_OFFSET, + mesh_frame_correction_euler_deg=PAN_MESH_FRAME_CORRECTION_EULER_DEG, + world_yaw_correction_deg=PAN_WORLD_YAW_CORRECTION_DEG, +) +PAN_INIT_POS = (PAN_INIT_POS[0], PAN_INIT_POS[1], TABLE_TOP_Z + 0.001) +PAN_TARGET_CENTER_XY = (-0.06, 0.0) +PAN_TARGET_Z_LIFT = 0.06 +BREAD_PLACE_TARGET_OFFSET_XY = (-0.06, -0.16) +BREAD_ON_PAN_CLEARANCE = 0.006 +BREAD_GRASP_Z_CLEARANCE = 0.018 +PAN_GRASP_Z_CLEARANCE = 0.0 +PAN_HANDLE_LOCAL_Z_MIN = 0.04 +PAN_BASIN_LOCAL_Z_MAX = 0.04 +PAN_HANDLE_ROOT_OFFSET = 0.035 +PAN_HANDLE_CLOSE_QPOS = 0.045 +PAN_PICK_SAMPLE_INTERVAL = 130 +PAN_PICK_HAND_INTERP_STEPS = 32 +BREAD_TARGET_WORLD_YAW_DEG = 0.0 +BREAD_TARGET_HEIGHT_OFFSET = 0.1 +SUPPORT_TARGET_HEIGHT_OFFSET = 0.0 +PICK_APPROACH_DISTANCE = 0.12 +PLACE_LIFT_HEIGHT = 0.10 +TRAJECTORY_SIM_STEPS = 8 + + +def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the demo.""" + parser = argparse.ArgumentParser(description="Dual-arm coordinated placement demo") + add_env_launcher_args_to_parser(parser) + parser.set_defaults(device="cuda", renderer="hybrid") + parser.add_argument( + "--diagnose_plan", + action="store_true", + help="Plan and print diagnostics without playing the trajectory.", + ) + parser.add_argument( + "--debug_state", + action="store_true", + help="Log hand targets and object poses during execution.", + ) + parser.add_argument( + "--auto_play", + action="store_true", + help="Run the viewer demo without waiting for keyboard input.", + ) + parser.add_argument( + "--no_vis_eef_axis", + action="store_true", + help="Do not draw coordinated placement target coordinate frames.", + ) + parser.add_argument( + "--headless_play", + action="store_true", + help="Execute planned trajectories without opening the viewer window.", + ) + return parser.parse_args() + + +def get_cached_data_path(data_path: str) -> str: + """Resolve an asset path from the local cache before importing data helpers.""" + if os.path.isabs(data_path): + return data_path + + data_root = Path( + os.environ.get( + "EMBODICHAIN_DATA_ROOT", + str(Path.home() / ".cache" / "embodichain_data"), + ) + ) + candidates = ( + data_root / data_path, + data_root / "extract" / data_path, + ) + for candidate in candidates: + if candidate.exists(): + return str(candidate) + + from embodichain.data import get_data_path + + return get_data_path(data_path) + + +def rotation_z(yaw: float) -> np.ndarray: + """Build a 3x3 yaw rotation matrix.""" + cos_yaw = math.cos(yaw) + sin_yaw = math.sin(yaw) + return np.array( + [ + [cos_yaw, -sin_yaw, 0.0], + [sin_yaw, cos_yaw, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + + +def make_transform(xyz: tuple[float, float, float], yaw: float) -> np.ndarray: + """Build a homogeneous transform from translation and yaw.""" + transform = np.eye(4, dtype=np.float32) + transform[:3, :3] = rotation_z(yaw) + transform[:3, 3] = np.asarray(xyz, dtype=np.float32) + return transform + + +def initialize_simulation(args: argparse.Namespace) -> SimulationManager: + """Create the simulation manager and a light.""" + width, height = get_tutorial_window_size(args) + sim = SimulationManager( + SimulationManagerCfg( + width=width, + height=height, + headless=True, + sim_device=args.device, + render_cfg=RenderCfg(renderer=args.renderer), + physics_dt=1.0 / 100.0, + arena_space=3.0, + ) + ) + sim.add_light( + cfg=LightCfg( + uid="main_light", + color=(0.6, 0.6, 0.6), + intensity=30.0, + init_pos=(0.0, -0.4, 3.0), + ) + ) + return sim + + +def make_prefixed_ur5_solver_cfg(prefix: str) -> URSolverCfg: + """Create a UR5 solver cfg for a prefixed arm in the assembled robot.""" + cfg = make_ur5_solver_cfg(GRIPPER_TCP_Z) + cfg.root_link_name = f"{prefix}_base_link" + cfg.end_link_name = f"{prefix}_ee_link" + return cfg + + +def create_dual_ur5_robot(sim: SimulationManager) -> Robot: + """Create a dual-UR5 robot with one PGI gripper on each arm.""" + arm_urdf_path = get_cached_data_path(ARM_URDF_PATH) + gripper_urdf_path = get_cached_data_path(GRIPPER_URDF_PATH) + cfg = RobotCfg( + uid="DualUR5CoordinatedPlacement", + urdf_cfg=URDFCfg( + components=[ + { + "component_type": "left_arm", + "urdf_path": arm_urdf_path, + "transform": make_transform((-0.3, -1.45, 0.4), np.pi / 2), + }, + { + "component_type": "right_arm", + "urdf_path": arm_urdf_path, + "transform": make_transform((0.3, -1.45, 0.4), np.pi / 2), + }, + {"component_type": "left_hand", "urdf_path": gripper_urdf_path}, + {"component_type": "right_hand", "urdf_path": gripper_urdf_path}, + ], + fname="dual_ur5_coordinated_placement", + ), + drive_pros=JointDrivePropertiesCfg( + stiffness={ + "left_joint[0-9]": 1e4, + "right_joint[0-9]": 1e4, + "left_gripper_finger[1-2]_joint_1": 1e3, + "right_gripper_finger[1-2]_joint_1": 1e3, + }, + damping={ + "left_joint[0-9]": 1e3, + "right_joint[0-9]": 1e3, + "left_gripper_finger[1-2]_joint_1": 1e2, + "right_gripper_finger[1-2]_joint_1": 1e2, + }, + max_effort={ + "left_joint[0-9]": 1e5, + "right_joint[0-9]": 1e5, + "left_gripper_finger[1-2]_joint_1": 1e4, + "right_gripper_finger[1-2]_joint_1": 1e4, + }, + drive_type="force", + ), + control_parts={ + "left_arm": ["left_joint[0-9]"], + "right_arm": ["right_joint[0-9]"], + "dual_arm": ["left_joint[0-9]", "right_joint[0-9]"], + "left_hand": ["left_gripper_finger1_joint_1"], + "right_hand": ["right_gripper_finger1_joint_1"], + }, + solver_cfg={ + "left_arm": make_prefixed_ur5_solver_cfg("left"), + "right_arm": make_prefixed_ur5_solver_cfg("right"), + }, + init_pos=list(ROBOT_INIT_POS), + init_rot=list(ROBOT_INIT_ROT), + init_qpos=list(LEFT_ARM_HOME) + list(RIGHT_ARM_HOME) + [0.0, 0.0, 0.0, 0.0], + ) + return sim.add_robot(cfg=cfg) + + +def create_table(sim: SimulationManager) -> RigidObject: + """Create the table mesh from the bread-pan gym project.""" + return sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="table", + shape=MeshCfg(fpath=str(TABLE_MESH_PATH)), + attrs=RigidBodyAttributesCfg( + mass=10.0, + dynamic_friction=0.9, + static_friction=0.95, + restitution=0.01, + ), + body_type="kinematic", + init_pos=list(TABLE_INIT_POS), + init_rot=list(TABLE_INIT_ROT), + ) + ) + + +def create_bread(sim: SimulationManager) -> RigidObject: + """Create the bread mesh to be placed by the left arm.""" + return sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="bread", + shape=MeshCfg(fpath=str(BREAD_MESH_PATH), compute_uv=False), + attrs=RigidBodyAttributesCfg( + mass=0.01, + contact_offset=0.003, + rest_offset=0.001, + restitution=0.01, + min_position_iters=32, + min_velocity_iters=8, + max_depenetration_velocity=10.0, + ), + body_scale=(1.75, 1.75, 1.75), + max_convex_hull_num=8, + init_pos=list(BREAD_INIT_POS), + init_rot=list(BREAD_INIT_ROT), + ) + ) + + +def create_pan(sim: SimulationManager) -> RigidObject: + """Create the pan mesh held below the bread by the right arm.""" + return sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="pan", + shape=MeshCfg(fpath=str(PAN_MESH_PATH), compute_uv=False), + attrs=RigidBodyAttributesCfg( + mass=0.01, + dynamic_friction=0.97, + static_friction=0.99, + angular_damping=2.0, + linear_damping=1.0, + contact_offset=0.001, + rest_offset=0.0, + restitution=0.01, + min_position_iters=32, + min_velocity_iters=8, + max_depenetration_velocity=2.0, + ), + body_scale=(1.75, 1.75, 1.75), + max_convex_hull_num=16, + init_pos=list(PAN_INIT_POS), + init_rot=list(PAN_INIT_ROT), + ) + ) + + +def settle_object(sim: SimulationManager, obj: RigidObject, step: int = 5) -> None: + """Settle an object before planning.""" + if sim.device.type == "cuda": + sim.init_gpu_physics() + obj.reset() + if step > 0: + sim.update(step=step) + obj.clear_dynamics() + + +def create_object_semantics(obj: RigidObject, label: str) -> ObjectSemantics: + """Create minimal object semantics for manually specified grasps.""" + return ObjectSemantics( + label=label, + geometry={}, + affordance=Affordance(object_label=label), + entity=obj, + ) + + +def get_hand_open_close_qpos( + robot: Robot, hand_control_part: str, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + """Get open and close qpos for a PGI gripper control part.""" + limits = robot.get_qpos_limits(name=hand_control_part)[0].to( + device=device, dtype=torch.float32 + ) + hand_open = limits[:, 0] + hand_close = torch.minimum(limits[:, 1], torch.full_like(limits[:, 1], 0.030)) + return hand_open, hand_close + + +def get_pan_handle_open_close_qpos( + robot: Robot, hand_control_part: str, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + """Get right hand qpos tuned for holding the thin pan handle.""" + limits = robot.get_qpos_limits(name=hand_control_part)[0].to( + device=device, dtype=torch.float32 + ) + hand_open = limits[:, 0] + hand_close = torch.clamp( + torch.full_like(limits[:, 1], PAN_HANDLE_CLOSE_QPOS), + min=limits[:, 0], + max=limits[:, 1], + ) + return hand_open, hand_close + + +def invert_pose(pose: torch.Tensor) -> torch.Tensor: + """Invert batched homogeneous transforms.""" + inv_pose = pose.clone() + rot_t = pose[:, :3, :3].transpose(1, 2) + inv_pose[:, :3, :3] = rot_t + inv_pose[:, :3, 3] = -torch.bmm(rot_t, pose[:, :3, 3:4]).squeeze(-1) + return inv_pose + + +def get_local_vertices(obj: RigidObject) -> torch.Tensor: + """Get scaled local mesh vertices.""" + return obj.get_vertices(env_ids=[0], scale=True)[0] + + +def compute_local_bounds(vertices: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Compute local mesh AABB from scaled vertices.""" + return vertices.min(dim=0).values, vertices.max(dim=0).values + + +def transform_points(pose: torch.Tensor, points: torch.Tensor) -> torch.Tensor: + """Transform local points by a homogeneous pose.""" + return points @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + + +def compute_world_bounds( + object_pose: torch.Tensor, + local_vertices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute world AABB from transformed local mesh vertices.""" + world_vertices = transform_points(object_pose, local_vertices) + return world_vertices.min(dim=0).values, world_vertices.max(dim=0).values + + +def rotate_pose_about_world_z(pose: torch.Tensor, yaw_deg: float) -> torch.Tensor: + """Rotate pose orientation about world Z while preserving translation.""" + yaw = math.radians(yaw_deg) + rot = torch.eye(3, dtype=pose.dtype, device=pose.device) + rot[0, 0] = math.cos(yaw) + rot[0, 1] = -math.sin(yaw) + rot[1, 0] = math.sin(yaw) + rot[1, 1] = math.cos(yaw) + rotated_pose = pose.clone() + rotated_pose[:3, :3] = rot @ pose[:3, :3] + return rotated_pose + + +def get_pan_basin_vertices(pan_vertices: torch.Tensor) -> torch.Tensor: + """Select pan vertices that belong to the basin instead of the long handle.""" + basin_vertices = pan_vertices[pan_vertices[:, 2] <= PAN_BASIN_LOCAL_Z_MAX] + if basin_vertices.numel() == 0: + logger.log_warning("Pan basin vertex mask is empty; falling back to full mesh.") + return pan_vertices + return basin_vertices + + +def build_top_down_tcp_pose( + position: torch.Tensor, device: torch.device +) -> torch.Tensor: + """Build a simple top-down TCP pose for manually grasping flat objects.""" + pose = torch.eye(4, dtype=torch.float32, device=device) + pose[:3, :3] = torch.tensor( + [ + [-0.0539, -0.9985, -0.0022], + [-0.9977, 0.0540, -0.0401], + [0.0401, 0.0000, -0.9992], + ], + dtype=torch.float32, + device=device, + ) + pose[:3, 3] = position + return pose + + +def build_flat_object_grasp_pose( + object_pose: torch.Tensor, + local_vertices: torch.Tensor, + local_min: torch.Tensor, + local_max: torch.Tensor, + device: torch.device, + *, + world_xy_offset: tuple[float, float] = (0.0, 0.0), + z_clearance: float = 0.02, +) -> torch.Tensor: + """Build a hand-tuned top-down grasp TCP pose over a flat object.""" + local_center = 0.5 * (local_min + local_max) + local_center = local_center.to(device=device, dtype=torch.float32) + grasp_position = object_pose[:3, 3] + object_pose[:3, :3] @ local_center + _, world_max = compute_world_bounds(object_pose, local_vertices) + grasp_position[0] += world_xy_offset[0] + grasp_position[1] += world_xy_offset[1] + grasp_position[2] = world_max[2] + z_clearance + return build_top_down_tcp_pose(grasp_position, device) + + +def normalize_vector(vector: torch.Tensor, fallback: torch.Tensor) -> torch.Tensor: + """Normalize a vector with a deterministic fallback for degenerate cases.""" + norm = torch.linalg.norm(vector) + if norm < 1e-6: + return fallback.to(device=vector.device, dtype=vector.dtype) + return vector / norm + + +def build_pan_handle_grasp_pose( + pan_pose: torch.Tensor, + pan_vertices: torch.Tensor, + device: torch.device, + *, + z_clearance: float = 0.006, +) -> torch.Tensor: + """Build a top-down TCP pose that pinches the pan handle.""" + handle_vertices = pan_vertices[pan_vertices[:, 2] > PAN_HANDLE_LOCAL_Z_MIN] + if handle_vertices.numel() == 0: + logger.log_warning( + "Pan handle vertex mask is empty; falling back to full mesh." + ) + handle_vertices = pan_vertices + + handle_world = transform_points(pan_pose, handle_vertices) + handle_min = handle_world.min(dim=0).values + handle_max = handle_world.max(dim=0).values + grasp_position = 0.5 * (handle_min + handle_max) + + pan_world = transform_points(pan_pose, pan_vertices) + pan_center_xy = pan_world[:, :2].mean(dim=0) + handle_dir_xy = grasp_position[:2] - pan_center_xy + handle_axis = normalize_vector( + torch.tensor( + [handle_dir_xy[0], handle_dir_xy[1], 0.0], + dtype=torch.float32, + device=device, + ), + torch.tensor([1.0, -0.2, 0.0], dtype=torch.float32, device=device), + ) + grasp_position[:2] -= handle_axis[:2] * PAN_HANDLE_ROOT_OFFSET + grasp_position[2] = handle_max[2] + z_clearance + + y_axis = handle_axis + z_axis = torch.tensor([0.0, 0.0, -1.0], dtype=torch.float32, device=device) + x_axis = normalize_vector( + torch.cross(y_axis, z_axis, dim=0), + torch.tensor([1.0, 0.0, 0.0], dtype=torch.float32, device=device), + ) + y_axis = normalize_vector( + torch.cross(z_axis, x_axis, dim=0), + y_axis, + ) + + pose = torch.eye(4, dtype=torch.float32, device=device) + pose[:3, 0] = x_axis + pose[:3, 1] = y_axis + pose[:3, 2] = z_axis + pose[:3, 3] = grasp_position + return pose + + +def build_support_object_target_pose( + pan_pose: torch.Tensor, + device: torch.device, +) -> torch.Tensor: + """Build target pose for the support pan.""" + pose = pan_pose.clone().to(device=device, dtype=torch.float32) + pose[0, 3] = PAN_TARGET_CENTER_XY[0] + pose[1, 3] = PAN_TARGET_CENTER_XY[1] + pose[2, 3] = TABLE_TOP_Z + 0.001 + PAN_TARGET_Z_LIFT + return pose + + +def build_placing_object_target_pose( + bread_pose: torch.Tensor, + bread_vertices: torch.Tensor, + pan_vertices: torch.Tensor, + support_target_pose: torch.Tensor, + device: torch.device, +) -> torch.Tensor: + """Build target bread pose aligned above the pan.""" + pose = rotate_pose_about_world_z( + bread_pose.clone().to(device=device, dtype=torch.float32), + BREAD_TARGET_WORLD_YAW_DEG, + ) + pan_basin_world = transform_points( + support_target_pose, get_pan_basin_vertices(pan_vertices) + ) + basin_center_xy = 0.5 * ( + pan_basin_world[:, :2].min(dim=0).values + + pan_basin_world[:, :2].max(dim=0).values + ) + pose[0, 3] = basin_center_xy[0] + BREAD_PLACE_TARGET_OFFSET_XY[0] + pose[1, 3] = basin_center_xy[1] + BREAD_PLACE_TARGET_OFFSET_XY[1] + pan_top_z = pan_basin_world[:, 2].max() + bread_bottom_z = compute_world_bounds(pose, bread_vertices)[0][2] + pose[2, 3] += pan_top_z + BREAD_ON_PAN_CLEARANCE - bread_bottom_z + return pose + + +def format_tensor(tensor: torch.Tensor) -> str: + """Format tensor values for compact logging.""" + rounded = (tensor.detach().cpu() * 10000.0).round() / 10000.0 + return str(rounded.tolist()) + + +def compute_actual_held_state( + robot: Robot, + semantics: ObjectSemantics, + object_pose: torch.Tensor, + arm_control_part: str, + device: torch.device, +) -> HeldObjectState: + """Build held-object state from current object pose and current TCP FK.""" + arm_qpos = robot.get_qpos(name=arm_control_part).to(device=device) + tcp_pose = robot.compute_fk( + arm_qpos, + name=arm_control_part, + to_matrix=True, + ).to(device=device, dtype=torch.float32) + object_pose = object_pose.to(device=device, dtype=torch.float32).unsqueeze(0) + object_to_eef = torch.bmm(invert_pose(object_pose), tcp_pose) + return HeldObjectState( + semantics=semantics, + object_to_eef=object_to_eef, + grasp_xpos=tcp_pose, + ) + + +def plan_manual_pick( + motion_gen: MotionGenerator, + robot: Robot, + *, + arm_control_part: str, + hand_control_part: str, + grasp_pose: torch.Tensor, + hand_open_qpos: torch.Tensor, + hand_close_qpos: torch.Tensor, + sample_interval: int, + hand_interp_steps: int, + pre_grasp_distance: float, + lift_height: float, + start_state: WorldState, +) -> tuple[bool, torch.Tensor, WorldState]: + """Plan a demo-local pick trajectory from a hand-authored TCP grasp pose.""" + builder = TrajectoryBuilder(motion_gen) + device = start_state.last_qpos.device + n_envs = robot.get_qpos().shape[0] + robot_dof = robot.dof + arm_joint_ids = robot.get_joint_ids(name=arm_control_part) + hand_joint_ids = robot.get_joint_ids(name=hand_control_part) + arm_dof = len(arm_joint_ids) + + grasp_xpos = builder.resolve_pose_target(grasp_pose, n_envs=n_envs) + grasp_z = grasp_xpos[:, :3, 2] + pre_grasp_xpos = builder.apply_local_offset( + grasp_xpos, -grasp_z * pre_grasp_distance + ) + lift_xpos = builder.apply_local_offset( + grasp_xpos, + torch.tensor([0.0, 0.0, lift_height], device=device), + ) + start_arm_qpos = builder.resolve_start_qpos( + start_state.last_qpos[:, arm_joint_ids], + n_envs=n_envs, + arm_dof=arm_dof, + control_part=arm_control_part, + ) + n_approach, n_close, n_lift = builder.split_three_phase( + sample_interval, + hand_interp_steps, + first_phase_name="approach", + third_phase_name="lift", + ) + + target_states_list = [ + [ + PlanState(xpos=pre_grasp_xpos[i], move_type=MoveType.EEF_MOVE), + PlanState(xpos=grasp_xpos[i], move_type=MoveType.EEF_MOVE), + ] + for i in range(n_envs) + ] + ok, approach_arm = builder.plan_arm_traj( + target_states_list, + start_arm_qpos, + n_approach, + control_part=arm_control_part, + arm_dof=arm_dof, + ) + if not ok: + logger.log_warning(f"Failed to plan {arm_control_part} manual pick approach.") + return ( + False, + torch.empty((n_envs, 0, robot_dof), dtype=torch.float32, device=device), + start_state, + ) + + grasp_arm_qpos = approach_arm[:, -1, :] + target_states_list = [ + [PlanState(xpos=lift_xpos[i], move_type=MoveType.EEF_MOVE)] + for i in range(n_envs) + ] + ok, lift_arm = builder.plan_arm_traj( + target_states_list, + grasp_arm_qpos, + n_lift, + control_part=arm_control_part, + arm_dof=arm_dof, + ) + if not ok: + logger.log_warning(f"Failed to plan {arm_control_part} manual pick lift.") + return ( + False, + torch.empty((n_envs, 0, robot_dof), dtype=torch.float32, device=device), + start_state, + ) + + hand_open = builder.expand_hand_qpos( + hand_open_qpos, n_envs=n_envs, hand_dof=len(hand_joint_ids) + ) + hand_close = builder.expand_hand_qpos( + hand_close_qpos, n_envs=n_envs, hand_dof=len(hand_joint_ids) + ) + hand_close_path = builder.interpolate_hand_qpos( + hand_open, hand_close, n_waypoints=n_close + ) + + full = start_state.last_qpos.unsqueeze(1).repeat(1, sample_interval, 1).clone() + full[:, :n_approach, arm_joint_ids] = approach_arm + full[:, :n_approach, hand_joint_ids] = hand_open + full[:, n_approach : n_approach + n_close, arm_joint_ids] = ( + grasp_arm_qpos.unsqueeze(1) + ) + full[:, n_approach : n_approach + n_close, hand_joint_ids] = hand_close_path + full[:, n_approach + n_close :, arm_joint_ids] = lift_arm + full[:, n_approach + n_close :, hand_joint_ids] = hand_close + return True, full, WorldState(last_qpos=full[:, -1, :].clone(), held_object=None) + + +def log_scene_targets( + bread_pose: torch.Tensor, + pan_pose: torch.Tensor, + support_target_pose: torch.Tensor | None = None, + placing_target_pose: torch.Tensor | None = None, +) -> None: + """Log compact object and target positions for diagnosis.""" + logger.log_info( + "scene objects: " + f"bread_origin={format_tensor(bread_pose[:3, 3])}, " + f"pan_origin={format_tensor(pan_pose[:3, 3])}" + ) + if support_target_pose is not None and placing_target_pose is not None: + logger.log_info( + "coordinated targets: " + f"support_pan_origin={format_tensor(support_target_pose[:3, 3])}, " + f"placing_bread_origin={format_tensor(placing_target_pose[:3, 3])}" + ) + + +def draw_coordinated_axes( + sim: SimulationManager, + support_target_pose: torch.Tensor, + placing_target_pose: torch.Tensor, +) -> None: + """Draw coordinate-frame markers for coordinated placement targets.""" + draw_axis_marker( + sim, + "support_pan_target_axis", + support_target_pose, + axis_len=0.08, + axis_size=0.004, + arena_index=0, + ) + draw_axis_marker( + sim, + "placing_bread_target_axis", + placing_target_pose, + axis_len=0.08, + axis_size=0.004, + arena_index=0, + ) + + +def log_action_plan( + robot: Robot, + action_name: str, + traj: torch.Tensor, + joint_ids: list[int], + segments: dict[str, int] | None = None, +) -> None: + """Log common action plan details.""" + joint_names = [robot.joint_names[joint_id] for joint_id in joint_ids] + logger.log_info(f"{action_name} joint ids: {joint_ids}") + logger.log_info(f"{action_name} joint names: {joint_names}") + logger.log_info(f"{action_name} trajectory shape: {tuple(traj.shape)}") + if segments is not None: + logger.log_info(f"{action_name} trajectory segments: {segments}") + + +def log_execution_state( + robot: Robot, + bread: RigidObject, + pan: RigidObject, + step_idx: int, + total_steps: int, +) -> None: + """Log hand and object state during execution.""" + bread_pose = bread.get_local_pose(to_matrix=True) + pan_pose = pan.get_local_pose(to_matrix=True) + left_hand = robot.get_qpos(name="left_hand") + right_hand = robot.get_qpos(name="right_hand") + logger.log_info( + f"step={step_idx}/{total_steps - 1}, " + f"left_hand={format_tensor(left_hand[0])}, " + f"right_hand={format_tensor(right_hand[0])}, " + f"bread_pos={format_tensor(bread_pose[0, :3, 3])}, " + f"pan_pos={format_tensor(pan_pose[0, :3, 3])}" + ) + + +def execute_trajectory( + sim: SimulationManager, + robot: Robot, + traj: torch.Tensor, + joint_ids: list[int], + bread: RigidObject, + pan: RigidObject, + debug_state: bool, +) -> None: + """Play a planned trajectory in simulation.""" + total_steps = traj.shape[1] + log_stride = max(1, total_steps // 10) + for i in range(total_steps): + robot.set_qpos(traj[:, i, :], joint_ids=joint_ids) + sim.update(step=TRAJECTORY_SIM_STEPS) + if debug_state and (i % log_stride == 0 or i == total_steps - 1): + log_execution_state(robot, bread, pan, i, total_steps) + time.sleep(1e-2) + + +def run_coordinated_placement_demo( + args: argparse.Namespace, sim: SimulationManager, robot: Robot +) -> None: + """Plan and optionally execute pick-up and coordinated placement.""" + create_table(sim) + bread = create_bread(sim) + pan = create_pan(sim) + settle_object(sim, bread, step=0) + settle_object(sim, pan, step=0) + bread_pose = bread.get_local_pose(to_matrix=True)[0].to( + device=sim.device, dtype=torch.float32 + ) + pan_pose = pan.get_local_pose(to_matrix=True)[0].to( + device=sim.device, dtype=torch.float32 + ) + bread_vertices = get_local_vertices(bread) + pan_vertices = get_local_vertices(pan) + bread_local_min, bread_local_max = compute_local_bounds(bread_vertices) + log_scene_targets(bread_pose, pan_pose) + bread_semantics = create_object_semantics(bread, BREAD_LABEL) + pan_semantics = create_object_semantics(pan, PAN_LABEL) + motion_gen = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=robot.uid)) + ) + + right_open, right_close = get_pan_handle_open_close_qpos( + robot, "right_hand", sim.device + ) + left_open, left_close = get_hand_open_close_qpos(robot, "left_hand", sim.device) + coordinated_action = CoordinatedPlacement( + motion_generator=motion_gen, + cfg=CoordinatedPlacementCfg( + control_part="dual_arm", + placing_arm_control_part="left_arm", + support_arm_control_part="right_arm", + placing_hand_control_part="left_hand", + support_hand_control_part="right_hand", + placing_hand_open_qpos=left_open, + placing_hand_close_qpos=left_close, + support_hand_close_qpos=right_close, + release=True, + placing_height_offset=BREAD_TARGET_HEIGHT_OFFSET, + support_height_offset=SUPPORT_TARGET_HEIGHT_OFFSET, + lift_height=PLACE_LIFT_HEIGHT, + sample_interval=COORDINATED_SAMPLE_INTERVAL, + hand_interp_steps=10, + hold_steps=6, + retreat_steps=18, + ), + ) + full_joint_ids = list(range(robot.dof)) + state = WorldState(last_qpos=robot.get_qpos().clone()) + + wait_for_user = not args.auto_play and not args.headless_play + if not args.diagnose_plan and not args.headless_play and not args.headless: + sim.open_window() + if wait_for_user: + input("Inspect the scene, then press Enter to plan left pick-up...") + + bread_grasp_pose = build_flat_object_grasp_pose( + bread_pose, + bread_vertices, + bread_local_min, + bread_local_max, + sim.device, + z_clearance=BREAD_GRASP_Z_CLEARANCE, + ) + start_time = time.time() + left_pick_success, left_pick_traj, state = plan_manual_pick( + motion_gen, + robot, + arm_control_part="left_arm", + hand_control_part="left_hand", + grasp_pose=bread_grasp_pose, + hand_open_qpos=left_open, + hand_close_qpos=left_close, + sample_interval=PICK_SAMPLE_INTERVAL, + hand_interp_steps=10, + pre_grasp_distance=PICK_APPROACH_DISTANCE, + lift_height=0.12, + start_state=state, + ) + logger.log_info( + f"Plan left bread pick-up cost time: {time.time() - start_time:.2f} seconds" + ) + if not left_pick_success: + logger.log_warning("Failed to plan left bread pick-up trajectory.") + return + log_action_plan(robot, "left_pick_up", left_pick_traj, full_joint_ids) + bread_object_to_eef = torch.bmm( + invert_pose(bread_pose.unsqueeze(0)), + bread_grasp_pose.unsqueeze(0), + ) + bread_held_state = HeldObjectState( + semantics=bread_semantics, + object_to_eef=bread_object_to_eef, + grasp_xpos=bread_grasp_pose.unsqueeze(0), + ) + + pan_grasp_pose = build_pan_handle_grasp_pose( + pan_pose, + pan_vertices, + sim.device, + z_clearance=PAN_GRASP_Z_CLEARANCE, + ) + start_time = time.time() + right_pick_success, right_pick_traj, state = plan_manual_pick( + motion_gen, + robot, + arm_control_part="right_arm", + hand_control_part="right_hand", + grasp_pose=pan_grasp_pose, + hand_open_qpos=right_open, + hand_close_qpos=right_close, + sample_interval=PAN_PICK_SAMPLE_INTERVAL, + hand_interp_steps=PAN_PICK_HAND_INTERP_STEPS, + pre_grasp_distance=PICK_APPROACH_DISTANCE, + lift_height=0.10, + start_state=state, + ) + logger.log_info( + f"Plan right pan pick-up cost time: {time.time() - start_time:.2f} seconds" + ) + if not right_pick_success: + logger.log_warning("Failed to plan right pan pick-up trajectory.") + return + log_action_plan(robot, "right_pick_up", right_pick_traj, full_joint_ids) + pan_object_to_eef = torch.bmm( + invert_pose(pan_pose.unsqueeze(0)), + pan_grasp_pose.unsqueeze(0), + ) + pan_held_state = HeldObjectState( + semantics=pan_semantics, + object_to_eef=pan_object_to_eef, + grasp_xpos=pan_grasp_pose.unsqueeze(0), + ) + + if args.diagnose_plan: + robot.set_qpos(state.last_qpos, joint_ids=full_joint_ids) + else: + if wait_for_user: + input("Press Enter to execute both pick-up trajectories...") + execute_trajectory( + sim, + robot, + left_pick_traj, + full_joint_ids, + bread, + pan, + args.debug_state, + ) + bread.clear_dynamics() + execute_trajectory( + sim, + robot, + right_pick_traj, + full_joint_ids, + bread, + pan, + args.debug_state, + ) + pan.clear_dynamics() + bread_pose = bread.get_local_pose(to_matrix=True)[0].to( + device=sim.device, dtype=torch.float32 + ) + actual_pan_pose = pan.get_local_pose(to_matrix=True)[0].to( + device=sim.device, dtype=torch.float32 + ) + bread_held_state = compute_actual_held_state( + robot, + bread_semantics, + bread_pose, + "left_arm", + sim.device, + ) + pan_held_state = compute_actual_held_state( + robot, + pan_semantics, + actual_pan_pose, + "right_arm", + sim.device, + ) + + support_target_pose = build_support_object_target_pose(pan_pose, sim.device) + placing_target_pose = build_placing_object_target_pose( + bread_pose, + bread_vertices, + pan_vertices, + support_target_pose, + sim.device, + ) + log_scene_targets( + bread_pose, + pan_pose, + support_target_pose, + placing_target_pose, + ) + if not args.auto_play and not args.no_vis_eef_axis: + draw_coordinated_axes( + sim, + support_target_pose, + placing_target_pose, + ) + coordinated_target = CoordinatedPlacementTarget( + placing_object_target_pose=placing_target_pose, + support_object_target_pose=support_target_pose, + placing_held_object=bread_held_state, + support_held_object=pan_held_state, + placing_height_offset=BREAD_TARGET_HEIGHT_OFFSET, + support_height_offset=SUPPORT_TARGET_HEIGHT_OFFSET, + release=True, + ) + start_time = time.time() + coordinated_result = coordinated_action.execute(coordinated_target, state) + coordinated_success = coordinated_result.success + coordinated_traj = coordinated_result.trajectory + state = coordinated_result.next_state + logger.log_info( + "Plan coordinated placement cost time: " + f"{time.time() - start_time:.2f} seconds" + ) + if not coordinated_success: + logger.log_warning("Failed to plan coordinated placement trajectory.") + return + log_action_plan( + robot, + "coordinated_placement", + coordinated_traj, + full_joint_ids, + coordinated_action._compute_segment_lengths(coordinated_action.cfg.release), + ) + + if args.diagnose_plan: + return + + recording_started = start_auto_play_recording( + sim, + args, + video_prefix="coordinated_placement_auto_play", + look_at=( + (-0.25, 0.0, 2.5), + (-0.05, 0.0, 0.72), + (0.0, 0.0, 1.0), + ), + ) + try: + if args.auto_play and not args.no_vis_eef_axis: + draw_coordinated_axes( + sim, + support_target_pose, + placing_target_pose, + ) + if wait_for_user: + input("Press Enter to execute coordinated placement...") + execute_trajectory( + sim, + robot, + coordinated_traj, + full_joint_ids, + bread, + pan, + args.debug_state, + ) + for _ in range(80): + robot.set_qpos(state.last_qpos, joint_ids=full_joint_ids) + sim.update(step=2) + time.sleep(1e-2) + finally: + stop_auto_play_recording(sim, recording_started) + if wait_for_user: + input("Press Enter to exit the simulation...") + + +def main() -> None: + """Run the coordinated placement demo.""" + args = parse_arguments() + sim = initialize_simulation(args) + robot = create_dual_ur5_robot(sim) + run_coordinated_placement_demo(args, sim, robot) + + +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py index 9190e524d..245239feb 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -53,11 +53,11 @@ from embodichain.lab.sim.objects import RigidObject, Robot from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg from embodichain.lab.sim.shapes import CubeCfg -from embodichain.lab.sim.solvers import PytorchSolverCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( draw_axis_marker, get_tutorial_window_size, + make_ur5_solver_cfg, start_auto_play_recording, stop_auto_play_recording, ) @@ -160,18 +160,7 @@ def create_robot(sim: SimulationManager, position=(0.0, 0.0, 0.0)) -> Robot: "arm": ["JOINT[0-9]"], "hand": [GRIPPER_HAND_JOINT_PATTERN], }, - solver_cfg={ - "arm": PytorchSolverCfg( - end_link_name="ee_link", - root_link_name="base_link", - tcp=[ - [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, 0.0], - [0.0, 0.0, 1.0, GRIPPER_TCP_Z], - [0.0, 0.0, 0.0, 1.0], - ], - ) - }, + solver_cfg={"arm": make_ur5_solver_cfg(GRIPPER_TCP_Z)}, init_qpos=[0.0, -1.57, 1.57, -1.57, -1.57, 0.0, 0.0, 0.0], init_pos=position, ) diff --git a/scripts/tutorials/atomic_action/tutorial_utils.py b/scripts/tutorials/atomic_action/tutorial_utils.py index 8b8ea6ba7..282b6191f 100644 --- a/scripts/tutorials/atomic_action/tutorial_utils.py +++ b/scripts/tutorials/atomic_action/tutorial_utils.py @@ -27,6 +27,7 @@ from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.cfg import MarkerCfg, RobotCfg from embodichain.lab.sim.robots import URRobotCfg +from embodichain.lab.sim.solvers import URSolverCfg RECORD_WIDTH = 640 RECORD_HEIGHT = 480 @@ -47,6 +48,23 @@ GRIPPER_TCP_Z = 0.15 +def make_ur5_solver_cfg(tcp_z: float) -> URSolverCfg: + """Create the UR5 arm solver cfg used by atomic-action tutorials.""" + cfg = URSolverCfg( + ur_type="ur5", + end_link_name="ee_link", + root_link_name="base_link", + tcp=[ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, tcp_z], + [0.0, 0.0, 0.0, 1.0], + ], + ) + cfg.urdf_path = None + return cfg + + def get_tutorial_window_size(args: argparse.Namespace) -> tuple[int, int]: """Return the viewer window size used by atomic-action tutorials.""" return VIEWER_WIDTH, VIEWER_HEIGHT @@ -104,6 +122,7 @@ def draw_axis_marker( xpos: torch.Tensor, axis_len: float = DEFAULT_AXIS_LEN, axis_size: float = DEFAULT_AXIS_SIZE, + arena_index: int = -1, ) -> None: """Draw a named coordinate-frame marker for a semantic tutorial target.""" sim.draw_marker( @@ -113,6 +132,7 @@ def draw_axis_marker( axis_xpos=xpos, axis_size=axis_size, axis_len=axis_len, + arena_index=arena_index, ) ) @@ -193,6 +213,7 @@ def create_ur5_gripper_robot_cfg( "GRIPPER_TCP_Z", "GRIPPER_URDF_PATH", "create_ur5_gripper_robot_cfg", + "make_ur5_solver_cfg", "get_tutorial_window_size", "start_auto_play_recording", "stop_auto_play_recording", diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index 051efeaf0..558d0b992 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -30,6 +30,7 @@ AtomicAction, CoordinatedHeldObjectState, CoordinatedPickmentTarget, + CoordinatedPlacementTarget, GraspTarget, HeldObjectState, HeldObjectPoseTarget, @@ -42,6 +43,8 @@ from embodichain.lab.sim.atomic_actions.actions import ( CoordinatedPickment, CoordinatedPickmentCfg, + CoordinatedPlacement, + CoordinatedPlacementCfg, MoveEndEffector, MoveEndEffectorCfg, MoveJoints, @@ -705,3 +708,145 @@ def test_execute_returns_full_dof_trajectory_and_dual_held_state(self): CoordinatedHeldObjectState, ) assert result.next_state.held_object is None + + +# --------------------------------------------------------------------------- +# CoordinatedPlacement +# --------------------------------------------------------------------------- + + +class TestCoordinatedPlacementAction: + def setup_method(self): + self.mg = _make_dual_arm_mock_motion_generator() + self.cfg = CoordinatedPlacementCfg( + placing_hand_open_qpos=_hand_open(), + placing_hand_close_qpos=_hand_close(), + support_hand_close_qpos=_hand_close(), + sample_interval=30, + hand_interp_steps=4, + hold_steps=3, + retreat_steps=5, + lift_height=0.08, + ) + self.action = CoordinatedPlacement(self.mg, cfg=self.cfg) + + def _make_target(self) -> CoordinatedPlacementTarget: + placing_pose = torch.eye(4) + placing_pose[0, 3] = 0.2 + support_pose = torch.eye(4) + support_pose[0, 3] = 0.2 + support_pose[2, 3] = -0.05 + + placing_object_to_eef = torch.eye(4) + placing_object_to_eef[2, 3] = 0.12 + support_object_to_eef = torch.eye(4) + support_object_to_eef[2, 3] = 0.10 + + placing_semantics = ObjectSemantics( + affordance=AntipodalAffordance(), geometry={}, label="placing" + ) + support_semantics = ObjectSemantics( + affordance=AntipodalAffordance(), geometry={}, label="support" + ) + return CoordinatedPlacementTarget( + placing_object_target_pose=placing_pose, + support_object_target_pose=support_pose, + placing_held_object=HeldObjectState( + semantics=placing_semantics, + object_to_eef=placing_object_to_eef, + grasp_xpos=torch.eye(4), + ), + support_held_object=HeldObjectState( + semantics=support_semantics, + object_to_eef=support_object_to_eef, + grasp_xpos=torch.eye(4), + ), + ) + + def test_target_type_is_coordinated_placement_target(self): + assert CoordinatedPlacement.TargetType is CoordinatedPlacementTarget + + def test_init_sets_dual_arm_and_hand_joint_ids(self): + assert self.action.dual_arm_joint_ids == list(range(DUAL_ARM_DOF)) + assert self.action.placing_arm_joint_ids == list(range(ARM_DOF)) + assert self.action.support_arm_joint_ids == list(range(ARM_DOF, DUAL_ARM_DOF)) + assert self.action.placing_hand_joint_ids == list( + range(DUAL_ARM_DOF, DUAL_ARM_DOF + HAND_DOF) + ) + assert self.action.support_hand_joint_ids == list( + range(DUAL_ARM_DOF + HAND_DOF, DUAL_TOTAL_DOF) + ) + assert self.action.joint_ids == list(range(DUAL_TOTAL_DOF)) + + def test_resolve_target_composes_object_and_tcp_poses(self): + target = self._make_target() + placing_xpos, support_xpos, release, held_state = self.action._resolve_target( + target + ) + assert placing_xpos.shape == (NUM_ENVS, 4, 4) + assert support_xpos.shape == (NUM_ENVS, 4, 4) + assert placing_xpos[0, 2, 3].item() == pytest.approx(0.12) + assert support_xpos[0, 2, 3].item() == pytest.approx(0.05) + assert release is True + assert held_state.semantics is target.support_held_object.semantics + assert held_state.object_to_eef.shape == (NUM_ENVS, 4, 4) + assert held_state.grasp_xpos.shape == (NUM_ENVS, 4, 4) + assert torch.allclose( + held_state.object_to_eef, + target.support_held_object.object_to_eef.unsqueeze(0).repeat( + NUM_ENVS, 1, 1 + ), + ) + + def test_segment_lengths_sum_to_sample_interval(self): + segments = self.action._compute_segment_lengths(self.cfg.release) + assert sum(segments.values()) == self.cfg.sample_interval + assert segments["approach"] >= 2 + assert segments["release"] == self.cfg.hand_interp_steps + assert segments["retreat"] == self.cfg.retreat_steps + + def test_execute_returns_full_dof_and_final_hand_states(self): + target = self._make_target() + state = WorldState(last_qpos=torch.zeros(NUM_ENVS, DUAL_TOTAL_DOF)) + + def interpolate(trajectory, interp_num, device): + weights = torch.linspace( + 0.0, + 1.0, + steps=interp_num, + dtype=trajectory.dtype, + device=trajectory.device, + ) + return torch.lerp( + trajectory[:, :1], + trajectory[:, -1:], + weights.view(1, -1, 1), + ) + + with patch( + "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", + side_effect=interpolate, + ): + result = self.action.execute(target, state) + + assert result.success is True + assert result.trajectory.shape == ( + NUM_ENVS, + self.cfg.sample_interval, + DUAL_TOTAL_DOF, + ) + assert torch.allclose( + result.trajectory[:, -1, self.action.placing_hand_joint_ids], + _hand_open().unsqueeze(0).repeat(NUM_ENVS, 1), + ) + assert torch.allclose( + result.trajectory[:, -1, self.action.support_hand_joint_ids], + _hand_close().unsqueeze(0).repeat(NUM_ENVS, 1), + ) + assert result.next_state.held_object is not None + assert ( + result.next_state.held_object.semantics + is target.support_held_object.semantics + ) + assert result.next_state.held_object.object_to_eef.shape == (NUM_ENVS, 4, 4) + assert result.next_state.held_object.grasp_xpos.shape == (NUM_ENVS, 4, 4)