Fix environment reset for full-body mode

bigym/action_modes.py

@@ -93,6 +93,14 @@ def step(self, action: np.ndarray):
         """
         pass
 
+    @abstractmethod
+    def reset(self, reset_state: np.ndarray):
+        """Reset state of the robot accordingly to the action mode.
+
+        :param reset_state: Target reset state of robot actuators.
+        """
+        pass
+
 
 class TorqueActionMode(ActionMode):
     """Control all joints through torque control.
@@ -140,6 +148,27 @@ def step(self, action: np.ndarray):
             self._robot.grippers[side].set_control(action)
         self._mojo.step()
 
+    def reset(self, reset_state: np.ndarray):
+        """See base."""
+        if len(reset_state) != len(self._robot.limb_actuators):
+            raise ValueError(
+                f"Mismatch between reset_state length "
+                f"({len(reset_state)}) "
+                f"and number of actuators ({len(self._robot.limb_actuators)}). "
+                f"Ensure reset_state matches the actuators count in the model."
+            )
+        for value, actuator in zip(reset_state, self._robot.limb_actuators):
+            if actuator.joint:
+                joint = self._mojo.physics.bind(actuator.joint)
+                joint.qpos = value
+                joint.qvel *= 0
+                joint.qacc *= 0
+            elif actuator.tendon:
+                warnings.warn(
+                    f"Tendon actuators are not fully supported "
+                    f"for {self.__class__.__name__} action mode."
+                )
+
 
 class JointPositionActionMode(ActionMode):
     """Control all joints through joint position.
@@ -220,6 +249,36 @@ def step(self, action: np.ndarray):
         else:
             self._mojo.step()
 
+    def reset(self, reset_state: np.ndarray):
+        """See base."""
+        if len(reset_state) != len(self._robot.limb_actuators):
+            raise ValueError(
+                f"Mismatch between reset_state length "
+                f"({len(reset_state)}) "
+                f"and number of actuators ({len(self._robot.limb_actuators)}). "
+                f"Ensure reset_state matches the actuators count in the model."
+            )
+        for value, actuator in zip(reset_state, self._robot.limb_actuators):
+            if actuator.joint:
+                bound_joint = self._mojo.physics.bind(actuator.joint)
+                bound_joint.qpos = value
+                bound_joint.qvel *= 0
+                bound_joint.qacc *= 0
+            elif actuator.tendon:
+                if actuator.tendon.joint is None or len(actuator.tendon.joint) == 0:
+                    raise RuntimeError(
+                        "Currently only fixed tendons with joints are supported."
+                    )
+                joint_value = value / len(actuator.tendon.joint)
+                for tendon_joint in actuator.tendon.joint:
+                    value_coefficient = tendon_joint.coef
+                    bound_joint = self._mojo.physics.bind(tendon_joint.joint)
+                    bound_joint.qpos = joint_value / value_coefficient
+                    bound_joint.qvel *= 0
+                    bound_joint.qacc *= 0
+            bound_actuator = self._mojo.physics.bind(actuator)
+            bound_actuator.ctrl = value
+
     def _step_until_reached(self):
         """Step physics until the target position is reached."""
         steps_counter = 0

bigym/bigym_env.py

@@ -401,7 +401,7 @@ def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None):
         self._update_seed(override_seed=seed)
         self._mojo.physics.reset()
         self._action = np.zeros_like(self._action)
-        self._robot.set_pose(self.RESET_ROBOT_POS, self.RESET_ROBOT_QUAT)
+        self._robot.reset(self.RESET_ROBOT_POS, self.RESET_ROBOT_QUAT)
         self._on_reset()
         return self.get_observation(), self.get_info()
 

bigym/robots/config.py

@@ -48,8 +48,8 @@ class ArmConfig:
         site: The site on the robot where the gripper could be attached.
         links: A list of body links of the hand.
         wrist_dof: Optional wrist DOF which could be added to the arm.
-        offset_position: Mounting positional offset.
-        offset_euler: Mounting euler offset.
+        offset_position: Gripper mounting positional offset.
+        offset_euler: Gripper mounting euler offset.
     """
 
     site: str
@@ -67,9 +67,18 @@ class FloatingBaseConfig:
     delta_range_position: tuple[float, float]
     delta_range_rotation: tuple[float, float]
     offset_position: np.ndarray = field(default_factory=lambda: np.zeros(3))
+    reset_state: Optional[np.ndarray] = None
     animated_legs_class: Optional[Type[AnimatedLegs]] = None
 
 
+@dataclass
+class FullBodyConfig:
+    """Configuration for full-body mode."""
+
+    offset_position: np.ndarray = field(default_factory=lambda: np.zeros(3))
+    reset_state: Optional[np.ndarray] = None
+
+
 @dataclass
 class RobotConfig:
     """Configuration for a robot.
@@ -79,7 +88,8 @@ class RobotConfig:
         delta_range: Action range for delta position action mode.
         position_kp: Stiffness of actuators for absolute position action mode.
         pelvis_body: Name of the pelvis body element.
-        floating_base: Configuration for the robot's floating base.
+        full_body: Config for full-body mode.
+        floating_base: Configuration for the floating base.
         gripper: Configuration for the robot's gripper.
         arms: Configuration for the robot's hands.
         actuators: Dictionary containing all actuators
@@ -92,6 +102,7 @@ class RobotConfig:
     delta_range: tuple[float, float]
     position_kp: float
     pelvis_body: str
+    full_body: FullBodyConfig
     floating_base: FloatingBaseConfig
     gripper: GripperConfig
     arms: dict[HandSide, ArmConfig]

bigym/robots/configs/google_robot.py

@@ -9,6 +9,7 @@
     FloatingBaseConfig,
     GripperConfig,
     RobotConfig,
+    FullBodyConfig,
 )
 from bigym.robots.robot import Robot
 from bigym.utils.dof import Dof
@@ -35,6 +36,10 @@
     ],
 )
 STIFFNESS = 1e4
+GOOGLE_ROBOT_FULL_BODY = FullBodyConfig(
+    offset_position=np.zeros(3),
+    reset_state=None,
+)
 GOOGLE_ROBOT_FLOATING_BASE = FloatingBaseConfig(
     dofs={
         PelvisDof.X: Dof(
@@ -69,6 +74,7 @@
     delta_range=(-0.1, 0.1),
     position_kp=300,
     pelvis_body="base_link",
+    full_body=GOOGLE_ROBOT_FULL_BODY,
     floating_base=GOOGLE_ROBOT_FLOATING_BASE,
     gripper=GOOGLE_ROBOT_GRIPPER,
     arms={HandSide.RIGHT: GOOGLE_ROBOT_HAND},

bigym/robots/configs/h1.py

@@ -6,7 +6,12 @@
 from bigym.const import ASSETS_PATH, HandSide
 from bigym.robots.animated_legs import H1AnimatedLegs
 
-from bigym.robots.config import ArmConfig, FloatingBaseConfig, RobotConfig
+from bigym.robots.config import (
+    ArmConfig,
+    FloatingBaseConfig,
+    RobotConfig,
+    FullBodyConfig,
+)
 from bigym.robots.configs.robotiq import ROBOTIQ_2F85, ROBOTIQ_2F85_FINE_MANIPULATION
 from bigym.robots.robot import Robot
 from bigym.utils.dof import Dof
@@ -65,6 +70,34 @@
 STIFFNESS_XY = 1e4
 STIFFNESS_Z = 1e6
 RANGE_DOF_Z = (0.4, 1.0)
+H1_FULL_BODY = FullBodyConfig(
+    offset_position=np.array([0, 0, 0.98]),
+    reset_state=np.array(
+        [
+            0,  # left_hip_yaw
+            0,
+            -0.4,
+            0.8,
+            -0.4,
+            0,  # right_hip_yaw
+            0,
+            -0.4,
+            0.8,
+            -0.4,
+            0,  # torso
+            0,  # left_shoulder_pitch
+            0,
+            0,
+            0,
+            0,
+            0,  # right_shoulder_pitch
+            0,
+            0,
+            0,
+            0,
+        ]
+    ),
+)
 H1_FLOATING_BASE = FloatingBaseConfig(
     dofs={
         PelvisDof.X: Dof(
@@ -93,13 +126,28 @@
     delta_range_position=(-0.01, 0.01),
     delta_range_rotation=(-0.05, 0.05),
     offset_position=np.array([0, 0, 1]),
+    reset_state=np.array(
+        [
+            0,  # left_shoulder_pitch
+            0,
+            0,
+            0,
+            0,
+            0,  # right_shoulder_pitch
+            0,
+            0,
+            0,
+            0,
+        ]
+    ),
     animated_legs_class=H1AnimatedLegs,
 )
 H1_CONFIG = RobotConfig(
     model=ASSETS_PATH / "h1/h1.xml",
     delta_range=(-0.1, 0.1),
     position_kp=300,
     pelvis_body="pelvis",
+    full_body=H1_FULL_BODY,
     floating_base=H1_FLOATING_BASE,
     gripper=ROBOTIQ_2F85,
     arms={HandSide.LEFT: H1_LEFT_ARM, HandSide.RIGHT: H1_RIGHT_ARM},
@@ -112,6 +160,7 @@
     delta_range=(-0.1, 0.1),
     position_kp=300,
     pelvis_body="pelvis",
+    full_body=H1_FULL_BODY,
     floating_base=H1_FLOATING_BASE,
     gripper=ROBOTIQ_2F85_FINE_MANIPULATION,
     arms={HandSide.LEFT: H1_LEFT_ARM, HandSide.RIGHT: H1_RIGHT_ARM},

bigym/robots/configs/stretch.py

@@ -9,6 +9,7 @@
     FloatingBaseConfig,
     GripperConfig,
     RobotConfig,
+    FullBodyConfig,
 )
 from bigym.robots.robot import Robot
 from bigym.utils.dof import Dof
@@ -36,6 +37,10 @@
     ],
 )
 STIFFNESS = 1e4
+STRETCH_FULL_BODY = FullBodyConfig(
+    offset_position=np.zeros(3),
+    reset_state=None,
+)
 STRETCH_FLOATING_BASE = FloatingBaseConfig(
     dofs={
         PelvisDof.X: Dof(
@@ -56,6 +61,7 @@
     },
     delta_range_position=(-0.01, 0.01),
     delta_range_rotation=(-0.05, 0.05),
+    reset_state=np.array([0, 0, 0, 0, 0, 0]),
 )
 STRETCH_GRIPPER = GripperConfig(
     body="stretch/link_gripper_slider",
@@ -68,6 +74,7 @@
     delta_range=(-0.1, 0.1),
     position_kp=300,
     pelvis_body="base_link",
+    full_body=STRETCH_FULL_BODY,
     floating_base=STRETCH_FLOATING_BASE,
     gripper=STRETCH_GRIPPER,
     arms={HandSide.RIGHT: STRETCH_HAND},

bigym/robots/floating_base.py

@@ -36,9 +36,7 @@ def __init__(
         self._config = config
         self._pelvis = pelvis
         self._mojo = mojo
-        self._offset_position = np.array(self._config.offset_position).astype(
-            np.float32
-        )
+        self._offset_position = np.array(self._config.offset_position)
         self._position_actuators: list[Optional[mjcf.Element]] = [None, None, None]
         self._rotation_actuators: list[Optional[mjcf.Element]] = [None, None, None]
 

bigym/robots/robot.py

@@ -156,12 +156,23 @@ def is_gripper_holding_object(
             return False
         return self._grippers[side].is_holding_object(other)
 
-    def set_pose(self, position: np.ndarray, orientation: np.ndarray):
+    def reset(self, position: np.ndarray, orientation: np.ndarray):
+        """Reset robot."""
+        self._set_pose(position, orientation)
+        reset_state = (
+            self.config.floating_base.reset_state
+            if self._action_mode.floating_base
+            else self.config.full_body.reset_state
+        )
+        if reset_state is not None:
+            self._action_mode.reset(reset_state)
+
+    def _set_pose(self, position: np.ndarray, orientation: np.ndarray):
         """Instantly set pose of the robot pelvis."""
         if self._action_mode.floating_base:
             self._floating_base.reset(position, orientation)
         else:
-            self._pelvis.set_position(position)
+            self._pelvis.set_position(position + self.config.full_body.offset_position)
             self._pelvis.set_quaternion(orientation)
 
     def get_limb_control_range(