Add press atomic action

docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst

@@ -33,6 +33,8 @@ embodichain.lab.sim.atomic_actions
       MoveHeldObject
       PlaceCfg
       Place
+      PressCfg
+      Press
       AtomicActionEngine
 
 .. currentmodule:: embodichain.lab.sim.atomic_actions
@@ -153,6 +155,15 @@ Actions
     :members:
     :show-inheritance:
 
+.. autoclass:: PressCfg
+    :members:
+    :exclude-members: __init__, copy, replace, to_dict
+    :show-inheritance:
+
+.. autoclass:: Press
+    :members:
+    :show-inheritance:
+
 Engine & Registry
 -----------------
 

docs/source/overview/sim/atomic_actions/builtin_actions.md

@@ -14,6 +14,7 @@ The following actions are available out of the box:
 | `PickUp` | Single | `GraspTarget` — object semantics | Approach → close gripper → lift | <img src="../../../_static/atomic_actions/pickup.gif" alt="PickUp" width="480" style="max-width: 100%;" /> |
 | `MoveHeldObject` | Single | `HeldObjectPoseTarget` — held-object pose | Move held object while keeping gripper closed | <img src="../../../_static/atomic_actions/move_held_object.gif" alt="MoveHeldObject" width="480" style="max-width: 100%;" /> |
 | `Place` | Single | `EndEffectorPoseTarget` — EEF release pose | Lower → open gripper → retract | <img src="../../../_static/atomic_actions/place.gif" alt="Place" width="480" style="max-width: 100%;" /> |
+| `Press` | Single | `EndEffectorPoseTarget` — EEF press pose | Close gripper → press down → return | <img src="../../../_static/atomic_actions/press.gif" alt="Press" width="480" style="max-width: 100%;" /> |
 
 ---
 
@@ -123,3 +124,26 @@ down to the target pose. On success, the returned `WorldState` clears `held_obje
 `(4, 4)` or `(n_envs, 4, 4)`.
 
 ![Place demo](../../../_static/atomic_actions/place.gif)
+
+---
+
+## `Press`
+
+Three-phase contact motion: *close gripper → press down → return*. This is useful
+for button-like or contact-based interactions where the end-effector should reach a
+target pose and then return to the pre-press arm pose.
+
+`Press` does not create or clear `WorldState.held_object`; it preserves the state
+threaded into it.
+
+| Config field | Default | Description |
+|---|---|---|
+| `hand_close_qpos` | `None` | **Required.** Gripper closed joint positions |
+| `hand_control_part` | `"hand"` | Robot control part for the gripper |
+| `hand_interp_steps` | `5` | Waypoints for the gripper close phase |
+| `sample_interval` | `80` | Total waypoints across all three phases |
+
+**Target:** `EndEffectorPoseTarget(xpos=...)` — the EEF pose to press, a `torch.Tensor`
+of shape `(4, 4)` or `(n_envs, 4, 4)`.
+
+![Press demo](../../../_static/atomic_actions/press.gif)

docs/source/overview/sim/atomic_actions/index.md

@@ -57,7 +57,7 @@ and every action declares the target type, or tuple of target types, it accepts
 
 | Target | Constructor | Used by |
 |---|---|---|
-| `EndEffectorPoseTarget` | `EndEffectorPoseTarget(xpos)` | `MoveEndEffector`, `Place` |
+| `EndEffectorPoseTarget` | `EndEffectorPoseTarget(xpos)` | `MoveEndEffector`, `Place`, `Press` |
 | `JointPositionTarget` | `JointPositionTarget(qpos)` | `MoveJoints` |
 | `NamedJointPositionTarget` | `NamedJointPositionTarget(name)` | `MoveJoints` |
 | `GraspTarget` | `GraspTarget(semantics)` | `PickUp` |
@@ -94,7 +94,7 @@ action's `TargetType` before calling `execute`:
 
 | Target | Holds | Accepted by |
 |---|---|---|
-| `EndEffectorPoseTarget(xpos)` | EEF pose tensor `(4,4)` or `(n_envs,4,4)` | `MoveEndEffector`, `Place` |
+| `EndEffectorPoseTarget(xpos)` | EEF pose tensor `(4,4)` or `(n_envs,4,4)` | `MoveEndEffector`, `Place`, `Press` |
 | `JointPositionTarget(qpos)` | Control-part qpos tensor `(control_dof,)` or `(n_envs, control_dof)` | `MoveJoints` |
 | `NamedJointPositionTarget(name)` | Name resolved from `MoveJointsCfg.named_joint_positions` | `MoveJoints` |
 | `GraspTarget(semantics)` | `ObjectSemantics` (affordance + entity) | `PickUp` |
@@ -110,6 +110,7 @@ action's `TargetType` before calling `execute`:
 | `Place` | Clears it to `None` |
 | `MoveEndEffector` | Leaves it unchanged |
 | `MoveJoints` | Leaves it unchanged |
+| `Press` | Leaves it unchanged |
 
 If a step fails, `run()` returns `success=False` with the partial trajectory concatenated up
 to (but not including) the failed step, and the `WorldState` going into that step.
@@ -136,6 +137,8 @@ from embodichain.lab.sim.atomic_actions import (
     MoveHeldObjectCfg,
     Place,
     PlaceCfg,
+    Press,
+    PressCfg,
     MoveEndEffector,
     MoveEndEffectorCfg,
 )
@@ -159,6 +162,11 @@ move_held_object_cfg = MoveHeldObjectCfg(
     hand_close_qpos=torch.tensor([0.025, 0.025]),
 )
 move_cfg = MoveEndEffectorCfg(control_part="arm")
+press_cfg = PressCfg(
+    control_part="arm",
+    hand_control_part="hand",
+    hand_close_qpos=torch.tensor([0.025, 0.025]),
+)
 move_joints_cfg = MoveJointsCfg(
     control_part="arm",
     named_joint_positions={"home": torch.zeros(6)},
@@ -170,6 +178,7 @@ engine.register(PickUp(motion_gen, cfg=pickup_cfg))
 engine.register(MoveHeldObject(motion_gen, cfg=move_held_object_cfg))
 engine.register(Place(motion_gen, cfg=place_cfg))
 engine.register(MoveEndEffector(motion_gen, cfg=move_cfg))
+engine.register(Press(motion_gen, cfg=press_cfg))
 engine.register(MoveJoints(motion_gen, cfg=move_joints_cfg))
 
 # 3. Describe the object to pick
@@ -297,6 +306,7 @@ builtin_actions
   - `scripts/tutorials/atomic_action/pickup.py`
   - `scripts/tutorials/atomic_action/move_held_object.py`
   - `scripts/tutorials/atomic_action/place.py`
+  - `scripts/tutorials/atomic_action/press.py`
 
 Run a demo in headless CPU mode with `--auto_play --headless --device cpu` to record
 an MP4 under `outputs/videos`. For example:

docs/source/tutorial/atomic_actions.rst

@@ -5,7 +5,7 @@ Atomic Actions
 
 EmbodiChain's **atomic action** layer provides a high-level, composable interface for common
 manipulation primitives such as *move end-effector*, *move joints*, *pick up*,
-*move held object*, and *place*.  Each action
+*move held object*, *place*, and *press*.  Each action
 encapsulates the full planning pipeline — grasp-pose estimation, IK, trajectory generation, and
 gripper interpolation — behind a single ``execute(target, state)`` call, making it straightforward
 to chain multiple actions together into complex robot behaviours.
@@ -18,7 +18,7 @@ Key Features
   ``GraspTarget`` (wrapping an ``ObjectSemantics``), or ``HeldObjectPoseTarget``. The
   engine checks each step's target against the action's declared ``TargetType`` before running.
 - **Built-in primitives** — ``MoveEndEffector``, ``MoveJoints``, ``PickUp``, ``MoveHeldObject``,
-  and ``Place``
+  ``Place``, and ``Press``
   cover the most common tabletop manipulation workflows out of the box.
   See :doc:`/overview/sim/atomic_actions/index` for configs and target types.
 - **Extensible registry** — custom action *classes* can be registered globally with
@@ -34,14 +34,15 @@ For the full design overview, architecture diagram, and extension guide see
 The Code
 --------
 
-Focused demo scripts are available for the five built-in primitives in the
+Focused demo scripts are available for the built-in primitives in the
 ``scripts/tutorials/atomic_action`` directory:
 
 - ``move_end_effector.py``
 - ``move_joints.py``
 - ``pickup.py``
 - ``move_held_object.py``
 - ``place.py``
+- ``press.py``
 
 Each script supports interactive inspection by default. Add ``--auto_play`` to skip
 keyboard prompts, and combine it with ``--headless --device cpu`` to record an MP4 under

embodichain/lab/sim/atomic_actions/__init__.py

@@ -18,7 +18,7 @@
 
 This module provides a unified interface for the atomic motion primitives
 (``move_end_effector``, ``move_joints``, ``pick_up``, ``move_held_object``,
-``place``), with typed targets, a ``WorldState`` threaded across sequenced
+``place``, ``press``), with typed targets, a ``WorldState`` threaded across sequenced
 actions, and extensible custom action registration.
 """
 
@@ -47,11 +47,13 @@
     MoveHeldObject,
     PickUp,
     Place,
+    Press,
     MoveEndEffectorCfg,
     MoveJointsCfg,
     MoveHeldObjectCfg,
     PickUpCfg,
     PlaceCfg,
+    PressCfg,
 )
 from .engine import (
     AtomicActionEngine,
@@ -84,11 +86,13 @@
     "MoveHeldObject",
     "PickUp",
     "Place",
+    "Press",
     "MoveEndEffectorCfg",
     "MoveJointsCfg",
     "MoveHeldObjectCfg",
     "PickUpCfg",
     "PlaceCfg",
+    "PressCfg",
     # Engine
     "AtomicActionEngine",
     "register_action",

embodichain/lab/sim/atomic_actions/actions.py

@@ -16,10 +16,11 @@
 
 """Concrete atomic actions built on :class:`AtomicAction` and :class:`TrajectoryBuilder`.
 
-Five sibling actions live here: :class:`MoveEndEffector`, :class:`MoveJoints`,
-:class:`PickUp`, :class:`MoveHeldObject`, and :class:`Place`. Each inherits
-:class:`AtomicAction` directly and composes a :class:`TrajectoryBuilder` for
-shared trajectory math. ``execute`` takes a typed target plus a
+Six sibling actions live here: :class:`MoveEndEffector`, :class:`MoveJoints`,
+:class:`PickUp`, :class:`MoveHeldObject`, :class:`Place`, and :class:`Press`.
+Each inherits :class:`AtomicAction` directly and composes a
+:class:`TrajectoryBuilder` for shared trajectory math. ``execute`` takes a typed
+target plus a
 :class:`WorldState` and returns an :class:`ActionResult` whose trajectory is
 full-robot DoF shaped ``(n_envs, n_waypoints, robot.dof)``.
 """
@@ -144,6 +145,24 @@ class PlaceCfg(ActionCfg):
     """Height (m) to retract the end-effector after opening the gripper."""
 
 
+@configclass
+class PressCfg(ActionCfg):
+    name: str = "press"
+    """Name of the action, used for identification and logging."""
+
+    sample_interval: int = 80
+    """Number of waypoints for the full trajectory (hand close + down + back)."""
+
+    hand_interp_steps: int = 5
+    """Number of waypoints for closing the gripper before pressing."""
+
+    hand_control_part: str = "hand"
+    """Name of the robot part that controls the hand joints."""
+
+    hand_close_qpos: torch.Tensor | None = None
+    """Joint positions for the closed hand state, shape ``[hand_dof,]``."""
+
+
 # =============================================================================
 # Shared helpers private to this module
 # =============================================================================
@@ -747,6 +766,135 @@ def _fail(self, state: WorldState) -> ActionResult:
         )
 
 
+# =============================================================================
+# Press
+# =============================================================================
+
+
+class Press(AtomicAction):
+    """Close the gripper, press down to a target pose, then return."""
+
+    TargetType: ClassVar[type] = EndEffectorPoseTarget
+
+    def __init__(
+        self,
+        motion_generator,
+        cfg: PressCfg | None = None,
+    ) -> None:
+        super().__init__(motion_generator, cfg or PressCfg())
+        self.builder = TrajectoryBuilder(motion_generator)
+        self.n_envs = self.robot.get_qpos().shape[0]
+        self.arm_joint_ids = self.robot.get_joint_ids(name=self.cfg.control_part)
+        self.hand_joint_ids = self.robot.get_joint_ids(name=self.cfg.hand_control_part)
+        self.arm_dof = len(self.arm_joint_ids)
+        self.hand_dof = len(self.hand_joint_ids)
+        self.robot_dof = self.robot.dof
+
+        if self.cfg.hand_close_qpos is None:
+            logger.log_error(
+                "hand_close_qpos must be specified in PressCfg", ValueError
+            )
+        self.hand_close_qpos = self.builder.expand_hand_qpos(
+            self.cfg.hand_close_qpos,
+            n_envs=self.n_envs,
+            hand_dof=self.hand_dof,
+        )
+
+    def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionResult:
+        press_xpos = self.builder.resolve_pose_target(target.xpos, n_envs=self.n_envs)
+        start_arm_qpos = self.builder.resolve_start_qpos(
+            _arm_qpos_from_state(state, self.arm_joint_ids, self.robot_dof),
+            n_envs=self.n_envs,
+            arm_dof=self.arm_dof,
+            control_part=self.cfg.control_part,
+        )
+        start_hand_qpos = state.last_qpos[:, self.hand_joint_ids]
+
+        n_close, n_down, n_back = self._compute_phase_waypoints()
+
+        # Phase 1: close the hand while holding the current EEF pose.
+        hand_close_path = self.builder.interpolate_hand_qpos(
+            start_hand_qpos,
+            self.hand_close_qpos,
+            n_waypoints=n_close,
+        )
+
+        # Phase 2: press down to the target pose.
+        target_states_list = [
+            [PlanState(xpos=press_xpos[i], move_type=MoveType.EEF_MOVE)]
+            for i in range(self.n_envs)
+        ]
+        ok, down_arm = self.builder.plan_arm_traj(
+            target_states_list,
+            start_arm_qpos,
+            n_down,
+            control_part=self.cfg.control_part,
+            arm_dof=self.arm_dof,
+        )
+        if not ok:
+            logger.log_warning("Press failed to plan the down trajectory.")
+            return self._fail(state)
+
+        # Phase 3: return to the arm pose from before pressing.
+        press_arm_qpos = down_arm[:, -1, :]
+        back_arm = self.builder.plan_joint_traj(press_arm_qpos, start_arm_qpos, n_back)
+
+        full = torch.empty(
+            (self.n_envs, n_close + n_down + n_back, self.robot_dof),
+            dtype=torch.float32,
+            device=self.device,
+        )
+        full[:, :, :] = state.last_qpos.unsqueeze(1)
+        full[:, :n_close, self.arm_joint_ids] = start_arm_qpos.unsqueeze(1)
+        full[:, :n_close, self.hand_joint_ids] = hand_close_path
+        full[:, n_close : n_close + n_down, self.arm_joint_ids] = down_arm
+        full[:, n_close : n_close + n_down, self.hand_joint_ids] = (
+            self.hand_close_qpos.unsqueeze(1)
+        )
+        full[:, n_close + n_down :, self.arm_joint_ids] = back_arm
+        full[:, n_close + n_down :, self.hand_joint_ids] = (
+            self.hand_close_qpos.unsqueeze(1)
+        )
+
+        return ActionResult(
+            success=True,
+            trajectory=full,
+            next_state=WorldState(
+                last_qpos=full[:, -1, :].clone(),
+                held_object=state.held_object,
+            ),
+        )
+
+    def _compute_phase_waypoints(self) -> tuple[int, int, int]:
+        n_close = self.cfg.hand_interp_steps
+        if n_close < 1:
+            logger.log_error(
+                "hand_interp_steps must be at least 1 for PressCfg.", ValueError
+            )
+
+        motion_waypoints = self.cfg.sample_interval - n_close
+        n_down = motion_waypoints // 2
+        n_back = motion_waypoints - n_down
+        if n_down < 2 or n_back < 2:
+            logger.log_error(
+                "Not enough waypoints for press trajectory. Increase "
+                "sample_interval or decrease hand_interp_steps.",
+                ValueError,
+            )
+        return n_close, n_down, n_back
+
+    def _fail(self, state: WorldState) -> ActionResult:
+        return ActionResult(
+            success=False,
+            trajectory=torch.empty(
+                (self.n_envs, 0, self.robot_dof),
+                dtype=torch.float32,
+                device=self.device,
+            ),
+            next_state=state,
+        )
+
+
 __all__ = [
     "MoveEndEffector",
     "MoveEndEffectorCfg",
@@ -758,4 +906,6 @@ def _fail(self, state: WorldState) -> ActionResult:
     "PickUpCfg",
     "Place",
     "PlaceCfg",
+    "Press",
+    "PressCfg",
 ]