Merge branch 'main' into feat/add-xvla

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")
+        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")
+        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):
+    def process_sample(sample: int, lock: Lock):
-        time_benchmark = TimeBenchmark()
+        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,11 +350,14 @@ 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 duet in [
-                    for value in tqdm(values, desc=f"encodings ({key})", leave=False):
+                    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"
@@ -409,9 +404,9 @@ if __name__ == "__main__":
         nargs="*",
         default=[
             "lerobot/pusht_image",
-            "aliberts/aloha_mobile_shrimp_image",
+            "lerobot/aloha_mobile_shrimp_image",
-            "aliberts/paris_street",
+            "lerobot/paris_street",
-            "aliberts/kitchen",
+            "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/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

@@ -139,7 +139,7 @@ from lerobot.teleoperators import (  # noqa: F401
     make_teleoperator_from_config,
     so101_leader,
 )
-from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import init_logging
 from lerobot.envs.factory import make_env
 
@@ -196,7 +196,7 @@ def teleop_loop(teleop: Teleoperator, env: gym.Env, fps: int):
             obs, info = env.reset()
 
         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)")
 

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()
 ```

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/dataset/load_lerobot_dataset.py

@@ -34,6 +34,8 @@ from huggingface_hub import HfApi
 import lerobot
 from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
 
+
+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)
@@ -132,7 +134,6 @@ 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,6 +33,8 @@ TASK_DESCRIPTION = "My task description"
 HF_MODEL_ID = "<hf_username>/<model_repo_id>"
 HF_DATASET_ID = "<hf_username>/<eval_dataset_repo_id>"
 
+
+def main():
     # Create the robot configuration & robot
     robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
 
@@ -136,3 +138,7 @@ listener.stop()
 
     dataset.finalize()
     dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()

examples/lekiwi/record.py

@@ -34,6 +34,8 @@ RESET_TIME_SEC = 10
 TASK_DESCRIPTION = "My task description"
 HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
 
+
+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")
@@ -133,3 +135,7 @@ listener.stop()
 
     dataset.finalize()
     dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()

examples/lekiwi/replay.py

@@ -20,11 +20,13 @@ 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
 
+
+def main():
     # Initialize the robot config
     robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
 
@@ -56,6 +58,10 @@ for idx in range(len(episode_frames)):
         # Send action to robot
         _ = robot.send_action(action)
 
-    busy_wait(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0))
+        precise_sleep(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0))
 
     robot.disconnect()
+
+
+if __name__ == "__main__":
+    main()

examples/lekiwi/teleoperate.py

@@ -19,11 +19,13 @@ 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
 
+
+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")
@@ -69,4 +71,8 @@ while True:
         # Visualize
         log_rerun_data(observation=observation, action=action)
 
-    busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
+        precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
+
+
+if __name__ == "__main__":
+    main()

examples/phone_to_so100/evaluate.py

@@ -52,6 +52,8 @@ TASK_DESCRIPTION = "My task description"
 HF_MODEL_ID = "<hf_username>/<model_repo_id>"
 HF_DATASET_ID = "<hf_username>/<dataset_repo_id>"
 
+
+def main():
     # Create the robot configuration & robot
     camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
     robot_config = SO100FollowerConfig(
@@ -89,7 +91,9 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
     # 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()))
+            ForwardKinematicsJointsToEE(
+                kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())
+            )
         ],
         to_transition=observation_to_transition,
         to_output=transition_to_observation,
@@ -197,3 +201,7 @@ listener.stop()
 
     dataset.finalize()
     dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()

examples/phone_to_so100/record.py

@@ -50,6 +50,8 @@ RESET_TIME_SEC = 30
 TASK_DESCRIPTION = "My task description"
 HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
 
+
+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(
@@ -72,7 +74,9 @@ kinematics_solver = RobotKinematics(
     )
 
     # Build pipeline to convert phone action to EE action
-phone_to_robot_ee_pose_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+    phone_to_robot_ee_pose_processor = RobotProcessorPipeline[
+        tuple[RobotAction, RobotObservation], RobotAction
+    ](
         steps=[
             MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
             EEReferenceAndDelta(
@@ -107,7 +111,9 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
     # 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()))
+            ForwardKinematicsJointsToEE(
+                kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())
+            )
         ],
         to_transition=observation_to_transition,
         to_output=transition_to_observation,
@@ -147,7 +153,6 @@ 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"]:
@@ -203,3 +208,7 @@ listener.stop()
 
     dataset.finalize()
     dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()

examples/phone_to_so100/replay.py

@@ -29,12 +29,14 @@ 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>"
 
+
+def main():
     # Initialize the robot config
     robot_config = SO100FollowerConfig(
         port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
@@ -94,7 +96,11 @@ for idx in range(len(episode_frames)):
         # Send action to robot
         _ = robot.send_action(joint_action)
 
-    busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))
+        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,11 +32,13 @@ 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
 
+
+def main():
     # Initialize the robot and teleoperator
     robot_config = SO100FollowerConfig(
         port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
@@ -55,7 +57,9 @@ kinematics_solver = RobotKinematics(
     )
 
     # Build pipeline to convert phone action to ee pose action to joint action
-phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+    phone_to_robot_joints_processor = RobotProcessorPipeline[
+        tuple[RobotAction, RobotObservation], RobotAction
+    ](
         steps=[
             MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
             EEReferenceAndDelta(
@@ -110,4 +114,8 @@ while True:
         # Visualize
         log_rerun_data(observation=phone_obs, action=joint_action)
 
-    busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
+        precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
+
+
+if __name__ == "__main__":
+    main()

examples/so100_to_so100_EE/evaluate.py

@@ -52,6 +52,8 @@ TASK_DESCRIPTION = "My task description"
 HF_MODEL_ID = "<hf_username>/<model_repo_id>"
 HF_DATASET_ID = "<hf_username>/<dataset_repo_id>"
 
+
+def main():
     # Create the robot configuration & robot
     camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
     robot_config = SO100FollowerConfig(
@@ -89,13 +91,14 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
     # 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()))
+            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,
@@ -198,3 +201,7 @@ listener.stop()
 
     dataset.finalize()
     dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()

examples/so100_to_so100_EE/record.py

@@ -48,10 +48,15 @@ RESET_TIME_SEC = 30
 TASK_DESCRIPTION = "My task description"
 HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
 
+
+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
+        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")
 
@@ -135,7 +140,6 @@ dataset = LeRobotDataset.create(
         image_writer_threads=4,
     )
 
-
     # Connect the robot and teleoperator
     leader.connect()
     follower.connect()
@@ -202,3 +206,7 @@ listener.stop()
 
     dataset.finalize()
     dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()

examples/so100_to_so100_EE/replay.py

@@ -30,12 +30,14 @@ 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>"
 
+
+def main():
     # Initialize the robot config
     robot_config = SO100FollowerConfig(
         port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
@@ -95,7 +97,11 @@ for idx in range(len(episode_frames)):
         # Send action to robot
         _ = robot.send_action(joint_action)
 
-    busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))
+        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,11 +32,13 @@ 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
 
+
+def main():
     # Initialize the robot and teleoperator config
     follower_config = SO100FollowerConfig(
         port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
@@ -118,4 +120,8 @@ while True:
         # Visualize
         log_rerun_data(observation=leader_ee_act, action=follower_joints_act)
 
-    busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
+        precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
+
+
+if __name__ == "__main__":
+    main()

examples/tutorial/act/act_training_example.py

@@ -19,6 +19,7 @@ def make_delta_timestamps(delta_indices: list[int] | None, fps: int) -> list[flo
     return [i / fps for i in delta_indices]
 
 
+def main():
     output_directory = Path("outputs/robot_learning_tutorial/act")
     output_directory.mkdir(parents=True, exist_ok=True)
 
@@ -48,7 +49,8 @@ delta_timestamps = {
 
     # 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
+        k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps)
+        for k in cfg.image_features
     }
 
     # Instantiate the dataset
@@ -93,6 +95,10 @@ 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")
+    policy.push_to_hub("<user>/robot_learning_tutorial_act")
-preprocessor.push_to_hub("fracapuano/robot_learning_tutorial_act")
+    preprocessor.push_to_hub("<user>/robot_learning_tutorial_act")
-postprocessor.push_to_hub("fracapuano/robot_learning_tutorial_act")
+    postprocessor.push_to_hub("<user>/robot_learning_tutorial_act")
+
+
+if __name__ == "__main__":
+    main()

examples/tutorial/act/act_using_example.py

@@ -8,8 +8,13 @@ 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
 
+MAX_EPISODES = 5
+MAX_STEPS_PER_EPISODE = 20
+
+
+def main():
     device = torch.device("mps")  # or "cuda" or "cpu"
-model_id = "fracapuano/robot_learning_tutorial_act"
+    model_id = "<user>/robot_learning_tutorial_act"
     model = ACTPolicy.from_pretrained(model_id)
 
     dataset_id = "lerobot/svla_so101_pickplace"
@@ -23,9 +28,6 @@ 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
@@ -55,3 +57,7 @@ for _ in range(MAX_EPISODES):
             robot.send_action(action)
 
         print("Episode finished! Starting new episode...")
+
+
+if __name__ == "__main__":
+    main()

examples/tutorial/async-inf/policy_server.py

@@ -1,6 +1,8 @@
 from lerobot.async_inference.configs import PolicyServerConfig
 from lerobot.async_inference.policy_server import serve
 
+
+def main():
     host = ...  # something like "127.0.0.1" if you're exposing to localhost
     port = ...  # something like 8080
 
@@ -9,3 +11,7 @@ config = PolicyServerConfig(
         port=port,
     )
     serve(config)
+
+
+if __name__ == "__main__":
+    main()

examples/tutorial/async-inf/robot_client.py

@@ -6,6 +6,8 @@ from lerobot.async_inference.robot_client import RobotClient
 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
 from lerobot.robots.so100_follower import SO100FollowerConfig
 
+
+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 = {
@@ -29,7 +31,7 @@ client_cfg = RobotClientConfig(
         server_address=server_address,
         policy_device="mps",
         policy_type="act",
-    pretrained_name_or_path="fracapuano/robot_learning_tutorial_act",
+        pretrained_name_or_path="<user>/robot_learning_tutorial_act",
         chunk_size_threshold=0.5,  # g
         actions_per_chunk=50,  # make sure this is less than the max actions of the policy
     )
@@ -53,3 +55,7 @@ if client.start():
             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,6 +19,7 @@ def make_delta_timestamps(delta_indices: list[int] | None, fps: int) -> list[flo
     return [i / fps for i in delta_indices]
 
 
+def main():
     output_directory = Path("outputs/robot_learning_tutorial/diffusion")
     output_directory.mkdir(parents=True, exist_ok=True)
 
@@ -49,7 +50,8 @@ delta_timestamps = {
 
     # 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
+        k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps)
+        for k in cfg.image_features
     }
 
     # Instantiate the dataset
@@ -94,6 +96,10 @@ 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")
+    policy.push_to_hub("<user>/robot_learning_tutorial_diffusion")
-preprocessor.push_to_hub("fracapuano/robot_learning_tutorial_diffusion")
+    preprocessor.push_to_hub("<user>/robot_learning_tutorial_diffusion")
-postprocessor.push_to_hub("fracapuano/robot_learning_tutorial_diffusion")
+    postprocessor.push_to_hub("<user>/robot_learning_tutorial_diffusion")
+
+
+if __name__ == "__main__":
+    main()

examples/tutorial/diffusion/diffusion_using_example.py

@@ -8,8 +8,13 @@ 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
 
+MAX_EPISODES = 5
+MAX_STEPS_PER_EPISODE = 20
+
+
+def main():
     device = torch.device("mps")  # or "cuda" or "cpu"
-model_id = "fracapuano/robot_learning_tutorial_diffusion"
+    model_id = "<user>/robot_learning_tutorial_diffusion"
 
     model = DiffusionPolicy.from_pretrained(model_id)
 
@@ -20,10 +25,6 @@ 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"
 
@@ -42,7 +43,6 @@ robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=came
     robot = SO100Follower(robot_cfg)
     robot.connect()
 
-
     for _ in range(MAX_EPISODES):
         for _ in range(MAX_STEPS_PER_EPISODE):
             obs = robot.get_observation()
@@ -58,3 +58,7 @@ for _ in range(MAX_EPISODES):
             robot.send_action(action)
 
         print("Episode finished! Starting new episode...")
+
+
+if __name__ == "__main__":
+    main()

examples/tutorial/pi0/using_pi0_example.py

@@ -11,6 +11,8 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower
 MAX_EPISODES = 5
 MAX_STEPS_PER_EPISODE = 20
 
+
+def main():
     device = torch.device("mps")  # or "cuda" or "cpu"
     model_id = "lerobot/pi0_base"
 
@@ -65,3 +67,7 @@ for _ in range(MAX_EPISODES):
             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,6 +225,7 @@ def make_policy_obs(obs, device: torch.device = "cpu"):
     }
 
 
+def main():
     """Main function - coordinates actor and learner processes."""
 
     device = "mps"  # or "cuda" or "cpu"
@@ -238,15 +241,12 @@ follower_id = ...
     leader_id = ...
 
     # A pretrained model (to be used in-distribution!)
-reward_classifier_id = "fracapuano/reward_classifier_hil_serl_example"
+    reward_classifier_id = "<user>/reward_classifier_hil_serl_example"
     reward_classifier = Classifier.from_pretrained(reward_classifier_id)
 
     reward_classifier.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)
@@ -285,13 +285,11 @@ 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)
 
@@ -343,3 +341,7 @@ finally:
             learner_process.terminate()
         if actor_process.is_alive():
             actor_process.terminate()
+
+
+if __name__ == "__main__":
+    main()

examples/tutorial/rl/reward_classifier_example.py

@@ -4,6 +4,8 @@ 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
 
+
+def main():
     # Device to use for training
     device = "mps"  # or "cuda", or "cpu"
 
@@ -26,8 +28,7 @@ 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"
-classifier_id = "fracapuano/reward_classifier_hil_serl_example"
 
     # Instantiate a dataloader
     dataloader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True)
@@ -60,3 +61,7 @@ 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,6 +11,8 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower
 MAX_EPISODES = 5
 MAX_STEPS_PER_EPISODE = 20
 
+
+def main():
     device = torch.device("mps")  # or "cuda" or "cpu"
     model_id = "lerobot/smolvla_base"
 
@@ -64,3 +66,7 @@ for _ in range(MAX_EPISODES):
             robot.send_action(action)
 
         print("Episode finished! Starting new episode...")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -25,7 +25,7 @@ discord = "https://discord.gg/s3KuuzsPFb"
 
 [project]
 name = "lerobot"
-version = "0.4.2"
+version = "0.4.3"
 description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch"
 readme = "README.md"
 license = { text = "Apache-2.0" }

src/lerobot/datasets/lerobot_dataset.py

@@ -16,7 +16,6 @@
 import concurrent.futures
 import contextlib
 import logging
-import platform
 import shutil
 import tempfile
 from collections.abc import Callable
@@ -1149,7 +1148,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
     def save_episode(
         self,
         episode_data: dict | None = None,
-        parallel_encoding: bool = platform.system() == "Linux",
+        parallel_encoding: bool = True,
     ) -> None:
         """
         This will save to disk the current episode in self.episode_buffer.

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):
+def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: float = 0.003):
-    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,
+    Wait for `seconds` with better precision than time.sleep alone at the expense of more CPU usage.
-        # but it consumes CPU cycles.
+
+    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:
+        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
+        # On Linux time.sleep is accurate enough for most uses
-        if seconds > 0:
         time.sleep(seconds)