From 7ce2823c56d33b66f350d9a1b5b8b20f3415dbf4 Mon Sep 17 00:00:00 2001 From: "hufei.hofee" Date: Tue, 24 Mar 2026 19:38:15 +0800 Subject: [PATCH] update --- benchmark.yaml | 26 +++++++++++++++++--- starvla_policy.py | 62 ++++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 81 insertions(+), 7 deletions(-) diff --git a/benchmark.yaml b/benchmark.yaml index e245ad2..18919bb 100644 --- a/benchmark.yaml +++ b/benchmark.yaml @@ -49,8 +49,8 @@ scene: - 0.001 - 0.001 position: - - 0.419859 - - -4.152430000000001 + - 0.219859 + - -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 diff --git a/starvla_policy.py b/starvla_policy.py index e7844f5..0f4c85a 100644 --- a/starvla_policy.py +++ b/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 \ No newline at end of file + 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() \ No newline at end of file