From b08a62af89cff55b9a3f757a1ac203ae80bcb81d Mon Sep 17 00:00:00 2001 From: Pepijn Date: Thu, 2 Apr 2026 13:01:46 +0200 Subject: [PATCH] feat(examples): adapt UMI pi0 evaluate script for OpenArm follower MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Switch from SO100 to a single right OpenArm follower with one camera (cam0 at 960x720). Strip dataset recording — just execute the policy. Filter out .vel/.torque observation features for the EE pipeline. Made-with: Cursor --- examples/umi_pi0_relative_ee/evaluate.py | 137 +++++++++-------------- 1 file changed, 56 insertions(+), 81 deletions(-) diff --git a/examples/umi_pi0_relative_ee/evaluate.py b/examples/umi_pi0_relative_ee/evaluate.py index d77e3213d..e3e73029a 100644 --- a/examples/umi_pi0_relative_ee/evaluate.py +++ b/examples/umi_pi0_relative_ee/evaluate.py @@ -15,7 +15,9 @@ # limitations under the License. """ -Inference script for a pi0 model trained with **relative EE actions**. +Inference script for a pi0 model trained with **relative EE actions** on an OpenArm robot. + +Single right OpenArm follower with one wrist camera. This uses the built-in ``DeriveStateFromActionStep`` (no-op at inference), ``RelativeActionsProcessorStep``, ``AbsoluteActionsProcessorStep``, and @@ -23,7 +25,7 @@ This uses the built-in ``DeriveStateFromActionStep`` (no-op at inference), pipeline. The inference loop: - 1. Reads joint positions from the robot. + 1. Reads joint positions from the robot (7-DOF arm + gripper). 2. Converts to EE pose via forward kinematics (FK). This produces ``observation.state`` with the current EE pose. 3. The pi0 preprocessor: @@ -37,8 +39,6 @@ The inference loop: b) ``AbsoluteActionsProcessorStep`` adds cached state → absolute EE. 6. IK converts absolute EE → joint targets → robot. -Based on the so100_to_so100_EE/evaluate.py example. - Usage: python evaluate.py """ @@ -64,7 +64,7 @@ from lerobot.processor.converters import ( transition_to_observation, transition_to_robot_action, ) -from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig +from lerobot.robots.openarm_follower import OpenArmFollower, OpenArmFollowerConfig from lerobot.robots.so_follower.robot_kinematic_processor import ( ForwardKinematicsJointsToEE, InverseKinematicsEEToJoints, @@ -75,13 +75,11 @@ from lerobot.utils.control_utils import init_keyboard_listener from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun -NUM_EPISODES = 5 -FPS = 10 +FPS = 30 EPISODE_TIME_SEC = 60 -TASK_DESCRIPTION = "manipulation task" +TASK_DESCRIPTION = "red cube" -HF_MODEL_ID = "/" -HF_DATASET_ID = "/" +HF_MODEL_ID = "pepijn223/grabette-umi-pi0" # Latency compensation: skip this many steps from the start of each predicted # action chunk. Formula: ceil(total_latency_ms / (1000 / FPS)). @@ -91,24 +89,30 @@ LATENCY_SKIP_STEPS = 0 # EE feature keys produced by ForwardKinematicsJointsToEE EE_KEYS = ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"] +URDF_PATH = "./SO101/so101_new_calib.urdf" +URDF_EE_FRAME = "gripper_frame_link" + def main(): - camera_config = {"wrist": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)} - robot_config = SO100FollowerConfig( - port="/dev/tty.usbmodem5A460814411", - id="my_awesome_follower_arm", + camera_config = {"cam0": OpenCVCameraConfig(index_or_path=0, width=960, height=720, fps=FPS)} + robot_config = OpenArmFollowerConfig( + port="can0", + id="right_openarm", + side="right", cameras=camera_config, - use_degrees=True, + max_relative_target=8.0, ) - robot = SO100Follower(robot_config) + robot = OpenArmFollower(robot_config) policy = PI0Policy.from_pretrained(HF_MODEL_ID) policy.config.latency_skip_steps = LATENCY_SKIP_STEPS + motor_names = list(robot.bus.motors.keys()) + kinematics_solver = RobotKinematics( - urdf_path="./SO101/so101_new_calib.urdf", - target_frame_name="gripper_frame_link", - joint_names=list(robot.bus.motors.keys()), + urdf_path=URDF_PATH, + target_frame_name=URDF_EE_FRAME, + joint_names=motor_names, ) # FK: joint observation → EE observation (produces observation.state) @@ -116,7 +120,7 @@ def main(): steps=[ ForwardKinematicsJointsToEE( kinematics=kinematics_solver, - motor_names=list(robot.bus.motors.keys()), + motor_names=motor_names, ) ], to_transition=observation_to_transition, @@ -128,7 +132,7 @@ def main(): steps=[ InverseKinematicsEEToJoints( kinematics=kinematics_solver, - motor_names=list(robot.bus.motors.keys()), + motor_names=motor_names, initial_guess_current_joints=True, ), ], @@ -136,14 +140,23 @@ def main(): to_output=transition_to_robot_action, ) - # Dataset handle for stats (used by preprocessor/postprocessor) + # OpenArm observations include .vel and .torque per motor; the EE policy + # pipeline only needs .pos (converted to EE by FK) and camera features. + obs_features = { + k: v + for k, v in robot.observation_features.items() + if not (k.endswith(".vel") or k.endswith(".torque")) + } + + # A dataset object is needed for its .features and .meta.stats even when + # not recording — record_loop uses them for building observation/action frames. dataset = LeRobotDataset.create( - repo_id=HF_DATASET_ID, + repo_id="tmp/openarm_eval_scratch", fps=FPS, features=combine_feature_dicts( aggregate_pipeline_dataset_features( pipeline=robot_joints_to_ee_processor, - initial_features=create_initial_features(observation=robot.observation_features), + initial_features=create_initial_features(observation=obs_features), use_videos=True, ), aggregate_pipeline_dataset_features( @@ -159,10 +172,6 @@ def main(): image_writer_threads=4, ) - # Build pre/post processors from the trained model. - # The pi0 processor pipeline already includes: - # pre: ... → RelativeStateProcessorStep → RelativeActionsProcessorStep → NormalizerProcessorStep - # post: UnnormalizerProcessorStep → AbsoluteActionsProcessorStep → ... preprocessor, postprocessor = make_pre_post_processors( policy_cfg=policy, pretrained_path=HF_MODEL_ID, @@ -170,74 +179,40 @@ def main(): preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}}, ) - # Find the relative state step (if present) so we can reset its buffer between episodes. _relative_state_steps = [s for s in preprocessor.steps if isinstance(s, RelativeStateProcessorStep)] robot.connect() listener, events = init_keyboard_listener() - init_rerun(session_name="umi_pi0_relative_ee_evaluate") + init_rerun(session_name="openarm_umi_pi0_relative_ee_evaluate") try: if not robot.is_connected: raise ValueError("Robot is not connected!") - print("Starting evaluate loop...") - for episode_idx in range(NUM_EPISODES): - log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}") + log_say("Starting policy execution") + for step in _relative_state_steps: + step.reset() - # Reset relative state buffer so velocity is zero at episode start - for step in _relative_state_steps: - step.reset() - - record_loop( - robot=robot, - events=events, - fps=FPS, - policy=policy, - preprocessor=preprocessor, - postprocessor=postprocessor, - dataset=dataset, - control_time_s=EPISODE_TIME_SEC, - single_task=TASK_DESCRIPTION, - display_data=True, - teleop_action_processor=make_default_teleop_action_processor(), - robot_action_processor=robot_ee_to_joints_processor, - robot_observation_processor=robot_joints_to_ee_processor, - ) - - if not events["stop_recording"] and ( - (episode_idx < NUM_EPISODES - 1) or events["rerecord_episode"] - ): - log_say("Reset the environment") - record_loop( - robot=robot, - events=events, - fps=FPS, - control_time_s=EPISODE_TIME_SEC, - single_task=TASK_DESCRIPTION, - display_data=True, - teleop_action_processor=make_default_teleop_action_processor(), - robot_action_processor=robot_ee_to_joints_processor, - robot_observation_processor=robot_joints_to_ee_processor, - ) - - if events["rerecord_episode"]: - log_say("Re-record episode") - events["rerecord_episode"] = False - events["exit_early"] = False - dataset.clear_episode_buffer() - continue - - dataset.save_episode() + record_loop( + robot=robot, + events=events, + fps=FPS, + policy=policy, + preprocessor=preprocessor, + postprocessor=postprocessor, + dataset=dataset, + control_time_s=EPISODE_TIME_SEC, + single_task=TASK_DESCRIPTION, + display_data=True, + teleop_action_processor=make_default_teleop_action_processor(), + robot_action_processor=robot_ee_to_joints_processor, + robot_observation_processor=robot_joints_to_ee_processor, + ) finally: - log_say("Stop recording") robot.disconnect() listener.stop() - dataset.finalize() - dataset.push_to_hub() - if __name__ == "__main__": main()