From 27a229ea648669b9ade6b8041eebe60ded030fa5 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Tue, 16 Sep 2025 14:56:36 +0200 Subject: [PATCH] chore(examples): homogenize style across example files (#1955) * chore(examples): homogenize style across example files * chore(examples): homogenize style across example files eval + replay * chore(examples): homogenize headers --- examples/lekiwi/evaluate.py | 66 +++++++++++++++++------ examples/lekiwi/record.py | 49 +++++++++++++---- examples/lekiwi/replay.py | 27 +++++++++- examples/lekiwi/teleoperate.py | 35 ++++++++++-- examples/phone_to_so100/evaluate.py | 64 ++++++++++++++++------ examples/phone_to_so100/record.py | 27 ++++++---- examples/phone_to_so100/replay.py | 28 +++++++--- examples/phone_to_so100/teleoperate.py | 30 ++++++++--- examples/so100_to_so100_EE/evaluate.py | 63 ++++++++++++++++------ examples/so100_to_so100_EE/record.py | 28 +++++++--- examples/so100_to_so100_EE/replay.py | 27 +++++++--- examples/so100_to_so100_EE/teleoperate.py | 43 ++++++++++----- 12 files changed, 373 insertions(+), 114 deletions(-) diff --git a/examples/lekiwi/evaluate.py b/examples/lekiwi/evaluate.py index c81130d4d..7c7f85f64 100644 --- a/examples/lekiwi/evaluate.py +++ b/examples/lekiwi/evaluate.py @@ -1,7 +1,24 @@ +# !/usr/bin/env python + +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.utils import hw_to_dataset_features from lerobot.policies.act.modeling_act import ACTPolicy from lerobot.policies.factory import make_pre_post_processors +from lerobot.processor import make_default_processors from lerobot.record import record_loop from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig from lerobot.utils.control_utils import init_keyboard_listener @@ -15,10 +32,12 @@ TASK_DESCRIPTION = "My task description" HF_MODEL_ID = "/" HF_DATASET_ID = "/" -# Create the robot and teleoperator configurations +# Create the robot configuration & robot robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi") + robot = LeKiwiClient(robot_config) +# Create policy policy = ACTPolicy.from_pretrained(HF_MODEL_ID) # Configure the dataset features @@ -36,41 +55,50 @@ dataset = LeRobotDataset.create( image_writer_threads=4, ) -# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi` -robot.connect() - -_init_rerun(session_name="recording") - -listener, events = init_keyboard_listener() - -if not robot.is_connected: - raise ValueError("Robot is not connected!") - +# Build Policy Processors preprocessor, postprocessor = make_pre_post_processors( policy_cfg=policy, pretrained_path=HF_MODEL_ID, dataset_stats=dataset.meta.stats, ) +# Connect the robot +# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi` +robot.connect() + +# TODO(Steven): Update this example to use pipelines +teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors() + +# Initialize the keyboard listener and rerun visualization +listener, events = init_keyboard_listener() +_init_rerun(session_name="lekiwi_evaluate") + +if not robot.is_connected: + raise ValueError("Robot is not connected!") + +print("Starting evaluate loop...") recorded_episodes = 0 while recorded_episodes < NUM_EPISODES and not events["stop_recording"]: log_say(f"Running inference, recording eval episode {recorded_episodes} of {NUM_EPISODES}") - # Run the policy inference loop + # Main record loop record_loop( robot=robot, events=events, fps=FPS, policy=policy, - preprocessor=preprocessor, + preprocessor=preprocessor, # Pass the pre and post policy processors postprocessor=postprocessor, dataset=dataset, control_time_s=EPISODE_TIME_SEC, single_task=TASK_DESCRIPTION, display_data=True, + teleop_action_processor=teleop_action_processor, + robot_action_processor=robot_action_processor, + robot_observation_processor=robot_observation_processor, ) - # Logic for reset env + # Reset the environment if not stopping or re-recording if not events["stop_recording"] and ( (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"] ): @@ -82,6 +110,9 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]: control_time_s=EPISODE_TIME_SEC, single_task=TASK_DESCRIPTION, display_data=True, + teleop_action_processor=teleop_action_processor, + robot_action_processor=robot_action_processor, + robot_observation_processor=robot_observation_processor, ) if events["rerecord_episode"]: @@ -91,11 +122,12 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]: dataset.clear_episode_buffer() continue + # Save episode dataset.save_episode() recorded_episodes += 1 -# Upload to hub and clean up -dataset.push_to_hub() - +# Clean up +log_say("Stop recording") robot.disconnect() listener.stop() +dataset.push_to_hub() diff --git a/examples/lekiwi/record.py b/examples/lekiwi/record.py index 11a716761..1344896ac 100644 --- a/examples/lekiwi/record.py +++ b/examples/lekiwi/record.py @@ -1,5 +1,22 @@ +# !/usr/bin/env python + +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.utils import hw_to_dataset_features +from lerobot.processor import make_default_processors from lerobot.record import record_loop from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient @@ -14,16 +31,21 @@ FPS = 30 EPISODE_TIME_SEC = 30 RESET_TIME_SEC = 10 TASK_DESCRIPTION = "My task description" +HF_REPO_ID = "/" # Create the robot and teleoperator configurations robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi") leader_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm") keyboard_config = KeyboardTeleopConfig() +# Initialize the robot and teleoperator robot = LeKiwiClient(robot_config) leader_arm = SO100Leader(leader_arm_config) keyboard = KeyboardTeleop(keyboard_config) +# TODO(Steven): Update this example to use pipelines +teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors() + # Configure the dataset features action_features = hw_to_dataset_features(robot.action_features, "action") obs_features = hw_to_dataset_features(robot.observation_features, "observation") @@ -31,7 +53,7 @@ dataset_features = {**action_features, **obs_features} # Create the dataset dataset = LeRobotDataset.create( - repo_id="/", + repo_id=HF_REPO_ID, fps=FPS, features=dataset_features, robot_type=robot.name, @@ -39,23 +61,25 @@ dataset = LeRobotDataset.create( image_writer_threads=4, ) +# Connect the robot and teleoperator # To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi` robot.connect() leader_arm.connect() keyboard.connect() +# Initialize the keyboard listener and rerun visualization +listener, events = init_keyboard_listener() _init_rerun(session_name="lekiwi_record") -listener, events = init_keyboard_listener() - if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected: - raise ValueError("Robot, leader arm of keyboard is not connected!") + raise ValueError("Robot or teleop is not connected!") +print("Starting record loop...") recorded_episodes = 0 while recorded_episodes < NUM_EPISODES and not events["stop_recording"]: log_say(f"Recording episode {recorded_episodes}") - # Run the record loop + # Main record loop record_loop( robot=robot, events=events, @@ -65,9 +89,12 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]: control_time_s=EPISODE_TIME_SEC, single_task=TASK_DESCRIPTION, display_data=True, + teleop_action_processor=teleop_action_processor, + robot_action_processor=robot_action_processor, + robot_observation_processor=robot_observation_processor, ) - # Logic for reset env + # Reset the environment if not stopping or re-recording if not events["stop_recording"] and ( (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"] ): @@ -80,6 +107,9 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]: control_time_s=RESET_TIME_SEC, single_task=TASK_DESCRIPTION, display_data=True, + teleop_action_processor=teleop_action_processor, + robot_action_processor=robot_action_processor, + robot_observation_processor=robot_observation_processor, ) if events["rerecord_episode"]: @@ -89,13 +119,14 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]: dataset.clear_episode_buffer() continue + # Save episode dataset.save_episode() recorded_episodes += 1 -# Upload to hub and clean up -dataset.push_to_hub() - +# Clean up +log_say("Stop recording") robot.disconnect() leader_arm.disconnect() keyboard.disconnect() listener.stop() +dataset.push_to_hub() diff --git a/examples/lekiwi/replay.py b/examples/lekiwi/replay.py index 248354df9..0ff549713 100644 --- a/examples/lekiwi/replay.py +++ b/examples/lekiwi/replay.py @@ -1,3 +1,19 @@ +# !/usr/bin/env python + +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import time from lerobot.datasets.lerobot_dataset import LeRobotDataset @@ -8,25 +24,34 @@ from lerobot.utils.utils import log_say EPISODE_IDX = 0 +# Initialize the robot config robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi") + +# Initialize the robot robot = LeKiwiClient(robot_config) +# Fetch the dataset to replay dataset = LeRobotDataset("/", episodes=[EPISODE_IDX]) actions = dataset.hf_dataset.select_columns("action") +# Connect to the robot robot.connect() if not robot.is_connected: raise ValueError("Robot is not connected!") +print("Starting replay loop...") log_say(f"Replaying episode {EPISODE_IDX}") for idx in range(dataset.num_frames): t0 = time.perf_counter() + # Get recorded action from dataset action = { name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"]) } - robot.send_action(action) + + # Send action to robot + _ = robot.send_action(action) busy_wait(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0)) diff --git a/examples/lekiwi/teleoperate.py b/examples/lekiwi/teleoperate.py index 45afca0cf..cde4000df 100644 --- a/examples/lekiwi/teleoperate.py +++ b/examples/lekiwi/teleoperate.py @@ -1,3 +1,19 @@ +# !/usr/bin/env python + +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import time from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig @@ -13,35 +29,44 @@ robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="my_lekiwi") teleop_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm") keyboard_config = KeyboardTeleopConfig(id="my_laptop_keyboard") +# Initialize the robot and teleoperator robot = LeKiwiClient(robot_config) leader_arm = SO100Leader(teleop_arm_config) keyboard = KeyboardTeleop(keyboard_config) +# Connect to the robot and teleoperator # To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi` robot.connect() leader_arm.connect() keyboard.connect() +# Init rerun viewer _init_rerun(session_name="lekiwi_teleop") if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected: - raise ValueError("Robot, leader arm of keyboard is not connected!") + raise ValueError("Robot or teleop is not connected!") +print("Starting teleop loop...") while True: t0 = time.perf_counter() + # Get robot observation observation = robot.get_observation() + # Get teleop action + # Arm arm_action = leader_arm.get_action() arm_action = {f"arm_{k}": v for k, v in arm_action.items()} - + # Keyboard keyboard_keys = keyboard.get_action() base_action = robot._from_keyboard_to_base_action(keyboard_keys) - log_rerun_data(observation=observation, action={**arm_action, **base_action}) - action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action - robot.send_action(action) + # Send action to robot + _ = robot.send_action(action) + + # Visualize + log_rerun_data(observation=observation, action=action) busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) diff --git a/examples/phone_to_so100/evaluate.py b/examples/phone_to_so100/evaluate.py index 30f098101..0703e01af 100644 --- a/examples/phone_to_so100/evaluate.py +++ b/examples/phone_to_so100/evaluate.py @@ -14,7 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.datasets.lerobot_dataset import LeRobotDataset @@ -54,7 +53,7 @@ TASK_DESCRIPTION = "My task description" HF_MODEL_ID = "/" HF_DATASET_ID = "/" -# Initialize the robot with degrees +# Create the robot configuration & robot camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)} robot_config = SO100FollowerConfig( port="/dev/tty.usbmodem58760434471", @@ -63,9 +62,11 @@ robot_config = SO100FollowerConfig( use_degrees=True, ) -# Initialize the robot robot = SO100Follower(robot_config) +# Create policy +policy = ACTPolicy.from_pretrained(HF_MODEL_ID) + # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf kinematics_solver = RobotKinematics( urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf", @@ -73,7 +74,7 @@ kinematics_solver = RobotKinematics( joint_names=list(robot.bus.motors.keys()), ) -# Build pipeline to convert ee pose action to joint action +# Build pipeline to convert EE action to joints action robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction]( steps=[ AddRobotObservationAsComplimentaryData(robot=robot), @@ -87,7 +88,7 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction]( to_output=transition_to_robot_action, ) -# Build pipeline to convert joint observation to ee pose observation +# Build pipeline to convert joints observation to EE observation robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation]( steps=[ ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) @@ -96,7 +97,6 @@ robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, Rob to_output=transition_to_observation, ) - # Create the dataset dataset = LeRobotDataset.create( repo_id=HF_DATASET_ID, @@ -125,31 +125,35 @@ dataset = LeRobotDataset.create( image_writer_threads=4, ) -# Initialize the keyboard listener and rerun visualization -_, events = init_keyboard_listener() -_init_rerun(session_name="recording_phone") - -# Connect the robot and teleoperator -robot.connect() - -episode_idx = 0 - -policy = ACTPolicy.from_pretrained(HF_MODEL_ID) +# Build Policy Processors preprocessor, postprocessor = make_pre_post_processors( policy_cfg=policy, pretrained_path=HF_MODEL_ID, dataset_stats=dataset.meta.stats, ) +# Connect the robot +robot.connect() + +# Initialize the keyboard listener and rerun visualization +listener, events = init_keyboard_listener() +_init_rerun(session_name="phone_so100_evaluate") + +if not robot.is_connected: + raise ValueError("Robot is not connected!") + +print("Starting evaluate loop...") +episode_idx = 0 for episode_idx in range(NUM_EPISODES): log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}") + # Main record loop record_loop( robot=robot, events=events, fps=FPS, policy=policy, - preprocessor=preprocessor, + preprocessor=preprocessor, # Pass the pre and post policy processors postprocessor=postprocessor, dataset=dataset, control_time_s=EPISODE_TIME_SEC, @@ -159,9 +163,35 @@ for episode_idx in range(NUM_EPISODES): robot_action_processor=robot_ee_to_joints_processor, robot_observation_processor=robot_joints_to_ee_pose_processor, ) + + # Reset the environment if not stopping or re-recording + 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_pose_processor, + ) + + if events["rerecord_episode"]: + log_say("Re-record episode") + events["rerecord_episode"] = False + events["exit_early"] = False + dataset.clear_episode_buffer() + continue + + # Save episode dataset.save_episode() + episode_idx += 1 # Clean up log_say("Stop recording") robot.disconnect() +listener.stop() dataset.push_to_hub() diff --git a/examples/phone_to_so100/record.py b/examples/phone_to_so100/record.py index 08f190e70..a205f6cc8 100644 --- a/examples/phone_to_so100/record.py +++ b/examples/phone_to_so100/record.py @@ -14,7 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features @@ -52,7 +51,7 @@ RESET_TIME_SEC = 30 TASK_DESCRIPTION = "My task description" HF_REPO_ID = "/" -# Initialize the robot and teleoperator +# Create the robot and teleoperator configurations camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)} robot_config = SO100FollowerConfig( port="/dev/tty.usbmodem58760434471", @@ -73,7 +72,7 @@ kinematics_solver = RobotKinematics( joint_names=list(robot.bus.motors.keys()), ) -# Build pipeline to convert phone action to ee pose action +# Build pipeline to convert phone action to EE action phone_to_robot_ee_pose_processor = RobotProcessorPipeline[RobotAction, RobotAction]( steps=[ MapPhoneActionToRobotAction(platform=teleop_config.phone_os), @@ -94,7 +93,7 @@ phone_to_robot_ee_pose_processor = RobotProcessorPipeline[RobotAction, RobotActi to_output=transition_to_robot_action, ) -# Build pipeline to convert ee pose action to joint action +# Build pipeline to convert EE action to joints action robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction]( steps=[ InverseKinematicsEEToJoints( @@ -107,7 +106,7 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction]( to_output=transition_to_robot_action, ) -# Build pipeline to convert joint observation to ee pose observation +# Build pipeline to convert joint observation to EE observation robot_joints_to_ee_pose = RobotProcessorPipeline[RobotObservation, RobotObservation]( steps=[ ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) @@ -121,6 +120,8 @@ dataset = LeRobotDataset.create( repo_id=HF_REPO_ID, fps=FPS, features=combine_feature_dicts( + # Run the feature contract of the pipelines + # This tells you how the features would look like after the pipeline steps aggregate_pipeline_dataset_features( pipeline=phone_to_robot_ee_pose_processor, initial_features=create_initial_features(action=phone.action_features), @@ -137,18 +138,24 @@ dataset = LeRobotDataset.create( image_writer_threads=4, ) -# Initialize the keyboard listener and rerun visualization -_, events = init_keyboard_listener() -_init_rerun(session_name="recording_phone") - # Connect the robot and teleoperator robot.connect() phone.connect() +# Initialize the keyboard listener and rerun visualization +listener, events = init_keyboard_listener() +_init_rerun(session_name="phone_so100_record") + +if not robot.is_connected or not phone.is_connected: + raise ValueError("Robot or teleop is not connected!") + + +print("Starting record loop. Move your phone to teleoperate the robot...") episode_idx = 0 while episode_idx < NUM_EPISODES and not events["stop_recording"]: log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}") + # Main record loop record_loop( robot=robot, events=events, @@ -186,6 +193,7 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]: dataset.clear_episode_buffer() continue + # Save episode dataset.save_episode() episode_idx += 1 @@ -193,4 +201,5 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]: log_say("Stop recording") robot.disconnect() phone.disconnect() +listener.stop() dataset.push_to_hub() diff --git a/examples/phone_to_so100/replay.py b/examples/phone_to_so100/replay.py index af28b0def..12485ecc1 100644 --- a/examples/phone_to_so100/replay.py +++ b/examples/phone_to_so100/replay.py @@ -14,7 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - import time from lerobot.datasets.lerobot_dataset import LeRobotDataset @@ -33,14 +32,13 @@ from lerobot.utils.utils import log_say EPISODE_IDX = 0 HF_REPO_ID = "/" +# Initialize the robot config robot_config = SO100FollowerConfig( port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", use_degrees=True ) -robot = SO100Follower(robot_config) -robot.connect() -dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX]) -actions = dataset.hf_dataset.select_columns("action") +# Initialize the robot +robot = SO100Follower(robot_config) # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf kinematics_solver = RobotKinematics( @@ -49,7 +47,7 @@ kinematics_solver = RobotKinematics( joint_names=list(robot.bus.motors.keys()), ) -# Build pipeline to convert ee pose action to joint action +# Build pipeline to convert EE action to joints action robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction]( steps=[ AddRobotObservationAsComplimentaryData(robot=robot), @@ -63,17 +61,33 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction]( to_output=transition_to_robot_action, ) +# Fetch the dataset to replay +dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX]) +actions = dataset.hf_dataset.select_columns("action") + +# Connect to the robot +robot.connect() + +if not robot.is_connected: + raise ValueError("Robot is not connected!") + +print("Starting replay loop...") log_say(f"Replaying episode {EPISODE_IDX}") for idx in range(dataset.num_frames): t0 = time.perf_counter() + # Get recorded action from dataset ee_action = { name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"]) } + # Dataset EE -> robot joints joint_action = robot_ee_to_joints_processor(ee_action) - action_sent = robot.send_action(joint_action) + + # Send action to robot + _ = robot.send_action(joint_action) busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0)) +# Clean up robot.disconnect() diff --git a/examples/phone_to_so100/teleoperate.py b/examples/phone_to_so100/teleoperate.py index f9aceeaa3..43862f1b5 100644 --- a/examples/phone_to_so100/teleoperate.py +++ b/examples/phone_to_so100/teleoperate.py @@ -1,6 +1,6 @@ -#!/usr/bin/env python +# !/usr/bin/env python -# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -30,6 +30,10 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction from lerobot.teleoperators.phone.teleop_phone import Phone +from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data + +FPS = 30 # Initialize the robot and teleoperator robot_config = SO100FollowerConfig( @@ -75,18 +79,30 @@ phone_to_robot_joints_processor = RobotProcessorPipeline[RobotAction, RobotActio to_output=transition_to_robot_action, ) +# Connect to the robot and teleoperator robot.connect() teleop_device.connect() -print("Starting teleop loop. Move your phone to teleoperate the robot.") +# Init rerun viewer +_init_rerun(session_name="phone_so100_teleop") + +if not robot.is_connected or not teleop_device.is_connected: + raise ValueError("Robot or teleop is not connected!") + +print("Starting teleop loop. Move your phone to teleoperate the robot...") while True: - # Get teleop observation + t0 = time.perf_counter() + + # Get teleop action phone_obs = teleop_device.get_action() # Phone -> EE pose -> Joints transition joint_action = phone_to_robot_joints_processor(phone_obs) - if joint_action: - robot.send_action(joint_action) + # Send action to robot + _ = robot.send_action(joint_action) - time.sleep(0.01) + # Visualize + log_rerun_data(observation=phone_obs, action=joint_action) + + busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) diff --git a/examples/so100_to_so100_EE/evaluate.py b/examples/so100_to_so100_EE/evaluate.py index 1ab28e575..155761f10 100644 --- a/examples/so100_to_so100_EE/evaluate.py +++ b/examples/so100_to_so100_EE/evaluate.py @@ -14,7 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.datasets.lerobot_dataset import LeRobotDataset @@ -54,7 +53,7 @@ TASK_DESCRIPTION = "My task description" HF_MODEL_ID = "/" HF_DATASET_ID = "/" -# Initialize the robot with degrees +# Create the robot configuration & robot camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)} robot_config = SO100FollowerConfig( port="/dev/tty.usbmodem5A460814411", @@ -63,9 +62,11 @@ robot_config = SO100FollowerConfig( use_degrees=True, ) -# Initialize the robot robot = SO100Follower(robot_config) +# Create policy +policy = ACTPolicy.from_pretrained(HF_MODEL_ID) + # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf kinematics_solver = RobotKinematics( urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf", @@ -73,7 +74,7 @@ kinematics_solver = RobotKinematics( joint_names=list(robot.bus.motors.keys()), ) -# Build pipeline to convert ee pose action to joint action +# Build pipeline to convert EE action to joints action robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction]( steps=[ AddRobotObservationAsComplimentaryData(robot=robot), @@ -87,7 +88,7 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction]( to_output=transition_to_robot_action, ) -# Build pipeline to convert joint observation to ee pose observation +# Build pipeline to convert joints observation to EE observation robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation]( steps=[ ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) @@ -125,31 +126,35 @@ dataset = LeRobotDataset.create( image_writer_threads=4, ) -# Initialize the keyboard listener and rerun visualization -_, events = init_keyboard_listener() -_init_rerun(session_name="recording_phone") - -# Connect the robot and teleoperator -robot.connect() - -episode_idx = 0 - -policy = ACTPolicy.from_pretrained(HF_MODEL_ID) +# Build Policy Processors preprocessor, postprocessor = make_pre_post_processors( policy_cfg=policy, pretrained_path=HF_MODEL_ID, dataset_stats=dataset.meta.stats, ) +# Connect the robot and teleoperator +robot.connect() + +# Initialize the keyboard listener and rerun visualization +listener, events = init_keyboard_listener() +_init_rerun(session_name="so100_so100_evaluate") + +if not robot.is_connected: + raise ValueError("Robot is not connected!") + +print("Starting evaluate loop...") +episode_idx = 0 for episode_idx in range(NUM_EPISODES): log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}") + # Main record loop record_loop( robot=robot, events=events, fps=FPS, policy=policy, - preprocessor=preprocessor, + preprocessor=preprocessor, # Pass the pre and post policy processors postprocessor=postprocessor, dataset=dataset, control_time_s=EPISODE_TIME_SEC, @@ -159,9 +164,35 @@ for episode_idx in range(NUM_EPISODES): robot_action_processor=robot_ee_to_joints_processor, robot_observation_processor=robot_joints_to_ee_pose_processor, ) + + # Reset the environment if not stopping or re-recording + 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_pose_processor, + ) + + if events["rerecord_episode"]: + log_say("Re-record episode") + events["rerecord_episode"] = False + events["exit_early"] = False + dataset.clear_episode_buffer() + continue + + # Save episode dataset.save_episode() + episode_idx += 1 # Clean up log_say("Stop recording") robot.disconnect() +listener.stop() dataset.push_to_hub() diff --git a/examples/so100_to_so100_EE/record.py b/examples/so100_to_so100_EE/record.py index f46b1c7f0..b6c137d6d 100644 --- a/examples/so100_to_so100_EE/record.py +++ b/examples/so100_to_so100_EE/record.py @@ -14,7 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import time from typing import Any from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig @@ -49,30 +48,34 @@ FPS = 30 EPISODE_TIME_SEC = 60 RESET_TIME_SEC = 30 TASK_DESCRIPTION = "My task description" -HF_REPO_ID = "imstevenpmwork/kin_chadil" + str(time.time()) +HF_REPO_ID = "/" +# Create the robot and teleoperator configurations camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)} follower_config = SO100FollowerConfig( port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", cameras=camera_config, use_degrees=True ) - leader_config = SO100LeaderConfig(port="/dev/tty.usbmodem5A460819811", id="my_awesome_leader_arm") +# Initialize the robot and teleoperator follower = SO100Follower(follower_config) leader = SO100Leader(leader_config) +# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf follower_kinematics_solver = RobotKinematics( urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf", target_frame_name="gripper_frame_link", joint_names=list(follower.bus.motors.keys()), ) +# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf leader_kinematics_solver = RobotKinematics( urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf", target_frame_name="gripper_frame_link", joint_names=list(leader.bus.motors.keys()), ) +# Build pipeline to convert follower joints to EE observation follower_joints_to_ee = RobotProcessorPipeline[dict[str, Any], dict[str, Any]]( steps=[ ForwardKinematicsJointsToEE( @@ -83,6 +86,7 @@ follower_joints_to_ee = RobotProcessorPipeline[dict[str, Any], dict[str, Any]]( to_output=transition_to_observation, ) +# Build pipeline to convert leader joints to EE action leader_joints_to_ee = RobotProcessorPipeline[RobotAction, RobotAction]( steps=[ ForwardKinematicsJointsToEE( @@ -93,7 +97,7 @@ leader_joints_to_ee = RobotProcessorPipeline[RobotAction, RobotAction]( to_output=transition_to_robot_action, ) - +# Build pipeline to convert EE action to follower joints ee_to_follower_joints = RobotProcessorPipeline[RobotAction, RobotAction]( [ AddRobotObservationAsComplimentaryData(robot=follower), @@ -116,6 +120,8 @@ dataset = LeRobotDataset.create( repo_id=HF_REPO_ID, fps=FPS, features=combine_feature_dicts( + # Run the feature contract of the pipelines + # This tells you how the features would look like after the pipeline steps aggregate_pipeline_dataset_features( pipeline=leader_joints_to_ee, initial_features=create_initial_features(action=leader.action_features), @@ -132,18 +138,24 @@ dataset = LeRobotDataset.create( image_writer_threads=4, ) -# Initialize the keyboard listener and rerun visualization -_, events = init_keyboard_listener() -_init_rerun(session_name="recording_phone") # Connect the robot and teleoperator leader.connect() follower.connect() +# Initialize the keyboard listener and rerun visualization +listener, events = init_keyboard_listener() +_init_rerun(session_name="recording_phone") + +if not leader.is_connected or not follower.is_connected: + raise ValueError("Robot or teleop is not connected!") + +print("Starting record loop...") episode_idx = 0 while episode_idx < NUM_EPISODES and not events["stop_recording"]: log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}") + # Main record loop record_loop( robot=follower, events=events, @@ -181,6 +193,7 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]: dataset.clear_episode_buffer() continue + # Save episode dataset.save_episode() episode_idx += 1 @@ -188,4 +201,5 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]: log_say("Stop recording") leader.disconnect() follower.disconnect() +listener.stop() dataset.push_to_hub() diff --git a/examples/so100_to_so100_EE/replay.py b/examples/so100_to_so100_EE/replay.py index 4954debc9..f8adf8d8a 100644 --- a/examples/so100_to_so100_EE/replay.py +++ b/examples/so100_to_so100_EE/replay.py @@ -33,14 +33,13 @@ from lerobot.utils.utils import log_say EPISODE_IDX = 0 HF_REPO_ID = "/" +# Initialize the robot config robot_config = SO100FollowerConfig( port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True ) -robot = SO100Follower(robot_config) -robot.connect() -dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX]) -actions = dataset.hf_dataset.select_columns("action") +# Initialize the robot +robot = SO100Follower(robot_config) # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf kinematics_solver = RobotKinematics( @@ -49,7 +48,7 @@ kinematics_solver = RobotKinematics( joint_names=list(robot.bus.motors.keys()), ) -# Build pipeline to convert ee pose action to joint action +# Build pipeline to convert EE action to joints action robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction]( steps=[ AddRobotObservationAsComplimentaryData(robot=robot), @@ -63,17 +62,33 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction]( to_output=transition_to_robot_action, ) +# Fetch the dataset to replay +dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX]) +actions = dataset.hf_dataset.select_columns("action") + +# Connect to the robot +robot.connect() + +if not robot.is_connected: + raise ValueError("Robot is not connected!") + +print("Starting replay loop...") log_say(f"Replaying episode {EPISODE_IDX}") for idx in range(dataset.num_frames): t0 = time.perf_counter() + # Get recorded action from dataset ee_action = { name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"]) } + # Dataset EE -> robot joints joint_action = robot_ee_to_joints_processor(ee_action) - action_sent = robot.send_action(joint_action) + + # Send action to robot + _ = robot.send_action(joint_action) busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0)) +# Clean up robot.disconnect() diff --git a/examples/so100_to_so100_EE/teleoperate.py b/examples/so100_to_so100_EE/teleoperate.py index db6da7444..00569bf46 100644 --- a/examples/so100_to_so100_EE/teleoperate.py +++ b/examples/so100_to_so100_EE/teleoperate.py @@ -1,6 +1,6 @@ -#!/usr/bin/env python +# !/usr/bin/env python -# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,7 +11,8 @@ # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specif +# See the License for the specific language governing permissions and +# limitations under the License. import time @@ -31,29 +32,36 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import ( from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader +from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data +FPS = 30 + +# Initialize the robot and teleoperator config follower_config = SO100FollowerConfig( port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True ) - leader_config = SO100LeaderConfig(port="/dev/tty.usbmodem5A460819811", id="my_awesome_leader_arm") +# Initialize the robot and teleoperator follower = SO100Follower(follower_config) leader = SO100Leader(leader_config) +# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf follower_kinematics_solver = RobotKinematics( urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf", target_frame_name="gripper_frame_link", joint_names=list(follower.bus.motors.keys()), ) +# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf leader_kinematics_solver = RobotKinematics( urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf", target_frame_name="gripper_frame_link", joint_names=list(leader.bus.motors.keys()), ) - +# Build pipeline to convert teleop joints to EE action leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction]( steps=[ ForwardKinematicsJointsToEE( @@ -64,6 +72,7 @@ leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction]( to_output=transition_to_robot_action, ) +# build pipeline to convert EE action to robot joints ee_to_follower_joints = RobotProcessorPipeline[RobotAction, RobotAction]( [ AddRobotObservationAsComplimentaryData(robot=follower), @@ -75,28 +84,36 @@ ee_to_follower_joints = RobotProcessorPipeline[RobotAction, RobotAction]( InverseKinematicsEEToJoints( kinematics=follower_kinematics_solver, motor_names=list(follower.bus.motors.keys()), - initial_guess_current_joints=False, ), ], to_transition=robot_action_to_transition, to_output=transition_to_robot_action, ) +# Connect to the robot and teleoperator follower.connect() leader.connect() -print("Starting teleop loop. Move your phone to teleoperate the robot.") +# Init rerun viewer +_init_rerun(session_name="so100_so100_EE_teleop") + +print("Starting teleop loop...") while True: - # Get leader joints observations + t0 = time.perf_counter() + + # Get teleop observation leader_joints_obs = leader.get_action() - # Convert them to EE + # teleop joints -> teleop EE action leader_ee_act = leader_to_ee(leader_joints_obs) - # Convert EE to follower joints actions + # teleop EE -> robot joints follower_joints_act = ee_to_follower_joints(leader_ee_act) - if follower_joints_act: - follower.send_action(follower_joints_act) + # Send action to robot + _ = follower.send_action(follower_joints_act) - time.sleep(0.01) + # Visualize + log_rerun_data(observation=leader_ee_act, action=follower_joints_act) + + busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))