From 8f47f2ba30823238674769fd4d103babd87c0d36 Mon Sep 17 00:00:00 2001
From: ZidongChen25 <tsztungchen25@gmail.com>
Date: Mon, 23 Feb 2026 13:34:48 +0000
Subject: [PATCH] Add gripper blocking option and enable default wait behavior

---
 bigym/action_modes.py   | 27 ++++++++++++++++++++-------
 bigym/robots/gripper.py | 12 ++++++++++--
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/bigym/action_modes.py b/bigym/action_modes.py
index b432b21..847c87c 100644
--- a/bigym/action_modes.py
+++ b/bigym/action_modes.py
@@ -189,6 +189,8 @@ def __init__(
         self,
         absolute: bool = False,
         block_until_reached: bool = False,
+        block_gripper_until_reached: bool = True,
+        gripper_tolerance: float = 0.05,
         floating_base: bool = True,
         floating_dofs: list[PelvisDof] = None,
     ):
@@ -197,6 +199,10 @@ def __init__(
         :param absolute: Use absolute or delta joint positions.
         :param block_until_reached: Continue stepping until the target
             position is reached or the step threshold is exceeded.
+        :param block_gripper_until_reached: Continue stepping until gripper
+            target state is reached or the step threshold is exceeded.
+        :param gripper_tolerance: Allowed gripper absolute error in normalized
+            gripper range before considering the target reached.
         """
         super().__init__(
             floating_base=floating_base,
@@ -204,6 +210,8 @@ def __init__(
         )
         self.absolute = absolute
         self.block_until_reached = block_until_reached
+        self.block_gripper_until_reached = block_gripper_until_reached
+        self.gripper_tolerance = gripper_tolerance
 
     def action_space(
         self, action_scale: float, seed: Optional[int] = None
@@ -244,7 +252,7 @@ def step(self, action: np.ndarray):
         for side, action in zip(self._robot.grippers, gripper_actions):
             self._robot.grippers[side].set_control(action)
 
-        if self.block_until_reached:
+        if self.block_until_reached or self.block_gripper_until_reached:
             self._step_until_reached()
         else:
             self._mojo.step()
@@ -294,10 +302,15 @@ def _step_until_reached(self):
                 break
 
     def _is_target_state_reached(self):
-        if self.floating_base:
-            if not self._robot.floating_base.is_target_reached:
-                return False
-        for actuator in self._robot.limb_actuators:
-            if not is_target_reached(actuator, self._mojo.physics, TOLERANCE_ANGULAR):
-                return False
+        if self.block_until_reached:
+            if self.floating_base:
+                if not self._robot.floating_base.is_target_reached:
+                    return False
+            for actuator in self._robot.limb_actuators:
+                if not is_target_reached(actuator, self._mojo.physics, TOLERANCE_ANGULAR):
+                    return False
+        if self.block_gripper_until_reached:
+            for _, gripper in self._robot.grippers.items():
+                if not gripper.is_target_reached(self.gripper_tolerance):
+                    return False
         return True
diff --git a/bigym/robots/gripper.py b/bigym/robots/gripper.py
index 58ec647..44cb8a5 100644
--- a/bigym/robots/gripper.py
+++ b/bigym/robots/gripper.py
@@ -38,6 +38,7 @@ def __init__(
         self._actuators: list[mjcf.Element] = []
         self._actuated_joints: list[mjcf.Element] = []
         self._pinch_site: Optional[Site] = None
+        self._target_ctrl: Optional[float] = None
 
         if self._config.model:
             self._body: Body = self._mojo.load_model(
@@ -115,10 +116,17 @@ def set_control(self, ctrl: float):
         if self._config.discrete:
             ctrl = np.round(ctrl)
         ctrl = np.interp(ctrl, self._NORMAL_RANGE, self._config.range)
+        self._target_ctrl = float(ctrl)
         for actuator in self._actuators:
             actuator = self._mojo.physics.bind(actuator)
-            ctrl = np.interp(ctrl, self._config.range, actuator.ctrlrange)
-            actuator.ctrl = ctrl
+            actuator_ctrl = np.interp(ctrl, self._config.range, actuator.ctrlrange)
+            actuator.ctrl = actuator_ctrl
+
+    def is_target_reached(self, tolerance: float) -> bool:
+        """Check if the gripper is sufficiently close to the last control target."""
+        if self._target_ctrl is None:
+            return True
+        return abs(self.qpos - self._target_ctrl) <= tolerance
 
     def _on_loaded(self, model: mjcf.RootElement):
         model.model += f"_{self._side.value.lower()}"