fix replay

benchmark.yaml

@@ -1,47 +1,42 @@
 general:
-  scan_project: true
   root_paths:
-    asset: /home/ubuntu/xionghao/sim_hofee/sim_hofee/assets
-    checkpoints: /home/ubuntu/xionghao/starVLA-starVLA/playground/Checkpoints
-    output: /home/ubuntu/xionghao/sim_hofee
+    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
-  intiailize_steps: 300
   launch_config:
     device: cuda
     enable_cameras: true
     headless: false
     livestream: 0
-    
-    
 scene:
-  name: 827313_home
+  name: kujiale_multispace
   base_config:
     stereotype: usd
-    name: _827313_home_workspace_01
-    source: local
-    asset_path: asset://scenes/kujiale_multispace/827313_home/workspace_01.usd
+    name: _827313_home_workspace_00
+    source: platform
+    asset_path: platform://scenes/kujiale_multispace/827313_home/workspace_00.usd
   object_cfg_dict:
-    omni6DPose_timer_017:
-      name: omni6DPose_timer_017
+    omni6DPose_can_016:
+      name: omni6DPose_can_016
       stereotype: rigid
       source: local
-      asset_path: asset://objects/omni6DPose/timer/omni6DPose_timer_017/Aligned.usd
+      asset_path: asset://objects/omni6DPose/can/omni6DPose_can_016/Aligned.usd
       scale:
       - 0.001
       - 0.001
       - 0.001
       position:
-      - 0.552364
-      - -4.0582599999999995
-      - 0.524713118
-      quaternion:
-      - 0.166210542394157
-      - 0.166210542394157
-      - 0.6872947370648492
-      - 0.6872947370648491
-      axis_y_up: true
+      - 0.15
+      - -4.02430000000001
+      - 0.510259093
+      rotation:
+      - -0.304408012043137
+      - -0.304408012043137
+      - 0.638228612805745
+      - 0.6382286128057448
     omni6DPose_book_031:
       name: omni6DPose_book_031
       stereotype: rigid
@@ -52,160 +47,169 @@ scene:
       - 0.001
       - 0.001
       position:
-      - 0.6623640000000001
-      - -3.7882599999999997
-      - 0.5101601435
-      quaternion:
-      - 0.7063055546421202
-      - 0.7063055546421203
-      - -0.03365209475927027
-      - -0.033652094759270265
+      - 0.15
+      - -4.152430000000001
+      - 0.510259093
+      quaternion: [1, 0, 0, 0]
       axis_y_up: true
+
   robot_cfg_dict:
-    Franka_Robotiq_2f85:
-      name: Franka_Robotiq_2f85
-      asset_path: asset://Franka/franka_robotiq_2f85_zedmini.usd
-      position:
-      - 1.082364
-      - -3.92826
-      - 0.47629299999999997
-      rotation: 
-      - 7.549799991308018e-08
-      - 0.0
-      - 0.0
-      - 0.9999999999999973
+    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
-      ee_link_name: panda_link8
-      ik_joint_names:
-        - panda_joint1
-        - panda_joint2
-        - panda_joint3
-        - panda_joint4
-        - panda_joint5
-        - panda_joint6
-        - panda_joint7
+
       init_joint_position:
-        # panda_joint1: 0.18641542
-        # panda_joint2: 0.47660449
-        # panda_joint3: -0.03320411
-        # panda_joint4: -2.27693725
-        # panda_joint5: 0.98161776
-        # panda_joint6: 2.20247197
-        # panda_joint7: 0.71794897
-        panda_joint2: -0.1633
-        panda_joint4: -1.07
-        panda_joint6: 0.8933
-        panda_joint7: 0.785
-      arm_modules:
-        main_arm:
-          arm_actuator_name: franka_arm
-          ee_link_name: panda_link8
-          ee_type: gripper
-          ee_actuator_name: robotiq_gripper
+        torso_joint1: 0.0
+        torso_joint2: 0.0
+        torso_joint3: 0.0
+        torso_joint4: 0.0
+        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.2
+        right_arm_joint2: -0.05
+        right_arm_joint3: 0.0
+        right_arm_joint4: -1.0
+        right_arm_joint5: 0.1
+        right_arm_joint6: 0.0
+        right_arm_joint7: 0.0
+
       actuator_cfg_dict:
-        franka_arm:
+        left_arm:
           stereotype: arm
-          joint_names_expr: [panda_joint1, panda_joint2, panda_joint3, panda_joint4, panda_joint5, panda_joint6, panda_joint7]
-          stiffness: 3000.0
-          damping: 800.0
-        robotiq_gripper:
-          stereotype: gripper
-          joint_names_expr: [robotiq_85_left_knuckle_joint]
-          stiffness: 10000
-          damping: 500.0
-          close_control_type: velocity
-          open_control_type: position
-          drive_joints:
-            robotiq_85_left_knuckle_joint:
-              close_velocity: 5.0
-              open_velocity: -5.0
-              close_position: 0.8
-              open_position: 0.0
+          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:
-    Hand_Camera:
-      name: Hand_Camera
+
+    head_camera:
+      name: head_camera
       stereotype: camera
-      data_types:
-      - rgb
-      - depth
-      - normals
+      data_types: [rgb]
       width: 1280
       height: 720
       camera_model: pinhole
       fix_camera: true
       focal_length: 2.8
-      horizontal_aperture: 4.893416860031241
+      horizontal_aperture: 4.890881131191918
       vertical_aperture: 2.7608816125932627
       convention: opengl
       attach_to:
-        target_name: Franka_Robotiq_2f85
+        target_name: r1pro_dex
         is_articulation_part: true
-        articulation_part_name: panda_link8
+        articulation_part_name: zed_link
         create_fixed_joint: true
-        local_position:
-        - -0.07128738160694643
-        - 0.03551506300731732
-        - 0.018927748370281355
+        local_position: [0.0, 0.0, 0.0]
         local_rotation:
-        - -0.12117023430710862
-        - -0.6862313269668
-        - 0.7070213671685396
-        - 0.12052023305019997
-    Left_Camera:
-      name: Left_Camera
+        - 0.33
+        - 1.0
+        - -0.0
+        - 0.0
+    front_camera:
+      name: front_camera
       stereotype: camera
-      data_types:
-      - rgb
-      - depth
-      - normals
+      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: false
-      focal_length: 2.1
-      horizontal_aperture: 5.030789363390793
-      vertical_aperture: 2.833796298140747
-      convention: opengl
-      attach_to:
-        target_name: Franka_Robotiq_2f85
-        local_position:
-        - 0.31702696813014064
-        - -0.3844238699868664
-        - 0.6551552990137672
-        local_rotation:
-        - 0.8742457685173938
-        - 0.38378563025938384
-        - -0.11951449178007277
-        - -0.27224843891267797
-    Right_Camera:
-      name: Right_Camera
+      fix_camera: true
+    left_camera:
+      name: left_camera
       stereotype: camera
-      data_types:
-      - rgb
-      - depth
-      - normals
+      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: false
-      focal_length: 2.1
-      horizontal_aperture: 5.050364265142387
-      vertical_aperture: 2.833796298140747
-      convention: opengl
-      attach_to:
-        target_name: Franka_Robotiq_2f85
-        local_position:
-        - 0.21844487914880717
-        - 0.20172329179193413
-        - 0.30108042236545296
-        local_rotation:
-        - -0.5316249212230874
-        - -0.38697158527836417
-        - 0.44338617110944967
-        - 0.6091277686910994
+      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:
     benchmark_data_collect:
@@ -213,8 +217,52 @@ extension:
       stereotype: data_collect
       observer_cfgs:
         - stereotype: robot_observer
-          name: Franka_Robotiq_2f85
-          target_joint_names: [panda_joint1, panda_joint2, panda_joint3, panda_joint4, panda_joint5, panda_joint6, panda_joint7, robotiq_85_left_knuckle_joint]
+          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
@@ -225,13 +273,16 @@ extension:
           observe_joint_position_targets: true
           observe_joint_velocity_targets: true
         - stereotype: sensor_observer
-          name: Hand_Camera
+          name: head_camera
           observe_rgb: true
         - stereotype: sensor_observer
-          name: Left_Camera
+          name: front_camera
           observe_rgb: true
         - stereotype: sensor_observer
-          name: Right_Camera
+          name: left_camera
+          observe_rgb: true
+        - stereotype: sensor_observer
+          name: right_camera
           observe_rgb: true
           
     starvla_benchmark:
@@ -241,19 +292,18 @@ extension:
       action_frequency: 15.0
       timeout_per_episode: 300
       goals:
-        - name: cola on top of book
-          description: check if the cola bottle is on the book
+        - name: can on top of book
+          description: check if the can is on the book
           stereotype: on_top
           object_A_name: omni6DPose_book_031
-          object_B_name: omni6DPose_timer_017
+          object_B_name: omni6DPose_can_016
       policy:
         stereotype: starvla
-        robot_name: Franka_Robotiq_2f85
-        arm_name: main_arm
-        sensor_names: [Hand_Camera, Left_Camera, Right_Camera]
-        prompt: pick up the timer and put on the book
+        robot_name: r1pro_dex
+        arm_name: right_arm
+        sensor_names: [head_camera]
+        prompt: pick up the can and put on the book
         run_trunk_size: 16
-        gripper_width_mapper_file: ./gripper_width_robotiq_2f85_fixed.json
         visualize_action_ee_pose: true
         visualize_state_ee_pose: true
         visualize_bounding_box_targets: [] # [omni6DPose_plug_001, omni6DPose_can_016] # 打开会被policy看到，会影响policy的推理结果
@@ -264,16 +314,11 @@ extension:
       data_collector_name: benchmark_data_collect
       record_fps: 30
       backend_root_path: output://benchmark_record
-      postprocess_list: ["hdf5", "video"]
+      postprocess_list: ["hdf5", "video", "preview_video"]
 
 policy_server:
-  # ckpt_path: checkpoints://0324_qwenpi_droid_pretrain_8node/checkpoints/steps_30000_pytorch_model.pt
-  # ckpt_path: checkpoints://0405_qwenpi_droid_norm_pretrain_8node/checkpoints/steps_60000_pytorch_model.pt
-  # ckpt_path: checkpoints://0407_qwenpi_droid_postrain/final_model/pytorch_model.pt
-  ckpt_path: checkpoints://0407_qwenpi_droid_from_scratch/final_model/pytorch_model.pt
+  ckpt_path: checkpoints://egodex_part1_restats_gbs1024/checkpoints/steps_70000_pytorch_model.pt
   ckpt_source: local
   host: 0.0.0.0
   port: 5000
   use_bf16: true
-  unnorm_key: oxe_bridge
-  state_mode: ee_pose7

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

@@ -21,6 +21,39 @@ def pad_to_dim(x: np.ndarray, target_dim: int, axis: int = -1, value: float = 0.
         return np.pad(x, pad_width, constant_values=value)
     return x
 
+
+def normalize_states(states, statistics):
+    stats = statistics["new_embodiment"]["state"]
+    q01 = np.array(stats["q01"]).astype(states.dtype)
+    q99 = np.array(stats["q99"]).astype(states.dtype)
+
+    # In the case of q01 == q99, the normalization will be undefined
+    # So we set the normalized values to the original values
+    mask = q01 != q99
+    normalized = np.zeros_like(states)
+
+    # Normalize the values where q01 != q99
+    # Formula: 2 * (x - q01) / (q99 - q01) - 1
+    normalized[..., mask] = (states[..., mask] - q01[..., mask]) / (
+        q99[..., mask] - q01[..., mask]
+    )
+    normalized[..., mask] = 2 * normalized[..., mask] - 1
+
+    # Set the normalized values to the original values where q01 == q99
+    normalized[..., ~mask] = states[..., ~mask]
+
+    # Clip the normalized values to be between -1 and 1
+    normalized = np.clip(normalized, -1, 1)
+    return normalized
+
+
+def unnormalize_actions(normalized_actions, statistics):
+    stats = statistics["new_embodiment"]["action"]
+    q01 = np.array(stats["q01"]).astype(normalized_actions.dtype)
+    q99 = np.array(stats["q99"]).astype(normalized_actions.dtype)
+
+    return (normalized_actions + 1) / 2 * (q99 - q01) + q01
+
 class StarvlaInferenceServer:
 
     def __init__(self, config_path: str):
@@ -38,8 +71,6 @@ class StarvlaInferenceServer:
         self.host = policy_server_cfg.get("host", "0.0.0.0")
         self.port = policy_server_cfg.get("port", 5000)
         self.use_bf16 = policy_server_cfg.get("use_bf16", True)
-        self.unnorm_key = policy_server_cfg.get("unnorm_key", "oxe_bridge")
-        self.state_mode = policy_server_cfg.get("state_mode", "ee_pose7")
 
         print("Loading StarVLA model...")
         self.model = self.load_model()
@@ -74,45 +105,40 @@ class StarvlaInferenceServer:
         model = build_framework(cfg=cfg)
         model.norm_stats = norm_stats
 
-        state_dict = torch.load(self.ckpt_path, map_location="cpu")
-        model.load_state_dict(state_dict, strict=True)
-
         if self.use_bf16:
             model = model.to(torch.bfloat16)
+        model = model.eval()
+
+        state_dict = torch.load(self.ckpt_path, map_location="cpu")
+        model.load_state_dict(state_dict, strict=True)
+        model = model.to("cuda")
 
-        model = model.to("cuda").eval()
 
         self.norm_stats = norm_stats
-        self.action_norm_stats = norm_stats.get(self.unnorm_key, {}).get("action", None)
-
         return model
 
     def parse_observation(self, obs, target_size=(320, 180)):
 
-        left_rgb, right_rgb, wrist_rgb = obs["rgb"]["Left_Camera"], obs["rgb"]["Right_Camera"], obs["rgb"]["Hand_Camera"]
+        head_rgb = obs["rgb"]["head_camera"]
 
-        img_left = Image.fromarray(cv2.resize(left_rgb, target_size))
-        img_right = Image.fromarray(cv2.resize(right_rgb, target_size))
-        img_wrist = Image.fromarray(cv2.resize(wrist_rgb, target_size))
-
-        state_vec = obs["state"]
-        # import ipdb;ipdb.set_trace()
+        img_head = Image.fromarray(cv2.resize(head_rgb, target_size))
+        state_vec = normalize_states(obs["state"], self.norm_stats)
         # state_vec = pad_to_dim(np.array(state_vec), 100, axis=-1)
-        return img_left, img_right, img_wrist, state_vec, obs["prompt"]
+        return img_head, state_vec, obs["prompt"]
 
     def inference(self, observation: dict) -> dict:
         
-        img_left, img_right, img_wrist, state_vec, prompt = \
-            self.parse_observation(observation)
-        print(f"{state_vec.shape}")
+        img_head, state_vec, prompt = \
+            self.parse_observation(observation, target_size=(410, 224))
         vla_input = {
-            "batch_images": [[img_left, img_right, img_wrist]],
-            "instructions": [prompt],
-            "state": [state_vec]
+            # "batch_images": [[img_left, img_right, img_wrist]],
+            "image": [img_head],
+            "lang": prompt,
+            "state": state_vec[None, :], # (1, 62)
         }
 
         with torch.no_grad():
-            output = self.model.predict_action(**vla_input)
+            output = self.model.predict_action(examples=vla_input)
 
         actions = output.get("normalized_actions")
 
@@ -121,9 +147,16 @@ class StarvlaInferenceServer:
 
         if actions.ndim == 3:
             actions = actions[0] # (16, 10)
-        return {"ee_delta_position_chunks": actions[:, :3].tolist(), 
+        actions = unnormalize_actions(actions, self.norm_stats)
+        return {"left_arm": {
+                    "ee_delta_position_chunks": actions[:, :3].tolist(),
                     "ee_delta_rot6d_chunks": actions[:, 3:9].tolist(),
-                "gripper_width_chunks": actions[:, 9:10].tolist()}
+                    "finger_chunks": actions[:, 9:31].tolist()},
+                "right_arm": {
+                    "ee_delta_position_chunks": actions[:, 31:34].tolist(),
+                    "ee_delta_rot6d_chunks": actions[:, 34:40].tolist(),
+                    "finger_chunks": actions[:, 40:62].tolist()}
+                }
 
     def register_routes(self):
 

starvla_policy.py

@@ -2,33 +2,31 @@ import pickle
 import time
 import json
 import numpy as np
+from scipy.spatial.transform import Rotation as R
 import requests
 
-from joysim.annotations.config_class import configclass, field
-from joysim.annotations.stereotype import stereotype
-from joysim.controllers.spawnable_controller import SpawnableController
-from joysim.controllers.visualize_controller import VisualizeController
-from joysim.unisim.robots.models.modular_robot import ModularRobot
-from joysim.utils.namespace import PoseVisualType, SimulatorType
-from joysim.unisim.robots.actuator_configs.grippers import GripperDriveJointConfig
-from joysim.extensions.benchmark.action import RobotAction
-from joysim.extensions.benchmark.benchmark import (
+from fastsim.annotations.config_class import configclass, field
+from fastsim.annotations.stereotype import stereotype
+from fastsim.controllers.spawnable_controller import SpawnableController
+from fastsim.controllers.visualize_controller import VisualizeController
+from fastsim.unisim.robots.models.modular_robot import ModularRobot
+from fastsim.utils.namespace import PoseVisualType, SimulatorType
+from fastsim.unisim.robots.actuator_configs.grippers import GripperDriveJointConfig
+from fastsim.extensions.benchmark.action import RobotAction
+from fastsim.extensions.benchmark.benchmark import (
     BenchmarkAction,
     BenchmarkObservation,
     ControlMode,
 )
-from joysim.extensions.benchmark.policy import Policy, PolicyConfig
-from joysim.utils.log import Log
-from joysim.utils.pose import Pose
+from fastsim.extensions.benchmark.policy import Policy, PolicyConfig
+from fastsim.utils.log import Log
+from fastsim.utils.pose import Pose
 
 @configclass
 @stereotype.register_config("starvla")
 class StarvlaPolicyConfig(PolicyConfig):
 
     robot_name: str = field(default="None", required=True, comment="The name of the robot")
-    arm_name: str = field(default="main_arm", required=True, comment="The name of the arm module to control")
-    drive_name: str = field(default="robotiq_85_left_knuckle_joint", required=True, comment="The name of the drive module to control")
-    gripper_width_mapper_file: str = field(default="", required=True, comment="The file path to the gripper width mapper")
     visualize_action_ee_pose: bool = field(default=False, required=True, comment="Whether to visualize the action end effector pose")
     visualize_state_ee_pose: bool = field(default=False, required=True, comment="Whether to visualize the state end effector pose")
     visualize_bounding_box_targets: list[str] = field(
@@ -67,37 +65,34 @@ class StarvlaPolicy(Policy):
         super().__init__(config)
 
         self.robot_name = config.robot_name
-        self.arm_name = config.arm_name
-        self.drive_name = config.drive_name
         self.sensor_names = config.sensor_names
         self.server_url = config.server_url
         self.prompt = config.prompt
-        self.gripper_width_mapper = json.load(open(config.gripper_width_mapper_file, "r"))
         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_ee_position_state = None
-        self.current_ee_rot6d_state = None
-        self.current_gripper_width = None
+        self.current_state = {}
         self.current_chunk_id = 0
         self.current_chunk_result = None
         self.run_trunk_size = self.config.run_trunk_size
         self.robot: ModularRobot = SpawnableController.get_spawnable_data(self.robot_name).unwrap()
-        self.drive_joints: dict[str, GripperDriveJointConfig] = self.robot.get_arm(self.arm_name).get_ee().get_drive_joints()
-        self.robot_drive_name = list(self.drive_joints.keys())[0]
-        
-        for joint_name, joint_config in self.drive_joints.items():
-            SpawnableController.control_robot(self.robot_name, "set_joint_stiffness", parameters={"joint_names": [joint_name], "stiffness": joint_config.position_control_stiffness}).unwrap()
-            SpawnableController.control_robot(self.robot_name, "set_joint_damping", parameters={"joint_names": [joint_name], "damping": joint_config.position_control_damping}).unwrap()
-            SpawnableController.control_robot(self.robot_name, "set_joint_effort_limit", parameters={"joint_names": [joint_name], "effort_limit": 5000}).unwrap()
-            SpawnableController.control_robot(self.robot_name, "set_joint_effort_limit", parameters={"joint_names": [self.robot_drive_name], "effort_limit": 5000}).unwrap()
-        self.max_width = float("-inf")
-        self.min_width = float("inf")
-        for entry in self.gripper_width_mapper:
-            self.max_width = max(self.max_width, entry["width"])
-            self.min_width = min(self.min_width, entry["width"])
+        self.left_hand_joints = SpawnableController.control_robot(
+            self.robot_name,
+            "get_actuator_joint_names",
+            parameters={"actuator_name": "left_hand"},
+        ).unwrap()
+        self.right_hand_joints = SpawnableController.control_robot(
+            self.robot_name,
+            "get_actuator_joint_names",
+            parameters={"actuator_name": "right_hand"},
+        ).unwrap()
 
     def warmup(self, benchmark_observation: BenchmarkObservation) -> None:
         Log.info(f"Waiting for StarVLA inference server to be ready...")
@@ -120,28 +115,41 @@ class StarvlaPolicy(Policy):
         elif response.status_code != 200:
             Log.error(f"StarVLA server error with status code <{response.status_code}> : {response.text}", exit=True)
 
+    def split_joints(self, state_or_action, keys=None) -> list[dict]:
+        if keys is None:
+            keys = ["left_arm", "right_arm"]
+        total_dim = 31 * len(keys)
+        assert state_or_action.shape[-1] == total_dim, f"Expected last dimension to be {total_dim}, got {state_or_action.shape[-1]}"
+        joints_all = np.split(state_or_action, [31], axis=-1)
+        return_dict = {}
+        for key, joints in zip(keys, joints_all):
+            ee_pos, ee_rot6d, finger_qpos = np.split(joints, [3, 9], axis=-1)
+            return_dict[key] = {
+                "ee_pos": ee_pos,
+                "ee_rot6d": ee_rot6d,
+                "finger_qpos": finger_qpos
+            }
+        return return_dict
+
     def preprocess_observation(self, benchmark_observation: BenchmarkObservation) -> dict:
         robot_obs = benchmark_observation.get_robot_observations(self.robot_name)["robot_data"]
-        ee_pose_base = robot_obs["ee_pose"][self.arm_name]["base_frame"]
-        ee_position, ee_rot6d = ee_pose_base["position"],ee_pose_base["rot6d"]
-        arm_joint_positions = robot_obs["joint_positions"][:7] # 临时多加了一个drive的位置，现在读的最后一个joint值是drive
-        drive_joint_positions = robot_obs["joint_positions"][-1]
-        normalized_gripper_width = self.__map_joint_position_to_normalized_width(drive_joint_positions)
-        Log.debug(f"input normalized_gripper_width state: {round(normalized_gripper_width, 2)}")
-        state = np.concatenate([ee_position,ee_rot6d,np.array([normalized_gripper_width]), [0]*10, np.array(arm_joint_positions)])
+        left_ee_pose_base = robot_obs["ee_pose"]["left_arm"]["base_frame"]
+        left_ee_position, left_ee_rot6d = left_ee_pose_base["position"], left_ee_pose_base["rot6d"]
+        right_ee_pose_base = robot_obs["ee_pose"]["right_arm"]["base_frame"]
+        right_ee_position, right_ee_rot6d = right_ee_pose_base["position"], right_ee_pose_base["rot6d"]
+        finger_positions = robot_obs["joint_positions"] # use finger joints(44) only
+        state = np.concatenate([left_ee_position, left_ee_rot6d, finger_positions[:22],
+                                right_ee_position, right_ee_rot6d, finger_positions[22:]], axis=-1) # (62,)
         rgb_data = {}
         for sensor_name in self.sensor_names:
             sensor_obs = benchmark_observation.get_sensor_observations(sensor_name)
             rgb_data[sensor_name] = sensor_obs["rgb"].data.cpu().numpy().astype(np.uint8)
         obs = {"state": state,"rgb": rgb_data,"prompt": self.prompt}
-
         return obs
 
     def compute_action(self, observation: dict) -> dict:
         if self.current_chunk_result is None:
-            self.current_ee_position_state = np.array(observation["state"][:3]).astype(np.float64)
-            self.current_ee_rot6d_state = np.array(observation["state"][3:9]).astype(np.float64)
-            self.current_gripper_width = np.array([observation["state"][9]])
+            self.current_state.update(self.split_joints(observation["state"]))
             payload = pickle.dumps(observation)
             response = requests.post(
                 f"{self.server_url}/inference",
@@ -151,7 +159,7 @@ class StarvlaPolicy(Policy):
             self.test_obs = observation["state"] #TODO
             self._handle_server_error(response)
             result = pickle.loads(response.content)
-            max_trunk_size = len(result["ee_delta_position_chunks"])
+            max_trunk_size = len(result["right_arm"]["ee_delta_position_chunks"])
             if self.run_trunk_size > max_trunk_size:
                 Log.warning(f"Run trunk size {self.run_trunk_size} is greater than the number of chunks {max_trunk_size}. Set run trunk size to {max_trunk_size}.")
                 self.run_trunk_size = max_trunk_size
@@ -161,57 +169,63 @@ class StarvlaPolicy(Policy):
             result = self.current_chunk_result
 
         return result
-    def __map_joint_position_to_normalized_width(self, joint_position: float) -> float:
-        if joint_position < 0:
-            joint_position = 0
-        if joint_position > 0.8:
-            joint_position = 0.8
-        for entry in self.gripper_width_mapper:
-            if round(entry["angel"], 2) == round(joint_position, 2):
-                return 1-(entry["width"] - self.min_width) / (self.max_width - self.min_width)
 
 
-
-    def __map_gripper_joint_position(self, normalized_gripper_width: float) -> float:
-        
-        joint_positions = []
-        joint_names = []
-        if normalized_gripper_width > 0.5:
-            for joint_name, joint_config in self.drive_joints.items():
-                joint_positions.append(joint_config.close_position)
-                joint_names.append(joint_name)
-        else:
-            for joint_name, joint_config in self.drive_joints.items():
-                joint_positions.append(joint_config.open_position)
-                joint_names.append(joint_name)
-        return joint_positions, joint_names
-
     def postprocess_action(self, action: dict) -> BenchmarkAction:
         benchmark_action = BenchmarkAction()
-        Log.debug(f"observation: {self.test_obs}")
-        # import ipdb;ipdb.set_trace()
+        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)
 
-        # get base frame end-effector pose
-        delta_ee_pose = Pose(position=action["ee_delta_position_chunks"][self.current_chunk_id], rot6d=action["ee_delta_rot6d_chunks"][self.current_chunk_id])
-        curr_state_ee_pose = Pose(position=self.current_ee_position_state, rot6d=self.current_ee_rot6d_state)
+        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
-        curr_action_gripper_width = action["gripper_width_chunks"][self.current_chunk_id]
-        
-        gripper_joint_positions, gripper_joint_names = self.__map_gripper_joint_position(curr_action_gripper_width[0])
-        Log.debug(f"action_gripper_joint_positions: {gripper_joint_positions}, action_normalized_gripper_width: {round(curr_action_gripper_width[0], 2)}")
+            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
+            state_arm = read_state[arm_key]
             benchmark_action.add_robot_action(
                 RobotAction(
                     control_mode=ControlMode.POSITION,
                     robot_name=self.robot_name,
-                joint_names=gripper_joint_names,
-                joint_positions=gripper_joint_positions
+                    joint_names=joint_names,
+                    # 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
                 )
             )
         self._visualize_base_frame_ee_poses(curr_state_ee_pose, curr_action_ee_pose)