Skip to content

Commit 2afbf62

Browse files
chernyadevchernyadev
authored
Add reset offset for full-body mode (#38)
1 parent 72d3054 commit 2afbf62

File tree

8 files changed

+151
-10
lines changed

8 files changed

+151
-10
lines changed

‎bigym/action_modes.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,14 @@ def step(self, action: np.ndarray):
9393
"""
9494
pass
9595

96+
@abstractmethod
97+
def reset(self, reset_state: np.ndarray):
98+
"""Reset state of the robot accordingly to the action mode.
99+
100+
:param reset_state: Target reset state of robot actuators.
101+
"""
102+
pass
103+
96104

97105
class TorqueActionMode(ActionMode):
98106
"""Control all joints through torque control.
@@ -140,6 +148,27 @@ def step(self, action: np.ndarray):
140148
self._robot.grippers[side].set_control(action)
141149
self._mojo.step()
142150

151+
def reset(self, reset_state: np.ndarray):
152+
"""See base."""
153+
if len(reset_state) != len(self._robot.limb_actuators):
154+
raise ValueError(
155+
f"Mismatch between reset_state length "
156+
f"({len(reset_state)}) "
157+
f"and number of actuators ({len(self._robot.limb_actuators)}). "
158+
f"Ensure reset_state matches the actuators count in the model."
159+
)
160+
for value, actuator in zip(reset_state, self._robot.limb_actuators):
161+
if actuator.joint:
162+
joint = self._mojo.physics.bind(actuator.joint)
163+
joint.qpos = value
164+
joint.qvel *= 0
165+
joint.qacc *= 0
166+
elif actuator.tendon:
167+
warnings.warn(
168+
f"Tendon actuators are not fully supported "
169+
f"for {self.__class__.__name__} action mode."
170+
)
171+
143172

144173
class JointPositionActionMode(ActionMode):
145174
"""Control all joints through joint position.
@@ -220,6 +249,36 @@ def step(self, action: np.ndarray):
220249
else:
221250
self._mojo.step()
222251

252+
def reset(self, reset_state: np.ndarray):
253+
"""See base."""
254+
if len(reset_state) != len(self._robot.limb_actuators):
255+
raise ValueError(
256+
f"Mismatch between reset_state length "
257+
f"({len(reset_state)}) "
258+
f"and number of actuators ({len(self._robot.limb_actuators)}). "
259+
f"Ensure reset_state matches the actuators count in the model."
260+
)
261+
for value, actuator in zip(reset_state, self._robot.limb_actuators):
262+
if actuator.joint:
263+
bound_joint = self._mojo.physics.bind(actuator.joint)
264+
bound_joint.qpos = value
265+
bound_joint.qvel *= 0
266+
bound_joint.qacc *= 0
267+
elif actuator.tendon:
268+
if actuator.tendon.joint is None or len(actuator.tendon.joint) == 0:
269+
raise RuntimeError(
270+
"Currently only fixed tendons with joints are supported."
271+
)
272+
joint_value = value / len(actuator.tendon.joint)
273+
for tendon_joint in actuator.tendon.joint:
274+
value_coefficient = tendon_joint.coef
275+
bound_joint = self._mojo.physics.bind(tendon_joint.joint)
276+
bound_joint.qpos = joint_value / value_coefficient
277+
bound_joint.qvel *= 0
278+
bound_joint.qacc *= 0
279+
bound_actuator = self._mojo.physics.bind(actuator)
280+
bound_actuator.ctrl = value
281+
223282
def _step_until_reached(self):
224283
"""Step physics until the target position is reached."""
225284
steps_counter = 0

‎bigym/bigym_env.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -401,7 +401,7 @@ def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None):
401401
self._update_seed(override_seed=seed)
402402
self._mojo.physics.reset()
403403
self._action = np.zeros_like(self._action)
404-
self._robot.set_pose(self.RESET_ROBOT_POS, self.RESET_ROBOT_QUAT)
404+
self._robot.reset(self.RESET_ROBOT_POS, self.RESET_ROBOT_QUAT)
405405
self._on_reset()
406406
return self.get_observation(), self.get_info()
407407

‎bigym/robots/config.py‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,8 @@ class ArmConfig:
4848
site: The site on the robot where the gripper could be attached.
4949
links: A list of body links of the hand.
5050
wrist_dof: Optional wrist DOF which could be added to the arm.
51-
offset_position: Mounting positional offset.
52-
offset_euler: Mounting euler offset.
51+
offset_position: Gripper mounting positional offset.
52+
offset_euler: Gripper mounting euler offset.
5353
"""
5454

5555
site: str
@@ -67,9 +67,18 @@ class FloatingBaseConfig:
6767
delta_range_position: tuple[float, float]
6868
delta_range_rotation: tuple[float, float]
6969
offset_position: np.ndarray = field(default_factory=lambda: np.zeros(3))
70+
reset_state: Optional[np.ndarray] = None
7071
animated_legs_class: Optional[Type[AnimatedLegs]] = None
7172

7273

74+
@dataclass
75+
class FullBodyConfig:
76+
"""Configuration for full-body mode."""
77+
78+
offset_position: np.ndarray = field(default_factory=lambda: np.zeros(3))
79+
reset_state: Optional[np.ndarray] = None
80+
81+
7382
@dataclass
7483
class RobotConfig:
7584
"""Configuration for a robot.
@@ -79,7 +88,8 @@ class RobotConfig:
7988
delta_range: Action range for delta position action mode.
8089
position_kp: Stiffness of actuators for absolute position action mode.
8190
pelvis_body: Name of the pelvis body element.
82-
floating_base: Configuration for the robot's floating base.
91+
full_body: Config for full-body mode.
92+
floating_base: Configuration for the floating base.
8393
gripper: Configuration for the robot's gripper.
8494
arms: Configuration for the robot's hands.
8595
actuators: Dictionary containing all actuators
@@ -92,6 +102,7 @@ class RobotConfig:
92102
delta_range: tuple[float, float]
93103
position_kp: float
94104
pelvis_body: str
105+
full_body: FullBodyConfig
95106
floating_base: FloatingBaseConfig
96107
gripper: GripperConfig
97108
arms: dict[HandSide, ArmConfig]

‎bigym/robots/configs/google_robot.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
FloatingBaseConfig,
1010
GripperConfig,
1111
RobotConfig,
12+
FullBodyConfig,
1213
)
1314
from bigym.robots.robot import Robot
1415
from bigym.utils.dof import Dof
@@ -35,6 +36,10 @@
3536
],
3637
)
3738
STIFFNESS = 1e4
39+
GOOGLE_ROBOT_FULL_BODY = FullBodyConfig(
40+
offset_position=np.zeros(3),
41+
reset_state=None,
42+
)
3843
GOOGLE_ROBOT_FLOATING_BASE = FloatingBaseConfig(
3944
dofs={
4045
PelvisDof.X: Dof(
@@ -69,6 +74,7 @@
6974
delta_range=(-0.1, 0.1),
7075
position_kp=300,
7176
pelvis_body="base_link",
77+
full_body=GOOGLE_ROBOT_FULL_BODY,
7278
floating_base=GOOGLE_ROBOT_FLOATING_BASE,
7379
gripper=GOOGLE_ROBOT_GRIPPER,
7480
arms={HandSide.RIGHT: GOOGLE_ROBOT_HAND},

‎bigym/robots/configs/h1.py‎

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,12 @@
66
from bigym.const import ASSETS_PATH, HandSide
77
from bigym.robots.animated_legs import H1AnimatedLegs
88

9-
from bigym.robots.config import ArmConfig, FloatingBaseConfig, RobotConfig
9+
from bigym.robots.config import (
10+
ArmConfig,
11+
FloatingBaseConfig,
12+
RobotConfig,
13+
FullBodyConfig,
14+
)
1015
from bigym.robots.configs.robotiq import ROBOTIQ_2F85, ROBOTIQ_2F85_FINE_MANIPULATION
1116
from bigym.robots.robot import Robot
1217
from bigym.utils.dof import Dof
@@ -65,6 +70,34 @@
6570
STIFFNESS_XY = 1e4
6671
STIFFNESS_Z = 1e6
6772
RANGE_DOF_Z = (0.4, 1.0)
73+
H1_FULL_BODY = FullBodyConfig(
74+
offset_position=np.array([0, 0, 0.98]),
75+
reset_state=np.array(
76+
[
77+
0, # left_hip_yaw
78+
0,
79+
-0.4,
80+
0.8,
81+
-0.4,
82+
0, # right_hip_yaw
83+
0,
84+
-0.4,
85+
0.8,
86+
-0.4,
87+
0, # torso
88+
0, # left_shoulder_pitch
89+
0,
90+
0,
91+
0,
92+
0,
93+
0, # right_shoulder_pitch
94+
0,
95+
0,
96+
0,
97+
0,
98+
]
99+
),
100+
)
68101
H1_FLOATING_BASE = FloatingBaseConfig(
69102
dofs={
70103
PelvisDof.X: Dof(
@@ -93,13 +126,28 @@
93126
delta_range_position=(-0.01, 0.01),
94127
delta_range_rotation=(-0.05, 0.05),
95128
offset_position=np.array([0, 0, 1]),
129+
reset_state=np.array(
130+
[
131+
0, # left_shoulder_pitch
132+
0,
133+
0,
134+
0,
135+
0,
136+
0, # right_shoulder_pitch
137+
0,
138+
0,
139+
0,
140+
0,
141+
]
142+
),
96143
animated_legs_class=H1AnimatedLegs,
97144
)
98145
H1_CONFIG = RobotConfig(
99146
model=ASSETS_PATH / "h1/h1.xml",
100147
delta_range=(-0.1, 0.1),
101148
position_kp=300,
102149
pelvis_body="pelvis",
150+
full_body=H1_FULL_BODY,
103151
floating_base=H1_FLOATING_BASE,
104152
gripper=ROBOTIQ_2F85,
105153
arms={HandSide.LEFT: H1_LEFT_ARM, HandSide.RIGHT: H1_RIGHT_ARM},
@@ -112,6 +160,7 @@
112160
delta_range=(-0.1, 0.1),
113161
position_kp=300,
114162
pelvis_body="pelvis",
163+
full_body=H1_FULL_BODY,
115164
floating_base=H1_FLOATING_BASE,
116165
gripper=ROBOTIQ_2F85_FINE_MANIPULATION,
117166
arms={HandSide.LEFT: H1_LEFT_ARM, HandSide.RIGHT: H1_RIGHT_ARM},

‎bigym/robots/configs/stretch.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
FloatingBaseConfig,
1010
GripperConfig,
1111
RobotConfig,
12+
FullBodyConfig,
1213
)
1314
from bigym.robots.robot import Robot
1415
from bigym.utils.dof import Dof
@@ -36,6 +37,10 @@
3637
],
3738
)
3839
STIFFNESS = 1e4
40+
STRETCH_FULL_BODY = FullBodyConfig(
41+
offset_position=np.zeros(3),
42+
reset_state=None,
43+
)
3944
STRETCH_FLOATING_BASE = FloatingBaseConfig(
4045
dofs={
4146
PelvisDof.X: Dof(
@@ -56,6 +61,7 @@
5661
},
5762
delta_range_position=(-0.01, 0.01),
5863
delta_range_rotation=(-0.05, 0.05),
64+
reset_state=np.array([0, 0, 0, 0, 0, 0]),
5965
)
6066
STRETCH_GRIPPER = GripperConfig(
6167
body="stretch/link_gripper_slider",
@@ -68,6 +74,7 @@
6874
delta_range=(-0.1, 0.1),
6975
position_kp=300,
7076
pelvis_body="base_link",
77+
full_body=STRETCH_FULL_BODY,
7178
floating_base=STRETCH_FLOATING_BASE,
7279
gripper=STRETCH_GRIPPER,
7380
arms={HandSide.RIGHT: STRETCH_HAND},

‎bigym/robots/floating_base.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,7 @@ def __init__(
3636
self._config = config
3737
self._pelvis = pelvis
3838
self._mojo = mojo
39-
self._offset_position = np.array(self._config.offset_position).astype(
40-
np.float32
41-
)
39+
self._offset_position = np.array(self._config.offset_position)
4240
self._position_actuators: list[Optional[mjcf.Element]] = [None, None, None]
4341
self._rotation_actuators: list[Optional[mjcf.Element]] = [None, None, None]
4442

‎bigym/robots/robot.py‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -156,12 +156,23 @@ def is_gripper_holding_object(
156156
return False
157157
return self._grippers[side].is_holding_object(other)
158158

159-
def set_pose(self, position: np.ndarray, orientation: np.ndarray):
159+
def reset(self, position: np.ndarray, orientation: np.ndarray):
160+
"""Reset robot."""
161+
self._set_pose(position, orientation)
162+
reset_state = (
163+
self.config.floating_base.reset_state
164+
if self._action_mode.floating_base
165+
else self.config.full_body.reset_state
166+
)
167+
if reset_state is not None:
168+
self._action_mode.reset(reset_state)
169+
170+
def _set_pose(self, position: np.ndarray, orientation: np.ndarray):
160171
"""Instantly set pose of the robot pelvis."""
161172
if self._action_mode.floating_base:
162173
self._floating_base.reset(position, orientation)
163174
else:
164-
self._pelvis.set_position(position)
175+
self._pelvis.set_position(position + self.config.full_body.offset_position)
165176
self._pelvis.set_quaternion(orientation)
166177

167178
def get_limb_control_range(

0 commit comments

Comments
 (0)