admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-24 18:26:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(examples): wrap all of them into a main function (#2524)

Browse Source
This commit is contained in:
Steven Palma
2025-11-26 14:28:04 +01:00
committed by GitHub
parent 87bee86640
commit 17581a9449
24 changed files with 1672 additions and 1532 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/async.mdx
+1 -1
View File
@@ -196,7 +196,7 @@ client_cfg = RobotClientConfig(
server_address="localhost:8080", server_address="localhost:8080",
policy_device="mps", policy_device="mps",
policy_type="smolvla", policy_type="smolvla",
pretrained_name_or_path="fracapuano/smolvla_async", pretrained_name_or_path="<user>/smolvla_async",
chunk_size_threshold=0.5, chunk_size_threshold=0.5,
actions_per_chunk=50, # make sure this is less than the max actions of the policy actions_per_chunk=50, # make sure this is less than the max actions of the policy
) )
examples/dataset/load_lerobot_dataset.py
+80 -75
View File
@@ -34,105 +34,106 @@ from huggingface_hub import HfApi
import lerobot import lerobot
from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
# We ported a number of existing datasets ourselves, use this to see the list:
print("List of available datasets:")
pprint(lerobot.available_datasets)
# You can also browse through the datasets created/ported by the community on the hub using the hub api: def main():
hub_api = HfApi() # We ported a number of existing datasets ourselves, use this to see the list:
repo_ids = [info.id for info in hub_api.list_datasets(task_categories="robotics", tags=["LeRobot"])] print("List of available datasets:")
pprint(repo_ids) pprint(lerobot.available_datasets)
# Or simply explore them in your web browser directly at: # You can also browse through the datasets created/ported by the community on the hub using the hub api:
# https://huggingface.co/datasets?other=LeRobot hub_api = HfApi()
repo_ids = [info.id for info in hub_api.list_datasets(task_categories="robotics", tags=["LeRobot"])]
pprint(repo_ids)
# Let's take this one for this example # Or simply explore them in your web browser directly at:
repo_id = "lerobot/aloha_mobile_cabinet" # https://huggingface.co/datasets?other=LeRobot
# We can have a look and fetch its metadata to know more about it:
ds_meta = LeRobotDatasetMetadata(repo_id)
# By instantiating just this class, you can quickly access useful information about the content and the # Let's take this one for this example
# structure of the dataset without downloading the actual data yet (only metadata files — which are repo_id = "lerobot/aloha_mobile_cabinet"
# lightweight). # We can have a look and fetch its metadata to know more about it:
print(f"Total number of episodes: {ds_meta.total_episodes}") ds_meta = LeRobotDatasetMetadata(repo_id)
print(f"Average number of frames per episode: {ds_meta.total_frames / ds_meta.total_episodes:.3f}")
print(f"Frames per second used during data collection: {ds_meta.fps}")
print(f"Robot type: {ds_meta.robot_type}")
print(f"keys to access images from cameras: {ds_meta.camera_keys=}\n")
print("Tasks:") # By instantiating just this class, you can quickly access useful information about the content and the
print(ds_meta.tasks) # structure of the dataset without downloading the actual data yet (only metadata files — which are
print("Features:") # lightweight).
pprint(ds_meta.features) print(f"Total number of episodes: {ds_meta.total_episodes}")
print(f"Average number of frames per episode: {ds_meta.total_frames / ds_meta.total_episodes:.3f}")
print(f"Frames per second used during data collection: {ds_meta.fps}")
print(f"Robot type: {ds_meta.robot_type}")
print(f"keys to access images from cameras: {ds_meta.camera_keys=}\n")
# You can also get a short summary by simply printing the object: print("Tasks:")
print(ds_meta) print(ds_meta.tasks)
print("Features:")
pprint(ds_meta.features)
# You can then load the actual dataset from the hub. # You can also get a short summary by simply printing the object:
# Either load any subset of episodes: print(ds_meta)
dataset = LeRobotDataset(repo_id, episodes=[0, 10, 11, 23])
# And see how many frames you have: # You can then load the actual dataset from the hub.
print(f"Selected episodes: {dataset.episodes}") # Either load any subset of episodes:
print(f"Number of episodes selected: {dataset.num_episodes}") dataset = LeRobotDataset(repo_id, episodes=[0, 10, 11, 23])
print(f"Number of frames selected: {dataset.num_frames}")
# Or simply load the entire dataset: # And see how many frames you have:
dataset = LeRobotDataset(repo_id) print(f"Selected episodes: {dataset.episodes}")
print(f"Number of episodes selected: {dataset.num_episodes}") print(f"Number of episodes selected: {dataset.num_episodes}")
print(f"Number of frames selected: {dataset.num_frames}") print(f"Number of frames selected: {dataset.num_frames}")
# The previous metadata class is contained in the 'meta' attribute of the dataset: # Or simply load the entire dataset:
print(dataset.meta) dataset = LeRobotDataset(repo_id)
print(f"Number of episodes selected: {dataset.num_episodes}")
print(f"Number of frames selected: {dataset.num_frames}")
# LeRobotDataset actually wraps an underlying Hugging Face dataset # The previous metadata class is contained in the 'meta' attribute of the dataset:
# (see https://huggingface.co/docs/datasets for more information). print(dataset.meta)
print(dataset.hf_dataset)
# LeRobot datasets also subclasses PyTorch datasets so you can do everything you know and love from working # LeRobotDataset actually wraps an underlying Hugging Face dataset
# with the latter, like iterating through the dataset. # (see https://huggingface.co/docs/datasets for more information).
# The __getitem__ iterates over the frames of the dataset. Since our datasets are also structured by print(dataset.hf_dataset)
# episodes, you can access the frame indices of any episode using dataset.meta.episodes. Here, we access
# frame indices associated to the first episode:
episode_index = 0
from_idx = dataset.meta.episodes["dataset_from_index"][episode_index]
to_idx = dataset.meta.episodes["dataset_to_index"][episode_index]
# Then we grab all the image frames from the first camera: # LeRobot datasets also subclasses PyTorch datasets so you can do everything you know and love from working
camera_key = dataset.meta.camera_keys[0] # with the latter, like iterating through the dataset.
frames = [dataset[idx][camera_key] for idx in range(from_idx, to_idx)] # The __getitem__ iterates over the frames of the dataset. Since our datasets are also structured by
# episodes, you can access the frame indices of any episode using dataset.meta.episodes. Here, we access
# frame indices associated to the first episode:
episode_index = 0
from_idx = dataset.meta.episodes["dataset_from_index"][episode_index]
to_idx = dataset.meta.episodes["dataset_to_index"][episode_index]
# The objects returned by the dataset are all torch.Tensors # Then we grab all the image frames from the first camera:
print(type(frames[0])) camera_key = dataset.meta.camera_keys[0]
print(frames[0].shape) frames = [dataset[idx][camera_key] for idx in range(from_idx, to_idx)]
# Since we're using pytorch, the shape is in pytorch, channel-first convention (c, h, w). # The objects returned by the dataset are all torch.Tensors
# We can compare this shape with the information available for that feature print(type(frames[0]))
pprint(dataset.features[camera_key]) print(frames[0].shape)
# In particular:
print(dataset.features[camera_key]["shape"])
# The shape is in (h, w, c) which is a more universal format.
# For many machine learning applications we need to load the history of past observations or trajectories of # Since we're using pytorch, the shape is in pytorch, channel-first convention (c, h, w).
# future actions. Our datasets can load previous and future frames for each key/modality, using timestamps # We can compare this shape with the information available for that feature
# differences with the current loaded frame. For instance: pprint(dataset.features[camera_key])
delta_timestamps = { # In particular:
print(dataset.features[camera_key]["shape"])
# The shape is in (h, w, c) which is a more universal format.
# For many machine learning applications we need to load the history of past observations or trajectories of
# future actions. Our datasets can load previous and future frames for each key/modality, using timestamps
# differences with the current loaded frame. For instance:
delta_timestamps = {
# loads 4 images: 1 second before current frame, 500 ms before, 200 ms before, and current frame # loads 4 images: 1 second before current frame, 500 ms before, 200 ms before, and current frame
camera_key: [-1, -0.5, -0.20, 0], camera_key: [-1, -0.5, -0.20, 0],
# loads 6 state vectors: 1.5 seconds before, 1 second before, ... 200 ms, 100 ms, and current frame # loads 6 state vectors: 1.5 seconds before, 1 second before, ... 200 ms, 100 ms, and current frame
"observation.state": [-1.5, -1, -0.5, -0.20, -0.10, 0], "observation.state": [-1.5, -1, -0.5, -0.20, -0.10, 0],
# loads 64 action vectors: current frame, 1 frame in the future, 2 frames, ... 63 frames in the future # loads 64 action vectors: current frame, 1 frame in the future, 2 frames, ... 63 frames in the future
"action": [t / dataset.fps for t in range(64)], "action": [t / dataset.fps for t in range(64)],
} }
# Note that in any case, these delta_timestamps values need to be multiples of (1/fps) so that added to any # Note that in any case, these delta_timestamps values need to be multiples of (1/fps) so that added to any
# timestamp, you still get a valid timestamp. # timestamp, you still get a valid timestamp.
dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps) dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps)
print(f"\n{dataset[0][camera_key].shape=}") # (4, c, h, w) print(f"\n{dataset[0][camera_key].shape=}") # (4, c, h, w)
print(f"{dataset[0]['observation.state'].shape=}") # (6, c) print(f"{dataset[0]['observation.state'].shape=}") # (6, c)
print(f"{dataset[0]['action'].shape=}\n") # (64, c) print(f"{dataset[0]['action'].shape=}\n") # (64, c)
if __name__ == "__main__":
dataloader = torch.utils.data.DataLoader( dataloader = torch.utils.data.DataLoader(
dataset, dataset,
num_workers=4, num_workers=4,
@@ -144,3 +145,7 @@ if __name__ == "__main__":
print(f"{batch['observation.state'].shape=}") # (32, 6, c) print(f"{batch['observation.state'].shape=}") # (32, 6, c)
print(f"{batch['action'].shape=}") # (32, 64, c) print(f"{batch['action'].shape=}") # (32, 64, c)
break break
if __name__ == "__main__":
main()
examples/lekiwi/evaluate.py
+39 -33
View File
@@ -33,55 +33,57 @@ TASK_DESCRIPTION = "My task description"
HF_MODEL_ID = "<hf_username>/<model_repo_id>" HF_MODEL_ID = "<hf_username>/<model_repo_id>"
HF_DATASET_ID = "<hf_username>/<eval_dataset_repo_id>" HF_DATASET_ID = "<hf_username>/<eval_dataset_repo_id>"
# Create the robot configuration & robot
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
robot = LeKiwiClient(robot_config) def main():
# Create the robot configuration & robot
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
# Create policy robot = LeKiwiClient(robot_config)
policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
# Configure the dataset features # Create policy
action_features = hw_to_dataset_features(robot.action_features, ACTION) policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
dataset_features = {**action_features, **obs_features}
# Create the dataset # Configure the dataset features
dataset = LeRobotDataset.create( action_features = hw_to_dataset_features(robot.action_features, ACTION)
obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
dataset_features = {**action_features, **obs_features}
# Create the dataset
dataset = LeRobotDataset.create(
repo_id=HF_DATASET_ID, repo_id=HF_DATASET_ID,
fps=FPS, fps=FPS,
features=dataset_features, features=dataset_features,
robot_type=robot.name, robot_type=robot.name,
use_videos=True, use_videos=True,
image_writer_threads=4, image_writer_threads=4,
) )
# Build Policy Processors # Build Policy Processors
preprocessor, postprocessor = make_pre_post_processors( preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=policy, policy_cfg=policy,
pretrained_path=HF_MODEL_ID, pretrained_path=HF_MODEL_ID,
dataset_stats=dataset.meta.stats, dataset_stats=dataset.meta.stats,
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility. # The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}}, preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
) )
# Connect the robot # 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` # 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() robot.connect()
# TODO(Steven): Update this example to use pipelines # TODO(Steven): Update this example to use pipelines
teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors() teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
# Initialize the keyboard listener and rerun visualization # Initialize the keyboard listener and rerun visualization
listener, events = init_keyboard_listener() listener, events = init_keyboard_listener()
init_rerun(session_name="lekiwi_evaluate") init_rerun(session_name="lekiwi_evaluate")
if not robot.is_connected: if not robot.is_connected:
raise ValueError("Robot is not connected!") raise ValueError("Robot is not connected!")
print("Starting evaluate loop...") print("Starting evaluate loop...")
recorded_episodes = 0 recorded_episodes = 0
while recorded_episodes < NUM_EPISODES and not events["stop_recording"]: while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
log_say(f"Running inference, recording eval episode {recorded_episodes} of {NUM_EPISODES}") log_say(f"Running inference, recording eval episode {recorded_episodes} of {NUM_EPISODES}")
# Main record loop # Main record loop
@@ -129,10 +131,14 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
dataset.save_episode() dataset.save_episode()
recorded_episodes += 1 recorded_episodes += 1
# Clean up # Clean up
log_say("Stop recording") log_say("Stop recording")
robot.disconnect() robot.disconnect()
listener.stop() listener.stop()
dataset.finalize() dataset.finalize()
dataset.push_to_hub() dataset.push_to_hub()
if __name__ == "__main__":
main()
examples/lekiwi/record.py
+43 -37
View File
@@ -34,50 +34,52 @@ RESET_TIME_SEC = 10
TASK_DESCRIPTION = "My task description" TASK_DESCRIPTION = "My task description"
HF_REPO_ID = "<hf_username>/<dataset_repo_id>" HF_REPO_ID = "<hf_username>/<dataset_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 def main():
robot = LeKiwiClient(robot_config) # Create the robot and teleoperator configurations
leader_arm = SO100Leader(leader_arm_config) robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
keyboard = KeyboardTeleop(keyboard_config) leader_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm")
keyboard_config = KeyboardTeleopConfig()
# TODO(Steven): Update this example to use pipelines # Initialize the robot and teleoperator
teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors() robot = LeKiwiClient(robot_config)
leader_arm = SO100Leader(leader_arm_config)
keyboard = KeyboardTeleop(keyboard_config)
# Configure the dataset features # TODO(Steven): Update this example to use pipelines
action_features = hw_to_dataset_features(robot.action_features, ACTION) teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
dataset_features = {**action_features, **obs_features}
# Create the dataset # Configure the dataset features
dataset = LeRobotDataset.create( action_features = hw_to_dataset_features(robot.action_features, ACTION)
obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
dataset_features = {**action_features, **obs_features}
# Create the dataset
dataset = LeRobotDataset.create(
repo_id=HF_REPO_ID, repo_id=HF_REPO_ID,
fps=FPS, fps=FPS,
features=dataset_features, features=dataset_features,
robot_type=robot.name, robot_type=robot.name,
use_videos=True, use_videos=True,
image_writer_threads=4, image_writer_threads=4,
) )
# Connect the robot and teleoperator # 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` # 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() robot.connect()
leader_arm.connect() leader_arm.connect()
keyboard.connect() keyboard.connect()
# Initialize the keyboard listener and rerun visualization # Initialize the keyboard listener and rerun visualization
listener, events = init_keyboard_listener() listener, events = init_keyboard_listener()
init_rerun(session_name="lekiwi_record") init_rerun(session_name="lekiwi_record")
if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected: if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
raise ValueError("Robot or teleop is not connected!") raise ValueError("Robot or teleop is not connected!")
print("Starting record loop...") print("Starting record loop...")
recorded_episodes = 0 recorded_episodes = 0
while recorded_episodes < NUM_EPISODES and not events["stop_recording"]: while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
log_say(f"Recording episode {recorded_episodes}") log_say(f"Recording episode {recorded_episodes}")
# Main record loop # Main record loop
@@ -124,12 +126,16 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
dataset.save_episode() dataset.save_episode()
recorded_episodes += 1 recorded_episodes += 1
# Clean up # Clean up
log_say("Stop recording") log_say("Stop recording")
robot.disconnect() robot.disconnect()
leader_arm.disconnect() leader_arm.disconnect()
keyboard.disconnect() keyboard.disconnect()
listener.stop() listener.stop()
dataset.finalize() dataset.finalize()
dataset.push_to_hub() dataset.push_to_hub()
if __name__ == "__main__":
main()
examples/lekiwi/replay.py
+22 -16
View File
@@ -25,27 +25,29 @@ from lerobot.utils.utils import log_say
EPISODE_IDX = 0 EPISODE_IDX = 0
# Initialize the robot config
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
# Initialize the robot def main():
robot = LeKiwiClient(robot_config) # Initialize the robot config
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
# Fetch the dataset to replay # Initialize the robot
dataset = LeRobotDataset("<hf_username>/<dataset_repo_id>", episodes=[EPISODE_IDX]) robot = LeKiwiClient(robot_config)
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
actions = episode_frames.select_columns(ACTION)
# Connect to the robot # Fetch the dataset to replay
robot.connect() dataset = LeRobotDataset("<hf_username>/<dataset_repo_id>", episodes=[EPISODE_IDX])
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
actions = episode_frames.select_columns(ACTION)
if not robot.is_connected: # Connect to the robot
robot.connect()
if not robot.is_connected:
raise ValueError("Robot is not connected!") raise ValueError("Robot is not connected!")
print("Starting replay loop...") print("Starting replay loop...")
log_say(f"Replaying episode {EPISODE_IDX}") log_say(f"Replaying episode {EPISODE_IDX}")
for idx in range(len(episode_frames)): for idx in range(len(episode_frames)):
t0 = time.perf_counter() t0 = time.perf_counter()
# Get recorded action from dataset # Get recorded action from dataset
@@ -58,4 +60,8 @@ for idx in range(len(episode_frames)):
busy_wait(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0)) busy_wait(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0))
robot.disconnect() robot.disconnect()
if __name__ == "__main__":
main()
examples/lekiwi/teleoperate.py
+24 -18
View File
@@ -24,30 +24,32 @@ from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
FPS = 30 FPS = 30
# Create the robot and teleoperator configurations
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 def main():
robot = LeKiwiClient(robot_config) # Create the robot and teleoperator configurations
leader_arm = SO100Leader(teleop_arm_config) robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="my_lekiwi")
keyboard = KeyboardTeleop(keyboard_config) teleop_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm")
keyboard_config = KeyboardTeleopConfig(id="my_laptop_keyboard")
# Connect to the robot and teleoperator # Initialize 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 = LeKiwiClient(robot_config)
robot.connect() leader_arm = SO100Leader(teleop_arm_config)
leader_arm.connect() keyboard = KeyboardTeleop(keyboard_config)
keyboard.connect()
# Init rerun viewer # Connect to the robot and teleoperator
init_rerun(session_name="lekiwi_teleop") # 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()
if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected: # 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 or teleop is not connected!") raise ValueError("Robot or teleop is not connected!")
print("Starting teleop loop...") print("Starting teleop loop...")
while True: while True:
t0 = time.perf_counter() t0 = time.perf_counter()
# Get robot observation # Get robot observation
@@ -70,3 +72,7 @@ while True:
log_rerun_data(observation=observation, action=action) log_rerun_data(observation=observation, action=action)
busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
if __name__ == "__main__":
main()
examples/phone_to_so100/evaluate.py
+46 -38
View File
@@ -52,29 +52,31 @@ TASK_DESCRIPTION = "My task description"
HF_MODEL_ID = "<hf_username>/<model_repo_id>" HF_MODEL_ID = "<hf_username>/<model_repo_id>"
HF_DATASET_ID = "<hf_username>/<dataset_repo_id>" HF_DATASET_ID = "<hf_username>/<dataset_repo_id>"
# Create the robot configuration & robot
camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)} def main():
robot_config = SO100FollowerConfig( # 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", port="/dev/tty.usbmodem58760434471",
id="my_awesome_follower_arm", id="my_awesome_follower_arm",
cameras=camera_config, cameras=camera_config,
use_degrees=True, use_degrees=True,
) )
robot = SO100Follower(robot_config) robot = SO100Follower(robot_config)
# Create policy # Create policy
policy = ACTPolicy.from_pretrained(HF_MODEL_ID) 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 # 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( kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()), joint_names=list(robot.bus.motors.keys()),
) )
# Build pipeline to convert EE action to joints action # Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
steps=[ steps=[
InverseKinematicsEEToJoints( InverseKinematicsEEToJoints(
kinematics=kinematics_solver, kinematics=kinematics_solver,
@@ -84,19 +86,21 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Build pipeline to convert joints observation to EE observation # Build pipeline to convert joints observation to EE observation
robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation]( robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
steps=[ steps=[
ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) ForwardKinematicsJointsToEE(
kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())
)
], ],
to_transition=observation_to_transition, to_transition=observation_to_transition,
to_output=transition_to_observation, to_output=transition_to_observation,
) )
# Create the dataset # Create the dataset
dataset = LeRobotDataset.create( dataset = LeRobotDataset.create(
repo_id=HF_DATASET_ID, repo_id=HF_DATASET_ID,
fps=FPS, fps=FPS,
features=combine_feature_dicts( features=combine_feature_dicts(
@@ -121,30 +125,30 @@ dataset = LeRobotDataset.create(
robot_type=robot.name, robot_type=robot.name,
use_videos=True, use_videos=True,
image_writer_threads=4, image_writer_threads=4,
) )
# Build Policy Processors # Build Policy Processors
preprocessor, postprocessor = make_pre_post_processors( preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=policy, policy_cfg=policy,
pretrained_path=HF_MODEL_ID, pretrained_path=HF_MODEL_ID,
dataset_stats=dataset.meta.stats, dataset_stats=dataset.meta.stats,
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility. # The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}}, preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
) )
# Connect the robot # Connect the robot
robot.connect() robot.connect()
# Initialize the keyboard listener and rerun visualization # Initialize the keyboard listener and rerun visualization
listener, events = init_keyboard_listener() listener, events = init_keyboard_listener()
init_rerun(session_name="phone_so100_evaluate") init_rerun(session_name="phone_so100_evaluate")
if not robot.is_connected: if not robot.is_connected:
raise ValueError("Robot is not connected!") raise ValueError("Robot is not connected!")
print("Starting evaluate loop...") print("Starting evaluate loop...")
episode_idx = 0 episode_idx = 0
for episode_idx in range(NUM_EPISODES): for episode_idx in range(NUM_EPISODES):
log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}") log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
# Main record loop # Main record loop
@@ -190,10 +194,14 @@ for episode_idx in range(NUM_EPISODES):
dataset.save_episode() dataset.save_episode()
episode_idx += 1 episode_idx += 1
# Clean up # Clean up
log_say("Stop recording") log_say("Stop recording")
robot.disconnect() robot.disconnect()
listener.stop() listener.stop()
dataset.finalize() dataset.finalize()
dataset.push_to_hub() dataset.push_to_hub()
if __name__ == "__main__":
main()
examples/phone_to_so100/record.py
+51 -42
View File
@@ -50,29 +50,33 @@ RESET_TIME_SEC = 30
TASK_DESCRIPTION = "My task description" TASK_DESCRIPTION = "My task description"
HF_REPO_ID = "<hf_username>/<dataset_repo_id>" HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
# Create the robot and teleoperator configurations
camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)} def main():
robot_config = SO100FollowerConfig( # 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.usbmodem5A460814411", port="/dev/tty.usbmodem5A460814411",
id="my_awesome_follower_arm", id="my_awesome_follower_arm",
cameras=camera_config, cameras=camera_config,
use_degrees=True, use_degrees=True,
) )
teleop_config = PhoneConfig(phone_os=PhoneOS.IOS) # or PhoneOS.ANDROID teleop_config = PhoneConfig(phone_os=PhoneOS.IOS) # or PhoneOS.ANDROID
# Initialize the robot and teleoperator # Initialize the robot and teleoperator
robot = SO100Follower(robot_config) robot = SO100Follower(robot_config)
phone = Phone(teleop_config) phone = Phone(teleop_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 # 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( kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()), joint_names=list(robot.bus.motors.keys()),
) )
# Build pipeline to convert phone action to EE action # Build pipeline to convert phone action to EE action
phone_to_robot_ee_pose_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( phone_to_robot_ee_pose_processor = RobotProcessorPipeline[
tuple[RobotAction, RobotObservation], RobotAction
](
steps=[ steps=[
MapPhoneActionToRobotAction(platform=teleop_config.phone_os), MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
EEReferenceAndDelta( EEReferenceAndDelta(
@@ -89,10 +93,10 @@ phone_to_robot_ee_pose_processor = RobotProcessorPipeline[tuple[RobotAction, Rob
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Build pipeline to convert EE action to joints action # Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
steps=[ steps=[
InverseKinematicsEEToJoints( InverseKinematicsEEToJoints(
kinematics=kinematics_solver, kinematics=kinematics_solver,
@@ -102,19 +106,21 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Build pipeline to convert joint observation to EE observation # Build pipeline to convert joint observation to EE observation
robot_joints_to_ee_pose = RobotProcessorPipeline[RobotObservation, RobotObservation]( robot_joints_to_ee_pose = RobotProcessorPipeline[RobotObservation, RobotObservation](
steps=[ steps=[
ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) ForwardKinematicsJointsToEE(
kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())
)
], ],
to_transition=observation_to_transition, to_transition=observation_to_transition,
to_output=transition_to_observation, to_output=transition_to_observation,
) )
# Create the dataset # Create the dataset
dataset = LeRobotDataset.create( dataset = LeRobotDataset.create(
repo_id=HF_REPO_ID, repo_id=HF_REPO_ID,
fps=FPS, fps=FPS,
features=combine_feature_dicts( features=combine_feature_dicts(
@@ -134,23 +140,22 @@ dataset = LeRobotDataset.create(
robot_type=robot.name, robot_type=robot.name,
use_videos=True, use_videos=True,
image_writer_threads=4, image_writer_threads=4,
) )
# Connect the robot and teleoperator # Connect the robot and teleoperator
robot.connect() robot.connect()
phone.connect() phone.connect()
# Initialize the keyboard listener and rerun visualization # Initialize the keyboard listener and rerun visualization
listener, events = init_keyboard_listener() listener, events = init_keyboard_listener()
init_rerun(session_name="phone_so100_record") init_rerun(session_name="phone_so100_record")
if not robot.is_connected or not phone.is_connected: if not robot.is_connected or not phone.is_connected:
raise ValueError("Robot or teleop is not connected!") raise ValueError("Robot or teleop is not connected!")
print("Starting record loop. Move your phone to teleoperate the robot...")
print("Starting record loop. Move your phone to teleoperate the robot...") episode_idx = 0
episode_idx = 0 while episode_idx < NUM_EPISODES and not events["stop_recording"]:
while episode_idx < NUM_EPISODES and not events["stop_recording"]:
log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}") log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
# Main record loop # Main record loop
@@ -195,11 +200,15 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]:
dataset.save_episode() dataset.save_episode()
episode_idx += 1 episode_idx += 1
# Clean up # Clean up
log_say("Stop recording") log_say("Stop recording")
robot.disconnect() robot.disconnect()
phone.disconnect() phone.disconnect()
listener.stop() listener.stop()
dataset.finalize() dataset.finalize()
dataset.push_to_hub() dataset.push_to_hub()
if __name__ == "__main__":
main()
examples/phone_to_so100/replay.py
+30 -24
View File
@@ -35,23 +35,25 @@ from lerobot.utils.utils import log_say
EPISODE_IDX = 0 EPISODE_IDX = 0
HF_REPO_ID = "<hf_username>/<dataset_repo_id>" HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
# Initialize the robot config
robot_config = SO100FollowerConfig( def main():
# Initialize the robot config
robot_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
) )
# Initialize the robot # Initialize the robot
robot = SO100Follower(robot_config) 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 # 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( kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()), joint_names=list(robot.bus.motors.keys()),
) )
# Build pipeline to convert EE action to joints action # Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
steps=[ steps=[
InverseKinematicsEEToJoints( InverseKinematicsEEToJoints(
kinematics=kinematics_solver, kinematics=kinematics_solver,
@@ -61,23 +63,23 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Fetch the dataset to replay # Fetch the dataset to replay
dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX]) dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0 # Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX) episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
actions = episode_frames.select_columns(ACTION) actions = episode_frames.select_columns(ACTION)
# Connect to the robot # Connect to the robot
robot.connect() robot.connect()
if not robot.is_connected: if not robot.is_connected:
raise ValueError("Robot is not connected!") raise ValueError("Robot is not connected!")
print("Starting replay loop...") print("Starting replay loop...")
log_say(f"Replaying episode {EPISODE_IDX}") log_say(f"Replaying episode {EPISODE_IDX}")
for idx in range(len(episode_frames)): for idx in range(len(episode_frames)):
t0 = time.perf_counter() t0 = time.perf_counter()
# Get recorded action from dataset # Get recorded action from dataset
@@ -96,5 +98,9 @@ for idx in range(len(episode_frames)):
busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0)) busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))
# Clean up # Clean up
robot.disconnect() robot.disconnect()
if __name__ == "__main__":
main()
examples/phone_to_so100/teleoperate.py
+29 -21
View File
@@ -37,25 +37,29 @@ from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
FPS = 30 FPS = 30
# Initialize the robot and teleoperator
robot_config = SO100FollowerConfig( def main():
# Initialize the robot and teleoperator
robot_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
) )
teleop_config = PhoneConfig(phone_os=PhoneOS.IOS) # or PhoneOS.ANDROID teleop_config = PhoneConfig(phone_os=PhoneOS.IOS) # or PhoneOS.ANDROID
# Initialize the robot and teleoperator # Initialize the robot and teleoperator
robot = SO100Follower(robot_config) robot = SO100Follower(robot_config)
teleop_device = Phone(teleop_config) teleop_device = Phone(teleop_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 # 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( kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()), joint_names=list(robot.bus.motors.keys()),
) )
# Build pipeline to convert phone action to ee pose action to joint action # Build pipeline to convert phone action to ee pose action to joint action
phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( phone_to_robot_joints_processor = RobotProcessorPipeline[
tuple[RobotAction, RobotObservation], RobotAction
](
steps=[ steps=[
MapPhoneActionToRobotAction(platform=teleop_config.phone_os), MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
EEReferenceAndDelta( EEReferenceAndDelta(
@@ -79,20 +83,20 @@ phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, Robo
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Connect to the robot and teleoperator # Connect to the robot and teleoperator
robot.connect() robot.connect()
teleop_device.connect() teleop_device.connect()
# Init rerun viewer # Init rerun viewer
init_rerun(session_name="phone_so100_teleop") init_rerun(session_name="phone_so100_teleop")
if not robot.is_connected or not teleop_device.is_connected: if not robot.is_connected or not teleop_device.is_connected:
raise ValueError("Robot or teleop is not connected!") raise ValueError("Robot or teleop is not connected!")
print("Starting teleop loop. Move your phone to teleoperate the robot...") print("Starting teleop loop. Move your phone to teleoperate the robot...")
while True: while True:
t0 = time.perf_counter() t0 = time.perf_counter()
# Get robot observation # Get robot observation
@@ -111,3 +115,7 @@ while True:
log_rerun_data(observation=phone_obs, action=joint_action) log_rerun_data(observation=phone_obs, action=joint_action)
busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
if __name__ == "__main__":
main()
examples/so100_to_so100_EE/evaluate.py
+46 -39
View File
@@ -52,29 +52,31 @@ TASK_DESCRIPTION = "My task description"
HF_MODEL_ID = "<hf_username>/<model_repo_id>" HF_MODEL_ID = "<hf_username>/<model_repo_id>"
HF_DATASET_ID = "<hf_username>/<dataset_repo_id>" HF_DATASET_ID = "<hf_username>/<dataset_repo_id>"
# Create the robot configuration & robot
camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)} def main():
robot_config = SO100FollowerConfig( # 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", port="/dev/tty.usbmodem5A460814411",
id="my_awesome_follower_arm", id="my_awesome_follower_arm",
cameras=camera_config, cameras=camera_config,
use_degrees=True, use_degrees=True,
) )
robot = SO100Follower(robot_config) robot = SO100Follower(robot_config)
# Create policy # Create policy
policy = ACTPolicy.from_pretrained(HF_MODEL_ID) 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 # 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( kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()), joint_names=list(robot.bus.motors.keys()),
) )
# Build pipeline to convert EE action to joints action # Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
steps=[ steps=[
InverseKinematicsEEToJoints( InverseKinematicsEEToJoints(
kinematics=kinematics_solver, kinematics=kinematics_solver,
@@ -84,20 +86,21 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Build pipeline to convert joints observation to EE observation # Build pipeline to convert joints observation to EE observation
robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation]( robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
steps=[ steps=[
ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) ForwardKinematicsJointsToEE(
kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())
)
], ],
to_transition=observation_to_transition, to_transition=observation_to_transition,
to_output=transition_to_observation, to_output=transition_to_observation,
) )
# Create the dataset
# Create the dataset dataset = LeRobotDataset.create(
dataset = LeRobotDataset.create(
repo_id=HF_DATASET_ID, repo_id=HF_DATASET_ID,
fps=FPS, fps=FPS,
features=combine_feature_dicts( features=combine_feature_dicts(
@@ -122,30 +125,30 @@ dataset = LeRobotDataset.create(
robot_type=robot.name, robot_type=robot.name,
use_videos=True, use_videos=True,
image_writer_threads=4, image_writer_threads=4,
) )
# Build Policy Processors # Build Policy Processors
preprocessor, postprocessor = make_pre_post_processors( preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=policy, policy_cfg=policy,
pretrained_path=HF_MODEL_ID, pretrained_path=HF_MODEL_ID,
dataset_stats=dataset.meta.stats, dataset_stats=dataset.meta.stats,
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility. # The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}}, preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
) )
# Connect the robot and teleoperator # Connect the robot and teleoperator
robot.connect() robot.connect()
# Initialize the keyboard listener and rerun visualization # Initialize the keyboard listener and rerun visualization
listener, events = init_keyboard_listener() listener, events = init_keyboard_listener()
init_rerun(session_name="so100_so100_evaluate") init_rerun(session_name="so100_so100_evaluate")
if not robot.is_connected: if not robot.is_connected:
raise ValueError("Robot is not connected!") raise ValueError("Robot is not connected!")
print("Starting evaluate loop...") print("Starting evaluate loop...")
episode_idx = 0 episode_idx = 0
for episode_idx in range(NUM_EPISODES): for episode_idx in range(NUM_EPISODES):
log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}") log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
# Main record loop # Main record loop
@@ -191,10 +194,14 @@ for episode_idx in range(NUM_EPISODES):
dataset.save_episode() dataset.save_episode()
episode_idx += 1 episode_idx += 1
# Clean up # Clean up
log_say("Stop recording") log_say("Stop recording")
robot.disconnect() robot.disconnect()
listener.stop() listener.stop()
dataset.finalize() dataset.finalize()
dataset.push_to_hub() dataset.push_to_hub()
if __name__ == "__main__":
main()
examples/so100_to_so100_EE/record.py
+53 -45
View File
@@ -48,33 +48,38 @@ RESET_TIME_SEC = 30
TASK_DESCRIPTION = "My task description" TASK_DESCRIPTION = "My task description"
HF_REPO_ID = "<hf_username>/<dataset_repo_id>" HF_REPO_ID = "<hf_username>/<dataset_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 def main():
follower = SO100Follower(follower_config) # Create the robot and teleoperator configurations
leader = SO100Leader(leader_config) 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")
# 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 # Initialize the robot and teleoperator
follower_kinematics_solver = RobotKinematics( 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="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(follower.bus.motors.keys()), 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 # 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( leader_kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(leader.bus.motors.keys()), joint_names=list(leader.bus.motors.keys()),
) )
# Build pipeline to convert follower joints to EE observation # Build pipeline to convert follower joints to EE observation
follower_joints_to_ee = RobotProcessorPipeline[RobotObservation, RobotObservation]( follower_joints_to_ee = RobotProcessorPipeline[RobotObservation, RobotObservation](
steps=[ steps=[
ForwardKinematicsJointsToEE( ForwardKinematicsJointsToEE(
kinematics=follower_kinematics_solver, motor_names=list(follower.bus.motors.keys()) kinematics=follower_kinematics_solver, motor_names=list(follower.bus.motors.keys())
@@ -82,10 +87,10 @@ follower_joints_to_ee = RobotProcessorPipeline[RobotObservation, RobotObservatio
], ],
to_transition=observation_to_transition, to_transition=observation_to_transition,
to_output=transition_to_observation, to_output=transition_to_observation,
) )
# Build pipeline to convert leader joints to EE action # Build pipeline to convert leader joints to EE action
leader_joints_to_ee = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( leader_joints_to_ee = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
steps=[ steps=[
ForwardKinematicsJointsToEE( ForwardKinematicsJointsToEE(
kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys()) kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
@@ -93,10 +98,10 @@ leader_joints_to_ee = RobotProcessorPipeline[tuple[RobotAction, RobotObservation
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Build pipeline to convert EE action to follower joints # Build pipeline to convert EE action to follower joints
ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
[ [
EEBoundsAndSafety( EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]}, end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
@@ -110,10 +115,10 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Create the dataset # Create the dataset
dataset = LeRobotDataset.create( dataset = LeRobotDataset.create(
repo_id=HF_REPO_ID, repo_id=HF_REPO_ID,
fps=FPS, fps=FPS,
features=combine_feature_dicts( features=combine_feature_dicts(
@@ -133,23 +138,22 @@ dataset = LeRobotDataset.create(
robot_type=follower.name, robot_type=follower.name,
use_videos=True, use_videos=True,
image_writer_threads=4, image_writer_threads=4,
) )
# Connect the robot and teleoperator
leader.connect()
follower.connect()
# Connect the robot and teleoperator # Initialize the keyboard listener and rerun visualization
leader.connect() listener, events = init_keyboard_listener()
follower.connect() init_rerun(session_name="recording_phone")
# Initialize the keyboard listener and rerun visualization if not leader.is_connected or not follower.is_connected:
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!") raise ValueError("Robot or teleop is not connected!")
print("Starting record loop...") print("Starting record loop...")
episode_idx = 0 episode_idx = 0
while episode_idx < NUM_EPISODES and not events["stop_recording"]: while episode_idx < NUM_EPISODES and not events["stop_recording"]:
log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}") log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
# Main record loop # Main record loop
@@ -194,11 +198,15 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]:
dataset.save_episode() dataset.save_episode()
episode_idx += 1 episode_idx += 1
# Clean up # Clean up
log_say("Stop recording") log_say("Stop recording")
leader.disconnect() leader.disconnect()
follower.disconnect() follower.disconnect()
listener.stop() listener.stop()
dataset.finalize() dataset.finalize()
dataset.push_to_hub() dataset.push_to_hub()
if __name__ == "__main__":
main()
examples/so100_to_so100_EE/replay.py
+30 -24
View File
@@ -36,23 +36,25 @@ from lerobot.utils.utils import log_say
EPISODE_IDX = 0 EPISODE_IDX = 0
HF_REPO_ID = "<hf_username>/<dataset_repo_id>" HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
# Initialize the robot config
robot_config = SO100FollowerConfig( def main():
# Initialize the robot config
robot_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
) )
# Initialize the robot # Initialize the robot
robot = SO100Follower(robot_config) 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 # 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( kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()), joint_names=list(robot.bus.motors.keys()),
) )
# Build pipeline to convert EE action to joints action # Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
steps=[ steps=[
InverseKinematicsEEToJoints( InverseKinematicsEEToJoints(
kinematics=kinematics_solver, kinematics=kinematics_solver,
@@ -62,23 +64,23 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Fetch the dataset to replay # Fetch the dataset to replay
dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX]) dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0 # Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX) episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
actions = episode_frames.select_columns(ACTION) actions = episode_frames.select_columns(ACTION)
# Connect to the robot # Connect to the robot
robot.connect() robot.connect()
if not robot.is_connected: if not robot.is_connected:
raise ValueError("Robot is not connected!") raise ValueError("Robot is not connected!")
print("Starting replay loop...") print("Starting replay loop...")
log_say(f"Replaying episode {EPISODE_IDX}") log_say(f"Replaying episode {EPISODE_IDX}")
for idx in range(len(episode_frames)): for idx in range(len(episode_frames)):
t0 = time.perf_counter() t0 = time.perf_counter()
# Get recorded action from dataset # Get recorded action from dataset
@@ -97,5 +99,9 @@ for idx in range(len(episode_frames)):
busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0)) busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))
# Clean up # Clean up
robot.disconnect() robot.disconnect()
if __name__ == "__main__":
main()
examples/so100_to_so100_EE/teleoperate.py
+32 -26
View File
@@ -37,32 +37,34 @@ from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
FPS = 30 FPS = 30
# Initialize the robot and teleoperator config
follower_config = SO100FollowerConfig( def main():
# Initialize the robot and teleoperator config
follower_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
) )
leader_config = SO100LeaderConfig(port="/dev/tty.usbmodem5A460819811", id="my_awesome_leader_arm") leader_config = SO100LeaderConfig(port="/dev/tty.usbmodem5A460819811", id="my_awesome_leader_arm")
# Initialize the robot and teleoperator # Initialize the robot and teleoperator
follower = SO100Follower(follower_config) follower = SO100Follower(follower_config)
leader = SO100Leader(leader_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 # 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( follower_kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(follower.bus.motors.keys()), 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 # 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( leader_kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(leader.bus.motors.keys()), joint_names=list(leader.bus.motors.keys()),
) )
# Build pipeline to convert teleop joints to EE action # Build pipeline to convert teleop joints to EE action
leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction]( leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
steps=[ steps=[
ForwardKinematicsJointsToEE( ForwardKinematicsJointsToEE(
kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys()) kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
@@ -70,10 +72,10 @@ leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
], ],
to_transition=robot_action_to_transition, to_transition=robot_action_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# build pipeline to convert EE action to robot joints # build pipeline to convert EE action to robot joints
ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
[ [
EEBoundsAndSafety( EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]}, end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
@@ -87,17 +89,17 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Connect to the robot and teleoperator # Connect to the robot and teleoperator
follower.connect() follower.connect()
leader.connect() leader.connect()
# Init rerun viewer # Init rerun viewer
init_rerun(session_name="so100_so100_EE_teleop") init_rerun(session_name="so100_so100_EE_teleop")
print("Starting teleop loop...") print("Starting teleop loop...")
while True: while True:
t0 = time.perf_counter() t0 = time.perf_counter()
# Get robot observation # Get robot observation
@@ -119,3 +121,7 @@ while True:
log_rerun_data(observation=leader_ee_act, action=follower_joints_act) log_rerun_data(observation=leader_ee_act, action=follower_joints_act)
busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
if __name__ == "__main__":
main()
examples/tutorial/act/act_training_example.py
+50 -44
View File
@@ -19,60 +19,62 @@ def make_delta_timestamps(delta_indices: list[int] | None, fps: int) -> list[flo
return [i / fps for i in delta_indices] return [i / fps for i in delta_indices]
output_directory = Path("outputs/robot_learning_tutorial/act") def main():
output_directory.mkdir(parents=True, exist_ok=True) output_directory = Path("outputs/robot_learning_tutorial/act")
output_directory.mkdir(parents=True, exist_ok=True)
# Select your device # Select your device
device = torch.device("mps") # or "cuda" or "cpu" device = torch.device("mps") # or "cuda" or "cpu"
dataset_id = "lerobot/svla_so101_pickplace" dataset_id = "lerobot/svla_so101_pickplace"
# This specifies the inputs the model will be expecting and the outputs it will produce # This specifies the inputs the model will be expecting and the outputs it will produce
dataset_metadata = LeRobotDatasetMetadata(dataset_id) dataset_metadata = LeRobotDatasetMetadata(dataset_id)
features = dataset_to_policy_features(dataset_metadata.features) features = dataset_to_policy_features(dataset_metadata.features)
output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION} output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
input_features = {key: ft for key, ft in features.items() if key not in output_features} input_features = {key: ft for key, ft in features.items() if key not in output_features}
cfg = ACTConfig(input_features=input_features, output_features=output_features) cfg = ACTConfig(input_features=input_features, output_features=output_features)
policy = ACTPolicy(cfg) policy = ACTPolicy(cfg)
preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats) preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
policy.train() policy.train()
policy.to(device) policy.to(device)
# To perform action chunking, ACT expects a given number of actions as targets # To perform action chunking, ACT expects a given number of actions as targets
delta_timestamps = { delta_timestamps = {
"action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps), "action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps),
} }
# add image features if they are present # add image features if they are present
delta_timestamps |= { delta_timestamps |= {
k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps) for k in cfg.image_features k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps)
} for k in cfg.image_features
}
# Instantiate the dataset # Instantiate the dataset
dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps) dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps)
# Create the optimizer and dataloader for offline training # Create the optimizer and dataloader for offline training
optimizer = cfg.get_optimizer_preset().build(policy.parameters()) optimizer = cfg.get_optimizer_preset().build(policy.parameters())
batch_size = 32 batch_size = 32
dataloader = torch.utils.data.DataLoader( dataloader = torch.utils.data.DataLoader(
dataset, dataset,
batch_size=batch_size, batch_size=batch_size,
shuffle=True, shuffle=True,
pin_memory=device.type != "cpu", pin_memory=device.type != "cpu",
drop_last=True, drop_last=True,
) )
# Number of training steps and logging frequency # Number of training steps and logging frequency
training_steps = 1 training_steps = 1
log_freq = 1 log_freq = 1
# Run training loop # Run training loop
step = 0 step = 0
done = False done = False
while not done: while not done:
for batch in dataloader: for batch in dataloader:
batch = preprocessor(batch) batch = preprocessor(batch)
loss, _ = policy.forward(batch) loss, _ = policy.forward(batch)
@@ -87,12 +89,16 @@ while not done:
done = True done = True
break break
# Save the policy checkpoint, alongside the pre/post processors # Save the policy checkpoint, alongside the pre/post processors
policy.save_pretrained(output_directory) policy.save_pretrained(output_directory)
preprocessor.save_pretrained(output_directory) preprocessor.save_pretrained(output_directory)
postprocessor.save_pretrained(output_directory) postprocessor.save_pretrained(output_directory)
# Save all assets to the Hub # Save all assets to the Hub
policy.push_to_hub("fracapuano/robot_learning_tutorial_act") policy.push_to_hub("<user>/robot_learning_tutorial_act")
preprocessor.push_to_hub("fracapuano/robot_learning_tutorial_act") preprocessor.push_to_hub("<user>/robot_learning_tutorial_act")
postprocessor.push_to_hub("fracapuano/robot_learning_tutorial_act") postprocessor.push_to_hub("<user>/robot_learning_tutorial_act")
if __name__ == "__main__":
main()
examples/tutorial/act/act_using_example.py
+30 -24
View File
@@ -8,37 +8,39 @@ from lerobot.policies.utils import build_inference_frame, make_robot_action
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.robots.so100_follower.so100_follower import SO100Follower
device = torch.device("mps") # or "cuda" or "cpu"
model_id = "fracapuano/robot_learning_tutorial_act"
model = ACTPolicy.from_pretrained(model_id)
dataset_id = "lerobot/svla_so101_pickplace"
# This only downloads the metadata for the dataset, ~10s of MB even for large-scale datasets
dataset_metadata = LeRobotDatasetMetadata(dataset_id)
preprocess, postprocess = make_pre_post_processors(model.config, dataset_stats=dataset_metadata.stats)
# # find ports using lerobot-find-port
follower_port = ... # something like "/dev/tty.usbmodem58760431631"
# # the robot ids are used the load the right calibration files
follower_id = ... # something like "follower_so100"
MAX_EPISODES = 5 MAX_EPISODES = 5
MAX_STEPS_PER_EPISODE = 20 MAX_STEPS_PER_EPISODE = 20
# Robot and environment configuration
# Camera keys must match the name and resolutions of the ones used for training! def main():
# You can check the camera keys expected by a model in the info.json card on the model card on the Hub device = torch.device("mps") # or "cuda" or "cpu"
camera_config = { model_id = "<user>/robot_learning_tutorial_act"
model = ACTPolicy.from_pretrained(model_id)
dataset_id = "lerobot/svla_so101_pickplace"
# This only downloads the metadata for the dataset, ~10s of MB even for large-scale datasets
dataset_metadata = LeRobotDatasetMetadata(dataset_id)
preprocess, postprocess = make_pre_post_processors(model.config, dataset_stats=dataset_metadata.stats)
# # find ports using lerobot-find-port
follower_port = ... # something like "/dev/tty.usbmodem58760431631"
# # the robot ids are used the load the right calibration files
follower_id = ... # something like "follower_so100"
# Robot and environment configuration
# Camera keys must match the name and resolutions of the ones used for training!
# You can check the camera keys expected by a model in the info.json card on the model card on the Hub
camera_config = {
"side": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30), "side": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
"up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30), "up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
} }
robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config) robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
robot = SO100Follower(robot_cfg) robot = SO100Follower(robot_cfg)
robot.connect() robot.connect()
for _ in range(MAX_EPISODES): for _ in range(MAX_EPISODES):
for _ in range(MAX_STEPS_PER_EPISODE): for _ in range(MAX_STEPS_PER_EPISODE):
obs = robot.get_observation() obs = robot.get_observation()
obs_frame = build_inference_frame( obs_frame = build_inference_frame(
@@ -55,3 +57,7 @@ for _ in range(MAX_EPISODES):
robot.send_action(action) robot.send_action(action)
print("Episode finished! Starting new episode...") print("Episode finished! Starting new episode...")
if __name__ == "__main__":
main()
examples/tutorial/async-inf/policy_server.py
+11 -5
View File
@@ -1,11 +1,17 @@
from lerobot.async_inference.configs import PolicyServerConfig from lerobot.async_inference.configs import PolicyServerConfig
from lerobot.async_inference.policy_server import serve from lerobot.async_inference.policy_server import serve
host = ... # something like "127.0.0.1" if you're exposing to localhost
port = ... # something like 8080
config = PolicyServerConfig( def main():
host = ... # something like "127.0.0.1" if you're exposing to localhost
port = ... # something like 8080
config = PolicyServerConfig(
host=host, host=host,
port=port, port=port,
) )
serve(config) serve(config)
if __name__ == "__main__":
main()
examples/tutorial/async-inf/robot_client.py
+25 -19
View File
@@ -6,41 +6,43 @@ from lerobot.async_inference.robot_client import RobotClient
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
from lerobot.robots.so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower import SO100FollowerConfig
# these cameras must match the ones expected by the policy - find your cameras with lerobot-find-cameras
# check the config.json on the Hub for the policy you are using to see the expected camera specs def main():
camera_cfg = { # these cameras must match the ones expected by the policy - find your cameras with lerobot-find-cameras
# check the config.json on the Hub for the policy you are using to see the expected camera specs
camera_cfg = {
"up": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30), "up": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
"side": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30), "side": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
} }
# # find ports using lerobot-find-port # # find ports using lerobot-find-port
follower_port = ... # something like "/dev/tty.usbmodem58760431631" follower_port = ... # something like "/dev/tty.usbmodem58760431631"
# # the robot ids are used the load the right calibration files # # the robot ids are used the load the right calibration files
follower_id = ... # something like "follower_so100" follower_id = ... # something like "follower_so100"
robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_cfg) robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_cfg)
server_address = ... # something like "127.0.0.1:8080" if using localhost server_address = ... # something like "127.0.0.1:8080" if using localhost
# 3. Create client configuration # 3. Create client configuration
client_cfg = RobotClientConfig( client_cfg = RobotClientConfig(
robot=robot_cfg, robot=robot_cfg,
server_address=server_address, server_address=server_address,
policy_device="mps", policy_device="mps",
policy_type="act", policy_type="act",
pretrained_name_or_path="fracapuano/robot_learning_tutorial_act", pretrained_name_or_path="<user>/robot_learning_tutorial_act",
chunk_size_threshold=0.5, # g chunk_size_threshold=0.5, # g
actions_per_chunk=50, # make sure this is less than the max actions of the policy actions_per_chunk=50, # make sure this is less than the max actions of the policy
) )
# 4. Create and start client # 4. Create and start client
client = RobotClient(client_cfg) client = RobotClient(client_cfg)
# 5. Provide a textual description of the task # 5. Provide a textual description of the task
task = ... task = ...
if client.start(): if client.start():
# Start action receiver thread # Start action receiver thread
action_receiver_thread = threading.Thread(target=client.receive_actions, daemon=True) action_receiver_thread = threading.Thread(target=client.receive_actions, daemon=True)
action_receiver_thread.start() action_receiver_thread.start()
@@ -53,3 +55,7 @@ if client.start():
action_receiver_thread.join() action_receiver_thread.join()
# (Optionally) plot the action queue size # (Optionally) plot the action queue size
visualize_action_queue_size(client.action_queue_size) visualize_action_queue_size(client.action_queue_size)
if __name__ == "__main__":
main()
examples/tutorial/diffusion/diffusion_training_example.py
+50 -44
View File
@@ -19,61 +19,63 @@ def make_delta_timestamps(delta_indices: list[int] | None, fps: int) -> list[flo
return [i / fps for i in delta_indices] return [i / fps for i in delta_indices]
output_directory = Path("outputs/robot_learning_tutorial/diffusion") def main():
output_directory.mkdir(parents=True, exist_ok=True) output_directory = Path("outputs/robot_learning_tutorial/diffusion")
output_directory.mkdir(parents=True, exist_ok=True)
# Select your device # Select your device
device = torch.device("mps") # or "cuda" or "cpu" device = torch.device("mps") # or "cuda" or "cpu"
dataset_id = "lerobot/svla_so101_pickplace" dataset_id = "lerobot/svla_so101_pickplace"
# This specifies the inputs the model will be expecting and the outputs it will produce # This specifies the inputs the model will be expecting and the outputs it will produce
dataset_metadata = LeRobotDatasetMetadata(dataset_id) dataset_metadata = LeRobotDatasetMetadata(dataset_id)
features = dataset_to_policy_features(dataset_metadata.features) features = dataset_to_policy_features(dataset_metadata.features)
output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION} output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
input_features = {key: ft for key, ft in features.items() if key not in output_features} input_features = {key: ft for key, ft in features.items() if key not in output_features}
cfg = DiffusionConfig(input_features=input_features, output_features=output_features) cfg = DiffusionConfig(input_features=input_features, output_features=output_features)
policy = DiffusionPolicy(cfg) policy = DiffusionPolicy(cfg)
preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats) preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
policy.train() policy.train()
policy.to(device) policy.to(device)
# To perform action chunking, ACT expects a given number of actions as targets # To perform action chunking, ACT expects a given number of actions as targets
delta_timestamps = { delta_timestamps = {
"observation.state": make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps), "observation.state": make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps),
"action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps), "action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps),
} }
# add image features if they are present # add image features if they are present
delta_timestamps |= { delta_timestamps |= {
k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps) for k in cfg.image_features k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps)
} for k in cfg.image_features
}
# Instantiate the dataset # Instantiate the dataset
dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps) dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps)
# Create the optimizer and dataloader for offline training # Create the optimizer and dataloader for offline training
optimizer = cfg.get_optimizer_preset().build(policy.parameters()) optimizer = cfg.get_optimizer_preset().build(policy.parameters())
batch_size = 32 batch_size = 32
dataloader = torch.utils.data.DataLoader( dataloader = torch.utils.data.DataLoader(
dataset, dataset,
batch_size=batch_size, batch_size=batch_size,
shuffle=True, shuffle=True,
pin_memory=device.type != "cpu", pin_memory=device.type != "cpu",
drop_last=True, drop_last=True,
) )
# Number of training steps and logging frequency # Number of training steps and logging frequency
training_steps = 1 training_steps = 1
log_freq = 1 log_freq = 1
# Run training loop # Run training loop
step = 0 step = 0
done = False done = False
while not done: while not done:
for batch in dataloader: for batch in dataloader:
batch = preprocessor(batch) batch = preprocessor(batch)
loss, _ = policy.forward(batch) loss, _ = policy.forward(batch)
@@ -88,12 +90,16 @@ while not done:
done = True done = True
break break
# Save the policy checkpoint, alongside the pre/post processors # Save the policy checkpoint, alongside the pre/post processors
policy.save_pretrained(output_directory) policy.save_pretrained(output_directory)
preprocessor.save_pretrained(output_directory) preprocessor.save_pretrained(output_directory)
postprocessor.save_pretrained(output_directory) postprocessor.save_pretrained(output_directory)
# Save all assets to the Hub # Save all assets to the Hub
policy.push_to_hub("fracapuano/robot_learning_tutorial_diffusion") policy.push_to_hub("<user>/robot_learning_tutorial_diffusion")
preprocessor.push_to_hub("fracapuano/robot_learning_tutorial_diffusion") preprocessor.push_to_hub("<user>/robot_learning_tutorial_diffusion")
postprocessor.push_to_hub("fracapuano/robot_learning_tutorial_diffusion") postprocessor.push_to_hub("<user>/robot_learning_tutorial_diffusion")
if __name__ == "__main__":
main()
examples/tutorial/diffusion/diffusion_using_example.py
+30 -26
View File
@@ -8,42 +8,42 @@ from lerobot.policies.utils import build_inference_frame, make_robot_action
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.robots.so100_follower.so100_follower import SO100Follower
device = torch.device("mps") # or "cuda" or "cpu"
model_id = "fracapuano/robot_learning_tutorial_diffusion"
model = DiffusionPolicy.from_pretrained(model_id)
dataset_id = "lerobot/svla_so101_pickplace"
# This only downloads the metadata for the dataset, ~10s of MB even for large-scale datasets
dataset_metadata = LeRobotDatasetMetadata(dataset_id)
preprocess, postprocess = make_pre_post_processors(
model.config, model_id, dataset_stats=dataset_metadata.stats
)
MAX_EPISODES = 5 MAX_EPISODES = 5
MAX_STEPS_PER_EPISODE = 20 MAX_STEPS_PER_EPISODE = 20
# # find ports using lerobot-find-port def main():
follower_port = ... # something like "/dev/tty.usbmodem58760431631" device = torch.device("mps") # or "cuda" or "cpu"
model_id = "<user>/robot_learning_tutorial_diffusion"
# # the robot ids are used the load the right calibration files model = DiffusionPolicy.from_pretrained(model_id)
follower_id = ... # something like "follower_so100"
# Robot and environment configuration dataset_id = "lerobot/svla_so101_pickplace"
# Camera keys must match the name and resolutions of the ones used for training! # This only downloads the metadata for the dataset, ~10s of MB even for large-scale datasets
# You can check the camera keys expected by a model in the info.json card on the model card on the Hub dataset_metadata = LeRobotDatasetMetadata(dataset_id)
camera_config = { preprocess, postprocess = make_pre_post_processors(
model.config, model_id, dataset_stats=dataset_metadata.stats
)
# # find ports using lerobot-find-port
follower_port = ... # something like "/dev/tty.usbmodem58760431631"
# # the robot ids are used the load the right calibration files
follower_id = ... # something like "follower_so100"
# Robot and environment configuration
# Camera keys must match the name and resolutions of the ones used for training!
# You can check the camera keys expected by a model in the info.json card on the model card on the Hub
camera_config = {
"side": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30), "side": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
"up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30), "up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
} }
robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config) robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
robot = SO100Follower(robot_cfg) robot = SO100Follower(robot_cfg)
robot.connect() robot.connect()
for _ in range(MAX_EPISODES):
for _ in range(MAX_EPISODES):
for _ in range(MAX_STEPS_PER_EPISODE): for _ in range(MAX_STEPS_PER_EPISODE):
obs = robot.get_observation() obs = robot.get_observation()
obs_frame = build_inference_frame( obs_frame = build_inference_frame(
@@ -58,3 +58,7 @@ for _ in range(MAX_EPISODES):
robot.send_action(action) robot.send_action(action)
print("Episode finished! Starting new episode...") print("Episode finished! Starting new episode...")
if __name__ == "__main__":
main()
examples/tutorial/pi0/using_pi0_example.py
+30 -24
View File
@@ -11,46 +11,48 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower
MAX_EPISODES = 5 MAX_EPISODES = 5
MAX_STEPS_PER_EPISODE = 20 MAX_STEPS_PER_EPISODE = 20
device = torch.device("mps") # or "cuda" or "cpu"
model_id = "lerobot/pi0_base"
model = PI0Policy.from_pretrained(model_id) def main():
device = torch.device("mps") # or "cuda" or "cpu"
model_id = "lerobot/pi0_base"
preprocess, postprocess = make_pre_post_processors( model = PI0Policy.from_pretrained(model_id)
preprocess, postprocess = make_pre_post_processors(
model.config, model.config,
model_id, model_id,
# This overrides allows to run on MPS, otherwise defaults to CUDA (if available) # This overrides allows to run on MPS, otherwise defaults to CUDA (if available)
preprocessor_overrides={"device_processor": {"device": str(device)}}, preprocessor_overrides={"device_processor": {"device": str(device)}},
) )
# find ports using lerobot-find-port # find ports using lerobot-find-port
follower_port = ... # something like "/dev/tty.usbmodem58760431631" follower_port = ... # something like "/dev/tty.usbmodem58760431631"
# the robot ids are used the load the right calibration files # the robot ids are used the load the right calibration files
follower_id = ... # something like "follower_so100" follower_id = ... # something like "follower_so100"
# Robot and environment configuration # Robot and environment configuration
# Camera keys must match the name and resolutions of the ones used for training! # Camera keys must match the name and resolutions of the ones used for training!
# You can check the camera keys expected by a model in the info.json card on the model card on the Hub # You can check the camera keys expected by a model in the info.json card on the model card on the Hub
camera_config = { camera_config = {
"base_0_rgb": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30), "base_0_rgb": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
"left_wrist_0_rgb": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30), "left_wrist_0_rgb": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
"right_wrist_0_rgb": OpenCVCameraConfig(index_or_path=2, width=640, height=480, fps=30), "right_wrist_0_rgb": OpenCVCameraConfig(index_or_path=2, width=640, height=480, fps=30),
} }
robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config) robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
robot = SO100Follower(robot_cfg) robot = SO100Follower(robot_cfg)
robot.connect() robot.connect()
task = "" # something like "pick the red block" task = "" # something like "pick the red block"
robot_type = "" # something like "so100_follower" for multi-embodiment datasets robot_type = "" # something like "so100_follower" for multi-embodiment datasets
# This is used to match the raw observation keys to the keys expected by the policy # This is used to match the raw observation keys to the keys expected by the policy
action_features = hw_to_dataset_features(robot.action_features, "action") action_features = hw_to_dataset_features(robot.action_features, "action")
obs_features = hw_to_dataset_features(robot.observation_features, "observation") obs_features = hw_to_dataset_features(robot.observation_features, "observation")
dataset_features = {**action_features, **obs_features} dataset_features = {**action_features, **obs_features}
for _ in range(MAX_EPISODES): for _ in range(MAX_EPISODES):
for _ in range(MAX_STEPS_PER_EPISODE): for _ in range(MAX_STEPS_PER_EPISODE):
obs = robot.get_observation() obs = robot.get_observation()
obs_frame = build_inference_frame( obs_frame = build_inference_frame(
@@ -65,3 +67,7 @@ for _ in range(MAX_EPISODES):
robot.send_action(action) robot.send_action(action)
print("Episode finished! Starting new episode...") print("Episode finished! Starting new episode...")
if __name__ == "__main__":
main()
examples/tutorial/rl/hilserl_example.py
+61 -59
View File
@@ -20,6 +20,8 @@ from lerobot.teleoperators.utils import TeleopEvents
LOG_EVERY = 10 LOG_EVERY = 10
SEND_EVERY = 10 SEND_EVERY = 10
MAX_EPISODES = 5
MAX_STEPS_PER_EPISODE = 20
def run_learner( def run_learner(
@@ -223,80 +225,76 @@ def make_policy_obs(obs, device: torch.device = "cpu"):
} }
"""Main function - coordinates actor and learner processes.""" def main():
"""Main function - coordinates actor and learner processes."""
device = "mps" # or "cuda" or "cpu" device = "mps" # or "cuda" or "cpu"
output_directory = Path("outputs/robot_learning_tutorial/hil_serl") output_directory = Path("outputs/robot_learning_tutorial/hil_serl")
output_directory.mkdir(parents=True, exist_ok=True) output_directory.mkdir(parents=True, exist_ok=True)
# find ports using lerobot-find-port # find ports using lerobot-find-port
follower_port = ... follower_port = ...
leader_port = ... leader_port = ...
# the robot ids are used the load the right calibration files # the robot ids are used the load the right calibration files
follower_id = ... follower_id = ...
leader_id = ... leader_id = ...
# A pretrained model (to be used in-distribution!) # A pretrained model (to be used in-distribution!)
reward_classifier_id = "fracapuano/reward_classifier_hil_serl_example" reward_classifier_id = "<user>/reward_classifier_hil_serl_example"
reward_classifier = Classifier.from_pretrained(reward_classifier_id) reward_classifier = Classifier.from_pretrained(reward_classifier_id)
reward_classifier.to(device) reward_classifier.to(device)
reward_classifier.eval() reward_classifier.eval()
MAX_EPISODES = 5 # Robot and environment configuration
MAX_STEPS_PER_EPISODE = 20 robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id)
teleop_cfg = SO100LeaderConfig(port=leader_port, id=leader_id)
processor_cfg = HILSerlProcessorConfig(control_mode="leader")
# Robot and environment configuration env_cfg = HILSerlRobotEnvConfig(robot=robot_cfg, teleop=teleop_cfg, processor=processor_cfg)
robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id)
teleop_cfg = SO100LeaderConfig(port=leader_port, id=leader_id)
processor_cfg = HILSerlProcessorConfig(control_mode="leader")
env_cfg = HILSerlRobotEnvConfig(robot=robot_cfg, teleop=teleop_cfg, processor=processor_cfg) # Create robot environment
env, teleop_device = make_robot_env(env_cfg)
# Create robot environment obs_features = hw_to_dataset_features(env.robot.observation_features, "observation")
env, teleop_device = make_robot_env(env_cfg) action_features = hw_to_dataset_features(env.robot.action_features, "action")
obs_features = hw_to_dataset_features(env.robot.observation_features, "observation") # Create SAC policy for action selection
action_features = hw_to_dataset_features(env.robot.action_features, "action") policy_cfg = SACConfig(
# Create SAC policy for action selection
policy_cfg = SACConfig(
device=device, device=device,
input_features=obs_features, input_features=obs_features,
output_features=action_features, output_features=action_features,
) )
policy_actor = SACPolicy(policy_cfg) policy_actor = SACPolicy(policy_cfg)
policy_learner = SACPolicy(policy_cfg) policy_learner = SACPolicy(policy_cfg)
demonstrations_repo_id = "lerobot/example_hil_serl_dataset" demonstrations_repo_id = "lerobot/example_hil_serl_dataset"
offline_dataset = LeRobotDataset(repo_id=demonstrations_repo_id) offline_dataset = LeRobotDataset(repo_id=demonstrations_repo_id)
# Online buffer: initialized from scratch # Online buffer: initialized from scratch
online_replay_buffer = ReplayBuffer(device=device, state_keys=list(obs_features.keys())) online_replay_buffer = ReplayBuffer(device=device, state_keys=list(obs_features.keys()))
# Offline buffer: Created from dataset (pre-populated it with demonstrations) # Offline buffer: Created from dataset (pre-populated it with demonstrations)
offline_replay_buffer = ReplayBuffer.from_lerobot_dataset( offline_replay_buffer = ReplayBuffer.from_lerobot_dataset(
lerobot_dataset=offline_dataset, device=device, state_keys=list(obs_features.keys()) lerobot_dataset=offline_dataset, device=device, state_keys=list(obs_features.keys())
) )
# Create communication channels between learner and actor processes # Create communication channels between learner and actor processes
transitions_queue = mp.Queue(maxsize=10) transitions_queue = mp.Queue(maxsize=10)
parameters_queue = mp.Queue(maxsize=2) parameters_queue = mp.Queue(maxsize=2)
shutdown_event = mp.Event() shutdown_event = mp.Event()
# Signal handler for graceful shutdown
# Signal handler for graceful shutdown def signal_handler(sig):
def signal_handler(sig):
print(f"\nSignal {sig} received, shutting down...") print(f"\nSignal {sig} received, shutting down...")
shutdown_event.set() shutdown_event.set()
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
signal.signal(signal.SIGINT, signal_handler) # Create processes
signal.signal(signal.SIGTERM, signal_handler) learner_process = mp.Process(
# Create processes
learner_process = mp.Process(
target=run_learner, target=run_learner,
args=( args=(
transitions_queue, transitions_queue,
@@ -307,9 +305,9 @@ learner_process = mp.Process(
offline_replay_buffer, offline_replay_buffer,
), ),
kwargs={"device": device}, # can run on accelerated hardware for training kwargs={"device": device}, # can run on accelerated hardware for training
) )
actor_process = mp.Process( actor_process = mp.Process(
target=run_actor, target=run_actor,
args=( args=(
transitions_queue, transitions_queue,
@@ -321,25 +319,29 @@ actor_process = mp.Process(
output_directory, output_directory,
), ),
kwargs={"device": "cpu"}, # actor is frozen, can run on CPU or accelerate for inference kwargs={"device": "cpu"}, # actor is frozen, can run on CPU or accelerate for inference
) )
learner_process.start() learner_process.start()
actor_process.start() actor_process.start()
try: try:
# Wait for actor to finish (it controls the episode loop) # Wait for actor to finish (it controls the episode loop)
actor_process.join() actor_process.join()
shutdown_event.set() shutdown_event.set()
learner_process.join(timeout=10) learner_process.join(timeout=10)
except KeyboardInterrupt: except KeyboardInterrupt:
print("Main process interrupted") print("Main process interrupted")
shutdown_event.set() shutdown_event.set()
actor_process.join(timeout=5) actor_process.join(timeout=5)
learner_process.join(timeout=10) learner_process.join(timeout=10)
finally: finally:
if learner_process.is_alive(): if learner_process.is_alive():
learner_process.terminate() learner_process.terminate()
if actor_process.is_alive(): if actor_process.is_alive():
actor_process.terminate() actor_process.terminate()
if __name__ == "__main__":
main()
examples/tutorial/rl/reward_classifier_example.py
+28 -23
View File
@@ -4,37 +4,38 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.policies.factory import make_policy, make_pre_post_processors from lerobot.policies.factory import make_policy, make_pre_post_processors
from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
# Device to use for training
device = "mps" # or "cuda", or "cpu"
# Load the dataset used for training def main():
repo_id = "lerobot/example_hil_serl_dataset" # Device to use for training
dataset = LeRobotDataset(repo_id) device = "mps" # or "cuda", or "cpu"
# Configure the policy to extract features from the image frames # Load the dataset used for training
camera_keys = dataset.meta.camera_keys repo_id = "lerobot/example_hil_serl_dataset"
dataset = LeRobotDataset(repo_id)
config = RewardClassifierConfig( # Configure the policy to extract features from the image frames
camera_keys = dataset.meta.camera_keys
config = RewardClassifierConfig(
num_cameras=len(camera_keys), num_cameras=len(camera_keys),
device=device, device=device,
# backbone model to extract features from the image frames # backbone model to extract features from the image frames
model_name="microsoft/resnet-18", model_name="microsoft/resnet-18",
) )
# Make policy, preprocessor, and optimizer # Make policy, preprocessor, and optimizer
policy = make_policy(config, ds_meta=dataset.meta) policy = make_policy(config, ds_meta=dataset.meta)
optimizer = config.get_optimizer_preset().build(policy.parameters()) optimizer = config.get_optimizer_preset().build(policy.parameters())
preprocessor, _ = make_pre_post_processors(policy_cfg=config, dataset_stats=dataset.meta.stats) preprocessor, _ = make_pre_post_processors(policy_cfg=config, dataset_stats=dataset.meta.stats)
classifier_id = "<user>/reward_classifier_hil_serl_example"
classifier_id = "fracapuano/reward_classifier_hil_serl_example" # Instantiate a dataloader
dataloader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True)
# Instantiate a dataloader # Training loop
dataloader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True) num_epochs = 5
for epoch in range(num_epochs):
# Training loop
num_epochs = 5
for epoch in range(num_epochs):
total_loss = 0 total_loss = 0
total_accuracy = 0 total_accuracy = 0
for batch in dataloader: for batch in dataloader:
@@ -56,7 +57,11 @@ for epoch in range(num_epochs):
avg_accuracy = total_accuracy / len(dataloader) avg_accuracy = total_accuracy / len(dataloader)
print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}, Accuracy: {avg_accuracy:.2f}%") print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}, Accuracy: {avg_accuracy:.2f}%")
print("Training finished!") print("Training finished!")
# You can now save the trained policy. # You can now save the trained policy.
policy.push_to_hub(classifier_id) policy.push_to_hub(classifier_id)
if __name__ == "__main__":
main()
examples/tutorial/smolvla/using_smolvla_example.py
+30 -24
View File
@@ -11,45 +11,47 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower
MAX_EPISODES = 5 MAX_EPISODES = 5
MAX_STEPS_PER_EPISODE = 20 MAX_STEPS_PER_EPISODE = 20
device = torch.device("mps") # or "cuda" or "cpu"
model_id = "lerobot/smolvla_base"
model = SmolVLAPolicy.from_pretrained(model_id) def main():
device = torch.device("mps") # or "cuda" or "cpu"
model_id = "lerobot/smolvla_base"
preprocess, postprocess = make_pre_post_processors( model = SmolVLAPolicy.from_pretrained(model_id)
preprocess, postprocess = make_pre_post_processors(
model.config, model.config,
model_id, model_id,
# This overrides allows to run on MPS, otherwise defaults to CUDA (if available) # This overrides allows to run on MPS, otherwise defaults to CUDA (if available)
preprocessor_overrides={"device_processor": {"device": str(device)}}, preprocessor_overrides={"device_processor": {"device": str(device)}},
) )
# find ports using lerobot-find-port # find ports using lerobot-find-port
follower_port = ... # something like "/dev/tty.usbmodem58760431631" follower_port = ... # something like "/dev/tty.usbmodem58760431631"
# the robot ids are used the load the right calibration files # the robot ids are used the load the right calibration files
follower_id = ... # something like "follower_so100" follower_id = ... # something like "follower_so100"
# Robot and environment configuration # Robot and environment configuration
# Camera keys must match the name and resolutions of the ones used for training! # Camera keys must match the name and resolutions of the ones used for training!
# You can check the camera keys expected by a model in the info.json card on the model card on the Hub # You can check the camera keys expected by a model in the info.json card on the model card on the Hub
camera_config = { camera_config = {
"camera1": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30), "camera1": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
"camera2": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30), "camera2": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
} }
robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config) robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
robot = SO100Follower(robot_cfg) robot = SO100Follower(robot_cfg)
robot.connect() robot.connect()
task = "" # something like "pick the red block" task = "" # something like "pick the red block"
robot_type = "" # something like "so100_follower" for multi-embodiment datasets robot_type = "" # something like "so100_follower" for multi-embodiment datasets
# This is used to match the raw observation keys to the keys expected by the policy # This is used to match the raw observation keys to the keys expected by the policy
action_features = hw_to_dataset_features(robot.action_features, "action") action_features = hw_to_dataset_features(robot.action_features, "action")
obs_features = hw_to_dataset_features(robot.observation_features, "observation") obs_features = hw_to_dataset_features(robot.observation_features, "observation")
dataset_features = {**action_features, **obs_features} dataset_features = {**action_features, **obs_features}
for _ in range(MAX_EPISODES): for _ in range(MAX_EPISODES):
for _ in range(MAX_STEPS_PER_EPISODE): for _ in range(MAX_STEPS_PER_EPISODE):
obs = robot.get_observation() obs = robot.get_observation()
obs_frame = build_inference_frame( obs_frame = build_inference_frame(
@@ -64,3 +66,7 @@ for _ in range(MAX_EPISODES):
robot.send_action(action) robot.send_action(action)
print("Episode finished! Starting new episode...") print("Episode finished! Starting new episode...")
if __name__ == "__main__":
main()