fix replay

benchmark.yaml

@@ -29,7 +29,7 @@ scene:
       - 0.001
       - 0.001
       position:
-      - 0.419859
+      - 0.15
       - -4.02430000000001
       - 0.510259093
       rotation:
@@ -47,7 +47,7 @@ scene:
       - 0.001
       - 0.001
       position:
-      - 0.419859
+      - 0.15
       - -4.152430000000001
       - 0.510259093
       quaternion: [1, 0, 0, 0]
@@ -57,7 +57,7 @@ scene:
     r1pro_dex:
       name: r1pro_dex
       asset_path: asset://robots/r1pro/r1pro_dex.usd
-      position: [-0.5, -4.0, 0.0]
+      position: [-0.2, -4.1, 0.0]
       rotation: [1, 0, 0, 0]
       stereotype: modular_robot
       source: local
@@ -67,18 +67,18 @@ scene:
         torso_joint2: 0.0
         torso_joint3: 0.0
         torso_joint4: 0.0
-        left_arm_joint1: 0.0
-        left_arm_joint2: 0.5
+        left_arm_joint1: -0.2
+        left_arm_joint2: 0.05
         left_arm_joint3: 0.0
         left_arm_joint4: -1.0
         left_arm_joint5: 0.0
         left_arm_joint6: 0.0
         left_arm_joint7: 0.0
-        right_arm_joint1: 0.0
-        right_arm_joint2: -0.5
+        right_arm_joint1: -0.2
+        right_arm_joint2: -0.05
         right_arm_joint3: 0.0
         right_arm_joint4: -1.0
-        right_arm_joint5: 0.0
+        right_arm_joint5: 0.1
         right_arm_joint6: 0.0
         right_arm_joint7: 0.0
 
@@ -137,22 +137,33 @@ scene:
     head_camera:
       name: head_camera
       stereotype: camera
-      position: [-0.4, -4.0, 1.2]
-      look_at:
-        is_point: true
-        look_at_point: [0.4200, -4.1530, 0.4885]
       data_types: [rgb]
       width: 1280
       height: 720
       camera_model: pinhole
       fix_camera: true
+      focal_length: 2.8
+      horizontal_aperture: 4.890881131191918
+      vertical_aperture: 2.7608816125932627
+      convention: opengl
+      attach_to:
+        target_name: r1pro_dex
+        is_articulation_part: true
+        articulation_part_name: zed_link
+        create_fixed_joint: true
+        local_position: [0.0, 0.0, 0.0]
+        local_rotation:
+        - 0.33
+        - 1.0
+        - -0.0
+        - 0.0
     front_camera:
       name: front_camera
       stereotype: camera
-      position: [1.0, -4.0, 1.5]
+      position: [2, -4.1, 1.8]
       look_at:
         is_point: true
-        look_at_point: [-1.0, -4.0, 1.2]
+        look_at_point: [0.0, -4.1, 1.2]
       data_types: [rgb]
       width: 1280
       height: 720
@@ -161,7 +172,7 @@ scene:
     left_camera:
       name: left_camera
       stereotype: camera
-      position: [-1.0, -1.0, 1.2]
+      position: [-0.58554, -2.0, 1.8]
       look_at:
         is_point: true
         look_at_point: [0.0, -4.1, 1.2]
@@ -173,7 +184,7 @@ scene:
     right_camera:
       name: right_camera
       stereotype: camera
-      position: [-1.0, -6.5, 1.2]
+      position: [0.36816, -5.36, 1.8]
       look_at:
         is_point: true
         look_at_point: [0.0, -4.1, 1.2]

benchmark_replay.yaml

@@ -0,0 +1,297 @@
+general:
+  root_paths:
+    asset: /home/zhiyuan/zhujuan/joysim_exp/gen_data/data    # Root directory for assets (robots, objects, scene USDs, etc.)
+    checkpoints: /home/zhiyuan/zhujuan/checkpoints
+    output: /home/zhiyuan/zhujuan/joysim_exp/output                   # Root directory for outputs (recorded data, logs, etc.)
+
+simulation:
+  stereotype: isaaclab
+  launch_config:
+    device: cuda
+    enable_cameras: true
+    headless: true
+    livestream: 0
+scene:
+  name: kujiale_multispace
+  base_config:
+    stereotype: usd
+    name: _827313_home_workspace_00
+    source: platform
+    asset_path: platform://scenes/kujiale_multispace/827313_home/workspace_00.usd
+  object_cfg_dict:
+    omni6DPose_can_016:
+      name: omni6DPose_can_016
+      stereotype: rigid
+      source: local
+      asset_path: asset://objects/omni6DPose/can/omni6DPose_can_016/Aligned.usd
+      scale:
+      - 0.001
+      - 0.001
+      - 0.001
+      position:
+      - 0.15
+      - -4.02430000000001
+      - 0.510259093
+      rotation:
+      - -0.304408012043137
+      - -0.304408012043137
+      - 0.638228612805745
+      - 0.6382286128057448
+    omni6DPose_book_031:
+      name: omni6DPose_book_031
+      stereotype: rigid
+      source: local
+      asset_path: asset://objects/omni6DPose/book/omni6DPose_book_031/Aligned.usd
+      scale:
+      - 0.001
+      - 0.001
+      - 0.001
+      position:
+      - 0.15
+      - -4.152430000000001
+      - 0.510259093
+      quaternion: [1, 0, 0, 0]
+      axis_y_up: true
+
+  robot_cfg_dict:
+    r1pro_dex:
+      name: r1pro_dex
+      asset_path: asset://robots/r1pro/r1pro_dex.usd
+      position: [-0.2, -4.1, 0.0]
+      rotation: [1, 0, 0, 0]
+      stereotype: modular_robot
+      source: local
+
+      init_joint_position:
+        torso_joint1: 0.0
+        torso_joint2: 0.0
+        torso_joint3: 0.0
+        torso_joint4: 0.0
+        # palms-down init (back of hand up): verified ee_link +X palm-normal -> ~[0,0,-1]
+        # (down-score 0.97 on both hands), within joint limits.
+        left_arm_joint1: -0.2
+        left_arm_joint2: 0.6
+        left_arm_joint3: 0.0
+        left_arm_joint4: -1.6
+        left_arm_joint5: -2.3
+        left_arm_joint6: -0.33
+        left_arm_joint7: 0.0
+        right_arm_joint1: -0.2
+        right_arm_joint2: -0.6
+        right_arm_joint3: 0.0
+        right_arm_joint4: -1.6
+        right_arm_joint5: 2.3
+        right_arm_joint6: -0.33
+        right_arm_joint7: 0.0
+
+      actuator_cfg_dict:
+        left_arm:
+          stereotype: arm
+          joint_names_expr: [left_arm_joint1, left_arm_joint2, left_arm_joint3, left_arm_joint4, left_arm_joint5, left_arm_joint6, left_arm_joint7]
+          stiffness: 60000.0
+          damping: 4000.0
+        right_arm:
+          stereotype: arm
+          joint_names_expr: [right_arm_joint1, right_arm_joint2, right_arm_joint3, right_arm_joint4, right_arm_joint5, right_arm_joint6, right_arm_joint7]
+          stiffness: 60000.0
+          damping: 4000.0
+        left_hand:
+          stereotype: arm
+          joint_names_expr: [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]
+          stiffness: 50.0
+          damping: 5.0
+        right_hand:
+          stereotype: arm
+          joint_names_expr: [right_thumb_CMC_FE, right_thumb_CMC_AA, right_thumb_MCP_FE, right_thumb_MCP_AA, right_thumb_IP, right_index_MCP_FE, right_index_MCP_AA, right_index_PIP, right_index_DIP, right_middle_MCP_FE, right_middle_MCP_AA, right_middle_PIP, right_middle_DIP, right_ring_MCP_FE, right_ring_MCP_AA, right_ring_PIP, right_ring_DIP, right_pinky_CMC, right_pinky_MCP_FE, right_pinky_MCP_AA, right_pinky_PIP, right_pinky_DIP]
+          stiffness: 50.0
+          damping: 5.0
+        torso:
+          stereotype: arm
+          joint_names_expr: [torso_joint1, torso_joint2, torso_joint3, torso_joint4]
+          stiffness: 100000.0
+          damping: 8000.0
+        base_lock:
+          stereotype: arm
+          joint_names_expr: [steer_motor_joint1, steer_motor_joint2, steer_motor_joint3, wheel_motor_joint1, wheel_motor_joint2, wheel_motor_joint3]
+          stiffness: 100000.0
+          damping: 5000.0
+
+      arm_modules:
+        left_arm:
+          arm_actuator_name: left_arm
+          ee_link_name: left_hand_C_MC
+          ee_type: dexterous_hand
+          ee_actuator_name: left_hand
+        right_arm:
+          arm_actuator_name: right_arm
+          ee_link_name: right_hand_C_MC
+          ee_type: dexterous_hand
+          ee_actuator_name: right_hand
+
+      extra_modules:
+        torso:
+          actuator_name: torso
+
+      use_planner: false
+
+  sensor_cfg_dict:
+
+    head_camera:
+      name: head_camera
+      stereotype: camera
+      data_types: [rgb]
+      width: 1280
+      height: 720
+      camera_model: pinhole
+      fix_camera: true
+      focal_length: 2.8
+      horizontal_aperture: 4.890881131191918
+      vertical_aperture: 2.7608816125932627
+      convention: opengl
+      attach_to:
+        target_name: r1pro_dex
+        is_articulation_part: true
+        articulation_part_name: zed_link
+        create_fixed_joint: true
+        local_position: [0.0, 0.0, 0.0]
+        local_rotation:
+        - 0.33
+        - 1.0
+        - -0.0
+        - 0.0
+    front_camera:
+      name: front_camera
+      stereotype: camera
+      position: [2, -4.1, 1.8]
+      look_at:
+        is_point: true
+        look_at_point: [0.0, -4.1, 1.2]
+      data_types: [rgb]
+      width: 1280
+      height: 720
+      camera_model: pinhole
+      fix_camera: true
+    left_camera:
+      name: left_camera
+      stereotype: camera
+      position: [-0.58554, -2.0, 1.8]
+      look_at:
+        is_point: true
+        look_at_point: [0.0, -4.1, 1.2]
+      data_types: [rgb]
+      width: 1280
+      height: 720
+      camera_model: pinhole
+      fix_camera: true
+    right_camera:
+      name: right_camera
+      stereotype: camera
+      position: [0.36816, -5.36, 1.8]
+      look_at:
+        is_point: true
+        look_at_point: [0.0, -4.1, 1.2]
+      data_types: [rgb]
+      width: 1280
+      height: 720
+      camera_model: pinhole
+      fix_camera: true
+
+  light_cfg_dict:
+    sun:
+      name: sun
+      stereotype: general_light
+      light_type: distant
+      position: [0, 0, 5]
+      rotation: [1, 0, 0, 0]
+      intensity: 1000
+      angle: 0.53
+      color: [1.0, 1.0, 1.0]
+    sky:
+      name: sky
+      stereotype: general_light
+      light_type: dome
+      intensity: 10.0
+      color: [1.0, 1.0, 1.0]
+extension:
+  extension_cfg_dict:
+    replay_data_collect:
+      enable: true
+      stereotype: data_collect
+      observer_cfgs:
+        - stereotype: robot_observer
+          name: r1pro_dex
+          target_joint_names:
+              - 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_thumb_CMC_FE
+              - right_thumb_CMC_AA
+              - right_thumb_MCP_FE
+              - right_thumb_MCP_AA
+              - right_thumb_IP
+              - right_index_MCP_FE
+              - right_index_MCP_AA
+              - right_index_PIP
+              - right_index_DIP
+              - right_middle_MCP_FE
+              - right_middle_MCP_AA
+              - right_middle_PIP
+              - right_middle_DIP
+              - right_ring_MCP_FE
+              - right_ring_MCP_AA
+              - right_ring_PIP
+              - right_ring_DIP
+              - right_pinky_CMC
+              - right_pinky_MCP_FE
+              - right_pinky_MCP_AA
+              - right_pinky_PIP
+              - right_pinky_DIP
+          observe_ee_pose: true
+          observe_ee_state: true
+          observe_joint_position: true
+          observe_joint_velocity: true
+          observe_joint_positions: true
+          observe_joint_velocities: true
+          observe_joint_accelerations: true
+          observe_joint_position_targets: true
+          observe_joint_velocity_targets: true
+        - stereotype: sensor_observer
+          name: head_camera
+          observe_rgb: true
+        - stereotype: sensor_observer
+          name: front_camera
+          observe_rgb: true
+        - stereotype: sensor_observer
+          name: left_camera
+          observe_rgb: true
+        - stereotype: sensor_observer
+          name: right_camera
+          observe_rgb: true
+    record:
+      enable: true # set to true to record the data
+      stereotype: record
+      data_collector_name: replay_data_collect
+      record_backend: lfs
+      record_fps: 30
+      use_recorder_step: true
+      backend_root_path: output://replay_egodex
+      postprocess_list: ["hdf5", "video", "preview_video"]

replay.py

@@ -0,0 +1,227 @@
+#!/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()

replay_config.yaml

@@ -0,0 +1,202 @@
+# ============================================================
+#  R1Pro-Dex LeRobot replay config (fastsim)
+#  Derived from benchmark.yaml, adapted for offline trajectory
+#  replay driven by replay_lerobot.py.
+#
+#  Differences vs benchmark.yaml (intentional):
+#   - root_paths fixed for THIS machine.
+#   - kujiale scene + can/book objects dropped (objects are NOT
+#     replayed; the lid-on-cup trajectory is unrelated to that room),
+#     replaced by a ground plane for fast load & clear viewing.
+#   - NO replay / benchmark / policy_server extensions: the motion is
+#     driven from Python (hands set directly, arms solved via IK).
+#   - only data_collect + record extensions remain, to record the
+#     replayed motion back out (hdf5 + video).
+# ============================================================
+general:
+  root_paths:
+    asset: /home/zhiyuan/zhujuan/joysim_exp/gen_data/data
+    output: /home/zhiyuan/zhujuan/joysim_exp/output
+
+simulation:
+  stereotype: isaaclab
+  dt: 0.008333333333333333          # 1/120 -> 8 sim steps per 15fps data frame (more IK convergence)
+  launch_config:
+    device: cuda
+    enable_cameras: true
+    headless: true                  # set false to watch live; true records faster
+    livestream: 0
+
+scene:
+  name: r1pro_dex_lerobot_replay
+  position: [0, 0, 0]
+  rotation: [1, 0, 0, 0]
+  base_config:
+    name: default_base
+    stereotype: ground_plane
+    ground_plane_size: [100, 100]
+
+  robot_cfg_dict:
+    r1pro_dex:
+      name: r1pro_dex
+      asset_path: asset://robots/r1pro/r1pro_dex.usd
+      position: [0, 0, 0]
+      rotation: [1, 0, 0, 0]
+      stereotype: modular_robot
+      source: local
+
+      # init pose copied from benchmark.yaml (defines the starting EE pose
+      # that the delta-pose arm replay is anchored to).
+      init_joint_position:
+        torso_joint1: 0.0
+        torso_joint2: 0.0
+        torso_joint3: 0.0
+        torso_joint4: 0.0
+        # palms-down init (back of hand up): verified ee_link +X palm-normal -> ~[0,0,-1]
+        # (down-score 0.97 on both hands), within joint limits.
+        left_arm_joint1: -0.2
+        left_arm_joint2: 0.6
+        left_arm_joint3: 0.0
+        left_arm_joint4: -1.6
+        left_arm_joint5: -2.3
+        left_arm_joint6: -0.33
+        left_arm_joint7: 0.0
+        right_arm_joint1: -0.2
+        right_arm_joint2: -0.6
+        right_arm_joint3: 0.0
+        right_arm_joint4: -1.6
+        right_arm_joint5: 2.3
+        right_arm_joint6: -0.33
+        right_arm_joint7: 0.0
+
+      actuator_cfg_dict:
+        left_arm:
+          stereotype: arm
+          joint_names_expr: [left_arm_joint1, left_arm_joint2, left_arm_joint3, left_arm_joint4, left_arm_joint5, left_arm_joint6, left_arm_joint7]
+          stiffness: 60000.0
+          damping: 4000.0
+        right_arm:
+          stereotype: arm
+          joint_names_expr: [right_arm_joint1, right_arm_joint2, right_arm_joint3, right_arm_joint4, right_arm_joint5, right_arm_joint6, right_arm_joint7]
+          stiffness: 60000.0
+          damping: 4000.0
+        left_hand:
+          stereotype: arm
+          joint_names_expr: [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]
+          stiffness: 50.0
+          damping: 5.0
+        right_hand:
+          stereotype: arm
+          joint_names_expr: [right_thumb_CMC_FE, right_thumb_CMC_AA, right_thumb_MCP_FE, right_thumb_MCP_AA, right_thumb_IP, right_index_MCP_FE, right_index_MCP_AA, right_index_PIP, right_index_DIP, right_middle_MCP_FE, right_middle_MCP_AA, right_middle_PIP, right_middle_DIP, right_ring_MCP_FE, right_ring_MCP_AA, right_ring_PIP, right_ring_DIP, right_pinky_CMC, right_pinky_MCP_FE, right_pinky_MCP_AA, right_pinky_PIP, right_pinky_DIP]
+          stiffness: 50.0
+          damping: 5.0
+        torso:
+          stereotype: arm
+          joint_names_expr: [torso_joint1, torso_joint2, torso_joint3, torso_joint4]
+          stiffness: 100000.0
+          damping: 8000.0
+        base_lock:
+          stereotype: arm
+          joint_names_expr: [steer_motor_joint1, steer_motor_joint2, steer_motor_joint3, wheel_motor_joint1, wheel_motor_joint2, wheel_motor_joint3]
+          stiffness: 100000.0
+          damping: 5000.0
+
+      arm_modules:
+        left_arm:
+          arm_actuator_name: left_arm
+          ee_link_name: left_hand_C_MC
+          ee_type: dexterous_hand
+          ee_actuator_name: left_hand
+        right_arm:
+          arm_actuator_name: right_arm
+          ee_link_name: right_hand_C_MC
+          ee_type: dexterous_hand
+          ee_actuator_name: right_hand
+
+      extra_modules:
+        torso:
+          actuator_name: torso
+
+      use_planner: false
+
+  sensor_cfg_dict:
+    # Ego/head view, attached to the robot head — comparable to the source
+    # dataset's observation.images.ego_view.
+    head_camera:
+      name: head_camera
+      stereotype: camera
+      data_types: [rgb]
+      width: 1280
+      height: 720
+      camera_model: pinhole
+      fix_camera: true
+      focal_length: 2.8
+      horizontal_aperture: 4.890881131191918
+      vertical_aperture: 2.7608816125932627
+      convention: opengl
+      attach_to:
+        target_name: r1pro_dex
+        is_articulation_part: true
+        articulation_part_name: zed_link
+        create_fixed_joint: true
+        local_position: [0.0, 0.0, 0.0]
+        local_rotation: [0.33, 1.0, -0.0, 0.0]
+    # Third-person view of the whole robot.
+    front_camera:
+      name: front_camera
+      stereotype: camera
+      position: [2.2, 0.0, 1.4]
+      look_at:
+        is_point: true
+        look_at_point: [0.0, 0.0, 1.0]
+      data_types: [rgb]
+      width: 1280
+      height: 720
+      camera_model: pinhole
+      fix_camera: true
+
+  light_cfg_dict:
+    sun:
+      name: sun
+      stereotype: general_light
+      light_type: distant
+      position: [0, 0, 5]
+      rotation: [1, 0, 0, 0]
+      intensity: 1000
+      angle: 0.53
+      color: [1.0, 1.0, 1.0]
+    sky:
+      name: sky
+      stereotype: general_light
+      light_type: dome
+      intensity: 10.0
+      color: [1.0, 1.0, 1.0]
+
+extension:
+  extension_cfg_dict:
+    data_collect:
+      enable: true
+      stereotype: data_collect
+      observer_cfgs:
+        - stereotype: robot_observer
+          name: r1pro_dex
+          observe_joint_positions: true
+          observe_joint_velocities: true
+          observe_joint_position_targets: true
+          observe_ee_pose: true
+        - stereotype: sensor_observer
+          name: head_camera
+          observe_rgb: true
+        - stereotype: sensor_observer
+          name: front_camera
+          observe_rgb: true
+    record:
+      enable: true
+      stereotype: record
+      backend_root_path: output://benchmark_replay_record
+      data_collector_name: data_collect
+      record_backend: lfs
+      record_fps: 15            # matches the source dataset fps
+      use_recorder_step: true
+      record_scene: true
+      postprocess_list: [hdf5, video, preview_video]

starvla_inference_server.py

@@ -129,7 +129,7 @@ class StarvlaInferenceServer:
     def inference(self, observation: dict) -> dict:
         
         img_head, state_vec, prompt = \
-            self.parse_observation(observation)
+            self.parse_observation(observation, target_size=(410, 224))
         vla_input = {
             # "batch_images": [[img_left, img_right, img_wrist]],
             "image": [img_head],

starvla_policy.py

@@ -2,6 +2,7 @@ import pickle
 import time
 import json
 import numpy as np
+from scipy.spatial.transform import Rotation as R
 import requests
 
 from fastsim.annotations.config_class import configclass, field
@@ -70,6 +71,11 @@ class StarvlaPolicy(Policy):
         self.visualize_action_ee_pose = config.visualize_action_ee_pose
         self.visualize_state_ee_pose = config.visualize_state_ee_pose
         self.visualize_bounding_box_targets = list(config.visualize_bounding_box_targets or [])
+        # prevent circular import
+        import pandas as pd
+        df_data = pd.read_parquet("/home/zhiyuan/zhujuan/datasets/add_remove_lid_15fps_10epi/data/chunk-000/file-000.parquet")
+        self.dummy_data = np.array(df_data.groupby('episode_index')['observation.state'].apply(list).to_dict()[0])
+        self.dummy_data_idx = 0
 
     def reset(self) -> None:
         self.current_state = {}
@@ -167,26 +173,58 @@ class StarvlaPolicy(Policy):
 
     def postprocess_action(self, action: dict) -> BenchmarkAction:
         benchmark_action = BenchmarkAction()
+        read_chunk_size = 1
+        dummy_action = self.dummy_data[self.dummy_data_idx:(self.dummy_data_idx + read_chunk_size)]
+        if self.dummy_data_idx + read_chunk_size >= self.dummy_data.shape[0]:
+            self.dummy_data_idx = 0
+            exit(0)
+        else:
+            self.dummy_data_idx += read_chunk_size
+        read_chunk_id = 0
+        print(f'{self.current_chunk_id=}, {self.dummy_data_idx = }, {read_chunk_id=}')
+        time.sleep(1.0)
+
+        left_rpy_state = dummy_action[:, 3:6]         # (3,)
+        right_rpy_state = dummy_action[:, 31:34]         # (3,)
+
+        left_rot_state = R.from_euler('xyz', left_rpy_state).as_matrix()
+        right_rot_state = R.from_euler('xyz', right_rpy_state).as_matrix()
+
+        left_state_rot6d = np.concatenate([left_rot_state[:, 0], left_rot_state[:, 1]], axis=-1) # (6,)
+        right_state_rot6d = np.concatenate([right_rot_state[:, 0], right_rot_state[:, 1]], axis=-1) # (6,)
+
+        read_state = {"left_arm": {
+                    "ee_position_chunks": dummy_action[:, :3].tolist(),
+                    "ee_rot6d_chunks": left_state_rot6d.tolist(),
+                    "finger_chunks": dummy_action[:, 6:28].tolist()},
+                "right_arm": {
+                    "ee_position_chunks": dummy_action[:, 28:31].tolist(),
+                    "ee_rot6d_chunks": right_state_rot6d.tolist(),
+                    "finger_chunks": dummy_action[:, 34:56].tolist()}
+                }
         for arm_key in self.robot['arms'].keys():
             action_arm = action[arm_key]
             delta_ee_pose = Pose(position=action_arm["ee_delta_position_chunks"][self.current_chunk_id], rot6d=action_arm["ee_delta_rot6d_chunks"][self.current_chunk_id])
             curr_state_ee_pose = Pose(position=self.current_state[arm_key]["ee_pos"], rot6d=self.current_state[arm_key]["ee_rot6d"])
             curr_action_ee_pose = curr_state_ee_pose * delta_ee_pose # action2base = state2base * action2state
             finger_joint_qpos = action_arm["finger_chunks"][self.current_chunk_id] + self.current_state[arm_key]["finger_qpos"] 
-            joint_names = self.left_hand_joints if arm_key == "left_arm" else self.right_hand_joints    
+            joint_names = self.left_hand_joints if arm_key == "left_arm" else self.right_hand_joints
+            state_arm = read_state[arm_key]
             benchmark_action.add_robot_action(
                 RobotAction(
                     control_mode=ControlMode.POSITION,
                     robot_name=self.robot_name,
                     joint_names=joint_names,
-                    joint_positions=finger_joint_qpos
+                    # joint_positions=finger_joint_qpos
+                    joint_positions=state_arm["finger_chunks"][read_chunk_id]
                 )
             )
             benchmark_action.add_robot_action(
                 RobotAction(
                     control_mode=ControlMode.EE_POSE,
                     robot_name=self.robot_name,
-                    ee_pose=curr_action_ee_pose,
+                    # ee_pose=curr_action_ee_pose,
+                    ee_pose=Pose(position=state_arm["ee_position_chunks"][read_chunk_id], rot6d=state_arm["ee_rot6d_chunks"][read_chunk_id]),
                     arm_name=arm_key
                 )
             )