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

docs/source/async.mdx

@@ -196,7 +196,7 @@ client_cfg = RobotClientConfig(
     server_address="localhost:8080",
     policy_device="mps",
     policy_type="smolvla",
-    pretrained_name_or_path="fracapuano/smolvla_async",
+    pretrained_name_or_path="<user>/smolvla_async",
     chunk_size_threshold=0.5,
     actions_per_chunk=50,  # make sure this is less than the max actions of the policy
 )

examples/dataset/load_lerobot_dataset.py

@@ -34,105 +34,106 @@ from huggingface_hub import HfApi
 import lerobot
 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:
-    # loads 4 images: 1 second before current frame, 500 ms before, 200 ms before, and current frame
+    print(dataset.features[camera_key]["shape"])
-    camera_key: [-1, -0.5, -0.20, 0],
+    # The shape is in (h, w, c) which is a more universal format.
-    # 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],
-    # 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)],
-}
-# 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.
 
-dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps)
+    # For many machine learning applications we need to load the history of past observations or trajectories of
-print(f"\n{dataset[0][camera_key].shape=}")  # (4, c, h, w)
+    # future actions. Our datasets can load previous and future frames for each key/modality, using timestamps
-print(f"{dataset[0]['observation.state'].shape=}")  # (6, c)
+    # differences with the current loaded frame. For instance:
-print(f"{dataset[0]['action'].shape=}\n")  # (64, c)
+    delta_timestamps = {
+        # 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],
+        # 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],
+        # 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)],
+    }
+    # 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.
+
+    dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps)
+    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]['action'].shape=}\n")  # (64, c)
 
-if __name__ == "__main__":
     dataloader = torch.utils.data.DataLoader(
         dataset,
         num_workers=4,
@@ -144,3 +145,7 @@ if __name__ == "__main__":
         print(f"{batch['observation.state'].shape=}")  # (32, 6, c)
         print(f"{batch['action'].shape=}")  # (32, 64, c)
         break
+
+
+if __name__ == "__main__":
+    main()

examples/lekiwi/evaluate.py

@@ -33,83 +33,68 @@ TASK_DESCRIPTION = "My task description"
 HF_MODEL_ID = "<hf_username>/<model_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)
-    repo_id=HF_DATASET_ID,
+    obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
-    fps=FPS,
+    dataset_features = {**action_features, **obs_features}
-    features=dataset_features,
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
 
-# Build Policy Processors
+    # Create the dataset
-preprocessor, postprocessor = make_pre_post_processors(
+    dataset = LeRobotDataset.create(
-    policy_cfg=policy,
+        repo_id=HF_DATASET_ID,
-    pretrained_path=HF_MODEL_ID,
-    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.
-    preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
-)
-
-# Connect the robot
-# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
-robot.connect()
-
-# TODO(Steven): Update this example to use pipelines
-teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
-
-# Initialize the keyboard listener and rerun visualization
-listener, events = init_keyboard_listener()
-init_rerun(session_name="lekiwi_evaluate")
-
-if not robot.is_connected:
-    raise ValueError("Robot is not connected!")
-
-print("Starting evaluate loop...")
-recorded_episodes = 0
-while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
-    log_say(f"Running inference, recording eval episode {recorded_episodes} of {NUM_EPISODES}")
-
-    # Main record loop
-    record_loop(
-        robot=robot,
-        events=events,
         fps=FPS,
-        policy=policy,
+        features=dataset_features,
-        preprocessor=preprocessor,  # Pass the pre and post policy processors
+        robot_type=robot.name,
-        postprocessor=postprocessor,
+        use_videos=True,
-        dataset=dataset,
+        image_writer_threads=4,
-        control_time_s=EPISODE_TIME_SEC,
-        single_task=TASK_DESCRIPTION,
-        display_data=True,
-        teleop_action_processor=teleop_action_processor,
-        robot_action_processor=robot_action_processor,
-        robot_observation_processor=robot_observation_processor,
     )
 
-    # Reset the environment if not stopping or re-recording
+    # Build Policy Processors
-    if not events["stop_recording"] and (
+    preprocessor, postprocessor = make_pre_post_processors(
-        (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
+        policy_cfg=policy,
-    ):
+        pretrained_path=HF_MODEL_ID,
-        log_say("Reset the environment")
+        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.
+        preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
+    )
+
+    # Connect the robot
+    # To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
+    robot.connect()
+
+    # TODO(Steven): Update this example to use pipelines
+    teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
+
+    # Initialize the keyboard listener and rerun visualization
+    listener, events = init_keyboard_listener()
+    init_rerun(session_name="lekiwi_evaluate")
+
+    if not robot.is_connected:
+        raise ValueError("Robot is not connected!")
+
+    print("Starting evaluate loop...")
+    recorded_episodes = 0
+    while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
+        log_say(f"Running inference, recording eval episode {recorded_episodes} of {NUM_EPISODES}")
+
+        # Main record loop
         record_loop(
             robot=robot,
             events=events,
             fps=FPS,
+            policy=policy,
+            preprocessor=preprocessor,  # Pass the pre and post policy processors
+            postprocessor=postprocessor,
+            dataset=dataset,
             control_time_s=EPISODE_TIME_SEC,
             single_task=TASK_DESCRIPTION,
             display_data=True,
@@ -118,21 +103,42 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
             robot_observation_processor=robot_observation_processor,
         )
 
-    if events["rerecord_episode"]:
+        # Reset the environment if not stopping or re-recording
-        log_say("Re-record episode")
+        if not events["stop_recording"] and (
-        events["rerecord_episode"] = False
+            (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
-        events["exit_early"] = False
+        ):
-        dataset.clear_episode_buffer()
+            log_say("Reset the environment")
-        continue
+            record_loop(
+                robot=robot,
+                events=events,
+                fps=FPS,
+                control_time_s=EPISODE_TIME_SEC,
+                single_task=TASK_DESCRIPTION,
+                display_data=True,
+                teleop_action_processor=teleop_action_processor,
+                robot_action_processor=robot_action_processor,
+                robot_observation_processor=robot_observation_processor,
+            )
 
-    # Save episode
+        if events["rerecord_episode"]:
-    dataset.save_episode()
+            log_say("Re-record episode")
-    recorded_episodes += 1
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-# Clean up
+        # Save episode
-log_say("Stop recording")
+        dataset.save_episode()
-robot.disconnect()
+        recorded_episodes += 1
-listener.stop()
 
-dataset.finalize()
+    # Clean up
-dataset.push_to_hub()
+    log_say("Stop recording")
+    robot.disconnect()
+    listener.stop()
+
+    dataset.finalize()
+    dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()

examples/lekiwi/record.py

@@ -34,78 +34,62 @@ RESET_TIME_SEC = 10
 TASK_DESCRIPTION = "My task description"
 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)
-    repo_id=HF_REPO_ID,
+    obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
-    fps=FPS,
+    dataset_features = {**action_features, **obs_features}
-    features=dataset_features,
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
 
-# Connect the robot and teleoperator
+    # Create the dataset
-# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
+    dataset = LeRobotDataset.create(
-robot.connect()
+        repo_id=HF_REPO_ID,
-leader_arm.connect()
-keyboard.connect()
-
-# Initialize the keyboard listener and rerun visualization
-listener, events = init_keyboard_listener()
-init_rerun(session_name="lekiwi_record")
-
-if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
-    raise ValueError("Robot or teleop is not connected!")
-
-print("Starting record loop...")
-recorded_episodes = 0
-while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
-    log_say(f"Recording episode {recorded_episodes}")
-
-    # Main record loop
-    record_loop(
-        robot=robot,
-        events=events,
         fps=FPS,
-        dataset=dataset,
+        features=dataset_features,
-        teleop=[leader_arm, keyboard],
+        robot_type=robot.name,
-        control_time_s=EPISODE_TIME_SEC,
+        use_videos=True,
-        single_task=TASK_DESCRIPTION,
+        image_writer_threads=4,
-        display_data=True,
-        teleop_action_processor=teleop_action_processor,
-        robot_action_processor=robot_action_processor,
-        robot_observation_processor=robot_observation_processor,
     )
 
-    # Reset the environment if not stopping or re-recording
+    # Connect the robot and teleoperator
-    if not events["stop_recording"] and (
+    # To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
-        (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
+    robot.connect()
-    ):
+    leader_arm.connect()
-        log_say("Reset the environment")
+    keyboard.connect()
+
+    # Initialize the keyboard listener and rerun visualization
+    listener, events = init_keyboard_listener()
+    init_rerun(session_name="lekiwi_record")
+
+    if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
+        raise ValueError("Robot or teleop is not connected!")
+
+    print("Starting record loop...")
+    recorded_episodes = 0
+    while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
+        log_say(f"Recording episode {recorded_episodes}")
+
+        # Main record loop
         record_loop(
             robot=robot,
             events=events,
             fps=FPS,
+            dataset=dataset,
             teleop=[leader_arm, keyboard],
-            control_time_s=RESET_TIME_SEC,
+            control_time_s=EPISODE_TIME_SEC,
             single_task=TASK_DESCRIPTION,
             display_data=True,
             teleop_action_processor=teleop_action_processor,
@@ -113,23 +97,45 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
             robot_observation_processor=robot_observation_processor,
         )
 
-    if events["rerecord_episode"]:
+        # Reset the environment if not stopping or re-recording
-        log_say("Re-record episode")
+        if not events["stop_recording"] and (
-        events["rerecord_episode"] = False
+            (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
-        events["exit_early"] = False
+        ):
-        dataset.clear_episode_buffer()
+            log_say("Reset the environment")
-        continue
+            record_loop(
+                robot=robot,
+                events=events,
+                fps=FPS,
+                teleop=[leader_arm, keyboard],
+                control_time_s=RESET_TIME_SEC,
+                single_task=TASK_DESCRIPTION,
+                display_data=True,
+                teleop_action_processor=teleop_action_processor,
+                robot_action_processor=robot_action_processor,
+                robot_observation_processor=robot_observation_processor,
+            )
 
-    # Save episode
+        if events["rerecord_episode"]:
-    dataset.save_episode()
+            log_say("Re-record episode")
-    recorded_episodes += 1
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-# Clean up
+        # Save episode
-log_say("Stop recording")
+        dataset.save_episode()
-robot.disconnect()
+        recorded_episodes += 1
-leader_arm.disconnect()
-keyboard.disconnect()
-listener.stop()
 
-dataset.finalize()
+    # Clean up
-dataset.push_to_hub()
+    log_say("Stop recording")
+    robot.disconnect()
+    leader_arm.disconnect()
+    keyboard.disconnect()
+    listener.stop()
+
+    dataset.finalize()
+    dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()

examples/lekiwi/replay.py

@@ -25,37 +25,43 @@ from lerobot.utils.utils import log_say
 
 EPISODE_IDX = 0
 
-# Initialize the robot config
-robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
 
-# Initialize the robot
+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
-    raise ValueError("Robot is not connected!")
+    robot.connect()
 
-print("Starting replay loop...")
+    if not robot.is_connected:
-log_say(f"Replaying episode {EPISODE_IDX}")
+        raise ValueError("Robot is not connected!")
-for idx in range(len(episode_frames)):
-    t0 = time.perf_counter()
 
-    # Get recorded action from dataset
+    print("Starting replay loop...")
-    action = {
+    log_say(f"Replaying episode {EPISODE_IDX}")
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+    for idx in range(len(episode_frames)):
-    }
+        t0 = time.perf_counter()
 
-    # Send action to robot
+        # Get recorded action from dataset
-    _ = robot.send_action(action)
+        action = {
+            name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        }
 
-    busy_wait(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0))
+        # Send action to robot
+        _ = robot.send_action(action)
 
-robot.disconnect()
+        busy_wait(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0))
+
+    robot.disconnect()
+
+
+if __name__ == "__main__":
+    main()

examples/lekiwi/teleoperate.py

@@ -24,49 +24,55 @@ from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 
 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
-    raise ValueError("Robot or teleop is not connected!")
+    init_rerun(session_name="lekiwi_teleop")
 
-print("Starting teleop loop...")
+    if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
-while True:
+        raise ValueError("Robot or teleop is not connected!")
-    t0 = time.perf_counter()
 
-    # Get robot observation
+    print("Starting teleop loop...")
-    observation = robot.get_observation()
+    while True:
+        t0 = time.perf_counter()
 
-    # Get teleop action
+        # Get robot observation
-    # Arm
+        observation = robot.get_observation()
-    arm_action = leader_arm.get_action()
-    arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
-    # Keyboard
-    keyboard_keys = keyboard.get_action()
-    base_action = robot._from_keyboard_to_base_action(keyboard_keys)
 
-    action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
+        # Get teleop action
+        # Arm
+        arm_action = leader_arm.get_action()
+        arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
+        # Keyboard
+        keyboard_keys = keyboard.get_action()
+        base_action = robot._from_keyboard_to_base_action(keyboard_keys)
 
-    # Send action to robot
+        action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
-    _ = robot.send_action(action)
 
-    # Visualize
+        # Send action to robot
-    log_rerun_data(observation=observation, action=action)
+        _ = robot.send_action(action)
 
-    busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
+        # Visualize
+        log_rerun_data(observation=observation, action=action)
+
+        busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
+
+
+if __name__ == "__main__":
+    main()

examples/phone_to_so100/evaluate.py

@@ -52,125 +52,114 @@ TASK_DESCRIPTION = "My task description"
 HF_MODEL_ID = "<hf_username>/<model_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)}
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem58760434471",
-    id="my_awesome_follower_arm",
-    cameras=camera_config,
-    use_degrees=True,
-)
 
-robot = SO100Follower(robot_config)
+def main():
-
+    # Create the robot configuration & robot
-# Create policy
+    camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
-policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
+    robot_config = SO100FollowerConfig(
-
+        port="/dev/tty.usbmodem58760434471",
-# 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
+        id="my_awesome_follower_arm",
-kinematics_solver = RobotKinematics(
+        cameras=camera_config,
-    urdf_path="./SO101/so101_new_calib.urdf",
+        use_degrees=True,
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
-
-# Build pipeline to convert EE action to joints action
-robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
-    steps=[
-        InverseKinematicsEEToJoints(
-            kinematics=kinematics_solver,
-            motor_names=list(robot.bus.motors.keys()),
-            initial_guess_current_joints=True,
-        ),
-    ],
-    to_transition=robot_action_observation_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Build pipeline to convert joints observation to EE observation
-robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
-    steps=[
-        ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
-    ],
-    to_transition=observation_to_transition,
-    to_output=transition_to_observation,
-)
-
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id=HF_DATASET_ID,
-    fps=FPS,
-    features=combine_feature_dicts(
-        aggregate_pipeline_dataset_features(
-            pipeline=robot_joints_to_ee_pose_processor,
-            initial_features=create_initial_features(observation=robot.observation_features),
-            use_videos=True,
-        ),
-        # User for now should be explicit on the feature keys that were used for record
-        # Alternatively, the user can pass the processor step that has the right features
-        aggregate_pipeline_dataset_features(
-            pipeline=make_default_teleop_action_processor(),
-            initial_features=create_initial_features(
-                action={
-                    f"ee.{k}": PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-                    for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]
-                }
-            ),
-            use_videos=True,
-        ),
-    ),
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
-
-# Build Policy Processors
-preprocessor, postprocessor = make_pre_post_processors(
-    policy_cfg=policy,
-    pretrained_path=HF_MODEL_ID,
-    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.
-    preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
-)
-
-# Connect the robot
-robot.connect()
-
-# Initialize the keyboard listener and rerun visualization
-listener, events = init_keyboard_listener()
-init_rerun(session_name="phone_so100_evaluate")
-
-if not robot.is_connected:
-    raise ValueError("Robot is not connected!")
-
-print("Starting evaluate loop...")
-episode_idx = 0
-for episode_idx in range(NUM_EPISODES):
-    log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
-
-    # Main record loop
-    record_loop(
-        robot=robot,
-        events=events,
-        fps=FPS,
-        policy=policy,
-        preprocessor=preprocessor,  # Pass the pre and post policy processors
-        postprocessor=postprocessor,
-        dataset=dataset,
-        control_time_s=EPISODE_TIME_SEC,
-        single_task=TASK_DESCRIPTION,
-        display_data=True,
-        teleop_action_processor=make_default_teleop_action_processor(),
-        robot_action_processor=robot_ee_to_joints_processor,
-        robot_observation_processor=robot_joints_to_ee_pose_processor,
     )
 
-    # Reset the environment if not stopping or re-recording
+    robot = SO100Follower(robot_config)
-    if not events["stop_recording"] and ((episode_idx < NUM_EPISODES - 1) or events["rerecord_episode"]):
+
-        log_say("Reset the environment")
+    # Create policy
+    policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
+
+    # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
+    kinematics_solver = RobotKinematics(
+        urdf_path="./SO101/so101_new_calib.urdf",
+        target_frame_name="gripper_frame_link",
+        joint_names=list(robot.bus.motors.keys()),
+    )
+
+    # Build pipeline to convert EE action to joints action
+    robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+        steps=[
+            InverseKinematicsEEToJoints(
+                kinematics=kinematics_solver,
+                motor_names=list(robot.bus.motors.keys()),
+                initial_guess_current_joints=True,
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
+
+    # Build pipeline to convert joints observation to EE observation
+    robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
+        steps=[
+            ForwardKinematicsJointsToEE(
+                kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())
+            )
+        ],
+        to_transition=observation_to_transition,
+        to_output=transition_to_observation,
+    )
+
+    # Create the dataset
+    dataset = LeRobotDataset.create(
+        repo_id=HF_DATASET_ID,
+        fps=FPS,
+        features=combine_feature_dicts(
+            aggregate_pipeline_dataset_features(
+                pipeline=robot_joints_to_ee_pose_processor,
+                initial_features=create_initial_features(observation=robot.observation_features),
+                use_videos=True,
+            ),
+            # User for now should be explicit on the feature keys that were used for record
+            # Alternatively, the user can pass the processor step that has the right features
+            aggregate_pipeline_dataset_features(
+                pipeline=make_default_teleop_action_processor(),
+                initial_features=create_initial_features(
+                    action={
+                        f"ee.{k}": PolicyFeature(type=FeatureType.ACTION, shape=(1,))
+                        for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]
+                    }
+                ),
+                use_videos=True,
+            ),
+        ),
+        robot_type=robot.name,
+        use_videos=True,
+        image_writer_threads=4,
+    )
+
+    # Build Policy Processors
+    preprocessor, postprocessor = make_pre_post_processors(
+        policy_cfg=policy,
+        pretrained_path=HF_MODEL_ID,
+        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.
+        preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
+    )
+
+    # Connect the robot
+    robot.connect()
+
+    # Initialize the keyboard listener and rerun visualization
+    listener, events = init_keyboard_listener()
+    init_rerun(session_name="phone_so100_evaluate")
+
+    if not robot.is_connected:
+        raise ValueError("Robot is not connected!")
+
+    print("Starting evaluate loop...")
+    episode_idx = 0
+    for episode_idx in range(NUM_EPISODES):
+        log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
+
+        # Main record loop
         record_loop(
             robot=robot,
             events=events,
             fps=FPS,
+            policy=policy,
+            preprocessor=preprocessor,  # Pass the pre and post policy processors
+            postprocessor=postprocessor,
+            dataset=dataset,
             control_time_s=EPISODE_TIME_SEC,
             single_task=TASK_DESCRIPTION,
             display_data=True,
@@ -179,21 +168,40 @@ for episode_idx in range(NUM_EPISODES):
             robot_observation_processor=robot_joints_to_ee_pose_processor,
         )
 
-    if events["rerecord_episode"]:
+        # Reset the environment if not stopping or re-recording
-        log_say("Re-record episode")
+        if not events["stop_recording"] and ((episode_idx < NUM_EPISODES - 1) or events["rerecord_episode"]):
-        events["rerecord_episode"] = False
+            log_say("Reset the environment")
-        events["exit_early"] = False
+            record_loop(
-        dataset.clear_episode_buffer()
+                robot=robot,
-        continue
+                events=events,
+                fps=FPS,
+                control_time_s=EPISODE_TIME_SEC,
+                single_task=TASK_DESCRIPTION,
+                display_data=True,
+                teleop_action_processor=make_default_teleop_action_processor(),
+                robot_action_processor=robot_ee_to_joints_processor,
+                robot_observation_processor=robot_joints_to_ee_pose_processor,
+            )
 
-    # Save episode
+        if events["rerecord_episode"]:
-    dataset.save_episode()
+            log_say("Re-record episode")
-    episode_idx += 1
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-# Clean up
+        # Save episode
-log_say("Stop recording")
+        dataset.save_episode()
-robot.disconnect()
+        episode_idx += 1
-listener.stop()
 
-dataset.finalize()
+    # Clean up
-dataset.push_to_hub()
+    log_say("Stop recording")
+    robot.disconnect()
+    listener.stop()
+
+    dataset.finalize()
+    dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()

examples/phone_to_so100/record.py

@@ -50,133 +50,122 @@ RESET_TIME_SEC = 30
 TASK_DESCRIPTION = "My task description"
 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)}
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411",
-    id="my_awesome_follower_arm",
-    cameras=camera_config,
-    use_degrees=True,
-)
-teleop_config = PhoneConfig(phone_os=PhoneOS.IOS)  # or PhoneOS.ANDROID
 
-# Initialize the robot and teleoperator
+def main():
-robot = SO100Follower(robot_config)
+    # Create the robot and teleoperator configurations
-phone = Phone(teleop_config)
+    camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
+    robot_config = SO100FollowerConfig(
+        port="/dev/tty.usbmodem5A460814411",
+        id="my_awesome_follower_arm",
+        cameras=camera_config,
+        use_degrees=True,
+    )
+    teleop_config = PhoneConfig(phone_os=PhoneOS.IOS)  # or PhoneOS.ANDROID
 
-# 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
-kinematics_solver = RobotKinematics(
+    robot = SO100Follower(robot_config)
-    urdf_path="./SO101/so101_new_calib.urdf",
+    phone = Phone(teleop_config)
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
 
-# Build pipeline to convert phone action to EE action
+    # 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
-phone_to_robot_ee_pose_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+    kinematics_solver = RobotKinematics(
-    steps=[
+        urdf_path="./SO101/so101_new_calib.urdf",
-        MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
+        target_frame_name="gripper_frame_link",
-        EEReferenceAndDelta(
+        joint_names=list(robot.bus.motors.keys()),
-            kinematics=kinematics_solver,
-            end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
-            motor_names=list(robot.bus.motors.keys()),
-            use_latched_reference=True,
-        ),
-        EEBoundsAndSafety(
-            end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
-            max_ee_step_m=0.20,
-        ),
-        GripperVelocityToJoint(speed_factor=20.0),
-    ],
-    to_transition=robot_action_observation_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Build pipeline to convert EE action to joints action
-robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
-    steps=[
-        InverseKinematicsEEToJoints(
-            kinematics=kinematics_solver,
-            motor_names=list(robot.bus.motors.keys()),
-            initial_guess_current_joints=True,
-        ),
-    ],
-    to_transition=robot_action_observation_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Build pipeline to convert joint observation to EE observation
-robot_joints_to_ee_pose = RobotProcessorPipeline[RobotObservation, RobotObservation](
-    steps=[
-        ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
-    ],
-    to_transition=observation_to_transition,
-    to_output=transition_to_observation,
-)
-
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id=HF_REPO_ID,
-    fps=FPS,
-    features=combine_feature_dicts(
-        # Run the feature contract of the pipelines
-        # This tells you how the features would look like after the pipeline steps
-        aggregate_pipeline_dataset_features(
-            pipeline=phone_to_robot_ee_pose_processor,
-            initial_features=create_initial_features(action=phone.action_features),
-            use_videos=True,
-        ),
-        aggregate_pipeline_dataset_features(
-            pipeline=robot_joints_to_ee_pose,
-            initial_features=create_initial_features(observation=robot.observation_features),
-            use_videos=True,
-        ),
-    ),
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
-
-# Connect the robot and teleoperator
-robot.connect()
-phone.connect()
-
-# Initialize the keyboard listener and rerun visualization
-listener, events = init_keyboard_listener()
-init_rerun(session_name="phone_so100_record")
-
-if not robot.is_connected or not phone.is_connected:
-    raise ValueError("Robot or teleop is not connected!")
-
-
-print("Starting record loop. Move your phone to teleoperate the robot...")
-episode_idx = 0
-while episode_idx < NUM_EPISODES and not events["stop_recording"]:
-    log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
-
-    # Main record loop
-    record_loop(
-        robot=robot,
-        events=events,
-        fps=FPS,
-        teleop=phone,
-        dataset=dataset,
-        control_time_s=EPISODE_TIME_SEC,
-        single_task=TASK_DESCRIPTION,
-        display_data=True,
-        teleop_action_processor=phone_to_robot_ee_pose_processor,
-        robot_action_processor=robot_ee_to_joints_processor,
-        robot_observation_processor=robot_joints_to_ee_pose,
     )
 
-    # Reset the environment if not stopping or re-recording
+    # Build pipeline to convert phone action to EE action
-    if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
+    phone_to_robot_ee_pose_processor = RobotProcessorPipeline[
-        log_say("Reset the environment")
+        tuple[RobotAction, RobotObservation], RobotAction
+    ](
+        steps=[
+            MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
+            EEReferenceAndDelta(
+                kinematics=kinematics_solver,
+                end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
+                motor_names=list(robot.bus.motors.keys()),
+                use_latched_reference=True,
+            ),
+            EEBoundsAndSafety(
+                end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
+                max_ee_step_m=0.20,
+            ),
+            GripperVelocityToJoint(speed_factor=20.0),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
+
+    # Build pipeline to convert EE action to joints action
+    robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+        steps=[
+            InverseKinematicsEEToJoints(
+                kinematics=kinematics_solver,
+                motor_names=list(robot.bus.motors.keys()),
+                initial_guess_current_joints=True,
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
+
+    # Build pipeline to convert joint observation to EE observation
+    robot_joints_to_ee_pose = RobotProcessorPipeline[RobotObservation, RobotObservation](
+        steps=[
+            ForwardKinematicsJointsToEE(
+                kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())
+            )
+        ],
+        to_transition=observation_to_transition,
+        to_output=transition_to_observation,
+    )
+
+    # Create the dataset
+    dataset = LeRobotDataset.create(
+        repo_id=HF_REPO_ID,
+        fps=FPS,
+        features=combine_feature_dicts(
+            # Run the feature contract of the pipelines
+            # This tells you how the features would look like after the pipeline steps
+            aggregate_pipeline_dataset_features(
+                pipeline=phone_to_robot_ee_pose_processor,
+                initial_features=create_initial_features(action=phone.action_features),
+                use_videos=True,
+            ),
+            aggregate_pipeline_dataset_features(
+                pipeline=robot_joints_to_ee_pose,
+                initial_features=create_initial_features(observation=robot.observation_features),
+                use_videos=True,
+            ),
+        ),
+        robot_type=robot.name,
+        use_videos=True,
+        image_writer_threads=4,
+    )
+
+    # Connect the robot and teleoperator
+    robot.connect()
+    phone.connect()
+
+    # Initialize the keyboard listener and rerun visualization
+    listener, events = init_keyboard_listener()
+    init_rerun(session_name="phone_so100_record")
+
+    if not robot.is_connected or not phone.is_connected:
+        raise ValueError("Robot or teleop is not connected!")
+
+    print("Starting record loop. Move your phone to teleoperate the robot...")
+    episode_idx = 0
+    while episode_idx < NUM_EPISODES and not events["stop_recording"]:
+        log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
+
+        # Main record loop
         record_loop(
             robot=robot,
             events=events,
             fps=FPS,
             teleop=phone,
-            control_time_s=RESET_TIME_SEC,
+            dataset=dataset,
+            control_time_s=EPISODE_TIME_SEC,
             single_task=TASK_DESCRIPTION,
             display_data=True,
             teleop_action_processor=phone_to_robot_ee_pose_processor,
@@ -184,22 +173,42 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]:
             robot_observation_processor=robot_joints_to_ee_pose,
         )
 
-    if events["rerecord_episode"]:
+        # Reset the environment if not stopping or re-recording
-        log_say("Re-recording episode")
+        if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
-        events["rerecord_episode"] = False
+            log_say("Reset the environment")
-        events["exit_early"] = False
+            record_loop(
-        dataset.clear_episode_buffer()
+                robot=robot,
-        continue
+                events=events,
+                fps=FPS,
+                teleop=phone,
+                control_time_s=RESET_TIME_SEC,
+                single_task=TASK_DESCRIPTION,
+                display_data=True,
+                teleop_action_processor=phone_to_robot_ee_pose_processor,
+                robot_action_processor=robot_ee_to_joints_processor,
+                robot_observation_processor=robot_joints_to_ee_pose,
+            )
 
-    # Save episode
+        if events["rerecord_episode"]:
-    dataset.save_episode()
+            log_say("Re-recording episode")
-    episode_idx += 1
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-# Clean up
+        # Save episode
-log_say("Stop recording")
+        dataset.save_episode()
-robot.disconnect()
+        episode_idx += 1
-phone.disconnect()
-listener.stop()
 
-dataset.finalize()
+    # Clean up
-dataset.push_to_hub()
+    log_say("Stop recording")
+    robot.disconnect()
+    phone.disconnect()
+    listener.stop()
+
+    dataset.finalize()
+    dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()

examples/phone_to_so100/replay.py

@@ -35,66 +35,72 @@ from lerobot.utils.utils import log_say
 EPISODE_IDX = 0
 HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
 
-# Initialize the robot config
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
-)
 
-# Initialize the robot
+def main():
-robot = SO100Follower(robot_config)
+    # Initialize the robot config
+    robot_config = SO100FollowerConfig(
+        port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
+    )
 
-# 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
-kinematics_solver = RobotKinematics(
+    robot = SO100Follower(robot_config)
-    urdf_path="./SO101/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
 
-# Build pipeline to convert EE action to joints action
+    # 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
-robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+    kinematics_solver = RobotKinematics(
-    steps=[
+        urdf_path="./SO101/so101_new_calib.urdf",
-        InverseKinematicsEEToJoints(
+        target_frame_name="gripper_frame_link",
-            kinematics=kinematics_solver,
+        joint_names=list(robot.bus.motors.keys()),
-            motor_names=list(robot.bus.motors.keys()),
+    )
-            initial_guess_current_joints=False,  # Because replay is open loop
-        ),
-    ],
-    to_transition=robot_action_observation_to_transition,
-    to_output=transition_to_robot_action,
-)
 
-# Fetch the dataset to replay
+    # Build pipeline to convert EE action to joints action
-dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
+    robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
-# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
+        steps=[
-episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
+            InverseKinematicsEEToJoints(
-actions = episode_frames.select_columns(ACTION)
+                kinematics=kinematics_solver,
+                motor_names=list(robot.bus.motors.keys()),
+                initial_guess_current_joints=False,  # Because replay is open loop
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
 
-# Connect to the robot
+    # Fetch the dataset to replay
-robot.connect()
+    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
+    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
-    raise ValueError("Robot is not connected!")
+    robot.connect()
 
-print("Starting replay loop...")
+    if not robot.is_connected:
-log_say(f"Replaying episode {EPISODE_IDX}")
+        raise ValueError("Robot is not connected!")
-for idx in range(len(episode_frames)):
-    t0 = time.perf_counter()
 
-    # Get recorded action from dataset
+    print("Starting replay loop...")
-    ee_action = {
+    log_say(f"Replaying episode {EPISODE_IDX}")
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+    for idx in range(len(episode_frames)):
-    }
+        t0 = time.perf_counter()
 
-    # Get robot observation
+        # Get recorded action from dataset
-    robot_obs = robot.get_observation()
+        ee_action = {
+            name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        }
 
-    # Dataset EE -> robot joints
+        # Get robot observation
-    joint_action = robot_ee_to_joints_processor((ee_action, robot_obs))
+        robot_obs = robot.get_observation()
 
-    # Send action to robot
+        # Dataset EE -> robot joints
-    _ = robot.send_action(joint_action)
+        joint_action = robot_ee_to_joints_processor((ee_action, robot_obs))
 
-    busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))
+        # Send action to robot
+        _ = robot.send_action(joint_action)
 
-# Clean up
+        busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))
-robot.disconnect()
+
+    # Clean up
+    robot.disconnect()
+
+
+if __name__ == "__main__":
+    main()

examples/phone_to_so100/teleoperate.py

@@ -37,77 +37,85 @@ from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 
 FPS = 30
 
-# Initialize the robot and teleoperator
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
-)
-teleop_config = PhoneConfig(phone_os=PhoneOS.IOS)  # or PhoneOS.ANDROID
 
-# Initialize the robot and teleoperator
+def main():
-robot = SO100Follower(robot_config)
+    # Initialize the robot and teleoperator
-teleop_device = Phone(teleop_config)
+    robot_config = SO100FollowerConfig(
+        port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
+    )
+    teleop_config = PhoneConfig(phone_os=PhoneOS.IOS)  # or PhoneOS.ANDROID
 
-# 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
-kinematics_solver = RobotKinematics(
+    robot = SO100Follower(robot_config)
-    urdf_path="./SO101/so101_new_calib.urdf",
+    teleop_device = Phone(teleop_config)
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
 
-# Build pipeline to convert phone action to ee pose action to joint action
+    # 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
-phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+    kinematics_solver = RobotKinematics(
-    steps=[
+        urdf_path="./SO101/so101_new_calib.urdf",
-        MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
+        target_frame_name="gripper_frame_link",
-        EEReferenceAndDelta(
+        joint_names=list(robot.bus.motors.keys()),
-            kinematics=kinematics_solver,
+    )
-            end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
-            motor_names=list(robot.bus.motors.keys()),
-            use_latched_reference=True,
-        ),
-        EEBoundsAndSafety(
-            end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
-            max_ee_step_m=0.10,
-        ),
-        GripperVelocityToJoint(
-            speed_factor=20.0,
-        ),
-        InverseKinematicsEEToJoints(
-            kinematics=kinematics_solver,
-            motor_names=list(robot.bus.motors.keys()),
-            initial_guess_current_joints=True,
-        ),
-    ],
-    to_transition=robot_action_observation_to_transition,
-    to_output=transition_to_robot_action,
-)
 
-# Connect to the robot and teleoperator
+    # Build pipeline to convert phone action to ee pose action to joint action
-robot.connect()
+    phone_to_robot_joints_processor = RobotProcessorPipeline[
-teleop_device.connect()
+        tuple[RobotAction, RobotObservation], RobotAction
+    ](
+        steps=[
+            MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
+            EEReferenceAndDelta(
+                kinematics=kinematics_solver,
+                end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
+                motor_names=list(robot.bus.motors.keys()),
+                use_latched_reference=True,
+            ),
+            EEBoundsAndSafety(
+                end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
+                max_ee_step_m=0.10,
+            ),
+            GripperVelocityToJoint(
+                speed_factor=20.0,
+            ),
+            InverseKinematicsEEToJoints(
+                kinematics=kinematics_solver,
+                motor_names=list(robot.bus.motors.keys()),
+                initial_guess_current_joints=True,
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
 
-# Init rerun viewer
+    # Connect to the robot and teleoperator
-init_rerun(session_name="phone_so100_teleop")
+    robot.connect()
+    teleop_device.connect()
 
-if not robot.is_connected or not teleop_device.is_connected:
+    # Init rerun viewer
-    raise ValueError("Robot or teleop is not connected!")
+    init_rerun(session_name="phone_so100_teleop")
 
-print("Starting teleop loop. Move your phone to teleoperate the robot...")
+    if not robot.is_connected or not teleop_device.is_connected:
-while True:
+        raise ValueError("Robot or teleop is not connected!")
-    t0 = time.perf_counter()
 
-    # Get robot observation
+    print("Starting teleop loop. Move your phone to teleoperate the robot...")
-    robot_obs = robot.get_observation()
+    while True:
+        t0 = time.perf_counter()
 
-    # Get teleop action
+        # Get robot observation
-    phone_obs = teleop_device.get_action()
+        robot_obs = robot.get_observation()
 
-    # Phone -> EE pose -> Joints transition
+        # Get teleop action
-    joint_action = phone_to_robot_joints_processor((phone_obs, robot_obs))
+        phone_obs = teleop_device.get_action()
 
-    # Send action to robot
+        # Phone -> EE pose -> Joints transition
-    _ = robot.send_action(joint_action)
+        joint_action = phone_to_robot_joints_processor((phone_obs, robot_obs))
 
-    # Visualize
+        # Send action to robot
-    log_rerun_data(observation=phone_obs, action=joint_action)
+        _ = robot.send_action(joint_action)
 
-    busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
+        # Visualize
+        log_rerun_data(observation=phone_obs, action=joint_action)
+
+        busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
+
+
+if __name__ == "__main__":
+    main()

examples/so100_to_so100_EE/evaluate.py

@@ -52,126 +52,114 @@ TASK_DESCRIPTION = "My task description"
 HF_MODEL_ID = "<hf_username>/<model_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)}
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411",
-    id="my_awesome_follower_arm",
-    cameras=camera_config,
-    use_degrees=True,
-)
 
-robot = SO100Follower(robot_config)
+def main():
-
+    # Create the robot configuration & robot
-# Create policy
+    camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
-policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
+    robot_config = SO100FollowerConfig(
-
+        port="/dev/tty.usbmodem5A460814411",
-# 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
+        id="my_awesome_follower_arm",
-kinematics_solver = RobotKinematics(
+        cameras=camera_config,
-    urdf_path="./SO101/so101_new_calib.urdf",
+        use_degrees=True,
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
-
-# Build pipeline to convert EE action to joints action
-robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
-    steps=[
-        InverseKinematicsEEToJoints(
-            kinematics=kinematics_solver,
-            motor_names=list(robot.bus.motors.keys()),
-            initial_guess_current_joints=True,
-        ),
-    ],
-    to_transition=robot_action_observation_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Build pipeline to convert joints observation to EE observation
-robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
-    steps=[
-        ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
-    ],
-    to_transition=observation_to_transition,
-    to_output=transition_to_observation,
-)
-
-
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id=HF_DATASET_ID,
-    fps=FPS,
-    features=combine_feature_dicts(
-        aggregate_pipeline_dataset_features(
-            pipeline=robot_joints_to_ee_pose_processor,
-            initial_features=create_initial_features(observation=robot.observation_features),
-            use_videos=True,
-        ),
-        # User for now should be explicit on the feature keys that were used for record
-        # Alternatively, the user can pass the processor step that has the right features
-        aggregate_pipeline_dataset_features(
-            pipeline=make_default_teleop_action_processor(),
-            initial_features=create_initial_features(
-                action={
-                    f"ee.{k}": PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-                    for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]
-                }
-            ),
-            use_videos=True,
-        ),
-    ),
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
-
-# Build Policy Processors
-preprocessor, postprocessor = make_pre_post_processors(
-    policy_cfg=policy,
-    pretrained_path=HF_MODEL_ID,
-    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.
-    preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
-)
-
-# Connect the robot and teleoperator
-robot.connect()
-
-# Initialize the keyboard listener and rerun visualization
-listener, events = init_keyboard_listener()
-init_rerun(session_name="so100_so100_evaluate")
-
-if not robot.is_connected:
-    raise ValueError("Robot is not connected!")
-
-print("Starting evaluate loop...")
-episode_idx = 0
-for episode_idx in range(NUM_EPISODES):
-    log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
-
-    # Main record loop
-    record_loop(
-        robot=robot,
-        events=events,
-        fps=FPS,
-        policy=policy,
-        preprocessor=preprocessor,  # Pass the pre and post policy processors
-        postprocessor=postprocessor,
-        dataset=dataset,
-        control_time_s=EPISODE_TIME_SEC,
-        single_task=TASK_DESCRIPTION,
-        display_data=True,
-        teleop_action_processor=make_default_teleop_action_processor(),
-        robot_action_processor=robot_ee_to_joints_processor,
-        robot_observation_processor=robot_joints_to_ee_pose_processor,
     )
 
-    # Reset the environment if not stopping or re-recording
+    robot = SO100Follower(robot_config)
-    if not events["stop_recording"] and ((episode_idx < NUM_EPISODES - 1) or events["rerecord_episode"]):
+
-        log_say("Reset the environment")
+    # Create policy
+    policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
+
+    # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
+    kinematics_solver = RobotKinematics(
+        urdf_path="./SO101/so101_new_calib.urdf",
+        target_frame_name="gripper_frame_link",
+        joint_names=list(robot.bus.motors.keys()),
+    )
+
+    # Build pipeline to convert EE action to joints action
+    robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+        steps=[
+            InverseKinematicsEEToJoints(
+                kinematics=kinematics_solver,
+                motor_names=list(robot.bus.motors.keys()),
+                initial_guess_current_joints=True,
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
+
+    # Build pipeline to convert joints observation to EE observation
+    robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
+        steps=[
+            ForwardKinematicsJointsToEE(
+                kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())
+            )
+        ],
+        to_transition=observation_to_transition,
+        to_output=transition_to_observation,
+    )
+
+    # Create the dataset
+    dataset = LeRobotDataset.create(
+        repo_id=HF_DATASET_ID,
+        fps=FPS,
+        features=combine_feature_dicts(
+            aggregate_pipeline_dataset_features(
+                pipeline=robot_joints_to_ee_pose_processor,
+                initial_features=create_initial_features(observation=robot.observation_features),
+                use_videos=True,
+            ),
+            # User for now should be explicit on the feature keys that were used for record
+            # Alternatively, the user can pass the processor step that has the right features
+            aggregate_pipeline_dataset_features(
+                pipeline=make_default_teleop_action_processor(),
+                initial_features=create_initial_features(
+                    action={
+                        f"ee.{k}": PolicyFeature(type=FeatureType.ACTION, shape=(1,))
+                        for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]
+                    }
+                ),
+                use_videos=True,
+            ),
+        ),
+        robot_type=robot.name,
+        use_videos=True,
+        image_writer_threads=4,
+    )
+
+    # Build Policy Processors
+    preprocessor, postprocessor = make_pre_post_processors(
+        policy_cfg=policy,
+        pretrained_path=HF_MODEL_ID,
+        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.
+        preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
+    )
+
+    # Connect the robot and teleoperator
+    robot.connect()
+
+    # Initialize the keyboard listener and rerun visualization
+    listener, events = init_keyboard_listener()
+    init_rerun(session_name="so100_so100_evaluate")
+
+    if not robot.is_connected:
+        raise ValueError("Robot is not connected!")
+
+    print("Starting evaluate loop...")
+    episode_idx = 0
+    for episode_idx in range(NUM_EPISODES):
+        log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
+
+        # Main record loop
         record_loop(
             robot=robot,
             events=events,
             fps=FPS,
+            policy=policy,
+            preprocessor=preprocessor,  # Pass the pre and post policy processors
+            postprocessor=postprocessor,
+            dataset=dataset,
             control_time_s=EPISODE_TIME_SEC,
             single_task=TASK_DESCRIPTION,
             display_data=True,
@@ -180,21 +168,40 @@ for episode_idx in range(NUM_EPISODES):
             robot_observation_processor=robot_joints_to_ee_pose_processor,
         )
 
-    if events["rerecord_episode"]:
+        # Reset the environment if not stopping or re-recording
-        log_say("Re-record episode")
+        if not events["stop_recording"] and ((episode_idx < NUM_EPISODES - 1) or events["rerecord_episode"]):
-        events["rerecord_episode"] = False
+            log_say("Reset the environment")
-        events["exit_early"] = False
+            record_loop(
-        dataset.clear_episode_buffer()
+                robot=robot,
-        continue
+                events=events,
+                fps=FPS,
+                control_time_s=EPISODE_TIME_SEC,
+                single_task=TASK_DESCRIPTION,
+                display_data=True,
+                teleop_action_processor=make_default_teleop_action_processor(),
+                robot_action_processor=robot_ee_to_joints_processor,
+                robot_observation_processor=robot_joints_to_ee_pose_processor,
+            )
 
-    # Save episode
+        if events["rerecord_episode"]:
-    dataset.save_episode()
+            log_say("Re-record episode")
-    episode_idx += 1
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-# Clean up
+        # Save episode
-log_say("Stop recording")
+        dataset.save_episode()
-robot.disconnect()
+        episode_idx += 1
-listener.stop()
 
-dataset.finalize()
+    # Clean up
-dataset.push_to_hub()
+    log_say("Stop recording")
+    robot.disconnect()
+    listener.stop()
+
+    dataset.finalize()
+    dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()

examples/so100_to_so100_EE/record.py

@@ -48,134 +48,122 @@ RESET_TIME_SEC = 30
 TASK_DESCRIPTION = "My task description"
 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)
-    urdf_path="./SO101/so101_new_calib.urdf",
+    leader = SO100Leader(leader_config)
-    target_frame_name="gripper_frame_link",
-    joint_names=list(follower.bus.motors.keys()),
-)
 
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
+    # 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(
+    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(leader.bus.motors.keys()),
+        joint_names=list(follower.bus.motors.keys()),
-)
-
-# Build pipeline to convert follower joints to EE observation
-follower_joints_to_ee = RobotProcessorPipeline[RobotObservation, RobotObservation](
-    steps=[
-        ForwardKinematicsJointsToEE(
-            kinematics=follower_kinematics_solver, motor_names=list(follower.bus.motors.keys())
-        ),
-    ],
-    to_transition=observation_to_transition,
-    to_output=transition_to_observation,
-)
-
-# Build pipeline to convert leader joints to EE action
-leader_joints_to_ee = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
-    steps=[
-        ForwardKinematicsJointsToEE(
-            kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
-        ),
-    ],
-    to_transition=robot_action_observation_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Build pipeline to convert EE action to follower joints
-ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
-    [
-        EEBoundsAndSafety(
-            end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
-            max_ee_step_m=0.10,
-        ),
-        InverseKinematicsEEToJoints(
-            kinematics=follower_kinematics_solver,
-            motor_names=list(follower.bus.motors.keys()),
-            initial_guess_current_joints=True,
-        ),
-    ],
-    to_transition=robot_action_observation_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id=HF_REPO_ID,
-    fps=FPS,
-    features=combine_feature_dicts(
-        # Run the feature contract of the pipelines
-        # This tells you how the features would look like after the pipeline steps
-        aggregate_pipeline_dataset_features(
-            pipeline=leader_joints_to_ee,
-            initial_features=create_initial_features(action=leader.action_features),
-            use_videos=True,
-        ),
-        aggregate_pipeline_dataset_features(
-            pipeline=follower_joints_to_ee,
-            initial_features=create_initial_features(observation=follower.observation_features),
-            use_videos=True,
-        ),
-    ),
-    robot_type=follower.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
-
-
-# Connect the robot and teleoperator
-leader.connect()
-follower.connect()
-
-# Initialize the keyboard listener and rerun visualization
-listener, events = init_keyboard_listener()
-init_rerun(session_name="recording_phone")
-
-if not leader.is_connected or not follower.is_connected:
-    raise ValueError("Robot or teleop is not connected!")
-
-print("Starting record loop...")
-episode_idx = 0
-while episode_idx < NUM_EPISODES and not events["stop_recording"]:
-    log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
-
-    # Main record loop
-    record_loop(
-        robot=follower,
-        events=events,
-        fps=FPS,
-        teleop=leader,
-        dataset=dataset,
-        control_time_s=EPISODE_TIME_SEC,
-        single_task=TASK_DESCRIPTION,
-        display_data=True,
-        teleop_action_processor=leader_joints_to_ee,
-        robot_action_processor=ee_to_follower_joints,
-        robot_observation_processor=follower_joints_to_ee,
     )
 
-    # Reset the environment if not stopping or re-recording
+    # 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
-    if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
+    leader_kinematics_solver = RobotKinematics(
-        log_say("Reset the environment")
+        urdf_path="./SO101/so101_new_calib.urdf",
+        target_frame_name="gripper_frame_link",
+        joint_names=list(leader.bus.motors.keys()),
+    )
+
+    # Build pipeline to convert follower joints to EE observation
+    follower_joints_to_ee = RobotProcessorPipeline[RobotObservation, RobotObservation](
+        steps=[
+            ForwardKinematicsJointsToEE(
+                kinematics=follower_kinematics_solver, motor_names=list(follower.bus.motors.keys())
+            ),
+        ],
+        to_transition=observation_to_transition,
+        to_output=transition_to_observation,
+    )
+
+    # Build pipeline to convert leader joints to EE action
+    leader_joints_to_ee = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+        steps=[
+            ForwardKinematicsJointsToEE(
+                kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
+
+    # Build pipeline to convert EE action to follower joints
+    ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+        [
+            EEBoundsAndSafety(
+                end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
+                max_ee_step_m=0.10,
+            ),
+            InverseKinematicsEEToJoints(
+                kinematics=follower_kinematics_solver,
+                motor_names=list(follower.bus.motors.keys()),
+                initial_guess_current_joints=True,
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
+
+    # Create the dataset
+    dataset = LeRobotDataset.create(
+        repo_id=HF_REPO_ID,
+        fps=FPS,
+        features=combine_feature_dicts(
+            # Run the feature contract of the pipelines
+            # This tells you how the features would look like after the pipeline steps
+            aggregate_pipeline_dataset_features(
+                pipeline=leader_joints_to_ee,
+                initial_features=create_initial_features(action=leader.action_features),
+                use_videos=True,
+            ),
+            aggregate_pipeline_dataset_features(
+                pipeline=follower_joints_to_ee,
+                initial_features=create_initial_features(observation=follower.observation_features),
+                use_videos=True,
+            ),
+        ),
+        robot_type=follower.name,
+        use_videos=True,
+        image_writer_threads=4,
+    )
+
+    # Connect the robot and teleoperator
+    leader.connect()
+    follower.connect()
+
+    # Initialize the keyboard listener and rerun visualization
+    listener, events = init_keyboard_listener()
+    init_rerun(session_name="recording_phone")
+
+    if not leader.is_connected or not follower.is_connected:
+        raise ValueError("Robot or teleop is not connected!")
+
+    print("Starting record loop...")
+    episode_idx = 0
+    while episode_idx < NUM_EPISODES and not events["stop_recording"]:
+        log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
+
+        # Main record loop
         record_loop(
             robot=follower,
             events=events,
             fps=FPS,
             teleop=leader,
-            control_time_s=RESET_TIME_SEC,
+            dataset=dataset,
+            control_time_s=EPISODE_TIME_SEC,
             single_task=TASK_DESCRIPTION,
             display_data=True,
             teleop_action_processor=leader_joints_to_ee,
@@ -183,22 +171,42 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]:
             robot_observation_processor=follower_joints_to_ee,
         )
 
-    if events["rerecord_episode"]:
+        # Reset the environment if not stopping or re-recording
-        log_say("Re-recording episode")
+        if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
-        events["rerecord_episode"] = False
+            log_say("Reset the environment")
-        events["exit_early"] = False
+            record_loop(
-        dataset.clear_episode_buffer()
+                robot=follower,
-        continue
+                events=events,
+                fps=FPS,
+                teleop=leader,
+                control_time_s=RESET_TIME_SEC,
+                single_task=TASK_DESCRIPTION,
+                display_data=True,
+                teleop_action_processor=leader_joints_to_ee,
+                robot_action_processor=ee_to_follower_joints,
+                robot_observation_processor=follower_joints_to_ee,
+            )
 
-    # Save episode
+        if events["rerecord_episode"]:
-    dataset.save_episode()
+            log_say("Re-recording episode")
-    episode_idx += 1
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-# Clean up
+        # Save episode
-log_say("Stop recording")
+        dataset.save_episode()
-leader.disconnect()
+        episode_idx += 1
-follower.disconnect()
-listener.stop()
 
-dataset.finalize()
+    # Clean up
-dataset.push_to_hub()
+    log_say("Stop recording")
+    leader.disconnect()
+    follower.disconnect()
+    listener.stop()
+
+    dataset.finalize()
+    dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()

examples/so100_to_so100_EE/replay.py

@@ -36,66 +36,72 @@ from lerobot.utils.utils import log_say
 EPISODE_IDX = 0
 HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
 
-# Initialize the robot config
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
-)
 
-# Initialize the robot
+def main():
-robot = SO100Follower(robot_config)
+    # Initialize the robot config
+    robot_config = SO100FollowerConfig(
+        port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
+    )
 
-# 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
-kinematics_solver = RobotKinematics(
+    robot = SO100Follower(robot_config)
-    urdf_path="./SO101/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
 
-# Build pipeline to convert EE action to joints action
+    # 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
-robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+    kinematics_solver = RobotKinematics(
-    steps=[
+        urdf_path="./SO101/so101_new_calib.urdf",
-        InverseKinematicsEEToJoints(
+        target_frame_name="gripper_frame_link",
-            kinematics=kinematics_solver,
+        joint_names=list(robot.bus.motors.keys()),
-            motor_names=list(robot.bus.motors.keys()),
+    )
-            initial_guess_current_joints=False,  # Because replay is open loop
-        ),
-    ],
-    to_transition=robot_action_observation_to_transition,
-    to_output=transition_to_robot_action,
-)
 
-# Fetch the dataset to replay
+    # Build pipeline to convert EE action to joints action
-dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
+    robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
-# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
+        steps=[
-episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
+            InverseKinematicsEEToJoints(
-actions = episode_frames.select_columns(ACTION)
+                kinematics=kinematics_solver,
+                motor_names=list(robot.bus.motors.keys()),
+                initial_guess_current_joints=False,  # Because replay is open loop
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
 
-# Connect to the robot
+    # Fetch the dataset to replay
-robot.connect()
+    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
+    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
-    raise ValueError("Robot is not connected!")
+    robot.connect()
 
-print("Starting replay loop...")
+    if not robot.is_connected:
-log_say(f"Replaying episode {EPISODE_IDX}")
+        raise ValueError("Robot is not connected!")
-for idx in range(len(episode_frames)):
-    t0 = time.perf_counter()
 
-    # Get recorded action from dataset
+    print("Starting replay loop...")
-    ee_action = {
+    log_say(f"Replaying episode {EPISODE_IDX}")
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+    for idx in range(len(episode_frames)):
-    }
+        t0 = time.perf_counter()
 
-    # Get robot observation
+        # Get recorded action from dataset
-    robot_obs = robot.get_observation()
+        ee_action = {
+            name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        }
 
-    # Dataset EE -> robot joints
+        # Get robot observation
-    joint_action = robot_ee_to_joints_processor((ee_action, robot_obs))
+        robot_obs = robot.get_observation()
 
-    # Send action to robot
+        # Dataset EE -> robot joints
-    _ = robot.send_action(joint_action)
+        joint_action = robot_ee_to_joints_processor((ee_action, robot_obs))
 
-    busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))
+        # Send action to robot
+        _ = robot.send_action(joint_action)
 
-# Clean up
+        busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))
-robot.disconnect()
+
+    # Clean up
+    robot.disconnect()
+
+
+if __name__ == "__main__":
+    main()

examples/so100_to_so100_EE/teleoperate.py

@@ -37,85 +37,91 @@ from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 
 FPS = 30
 
-# Initialize the robot and teleoperator config
-follower_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
-)
-leader_config = SO100LeaderConfig(port="/dev/tty.usbmodem5A460819811", id="my_awesome_leader_arm")
 
-# Initialize the robot and teleoperator
+def main():
-follower = SO100Follower(follower_config)
+    # Initialize the robot and teleoperator config
-leader = SO100Leader(leader_config)
+    follower_config = SO100FollowerConfig(
+        port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
+    )
+    leader_config = SO100LeaderConfig(port="/dev/tty.usbmodem5A460819811", id="my_awesome_leader_arm")
 
-# 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)
-    urdf_path="./SO101/so101_new_calib.urdf",
+    leader = SO100Leader(leader_config)
-    target_frame_name="gripper_frame_link",
-    joint_names=list(follower.bus.motors.keys()),
-)
 
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
+    # 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(
+    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(leader.bus.motors.keys()),
+        joint_names=list(follower.bus.motors.keys()),
-)
+    )
 
-# Build pipeline to convert teleop joints to EE action
+    # 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_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
+    leader_kinematics_solver = RobotKinematics(
-    steps=[
+        urdf_path="./SO101/so101_new_calib.urdf",
-        ForwardKinematicsJointsToEE(
+        target_frame_name="gripper_frame_link",
-            kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
+        joint_names=list(leader.bus.motors.keys()),
-        ),
+    )
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
-)
 
-# build pipeline to convert EE action to robot joints
+    # Build pipeline to convert teleop joints to EE action
-ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+    leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
-    [
+        steps=[
-        EEBoundsAndSafety(
+            ForwardKinematicsJointsToEE(
-            end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
+                kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
-            max_ee_step_m=0.10,
+            ),
-        ),
+        ],
-        InverseKinematicsEEToJoints(
+        to_transition=robot_action_to_transition,
-            kinematics=follower_kinematics_solver,
+        to_output=transition_to_robot_action,
-            motor_names=list(follower.bus.motors.keys()),
+    )
-            initial_guess_current_joints=False,
-        ),
-    ],
-    to_transition=robot_action_observation_to_transition,
-    to_output=transition_to_robot_action,
-)
 
-# Connect to the robot and teleoperator
+    # build pipeline to convert EE action to robot joints
-follower.connect()
+    ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
-leader.connect()
+        [
+            EEBoundsAndSafety(
+                end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
+                max_ee_step_m=0.10,
+            ),
+            InverseKinematicsEEToJoints(
+                kinematics=follower_kinematics_solver,
+                motor_names=list(follower.bus.motors.keys()),
+                initial_guess_current_joints=False,
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
 
-# Init rerun viewer
+    # Connect to the robot and teleoperator
-init_rerun(session_name="so100_so100_EE_teleop")
+    follower.connect()
+    leader.connect()
 
-print("Starting teleop loop...")
+    # Init rerun viewer
-while True:
+    init_rerun(session_name="so100_so100_EE_teleop")
-    t0 = time.perf_counter()
 
-    # Get robot observation
+    print("Starting teleop loop...")
-    robot_obs = follower.get_observation()
+    while True:
+        t0 = time.perf_counter()
 
-    # Get teleop observation
+        # Get robot observation
-    leader_joints_obs = leader.get_action()
+        robot_obs = follower.get_observation()
 
-    # teleop joints -> teleop EE action
+        # Get teleop observation
-    leader_ee_act = leader_to_ee(leader_joints_obs)
+        leader_joints_obs = leader.get_action()
 
-    # teleop EE -> robot joints
+        # teleop joints -> teleop EE action
-    follower_joints_act = ee_to_follower_joints((leader_ee_act, robot_obs))
+        leader_ee_act = leader_to_ee(leader_joints_obs)
 
-    # Send action to robot
+        # teleop EE -> robot joints
-    _ = follower.send_action(follower_joints_act)
+        follower_joints_act = ee_to_follower_joints((leader_ee_act, robot_obs))
 
-    # Visualize
+        # Send action to robot
-    log_rerun_data(observation=leader_ee_act, action=follower_joints_act)
+        _ = follower.send_action(follower_joints_act)
 
-    busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
+        # Visualize
+        log_rerun_data(observation=leader_ee_act, action=follower_joints_act)
+
+        busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
+
+
+if __name__ == "__main__":
+    main()

examples/tutorial/act/act_training_example.py

@@ -19,80 +19,86 @@ def make_delta_timestamps(delta_indices: list[int] | None, fps: int) -> list[flo
     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)
-        loss.backward()
+            loss.backward()
-        optimizer.step()
+            optimizer.step()
-        optimizer.zero_grad()
+            optimizer.zero_grad()
 
-        if step % log_freq == 0:
+            if step % log_freq == 0:
-            print(f"step: {step} loss: {loss.item():.3f}")
+                print(f"step: {step} loss: {loss.item():.3f}")
-        step += 1
+            step += 1
-        if step >= training_steps:
+            if step >= training_steps:
-            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

@@ -8,50 +8,56 @@ 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.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_STEPS_PER_EPISODE = 20
 
-# 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),
-    "up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
-}
 
-robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
+def main():
-robot = SO100Follower(robot_cfg)
+    device = torch.device("mps")  # or "cuda" or "cpu"
-robot.connect()
+    model_id = "<user>/robot_learning_tutorial_act"
+    model = ACTPolicy.from_pretrained(model_id)
 
-for _ in range(MAX_EPISODES):
+    dataset_id = "lerobot/svla_so101_pickplace"
-    for _ in range(MAX_STEPS_PER_EPISODE):
+    # This only downloads the metadata for the dataset, ~10s of MB even for large-scale datasets
-        obs = robot.get_observation()
+    dataset_metadata = LeRobotDatasetMetadata(dataset_id)
-        obs_frame = build_inference_frame(
+    preprocess, postprocess = make_pre_post_processors(model.config, dataset_stats=dataset_metadata.stats)
-            observation=obs, ds_features=dataset_metadata.features, device=device
-        )
 
-        obs = preprocess(obs_frame)
+    # # find ports using lerobot-find-port
+    follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
 
-        action = model.select_action(obs)
+    # # the robot ids are used the load the right calibration files
-        action = postprocess(action)
+    follower_id = ...  # something like "follower_so100"
 
-        action = make_robot_action(action, dataset_metadata.features)
+    # 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),
+        "up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
+    }
 
-        robot.send_action(action)
+    robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
+    robot = SO100Follower(robot_cfg)
+    robot.connect()
 
-    print("Episode finished! Starting new episode...")
+    for _ in range(MAX_EPISODES):
+        for _ in range(MAX_STEPS_PER_EPISODE):
+            obs = robot.get_observation()
+            obs_frame = build_inference_frame(
+                observation=obs, ds_features=dataset_metadata.features, device=device
+            )
+
+            obs = preprocess(obs_frame)
+
+            action = model.select_action(obs)
+            action = postprocess(action)
+
+            action = make_robot_action(action, dataset_metadata.features)
+
+            robot.send_action(action)
+
+        print("Episode finished! Starting new episode...")
+
+
+if __name__ == "__main__":
+    main()

examples/tutorial/async-inf/policy_server.py

@@ -1,11 +1,17 @@
 from lerobot.async_inference.configs import PolicyServerConfig
 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=host,
+    host = ...  # something like "127.0.0.1" if you're exposing to localhost
-    port=port,
+    port = ...  # something like 8080
-)
+
-serve(config)
+    config = PolicyServerConfig(
+        host=host,
+        port=port,
+    )
+    serve(config)
+
+
+if __name__ == "__main__":
+    main()

examples/tutorial/async-inf/robot_client.py

@@ -6,50 +6,56 @@ from lerobot.async_inference.robot_client import RobotClient
 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
 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
-camera_cfg = {
-    "up": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
-    "side": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
-}
 
-# # find ports using lerobot-find-port
+def main():
-follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
+    # 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),
+        "side": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
+    }
 
-# # the robot ids are used the load the right calibration files
+    # # find ports using lerobot-find-port
-follower_id = ...  # something like "follower_so100"
+    follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
 
-robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_cfg)
+    # # the robot ids are used the load the right calibration files
+    follower_id = ...  # something like "follower_so100"
 
-server_address = ...  # something like "127.0.0.1:8080" if using localhost
+    robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_cfg)
 
-# 3. Create client configuration
+    server_address = ...  # something like "127.0.0.1:8080" if using localhost
-client_cfg = RobotClientConfig(
-    robot=robot_cfg,
-    server_address=server_address,
-    policy_device="mps",
-    policy_type="act",
-    pretrained_name_or_path="fracapuano/robot_learning_tutorial_act",
-    chunk_size_threshold=0.5,  # g
-    actions_per_chunk=50,  # make sure this is less than the max actions of the policy
-)
 
-# 4. Create and start client
+    # 3. Create client configuration
-client = RobotClient(client_cfg)
+    client_cfg = RobotClientConfig(
+        robot=robot_cfg,
+        server_address=server_address,
+        policy_device="mps",
+        policy_type="act",
+        pretrained_name_or_path="<user>/robot_learning_tutorial_act",
+        chunk_size_threshold=0.5,  # g
+        actions_per_chunk=50,  # make sure this is less than the max actions of the policy
+    )
 
-# 5. Provide a textual description of the task
+    # 4. Create and start client
-task = ...
+    client = RobotClient(client_cfg)
 
-if client.start():
+    # 5. Provide a textual description of the task
-    # Start action receiver thread
+    task = ...
-    action_receiver_thread = threading.Thread(target=client.receive_actions, daemon=True)
-    action_receiver_thread.start()
 
-    try:
+    if client.start():
-        # Run the control loop
+        # Start action receiver thread
-        client.control_loop(task)
+        action_receiver_thread = threading.Thread(target=client.receive_actions, daemon=True)
-    except KeyboardInterrupt:
+        action_receiver_thread.start()
-        client.stop()
+
-        action_receiver_thread.join()
+        try:
-        # (Optionally) plot the action queue size
+            # Run the control loop
-        visualize_action_queue_size(client.action_queue_size)
+            client.control_loop(task)
+        except KeyboardInterrupt:
+            client.stop()
+            action_receiver_thread.join()
+            # (Optionally) plot the action queue size
+            visualize_action_queue_size(client.action_queue_size)
+
+
+if __name__ == "__main__":
+    main()

examples/tutorial/diffusion/diffusion_training_example.py

@@ -19,81 +19,87 @@ def make_delta_timestamps(delta_indices: list[int] | None, fps: int) -> list[flo
     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)
-        loss.backward()
+            loss.backward()
-        optimizer.step()
+            optimizer.step()
-        optimizer.zero_grad()
+            optimizer.zero_grad()
 
-        if step % log_freq == 0:
+            if step % log_freq == 0:
-            print(f"step: {step} loss: {loss.item():.3f}")
+                print(f"step: {step} loss: {loss.item():.3f}")
-        step += 1
+            step += 1
-        if step >= training_steps:
+            if step >= training_steps:
-            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

@@ -8,53 +8,57 @@ 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.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_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(
-    "side": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
+        model.config, model_id, dataset_stats=dataset_metadata.stats
-    "up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
+    )
-}
 
-robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
+    # # find ports using lerobot-find-port
-robot = SO100Follower(robot_cfg)
+    follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
-robot.connect()
+
+    # # 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),
+        "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 = SO100Follower(robot_cfg)
+    robot.connect()
+
+    for _ in range(MAX_EPISODES):
+        for _ in range(MAX_STEPS_PER_EPISODE):
+            obs = robot.get_observation()
+            obs_frame = build_inference_frame(
+                observation=obs, ds_features=dataset_metadata.features, device=device
+            )
+
+            obs = preprocess(obs_frame)
+
+            action = model.select_action(obs)
+            action = postprocess(action)
+            action = make_robot_action(action, dataset_metadata.features)
+            robot.send_action(action)
+
+        print("Episode finished! Starting new episode...")
 
 
-for _ in range(MAX_EPISODES):
+if __name__ == "__main__":
-    for _ in range(MAX_STEPS_PER_EPISODE):
+    main()
-        obs = robot.get_observation()
-        obs_frame = build_inference_frame(
-            observation=obs, ds_features=dataset_metadata.features, device=device
-        )
-
-        obs = preprocess(obs_frame)
-
-        action = model.select_action(obs)
-        action = postprocess(action)
-        action = make_robot_action(action, dataset_metadata.features)
-        robot.send_action(action)
-
-    print("Episode finished! Starting new episode...")

examples/tutorial/pi0/using_pi0_example.py

@@ -11,57 +11,63 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower
 MAX_EPISODES = 5
 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)
-    model.config,
-    model_id,
-    # This overrides allows to run on MPS, otherwise defaults to CUDA (if available)
-    preprocessor_overrides={"device_processor": {"device": str(device)}},
-)
 
-# find ports using lerobot-find-port
+    preprocess, postprocess = make_pre_post_processors(
-follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
+        model.config,
+        model_id,
+        # This overrides allows to run on MPS, otherwise defaults to CUDA (if available)
+        preprocessor_overrides={"device_processor": {"device": str(device)}},
+    )
 
-# the robot ids are used the load the right calibration files
+    # find ports using lerobot-find-port
-follower_id = ...  # something like "follower_so100"
+    follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
 
-# Robot and environment configuration
+    # the robot ids are used the load the right calibration files
-# Camera keys must match the name and resolutions of the ones used for training!
+    follower_id = ...  # something like "follower_so100"
-# You can check the camera keys expected by a model in the info.json card on the model card on the Hub
-camera_config = {
-    "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),
-    "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 and environment configuration
-robot = SO100Follower(robot_cfg)
+    # Camera keys must match the name and resolutions of the ones used for training!
-robot.connect()
+    # You can check the camera keys expected by a model in the info.json card on the model card on the Hub
+    camera_config = {
+        "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),
+        "right_wrist_0_rgb": OpenCVCameraConfig(index_or_path=2, width=640, height=480, fps=30),
+    }
 
-task = ""  # something like "pick the red block"
+    robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
-robot_type = ""  # something like "so100_follower" for multi-embodiment datasets
+    robot = SO100Follower(robot_cfg)
+    robot.connect()
 
-# This is used to match the raw observation keys to the keys expected by the policy
+    task = ""  # something like "pick the red block"
-action_features = hw_to_dataset_features(robot.action_features, "action")
+    robot_type = ""  # something like "so100_follower" for multi-embodiment datasets
-obs_features = hw_to_dataset_features(robot.observation_features, "observation")
-dataset_features = {**action_features, **obs_features}
 
-for _ in range(MAX_EPISODES):
+    # This is used to match the raw observation keys to the keys expected by the policy
-    for _ in range(MAX_STEPS_PER_EPISODE):
+    action_features = hw_to_dataset_features(robot.action_features, "action")
-        obs = robot.get_observation()
+    obs_features = hw_to_dataset_features(robot.observation_features, "observation")
-        obs_frame = build_inference_frame(
+    dataset_features = {**action_features, **obs_features}
-            observation=obs, ds_features=dataset_features, device=device, task=task, robot_type=robot_type
-        )
 
-        obs = preprocess(obs_frame)
+    for _ in range(MAX_EPISODES):
+        for _ in range(MAX_STEPS_PER_EPISODE):
+            obs = robot.get_observation()
+            obs_frame = build_inference_frame(
+                observation=obs, ds_features=dataset_features, device=device, task=task, robot_type=robot_type
+            )
 
-        action = model.select_action(obs)
+            obs = preprocess(obs_frame)
-        action = postprocess(action)
-        action = make_robot_action(action, dataset_features)
-        robot.send_action(action)
 
-    print("Episode finished! Starting new episode...")
+            action = model.select_action(obs)
+            action = postprocess(action)
+            action = make_robot_action(action, dataset_features)
+            robot.send_action(action)
+
+        print("Episode finished! Starting new episode...")
+
+
+if __name__ == "__main__":
+    main()

examples/tutorial/rl/hilserl_example.py

@@ -20,6 +20,8 @@ from lerobot.teleoperators.utils import TeleopEvents
 
 LOG_EVERY = 10
 SEND_EVERY = 10
+MAX_EPISODES = 5
+MAX_STEPS_PER_EPISODE = 20
 
 
 def run_learner(
@@ -223,123 +225,123 @@ 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(
+        device=device,
+        input_features=obs_features,
+        output_features=action_features,
+    )
 
-# Create SAC policy for action selection
+    policy_actor = SACPolicy(policy_cfg)
-policy_cfg = SACConfig(
+    policy_learner = SACPolicy(policy_cfg)
-    device=device,
-    input_features=obs_features,
-    output_features=action_features,
-)
 
-policy_actor = SACPolicy(policy_cfg)
+    demonstrations_repo_id = "lerobot/example_hil_serl_dataset"
-policy_learner = SACPolicy(policy_cfg)
+    offline_dataset = LeRobotDataset(repo_id=demonstrations_repo_id)
 
-demonstrations_repo_id = "lerobot/example_hil_serl_dataset"
+    # Online buffer: initialized from scratch
-offline_dataset = LeRobotDataset(repo_id=demonstrations_repo_id)
+    online_replay_buffer = ReplayBuffer(device=device, state_keys=list(obs_features.keys()))
+    # Offline buffer: Created from dataset (pre-populated it with demonstrations)
+    offline_replay_buffer = ReplayBuffer.from_lerobot_dataset(
+        lerobot_dataset=offline_dataset, device=device, state_keys=list(obs_features.keys())
+    )
 
-# Online buffer: initialized from scratch
+    # Create communication channels between learner and actor processes
-online_replay_buffer = ReplayBuffer(device=device, state_keys=list(obs_features.keys()))
+    transitions_queue = mp.Queue(maxsize=10)
-# Offline buffer: Created from dataset (pre-populated it with demonstrations)
+    parameters_queue = mp.Queue(maxsize=2)
-offline_replay_buffer = ReplayBuffer.from_lerobot_dataset(
+    shutdown_event = mp.Event()
-    lerobot_dataset=offline_dataset, device=device, state_keys=list(obs_features.keys())
-)
 
-# Create communication channels between learner and actor processes
+    # Signal handler for graceful shutdown
-transitions_queue = mp.Queue(maxsize=10)
+    def signal_handler(sig):
-parameters_queue = mp.Queue(maxsize=2)
+        print(f"\nSignal {sig} received, shutting down...")
-shutdown_event = mp.Event()
+        shutdown_event.set()
+
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
+
+    # Create processes
+    learner_process = mp.Process(
+        target=run_learner,
+        args=(
+            transitions_queue,
+            parameters_queue,
+            shutdown_event,
+            policy_learner,
+            online_replay_buffer,
+            offline_replay_buffer,
+        ),
+        kwargs={"device": device},  # can run on accelerated hardware for training
+    )
+
+    actor_process = mp.Process(
+        target=run_actor,
+        args=(
+            transitions_queue,
+            parameters_queue,
+            shutdown_event,
+            policy_actor,
+            reward_classifier,
+            env_cfg,
+            output_directory,
+        ),
+        kwargs={"device": "cpu"},  # actor is frozen, can run on CPU or accelerate for inference
+    )
+
+    learner_process.start()
+    actor_process.start()
+
+    try:
+        # Wait for actor to finish (it controls the episode loop)
+        actor_process.join()
+        shutdown_event.set()
+        learner_process.join(timeout=10)
+
+    except KeyboardInterrupt:
+        print("Main process interrupted")
+        shutdown_event.set()
+        actor_process.join(timeout=5)
+        learner_process.join(timeout=10)
+
+    finally:
+        if learner_process.is_alive():
+            learner_process.terminate()
+        if actor_process.is_alive():
+            actor_process.terminate()
 
 
-# Signal handler for graceful shutdown
+if __name__ == "__main__":
-def signal_handler(sig):
+    main()
-    print(f"\nSignal {sig} received, shutting down...")
-    shutdown_event.set()
-
-
-signal.signal(signal.SIGINT, signal_handler)
-signal.signal(signal.SIGTERM, signal_handler)
-
-# Create processes
-learner_process = mp.Process(
-    target=run_learner,
-    args=(
-        transitions_queue,
-        parameters_queue,
-        shutdown_event,
-        policy_learner,
-        online_replay_buffer,
-        offline_replay_buffer,
-    ),
-    kwargs={"device": device},  # can run on accelerated hardware for training
-)
-
-actor_process = mp.Process(
-    target=run_actor,
-    args=(
-        transitions_queue,
-        parameters_queue,
-        shutdown_event,
-        policy_actor,
-        reward_classifier,
-        env_cfg,
-        output_directory,
-    ),
-    kwargs={"device": "cpu"},  # actor is frozen, can run on CPU or accelerate for inference
-)
-
-learner_process.start()
-actor_process.start()
-
-try:
-    # Wait for actor to finish (it controls the episode loop)
-    actor_process.join()
-    shutdown_event.set()
-    learner_process.join(timeout=10)
-
-except KeyboardInterrupt:
-    print("Main process interrupted")
-    shutdown_event.set()
-    actor_process.join(timeout=5)
-    learner_process.join(timeout=10)
-
-finally:
-    if learner_process.is_alive():
-        learner_process.terminate()
-    if actor_process.is_alive():
-        actor_process.terminate()

examples/tutorial/rl/reward_classifier_example.py

@@ -4,59 +4,64 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.policies.factory import make_policy, make_pre_post_processors
 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
-    num_cameras=len(camera_keys),
+    camera_keys = dataset.meta.camera_keys
-    device=device,
-    # backbone model to extract features from the image frames
-    model_name="microsoft/resnet-18",
-)
 
-# Make policy, preprocessor, and optimizer
+    config = RewardClassifierConfig(
-policy = make_policy(config, ds_meta=dataset.meta)
+        num_cameras=len(camera_keys),
-optimizer = config.get_optimizer_preset().build(policy.parameters())
+        device=device,
-preprocessor, _ = make_pre_post_processors(policy_cfg=config, dataset_stats=dataset.meta.stats)
+        # backbone model to extract features from the image frames
+        model_name="microsoft/resnet-18",
+    )
+
+    # Make policy, preprocessor, and optimizer
+    policy = make_policy(config, ds_meta=dataset.meta)
+    optimizer = config.get_optimizer_preset().build(policy.parameters())
+    preprocessor, _ = make_pre_post_processors(policy_cfg=config, dataset_stats=dataset.meta.stats)
+
+    classifier_id = "<user>/reward_classifier_hil_serl_example"
+
+    # Instantiate a dataloader
+    dataloader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True)
+
+    # Training loop
+    num_epochs = 5
+    for epoch in range(num_epochs):
+        total_loss = 0
+        total_accuracy = 0
+        for batch in dataloader:
+            # Preprocess the batch and move it to the correct device.
+            batch = preprocessor(batch)
+
+            # Forward pass
+            loss, output_dict = policy.forward(batch)
+
+            # Backward pass and optimization
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            total_loss += loss.item()
+            total_accuracy += output_dict["accuracy"]
+
+        avg_loss = total_loss / len(dataloader)
+        avg_accuracy = total_accuracy / len(dataloader)
+        print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}, Accuracy: {avg_accuracy:.2f}%")
+
+    print("Training finished!")
+
+    # You can now save the trained policy.
+    policy.push_to_hub(classifier_id)
 
 
-classifier_id = "fracapuano/reward_classifier_hil_serl_example"
+if __name__ == "__main__":
-
+    main()
-# Instantiate a dataloader
-dataloader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True)
-
-# Training loop
-num_epochs = 5
-for epoch in range(num_epochs):
-    total_loss = 0
-    total_accuracy = 0
-    for batch in dataloader:
-        # Preprocess the batch and move it to the correct device.
-        batch = preprocessor(batch)
-
-        # Forward pass
-        loss, output_dict = policy.forward(batch)
-
-        # Backward pass and optimization
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-
-        total_loss += loss.item()
-        total_accuracy += output_dict["accuracy"]
-
-    avg_loss = total_loss / len(dataloader)
-    avg_accuracy = total_accuracy / len(dataloader)
-    print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}, Accuracy: {avg_accuracy:.2f}%")
-
-print("Training finished!")
-
-# You can now save the trained policy.
-policy.push_to_hub(classifier_id)

examples/tutorial/smolvla/using_smolvla_example.py

@@ -11,56 +11,62 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower
 MAX_EPISODES = 5
 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)
-    model.config,
-    model_id,
-    # This overrides allows to run on MPS, otherwise defaults to CUDA (if available)
-    preprocessor_overrides={"device_processor": {"device": str(device)}},
-)
 
-# find ports using lerobot-find-port
+    preprocess, postprocess = make_pre_post_processors(
-follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
+        model.config,
+        model_id,
+        # This overrides allows to run on MPS, otherwise defaults to CUDA (if available)
+        preprocessor_overrides={"device_processor": {"device": str(device)}},
+    )
 
-# the robot ids are used the load the right calibration files
+    # find ports using lerobot-find-port
-follower_id = ...  # something like "follower_so100"
+    follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
 
-# Robot and environment configuration
+    # the robot ids are used the load the right calibration files
-# Camera keys must match the name and resolutions of the ones used for training!
+    follower_id = ...  # something like "follower_so100"
-# You can check the camera keys expected by a model in the info.json card on the model card on the Hub
-camera_config = {
-    "camera1": OpenCVCameraConfig(index_or_path=0, 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 and environment configuration
-robot = SO100Follower(robot_cfg)
+    # Camera keys must match the name and resolutions of the ones used for training!
-robot.connect()
+    # You can check the camera keys expected by a model in the info.json card on the model card on the Hub
+    camera_config = {
+        "camera1": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
+        "camera2": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
+    }
 
-task = ""  # something like "pick the red block"
+    robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
-robot_type = ""  # something like "so100_follower" for multi-embodiment datasets
+    robot = SO100Follower(robot_cfg)
+    robot.connect()
 
-# This is used to match the raw observation keys to the keys expected by the policy
+    task = ""  # something like "pick the red block"
-action_features = hw_to_dataset_features(robot.action_features, "action")
+    robot_type = ""  # something like "so100_follower" for multi-embodiment datasets
-obs_features = hw_to_dataset_features(robot.observation_features, "observation")
-dataset_features = {**action_features, **obs_features}
 
-for _ in range(MAX_EPISODES):
+    # This is used to match the raw observation keys to the keys expected by the policy
-    for _ in range(MAX_STEPS_PER_EPISODE):
+    action_features = hw_to_dataset_features(robot.action_features, "action")
-        obs = robot.get_observation()
+    obs_features = hw_to_dataset_features(robot.observation_features, "observation")
-        obs_frame = build_inference_frame(
+    dataset_features = {**action_features, **obs_features}
-            observation=obs, ds_features=dataset_features, device=device, task=task, robot_type=robot_type
-        )
 
-        obs = preprocess(obs_frame)
+    for _ in range(MAX_EPISODES):
+        for _ in range(MAX_STEPS_PER_EPISODE):
+            obs = robot.get_observation()
+            obs_frame = build_inference_frame(
+                observation=obs, ds_features=dataset_features, device=device, task=task, robot_type=robot_type
+            )
 
-        action = model.select_action(obs)
+            obs = preprocess(obs_frame)
-        action = postprocess(action)
-        action = make_robot_action(action, dataset_features)
-        robot.send_action(action)
 
-    print("Episode finished! Starting new episode...")
+            action = model.select_action(obs)
+            action = postprocess(action)
+            action = make_robot_action(action, dataset_features)
+            robot.send_action(action)
+
+        print("Episode finished! Starting new episode...")
+
+
+if __name__ == "__main__":
+    main()