diff --git a/bigym/action_modes.py b/bigym/action_modes.py index 6823e78..b432b21 100644 --- a/bigym/action_modes.py +++ b/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 diff --git a/bigym/bigym_env.py b/bigym/bigym_env.py index a909a28..03c6fc8 100644 --- a/bigym/bigym_env.py +++ b/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() diff --git a/bigym/robots/config.py b/bigym/robots/config.py index 4ee6183..67ec695 100644 --- a/bigym/robots/config.py +++ b/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] diff --git a/bigym/robots/configs/google_robot.py b/bigym/robots/configs/google_robot.py index 8cfc83b..79bac9d 100644 --- a/bigym/robots/configs/google_robot.py +++ b/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}, diff --git a/bigym/robots/configs/h1.py b/bigym/robots/configs/h1.py index eaa6107..698eef6 100644 --- a/bigym/robots/configs/h1.py +++ b/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,6 +126,20 @@ 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( @@ -100,6 +147,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, 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}, diff --git a/bigym/robots/configs/stretch.py b/bigym/robots/configs/stretch.py index a4e6bb0..14af637 100644 --- a/bigym/robots/configs/stretch.py +++ b/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}, diff --git a/bigym/robots/floating_base.py b/bigym/robots/floating_base.py index bcc990e..f3dc223 100644 --- a/bigym/robots/floating_base.py +++ b/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] diff --git a/bigym/robots/robot.py b/bigym/robots/robot.py index 5334678..660def6 100644 --- a/bigym/robots/robot.py +++ b/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(