update

benchmark.yaml

@@ -49,8 +49,8 @@ scene:
       - 0.001
       - 0.001
       position:
-      - 0.419859
+      - 0.219859
-      - -4.152430000000001
+      - -3.852430000000001
       - 0.510259093
       quaternion:
       - 0.7070997233636068
@@ -74,6 +74,14 @@ scene:
       stereotype: single_gripper_arm_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_joint2: -0.1633
         panda_joint4: -1.07
@@ -180,6 +188,13 @@ extension:
           name: Franka
           observe_ee_pose: true
           observe_gripper_drive_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: Hand_Camera
           observe_rgb: true
@@ -195,6 +210,7 @@ extension:
       stereotype: benchmark
       data_collector_name: benchmark_data_collect
       action_frequency: 15.0
+      timeout_per_episode: 30
       goals:
         - name: cola on top of book
           description: check if the cola bottle is on the book
@@ -208,13 +224,17 @@ extension:
         prompt: pick up the can
         run_trunk_size: 16
         gripper_width_mapper_file: ./gripper_width_robotiq_2f85_fixed.json
+        visualize_action_ee_pose: false
+        visualize_state_ee_pose: false
+        visualize_bounding_box_targets: [] # [omni6DPose_plug_001, omni6DPose_can_016] # 打开会被policy看到，会影响policy的推理结果
 
     recorder:
-      enable: false # set to true to record the data
+      enable: true # set to true to record the data
       stereotype: record
       data_collector_name: benchmark_data_collect
       record_fps: 30
       backend_root_path: output://benchmark_record
+      postprocess_list: ["hdf5", "video"]
 
 policy_server:
   ckpt_path: checkpoints://0318_qwenpi_droid_pretrain_8node/checkpoints/steps_30000_pytorch_model.pt

starvla_policy.py

@@ -7,7 +7,9 @@ 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.core.robots.configs.actuator_configs.grippers import GripperDriveJointConfig
+from joysim.controllers.visualize_controller import VisualizeController
+from joysim.utils.namespace import PoseVisualType, SimulatorType
+from joysim.unisim.robots.configs.actuator_configs.grippers import GripperDriveJointConfig
 from joysim.extensions.benchmark.action import RobotAction
 from joysim.extensions.benchmark.benchmark import (
     BenchmarkAction,
@@ -24,6 +26,13 @@ class StarvlaPolicyConfig(PolicyConfig):
 
     robot_name: str = field(default="None", required=True, comment="The name of the robot")
     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(
+        default_factory=list,
+        required=False,
+        comment="Spawnable object names to draw bounding boxes for (empty = off)",
+    )
     sensor_names: list[str] = field(
         default=["Hand_Camera", "Left_Camera", "Right_Camera"],
         required=True,
@@ -59,6 +68,10 @@ class StarvlaPolicy(Policy):
         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 [])
+
     def reset(self) -> None:
         self.current_ee_position_state = None
         self.current_ee_rot6d_state = None
@@ -70,7 +83,7 @@ class StarvlaPolicy(Policy):
         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": joint_config.position_control_effort_limit}).unwrap()
+            SpawnableController.control_robot(self.robot_name, "set_joint_effort_limit", parameters={"joint_names": [joint_name], "effort_limit": 50}).unwrap()
         self.max_width = float("-inf")
         self.min_width = float("inf")
         for entry in self.gripper_width_mapper:
@@ -103,6 +116,7 @@ class StarvlaPolicy(Policy):
         ee_pose_base = robot_obs["ee_pose"]["base_frame"]
         ee_position, ee_rot6d = ee_pose_base["position"],ee_pose_base["rot6d"]
         normalized_gripper_width = self.__map_joint_position_to_normalized_width(robot_obs["gripper_drive_state"]["position"][0])
+        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])])
         rgb_data = {}
         for sensor_name in self.sensor_names:
@@ -136,11 +150,15 @@ class StarvlaPolicy(Policy):
 
         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)
             
-        raise ValueError(f"Joint position {joint_position} not found in gripper width mapper")
+        
 
     def __map_gripper_joint_position(self, normalized_gripper_width: float) -> float:
         
@@ -182,8 +200,44 @@ class StarvlaPolicy(Policy):
                 ee_pose=curr_action_ee_pose
             )
         )
+        self._visualize_base_frame_ee_poses(curr_state_ee_pose, curr_action_ee_pose)
+        self._visualize_bounding_boxes()
         self.current_chunk_id += 1
         if self.current_chunk_id == self.run_trunk_size:
             self.current_chunk_id = 0
             self.current_chunk_result = None
         return benchmark_action
+
+    # ------------------- Visualization -------------------
+    def _visualize_base_frame_ee_poses(
+        self, pose_state_base: Pose, pose_action_base: Pose
+    ) -> None:
+        if not self.visualize_action_ee_pose and not self.visualize_state_ee_pose:
+            return
+        robot_base_world = SpawnableController.control_robot(
+            self.robot_name, "get_pose"
+        ).unwrap()
+        if self.visualize_state_ee_pose:
+            VisualizeController.create_pose_visualization(
+                robot_base_world * pose_state_base,
+                name=f"{self.robot_name}/starvla_state_ee",
+                simulator=SimulatorType.ISAACSIM,
+                pose_type=PoseVisualType.COORDINATE,
+                extra_params={"axis_length": 0.08, "thickness": 0.006},
+            ).unwrap()
+        if self.visualize_action_ee_pose:
+            VisualizeController.create_pose_visualization(
+                robot_base_world * pose_action_base,
+                name=f"{self.robot_name}/starvla_action_ee",
+                simulator=SimulatorType.ISAACSIM,
+                pose_type=PoseVisualType.COORDINATE,
+                extra_params={"axis_length": 0.1, "thickness": 0.006},
+            ).unwrap()
+
+    def _visualize_bounding_boxes(self) -> None:
+        if not self.visualize_bounding_box_targets:
+            return
+        for target_name in self.visualize_bounding_box_targets:
+            VisualizeController.visualize_target_bounding_box(
+                target_name, simulator=SimulatorType.ISAACSIM
+            ).unwrap()