fix(scripts): better prints teleop (#2538)

* fix(time benchmark): removing deprecated TimeBenchmark dependency

* fix(typo): renaming frames in an up-to-date fashion

* feat(duets): rearanging crf and g parameters in a proper unique combination manner

* fix(segfault): fixing segfault by adding a lock in ThreadPoolExecutor

* chore(update) : update datasets, codecs and backends to the latest versions

* chore(unused files): removing unused files

* fix(dataset paths): fix datasets paths to live among lerobot datasets

benchmarks/video/benchmark.py

@@ -1,94 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import threading
-import time
-from contextlib import ContextDecorator
-
-
-class TimeBenchmark(ContextDecorator):
-    """
-    Measures execution time using a context manager or decorator.
-
-    This class supports both context manager and decorator usage, and is thread-safe for multithreaded
-    environments.
-
-    Args:
-        print: If True, prints the elapsed time upon exiting the context or completing the function. Defaults
-        to False.
-
-    Examples:
-
-        Using as a context manager:
-
-        >>> benchmark = TimeBenchmark()
-        >>> with benchmark:
-        ...     time.sleep(1)
-        >>> print(f"Block took {benchmark.result:.4f} seconds")
-        Block took approximately 1.0000 seconds
-
-        Using with multithreading:
-
-        ```python
-        import threading
-
-        benchmark = TimeBenchmark()
-
-
-        def context_manager_example():
-            with benchmark:
-                time.sleep(0.01)
-            print(f"Block took {benchmark.result_ms:.2f} milliseconds")
-
-
-        threads = []
-        for _ in range(3):
-            t1 = threading.Thread(target=context_manager_example)
-            threads.append(t1)
-
-        for t in threads:
-            t.start()
-
-        for t in threads:
-            t.join()
-        ```
-        Expected output:
-        Block took approximately 10.00 milliseconds
-        Block took approximately 10.00 milliseconds
-        Block took approximately 10.00 milliseconds
-    """
-
-    def __init__(self, print=False):
-        self.local = threading.local()
-        self.print_time = print
-
-    def __enter__(self):
-        self.local.start_time = time.perf_counter()
-        return self
-
-    def __exit__(self, *exc):
-        self.local.end_time = time.perf_counter()
-        self.local.elapsed_time = self.local.end_time - self.local.start_time
-        if self.print_time:
-            print(f"Elapsed time: {self.local.elapsed_time:.4f} seconds")
-        return False
-
-    @property
-    def result(self):
-        return getattr(self.local, "elapsed_time", None)
-
-    @property
-    def result_ms(self):
-        return self.result * 1e3

benchmarks/video/capture_camera_feed.py

@@ -1,102 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Capture video feed from a camera as raw images."""
-
-import argparse
-import datetime as dt
-import os
-import time
-from pathlib import Path
-
-import cv2
-import rerun as rr
-
-# see https://rerun.io/docs/howto/visualization/limit-ram
-RERUN_MEMORY_LIMIT = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "5%")
-
-
-def display_and_save_video_stream(output_dir: Path, fps: int, width: int, height: int, duration: int):
-    rr.init("lerobot_capture_camera_feed")
-    rr.spawn(memory_limit=RERUN_MEMORY_LIMIT)
-
-    now = dt.datetime.now()
-    capture_dir = output_dir / f"{now:%Y-%m-%d}" / f"{now:%H-%M-%S}"
-    if not capture_dir.exists():
-        capture_dir.mkdir(parents=True, exist_ok=True)
-
-    # Opens the default webcam
-    cap = cv2.VideoCapture(0)
-    if not cap.isOpened():
-        print("Error: Could not open video stream.")
-        return
-
-    cap.set(cv2.CAP_PROP_FPS, fps)
-    cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
-    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
-
-    frame_index = 0
-    start_time = time.time()
-    while time.time() - start_time < duration:
-        ret, frame = cap.read()
-
-        if not ret:
-            print("Error: Could not read frame.")
-            break
-        rr.log("video/stream", rr.Image(frame), static=True)
-        cv2.imwrite(str(capture_dir / f"frame_{frame_index:06d}.png"), frame)
-        frame_index += 1
-
-    # Release the capture
-    cap.release()
-
-    # TODO(Steven): Add a graceful shutdown via a close() method for the Viewer context, though not currently supported in the Rerun API.
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--output-dir",
-        type=Path,
-        default=Path("outputs/cam_capture/"),
-        help="Directory where the capture images are written. A subfolder named with the current date & time will be created inside it for each capture.",
-    )
-    parser.add_argument(
-        "--fps",
-        type=int,
-        default=30,
-        help="Frames Per Second of the capture.",
-    )
-    parser.add_argument(
-        "--width",
-        type=int,
-        default=1280,
-        help="Width of the captured images.",
-    )
-    parser.add_argument(
-        "--height",
-        type=int,
-        default=720,
-        help="Height of the captured images.",
-    )
-    parser.add_argument(
-        "--duration",
-        type=int,
-        default=20,
-        help="Duration in seconds for which the video stream should be captured.",
-    )
-    args = parser.parse_args()
-    display_and_save_video_stream(**vars(args))

benchmarks/video/run_video_benchmark.py

@@ -21,11 +21,13 @@ See the provided README.md or run `python benchmark/video/run_video_benchmark.py
 
 import argparse
 import datetime as dt
+import itertools
 import random
 import shutil
 from collections import OrderedDict
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
+from threading import Lock
 
 import einops
 import numpy as np
@@ -35,13 +37,13 @@ import torch
 from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity
 from tqdm import tqdm
 
-from benchmarks.video.benchmark import TimeBenchmark
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.video_utils import (
-    decode_video_frames_torchvision,
+    decode_video_frames,
     encode_video_frames,
 )
 from lerobot.utils.constants import OBS_IMAGE
+from lerobot.utils.utils import TimerManager
 
 BASE_ENCODING = OrderedDict(
     [
@@ -86,7 +88,7 @@ def load_original_frames(imgs_dir: Path, timestamps: list[float], fps: int) -> t
     frames = []
     for ts in timestamps:
         idx = int(ts * fps)
-        frame = PIL.Image.open(imgs_dir / f"frame_{idx:06d}.png")
+        frame = PIL.Image.open(imgs_dir / f"frame-{idx:06d}.png")
         frame = torch.from_numpy(np.array(frame))
         frame = frame.type(torch.float32) / 255
         frame = einops.rearrange(frame, "h w c -> c h w")
@@ -97,21 +99,21 @@ def load_original_frames(imgs_dir: Path, timestamps: list[float], fps: int) -> t
 def save_decoded_frames(
     imgs_dir: Path, save_dir: Path, frames: torch.Tensor, timestamps: list[float], fps: int
 ) -> None:
-    if save_dir.exists() and len(list(save_dir.glob("frame_*.png"))) == len(timestamps):
+    if save_dir.exists() and len(list(save_dir.glob("frame-*.png"))) == len(timestamps):
         return
 
     save_dir.mkdir(parents=True, exist_ok=True)
     for i, ts in enumerate(timestamps):
         idx = int(ts * fps)
         frame_hwc = (frames[i].permute((1, 2, 0)) * 255).type(torch.uint8).cpu().numpy()
-        PIL.Image.fromarray(frame_hwc).save(save_dir / f"frame_{idx:06d}_decoded.png")
-        shutil.copyfile(imgs_dir / f"frame_{idx:06d}.png", save_dir / f"frame_{idx:06d}_original.png")
+        PIL.Image.fromarray(frame_hwc).save(save_dir / f"frame-{idx:06d}_decoded.png")
+        shutil.copyfile(imgs_dir / f"frame-{idx:06d}.png", save_dir / f"frame-{idx:06d}_original.png")
 
 
 def save_first_episode(imgs_dir: Path, dataset: LeRobotDataset) -> None:
     episode_index = 0
     ep_num_images = dataset.meta.episodes["length"][episode_index]
-    if imgs_dir.exists() and len(list(imgs_dir.glob("frame_*.png"))) == ep_num_images:
+    if imgs_dir.exists() and len(list(imgs_dir.glob("frame-*.png"))) == ep_num_images:
         return
 
     imgs_dir.mkdir(parents=True, exist_ok=True)
@@ -125,7 +127,7 @@ def save_first_episode(imgs_dir: Path, dataset: LeRobotDataset) -> None:
         tqdm(imgs_dataset, desc=f"saving {dataset.repo_id} first episode images", leave=False)
     ):
         img = item[img_keys[0]]
-        img.save(str(imgs_dir / f"frame_{i:06d}.png"), quality=100)
+        img.save(str(imgs_dir / f"frame-{i:06d}.png"), quality=100)
 
         if i >= ep_num_images - 1:
             break
@@ -149,18 +151,6 @@ def sample_timestamps(timestamps_mode: str, ep_num_images: int, fps: int) -> lis
     return [idx / fps for idx in frame_indexes]
 
 
-def decode_video_frames(
-    video_path: str,
-    timestamps: list[float],
-    tolerance_s: float,
-    backend: str,
-) -> torch.Tensor:
-    if backend in ["pyav", "video_reader"]:
-        return decode_video_frames_torchvision(video_path, timestamps, tolerance_s, backend)
-    else:
-        raise NotImplementedError(backend)
-
-
 def benchmark_decoding(
     imgs_dir: Path,
     video_path: Path,
@@ -172,8 +162,8 @@ def benchmark_decoding(
     num_workers: int = 4,
     save_frames: bool = False,
 ) -> dict:
-    def process_sample(sample: int):
-        time_benchmark = TimeBenchmark()
+    def process_sample(sample: int, lock: Lock):
+        time_benchmark = TimerManager(log=False)
         timestamps = sample_timestamps(timestamps_mode, ep_num_images, fps)
         num_frames = len(timestamps)
         result = {
@@ -182,13 +172,13 @@ def benchmark_decoding(
             "mse_values": [],
         }
 
-        with time_benchmark:
+        with time_benchmark, lock:
             frames = decode_video_frames(video_path, timestamps=timestamps, tolerance_s=5e-1, backend=backend)
-        result["load_time_video_ms"] = time_benchmark.result_ms / num_frames
+        result["load_time_video_ms"] = (time_benchmark.last * 1000) / num_frames
 
         with time_benchmark:
             original_frames = load_original_frames(imgs_dir, timestamps, fps)
-        result["load_time_images_ms"] = time_benchmark.result_ms / num_frames
+        result["load_time_images_ms"] = (time_benchmark.last * 1000) / num_frames
 
         frames_np, original_frames_np = frames.numpy(), original_frames.numpy()
         for i in range(num_frames):
@@ -215,8 +205,10 @@ def benchmark_decoding(
     # A sample is a single set of decoded frames specified by timestamps_mode (e.g. a single frame, 2 frames, etc.).
     # For each sample, we record metrics (loading time and quality metrics) which are then averaged over all samples.
     # As these samples are independent, we run them in parallel threads to speed up the benchmark.
+    # Use a single shared lock for all worker threads
+    shared_lock = Lock()
     with ThreadPoolExecutor(max_workers=num_workers) as executor:
-        futures = [executor.submit(process_sample, i) for i in range(num_samples)]
+        futures = [executor.submit(process_sample, i, shared_lock) for i in range(num_samples)]
         for future in tqdm(as_completed(futures), total=num_samples, desc="samples", leave=False):
             result = future.result()
             load_times_video_ms.append(result["load_time_video_ms"])
@@ -358,24 +350,27 @@ def main(
                 imgs_dir = output_dir / "images" / dataset.repo_id.replace("/", "_")
                 # We only use the first episode
                 save_first_episode(imgs_dir, dataset)
-                for key, values in tqdm(encoding_benchmarks.items(), desc="encodings (g, crf)", leave=False):
-                    for value in tqdm(values, desc=f"encodings ({key})", leave=False):
-                        encoding_cfg = BASE_ENCODING.copy()
-                        encoding_cfg["vcodec"] = video_codec
-                        encoding_cfg["pix_fmt"] = pixel_format
+                for duet in [
+                    dict(zip(encoding_benchmarks.keys(), unique_combination, strict=False))
+                    for unique_combination in itertools.product(*encoding_benchmarks.values())
+                ]:
+                    encoding_cfg = BASE_ENCODING.copy()
+                    encoding_cfg["vcodec"] = video_codec
+                    encoding_cfg["pix_fmt"] = pixel_format
+                    for key, value in duet.items():
                         encoding_cfg[key] = value
-                        args_path = Path("_".join(str(value) for value in encoding_cfg.values()))
-                        video_path = output_dir / "videos" / args_path / f"{repo_id.replace('/', '_')}.mp4"
-                        benchmark_table += benchmark_encoding_decoding(
-                            dataset,
-                            video_path,
-                            imgs_dir,
-                            encoding_cfg,
-                            decoding_benchmarks,
-                            num_samples,
-                            num_workers,
-                            save_frames,
-                        )
+                    args_path = Path("_".join(str(value) for value in encoding_cfg.values()))
+                    video_path = output_dir / "videos" / args_path / f"{repo_id.replace('/', '_')}.mp4"
+                    benchmark_table += benchmark_encoding_decoding(
+                        dataset,
+                        video_path,
+                        imgs_dir,
+                        encoding_cfg,
+                        decoding_benchmarks,
+                        num_samples,
+                        num_workers,
+                        save_frames,
+                    )
 
             # Save intermediate results
             benchmark_df = pd.DataFrame(benchmark_table, columns=headers)
@@ -409,9 +404,9 @@ if __name__ == "__main__":
         nargs="*",
         default=[
             "lerobot/pusht_image",
-            "aliberts/aloha_mobile_shrimp_image",
-            "aliberts/paris_street",
-            "aliberts/kitchen",
+            "lerobot/aloha_mobile_shrimp_image",
+            "lerobot/paris_street",
+            "lerobot/kitchen",
         ],
         help="Datasets repo-ids to test against. First episodes only are used. Must be images.",
     )
@@ -419,7 +414,7 @@ if __name__ == "__main__":
         "--vcodec",
         type=str,
         nargs="*",
-        default=["libx264", "hevc", "libsvtav1"],
+        default=["h264", "hevc", "libsvtav1"],
         help="Video codecs to be tested",
     )
     parser.add_argument(
@@ -468,7 +463,7 @@ if __name__ == "__main__":
         "--backends",
         type=str,
         nargs="*",
-        default=["pyav", "video_reader"],
+        default=["torchcodec", "pyav"],
         help="Torchvision decoding backend to be tested.",
     )
     parser.add_argument(

docs/source/_toctree.yml

@@ -47,8 +47,8 @@
 - sections:
   - local: envhub
     title: Environments from the Hub
-  - local: il_sim
-    title: Imitation Learning in Sim
+  - local: envhub_leisaac
+    title: Control & Train Robots in Sim (LeIsaac)
   - local: libero
     title: Using Libero
   - local: metaworld

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
 )

docs/source/envhub_leisaac.mdx

@@ -0,0 +1,301 @@
+# LeIsaac × LeRobot EnvHub
+
+LeRobot EnvHub now supports **imitation learning in simulation** with LeIsaac.
+Spin up everyday manipulation tasks, teleoperate the robot, collect demos, push them to the Hub, and train policies in LeRobot — all in one loop.
+
+[LeIsaac](https://github.com/LightwheelAI/leisaac) integrates with IsaacLab and the SO101 Leader/Follower setup to provide:
+
+- 🕹️ **Teleoperation-first workflows** for data collection
+- 📦 **Built-in data conversion** ready for LeRobot training
+- 🤖 **Everyday skills** like picking oranges, lifting cubes, cleaning tables, and folding cloth
+- ☁️ **Ongoing upgrades** from [LightWheel](https://lightwheel.ai/): cloud simulation, EnvHub support, Sim2Real tooling, and more
+
+Below you’ll find the currently supported LeIsaac tasks exposed through LeRobot EnvHub.
+
+# Available Environments
+
+The following table lists all available tasks and environments in LeIsaac x LeRobot Envhub. You can also get the latest list of environments by running the following command:
+
+```bash
+python scripts/environments/list_envs.py
+```
+
+| Task                                                                                                                                                            | Environment ID                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                | Task Description                                                                                                           | Related Robot                                              |
+| :-------------------------------------------------------------------------------------------------------------------------------------------------------------- | :-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | :------------------------------------------------------------------------------------------------------------------------- | :--------------------------------------------------------- |
+| <video src="https://github.com/user-attachments/assets/466eddff-f720-4f99-94d5-5e123e4c302c" autoplay loop muted playsinline style="max-width: 300px;"></video> | [LeIsaac-SO101-PickOrange-v0](https://github.com/LightwheelAI/leisaac/blob/main/source/leisaac/leisaac/tasks/pick_orange/pick_orange_env_cfg.py)<br /><br />[LeIsaac-SO101-PickOrange-Direct-v0](https://github.com/LightwheelAI/leisaac/blob/main/source/leisaac/leisaac/tasks/pick_orange/direct/pick_orange_env.py)                                                                                                                                                                                                                        | Pick three oranges and put them into the plate, then reset the arm to rest state.                                          | Single-Arm SO101 Follower                                  |
+| <video src="https://github.com/user-attachments/assets/1e4eb83a-0b38-40fb-a0b2-ddb0fe201e6d" autoplay loop muted playsinline style="max-width: 300px;"></video> | [LeIsaac-SO101-LiftCube-v0](https://github.com/LightwheelAI/leisaac/blob/main/source/leisaac/leisaac/tasks/lift_cube/lift_cube_env_cfg.py)<br /><br />[LeIsaac-SO101-LiftCube-Direct-v0](https://github.com/LightwheelAI/leisaac/blob/main/source/leisaac/leisaac/tasks/lift_cube/direct/lift_cube_env.py)                                                                                                                                                                                                                                    | Lift the red cube up.                                                                                                      | Single-Arm SO101 Follower                                  |
+| <video src="https://github.com/user-attachments/assets/e49d8f1c-dcc9-412b-a88f-100680d8a45b" autoplay loop muted playsinline style="max-width: 300px;"></video> | [LeIsaac-SO101-CleanToyTable-v0](https://github.com/LightwheelAI/leisaac/blob/main/source/leisaac/leisaac/tasks/clean_toy_table/clean_toy_table_env_cfg.py)<br /><br />[LeIsaac-SO101-CleanToyTable-BiArm-v0](https://github.com/LightwheelAI/leisaac/blob/main/source/leisaac/leisaac/tasks/clean_toy_table/clean_toy_table_bi_arm_env_cfg.py)<br /><br />[LeIsaac-SO101-CleanToyTable-BiArm-Direct-v0](https://github.com/LightwheelAI/leisaac/blob/main/source/leisaac/leisaac/tasks/clean_toy_table/direct/clean_toy_table_bi_arm_env.py) | Pick two letter e objects into the box, and reset the arm to rest state.                                                   | Single-Arm SO101 Follower<br /><br />Bi-Arm SO101 Follower |
+| <video src="https://github.com/user-attachments/assets/e29a0f8a-9286-4ce6-b45d-342c3d3ba754" autoplay loop muted playsinline style="max-width: 300px;"></video> | [LeIsaac-SO101-FoldCloth-BiArm-v0](https://github.com/LightwheelAI/leisaac/blob/main/source/leisaac/leisaac/tasks/fold_cloth/fold_cloth_bi_arm_env_cfg.py)<br /><br />[LeIsaac-SO101-FoldCloth-BiArm-Direct-v0](https://github.com/LightwheelAI/leisaac/blob/main/source/leisaac/leisaac/tasks/fold_cloth/direct/fold_cloth_bi_arm_env.py)                                                                                                                                                                                                    | Fold the cloth, and reset the arm to rest state.<br /><br />_Note: Only the DirectEnv support check_success in this task._ | Bi-Arm SO101 Follower                                      |
+
+# Load LeIsaac directly in LeRobot with one line of code
+
+> EnvHub: Share LeIsaac environments through HuggingFace
+
+[EnvHub](https://huggingface.co/docs/lerobot/envhub) is our reproducible environment hub, spin up a packaged simulation with one line, experiment immediately, and publish your own tasks for the community.
+
+LeIsaac offers EnvHub support so you can consume or share tasks with only a few commands.
+
+<video
+  controls
+  src="https://github.com/user-attachments/assets/687666f5-ebe0-421d-84a0-eb86116ac5f8"
+  style={{ width: "100%", maxWidth: "960px", borderRadius: "8px" }}
+/>
+
+## How to get started, environment Setup
+
+Run the following commands to setup your code environments:
+
+```bash
+# Refer to Getting Started/Installation to install leisaac firstly
+conda create -n leisaac_envhub python=3.11
+conda activate leisaac_envhub
+
+conda install -c "nvidia/label/cuda-12.8.1" cuda-toolkit
+pip install -U torch==2.7.0 torchvision==0.22.0 --index-url https://download.pytorch.org/whl/cu128
+pip install 'leisaac[isaaclab] @ git+https://github.com/LightwheelAI/leisaac.git#subdirectory=source/leisaac' --extra-index-url https://pypi.nvidia.com
+
+# Install lerobot
+pip install lerobot==0.4.1
+
+# Fix numpy version
+pip install numpy==1.26.0
+```
+
+## Usage Example
+
+EnvHub exposes every LeIsaac-supported task in a uniform interface. The examples below load `so101_pick_orange` and demonstrate a random-action rollout and an interactive teleoperation.
+
+### Random Action
+
+<details>
+<summary>Click to expand code example</summary>
+
+```python
+# envhub_random_action.py
+
+import torch
+from lerobot.envs.factory import make_env
+
+# Load from the hub
+envs_dict = make_env("LightwheelAI/leisaac_env:envs/so101_pick_orange.py", n_envs=1, trust_remote_code=True)
+
+# Access the environment
+suite_name = next(iter(envs_dict))
+sync_vector_env = envs_dict[suite_name][0]
+# retrieve the isaac environment from the sync vector env
+env = sync_vector_env.envs[0].unwrapped
+
+# Use it like any gym environment
+obs, info = env.reset()
+
+while True:
+    action = torch.tensor(env.action_space.sample())
+    obs, reward, terminated, truncated, info = env.step(action)
+    if terminated or truncated:
+        obs, info = env.reset()
+
+env.close()
+```
+
+</details>
+
+```bash
+python envhub_random_action.py
+```
+
+You should see the SO101 arm swinging under purely random commands.
+
+### Teleoperation
+
+LeRobot’s teleoperation stack can drive the simulated arm.
+
+Connect the SO101 Leader controller, run the calibration command below.
+
+```bash
+lerobot-calibrate \
+    --teleop.type=so101_leader \
+    --teleop.port=/dev/ttyACM0 \
+    --teleop.id=leader
+```
+
+And then launch the teleop script.
+
+<details>
+<summary>Click to expand code example</summary>
+
+```python
+# envhub_teleop_example.py
+
+import logging
+import time
+import gymnasium as gym
+
+from dataclasses import asdict, dataclass
+from pprint import pformat
+
+from lerobot.teleoperators import (  # noqa: F401
+    Teleoperator,
+    TeleoperatorConfig,
+    make_teleoperator_from_config,
+    so101_leader,
+)
+from lerobot.utils.robot_utils import precise_sleep
+from lerobot.utils.utils import init_logging
+from lerobot.envs.factory import make_env
+
+
+@dataclass
+class TeleoperateConfig:
+    teleop: TeleoperatorConfig
+    env_name: str = "so101_pick_orange"
+    fps: int = 60
+
+
+@dataclass
+class EnvWrap:
+    env: gym.Env
+
+
+def make_env_from_leisaac(env_name: str = "so101_pick_orange"):
+    envs_dict = make_env(
+        f'LightwheelAI/leisaac_env:envs/{env_name}.py',
+        n_envs=1,
+        trust_remote_code=True
+    )
+    suite_name = next(iter(envs_dict))
+    sync_vector_env = envs_dict[suite_name][0]
+    env = sync_vector_env.envs[0].unwrapped
+
+    return env
+
+
+def teleop_loop(teleop: Teleoperator, env: gym.Env, fps: int):
+    from leisaac.devices.action_process import preprocess_device_action
+    from leisaac.assets.robots.lerobot import SO101_FOLLOWER_MOTOR_LIMITS
+    from leisaac.utils.env_utils import dynamic_reset_gripper_effort_limit_sim
+
+    env_wrap = EnvWrap(env=env)
+
+    obs, info = env.reset()
+    while True:
+        loop_start = time.perf_counter()
+        if env.cfg.dynamic_reset_gripper_effort_limit:
+            dynamic_reset_gripper_effort_limit_sim(env, 'so101leader')
+
+        raw_action = teleop.get_action()
+        processed_action = preprocess_device_action(
+            dict(
+                so101_leader=True,
+                joint_state={
+                    k.removesuffix(".pos"): v for k, v in raw_action.items()},
+                motor_limits=SO101_FOLLOWER_MOTOR_LIMITS),
+            env_wrap
+        )
+        obs, reward, terminated, truncated, info = env.step(processed_action)
+        if terminated or truncated:
+            obs, info = env.reset()
+
+        dt_s = time.perf_counter() - loop_start
+        precise_sleep(1 / fps - dt_s)
+        loop_s = time.perf_counter() - loop_start
+        print(f"\ntime: {loop_s * 1e3:.2f}ms ({1 / loop_s:.0f} Hz)")
+
+
+def teleoperate(cfg: TeleoperateConfig):
+    init_logging()
+    logging.info(pformat(asdict(cfg)))
+
+    teleop = make_teleoperator_from_config(cfg.teleop)
+    env = make_env_from_leisaac(cfg.env_name)
+
+    teleop.connect()
+    if hasattr(env, 'initialize'):
+        env.initialize()
+    try:
+        teleop_loop(teleop=teleop, env=env, fps=cfg.fps)
+    except KeyboardInterrupt:
+        pass
+    finally:
+        teleop.disconnect()
+        env.close()
+
+
+def main():
+    teleoperate(TeleoperateConfig(
+        teleop=so101_leader.SO101LeaderConfig(
+            port="/dev/ttyACM0",
+            id='leader',
+            use_degrees=False,
+        ),
+        env_name="so101_pick_orange",
+        fps=60,
+    ))
+
+
+if __name__ == "__main__":
+    main()
+
+```
+
+</details>
+
+```bash
+python envhub_teleop_example.py
+```
+
+Running the script lets you operate the simulated arm using the physical Leader device.
+
+## ☁️ Cloud Simulation (No GPU Required)
+
+Don’t have a local GPU or the right drivers? No problem! You can run LeIsaac entirely in the cloud with zero setup.
+LeIsaac works out-of-the-box on **NVIDIA Brev**, giving you a fully configured environment directly in your browser.
+
+👉 **Start here:** [https://lightwheelai.github.io/leisaac/docs/cloud_simulation/nvidia_brev](https://lightwheelai.github.io/leisaac/docs/cloud_simulation/nvidia_brev)
+
+Once your instance is deployed, simply open the link for **port 80 (HTTP)** to launch **Visual Studio Code Server** (default password: `password`). From there, you can run simulations, edit code, and visualize IsaacLab environments — all from your web browser.
+
+**No GPU, no drivers, no local installation. Just click and run.**
+
+## Additional Notes
+
+We keep EnvHub coverage aligned with the LeIsaac task. Currently supported:
+
+- `so101_pick_orange`
+- `so101_lift_cube`
+- `so101_clean_toytable`
+- `bi_so101_fold_cloth`
+
+Switch tasks by targeting a different script when calling `make_env`, for example:
+
+```python
+envs_dict_pick_orange = make_env("LightwheelAI/leisaac_env:envs/so101_pick_orange.py", n_envs=1, trust_remote_code=True)
+envs_dict_lift_cube = make_env("LightwheelAI/leisaac_env:envs/so101_lift_cube.py", n_envs=1, trust_remote_code=True)
+envs_dict_clean_toytable = make_env("LightwheelAI/leisaac_env:envs/so101_clean_toytable.py", n_envs=1, trust_remote_code=True)
+envs_dict_fold_cloth = make_env("LightwheelAI/leisaac_env:envs/bi_so101_fold_cloth.py", n_envs=1, trust_remote_code=True)
+```
+
+Note: when working with `bi_so101_fold_cloth`, call `initialize()` immediately after retrieving the env before performing any other operations:
+
+<details>
+<summary>Click to expand code example</summary>
+
+```python
+import torch
+from lerobot.envs.factory import make_env
+
+# Load from the hub
+envs_dict = make_env("LightwheelAI/leisaac_env:envs/bi_so101_fold_cloth.py", n_envs=1, trust_remote_code=True)
+
+# Access the environment
+suite_name = next(iter(envs_dict))
+sync_vector_env = envs_dict[suite_name][0]
+# retrieve the isaac environment from the sync vector env
+env = sync_vector_env.envs[0].unwrapped
+
+# NOTE: initialize() first
+env.initialize()
+
+# other operation with env...
+```
+
+</details>

docs/source/il_robots.mdx

@@ -393,7 +393,7 @@ import time
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say
 
 episode_idx = 0
@@ -415,7 +415,7 @@ for idx in range(dataset.num_frames):
     }
     robot.send_action(action)
 
-    busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))
+    precise_sleep(1.0 / dataset.fps - (time.perf_counter() - t0))
 
 robot.disconnect()
 ```

docs/source/il_sim.mdx

@@ -1,220 +0,0 @@
-# Imitation Learning in Sim
-
-This tutorial will explain how to train a neural network to control a robot in simulation with imitation learning.
-
-**You'll learn:**
-
-1. How to record a dataset in simulation with [gym-hil](https://github.com/huggingface/gym-hil) and visualize the dataset.
-2. How to train a policy using your data.
-3. How to evaluate your policy in simulation and visualize the results.
-
-For the simulation environment we use the same [repo](https://github.com/huggingface/gym-hil) that is also being used by the Human-In-the-Loop (HIL) reinforcement learning algorithm.
-This environment is based on [MuJoCo](https://mujoco.org) and allows you to record datasets in LeRobotDataset format.
-Teleoperation is easiest with a controller like the Logitech F710, but you can also use your keyboard if you are up for the challenge.
-
-## Installation
-
-First, install the `gym_hil` package within the LeRobot environment, go to your LeRobot folder and run this command:
-
-```bash
-pip install -e ".[hilserl]"
-```
-
-## Teleoperate and Record a Dataset
-
-To use `gym_hil` with LeRobot, you need to use a configuration file. An example config file can be found [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/sim_il/env_config.json).
-
-To teleoperate and collect a dataset, we need to modify this config file. Here's an example configuration for imitation learning data collection:
-
-```json
-{
-  "env": {
-    "type": "gym_manipulator",
-    "name": "gym_hil",
-    "task": "PandaPickCubeGamepad-v0",
-    "fps": 10
-  },
-  "dataset": {
-    "repo_id": "your_username/il_gym",
-    "root": null,
-    "task": "pick_cube",
-    "num_episodes_to_record": 30,
-    "replay_episode": null,
-    "push_to_hub": true
-  },
-  "mode": "record",
-  "device": "cuda"
-}
-```
-
-Key configuration points:
-
-- Set your `repo_id` in the `dataset` section: `"repo_id": "your_username/il_gym"`
-- Set `num_episodes_to_record: 30` to collect 30 demonstration episodes
-- Ensure `mode` is set to `"record"`
-- If you don't have an NVIDIA GPU, change `"device": "cuda"` to `"mps"` for macOS or `"cpu"`
-- To use keyboard instead of gamepad, change `"task"` to `"PandaPickCubeKeyboard-v0"`
-
-Then we can run this command to start:
-
-<hfoptions id="teleop_sim">
-<hfoption id="Linux">
-
-```bash
-python -m lerobot.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
-```
-
-</hfoption>
-<hfoption id="MacOS">
-
-```bash
-mjpython -m lerobot.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
-```
-
-</hfoption>
-</hfoptions>
-
-Once rendered you can teleoperate the robot with the gamepad or keyboard, below you can find the gamepad/keyboard controls.
-
-Note that to teleoperate the robot you have to hold the "Human Take Over Pause Policy" Button `RB` to enable control!
-
-**Gamepad Controls**
-
-<p align="center">
-  <img
-    src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/gamepad_guide.jpg?raw=true"
-    alt="Figure shows the control mappings on a Logitech gamepad."
-    title="Gamepad Control Mapping"
-    width="100%"
-  ></img>
-</p>
-<p align="center">
-  <i>Gamepad button mapping for robot control and episode management</i>
-</p>
-
-**Keyboard controls**
-
-For keyboard controls use the `spacebar` to enable control and the following keys to move the robot:
-
-```bash
-  Arrow keys: Move in X-Y plane
-  Shift and Shift_R: Move in Z axis
-  Right Ctrl and Left Ctrl: Open and close gripper
-  ESC: Exit
-```
-
-## Visualize a dataset
-
-If you uploaded your dataset to the hub you can [visualize your dataset online](https://huggingface.co/spaces/lerobot/visualize_dataset) by copy pasting your repo id.
-
-<p align="center">
-  <img
-    src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/dataset_visualizer_sim.png"
-    alt="Figure shows the dataset visualizer"
-    title="Dataset visualization"
-    width="100%"
-  ></img>
-</p>
-<p align="center">
-  <i>Dataset visualizer</i>
-</p>
-
-## Train a policy
-
-To train a policy to control your robot, use the [`lerobot-train`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
-
-```bash
-lerobot-train \
-  --dataset.repo_id=${HF_USER}/il_gym \
-  --policy.type=act \
-  --output_dir=outputs/train/il_sim_test \
-  --job_name=il_sim_test \
-  --policy.device=cuda \
-  --wandb.enable=true
-```
-
-Let's explain the command:
-
-1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/il_gym`.
-2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor states, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
-3. We provided `policy.device=cuda` since we are training on a Nvidia GPU, but you could use `policy.device=mps` to train on Apple silicon.
-4. We provided `wandb.enable=true` to use [Weights and Biases](https://docs.wandb.ai/quickstart) for visualizing training plots. This is optional but if you use it, make sure you are logged in by running `wandb login`.
-
-Training should take several hours, 100k steps (which is the default) will take about 1h on Nvidia A100. You will find checkpoints in `outputs/train/il_sim_test/checkpoints`.
-
-#### Train using Collab
-
-If your local computer doesn't have a powerful GPU you could utilize Google Collab to train your model by following the [ACT training notebook](./notebooks#training-act).
-
-#### Upload policy checkpoints
-
-Once training is done, upload the latest checkpoint with:
-
-```bash
-huggingface-cli upload ${HF_USER}/il_sim_test \
-  outputs/train/il_sim_test/checkpoints/last/pretrained_model
-```
-
-You can also upload intermediate checkpoints with:
-
-```bash
-CKPT=010000
-huggingface-cli upload ${HF_USER}/il_sim_test${CKPT} \
-  outputs/train/il_sim_test/checkpoints/${CKPT}/pretrained_model
-```
-
-## Evaluate your policy in Sim
-
-To evaluate your policy we have to use a configuration file. An example can be found [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/sim_il/eval_config.json).
-
-Here's an example evaluation configuration:
-
-```json
-{
-  "env": {
-    "type": "gym_manipulator",
-    "name": "gym_hil",
-    "task": "PandaPickCubeGamepad-v0",
-    "fps": 10
-  },
-  "dataset": {
-    "repo_id": "your_username/il_sim_dataset",
-    "dataset_root": null,
-    "task": "pick_cube"
-  },
-  "pretrained_policy_name_or_path": "your_username/il_sim_model",
-  "device": "cuda"
-}
-```
-
-Make sure to replace:
-
-- `repo_id` with the dataset you trained on (e.g., `your_username/il_sim_dataset`)
-- `pretrained_policy_name_or_path` with your model ID (e.g., `your_username/il_sim_model`)
-
-Then you can run this command to visualize your trained policy
-
-<hfoptions id="eval_policy">
-<hfoption id="Linux">
-
-```bash
-python -m lerobot.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
-```
-
-</hfoption>
-<hfoption id="MacOS">
-
-```bash
-mjpython -m lerobot.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
-```
-
-</hfoption>
-</hfoptions>
-
-> [!WARNING]
-> While the main workflow of training ACT in simulation is straightforward, there is significant room for exploring how to set up the task, define the initial state of the environment, and determine the type of data required during collection to learn the most effective policy. If your trained policy doesn't perform well, investigate the quality of the dataset it was trained on using our visualizers, as well as the action values and various hyperparameters related to ACT and the simulation.
-
-Congrats 🎉, you have finished this tutorial. If you want to continue with using LeRobot in simulation follow this [Tutorial on reinforcement learning in sim with HIL-SERL](https://huggingface.co/docs/lerobot/hilserl_sim)
-
-> [!TIP]
-> If you have any questions or need help, please reach out on [Discord](https://discord.com/invite/s3KuuzsPFb).

examples/backward_compatibility/replay.py

@@ -45,7 +45,7 @@ from lerobot.robots import (  # noqa: F401
     so101_follower,
 )
 from lerobot.utils.constants import ACTION
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import (
     init_logging,
     log_say,
@@ -97,7 +97,7 @@ def replay(cfg: ReplayConfig):
         robot.send_action(action)
 
         dt_s = time.perf_counter() - start_episode_t
-        busy_wait(1 / dataset.fps - dt_s)
+        precise_sleep(1 / dataset.fps - dt_s)
 
     robot.disconnect()
 

examples/behavior_1k/aggregate.slurm

@@ -1,29 +0,0 @@
-#!/bin/bash
-#SBATCH -J b1k-aggregate
-#SBATCH -p hopper-cpu
-#SBATCH --qos=high
-#SBATCH -c 2
-#SBATCH -t 20:00:00
-#SBATCH --mem=4G
-#SBATCH -D /admin/home/francesco_capuano/lerobot
-#SBATCH -o /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%j.out
-#SBATCH -e /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%j.err
-
-set -euo pipefail
-set -x
-export PYTHONUNBUFFERED=1
-export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK:-1}
-
-source "$HOME/.bashrc" 2>/dev/null || true
-if ! command -v conda >/dev/null 2>&1; then
-  source "$HOME/miniconda3/etc/profile.d/conda.sh" 2>/dev/null || true
-  source "$HOME/anaconda3/etc/profile.d/conda.sh" 2>/dev/null || true
-fi
-conda activate lerobot
-
-python examples/behavior_1k/aggregate_tasks_datasets.py \
-  --task-datasets-dir /fsx/francesco_capuano/behavior1k-v3 \
-  --aggregated-root /fsx/francesco_capuano/behavior1k-v3/behavior1k \
-  --num-tasks 50 \
-  --hf-user fracapuano \
-  --push-to-hub

examples/behavior_1k/aggregate_tasks_datasets.py

@@ -1,100 +0,0 @@
-"""Aggregate multiple task-specific LeRobot datasets into a single combined dataset."""
-
-import argparse
-import os
-from pathlib import Path
-
-from lerobot.datasets.aggregate import aggregate_datasets
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Aggregate multiple task-specific datasets into a single LeRobot dataset"
-    )
-    parser.add_argument(
-        "--task-datasets-dir",
-        type=str,
-        required=True,
-        help="Directory containing individual task datasets (e.g., /path/to/behavior1k/)",
-    )
-    parser.add_argument(
-        "--aggregated-root",
-        type=str,
-        required=True,
-        help="Path where the aggregated dataset will be written",
-    )
-    parser.add_argument(
-        "--num-tasks",
-        type=int,
-        default=50,
-        help="Number of tasks to aggregate (default: 50)",
-    )
-    parser.add_argument(
-        "--task-start-idx",
-        type=int,
-        default=0,
-        help="Starting task index (default: 0)",
-    )
-    parser.add_argument(
-        "--hf-user",
-        type=str,
-        default=None,
-        help="HuggingFace username for repo IDs (defaults to HF_USER env var or 'lerobot')",
-    )
-    parser.add_argument(
-        "--aggregated-repo-id",
-        type=str,
-        default=None,
-        help="Repository ID for the aggregated dataset (defaults to {hf_user}/behavior1k)",
-    )
-    parser.add_argument(
-        "--push-to-hub",
-        action="store_true",
-        help="Push the aggregated dataset to the Hugging Face Hub",
-    )
-
-    args = parser.parse_args()
-
-    # Determine HF user
-    hf_user = args.hf_user or os.environ.get("HF_USER", "lerobot")
-
-    # Set default aggregated repo ID if not provided
-    aggregated_repo_id = args.aggregated_repo_id or f"{hf_user}/behavior1k"
-
-    # Generate task indices
-    task_indices = range(args.task_start_idx, args.task_start_idx + args.num_tasks)
-
-    # Generate repo IDs for individual tasks
-    repo_ids = [f"{hf_user}/behavior1k-task{i:04d}" for i in task_indices]
-
-    # Generate local paths for individual task datasets
-    task_datasets_dir = Path(args.task_datasets_dir)
-    roots = [task_datasets_dir / f"behavior1k-task{i:04d}" for i in task_indices]
-
-    # Aggregated dataset path
-    aggregated_root = Path(args.aggregated_root)
-
-    print(f"🔹 Aggregating {args.num_tasks} task datasets")
-    print(f"Task datasets directory: {task_datasets_dir}")
-    print(f"Aggregated output: {aggregated_root}")
-    print(f"Aggregated repo ID: {aggregated_repo_id}")
-
-    aggregate_datasets(
-        repo_ids=repo_ids,
-        roots=roots,
-        aggr_repo_id=aggregated_repo_id,
-        aggr_root=aggregated_root,
-    )
-
-    print("✅ Aggregation complete")
-
-    if args.push_to_hub:
-        print(f"📤 Pushing aggregated dataset to {aggregated_repo_id}")
-        ds = LeRobotDataset(repo_id=aggregated_repo_id, root=aggregated_root)
-        ds.push_to_hub()
-        print("✅ Successfully pushed to hub")
-
-
-if __name__ == "__main__":
-    main()

examples/behavior_1k/convert.slurm

@@ -1,38 +0,0 @@
-#!/bin/bash
-#SBATCH -J b1k-convert
-#SBATCH -p hopper-cpu            # pick your partition
-#SBATCH --qos=high
-#SBATCH --array=0-49%8           # 50 tasks, max 8 running concurrently (conversion is I/O bound)
-#SBATCH -c 1                     # CPUs per conversion (tune as needed)
-#SBATCH -t 2:00:00               # Time per conversion
-#SBATCH --mem=3G                 # ~1.75GB for task 0, ~doubled for safety
-#SBATCH -D /admin/home/francesco_capuano/lerobot
-#SBATCH -o /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%A_%a.out
-#SBATCH -e /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%A_%a.err
-
-set -euo pipefail
-set -x
-export PYTHONUNBUFFERED=1
-export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK:-1}  # avoid BLAS oversubscription
-
-DATA_PATH="/fsx/francesco_capuano/behavior1k-2025-v21"
-BASE_OUT="/fsx/francesco_capuano/behavior1k-v3"
-mkdir -p "$BASE_OUT" logs
-
-i="${SLURM_ARRAY_TASK_ID}"
-OUT_DIR="$(printf "%s/behavior1k-task%04d" "$BASE_OUT" "$i")"
-
-# activate your env if needed
-source "$HOME/.bashrc" 2>/dev/null || true
-if ! command -v conda >/dev/null 2>&1; then
-  source "$HOME/miniconda3/etc/profile.d/conda.sh" 2>/dev/null || true
-  source "$HOME/anaconda3/etc/profile.d/conda.sh" 2>/dev/null || true
-fi
-conda activate lerobot
-
-python examples/behavior_1k/convert_to_lerobot_v3.py \
-  --data-path "$DATA_PATH" \
-  --new-repo "$OUT_DIR" \
-  --task-id "$i" \
-  --force-conversion \
-  --push-to-hub

examples/behavior_1k/convert_to_lerobot_v3.py

@@ -1,667 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert Behavior Dataset to LeRobotDataset v3.0 format"""
-
-import argparse
-import json
-import logging
-import os
-import shutil
-from pathlib import Path
-
-import jsonlines
-import pandas as pd
-import pyarrow as pa
-import tqdm
-from datasets import Dataset, Features, Image
-
-from lerobot.datasets.compute_stats import aggregate_stats
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.utils import (
-    DEFAULT_CHUNK_SIZE,
-    DEFAULT_DATA_FILE_SIZE_IN_MB,
-    DEFAULT_DATA_PATH,
-    DEFAULT_FEATURES,
-    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
-    DEFAULT_VIDEO_PATH,
-    LEGACY_EPISODES_PATH,
-    LEGACY_EPISODES_STATS_PATH,
-    LEGACY_TASKS_PATH,
-    cast_stats_to_numpy,
-    flatten_dict,
-    get_file_size_in_mb,
-    get_parquet_file_size_in_mb,
-    get_parquet_num_frames,
-    load_info,
-    update_chunk_file_indices,
-    write_episodes,
-    write_info,
-    write_stats,
-    write_tasks,
-)
-from lerobot.datasets.video_utils import concatenate_video_files, get_video_duration_in_s
-from lerobot.utils.utils import init_logging
-
-# script to convert one single task to v3.1
-# TASK = 1
-NEW_ROOT = Path("/fsx/jade_choghari/tmp/bb")
-
-
-def fix_episode_dataframe(df: pd.DataFrame) -> pd.DataFrame:
-    """Performs several fixes to an underlying dataframe to make it LeRobotDataset-v3 compatible"""
-    # Inject per-episode frame_index if missing (0..N-1 within each episode)
-    if "frame_index" not in df.columns:
-        df["frame_index"] = range(len(df))
-
-    # Remove variable-length task_info feature (NOTE(fracapuano): change to padding at some point?)
-    if "observation.task_info" in df.columns:
-        df = df.drop(columns=["observation.task_info"])
-
-    # NOTE(fracapuano): tasks are ordered (and there is one task per file/dataset)
-    if "task_index" in df.columns:
-        df["task_index"] = 0
-
-    return df
-
-
-def get_total_episodes_task(local_dir: Path, task_id: int, task_ranges: dict, step) -> int:
-    """
-    Calculates the total number of episodes for a single, specified task.
-    """
-    # Simply load the episodes for the task and count them.
-    episodes = legacy_load_episodes_task(
-        local_dir=local_dir, task_id=task_id, task_ranges=task_ranges, step=step
-    )
-    return len(episodes)
-
-
-NUM_CAMERAS = 9
-
-
-def get_total_frames_task(local_dir, meta_path, task_id: int, task_ranges: dict, step: int) -> int:
-    episodes_metadata = legacy_load_episodes_task(
-        local_dir=local_dir, task_id=task_id, task_ranges=task_ranges, step=step
-    )
-    total_frames = 0
-    # like 'duration'
-    for ep in episodes_metadata.values():
-        duration_s = ep["length"]
-        total_frames += int(duration_s)
-    return total_frames
-
-
-def convert_info(
-    root, new_root, data_file_size_in_mb, video_file_size_in_mb, meta_path, task_id: int, task_ranges, step
-):
-    info = load_info(root)
-    features = {**info["features"], **DEFAULT_FEATURES}
-    del features[
-        "observation.task_info"
-    ]  # variable-length task_info is not supported in LeRobotDataset v3.0!
-
-    info["codebase_version"] = "v3.0"
-    info["features"] = features
-    del info["total_videos"]
-    info["data_files_size_in_mb"] = data_file_size_in_mb
-    info["video_files_size_in_mb"] = video_file_size_in_mb
-    info["data_path"] = DEFAULT_DATA_PATH
-    info["video_path"] = DEFAULT_VIDEO_PATH if info["video_path"] is not None else None
-    info["fps"] = int(info["fps"])
-    for key in info["features"]:
-        if info["features"][key]["dtype"] == "video":
-            # already has fps in video_info
-            continue
-        info["features"][key]["fps"] = info["fps"]
-
-    info["total_episodes"] = get_total_episodes_task(root, task_id, task_ranges, step)
-    info["total_videos"] = info["total_episodes"] * NUM_CAMERAS
-    info["total_frames"] = get_total_frames_task(root, meta_path, task_id, task_ranges, step)
-    info["total_tasks"] = 1
-    write_info(info, new_root)
-
-
-def load_jsonlines(fpath: Path) -> list[any]:
-    with jsonlines.open(fpath, "r") as reader:
-        return list(reader)
-
-
-def legacy_load_tasks(local_dir: Path) -> tuple[dict, dict]:
-    tasks = load_jsonlines(local_dir / LEGACY_TASKS_PATH)
-    # return tasks dict such that
-    tasks = {item["task_index"]: item["task"] for item in sorted(tasks, key=lambda x: x["task_index"])}
-    task_to_task_index = {task: task_index for task_index, task in tasks.items()}
-    return tasks, task_to_task_index
-
-
-def convert_tasks(root, new_root, task_id: int):
-    tasks, _ = legacy_load_tasks(root)
-    if task_id not in tasks:
-        raise ValueError(f"Task ID {task_id} not found in tasks (available: {list(tasks.keys())})")
-    tasks = {task_id: tasks[task_id]}
-    # Tasks are ordered with 0..ntasks-1 in the converted dataset
-    task_indices = range(len(tasks.keys()))
-    task_strings = tasks.values()
-    df_tasks = pd.DataFrame({"task_index": task_indices}, index=task_strings)
-    write_tasks(df_tasks, new_root)
-
-
-def concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys):
-    # TODO(rcadene): to save RAM use Dataset.from_parquet(file) and concatenate_datasets
-    dataframes = []
-    for file in paths_to_cat:
-        df = pd.read_parquet(file)
-        df = fix_episode_dataframe(df)
-        dataframes.append(df)
-
-    # Concatenate all DataFrames along rows
-    concatenated_df = pd.concat(dataframes, ignore_index=True)
-
-    path = new_root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    if len(image_keys) > 0:
-        schema = pa.Schema.from_pandas(concatenated_df)
-        features = Features.from_arrow_schema(schema)
-        for key in image_keys:
-            features[key] = Image()
-        schema = features.arrow_schema
-    else:
-        schema = None
-
-    concatenated_df.to_parquet(path, index=False, schema=schema)
-
-
-def get_image_keys(root):
-    info = load_info(root)
-    features = info["features"]
-    image_keys = [key for key, ft in features.items() if ft["dtype"] == "image"]
-    return image_keys
-
-
-def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int, task_index: int):
-    task_dir_name = f"task-{task_index:04d}"
-    data_dir = root / "data" / task_dir_name
-    ep_paths = sorted(data_dir.glob("*.parquet"))
-    image_keys = get_image_keys(root)
-
-    ep_idx = 0
-    chunk_idx = 0
-    file_idx = 0
-    size_in_mb = 0
-    num_frames = 0
-    paths_to_cat = []
-    episodes_metadata = []
-
-    logging.info(f"Converting data files from {len(ep_paths)} episodes")
-
-    for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"):
-        ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
-        ep_num_frames = get_parquet_num_frames(ep_path)
-        ep_metadata = {
-            "episode_index": ep_idx,
-            "data/chunk_index": chunk_idx,
-            "data/file_index": file_idx,
-            "dataset_from_index": num_frames,
-            "dataset_to_index": num_frames + ep_num_frames,
-        }
-        size_in_mb += ep_size_in_mb
-        num_frames += ep_num_frames
-        episodes_metadata.append(ep_metadata)
-
-        # write 0-based episode index instead of custom episode index (otherwise breaks compatibility with LeRobotDataset)
-        tmp_df = pd.read_parquet(ep_path)
-        tmp_df["episode_index"] = ep_idx
-        tmp_df.to_parquet(ep_path)
-
-        ep_idx += 1
-
-        if size_in_mb < data_file_size_in_mb:
-            paths_to_cat.append(ep_path)
-            continue
-
-        if paths_to_cat:
-            concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
-
-        # Reset for the next file
-        size_in_mb = ep_size_in_mb
-        paths_to_cat = [ep_path]
-
-        chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
-
-    # Write remaining data if any
-    if paths_to_cat:
-        concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
-
-    return episodes_metadata
-
-
-def convert_videos_of_camera(
-    root: Path, new_root: Path, video_key: str, video_file_size_in_mb: int, task_index: int
-):
-    # Access old paths to mp4
-    # videos_dir = root / "videos"
-    # ep_paths = sorted(videos_dir.glob(f"*/{video_key}/*.mp4"))
-    task_dir_name = f"task-{task_index:04d}"
-    videos_dir = root / "videos" / task_dir_name / video_key
-    ep_paths = sorted(videos_dir.glob("*.mp4"))
-    ep_idx = 0
-    chunk_idx = 0
-    file_idx = 0
-    size_in_mb = 0
-    duration_in_s = 0.0
-    paths_to_cat = []
-    episodes_metadata = []
-
-    for ep_path in tqdm.tqdm(ep_paths, desc=f"convert videos of {video_key}"):
-        ep_size_in_mb = get_file_size_in_mb(ep_path)
-        ep_duration_in_s = get_video_duration_in_s(ep_path)
-
-        # Check if adding this episode would exceed the limit
-        if size_in_mb + ep_size_in_mb >= video_file_size_in_mb and len(paths_to_cat) > 0:
-            # Size limit would be exceeded, save current accumulation WITHOUT this episode
-            concatenate_video_files(
-                paths_to_cat,
-                new_root
-                / DEFAULT_VIDEO_PATH.format(video_key=video_key, chunk_index=chunk_idx, file_index=file_idx),
-            )
-
-            # Update episodes metadata for the file we just saved
-            for i, _ in enumerate(paths_to_cat):
-                past_ep_idx = ep_idx - len(paths_to_cat) + i
-                episodes_metadata[past_ep_idx][f"videos/{video_key}/chunk_index"] = chunk_idx
-                episodes_metadata[past_ep_idx][f"videos/{video_key}/file_index"] = file_idx
-
-            # Move to next file and start fresh with current episode
-            chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
-            size_in_mb = 0
-            duration_in_s = 0.0
-            paths_to_cat = []
-
-        # Add current episode metadata
-        ep_metadata = {
-            "episode_index": ep_idx,
-            f"videos/{video_key}/chunk_index": chunk_idx,  # Will be updated when file is saved
-            f"videos/{video_key}/file_index": file_idx,  # Will be updated when file is saved
-            f"videos/{video_key}/from_timestamp": duration_in_s,
-            f"videos/{video_key}/to_timestamp": duration_in_s + ep_duration_in_s,
-        }
-        episodes_metadata.append(ep_metadata)
-
-        # Add current episode to accumulation
-        paths_to_cat.append(ep_path)
-        size_in_mb += ep_size_in_mb
-        duration_in_s += ep_duration_in_s
-        ep_idx += 1
-
-    # Write remaining videos if any
-    if paths_to_cat:
-        concatenate_video_files(
-            paths_to_cat,
-            new_root
-            / DEFAULT_VIDEO_PATH.format(video_key=video_key, chunk_index=chunk_idx, file_index=file_idx),
-        )
-
-        # Update episodes metadata for the final file
-        for i, _ in enumerate(paths_to_cat):
-            past_ep_idx = ep_idx - len(paths_to_cat) + i
-            episodes_metadata[past_ep_idx][f"videos/{video_key}/chunk_index"] = chunk_idx
-            episodes_metadata[past_ep_idx][f"videos/{video_key}/file_index"] = file_idx
-
-    return episodes_metadata
-
-
-def get_video_keys(root):
-    info = load_info(root)
-    features = info["features"]
-    video_keys = [key for key, ft in features.items() if ft["dtype"] == "video"]
-    return video_keys
-
-
-def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int, task_id: int):
-    logging.info(f"Converting videos from {root} to {new_root}")
-
-    video_keys = get_video_keys(root)
-    if len(video_keys) == 0:
-        return None
-
-    video_keys = sorted(video_keys)
-
-    eps_metadata_per_cam = []
-    for camera in video_keys:
-        eps_metadata = convert_videos_of_camera(root, new_root, camera, video_file_size_in_mb, task_id)
-        eps_metadata_per_cam.append(eps_metadata)
-
-    num_eps_per_cam = [len(eps_cam_map) for eps_cam_map in eps_metadata_per_cam]
-    if len(set(num_eps_per_cam)) != 1:
-        raise ValueError(f"All cams dont have same number of episodes ({num_eps_per_cam}).")
-
-    episodes_metadata = []
-    num_cameras = len(video_keys)
-    num_episodes = num_eps_per_cam[0]
-    for ep_idx in tqdm.tqdm(range(num_episodes), desc="convert videos"):
-        # Sanity check
-        ep_ids = [eps_metadata_per_cam[cam_idx][ep_idx]["episode_index"] for cam_idx in range(num_cameras)]
-        ep_ids += [ep_idx]
-        if len(set(ep_ids)) != 1:
-            raise ValueError(f"All episode indices need to match ({ep_ids}).")
-
-        ep_dict = {}
-        for cam_idx in range(num_cameras):
-            ep_dict.update(eps_metadata_per_cam[cam_idx][ep_idx])
-        episodes_metadata.append(ep_dict)
-
-    return episodes_metadata
-
-
-def infer_task_episode_ranges(episodes_jsonl_path: Path) -> dict:
-    """
-    Parse the Behavior-1K episodes.jsonl metadata and infer contiguous episode ranges per unique task.
-    Returns a dict:
-      { task_id: { "task_string": ..., "ep_start": ..., "ep_end": ... } }
-    """
-    task_ranges = {}
-    task_id = 0
-    current_task_str = None
-    ep_start = None
-    ep_end = None
-
-    with open(episodes_jsonl_path) as f:
-        for line in f:
-            if not line.strip():
-                continue
-            ep = json.loads(line)
-            ep_idx = ep["episode_index"]
-            task_str = ep["tasks"][0] if ep["tasks"] else "UNKNOWN"
-
-            if current_task_str is None:
-                current_task_str = task_str
-                ep_start = ep_idx
-                ep_end = ep_idx
-            elif task_str == current_task_str:
-                ep_end = ep_idx
-            else:
-                # close previous task group
-                task_ranges[task_id] = {
-                    "task_string": current_task_str,
-                    "ep_start": ep_start,
-                    "ep_end": ep_end,
-                }
-                task_id += 1
-                # start new one
-                current_task_str = task_str
-                ep_start = ep_idx
-                ep_end = ep_idx
-
-    # store last task
-    if current_task_str is not None:
-        task_ranges[task_id] = {
-            "task_string": current_task_str,
-            "ep_start": ep_start,
-            "ep_end": ep_end,
-        }
-
-    return task_ranges
-
-
-def legacy_load_episodes_task(local_dir: Path, task_id: int, task_ranges: dict, step: int = 10) -> dict:
-    """
-    Load only the episodes belonging to a specific task, inferred automatically from episode ranges.
-
-    Args:
-        local_dir (Path): Root path containing legacy meta/episodes.jsonl
-        task_id (int): Which task to load (key from the inferred task_ranges dict)
-        task_ranges (dict): Mapping from infer_task_episode_ranges()
-        step (int): Episode index step (Behavior-1K = 10)
-    """
-    all_episodes = legacy_load_episodes(local_dir)
-
-    # get the range for this task
-    if task_id not in task_ranges:
-        raise ValueError(f"Task id {task_id} not found in task_ranges")
-
-    ep_start = task_ranges[task_id]["ep_start"]
-    ep_end = task_ranges[task_id]["ep_end"]
-
-    task_episode_indices = range(ep_start, ep_end + step, step)
-    return {i: all_episodes[i] for i in task_episode_indices if i in all_episodes}
-
-
-def legacy_load_episodes(local_dir: Path) -> dict:
-    episodes = load_jsonlines(local_dir / LEGACY_EPISODES_PATH)
-    return {item["episode_index"]: item for item in sorted(episodes, key=lambda x: x["episode_index"])}
-
-
-def legacy_load_episodes_stats(local_dir: Path) -> dict:
-    episodes_stats = load_jsonlines(local_dir / LEGACY_EPISODES_STATS_PATH)
-    return {
-        item["episode_index"]: cast_stats_to_numpy(item["stats"])
-        for item in sorted(episodes_stats, key=lambda x: x["episode_index"])
-    }
-
-
-def legacy_load_episodes_stats_task(local_dir: Path, task_id: int, task_ranges: dict, step: int = 10) -> dict:
-    all_stats = legacy_load_episodes_stats(local_dir)
-
-    if task_id not in task_ranges:
-        raise ValueError(f"Task id {task_id} not found in task_ranges")
-
-    ep_start = task_ranges[task_id]["ep_start"]
-    ep_end = task_ranges[task_id]["ep_end"]
-
-    task_episode_indices = range(ep_start, ep_end + step, step)
-    return {i: all_stats[i] for i in task_episode_indices if i in all_stats}
-
-
-def generate_episode_metadata_dict(
-    episodes_legacy_metadata, episodes_metadata, episodes_stats, episodes_videos=None
-):
-    num_episodes = len(episodes_metadata)
-    episodes_legacy_metadata_vals = list(episodes_legacy_metadata.values())
-    episodes_stats_vals = list(episodes_stats.values())
-    episodes_stats_keys = list(episodes_stats.keys())
-
-    for i in range(num_episodes):
-        ep_legacy_metadata = episodes_legacy_metadata_vals[i]
-        ep_metadata = episodes_metadata[i]
-        ep_stats = episodes_stats_vals[i]
-
-        ep_ids_set = {
-            ep_legacy_metadata["episode_index"],
-            ep_metadata["episode_index"],
-            episodes_stats_keys[i],
-        }
-
-        if episodes_videos is None:
-            ep_video = {}
-        else:
-            ep_video = episodes_videos[i]
-            ep_ids_set.add(ep_video["episode_index"])
-
-        ep_dict = {
-            **ep_legacy_metadata,
-            **ep_video,
-            **ep_metadata,
-            **flatten_dict({"stats": ep_stats}),
-        }
-
-        # enforce contiguous indexing 0..n-1, but also stores the legacy episode index
-        ep_dict["episode_index"] = i
-
-        yield ep_dict
-
-
-def convert_episodes_metadata(
-    root, new_root, episodes_metadata, task_id: int, task_ranges, episodes_video_metadata=None
-):
-    logging.info(f"Converting episodes metadata from {root} to {new_root}")
-
-    # filter by task
-    episodes_legacy_metadata = legacy_load_episodes_task(root, task_id=task_id, task_ranges=task_ranges)
-    episodes_stats = legacy_load_episodes_stats_task(root, task_id=task_id, task_ranges=task_ranges)
-
-    num_eps_set = {len(episodes_legacy_metadata), len(episodes_metadata)}
-    if episodes_video_metadata is not None:
-        num_eps_set.add(len(episodes_video_metadata))
-
-    if len(num_eps_set) != 1:
-        raise ValueError(f"Number of episodes is not the same ({num_eps_set}).")
-
-    # Single file approach: set meta indices to 0 for all rows and write once
-    ds_episodes = Dataset.from_generator(
-        lambda: generate_episode_metadata_dict(
-            episodes_legacy_metadata, episodes_metadata, episodes_stats, episodes_video_metadata
-        )
-    )
-    num_eps = len(ds_episodes)
-    # NOTE(fracapuano): for the size of the average dataset this is fine!
-    ds_episodes = ds_episodes.add_column("meta/episodes/chunk_index", [0] * num_eps)
-    ds_episodes = ds_episodes.add_column("meta/episodes/file_index", [0] * num_eps)
-    write_episodes(ds_episodes, new_root)
-
-    stats = aggregate_stats(list(episodes_stats.values()))
-    write_stats(stats, new_root)
-
-
-def convert_dataset_local(
-    data_path: Path,
-    new_repo: Path,
-    task_id: int,
-    data_file_size_in_mb: int = DEFAULT_DATA_FILE_SIZE_IN_MB,
-    video_file_size_in_mb: int = DEFAULT_VIDEO_FILE_SIZE_IN_MB,
-    force_conversion: bool = False,
-):
-    """
-    Convert a local dataset to v3.x format, task-by-task, without using the Hugging Face Hub.
-
-    Args:
-        data_path (Path): path to local dataset root (e.g. /fsx/.../2025-challenge-demos)
-        new_repo (Path): path where converted dataset will be written (e.g. /fsx/.../behavior1k_v3)
-        task_id (int): which task to convert (index)
-        data_file_size_in_mb (int): max size per data chunk
-        video_file_size_in_mb (int): max size per video chunk
-        force_conversion (bool): overwrite existing conversion if True
-    """
-
-    root = Path(data_path)
-    new_root = Path(new_repo)
-
-    # Clean up if needed
-    if new_root.exists() and force_conversion:
-        shutil.rmtree(new_root)
-    new_root.mkdir(parents=True, exist_ok=True)
-
-    print(f"🔹 Starting conversion for task {task_id}")
-    print(f"Input root: {root}")
-    print(f"Output root: {new_root}")
-    # Infer task episode ranges
-    episodes_meta_path = root / "meta" / "episodes.jsonl"
-    task_ranges = infer_task_episode_ranges(episodes_meta_path)
-    convert_info(
-        root,
-        new_root,
-        data_file_size_in_mb,
-        video_file_size_in_mb,
-        episodes_meta_path,
-        task_id,
-        task_ranges,
-        step=10,
-    )
-    convert_tasks(root, new_root, task_id)
-    episodes_metadata = convert_data(root, new_root, data_file_size_in_mb, task_index=task_id)
-    episodes_videos_metadata = convert_videos(root, new_root, video_file_size_in_mb, task_id=task_id)
-    convert_episodes_metadata(
-        root,
-        new_root,
-        episodes_metadata,
-        task_id=task_id,
-        task_ranges=task_ranges,
-        episodes_video_metadata=episodes_videos_metadata,
-    )
-
-    print(f"✅ Conversion complete for task {task_id}")
-    print(f"Converted dataset written to: {new_root}")
-
-
-if __name__ == "__main__":
-    import argparse
-    from pathlib import Path
-
-    init_logging()
-
-    parser = argparse.ArgumentParser(
-        description="Convert Behavior-1K tasks to LeRobot v3 format (local only)"
-    )
-    parser.add_argument(
-        "--data-path",
-        type=str,
-        required=True,
-        help="Path to the local Behavior-1K dataset (e.g. /fsx/francesco_capuano/.cache/behavior-1k/2025-challenge-demos)",
-    )
-    parser.add_argument(
-        "--new-repo",
-        type=str,
-        required=True,
-        help="Path to the output directory for the converted dataset",
-    )
-    parser.add_argument(
-        "--task-id",
-        type=int,
-        required=True,
-        help="Task index to convert (e.g. 0, 1, 2, ...)",
-    )
-    parser.add_argument(
-        "--data-file-size-in-mb",
-        type=int,
-        default=DEFAULT_DATA_FILE_SIZE_IN_MB,
-        help=f"Maximum size per data chunk (default: {DEFAULT_DATA_FILE_SIZE_IN_MB})",
-    )
-    parser.add_argument(
-        "--video-file-size-in-mb",
-        type=int,
-        default=DEFAULT_VIDEO_FILE_SIZE_IN_MB,
-        help=f"Maximum size per video chunk (default: {DEFAULT_VIDEO_FILE_SIZE_IN_MB})",
-    )
-    parser.add_argument(
-        "--force-conversion",
-        action="store_true",
-        help="Force overwrite of existing conversion output if present.",
-    )
-    parser.add_argument(
-        "--push-to-hub",
-        action="store_true",
-        help="Push the (converted) dataset to the hub.",
-    )
-
-    args = parser.parse_args()
-
-    if args.push_to_hub:
-        HF_USER = os.environ.get("HF_USER", "fracapuano")
-        if HF_USER is None:
-            raise ValueError(
-                "HF_USER environment variable is not set! Set before converting and pushing to hub."
-            )
-
-    convert_dataset_local(
-        data_path=Path(args.data_path),
-        new_repo=Path(args.new_repo),
-        task_id=args.task_id,
-        data_file_size_in_mb=args.data_file_size_in_mb,
-        video_file_size_in_mb=args.video_file_size_in_mb,
-        force_conversion=args.force_conversion,
-    )
-
-    if args.push_to_hub:
-        ds = LeRobotDataset(repo_id=f"{HF_USER}/behavior1k-task{args.task_id:04d}", root=args.new_repo)
-        ds.push_to_hub()

examples/behavior_1k/download.slurm

@@ -1,27 +0,0 @@
-#!/bin/bash
-#SBATCH -J b1k-download
-#SBATCH -p hopper-cpu
-#SBATCH --qos=high
-#SBATCH -c 32                    # CPUs per conversion (tune as needed)
-#SBATCH -t 20:00:00              # Time per conversion
-#SBATCH -D /admin/home/francesco_capuano/lerobot
-#SBATCH -o /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%A.out
-#SBATCH -e /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%A.err
-
-set -euo pipefail
-set -x
-export PYTHONUNBUFFERED=1
-export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK:-1}
-
-# activate your env if needed
-source "$HOME/.bashrc" 2>/dev/null || true
-if ! command -v conda >/dev/null 2>&1; then
-  source "$HOME/miniconda3/etc/profile.d/conda.sh" 2>/dev/null || true
-  source "$HOME/anaconda3/etc/profile.d/conda.sh" 2>/dev/null || true
-fi
-conda activate lerobot
-
-python examples/behavior_1k/download_data.py \
-  --repo-id "behavior-1k/2025-challenge-demos" \
-  --local-dir "/fsx/francesco_capuano/behavior1k-2025-v21" \
-  --max-workers 32

examples/behavior_1k/download_data.py

@@ -1,26 +0,0 @@
-import shutil
-
-from huggingface_hub import snapshot_download
-
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--repo-id", type=str, required=True)
-    parser.add_argument("--max-workers", type=int, default=8)
-    parser.add_argument("--local-dir", type=str, required=True)
-    parser.add_argument("--force-download", action="store_true")
-
-    args = parser.parse_args()
-
-    if args.force_download:
-        shutil.rmtree(args.local_dir, ignore_errors=True)
-
-    snapshot_download(
-        repo_id=args.repo_id,
-        repo_type="dataset",
-        force_download=args.force_download,
-        max_workers=args.max_workers,
-        local_dir=args.local_dir,
-        ignore_patterns=["annotations/*"],  # NOTE(fracapuano): Dropping textual annotations right now
-    )

examples/behavior_1k/upload.slurm

@@ -1,41 +0,0 @@
-#!/bin/bash
-#SBATCH -J b1k-upload
-#SBATCH -p hopper-cpu
-#SBATCH --qos=high
-#SBATCH -c 1
-#SBATCH -t 48:00:00
-#SBATCH --mem=4G
-#SBATCH --array=0-49%2
-#SBATCH -D /admin/home/francesco_capuano/lerobot
-#SBATCH -o /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%A_%a.out
-#SBATCH -e /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%A_%a.err
-
-set -euo pipefail
-set -x
-export PYTHONUNBUFFERED=1
-export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK:-1}
-
-source "$HOME/.bashrc" 2>/dev/null || true
-if ! command -v conda >/dev/null 2>&1; then
-  source "$HOME/miniconda3/etc/profile.d/conda.sh" 2>/dev/null || true
-  source "$HOME/anaconda3/etc/profile.d/conda.sh" 2>/dev/null || true
-fi
-conda activate lerobot
-
-# The SLURM_ARRAY_TASK_ID will be used as the task-id
-TASK_ID=${SLURM_ARRAY_TASK_ID}
-
-# Configuration
-ROOT_PATH="/fsx/francesco_capuano/behavior1k-v3"
-HF_USER="fracapuano"
-# Limit upload workers to reduce network contention (default in HF Hub is 4)
-# For I/O-bound uploads, 2-4 workers per task is optimal
-NUM_WORKERS=2
-
-echo "Task ${TASK_ID}: uploading with ${NUM_WORKERS} workers from ${ROOT_PATH}"
-
-python examples/behavior_1k/upload_folders.py \
-  --task-id ${TASK_ID} \
-  --root-path ${ROOT_PATH} \
-  --hf-user ${HF_USER} \
-  --num-workers ${NUM_WORKERS}

examples/behavior_1k/upload_folders.py

@@ -1,108 +0,0 @@
-import argparse
-from pathlib import Path
-
-from huggingface_hub import HfApi, upload_large_folder
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Upload a folder to Hugging Face Hub using upload_large_folder"
-    )
-    parser.add_argument(
-        "--folder-path",
-        type=str,
-        required=False,
-        help="Path to the folder to upload (used if task-id is not provided)",
-    )
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        required=False,
-        help="Repository ID on Hugging Face Hub (e.g., 'username/repo-name'). If task-id is provided, will be constructed as '{hf-user}/behavior1k-task{task_id:04d}'",
-    )
-    parser.add_argument(
-        "--task-id",
-        type=int,
-        required=False,
-        help="Task index to upload (e.g., 0, 1, 2, ...). When provided, folder-path is constructed from root-path.",
-    )
-    parser.add_argument(
-        "--root-path",
-        type=str,
-        required=False,
-        help="Root path containing task folders (e.g., /fsx/user/behavior1k-v3). Used with --task-id to construct folder path.",
-    )
-    parser.add_argument(
-        "--hf-user",
-        type=str,
-        default=None,
-        help="Hugging Face username for constructing repo-id with task-id (default: from HF_USER env var or 'fracapuano')",
-    )
-    parser.add_argument(
-        "--create-repo", action="store_true", help="Create the repository if it doesn't exist"
-    )
-    parser.add_argument(
-        "--num-workers",
-        type=int,
-        default=2,
-        help="Number of parallel workers for upload (default: 2). For I/O-bound uploads, use 1-4 to avoid network contention.",
-    )
-
-    args = parser.parse_args()
-
-    # Construct folder path and repo ID based on task-id or use provided values
-    if args.task_id is not None:
-        if not args.root_path:
-            raise ValueError("--root-path is required when --task-id is provided")
-
-        task_folder_name = f"behavior1k-task{args.task_id:04d}"
-        folder_path = Path(args.root_path) / task_folder_name
-        repo_id = f"{args.hf_user}/{task_folder_name}"
-
-        print(f"Task mode: uploading task {args.task_id}")
-    else:
-        if not args.folder_path or not args.repo_id:
-            raise ValueError(
-                "Either --task-id with --root-path, or both --folder-path and --repo-id must be provided"
-            )
-
-        folder_path = Path(args.folder_path)
-        repo_id = args.repo_id
-
-    # Validate folder path
-    if not folder_path.exists():
-        raise ValueError(f"Folder path does not exist: {folder_path}")
-    if not folder_path.is_dir():
-        raise ValueError(f"Path is not a directory: {folder_path}")
-
-    print(f"Uploading folder: {folder_path}")
-    print(f"Repository: {repo_id}")
-
-    # Create repository if requested
-    if args.create_repo:
-        api = HfApi()
-        print(f"Creating repository {repo_id}...")
-        try:
-            api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True)
-            print("Repository created or already exists. Updating its contents")
-        except Exception as e:
-            print(f"Warning: Could not create repository: {e}")
-
-    # Upload the folder
-    print(f"Starting upload with {args.num_workers} parallel workers...")
-    try:
-        result = upload_large_folder(
-            folder_path=str(folder_path),
-            repo_id=repo_id,
-            repo_type="dataset",
-            num_workers=args.num_workers,
-        )
-        print("✓ Upload completed successfully!")
-        print(f"Commit URL: {result}")
-    except Exception as e:
-        print(f"✗ Upload failed: {e}")
-        raise
-
-
-if __name__ == "__main__":
-    main()

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:
-hub_api = HfApi()
-repo_ids = [info.id for info in hub_api.list_datasets(task_categories="robotics", tags=["LeRobot"])]
-pprint(repo_ids)
+def main():
+    # We ported a number of existing datasets ourselves, use this to see the list:
+    print("List of available datasets:")
+    pprint(lerobot.available_datasets)
 
-# Or simply explore them in your web browser directly at:
-# https://huggingface.co/datasets?other=LeRobot
+    # You can also browse through the datasets created/ported by the community on the hub using the hub api:
+    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
-repo_id = "lerobot/aloha_mobile_cabinet"
-# We can have a look and fetch its metadata to know more about it:
-ds_meta = LeRobotDatasetMetadata(repo_id)
+    # Or simply explore them in your web browser directly at:
+    # https://huggingface.co/datasets?other=LeRobot
 
-# By instantiating just this class, you can quickly access useful information about the content and the
-# structure of the dataset without downloading the actual data yet (only metadata files — which are
-# lightweight).
-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")
+    # Let's take this one for this example
+    repo_id = "lerobot/aloha_mobile_cabinet"
+    # We can have a look and fetch its metadata to know more about it:
+    ds_meta = LeRobotDatasetMetadata(repo_id)
 
-print("Tasks:")
-print(ds_meta.tasks)
-print("Features:")
-pprint(ds_meta.features)
+    # By instantiating just this class, you can quickly access useful information about the content and the
+    # structure of the dataset without downloading the actual data yet (only metadata files — which are
+    # lightweight).
+    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(ds_meta)
+    print("Tasks:")
+    print(ds_meta.tasks)
+    print("Features:")
+    pprint(ds_meta.features)
 
-# You can then load the actual dataset from the hub.
-# Either load any subset of episodes:
-dataset = LeRobotDataset(repo_id, episodes=[0, 10, 11, 23])
+    # You can also get a short summary by simply printing the object:
+    print(ds_meta)
 
-# And see how many frames you have:
-print(f"Selected episodes: {dataset.episodes}")
-print(f"Number of episodes selected: {dataset.num_episodes}")
-print(f"Number of frames selected: {dataset.num_frames}")
+    # You can then load the actual dataset from the hub.
+    # Either load any subset of episodes:
+    dataset = LeRobotDataset(repo_id, episodes=[0, 10, 11, 23])
 
-# Or simply load the entire dataset:
-dataset = LeRobotDataset(repo_id)
-print(f"Number of episodes selected: {dataset.num_episodes}")
-print(f"Number of frames selected: {dataset.num_frames}")
+    # And see how many frames you have:
+    print(f"Selected episodes: {dataset.episodes}")
+    print(f"Number of episodes selected: {dataset.num_episodes}")
+    print(f"Number of frames selected: {dataset.num_frames}")
 
-# The previous metadata class is contained in the 'meta' attribute of the dataset:
-print(dataset.meta)
+    # Or simply load the entire dataset:
+    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
-# (see https://huggingface.co/docs/datasets for more information).
-print(dataset.hf_dataset)
+    # The previous metadata class is contained in the 'meta' attribute of the dataset:
+    print(dataset.meta)
 
-# LeRobot datasets also subclasses PyTorch datasets so you can do everything you know and love from working
-# with the latter, like iterating through the dataset.
-# 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]
+    # LeRobotDataset actually wraps an underlying Hugging Face dataset
+    # (see https://huggingface.co/docs/datasets for more information).
+    print(dataset.hf_dataset)
 
-# Then we grab all the image frames from the first camera:
-camera_key = dataset.meta.camera_keys[0]
-frames = [dataset[idx][camera_key] for idx in range(from_idx, to_idx)]
+    # LeRobot datasets also subclasses PyTorch datasets so you can do everything you know and love from working
+    # with the latter, like iterating through the dataset.
+    # 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
-print(type(frames[0]))
-print(frames[0].shape)
+    # Then we grab all the image frames from the first camera:
+    camera_key = dataset.meta.camera_keys[0]
+    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).
-# We can compare this shape with the information available for that feature
-pprint(dataset.features[camera_key])
-# In particular:
-print(dataset.features[camera_key]["shape"])
-# The shape is in (h, w, c) which is a more universal format.
+    # The objects returned by the dataset are all torch.Tensors
+    print(type(frames[0]))
+    print(frames[0].shape)
 
-# For many machine learning applications we need to load the history of past observations or trajectories of
-# future actions. Our datasets can load previous and future frames for each key/modality, using timestamps
-# differences with the current loaded frame. For instance:
-delta_timestamps = {
-    # loads 4 images: 1 second before current frame, 500 ms before, 200 ms before, and current frame
-    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.
+    # Since we're using pytorch, the shape is in pytorch, channel-first convention (c, h, w).
+    # We can compare this shape with the information available for that feature
+    pprint(dataset.features[camera_key])
+    # In particular:
+    print(dataset.features[camera_key]["shape"])
+    # The shape is in (h, w, c) which is a more universal format.
 
-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)
+    # For many machine learning applications we need to load the history of past observations or trajectories of
+    # future actions. Our datasets can load previous and future frames for each key/modality, using timestamps
+    # differences with the current loaded frame. For instance:
+    delta_timestamps = {
+        # loads 4 images: 1 second before current frame, 500 ms before, 200 ms before, and current frame
+        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
-policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
+    robot = LeKiwiClient(robot_config)
 
-# Configure the dataset features
-action_features = hw_to_dataset_features(robot.action_features, ACTION)
-obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
-dataset_features = {**action_features, **obs_features}
+    # Create policy
+    policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
 
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id=HF_DATASET_ID,
-    fps=FPS,
-    features=dataset_features,
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
+    # Configure the dataset features
+    action_features = hw_to_dataset_features(robot.action_features, ACTION)
+    obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
+    dataset_features = {**action_features, **obs_features}
 
-# 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
-# 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,
+    # Create the dataset
+    dataset = LeRobotDataset.create(
+        repo_id=HF_DATASET_ID,
         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=teleop_action_processor,
-        robot_action_processor=robot_action_processor,
-        robot_observation_processor=robot_observation_processor,
+        features=dataset_features,
+        robot_type=robot.name,
+        use_videos=True,
+        image_writer_threads=4,
     )
 
-    # Reset the environment if not stopping or re-recording
-    if not events["stop_recording"] and (
-        (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
-    ):
-        log_say("Reset the environment")
+    # 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
+    # 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"]:
-        log_say("Re-record episode")
-        events["rerecord_episode"] = False
-        events["exit_early"] = False
-        dataset.clear_episode_buffer()
-        continue
+        # Reset the environment if not stopping or re-recording
+        if not events["stop_recording"] and (
+            (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
+        ):
+            log_say("Reset the environment")
+            record_loop(
+                robot=robot,
+                events=events,
+                fps=FPS,
+                control_time_s=EPISODE_TIME_SEC,
+                single_task=TASK_DESCRIPTION,
+                display_data=True,
+                teleop_action_processor=teleop_action_processor,
+                robot_action_processor=robot_action_processor,
+                robot_observation_processor=robot_observation_processor,
+            )
 
-    # Save episode
-    dataset.save_episode()
-    recorded_episodes += 1
+        if events["rerecord_episode"]:
+            log_say("Re-record episode")
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-# Clean up
-log_say("Stop recording")
-robot.disconnect()
-listener.stop()
+        # Save episode
+        dataset.save_episode()
+        recorded_episodes += 1
 
-dataset.finalize()
-dataset.push_to_hub()
+    # Clean up
+    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
-robot = LeKiwiClient(robot_config)
-leader_arm = SO100Leader(leader_arm_config)
-keyboard = KeyboardTeleop(keyboard_config)
+def main():
+    # 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()
 
-# TODO(Steven): Update this example to use pipelines
-teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
+    # Initialize the robot and teleoperator
+    robot = LeKiwiClient(robot_config)
+    leader_arm = SO100Leader(leader_arm_config)
+    keyboard = KeyboardTeleop(keyboard_config)
 
-# Configure the dataset features
-action_features = hw_to_dataset_features(robot.action_features, ACTION)
-obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
-dataset_features = {**action_features, **obs_features}
+    # TODO(Steven): Update this example to use pipelines
+    teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
 
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id=HF_REPO_ID,
-    fps=FPS,
-    features=dataset_features,
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
+    # Configure the dataset features
+    action_features = hw_to_dataset_features(robot.action_features, ACTION)
+    obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
+    dataset_features = {**action_features, **obs_features}
 
-# Connect the robot and teleoperator
-# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
-robot.connect()
-leader_arm.connect()
-keyboard.connect()
-
-# Initialize the keyboard listener and rerun visualization
-listener, events = init_keyboard_listener()
-init_rerun(session_name="lekiwi_record")
-
-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,
+    # Create the dataset
+    dataset = LeRobotDataset.create(
+        repo_id=HF_REPO_ID,
         fps=FPS,
-        dataset=dataset,
-        teleop=[leader_arm, keyboard],
-        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,
+        features=dataset_features,
+        robot_type=robot.name,
+        use_videos=True,
+        image_writer_threads=4,
     )
 
-    # Reset the environment if not stopping or re-recording
-    if not events["stop_recording"] and (
-        (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
-    ):
-        log_say("Reset the environment")
+    # Connect the robot and teleoperator
+    # To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
+    robot.connect()
+    leader_arm.connect()
+    keyboard.connect()
+
+    # Initialize the keyboard listener and rerun visualization
+    listener, events = init_keyboard_listener()
+    init_rerun(session_name="lekiwi_record")
+
+    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"]:
-        log_say("Re-record episode")
-        events["rerecord_episode"] = False
-        events["exit_early"] = False
-        dataset.clear_episode_buffer()
-        continue
+        # Reset the environment if not stopping or re-recording
+        if not events["stop_recording"] and (
+            (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
+        ):
+            log_say("Reset the environment")
+            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
-    dataset.save_episode()
-    recorded_episodes += 1
+        if events["rerecord_episode"]:
+            log_say("Re-record episode")
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-# Clean up
-log_say("Stop recording")
-robot.disconnect()
-leader_arm.disconnect()
-keyboard.disconnect()
-listener.stop()
+        # Save episode
+        dataset.save_episode()
+        recorded_episodes += 1
 
-dataset.finalize()
-dataset.push_to_hub()
+    # Clean up
+    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

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

examples/lekiwi/teleoperate.py

@@ -19,54 +19,60 @@ import time
 from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
 from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop, KeyboardTeleopConfig
 from lerobot.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 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
-robot = LeKiwiClient(robot_config)
-leader_arm = SO100Leader(teleop_arm_config)
-keyboard = KeyboardTeleop(keyboard_config)
+def main():
+    # 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")
 
-# Connect to the robot and teleoperator
-# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
-robot.connect()
-leader_arm.connect()
-keyboard.connect()
+    # Initialize the robot and teleoperator
+    robot = LeKiwiClient(robot_config)
+    leader_arm = SO100Leader(teleop_arm_config)
+    keyboard = KeyboardTeleop(keyboard_config)
 
-# Init rerun viewer
-init_rerun(session_name="lekiwi_teleop")
+    # Connect to the robot and teleoperator
+    # To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
+    robot.connect()
+    leader_arm.connect()
+    keyboard.connect()
 
-if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
-    raise ValueError("Robot or teleop is not connected!")
+    # Init rerun viewer
+    init_rerun(session_name="lekiwi_teleop")
 
-print("Starting teleop loop...")
-while True:
-    t0 = time.perf_counter()
+    if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
+        raise ValueError("Robot or teleop is not connected!")
 
-    # Get robot observation
-    observation = robot.get_observation()
+    print("Starting teleop loop...")
+    while True:
+        t0 = time.perf_counter()
 
-    # 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)
+        # Get robot observation
+        observation = robot.get_observation()
 
-    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
-    _ = robot.send_action(action)
+        action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
 
-    # Visualize
-    log_rerun_data(observation=observation, action=action)
+        # Send action to robot
+        _ = robot.send_action(action)
 
-    busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
+        # Visualize
+        log_rerun_data(observation=observation, action=action)
+
+        precise_sleep(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)
-
-# 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,
-        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,
+def main():
+    # 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,
     )
 
-    # Reset the environment if not stopping or re-recording
-    if not events["stop_recording"] and ((episode_idx < NUM_EPISODES - 1) or events["rerecord_episode"]):
-        log_say("Reset the environment")
+    robot = SO100Follower(robot_config)
+
+    # Create policy
+    policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
+
+    # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
+    kinematics_solver = RobotKinematics(
+        urdf_path="./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"]:
-        log_say("Re-record episode")
-        events["rerecord_episode"] = False
-        events["exit_early"] = False
-        dataset.clear_episode_buffer()
-        continue
+        # Reset the environment if not stopping or re-recording
+        if not events["stop_recording"] and ((episode_idx < NUM_EPISODES - 1) or events["rerecord_episode"]):
+            log_say("Reset the environment")
+            record_loop(
+                robot=robot,
+                events=events,
+                fps=FPS,
+                control_time_s=EPISODE_TIME_SEC,
+                single_task=TASK_DESCRIPTION,
+                display_data=True,
+                teleop_action_processor=make_default_teleop_action_processor(),
+                robot_action_processor=robot_ee_to_joints_processor,
+                robot_observation_processor=robot_joints_to_ee_pose_processor,
+            )
 
-    # Save episode
-    dataset.save_episode()
-    episode_idx += 1
+        if events["rerecord_episode"]:
+            log_say("Re-record episode")
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-# Clean up
-log_say("Stop recording")
-robot.disconnect()
-listener.stop()
+        # Save episode
+        dataset.save_episode()
+        episode_idx += 1
 
-dataset.finalize()
-dataset.push_to_hub()
+    # Clean up
+    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
-robot = SO100Follower(robot_config)
-phone = Phone(teleop_config)
+def main():
+    # 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
 
-# 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()),
-)
+    # Initialize the robot and teleoperator
+    robot = SO100Follower(robot_config)
+    phone = Phone(teleop_config)
 
-# Build pipeline to convert phone action to EE action
-phone_to_robot_ee_pose_processor = RobotProcessorPipeline[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,
-        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,
+    # 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()),
     )
 
-    # Reset the environment if not stopping or re-recording
-    if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
-        log_say("Reset the environment")
+    # Build pipeline to convert phone action to EE action
+    phone_to_robot_ee_pose_processor = RobotProcessorPipeline[
+        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"]:
-        log_say("Re-recording episode")
-        events["rerecord_episode"] = False
-        events["exit_early"] = False
-        dataset.clear_episode_buffer()
-        continue
+        # Reset the environment if not stopping or re-recording
+        if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
+            log_say("Reset the environment")
+            record_loop(
+                robot=robot,
+                events=events,
+                fps=FPS,
+                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
-    dataset.save_episode()
-    episode_idx += 1
+        if events["rerecord_episode"]:
+            log_say("Re-recording episode")
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-# Clean up
-log_say("Stop recording")
-robot.disconnect()
-phone.disconnect()
-listener.stop()
+        # Save episode
+        dataset.save_episode()
+        episode_idx += 1
 
-dataset.finalize()
-dataset.push_to_hub()
+    # Clean up
+    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

@@ -29,72 +29,78 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
 from lerobot.utils.constants import ACTION
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 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
-robot = SO100Follower(robot_config)
+def main():
+    # 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
-kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
+    # Initialize the robot
+    robot = SO100Follower(robot_config)
 
-# 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=False,  # Because replay is open loop
-        ),
-    ],
-    to_transition=robot_action_observation_to_transition,
-    to_output=transition_to_robot_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
+    kinematics_solver = RobotKinematics(
+        urdf_path="./SO101/so101_new_calib.urdf",
+        target_frame_name="gripper_frame_link",
+        joint_names=list(robot.bus.motors.keys()),
+    )
 
-# Fetch the dataset to replay
-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)
+    # 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=False,  # Because replay is open loop
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
 
-# Connect to the robot
-robot.connect()
+    # Fetch the dataset to replay
+    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:
-    raise ValueError("Robot is not connected!")
+    # Connect to the robot
+    robot.connect()
 
-print("Starting replay loop...")
-log_say(f"Replaying episode {EPISODE_IDX}")
-for idx in range(len(episode_frames)):
-    t0 = time.perf_counter()
+    if not robot.is_connected:
+        raise ValueError("Robot is not connected!")
 
-    # Get recorded action from dataset
-    ee_action = {
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
-    }
+    print("Starting replay loop...")
+    log_say(f"Replaying episode {EPISODE_IDX}")
+    for idx in range(len(episode_frames)):
+        t0 = time.perf_counter()
 
-    # Get robot observation
-    robot_obs = robot.get_observation()
+        # Get recorded action from dataset
+        ee_action = {
+            name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        }
 
-    # Dataset EE -> robot joints
-    joint_action = robot_ee_to_joints_processor((ee_action, robot_obs))
+        # Get robot observation
+        robot_obs = robot.get_observation()
 
-    # Send action to robot
-    _ = robot.send_action(joint_action)
+        # Dataset EE -> robot joints
+        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
-robot.disconnect()
+        precise_sleep(1.0 / dataset.fps - (time.perf_counter() - t0))
+
+    # Clean up
+    robot.disconnect()
+
+
+if __name__ == "__main__":
+    main()

examples/phone_to_so100/teleoperate.py

@@ -32,82 +32,90 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower
 from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
 from lerobot.teleoperators.phone.teleop_phone import Phone
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 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
-robot = SO100Follower(robot_config)
-teleop_device = Phone(teleop_config)
+def main():
+    # 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
 
-# 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()),
-)
+    # Initialize the robot and teleoperator
+    robot = SO100Follower(robot_config)
+    teleop_device = Phone(teleop_config)
 
-# Build pipeline to convert phone action to ee pose action to joint action
-phone_to_robot_joints_processor = RobotProcessorPipeline[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,
-)
+    # 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()),
+    )
 
-# Connect to the robot and teleoperator
-robot.connect()
-teleop_device.connect()
+    # Build pipeline to convert phone action to ee pose action to joint action
+    phone_to_robot_joints_processor = RobotProcessorPipeline[
+        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
-init_rerun(session_name="phone_so100_teleop")
+    # Connect to the robot and teleoperator
+    robot.connect()
+    teleop_device.connect()
 
-if not robot.is_connected or not teleop_device.is_connected:
-    raise ValueError("Robot or teleop is not connected!")
+    # Init rerun viewer
+    init_rerun(session_name="phone_so100_teleop")
 
-print("Starting teleop loop. Move your phone to teleoperate the robot...")
-while True:
-    t0 = time.perf_counter()
+    if not robot.is_connected or not teleop_device.is_connected:
+        raise ValueError("Robot or teleop is not connected!")
 
-    # Get robot observation
-    robot_obs = robot.get_observation()
+    print("Starting teleop loop. Move your phone to teleoperate the robot...")
+    while True:
+        t0 = time.perf_counter()
 
-    # Get teleop action
-    phone_obs = teleop_device.get_action()
+        # Get robot observation
+        robot_obs = robot.get_observation()
 
-    # Phone -> EE pose -> Joints transition
-    joint_action = phone_to_robot_joints_processor((phone_obs, robot_obs))
+        # Get teleop action
+        phone_obs = teleop_device.get_action()
 
-    # Send action to robot
-    _ = robot.send_action(joint_action)
+        # Phone -> EE pose -> Joints transition
+        joint_action = phone_to_robot_joints_processor((phone_obs, robot_obs))
 
-    # Visualize
-    log_rerun_data(observation=phone_obs, action=joint_action)
+        # Send action to robot
+        _ = 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)
+
+        precise_sleep(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)
-
-# 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,
-        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,
+def main():
+    # 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,
     )
 
-    # Reset the environment if not stopping or re-recording
-    if not events["stop_recording"] and ((episode_idx < NUM_EPISODES - 1) or events["rerecord_episode"]):
-        log_say("Reset the environment")
+    robot = SO100Follower(robot_config)
+
+    # Create policy
+    policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
+
+    # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
+    kinematics_solver = RobotKinematics(
+        urdf_path="./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"]:
-        log_say("Re-record episode")
-        events["rerecord_episode"] = False
-        events["exit_early"] = False
-        dataset.clear_episode_buffer()
-        continue
+        # Reset the environment if not stopping or re-recording
+        if not events["stop_recording"] and ((episode_idx < NUM_EPISODES - 1) or events["rerecord_episode"]):
+            log_say("Reset the environment")
+            record_loop(
+                robot=robot,
+                events=events,
+                fps=FPS,
+                control_time_s=EPISODE_TIME_SEC,
+                single_task=TASK_DESCRIPTION,
+                display_data=True,
+                teleop_action_processor=make_default_teleop_action_processor(),
+                robot_action_processor=robot_ee_to_joints_processor,
+                robot_observation_processor=robot_joints_to_ee_pose_processor,
+            )
 
-    # Save episode
-    dataset.save_episode()
-    episode_idx += 1
+        if events["rerecord_episode"]:
+            log_say("Re-record episode")
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-# Clean up
-log_say("Stop recording")
-robot.disconnect()
-listener.stop()
+        # Save episode
+        dataset.save_episode()
+        episode_idx += 1
 
-dataset.finalize()
-dataset.push_to_hub()
+    # Clean up
+    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
-follower = SO100Follower(follower_config)
-leader = SO100Leader(leader_config)
+def main():
+    # 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")
 
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-follower_kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(follower.bus.motors.keys()),
-)
+    # Initialize the robot and teleoperator
+    follower = SO100Follower(follower_config)
+    leader = SO100Leader(leader_config)
 
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-leader_kinematics_solver = RobotKinematics(
-    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,
-        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,
+    # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
+    follower_kinematics_solver = RobotKinematics(
+        urdf_path="./SO101/so101_new_calib.urdf",
+        target_frame_name="gripper_frame_link",
+        joint_names=list(follower.bus.motors.keys()),
     )
 
-    # Reset the environment if not stopping or re-recording
-    if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
-        log_say("Reset the environment")
+    # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
+    leader_kinematics_solver = RobotKinematics(
+        urdf_path="./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"]:
-        log_say("Re-recording episode")
-        events["rerecord_episode"] = False
-        events["exit_early"] = False
-        dataset.clear_episode_buffer()
-        continue
+        # Reset the environment if not stopping or re-recording
+        if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
+            log_say("Reset the environment")
+            record_loop(
+                robot=follower,
+                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
-    dataset.save_episode()
-    episode_idx += 1
+        if events["rerecord_episode"]:
+            log_say("Re-recording episode")
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-# Clean up
-log_say("Stop recording")
-leader.disconnect()
-follower.disconnect()
-listener.stop()
+        # Save episode
+        dataset.save_episode()
+        episode_idx += 1
 
-dataset.finalize()
-dataset.push_to_hub()
+    # Clean up
+    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

@@ -30,72 +30,78 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
 from lerobot.utils.constants import ACTION
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 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
-robot = SO100Follower(robot_config)
+def main():
+    # 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
-kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
+    # Initialize the robot
+    robot = SO100Follower(robot_config)
 
-# 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=False,  # Because replay is open loop
-        ),
-    ],
-    to_transition=robot_action_observation_to_transition,
-    to_output=transition_to_robot_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
+    kinematics_solver = RobotKinematics(
+        urdf_path="./SO101/so101_new_calib.urdf",
+        target_frame_name="gripper_frame_link",
+        joint_names=list(robot.bus.motors.keys()),
+    )
 
-# Fetch the dataset to replay
-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)
+    # 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=False,  # Because replay is open loop
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
 
-# Connect to the robot
-robot.connect()
+    # Fetch the dataset to replay
+    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:
-    raise ValueError("Robot is not connected!")
+    # Connect to the robot
+    robot.connect()
 
-print("Starting replay loop...")
-log_say(f"Replaying episode {EPISODE_IDX}")
-for idx in range(len(episode_frames)):
-    t0 = time.perf_counter()
+    if not robot.is_connected:
+        raise ValueError("Robot is not connected!")
 
-    # Get recorded action from dataset
-    ee_action = {
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
-    }
+    print("Starting replay loop...")
+    log_say(f"Replaying episode {EPISODE_IDX}")
+    for idx in range(len(episode_frames)):
+        t0 = time.perf_counter()
 
-    # Get robot observation
-    robot_obs = robot.get_observation()
+        # Get recorded action from dataset
+        ee_action = {
+            name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        }
 
-    # Dataset EE -> robot joints
-    joint_action = robot_ee_to_joints_processor((ee_action, robot_obs))
+        # Get robot observation
+        robot_obs = robot.get_observation()
 
-    # Send action to robot
-    _ = robot.send_action(joint_action)
+        # Dataset EE -> robot joints
+        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
-robot.disconnect()
+        precise_sleep(1.0 / dataset.fps - (time.perf_counter() - t0))
+
+    # Clean up
+    robot.disconnect()
+
+
+if __name__ == "__main__":
+    main()

examples/so100_to_so100_EE/teleoperate.py

@@ -32,90 +32,96 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
 from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig
 from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 
 FPS = 30
 
-# Initialize the robot and teleoperator config
-follower_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
-)
-leader_config = SO100LeaderConfig(port="/dev/tty.usbmodem5A460819811", id="my_awesome_leader_arm")
 
-# Initialize the robot and teleoperator
-follower = SO100Follower(follower_config)
-leader = SO100Leader(leader_config)
+def main():
+    # 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")
 
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-follower_kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(follower.bus.motors.keys()),
-)
+    # Initialize the robot and teleoperator
+    follower = SO100Follower(follower_config)
+    leader = SO100Leader(leader_config)
 
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-leader_kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(leader.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
+    follower_kinematics_solver = RobotKinematics(
+        urdf_path="./SO101/so101_new_calib.urdf",
+        target_frame_name="gripper_frame_link",
+        joint_names=list(follower.bus.motors.keys()),
+    )
 
-# Build pipeline to convert teleop joints to EE action
-leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
-    steps=[
-        ForwardKinematicsJointsToEE(
-            kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
-        ),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_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_kinematics_solver = RobotKinematics(
+        urdf_path="./SO101/so101_new_calib.urdf",
+        target_frame_name="gripper_frame_link",
+        joint_names=list(leader.bus.motors.keys()),
+    )
 
-# build pipeline to convert EE action to robot 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=False,
-        ),
-    ],
-    to_transition=robot_action_observation_to_transition,
-    to_output=transition_to_robot_action,
-)
+    # Build pipeline to convert teleop joints to EE action
+    leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
+        steps=[
+            ForwardKinematicsJointsToEE(
+                kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
+            ),
+        ],
+        to_transition=robot_action_to_transition,
+        to_output=transition_to_robot_action,
+    )
 
-# Connect to the robot and teleoperator
-follower.connect()
-leader.connect()
+    # build pipeline to convert EE action to robot 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=False,
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
 
-# Init rerun viewer
-init_rerun(session_name="so100_so100_EE_teleop")
+    # Connect to the robot and teleoperator
+    follower.connect()
+    leader.connect()
 
-print("Starting teleop loop...")
-while True:
-    t0 = time.perf_counter()
+    # Init rerun viewer
+    init_rerun(session_name="so100_so100_EE_teleop")
 
-    # Get robot observation
-    robot_obs = follower.get_observation()
+    print("Starting teleop loop...")
+    while True:
+        t0 = time.perf_counter()
 
-    # Get teleop observation
-    leader_joints_obs = leader.get_action()
+        # Get robot observation
+        robot_obs = follower.get_observation()
 
-    # teleop joints -> teleop EE action
-    leader_ee_act = leader_to_ee(leader_joints_obs)
+        # Get teleop observation
+        leader_joints_obs = leader.get_action()
 
-    # teleop EE -> robot joints
-    follower_joints_act = ee_to_follower_joints((leader_ee_act, robot_obs))
+        # teleop joints -> teleop EE action
+        leader_ee_act = leader_to_ee(leader_joints_obs)
 
-    # Send action to robot
-    _ = follower.send_action(follower_joints_act)
+        # teleop EE -> robot joints
+        follower_joints_act = ee_to_follower_joints((leader_ee_act, robot_obs))
 
-    # Visualize
-    log_rerun_data(observation=leader_ee_act, action=follower_joints_act)
+        # Send action to robot
+        _ = 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)
+
+        precise_sleep(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")
-output_directory.mkdir(parents=True, exist_ok=True)
+def main():
+    output_directory = Path("outputs/robot_learning_tutorial/act")
+    output_directory.mkdir(parents=True, exist_ok=True)
 
-# Select your device
-device = torch.device("mps")  # or "cuda" or "cpu"
+    # Select your device
+    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
-dataset_metadata = LeRobotDatasetMetadata(dataset_id)
-features = dataset_to_policy_features(dataset_metadata.features)
+    # This specifies the inputs the model will be expecting and the outputs it will produce
+    dataset_metadata = LeRobotDatasetMetadata(dataset_id)
+    features = dataset_to_policy_features(dataset_metadata.features)
 
-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}
+    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}
 
-cfg = ACTConfig(input_features=input_features, output_features=output_features)
-policy = ACTPolicy(cfg)
-preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
+    cfg = ACTConfig(input_features=input_features, output_features=output_features)
+    policy = ACTPolicy(cfg)
+    preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
 
-policy.train()
-policy.to(device)
+    policy.train()
+    policy.to(device)
 
-# To perform action chunking, ACT expects a given number of actions as targets
-delta_timestamps = {
-    "action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps),
-}
+    # To perform action chunking, ACT expects a given number of actions as targets
+    delta_timestamps = {
+        "action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps),
+    }
 
-# add image features if they are present
-delta_timestamps |= {
-    k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps) for k in cfg.image_features
-}
+    # add image features if they are present
+    delta_timestamps |= {
+        k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps)
+        for k in cfg.image_features
+    }
 
-# Instantiate the dataset
-dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps)
+    # Instantiate the dataset
+    dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps)
 
-# Create the optimizer and dataloader for offline training
-optimizer = cfg.get_optimizer_preset().build(policy.parameters())
-batch_size = 32
-dataloader = torch.utils.data.DataLoader(
-    dataset,
-    batch_size=batch_size,
-    shuffle=True,
-    pin_memory=device.type != "cpu",
-    drop_last=True,
-)
+    # Create the optimizer and dataloader for offline training
+    optimizer = cfg.get_optimizer_preset().build(policy.parameters())
+    batch_size = 32
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        pin_memory=device.type != "cpu",
+        drop_last=True,
+    )
 
-# Number of training steps and logging frequency
-training_steps = 1
-log_freq = 1
+    # Number of training steps and logging frequency
+    training_steps = 1
+    log_freq = 1
 
-# Run training loop
-step = 0
-done = False
-while not done:
-    for batch in dataloader:
-        batch = preprocessor(batch)
-        loss, _ = policy.forward(batch)
-        loss.backward()
-        optimizer.step()
-        optimizer.zero_grad()
+    # Run training loop
+    step = 0
+    done = False
+    while not done:
+        for batch in dataloader:
+            batch = preprocessor(batch)
+            loss, _ = policy.forward(batch)
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
 
-        if step % log_freq == 0:
-            print(f"step: {step} loss: {loss.item():.3f}")
-        step += 1
-        if step >= training_steps:
-            done = True
-            break
+            if step % log_freq == 0:
+                print(f"step: {step} loss: {loss.item():.3f}")
+            step += 1
+            if step >= training_steps:
+                done = True
+                break
 
-# Save the policy checkpoint, alongside the pre/post processors
-policy.save_pretrained(output_directory)
-preprocessor.save_pretrained(output_directory)
-postprocessor.save_pretrained(output_directory)
+    # Save the policy checkpoint, alongside the pre/post processors
+    policy.save_pretrained(output_directory)
+    preprocessor.save_pretrained(output_directory)
+    postprocessor.save_pretrained(output_directory)
 
-# Save all assets to the Hub
-policy.push_to_hub("fracapuano/robot_learning_tutorial_act")
-preprocessor.push_to_hub("fracapuano/robot_learning_tutorial_act")
-postprocessor.push_to_hub("fracapuano/robot_learning_tutorial_act")
+    # Save all assets to the Hub
+    policy.push_to_hub("<user>/robot_learning_tutorial_act")
+    preprocessor.push_to_hub("<user>/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)
-robot = SO100Follower(robot_cfg)
-robot.connect()
+def main():
+    device = torch.device("mps")  # or "cuda" or "cpu"
+    model_id = "<user>/robot_learning_tutorial_act"
+    model = ACTPolicy.from_pretrained(model_id)
 
-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
-        )
+    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)
 
-        obs = preprocess(obs_frame)
+    # # find ports using lerobot-find-port
+    follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
 
-        action = model.select_action(obs)
-        action = postprocess(action)
+    # # the robot ids are used the load the right calibration files
+    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(
-    host=host,
-    port=port,
-)
-serve(config)
+def main():
+    host = ...  # something like "127.0.0.1" if you're exposing to localhost
+    port = ...  # something like 8080
+
+    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
-follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
+def main():
+    # 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
-follower_id = ...  # something like "follower_so100"
+    # # find ports using lerobot-find-port
+    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
-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
-)
+    server_address = ...  # something like "127.0.0.1:8080" if using localhost
 
-# 4. Create and start client
-client = RobotClient(client_cfg)
+    # 3. Create client configuration
+    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
-task = ...
+    # 4. Create and start client
+    client = RobotClient(client_cfg)
 
-if client.start():
-    # Start action receiver thread
-    action_receiver_thread = threading.Thread(target=client.receive_actions, daemon=True)
-    action_receiver_thread.start()
+    # 5. Provide a textual description of the task
+    task = ...
 
-    try:
-        # Run the control loop
-        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 client.start():
+        # Start action receiver thread
+        action_receiver_thread = threading.Thread(target=client.receive_actions, daemon=True)
+        action_receiver_thread.start()
+
+        try:
+            # Run the control loop
+            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")
-output_directory.mkdir(parents=True, exist_ok=True)
+def main():
+    output_directory = Path("outputs/robot_learning_tutorial/diffusion")
+    output_directory.mkdir(parents=True, exist_ok=True)
 
-# Select your device
-device = torch.device("mps")  # or "cuda" or "cpu"
+    # Select your device
+    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
-dataset_metadata = LeRobotDatasetMetadata(dataset_id)
-features = dataset_to_policy_features(dataset_metadata.features)
+    # This specifies the inputs the model will be expecting and the outputs it will produce
+    dataset_metadata = LeRobotDatasetMetadata(dataset_id)
+    features = dataset_to_policy_features(dataset_metadata.features)
 
-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}
+    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}
 
-cfg = DiffusionConfig(input_features=input_features, output_features=output_features)
-policy = DiffusionPolicy(cfg)
-preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
+    cfg = DiffusionConfig(input_features=input_features, output_features=output_features)
+    policy = DiffusionPolicy(cfg)
+    preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
 
-policy.train()
-policy.to(device)
+    policy.train()
+    policy.to(device)
 
-# To perform action chunking, ACT expects a given number of actions as targets
-delta_timestamps = {
-    "observation.state": make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps),
-    "action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps),
-}
+    # To perform action chunking, ACT expects a given number of actions as targets
+    delta_timestamps = {
+        "observation.state": make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps),
+        "action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps),
+    }
 
-# add image features if they are present
-delta_timestamps |= {
-    k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps) for k in cfg.image_features
-}
+    # add image features if they are present
+    delta_timestamps |= {
+        k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps)
+        for k in cfg.image_features
+    }
 
-# Instantiate the dataset
-dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps)
+    # Instantiate the dataset
+    dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps)
 
-# Create the optimizer and dataloader for offline training
-optimizer = cfg.get_optimizer_preset().build(policy.parameters())
-batch_size = 32
-dataloader = torch.utils.data.DataLoader(
-    dataset,
-    batch_size=batch_size,
-    shuffle=True,
-    pin_memory=device.type != "cpu",
-    drop_last=True,
-)
+    # Create the optimizer and dataloader for offline training
+    optimizer = cfg.get_optimizer_preset().build(policy.parameters())
+    batch_size = 32
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        pin_memory=device.type != "cpu",
+        drop_last=True,
+    )
 
-# Number of training steps and logging frequency
-training_steps = 1
-log_freq = 1
+    # Number of training steps and logging frequency
+    training_steps = 1
+    log_freq = 1
 
-# Run training loop
-step = 0
-done = False
-while not done:
-    for batch in dataloader:
-        batch = preprocessor(batch)
-        loss, _ = policy.forward(batch)
-        loss.backward()
-        optimizer.step()
-        optimizer.zero_grad()
+    # Run training loop
+    step = 0
+    done = False
+    while not done:
+        for batch in dataloader:
+            batch = preprocessor(batch)
+            loss, _ = policy.forward(batch)
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
 
-        if step % log_freq == 0:
-            print(f"step: {step} loss: {loss.item():.3f}")
-        step += 1
-        if step >= training_steps:
-            done = True
-            break
+            if step % log_freq == 0:
+                print(f"step: {step} loss: {loss.item():.3f}")
+            step += 1
+            if step >= training_steps:
+                done = True
+                break
 
-# Save the policy checkpoint, alongside the pre/post processors
-policy.save_pretrained(output_directory)
-preprocessor.save_pretrained(output_directory)
-postprocessor.save_pretrained(output_directory)
+    # Save the policy checkpoint, alongside the pre/post processors
+    policy.save_pretrained(output_directory)
+    preprocessor.save_pretrained(output_directory)
+    postprocessor.save_pretrained(output_directory)
 
-# Save all assets to the Hub
-policy.push_to_hub("fracapuano/robot_learning_tutorial_diffusion")
-preprocessor.push_to_hub("fracapuano/robot_learning_tutorial_diffusion")
-postprocessor.push_to_hub("fracapuano/robot_learning_tutorial_diffusion")
+    # Save all assets to the Hub
+    policy.push_to_hub("<user>/robot_learning_tutorial_diffusion")
+    preprocessor.push_to_hub("<user>/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
-follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
+def main():
+    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
-follower_id = ...  # something like "follower_so100"
+    model = DiffusionPolicy.from_pretrained(model_id)
 
-# 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),
-}
+    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
+    )
 
-robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
-robot = SO100Follower(robot_cfg)
-robot.connect()
+    # # find ports using lerobot-find-port
+    follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
+
+    # # the robot ids are used the load the right calibration files
+    follower_id = ...  # something like "follower_so100"
+
+    # Robot and environment configuration
+    # Camera keys must match the name and resolutions of the ones used for training!
+    # You can check the camera keys expected by a model in the info.json card on the model card on the Hub
+    camera_config = {
+        "side": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
+        "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):
-    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/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.config,
-    model_id,
-    # This overrides allows to run on MPS, otherwise defaults to CUDA (if available)
-    preprocessor_overrides={"device_processor": {"device": str(device)}},
-)
+    model = PI0Policy.from_pretrained(model_id)
 
-# find ports using lerobot-find-port
-follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
+    preprocess, postprocess = make_pre_post_processors(
+        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
-follower_id = ...  # something like "follower_so100"
+    # find ports using lerobot-find-port
+    follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
 
-# 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 = {
-    "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),
-}
+    # the robot ids are used the load the right calibration files
+    follower_id = ...  # something like "follower_so100"
 
-robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
-robot = SO100Follower(robot_cfg)
-robot.connect()
+    # 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 = {
+        "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_type = ""  # something like "so100_follower" for multi-embodiment datasets
+    robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
+    robot = SO100Follower(robot_cfg)
+    robot.connect()
 
-# This is used to match the raw observation keys to the keys expected by the policy
-action_features = hw_to_dataset_features(robot.action_features, "action")
-obs_features = hw_to_dataset_features(robot.observation_features, "observation")
-dataset_features = {**action_features, **obs_features}
+    task = ""  # something like "pick the red block"
+    robot_type = ""  # something like "so100_follower" for multi-embodiment datasets
 
-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
-        )
+    # This is used to match the raw observation keys to the keys expected by the policy
+    action_features = hw_to_dataset_features(robot.action_features, "action")
+    obs_features = hw_to_dataset_features(robot.observation_features, "observation")
+    dataset_features = {**action_features, **obs_features}
 
-        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)
-        action = postprocess(action)
-        action = make_robot_action(action, dataset_features)
-        robot.send_action(action)
+            obs = preprocess(obs_frame)
 
-    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"
-output_directory = Path("outputs/robot_learning_tutorial/hil_serl")
-output_directory.mkdir(parents=True, exist_ok=True)
+    device = "mps"  # or "cuda" or "cpu"
+    output_directory = Path("outputs/robot_learning_tutorial/hil_serl")
+    output_directory.mkdir(parents=True, exist_ok=True)
 
-# find ports using lerobot-find-port
-follower_port = ...
-leader_port = ...
+    # find ports using lerobot-find-port
+    follower_port = ...
+    leader_port = ...
 
-# the robot ids are used the load the right calibration files
-follower_id = ...
-leader_id = ...
+    # the robot ids are used the load the right calibration files
+    follower_id = ...
+    leader_id = ...
 
-# A pretrained model (to be used in-distribution!)
-reward_classifier_id = "fracapuano/reward_classifier_hil_serl_example"
-reward_classifier = Classifier.from_pretrained(reward_classifier_id)
+    # A pretrained model (to be used in-distribution!)
+    reward_classifier_id = "<user>/reward_classifier_hil_serl_example"
+    reward_classifier = Classifier.from_pretrained(reward_classifier_id)
 
-reward_classifier.to(device)
-reward_classifier.eval()
+    reward_classifier.to(device)
+    reward_classifier.eval()
 
-MAX_EPISODES = 5
-MAX_STEPS_PER_EPISODE = 20
+    # Robot and environment configuration
+    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
-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)
 
-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
-env, teleop_device = make_robot_env(env_cfg)
+    obs_features = hw_to_dataset_features(env.robot.observation_features, "observation")
+    action_features = hw_to_dataset_features(env.robot.action_features, "action")
 
-obs_features = hw_to_dataset_features(env.robot.observation_features, "observation")
-action_features = hw_to_dataset_features(env.robot.action_features, "action")
+    # Create SAC policy for action selection
+    policy_cfg = SACConfig(
+        device=device,
+        input_features=obs_features,
+        output_features=action_features,
+    )
 
-# Create SAC policy for action selection
-policy_cfg = SACConfig(
-    device=device,
-    input_features=obs_features,
-    output_features=action_features,
-)
+    policy_actor = SACPolicy(policy_cfg)
+    policy_learner = SACPolicy(policy_cfg)
 
-policy_actor = SACPolicy(policy_cfg)
-policy_learner = SACPolicy(policy_cfg)
+    demonstrations_repo_id = "lerobot/example_hil_serl_dataset"
+    offline_dataset = LeRobotDataset(repo_id=demonstrations_repo_id)
 
-demonstrations_repo_id = "lerobot/example_hil_serl_dataset"
-offline_dataset = LeRobotDataset(repo_id=demonstrations_repo_id)
+    # Online buffer: initialized from scratch
+    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
-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())
-)
+    # Create communication channels between learner and actor processes
+    transitions_queue = mp.Queue(maxsize=10)
+    parameters_queue = mp.Queue(maxsize=2)
+    shutdown_event = mp.Event()
 
-# Create communication channels between learner and actor processes
-transitions_queue = mp.Queue(maxsize=10)
-parameters_queue = mp.Queue(maxsize=2)
-shutdown_event = mp.Event()
+    # Signal handler for graceful shutdown
+    def signal_handler(sig):
+        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()
 
 
-# Signal handler for graceful shutdown
-def signal_handler(sig):
-    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()
+if __name__ == "__main__":
+    main()

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
-repo_id = "lerobot/example_hil_serl_dataset"
-dataset = LeRobotDataset(repo_id)
+def main():
+    # Device to use for training
+    device = "mps"  # or "cuda", or "cpu"
 
-# Configure the policy to extract features from the image frames
-camera_keys = dataset.meta.camera_keys
+    # Load the dataset used for training
+    repo_id = "lerobot/example_hil_serl_dataset"
+    dataset = LeRobotDataset(repo_id)
 
-config = RewardClassifierConfig(
-    num_cameras=len(camera_keys),
-    device=device,
-    # backbone model to extract features from the image frames
-    model_name="microsoft/resnet-18",
-)
+    # Configure the policy to extract features from the image frames
+    camera_keys = dataset.meta.camera_keys
 
-# 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)
+    config = RewardClassifierConfig(
+        num_cameras=len(camera_keys),
+        device=device,
+        # 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"
-
-# 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)
+if __name__ == "__main__":
+    main()

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.config,
-    model_id,
-    # This overrides allows to run on MPS, otherwise defaults to CUDA (if available)
-    preprocessor_overrides={"device_processor": {"device": str(device)}},
-)
+    model = SmolVLAPolicy.from_pretrained(model_id)
 
-# find ports using lerobot-find-port
-follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
+    preprocess, postprocess = make_pre_post_processors(
+        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
-follower_id = ...  # something like "follower_so100"
+    # find ports using lerobot-find-port
+    follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
 
-# 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 = {
-    "camera1": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
-    "camera2": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
-}
+    # the robot ids are used the load the right calibration files
+    follower_id = ...  # something like "follower_so100"
 
-robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
-robot = SO100Follower(robot_cfg)
-robot.connect()
+    # 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 = {
+        "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_type = ""  # something like "so100_follower" for multi-embodiment datasets
+    robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
+    robot = SO100Follower(robot_cfg)
+    robot.connect()
 
-# This is used to match the raw observation keys to the keys expected by the policy
-action_features = hw_to_dataset_features(robot.action_features, "action")
-obs_features = hw_to_dataset_features(robot.observation_features, "observation")
-dataset_features = {**action_features, **obs_features}
+    task = ""  # something like "pick the red block"
+    robot_type = ""  # something like "so100_follower" for multi-embodiment datasets
 
-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
-        )
+    # This is used to match the raw observation keys to the keys expected by the policy
+    action_features = hw_to_dataset_features(robot.action_features, "action")
+    obs_features = hw_to_dataset_features(robot.observation_features, "observation")
+    dataset_features = {**action_features, **obs_features}
 
-        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)
-        action = postprocess(action)
-        action = make_robot_action(action, dataset_features)
-        robot.send_action(action)
+            obs = preprocess(obs_frame)
 
-    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()

src/lerobot/datasets/image_writer.py

@@ -110,8 +110,8 @@ def worker_thread_loop(queue: queue.Queue):
         if item is None:
             queue.task_done()
             break
-        image_array, fpath = item
-        write_image(image_array, fpath)
+        image_array, fpath, compress_level = item
+        write_image(image_array, fpath, compress_level)
         queue.task_done()
 
 
@@ -169,11 +169,13 @@ class AsyncImageWriter:
                 p.start()
                 self.processes.append(p)
 
-    def save_image(self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path):
+    def save_image(
+        self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1
+    ):
         if isinstance(image, torch.Tensor):
             # Convert tensor to numpy array to minimize main process time
             image = image.cpu().numpy()
-        self.queue.put((image, fpath))
+        self.queue.put((image, fpath, compress_level))
 
     def wait_until_done(self):
         self.queue.join()

src/lerobot/datasets/lerobot_dataset.py

@@ -13,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import concurrent.futures
 import contextlib
 import logging
 import shutil
@@ -539,6 +540,15 @@ class LeRobotDatasetMetadata:
         return obj
 
 
+def _encode_video_worker(video_key: str, episode_index: int, root: Path, fps: int) -> Path:
+    temp_path = Path(tempfile.mkdtemp(dir=root)) / f"{video_key}_{episode_index:03d}.mp4"
+    fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=episode_index, frame_index=0)
+    img_dir = (root / fpath).parent
+    encode_video_frames(img_dir, temp_path, fps, overwrite=True)
+    shutil.rmtree(img_dir)
+    return temp_path
+
+
 class LeRobotDataset(torch.utils.data.Dataset):
     def __init__(
         self,
@@ -1071,6 +1081,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
             ep_buffer[key] = current_ep_idx if key == "episode_index" else []
         return ep_buffer
 
+    # TODO(Steven): consider move this to utils
     def _get_image_file_path(self, episode_index: int, image_key: str, frame_index: int) -> Path:
         fpath = DEFAULT_IMAGE_PATH.format(
             image_key=image_key, episode_index=episode_index, frame_index=frame_index
@@ -1080,13 +1091,15 @@ class LeRobotDataset(torch.utils.data.Dataset):
     def _get_image_file_dir(self, episode_index: int, image_key: str) -> Path:
         return self._get_image_file_path(episode_index, image_key, frame_index=0).parent
 
-    def _save_image(self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path) -> None:
+    def _save_image(
+        self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1
+    ) -> None:
         if self.image_writer is None:
             if isinstance(image, torch.Tensor):
                 image = image.cpu().numpy()
-            write_image(image, fpath)
+            write_image(image, fpath, compress_level=compress_level)
         else:
-            self.image_writer.save_image(image=image, fpath=fpath)
+            self.image_writer.save_image(image=image, fpath=fpath, compress_level=compress_level)
 
     def add_frame(self, frame: dict) -> None:
         """
@@ -1124,14 +1137,19 @@ class LeRobotDataset(torch.utils.data.Dataset):
                 )
                 if frame_index == 0:
                     img_path.parent.mkdir(parents=True, exist_ok=True)
-                self._save_image(frame[key], img_path)
+                compress_level = 1 if self.features[key]["dtype"] == "video" else 6
+                self._save_image(frame[key], img_path, compress_level)
                 self.episode_buffer[key].append(str(img_path))
             else:
                 self.episode_buffer[key].append(frame[key])
 
         self.episode_buffer["size"] += 1
 
-    def save_episode(self, episode_data: dict | None = None) -> None:
+    def save_episode(
+        self,
+        episode_data: dict | None = None,
+        parallel_encoding: bool = True,
+    ) -> None:
         """
         This will save to disk the current episode in self.episode_buffer.
 
@@ -1143,6 +1161,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
             episode_data (dict | None, optional): Dict containing the episode data to save. If None, this will
                 save the current episode in self.episode_buffer, which is filled with 'add_frame'. Defaults to
                 None.
+            parallel_encoding (bool, optional): If True, encode videos in parallel using ProcessPoolExecutor.
+                Defaults to True on Linux, False on macOS as it tends to use all the CPU available already.
         """
         episode_buffer = episode_data if episode_data is not None else self.episode_buffer
 
@@ -1179,8 +1199,40 @@ class LeRobotDataset(torch.utils.data.Dataset):
         use_batched_encoding = self.batch_encoding_size > 1
 
         if has_video_keys and not use_batched_encoding:
-            for video_key in self.meta.video_keys:
-                ep_metadata.update(self._save_episode_video(video_key, episode_index))
+            num_cameras = len(self.meta.video_keys)
+            if parallel_encoding and num_cameras > 1:
+                # TODO(Steven): Ideally we would like to control the number of threads per encoding such that:
+                # num_cameras * num_threads = (total_cpu -1)
+                with concurrent.futures.ProcessPoolExecutor(max_workers=num_cameras) as executor:
+                    future_to_key = {
+                        executor.submit(
+                            _encode_video_worker,
+                            video_key,
+                            episode_index,
+                            self.root,
+                            self.fps,
+                        ): video_key
+                        for video_key in self.meta.video_keys
+                    }
+
+                    results = {}
+                    for future in concurrent.futures.as_completed(future_to_key):
+                        video_key = future_to_key[future]
+                        try:
+                            temp_path = future.result()
+                            results[video_key] = temp_path
+                        except Exception as exc:
+                            logging.error(f"Video encoding failed for {video_key}: {exc}")
+                            raise exc
+
+                for video_key in self.meta.video_keys:
+                    temp_path = results[video_key]
+                    ep_metadata.update(
+                        self._save_episode_video(video_key, episode_index, temp_path=temp_path)
+                    )
+            else:
+                for video_key in self.meta.video_keys:
+                    ep_metadata.update(self._save_episode_video(video_key, episode_index))
 
         # `meta.save_episode` need to be executed after encoding the videos
         self.meta.save_episode(episode_index, episode_length, episode_tasks, ep_stats, ep_metadata)
@@ -1345,9 +1397,18 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
         return metadata
 
-    def _save_episode_video(self, video_key: str, episode_index: int) -> dict:
+    def _save_episode_video(
+        self,
+        video_key: str,
+        episode_index: int,
+        temp_path: Path | None = None,
+    ) -> dict:
         # Encode episode frames into a temporary video
-        ep_path = self._encode_temporary_episode_video(video_key, episode_index)
+        if temp_path is None:
+            ep_path = self._encode_temporary_episode_video(video_key, episode_index)
+        else:
+            ep_path = temp_path
+
         ep_size_in_mb = get_file_size_in_mb(ep_path)
         ep_duration_in_s = get_video_duration_in_s(ep_path)
 
@@ -1465,11 +1526,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         Note: `encode_video_frames` is a blocking call. Making it asynchronous shouldn't speedup encoding,
         since video encoding with ffmpeg is already using multithreading.
         """
-        temp_path = Path(tempfile.mkdtemp(dir=self.root)) / f"{video_key}_{episode_index:03d}.mp4"
-        img_dir = self._get_image_file_dir(episode_index, video_key)
-        encode_video_frames(img_dir, temp_path, self.fps, overwrite=True)
-        shutil.rmtree(img_dir)
-        return temp_path
+        return _encode_video_worker(video_key, episode_index, self.root, self.fps)
 
     @classmethod
     def create(

src/lerobot/datasets/utils.py

@@ -49,7 +49,7 @@ from lerobot.utils.utils import SuppressProgressBars, is_valid_numpy_dtype_strin
 
 DEFAULT_CHUNK_SIZE = 1000  # Max number of files per chunk
 DEFAULT_DATA_FILE_SIZE_IN_MB = 100  # Max size per file
-DEFAULT_VIDEO_FILE_SIZE_IN_MB = 500  # Max size per file
+DEFAULT_VIDEO_FILE_SIZE_IN_MB = 200  # Max size per file
 
 INFO_PATH = "meta/info.json"
 STATS_PATH = "meta/stats.json"

src/lerobot/datasets/video_utils.py

@@ -311,6 +311,7 @@ def encode_video_frames(
     fast_decode: int = 0,
     log_level: int | None = av.logging.ERROR,
     overwrite: bool = False,
+    preset: int | None = None,
 ) -> None:
     """More info on ffmpeg arguments tuning on `benchmark/video/README.md`"""
     # Check encoder availability
@@ -359,6 +360,9 @@ def encode_video_frames(
         value = f"fast-decode={fast_decode}" if vcodec == "libsvtav1" else "fastdecode"
         video_options[key] = value
 
+    if vcodec == "libsvtav1":
+        video_options["preset"] = str(preset) if preset is not None else "12"
+
     # Set logging level
     if log_level is not None:
         # "While less efficient, it is generally preferable to modify logging with Python's logging"

src/lerobot/rl/actor.py

@@ -78,7 +78,7 @@ from lerobot.transport.utils import (
     transitions_to_bytes,
 )
 from lerobot.utils.random_utils import set_seed
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.transition import (
     Transition,
     move_state_dict_to_device,
@@ -398,7 +398,7 @@ def act_with_policy(
 
         if cfg.env.fps is not None:
             dt_time = time.perf_counter() - start_time
-            busy_wait(1 / cfg.env.fps - dt_time)
+            precise_sleep(1 / cfg.env.fps - dt_time)
 
 
 #  Communication Functions - Group all gRPC/messaging functions

src/lerobot/rl/gym_manipulator.py

@@ -74,7 +74,7 @@ from lerobot.teleoperators import (
 from lerobot.teleoperators.teleoperator import Teleoperator
 from lerobot.teleoperators.utils import TeleopEvents
 from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, REWARD
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say
 
 logging.basicConfig(level=logging.INFO)
@@ -114,7 +114,7 @@ def reset_follower_position(robot_arm: Robot, target_position: np.ndarray) -> No
     for pose in trajectory:
         action_dict = dict(zip(current_position_dict, pose, strict=False))
         robot_arm.bus.sync_write("Goal_Position", action_dict)
-        busy_wait(0.015)
+        precise_sleep(0.015)
 
 
 class RobotEnv(gym.Env):
@@ -238,7 +238,7 @@ class RobotEnv(gym.Env):
             reset_follower_position(self.robot, np.array(self.reset_pose))
             log_say("Reset the environment done.", play_sounds=True)
 
-        busy_wait(self.reset_time_s - (time.perf_counter() - start_time))
+        precise_sleep(self.reset_time_s - (time.perf_counter() - start_time))
 
         super().reset(seed=seed, options=options)
 
@@ -713,7 +713,7 @@ def control_loop(
             transition = env_processor(transition)
 
         # Maintain fps timing
-        busy_wait(dt - (time.perf_counter() - step_start_time))
+        precise_sleep(dt - (time.perf_counter() - step_start_time))
 
     if dataset is not None and cfg.dataset.push_to_hub:
         logging.info("Pushing dataset to hub")
@@ -745,7 +745,7 @@ def replay_trajectory(
         )
         transition = action_processor(transition)
         env.step(transition[TransitionKey.ACTION])
-        busy_wait(1 / cfg.env.fps - (time.perf_counter() - start_time))
+        precise_sleep(1 / cfg.env.fps - (time.perf_counter() - start_time))
 
 
 @parser.wrap()

src/lerobot/scripts/lerobot_find_joint_limits.py

@@ -50,7 +50,7 @@ from lerobot.teleoperators import (  # noqa: F401
     make_teleoperator_from_config,
     so100_leader,
 )
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 
 
 @dataclass
@@ -114,7 +114,7 @@ def find_joint_and_ee_bounds(cfg: FindJointLimitsConfig):
             print(f"Min joint pos position {np.round(min_pos, 4).tolist()}")
             break
 
-        busy_wait(0.01)
+        precise_sleep(0.01)
 
 
 def main():

src/lerobot/scripts/lerobot_record.py

@@ -119,7 +119,7 @@ from lerobot.utils.control_utils import (
     sanity_check_dataset_robot_compatibility,
 )
 from lerobot.utils.import_utils import register_third_party_devices
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import (
     get_safe_torch_device,
     init_logging,
@@ -364,7 +364,7 @@ def record_loop(
             log_rerun_data(observation=obs_processed, action=action_values)
 
         dt_s = time.perf_counter() - start_loop_t
-        busy_wait(1 / fps - dt_s)
+        precise_sleep(1 / fps - dt_s)
 
         timestamp = time.perf_counter() - start_episode_t
 

src/lerobot/scripts/lerobot_replay.py

@@ -62,7 +62,7 @@ from lerobot.robots import (  # noqa: F401
 )
 from lerobot.utils.constants import ACTION
 from lerobot.utils.import_utils import register_third_party_devices
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import (
     init_logging,
     log_say,
@@ -121,7 +121,7 @@ def replay(cfg: ReplayConfig):
         _ = robot.send_action(processed_action)
 
         dt_s = time.perf_counter() - start_episode_t
-        busy_wait(1 / dataset.fps - dt_s)
+        precise_sleep(1 / dataset.fps - dt_s)
 
     robot.disconnect()
 

src/lerobot/scripts/lerobot_teleoperate.py

@@ -89,7 +89,7 @@ from lerobot.teleoperators import (  # noqa: F401
     so101_leader,
 )
 from lerobot.utils.import_utils import register_third_party_devices
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import init_logging, move_cursor_up
 from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 
@@ -170,12 +170,13 @@ def teleop_loop(
             # Display the final robot action that was sent
             for motor, value in robot_action_to_send.items():
                 print(f"{motor:<{display_len}} | {value:>7.2f}")
-            move_cursor_up(len(robot_action_to_send) + 5)
+            move_cursor_up(len(robot_action_to_send) + 3)
 
         dt_s = time.perf_counter() - loop_start
-        busy_wait(1 / fps - dt_s)
+        precise_sleep(1 / fps - dt_s)
         loop_s = time.perf_counter() - loop_start
-        print(f"\ntime: {loop_s * 1e3:.2f}ms ({1 / loop_s:.0f} Hz)")
+        print(f"Teleop loop time: {loop_s * 1e3:.2f}ms ({1 / loop_s:.0f} Hz)")
+        move_cursor_up(1)
 
         if duration is not None and time.perf_counter() - start >= duration:
             return

src/lerobot/utils/robot_utils.py

@@ -16,14 +16,40 @@ import platform
 import time
 
 
-def busy_wait(seconds):
-    if platform.system() == "Darwin" or platform.system() == "Windows":
-        # On Mac and Windows, `time.sleep` is not accurate and we need to use this while loop trick,
-        # but it consumes CPU cycles.
+def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: float = 0.003):
+    """
+    Wait for `seconds` with better precision than time.sleep alone at the expense of more CPU usage.
+
+    Parameters:
+      - seconds: duration to wait
+      - spin_threshold: if remaining <= spin_threshold -> spin; otherwise sleep (seconds). Default 10ms
+      - sleep_margin: when sleeping leave this much time before deadline to avoid oversleep. Default 3ms
+
+    Note:
+        The default parameters are chosen to prioritize timing accuracy over CPU usage for the common 30 FPS use case.
+    """
+    if seconds <= 0:
+        return
+
+    system = platform.system()
+    # On macOS and Windows the scheduler / sleep granularity can make
+    # short sleeps inaccurate. Instead of burning CPU for the whole
+    # duration, sleep for most of the time and spin for the final few
+    # milliseconds to achieve good accuracy with much lower CPU usage.
+    if system in ("Darwin", "Windows"):
         end_time = time.perf_counter() + seconds
-        while time.perf_counter() < end_time:
-            pass
+        while True:
+            remaining = end_time - time.perf_counter()
+            if remaining <= 0:
+                break
+            # If there's more than a couple milliseconds left, sleep most
+            # of the remaining time and leave a small margin for the final spin.
+            if remaining > spin_threshold:
+                # Sleep but avoid sleeping past the end by leaving a small margin.
+                time.sleep(max(remaining - sleep_margin, 0))
+            else:
+                # Final short spin to hit precise timing without long sleeps.
+                pass
     else:
-        # On Linux time.sleep is accurate
-        if seconds > 0:
-            time.sleep(seconds)
+        # On Linux time.sleep is accurate enough for most uses
+        time.sleep(seconds)