227 lines
9.9 KiB
Python
227 lines
9.9 KiB
Python
#!/usr/bin/env python
|
|
"""Replay a LeRobot (egodex-style) bimanual trajectory on r1pro_dex in fastsim,
|
|
and record the result via the `record` extension.
|
|
|
|
The source dataset stores, per frame, a 56-D ``observation.state``:
|
|
|
|
[ 0: 3] left wrist position (xyz) |
|
|
[ 3: 6] left wrist orientation (euler xyz)| left arm -> solved with IK
|
|
[ 6:28] left hand 22 joint angles -> set directly
|
|
[28:31] right wrist position (xyz) |
|
|
[31:34] right wrist orientation (euler xyz)| right arm -> solved with IK
|
|
[34:56] right hand 22 joint angles -> set directly
|
|
|
|
Because the wrist poses are expressed in the capture/camera frame (unknown
|
|
transform to the robot base), the arms are driven by *delta poses*: each frame's
|
|
target EE pose in the robot base frame is
|
|
|
|
target_base[t] = base_ee_init * ( cam_ee[0]^-1 * cam_ee[t] )
|
|
|
|
i.e. the trajectory's motion relative to its own first frame (expressed in the
|
|
EE-local frame, which is camera-frame-independent) applied on top of the robot's
|
|
actual initial EE pose. The 7-DoF arm joints are recovered with fastsim's
|
|
per-arm differential IK; the simulation runs physics-disabled (pure kinematic
|
|
playback) at dt=1/60 so each 15-fps data frame gets ~4 IK convergence steps.
|
|
|
|
Objects are NOT replayed (per requirement).
|
|
|
|
Usage (must use the fastsim conda env's python):
|
|
python replay_lerobot.py --episode 0
|
|
python replay_lerobot.py --all # one subprocess per episode
|
|
"""
|
|
import argparse
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Constants
|
|
# --------------------------------------------------------------------------- #
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
DEFAULT_DATASET = os.path.join(HERE, "extracted", "add_remove_lid_15fps_10epi")
|
|
DEFAULT_CONFIG = os.path.join(HERE, "replay_config.yaml")
|
|
DATA_FPS = 15
|
|
ROBOT_NAME = "r1pro_dex"
|
|
|
|
# 22 hand joints per side, in the exact order the dataset packs them
|
|
# (matches r1pro_dex.usd / benchmark.yaml actuator lists: thumb5, index4,
|
|
# middle4, ring4, pinky5).
|
|
LEFT_HAND_JOINTS = [
|
|
"left_thumb_CMC_FE", "left_thumb_CMC_AA", "left_thumb_MCP_FE", "left_thumb_MCP_AA", "left_thumb_IP",
|
|
"left_index_MCP_FE", "left_index_MCP_AA", "left_index_PIP", "left_index_DIP",
|
|
"left_middle_MCP_FE", "left_middle_MCP_AA", "left_middle_PIP", "left_middle_DIP",
|
|
"left_ring_MCP_FE", "left_ring_MCP_AA", "left_ring_PIP", "left_ring_DIP",
|
|
"left_pinky_CMC", "left_pinky_MCP_FE", "left_pinky_MCP_AA", "left_pinky_PIP", "left_pinky_DIP",
|
|
]
|
|
RIGHT_HAND_JOINTS = [n.replace("left_", "right_", 1) for n in LEFT_HAND_JOINTS]
|
|
|
|
# 56-D observation.state slices
|
|
SLICES = {
|
|
"left_arm": {"pos": slice(0, 3), "euler": slice(3, 6), "hand": slice(6, 28)},
|
|
"right_arm": {"pos": slice(28, 31), "euler": slice(31, 34), "hand": slice(34, 56)},
|
|
}
|
|
HAND_JOINTS = {"left_arm": LEFT_HAND_JOINTS, "right_arm": RIGHT_HAND_JOINTS}
|
|
ARMS = ["left_arm", "right_arm"]
|
|
|
|
# --- Arm targeting -------------------------------------------------------- #
|
|
# Delta-pose replay (camera pose is IGNORED): for each frame compute the relative
|
|
# transform from frame 0 in the data's own frame, then apply it on top of the
|
|
# robot's actual initial EE pose:
|
|
# delta_t = cam_ee[0]^-1 * cam_ee[t] # relative motion, in EE-local frame
|
|
# target_t = ee_init * delta_t # applied to the robot's init EE pose
|
|
# Frame 0 -> delta = identity -> robot starts exactly at its init EE pose.
|
|
# Optional fixed correction between the dataset wrist frame and r1pro's ee_link.
|
|
EE_OFFSET = {
|
|
"left_arm": np.eye(4),
|
|
"right_arm": np.eye(4),
|
|
}
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Data loading
|
|
# --------------------------------------------------------------------------- #
|
|
def load_episode(dataset_dir: str, episode: int) -> np.ndarray:
|
|
"""Return (N, 56) float32 states for one episode, ordered by frame_index."""
|
|
import pyarrow.parquet as pq
|
|
|
|
data_file = os.path.join(dataset_dir, "data", "chunk-000", "file-000.parquet")
|
|
df = pq.read_table(data_file).to_pandas()
|
|
ep = df[df["episode_index"] == episode].sort_values("frame_index")
|
|
if len(ep) == 0:
|
|
raise ValueError(f"episode {episode} not found in {data_file}")
|
|
states = np.stack(ep["observation.state"].to_numpy()).astype(np.float64)
|
|
assert states.shape[1] == 56, f"expected 56-D state, got {states.shape}"
|
|
return states
|
|
|
|
|
|
def build_temp_config(base_config: str, out_subdir: str) -> str:
|
|
"""Clone the YAML config, redirecting the recorder output to a per-episode dir."""
|
|
import yaml
|
|
|
|
with open(base_config) as f:
|
|
cfg = yaml.safe_load(f)
|
|
rec = cfg["extension"]["extension_cfg_dict"]["record"]
|
|
rec["backend_root_path"] = f"output://{out_subdir}"
|
|
fd, path = tempfile.mkstemp(suffix=".yaml", prefix="replay_cfg_")
|
|
with os.fdopen(fd, "w") as f:
|
|
yaml.safe_dump(cfg, f, sort_keys=False)
|
|
return path
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Single-episode replay (runs inside fastsim)
|
|
# --------------------------------------------------------------------------- #
|
|
def run_episode(dataset_dir: str, episode: int, base_config: str):
|
|
states = load_episode(dataset_dir, episode)
|
|
n_frames = len(states)
|
|
dataset_name = os.path.basename(dataset_dir.rstrip("/"))
|
|
out_subdir = f"benchmark_replay_record/{dataset_name}/episode_{episode:02d}"
|
|
temp_config = build_temp_config(base_config, out_subdir)
|
|
|
|
# Heavy sim imports happen only now (after the lightweight data load).
|
|
from fastsim.app import FastSim
|
|
from fastsim.unisim.scene_manager import SceneManager
|
|
from fastsim.utils.pose import Pose
|
|
from fastsim.utils.log import Log
|
|
|
|
# Pre-build camera-frame EE poses + hand targets per arm.
|
|
cam_ee = {arm: [] for arm in ARMS}
|
|
hand_targets = {arm: [] for arm in ARMS}
|
|
for s in states:
|
|
for arm in ARMS:
|
|
sl = SLICES[arm]
|
|
cam_ee[arm].append(
|
|
Pose.from_euler_xyz(position=s[sl["pos"]].tolist(),
|
|
euler_xyz=s[sl["euler"]].tolist())
|
|
)
|
|
hand_targets[arm].append(s[sl["hand"]].tolist())
|
|
|
|
sim = FastSim(temp_config)
|
|
sim.set_physics_disabled(True) # pure kinematic playback
|
|
|
|
ee_off = {arm: Pose.from_homogeneous_matrix(EE_OFFSET[arm]) for arm in ARMS}
|
|
state = {"base_ee_init": {}, "cam0_inv": {}, "ready": False}
|
|
|
|
def on_post_reset():
|
|
"""Anchor the delta-pose replay to the robot's actual init EE pose."""
|
|
robot = SceneManager.get_robot(ROBOT_NAME)
|
|
for arm in ARMS:
|
|
state["base_ee_init"][arm] = robot.get_ee_pose(arm_name=arm)
|
|
state["cam0_inv"][arm] = cam_ee[arm][0].inverse()
|
|
state["ready"] = True
|
|
Log.info(f"[replay] episode {episode}: {n_frames} frames @ {DATA_FPS}fps; "
|
|
f"delta-pose replay (camera pose ignored)", title="Replay")
|
|
|
|
def apply_frame(robot, idx):
|
|
for arm in ARMS:
|
|
# delta = cam_ee[0]^-1 * cam_ee[idx]; target = ee_init * delta * ee_off
|
|
delta = state["cam0_inv"][arm] * cam_ee[arm][idx]
|
|
target = state["base_ee_init"][arm] * (delta * ee_off[arm])
|
|
ik = robot.solve_ik(target, arm_name=arm)
|
|
jpos = np.asarray(ik["joint_position"]).reshape(-1).tolist()
|
|
robot.set_joint_position(jpos, ik["joint_names"])
|
|
robot.set_joint_position_target(jpos, ik["joint_names"])
|
|
# hand joints: exact angles, set directly
|
|
robot.set_joint_position(hand_targets[arm][idx], HAND_JOINTS[arm])
|
|
robot.set_joint_position_target(hand_targets[arm][idx], HAND_JOINTS[arm])
|
|
|
|
def on_step(ctx):
|
|
if not state["ready"]:
|
|
return
|
|
interval = max(1, round(1.0 / (ctx.dt * DATA_FPS)))
|
|
# Allow one extra interval so the recorder (which captures every
|
|
# `interval` steps) flushes the final converged frame.
|
|
if ctx.step > (n_frames + 1) * interval:
|
|
sim.request_terminate()
|
|
return
|
|
idx = min(ctx.step // interval, n_frames - 1)
|
|
robot = SceneManager.get_robot(ROBOT_NAME)
|
|
apply_frame(robot, idx)
|
|
|
|
sim.add_post_reset_callback(on_post_reset)
|
|
sim.add_step_callback(on_step)
|
|
Log.info(f"[replay] output -> output://{out_subdir}", title="Replay")
|
|
sim.start() # setup + loop; calls os._exit(0) on finish
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Multi-episode driver (subprocess per episode, since FastSim os._exit()s)
|
|
# --------------------------------------------------------------------------- #
|
|
def run_all(dataset_dir: str, base_config: str, episodes):
|
|
for ep in episodes:
|
|
print(f"\n========== EPISODE {ep} ==========", flush=True)
|
|
rc = subprocess.run(
|
|
[sys.executable, os.path.abspath(__file__),
|
|
"--dataset", dataset_dir, "--config", base_config, "--episode", str(ep)]
|
|
).returncode
|
|
# FastSim exits via os._exit(0); a non-zero code means a real failure.
|
|
if rc not in (0,):
|
|
print(f"[replay] episode {ep} exited with code {rc}", flush=True)
|
|
|
|
|
|
def list_episodes(dataset_dir: str):
|
|
import json
|
|
info = json.load(open(os.path.join(dataset_dir, "meta", "info.json")))
|
|
return list(range(int(info["total_episodes"])))
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description=__doc__)
|
|
ap.add_argument("--dataset", default=DEFAULT_DATASET, help="LeRobot dataset dir")
|
|
ap.add_argument("--config", default=DEFAULT_CONFIG, help="fastsim YAML config")
|
|
ap.add_argument("--episode", type=int, default=None, help="single episode index")
|
|
ap.add_argument("--all", action="store_true", help="replay every episode")
|
|
args = ap.parse_args()
|
|
|
|
if args.all:
|
|
run_all(args.dataset, args.config, list_episodes(args.dataset))
|
|
else:
|
|
ep = 0 if args.episode is None else args.episode
|
|
run_episode(args.dataset, ep, args.config)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |