Skip to content

Commit ac2ec9c

Browse files
tenfoldpapertenfoldpaper
committed
feat: arm mirroring functions
1 parent d53d078 commit ac2ec9c

1 file changed

Lines changed: 136 additions & 1 deletion

File tree

‎python/rcs/sim/replayer.py‎

Lines changed: 136 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,9 @@
1010
from rcs.envs.scenes import SimEnvCreator
1111
from rcs.envs.storage_wrapper import StorageWrapper
1212
from rcs.sim.sim import RAW_STATE_ENCODING
13-
13+
from rcs import DEFAULT_TRANSFORMS, ROBOTS, GRIPPER_OFFSETS
14+
from rcs import common
15+
import rcs_duobench.tasks.hinge_chest
1416

1517
def _normalize_sim_state_schema(value: Any) -> SimStateSchema:
1618
joint_names = [str(item) for item in value["joint_names"]]
@@ -22,6 +24,139 @@ def _normalize_sim_state_schema(value: Any) -> SimStateSchema:
2224
"encodings": [str(item) for item in value.get("encodings", [RAW_STATE_ENCODING] * len(joint_names))],
2325
}
2426

27+
def _mirror_arms_real(rows):
28+
prev_qpos = {
29+
'left': None,
30+
'right': None,
31+
}
32+
SHARED_TO_ROBOT = {
33+
'left': DEFAULT_TRANSFORMS["FR3_DUOMOUNT_LEFT_ROBOT"],
34+
'right': DEFAULT_TRANSFORMS["FR3_DUOMOUNT_RIGHT_ROBOT"],
35+
}
36+
ROBOT_TO_SHARED = {
37+
'left': SHARED_TO_ROBOT['left'].inverse(),
38+
'right': SHARED_TO_ROBOT['right'].inverse(),
39+
}
40+
PAIRS = {
41+
'left': 'right',
42+
'right': 'left',
43+
}
44+
for row in rows:
45+
for black in ['left', 'right']: # black and white to distinguish side vs mirror
46+
# First fetch the joint data in observation
47+
white_joint_obs = row[3][PAIRS[black]]['joints'].copy()
48+
black_tquat_obs = row[3][black]['tquat'].copy()
49+
50+
# Mirror the shared frame tquat along the XZ plane
51+
mirrored_black_pose = black_tquat_obs.copy()
52+
mirrored_black_pose[1] = -mirrored_black_pose[1]
53+
mirrored_black_pose[3] = -mirrored_black_pose[3]
54+
mirrored_black_pose[5] = -mirrored_black_pose[5]
55+
56+
# Place the mirrored pose in the white robot's frame
57+
white_target_pose = common.Pose(translation=mirrored_black_pose[:3], quaternion=mirrored_black_pose[3:])
58+
white_target_pose = ROBOT_TO_SHARED[PAIRS[black]] * white_target_pose
59+
60+
if prev_qpos[PAIRS[black]] is None:
61+
prev_qpos[PAIRS[black]] = white_joint_obs
62+
63+
# Calculate the IK for the white robot to achieve the mirrored pose
64+
robot_type = row[4][black]['robot_type']
65+
gripper_type = row[4][black]['gripper_type']
66+
tcp_offset = GRIPPER_OFFSETS[gripper_type]
67+
ik = common.Pin(
68+
ROBOTS[robot_type].mjcf_model_path,
69+
ROBOTS[robot_type].attachment_site,
70+
)
71+
white_target_qpos = ik.inverse(white_target_pose, prev_qpos[PAIRS[black]], tcp_offset=tcp_offset)
72+
prev_qpos[PAIRS[black]] = white_target_qpos
73+
# Write arm - make a copy of black and write to white_mirrored
74+
row[3][f'{PAIRS[black]}_mirrored'] = row[3][black].copy()
75+
row[3][f'{PAIRS[black]}_mirrored']['joints'] = white_target_qpos
76+
white_target_tquat = np.zeros(7)
77+
white_target_tquat[:3] = white_target_pose.translation()
78+
white_target_tquat[3:] = white_target_pose.rotation_q()
79+
row[3][f'{PAIRS[black]}_mirrored']['tquat'] = white_target_tquat
80+
return rows
81+
82+
def _mirror_arms_sim(rows):
83+
prev_qpos = {
84+
'left': None,
85+
'right': None,
86+
}
87+
SHARED_TO_ROBOT = {
88+
'left': DEFAULT_TRANSFORMS["FR3_DUOMOUNT_LEFT_ROBOT"],
89+
'right': DEFAULT_TRANSFORMS["FR3_DUOMOUNT_RIGHT_ROBOT"],
90+
}
91+
ROBOT_TO_SHARED = {
92+
'left': SHARED_TO_ROBOT['left'].inverse(),
93+
'right': SHARED_TO_ROBOT['right'].inverse(),
94+
}
95+
PAIRS = {
96+
'left': 'right',
97+
'right': 'left',
98+
}
99+
def get_qpos_start_and_size(joint_name_order, joint_qpos_sizes, target_prefix):
100+
start_index = None
101+
joint_size = 0
102+
for i, joint_name in enumerate(joint_name_order):
103+
if target_prefix in joint_name:
104+
if start_index is None:
105+
start_index = sum(joint_qpos_sizes[:i])
106+
joint_size += joint_qpos_sizes[i]
107+
return start_index, joint_size
108+
for row in rows:
109+
for black in ['left', 'right']: # black and white to distinguish side vs mirror
110+
# First fetch the joint data in observation
111+
white_joint_obs = row[3][PAIRS[black]]['joints'].copy()
112+
black_tquat_obs = row[3][black]['tquat'].copy()
113+
114+
# Identify which indices to write the mirrored joint data
115+
sim_state = row[4][black]['sim_state']
116+
sim_state_schema = row[4][black]['sim_state_schema']
117+
joint_name_order = sim_state_schema['joint_names']
118+
joint_qpos_sizes = sim_state_schema['qpos_sizes']
119+
white_robot_start_index, white_robot_joint_size = get_qpos_start_and_size(joint_name_order, joint_qpos_sizes, f"{PAIRS[black]}_fr3_joint")
120+
121+
# Get the gripper data for the current hand
122+
black_gripper_start_index, black_gripper_joint_size = get_qpos_start_and_size(joint_name_order, joint_qpos_sizes, f"gripper{black}")
123+
black_gripper_qpos = sim_state[black_gripper_start_index:black_gripper_start_index + black_gripper_joint_size]
124+
125+
# Get the gripper index range for the other hand
126+
white_gripper_start_index, white_gripper_joint_size = get_qpos_start_and_size(joint_name_order, joint_qpos_sizes, f"gripper{PAIRS[black]}")
127+
128+
# Mirror the shared frame tquat along the XZ plane
129+
mirrored_black_pose = black_tquat_obs.copy()
130+
mirrored_black_pose[1] = -mirrored_black_pose[1]
131+
mirrored_black_pose[3] = -mirrored_black_pose[3]
132+
mirrored_black_pose[5] = -mirrored_black_pose[5]
133+
134+
# Place the mirrored pose in the white robot's frame
135+
white_target_pose = common.Pose(translation=mirrored_black_pose[:3], quaternion=mirrored_black_pose[3:])
136+
white_target_pose = ROBOT_TO_SHARED[PAIRS[black]] * white_target_pose
137+
138+
if prev_qpos[PAIRS[black]] is None:
139+
prev_qpos[PAIRS[black]] = white_joint_obs
140+
141+
# Calculate the IK for the white robot to achieve the mirrored pose
142+
robot_type = row[4][black]['robot_type']
143+
gripper_type = row[4][black]['gripper_type']
144+
tcp_offset = GRIPPER_OFFSETS[gripper_type]
145+
ik = common.Pin(
146+
ROBOTS[robot_type].mjcf_model_path,
147+
ROBOTS[robot_type].attachment_site,
148+
)
149+
white_target_qpos = ik.inverse(white_target_pose, prev_qpos[PAIRS[black]], tcp_offset=tcp_offset)
150+
prev_qpos[PAIRS[black]] = white_target_qpos
151+
# Write arm
152+
row[4][black]['sim_state'][white_robot_start_index:white_robot_start_index+white_robot_joint_size] = white_target_qpos
153+
row[4][PAIRS[black]]['sim_state'][white_robot_start_index:white_robot_start_index+white_robot_joint_size] = white_target_qpos
154+
155+
# Write gripper
156+
row[4][black]['sim_state'][white_gripper_start_index:white_gripper_start_index + black_gripper_joint_size] = black_gripper_qpos
157+
row[4][PAIRS[black]]['sim_state'][white_gripper_start_index:white_gripper_start_index + white_gripper_joint_size] = black_gripper_qpos
158+
159+
return rows
25160

26161
@dataclass(frozen=True)
27162
class RecordedSimStep:

0 commit comments

Comments
 (0)