Remove register

.github/workflows/test.yml

@@ -126,7 +126,7 @@ jobs:
       # portaudio19-dev is needed to install pyaudio
         run: |
           sudo apt-get update && \
-          sudo apt-get install -y libegl1-mesa-dev ffmpeg portaudio19-dev
+          sudo apt-get install -y libegl1-mesa-dev portaudio19-dev
 
       - name: Install uv and python
         uses: astral-sh/setup-uv@v5

.pre-commit-config.yaml

@@ -16,13 +16,6 @@ exclude: ^(tests/data)
 default_language_version:
     python: python3.10
 repos:
-  ##### Meta #####
-  - repo: meta
-    hooks:
-      - id: check-useless-excludes
-      - id: check-hooks-apply
-
-
   ##### Style / Misc. #####
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v5.0.0
@@ -35,37 +28,31 @@ repos:
       - id: check-toml
       - id: end-of-file-fixer
       - id: trailing-whitespace
-
   - repo: https://github.com/crate-ci/typos
-    rev: v1.30.2
+    rev: v1.30.0
     hooks:
       - id: typos
         args: [--force-exclude]
-
   - repo: https://github.com/asottile/pyupgrade
     rev: v3.19.1
     hooks:
     -   id: pyupgrade
-
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.9.10
+    rev: v0.9.9
     hooks:
       - id: ruff
         args: [--fix]
       - id: ruff-format
 
-
   ##### Security #####
   - repo: https://github.com/gitleaks/gitleaks
     rev: v8.24.0
     hooks:
       - id: gitleaks
-
   - repo: https://github.com/woodruffw/zizmor-pre-commit
     rev: v1.4.1
     hooks:
       - id: zizmor
-
   - repo: https://github.com/PyCQA/bandit
     rev: 1.8.3
     hooks:

README.md

@@ -232,8 +232,8 @@ python lerobot/scripts/eval.py \
     --env.type=pusht \
     --eval.batch_size=10 \
     --eval.n_episodes=10 \
-    --policy.use_amp=false \
-    --policy.device=cuda
+    --use_amp=false \
+    --device=cuda
 ```
 
 Note: After training your own policy, you can re-evaluate the checkpoints with:

benchmarks/video/run_video_benchmark.py

@@ -67,7 +67,7 @@ def parse_int_or_none(value) -> int | None:
 def check_datasets_formats(repo_ids: list) -> None:
     for repo_id in repo_ids:
         dataset = LeRobotDataset(repo_id)
-        if len(dataset.meta.video_keys) > 0:
+        if dataset.video:
             raise ValueError(
                 f"Use only image dataset for running this benchmark. Video dataset provided: {repo_id}"
             )

examples/10_use_so100.md

@@ -454,8 +454,8 @@ Next, you'll need to calibrate your SO-100 robot to ensure that the leader and f
 
 You will need to move the follower arm to these positions sequentially:
 
-| 1. Zero position                                                                                                                                             | 2. Rotated position                                                                                                                                                   | 3. Rest position                                                                                                                                             |
-| ------------------------------------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| 1. Zero position | 2. Rotated position | 3. Rest position |
+|---|---|---|
 | <img src="../media/so100/follower_zero.webp?raw=true" alt="SO-100 follower arm zero position" title="SO-100 follower arm zero position" style="width:100%;"> | <img src="../media/so100/follower_rotated.webp?raw=true" alt="SO-100 follower arm rotated position" title="SO-100 follower arm rotated position" style="width:100%;"> | <img src="../media/so100/follower_rest.webp?raw=true" alt="SO-100 follower arm rest position" title="SO-100 follower arm rest position" style="width:100%;"> |
 
 Make sure both arms are connected and run this script to launch manual calibration:
@@ -470,8 +470,8 @@ python lerobot/scripts/control_robot.py \
 #### b. Manual calibration of leader arm
 Follow step 6 of the [assembly video](https://youtu.be/FioA2oeFZ5I?t=724) which illustrates the manual calibration. You will need to move the leader arm to these positions sequentially:
 
-| 1. Zero position                                                                                                                                       | 2. Rotated position                                                                                                                                             | 3. Rest position                                                                                                                                       |
-| ------------------------------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| 1. Zero position | 2. Rotated position | 3. Rest position |
+|---|---|---|
 | <img src="../media/so100/leader_zero.webp?raw=true" alt="SO-100 leader arm zero position" title="SO-100 leader arm zero position" style="width:100%;"> | <img src="../media/so100/leader_rotated.webp?raw=true" alt="SO-100 leader arm rotated position" title="SO-100 leader arm rotated position" style="width:100%;"> | <img src="../media/so100/leader_rest.webp?raw=true" alt="SO-100 leader arm rest position" title="SO-100 leader arm rest position" style="width:100%;"> |
 
 Run this script to launch manual calibration:
@@ -571,14 +571,14 @@ python lerobot/scripts/train.py \
   --policy.type=act \
   --output_dir=outputs/train/act_so100_test \
   --job_name=act_so100_test \
-  --policy.device=cuda \
+  --device=cuda \
   --wandb.enable=true
 ```
 
 Let's explain it:
 1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/so100_test`.
 2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../lerobot/common/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor sates, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
-4. 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 `device=cuda` since we are training on a Nvidia GPU, but you could use `device=mps` to train on Apple silicon.
 5. 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. You will find checkpoints in `outputs/train/act_so100_test/checkpoints`.

examples/11_use_lekiwi.md

@@ -366,8 +366,8 @@ Now we have to calibrate the leader arm and the follower arm. The wheel motors d
 
 You will need to move the follower arm to these positions sequentially:
 
-| 1. Zero position                                                                                                                                                  | 2. Rotated position                                                                                                                                                        | 3. Rest position                                                                                                                                                  |
-| ----------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 1. Zero position | 2. Rotated position | 3. Rest position |
+|---|---|---|
 | <img src="../media/lekiwi/mobile_calib_zero.webp?raw=true" alt="SO-100 follower arm zero position" title="SO-100 follower arm zero position" style="width:100%;"> | <img src="../media/lekiwi/mobile_calib_rotated.webp?raw=true" alt="SO-100 follower arm rotated position" title="SO-100 follower arm rotated position" style="width:100%;"> | <img src="../media/lekiwi/mobile_calib_rest.webp?raw=true" alt="SO-100 follower arm rest position" title="SO-100 follower arm rest position" style="width:100%;"> |
 
 Make sure the arm is connected to the Raspberry Pi and run this script (on the Raspberry Pi) to launch manual calibration:
@@ -385,8 +385,8 @@ If you have the **wired** LeKiwi version please run all commands including this
 ### Calibrate leader arm
 Then to calibrate the leader arm (which is attached to the laptop/pc). You will need to move the leader arm to these positions sequentially:
 
-| 1. Zero position                                                                                                                                       | 2. Rotated position                                                                                                                                             | 3. Rest position                                                                                                                                       |
-| ------------------------------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| 1. Zero position | 2. Rotated position | 3. Rest position |
+|---|---|---|
 | <img src="../media/so100/leader_zero.webp?raw=true" alt="SO-100 leader arm zero position" title="SO-100 leader arm zero position" style="width:100%;"> | <img src="../media/so100/leader_rotated.webp?raw=true" alt="SO-100 leader arm rotated position" title="SO-100 leader arm rotated position" style="width:100%;"> | <img src="../media/so100/leader_rest.webp?raw=true" alt="SO-100 leader arm rest position" title="SO-100 leader arm rest position" style="width:100%;"> |
 
 Run this script (on your laptop/pc) to launch manual calibration:
@@ -416,22 +416,22 @@ python lerobot/scripts/control_robot.py \
 
 You should see on your laptop something like this: ```[INFO] Connected to remote robot at tcp://172.17.133.91:5555 and video stream at tcp://172.17.133.91:5556.``` Now you can move the leader arm and use the keyboard (w,a,s,d) to drive forward, left, backwards, right. And use (z,x) to turn left or turn right. You can use (r,f) to increase and decrease the speed of the mobile robot. There are three speed modes, see the table below:
 | Speed Mode | Linear Speed (m/s) | Rotation Speed (deg/s) |
-| ---------- | ------------------ | ---------------------- |
-| Fast       | 0.4                | 90                     |
-| Medium     | 0.25               | 60                     |
-| Slow       | 0.1                | 30                     |
+|------------|-------------------|-----------------------|
+| Fast      | 0.4               | 90                    |
+| Medium    | 0.25              | 60                    |
+| Slow      | 0.1               | 30                    |
 
 
-| Key | Action         |
-| --- | -------------- |
-| W   | Move forward   |
-| A   | Move left      |
-| S   | Move backward  |
-| D   | Move right     |
-| Z   | Turn left      |
-| X   | Turn right     |
-| R   | Increase speed |
-| F   | Decrease speed |
+| Key  | Action                         |
+|------|--------------------------------|
+| W    | Move forward                   |
+| A    | Move left                       |
+| S    | Move backward                   |
+| D    | Move right                      |
+| Z    | Turn left                       |
+| X    | Turn right                      |
+| R    | Increase speed                  |
+| F    | Decrease speed                  |
 
 > [!TIP]
 >  If you use a different keyboard you can change the keys for each command in the [`LeKiwiRobotConfig`](../lerobot/common/robot_devices/robots/configs.py).
@@ -549,14 +549,14 @@ python lerobot/scripts/train.py \
   --policy.type=act \
   --output_dir=outputs/train/act_lekiwi_test \
   --job_name=act_lekiwi_test \
-  --policy.device=cuda \
+  --device=cuda \
   --wandb.enable=true
 ```
 
 Let's explain it:
 1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/lekiwi_test`.
 2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../lerobot/common/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor sates, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
-4. 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 `device=cuda` since we are training on a Nvidia GPU, but you could use `device=mps` to train on Apple silicon.
 5. 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. You will find checkpoints in `outputs/train/act_lekiwi_test/checkpoints`.

examples/11_use_moss.md

@@ -176,8 +176,8 @@ Next, you'll need to calibrate your Moss v1 robot to ensure that the leader and
 
 You will need to move the follower arm to these positions sequentially:
 
-| 1. Zero position                                                                                                                                              | 2. Rotated position                                                                                                                                                    | 3. Rest position                                                                                                                                              |
-| ------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 1. Zero position | 2. Rotated position | 3. Rest position |
+|---|---|---|
 | <img src="../media/moss/follower_zero.webp?raw=true" alt="Moss v1 follower arm zero position" title="Moss v1 follower arm zero position" style="width:100%;"> | <img src="../media/moss/follower_rotated.webp?raw=true" alt="Moss v1 follower arm rotated position" title="Moss v1 follower arm rotated position" style="width:100%;"> | <img src="../media/moss/follower_rest.webp?raw=true" alt="Moss v1 follower arm rest position" title="Moss v1 follower arm rest position" style="width:100%;"> |
 
 Make sure both arms are connected and run this script to launch manual calibration:
@@ -192,8 +192,8 @@ python lerobot/scripts/control_robot.py \
 **Manual calibration of leader arm**
 Follow step 6 of the [assembly video](https://www.youtube.com/watch?v=DA91NJOtMic) which illustrates the manual calibration. You will need to move the leader arm to these positions sequentially:
 
-| 1. Zero position                                                                                                                                        | 2. Rotated position                                                                                                                                              | 3. Rest position                                                                                                                                        |
-| ------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 1. Zero position | 2. Rotated position | 3. Rest position |
+|---|---|---|
 | <img src="../media/moss/leader_zero.webp?raw=true" alt="Moss v1 leader arm zero position" title="Moss v1 leader arm zero position" style="width:100%;"> | <img src="../media/moss/leader_rotated.webp?raw=true" alt="Moss v1 leader arm rotated position" title="Moss v1 leader arm rotated position" style="width:100%;"> | <img src="../media/moss/leader_rest.webp?raw=true" alt="Moss v1 leader arm rest position" title="Moss v1 leader arm rest position" style="width:100%;"> |
 
 Run this script to launch manual calibration:
@@ -293,14 +293,14 @@ python lerobot/scripts/train.py \
   --policy.type=act \
   --output_dir=outputs/train/act_moss_test \
   --job_name=act_moss_test \
-  --policy.device=cuda \
+  --device=cuda \
   --wandb.enable=true
 ```
 
 Let's explain it:
 1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/moss_test`.
 2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../lerobot/common/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor sates, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
-4. 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 `device=cuda` since we are training on a Nvidia GPU, but you could use `device=mps` to train on Apple silicon.
 5. 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. You will find checkpoints in `outputs/train/act_moss_test/checkpoints`.

examples/4_train_policy_with_script.md

@@ -1,5 +1,5 @@
 This tutorial will explain the training script, how to use it, and particularly how to configure everything needed for the training run.
-> **Note:** The following assume you're running these commands on a machine equipped with a cuda GPU. If you don't have one (or if you're using a Mac), you can add `--policy.device=cpu` (`--policy.device=mps` respectively). However, be advised that the code executes much slower on cpu.
+> **Note:** The following assume you're running these commands on a machine equipped with a cuda GPU. If you don't have one (or if you're using a Mac), you can add `--device=cpu` (`--device=mps` respectively). However, be advised that the code executes much slower on cpu.
 
 
 ## The training script

examples/7_get_started_with_real_robot.md

@@ -386,14 +386,14 @@ When you connect your robot for the first time, the [`ManipulatorRobot`](../lero
 
 Here are the positions you'll move the follower arm to:
 
-| 1. Zero position                                                                                                                                                  | 2. Rotated position                                                                                                                                                        | 3. Rest position                                                                                                                                                  |
-| ----------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 1. Zero position | 2. Rotated position | 3. Rest position |
+|---|---|---|
 | <img src="../media/koch/follower_zero.webp?raw=true" alt="Koch v1.1 follower arm zero position" title="Koch v1.1 follower arm zero position" style="width:100%;"> | <img src="../media/koch/follower_rotated.webp?raw=true" alt="Koch v1.1 follower arm rotated position" title="Koch v1.1 follower arm rotated position" style="width:100%;"> | <img src="../media/koch/follower_rest.webp?raw=true" alt="Koch v1.1 follower arm rest position" title="Koch v1.1 follower arm rest position" style="width:100%;"> |
 
 And here are the corresponding positions for the leader arm:
 
-| 1. Zero position                                                                                                                                            | 2. Rotated position                                                                                                                                                  | 3. Rest position                                                                                                                                            |
-| ----------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 1. Zero position | 2. Rotated position | 3. Rest position |
+|---|---|---|
 | <img src="../media/koch/leader_zero.webp?raw=true" alt="Koch v1.1 leader arm zero position" title="Koch v1.1 leader arm zero position" style="width:100%;"> | <img src="../media/koch/leader_rotated.webp?raw=true" alt="Koch v1.1 leader arm rotated position" title="Koch v1.1 leader arm rotated position" style="width:100%;"> | <img src="../media/koch/leader_rest.webp?raw=true" alt="Koch v1.1 leader arm rest position" title="Koch v1.1 leader arm rest position" style="width:100%;"> |
 
 You can watch a [video tutorial of the calibration procedure](https://youtu.be/8drnU9uRY24) for more details.
@@ -898,14 +898,14 @@ python lerobot/scripts/train.py \
   --policy.type=act \
   --output_dir=outputs/train/act_koch_test \
   --job_name=act_koch_test \
-  --policy.device=cuda \
+  --device=cuda \
   --wandb.enable=true
 ```
 
 Let's explain it:
 1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/koch_test`.
 2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../lerobot/common/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor sates, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
-4. 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 `device=cuda` since we are training on a Nvidia GPU, but you could use `device=mps` to train on Apple silicon.
 5. 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`.
 
 For more information on the `train` script see the previous tutorial: [`examples/4_train_policy_with_script.md`](../examples/4_train_policy_with_script.md)

examples/9_use_aloha.md

@@ -135,14 +135,14 @@ python lerobot/scripts/train.py \
   --policy.type=act \
   --output_dir=outputs/train/act_aloha_test \
   --job_name=act_aloha_test \
-  --policy.device=cuda \
+  --device=cuda \
   --wandb.enable=true
 ```
 
 Let's explain it:
 1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/aloha_test`.
 2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../lerobot/common/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor sates, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
-4. 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 `device=cuda` since we are training on a Nvidia GPU, but you could use `device=mps` to train on Apple silicon.
 5. 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`.
 
 For more information on the `train` script see the previous tutorial: [`examples/4_train_policy_with_script.md`](../examples/4_train_policy_with_script.md)

lerobot/common/datasets/factory.py

@@ -24,7 +24,7 @@ from lerobot.common.datasets.lerobot_dataset import (
     MultiLeRobotDataset,
 )
 from lerobot.common.datasets.transforms import ImageTransforms
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.common.policies import PreTrainedConfig
 from lerobot.configs.train import TrainPipelineConfig
 
 IMAGENET_STATS = {

lerobot/common/datasets/lerobot_dataset.py

@@ -67,7 +67,7 @@ from lerobot.common.datasets.utils import (
 )
 from lerobot.common.datasets.video_utils import (
     VideoFrame,
-    decode_video_frames,
+    decode_video_frames_torchvision,
     encode_video_frames,
     get_video_info,
 )
@@ -462,8 +462,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
             download_videos (bool, optional): Flag to download the videos. Note that when set to True but the
                 video files are already present on local disk, they won't be downloaded again. Defaults to
                 True.
-            video_backend (str | None, optional): Video backend to use for decoding videos. Defaults to torchcodec.
-                You can also use the 'pyav' decoder used by Torchvision, which used to be the default option, or 'video_reader' which is another decoder of Torchvision.
+            video_backend (str | None, optional): Video backend to use for decoding videos. There is currently
+                a single option which is the pyav decoder used by Torchvision. Defaults to pyav.
         """
         super().__init__()
         self.repo_id = repo_id
@@ -473,7 +473,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         self.episodes = episodes
         self.tolerance_s = tolerance_s
         self.revision = revision if revision else CODEBASE_VERSION
-        self.video_backend = video_backend if video_backend else "torchcodec"
+        self.video_backend = video_backend if video_backend else "pyav"
         self.delta_indices = None
 
         # Unused attributes
@@ -707,7 +707,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
         item = {}
         for vid_key, query_ts in query_timestamps.items():
             video_path = self.root / self.meta.get_video_file_path(ep_idx, vid_key)
-            frames = decode_video_frames(video_path, query_ts, self.tolerance_s, self.video_backend)
+            frames = decode_video_frames_torchvision(
+                video_path, query_ts, self.tolerance_s, self.video_backend
+            )
             item[vid_key] = frames.squeeze(0)
 
         return item
@@ -1027,7 +1029,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         obj.delta_timestamps = None
         obj.delta_indices = None
         obj.episode_data_index = None
-        obj.video_backend = video_backend if video_backend is not None else "torchcodec"
+        obj.video_backend = video_backend if video_backend is not None else "pyav"
         return obj
 
 

lerobot/common/datasets/video_utils.py

@@ -27,35 +27,6 @@ import torch
 import torchvision
 from datasets.features.features import register_feature
 from PIL import Image
-from torchcodec.decoders import VideoDecoder
-
-
-def decode_video_frames(
-    video_path: Path | str,
-    timestamps: list[float],
-    tolerance_s: float,
-    backend: str = "torchcodec",
-) -> torch.Tensor:
-    """
-    Decodes video frames using the specified backend.
-
-    Args:
-        video_path (Path): Path to the video file.
-        timestamps (list[float]): List of timestamps to extract frames.
-        tolerance_s (float): Allowed deviation in seconds for frame retrieval.
-        backend (str, optional): Backend to use for decoding. Defaults to "torchcodec".
-
-    Returns:
-        torch.Tensor: Decoded frames.
-
-    Currently supports torchcodec on cpu and pyav.
-    """
-    if backend == "torchcodec":
-        return decode_video_frames_torchcodec(video_path, timestamps, tolerance_s)
-    elif backend in ["pyav", "video_reader"]:
-        return decode_video_frames_torchvision(video_path, timestamps, tolerance_s, backend)
-    else:
-        raise ValueError(f"Unsupported video backend: {backend}")
 
 
 def decode_video_frames_torchvision(
@@ -156,76 +127,6 @@ def decode_video_frames_torchvision(
     return closest_frames
 
 
-def decode_video_frames_torchcodec(
-    video_path: Path | str,
-    timestamps: list[float],
-    tolerance_s: float,
-    device: str = "cpu",
-    log_loaded_timestamps: bool = False,
-) -> torch.Tensor:
-    """Loads frames associated with the requested timestamps of a video using torchcodec.
-
-    Note: Setting device="cuda" outside the main process, e.g. in data loader workers, will lead to CUDA initialization errors.
-
-    Note: Video benefits from inter-frame compression. Instead of storing every frame individually,
-    the encoder stores a reference frame (or a key frame) and subsequent frames as differences relative to
-    that key frame. As a consequence, to access a requested frame, we need to load the preceding key frame,
-    and all subsequent frames until reaching the requested frame. The number of key frames in a video
-    can be adjusted during encoding to take into account decoding time and video size in bytes.
-    """
-    video_path = str(video_path)
-    # initialize video decoder
-    decoder = VideoDecoder(video_path, device=device)
-    loaded_frames = []
-    loaded_ts = []
-    # get metadata for frame information
-    metadata = decoder.metadata
-    average_fps = metadata.average_fps
-
-    # convert timestamps to frame indices
-    frame_indices = [round(ts * average_fps) for ts in timestamps]
-
-    # retrieve frames based on indices
-    frames_batch = decoder.get_frames_at(indices=frame_indices)
-
-    for frame, pts in zip(frames_batch.data, frames_batch.pts_seconds, strict=False):
-        loaded_frames.append(frame)
-        loaded_ts.append(pts.item())
-        if log_loaded_timestamps:
-            logging.info(f"Frame loaded at timestamp={pts:.4f}")
-
-    query_ts = torch.tensor(timestamps)
-    loaded_ts = torch.tensor(loaded_ts)
-
-    # compute distances between each query timestamp and loaded timestamps
-    dist = torch.cdist(query_ts[:, None], loaded_ts[:, None], p=1)
-    min_, argmin_ = dist.min(1)
-
-    is_within_tol = min_ < tolerance_s
-    assert is_within_tol.all(), (
-        f"One or several query timestamps unexpectedly violate the tolerance ({min_[~is_within_tol]} > {tolerance_s=})."
-        "It means that the closest frame that can be loaded from the video is too far away in time."
-        "This might be due to synchronization issues with timestamps during data collection."
-        "To be safe, we advise to ignore this item during training."
-        f"\nqueried timestamps: {query_ts}"
-        f"\nloaded timestamps: {loaded_ts}"
-        f"\nvideo: {video_path}"
-    )
-
-    # get closest frames to the query timestamps
-    closest_frames = torch.stack([loaded_frames[idx] for idx in argmin_])
-    closest_ts = loaded_ts[argmin_]
-
-    if log_loaded_timestamps:
-        logging.info(f"{closest_ts=}")
-
-    # convert to float32 in [0,1] range (channel first)
-    closest_frames = closest_frames.type(torch.float32) / 255
-
-    assert len(timestamps) == len(closest_frames)
-    return closest_frames
-
-
 def encode_video_frames(
     imgs_dir: Path | str,
     video_path: Path | str,

lerobot/common/optim/__init__.py

@@ -12,4 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .optimizers import OptimizerConfig as OptimizerConfig
+from .optimizers import OptimizerConfig
+from .schedulers import LRSchedulerConfig
+
+__all__ = ["OptimizerConfig", "LRSchedulerConfig"]

lerobot/common/policies/__init__.py

@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .act.configuration_act import ACTConfig as ACTConfig
-from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
-from .pi0.configuration_pi0 import PI0Config as PI0Config
-from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
-from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
+from . import act, diffusion, pi0, tdmpc, vqbet
+from .config import PreTrainedConfig
+from .factory import make_policy
+from .pretrained import PreTrainedPolicy
+
+__all__ = ["act", "diffusion", "pi0", "tdmpc", "vqbet", "make_policy", "PreTrainedConfig", "PreTrainedPolicy"]

lerobot/common/policies/act/__init__.py

@@ -0,0 +1,4 @@
+from .configuration_act import ACTConfig
+from .modeling_act import ACTPolicy
+
+__all__ = ["ACTConfig", "ACTPolicy"]

lerobot/common/policies/act/configuration_act.py

@@ -16,9 +16,10 @@
 from dataclasses import dataclass, field
 
 from lerobot.common.optim.optimizers import AdamWConfig
-from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
 
+from ..config import PreTrainedConfig
+
 
 @PreTrainedConfig.register_subclass("act")
 @dataclass

lerobot/common/policies/config.py

@@ -23,8 +23,7 @@ from huggingface_hub import hf_hub_download
 from huggingface_hub.constants import CONFIG_NAME
 from huggingface_hub.errors import HfHubHTTPError
 
-from lerobot.common.optim.optimizers import OptimizerConfig
-from lerobot.common.optim.schedulers import LRSchedulerConfig
+from lerobot.common.optim import LRSchedulerConfig, OptimizerConfig
 from lerobot.common.utils.hub import HubMixin
 from lerobot.common.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature

lerobot/common/policies/diffusion/__init__.py

@@ -0,0 +1,4 @@
+from .configuration_diffusion import DiffusionConfig
+from .modeling_diffusion import DiffusionPolicy
+
+__all__ = ["DiffusionConfig", "DiffusionPolicy"]

lerobot/common/policies/diffusion/configuration_diffusion.py

@@ -18,9 +18,10 @@ from dataclasses import dataclass, field
 
 from lerobot.common.optim.optimizers import AdamConfig
 from lerobot.common.optim.schedulers import DiffuserSchedulerConfig
-from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
 
+from ..config import PreTrainedConfig
+
 
 @PreTrainedConfig.register_subclass("diffusion")
 @dataclass

lerobot/common/policies/factory.py

@@ -22,57 +22,38 @@ from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
 from lerobot.common.datasets.utils import dataset_to_policy_features
 from lerobot.common.envs.configs import EnvConfig
 from lerobot.common.envs.utils import env_to_policy_features
-from lerobot.common.policies.act.configuration_act import ACTConfig
-from lerobot.common.policies.diffusion.configuration_diffusion import DiffusionConfig
-from lerobot.common.policies.pi0.configuration_pi0 import PI0Config
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.policies.tdmpc.configuration_tdmpc import TDMPCConfig
-from lerobot.common.policies.vqbet.configuration_vqbet import VQBeTConfig
-from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType
 
+from .config import PreTrainedConfig
+from .pretrained import PreTrainedPolicy
+
 
 def get_policy_class(name: str) -> PreTrainedPolicy:
     """Get the policy's class and config class given a name (matching the policy class' `name` attribute)."""
-    if name == "tdmpc":
-        from lerobot.common.policies.tdmpc.modeling_tdmpc import TDMPCPolicy
-
-        return TDMPCPolicy
-    elif name == "diffusion":
-        from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
-
-        return DiffusionPolicy
-    elif name == "act":
-        from lerobot.common.policies.act.modeling_act import ACTPolicy
+    if name == "act":
+        from .act import ACTPolicy
 
         return ACTPolicy
-    elif name == "vqbet":
-        from lerobot.common.policies.vqbet.modeling_vqbet import VQBeTPolicy
+    elif name == "diffusion":
+        from .diffusion import DiffusionPolicy
 
-        return VQBeTPolicy
+        return DiffusionPolicy
     elif name == "pi0":
-        from lerobot.common.policies.pi0.modeling_pi0 import PI0Policy
+        from .pi0 import PI0Policy
 
         return PI0Policy
+    elif name == "tdmpc":
+        from .tdmpc import TDMPCPolicy
+
+        return TDMPCPolicy
+    elif name == "vqbet":
+        from .vqbet import VQBeTPolicy
+
+        return VQBeTPolicy
     else:
         raise NotImplementedError(f"Policy with name {name} is not implemented.")
 
 
-def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
-    if policy_type == "tdmpc":
-        return TDMPCConfig(**kwargs)
-    elif policy_type == "diffusion":
-        return DiffusionConfig(**kwargs)
-    elif policy_type == "act":
-        return ACTConfig(**kwargs)
-    elif policy_type == "vqbet":
-        return VQBeTConfig(**kwargs)
-    elif policy_type == "pi0":
-        return PI0Config(**kwargs)
-    else:
-        raise ValueError(f"Policy type '{policy_type}' is not available.")
-
-
 def make_policy(
     cfg: PreTrainedConfig,
     ds_meta: LeRobotDatasetMetadata | None = None,

lerobot/common/policies/pi0/__init__.py

@@ -0,0 +1,4 @@
+from .configuration_pi0 import PI0Config
+from .modeling_pi0 import PI0Policy
+
+__all__ = ["PI0Config", "PI0Policy"]

lerobot/common/policies/pi0/configuration_pi0.py

@@ -18,9 +18,10 @@ from lerobot.common.optim.optimizers import AdamWConfig
 from lerobot.common.optim.schedulers import (
     CosineDecayWithWarmupSchedulerConfig,
 )
-from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 
+from ..config import PreTrainedConfig
+
 
 @PreTrainedConfig.register_subclass("pi0")
 @dataclass

lerobot/common/policies/pi0/conversion_scripts/benchmark.py

@@ -15,8 +15,7 @@
 import torch
 
 from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.policies.factory import make_policy
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.common.policies import PreTrainedConfig, make_policy
 
 torch.backends.cudnn.benchmark = True
 

lerobot/common/policies/pi0/conversion_scripts/compare_with_jax.py

@@ -19,8 +19,7 @@ from pathlib import Path
 import torch
 
 from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
-from lerobot.common.policies.factory import make_policy
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.common.policies import PreTrainedConfig, make_policy
 
 
 def display(tensor: torch.Tensor):

lerobot/common/policies/pi0/modeling_pi0.py

@@ -22,11 +22,6 @@
 
 Designed by Physical Intelligence. Ported from Jax by Hugging Face.
 
-Install pi0 extra dependencies:
-```bash
-pip install -e ".[pi0]"
-```
-
 Example of finetuning the pi0 pretrained model (`pi0_base` in `openpi`):
 ```bash
 python lerobot/scripts/train.py \

lerobot/common/policies/pretrained.py

@@ -27,7 +27,8 @@ from safetensors.torch import save_model as save_model_as_safetensor
 from torch import Tensor, nn
 
 from lerobot.common.utils.hub import HubMixin
-from lerobot.configs.policies import PreTrainedConfig
+
+from .config import PreTrainedConfig
 
 T = TypeVar("T", bound="PreTrainedPolicy")
 

lerobot/common/policies/tdmpc/__init__.py

@@ -0,0 +1,4 @@
+from .configuration_tdmpc import TDMPCConfig
+from .modeling_tdmpc import TDMPCPolicy
+
+__all__ = ["TDMPCConfig", "TDMPCPolicy"]

lerobot/common/policies/tdmpc/configuration_tdmpc.py

@@ -17,9 +17,10 @@
 from dataclasses import dataclass, field
 
 from lerobot.common.optim.optimizers import AdamConfig
-from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
 
+from ..config import PreTrainedConfig
+
 
 @PreTrainedConfig.register_subclass("tdmpc")
 @dataclass

lerobot/common/policies/vqbet/__init__.py

@@ -0,0 +1,4 @@
+from .configuration_vqbet import VQBeTConfig
+from .modeling_vqbet import VQBeTPolicy
+
+__all__ = ["VQBeTConfig", "VQBeTPolicy"]

lerobot/common/policies/vqbet/configuration_vqbet.py

@@ -20,9 +20,10 @@ from dataclasses import dataclass, field
 
 from lerobot.common.optim.optimizers import AdamConfig
 from lerobot.common.optim.schedulers import VQBeTSchedulerConfig
-from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
 
+from ..config import PreTrainedConfig
+
 
 @PreTrainedConfig.register_subclass("vqbet")
 @dataclass

lerobot/common/robot_devices/cameras/intelrealsense.py

@@ -114,7 +114,7 @@ def save_images_from_cameras(
         camera = IntelRealSenseCamera(config)
         camera.connect()
         print(
-            f"IntelRealSenseCamera({camera.serial_number}, fps={camera.fps}, width={camera.capture_width}, height={camera.capture_height}, color_mode={camera.color_mode})"
+            f"IntelRealSenseCamera({camera.serial_number}, fps={camera.fps}, width={camera.width}, height={camera.height}, color_mode={camera.color_mode})"
         )
         cameras.append(camera)
 
@@ -224,20 +224,9 @@ class IntelRealSenseCamera:
             self.serial_number = self.find_serial_number_from_name(config.name)
         else:
             self.serial_number = config.serial_number
-
-        # Store the raw (capture) resolution from the config.
-        self.capture_width = config.width
-        self.capture_height = config.height
-
-        # If rotated by ±90, swap width and height.
-        if config.rotation in [-90, 90]:
-            self.width = config.height
-            self.height = config.width
-        else:
-            self.width = config.width
-            self.height = config.height
-
         self.fps = config.fps
+        self.width = config.width
+        self.height = config.height
         self.channels = config.channels
         self.color_mode = config.color_mode
         self.use_depth = config.use_depth
@@ -257,6 +246,7 @@ class IntelRealSenseCamera:
         else:
             import cv2
 
+        # TODO(alibets): Do we keep original width/height or do we define them after rotation?
         self.rotation = None
         if config.rotation == -90:
             self.rotation = cv2.ROTATE_90_COUNTERCLOCKWISE
@@ -294,19 +284,15 @@ class IntelRealSenseCamera:
         config = rs.config()
         config.enable_device(str(self.serial_number))
 
-        if self.fps and self.capture_width and self.capture_height:
+        if self.fps and self.width and self.height:
             # TODO(rcadene): can we set rgb8 directly?
-            config.enable_stream(
-                rs.stream.color, self.capture_width, self.capture_height, rs.format.rgb8, self.fps
-            )
+            config.enable_stream(rs.stream.color, self.width, self.height, rs.format.rgb8, self.fps)
         else:
             config.enable_stream(rs.stream.color)
 
         if self.use_depth:
-            if self.fps and self.capture_width and self.capture_height:
-                config.enable_stream(
-                    rs.stream.depth, self.capture_width, self.capture_height, rs.format.z16, self.fps
-                )
+            if self.fps and self.width and self.height:
+                config.enable_stream(rs.stream.depth, self.width, self.height, rs.format.z16, self.fps)
             else:
                 config.enable_stream(rs.stream.depth)
 
@@ -344,18 +330,18 @@ class IntelRealSenseCamera:
             raise OSError(
                 f"Can't set {self.fps=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_fps}."
             )
-        if self.capture_width is not None and self.capture_width != actual_width:
+        if self.width is not None and self.width != actual_width:
             raise OSError(
-                f"Can't set {self.capture_width=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_width}."
+                f"Can't set {self.width=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_width}."
             )
-        if self.capture_height is not None and self.capture_height != actual_height:
+        if self.height is not None and self.height != actual_height:
             raise OSError(
-                f"Can't set {self.capture_height=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_height}."
+                f"Can't set {self.height=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_height}."
             )
 
         self.fps = round(actual_fps)
-        self.capture_width = round(actual_width)
-        self.capture_height = round(actual_height)
+        self.width = round(actual_width)
+        self.height = round(actual_height)
 
         self.is_connected = True
 
@@ -401,7 +387,7 @@ class IntelRealSenseCamera:
             color_image = cv2.cvtColor(color_image, cv2.COLOR_RGB2BGR)
 
         h, w, _ = color_image.shape
-        if h != self.capture_height or w != self.capture_width:
+        if h != self.height or w != self.width:
             raise OSError(
                 f"Can't capture color image with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
             )
@@ -423,7 +409,7 @@ class IntelRealSenseCamera:
             depth_map = np.asanyarray(depth_frame.get_data())
 
             h, w = depth_map.shape
-            if h != self.capture_height or w != self.capture_width:
+            if h != self.height or w != self.width:
                 raise OSError(
                     f"Can't capture depth map with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
                 )

lerobot/common/robot_devices/cameras/opencv.py

@@ -144,8 +144,8 @@ def save_images_from_cameras(
         camera = OpenCVCamera(config)
         camera.connect()
         print(
-            f"OpenCVCamera({camera.camera_index}, fps={camera.fps}, width={camera.capture_width}, "
-            f"height={camera.capture_height}, color_mode={camera.color_mode})"
+            f"OpenCVCamera({camera.camera_index}, fps={camera.fps}, width={camera.width}, "
+            f"height={camera.height}, color_mode={camera.color_mode})"
         )
         cameras.append(camera)
 
@@ -244,19 +244,9 @@ class OpenCVCamera:
             else:
                 raise ValueError(f"Please check the provided camera_index: {self.camera_index}")
 
-        # Store the raw (capture) resolution from the config.
-        self.capture_width = config.width
-        self.capture_height = config.height
-
-        # If rotated by ±90, swap width and height.
-        if config.rotation in [-90, 90]:
-            self.width = config.height
-            self.height = config.width
-        else:
-            self.width = config.width
-            self.height = config.height
-
         self.fps = config.fps
+        self.width = config.width
+        self.height = config.height
         self.channels = config.channels
         self.color_mode = config.color_mode
         self.mock = config.mock
@@ -273,6 +263,7 @@ class OpenCVCamera:
         else:
             import cv2
 
+        # TODO(aliberts): Do we keep original width/height or do we define them after rotation?
         self.rotation = None
         if config.rotation == -90:
             self.rotation = cv2.ROTATE_90_COUNTERCLOCKWISE
@@ -294,20 +285,10 @@ class OpenCVCamera:
             # when other threads are used to save the images.
             cv2.setNumThreads(1)
 
-        backend = (
-            cv2.CAP_V4L2
-            if platform.system() == "Linux"
-            else cv2.CAP_DSHOW
-            if platform.system() == "Windows"
-            else cv2.CAP_AVFOUNDATION
-            if platform.system() == "Darwin"
-            else cv2.CAP_ANY
-        )
-
         camera_idx = f"/dev/video{self.camera_index}" if platform.system() == "Linux" else self.camera_index
         # First create a temporary camera trying to access `camera_index`,
         # and verify it is a valid camera by calling `isOpened`.
-        tmp_camera = cv2.VideoCapture(camera_idx, backend)
+        tmp_camera = cv2.VideoCapture(camera_idx)
         is_camera_open = tmp_camera.isOpened()
         # Release camera to make it accessible for `find_camera_indices`
         tmp_camera.release()
@@ -330,14 +311,14 @@ class OpenCVCamera:
         # Secondly, create the camera that will be used downstream.
         # Note: For some unknown reason, calling `isOpened` blocks the camera which then
         # needs to be re-created.
-        self.camera = cv2.VideoCapture(camera_idx, backend)
+        self.camera = cv2.VideoCapture(camera_idx)
 
         if self.fps is not None:
             self.camera.set(cv2.CAP_PROP_FPS, self.fps)
-        if self.capture_width is not None:
-            self.camera.set(cv2.CAP_PROP_FRAME_WIDTH, self.capture_width)
-        if self.capture_height is not None:
-            self.camera.set(cv2.CAP_PROP_FRAME_HEIGHT, self.capture_height)
+        if self.width is not None:
+            self.camera.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
+        if self.height is not None:
+            self.camera.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
 
         actual_fps = self.camera.get(cv2.CAP_PROP_FPS)
         actual_width = self.camera.get(cv2.CAP_PROP_FRAME_WIDTH)
@@ -349,22 +330,19 @@ class OpenCVCamera:
             raise OSError(
                 f"Can't set {self.fps=} for OpenCVCamera({self.camera_index}). Actual value is {actual_fps}."
             )
-        if self.capture_width is not None and not math.isclose(
-            self.capture_width, actual_width, rel_tol=1e-3
-        ):
+        if self.width is not None and not math.isclose(self.width, actual_width, rel_tol=1e-3):
             raise OSError(
-                f"Can't set {self.capture_width=} for OpenCVCamera({self.camera_index}). Actual value is {actual_width}."
+                f"Can't set {self.width=} for OpenCVCamera({self.camera_index}). Actual value is {actual_width}."
             )
-        if self.capture_height is not None and not math.isclose(
-            self.capture_height, actual_height, rel_tol=1e-3
-        ):
+        if self.height is not None and not math.isclose(self.height, actual_height, rel_tol=1e-3):
             raise OSError(
-                f"Can't set {self.capture_height=} for OpenCVCamera({self.camera_index}). Actual value is {actual_height}."
+                f"Can't set {self.height=} for OpenCVCamera({self.camera_index}). Actual value is {actual_height}."
             )
 
         self.fps = round(actual_fps)
-        self.capture_width = round(actual_width)
-        self.capture_height = round(actual_height)
+        self.width = round(actual_width)
+        self.height = round(actual_height)
+
         self.is_connected = True
 
     def read(self, temporary_color_mode: str | None = None) -> np.ndarray:
@@ -405,7 +383,7 @@ class OpenCVCamera:
             color_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2RGB)
 
         h, w, _ = color_image.shape
-        if h != self.capture_height or w != self.capture_width:
+        if h != self.height or w != self.width:
             raise OSError(
                 f"Can't capture color image with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
             )

lerobot/common/robot_devices/control_configs.py

@@ -17,9 +17,9 @@ from pathlib import Path
 
 import draccus
 
+from lerobot.common.policies import PreTrainedConfig
 from lerobot.common.robot_devices.robots.configs import RobotConfig
 from lerobot.configs import parser
-from lerobot.configs.policies import PreTrainedConfig
 
 
 @dataclass

lerobot/configs/eval.py

@@ -18,9 +18,9 @@ from dataclasses import dataclass, field
 from pathlib import Path
 
 from lerobot.common import envs, policies  # noqa: F401
+from lerobot.common.policies import PreTrainedConfig
 from lerobot.configs import parser
 from lerobot.configs.default import EvalConfig
-from lerobot.configs.policies import PreTrainedConfig
 
 
 @dataclass

lerobot/configs/parser.py

@@ -11,9 +11,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 importlib
 import inspect
-import pkgutil
 import sys
 from argparse import ArgumentError
 from functools import wraps
@@ -25,7 +23,6 @@ import draccus
 from lerobot.common.utils.utils import has_method
 
 PATH_KEY = "path"
-PLUGIN_DISCOVERY_SUFFIX = "discover_packages_path"
 draccus.set_config_type("json")
 
 
@@ -61,86 +58,6 @@ def parse_arg(arg_name: str, args: Sequence[str] | None = None) -> str | None:
     return None
 
 
-def parse_plugin_args(plugin_arg_suffix: str, args: Sequence[str]) -> dict:
-    """Parse plugin-related arguments from command-line arguments.
-
-    This function extracts arguments from command-line arguments that match a specified suffix pattern.
-    It processes arguments in the format '--key=value' and returns them as a dictionary.
-
-    Args:
-        plugin_arg_suffix (str): The suffix to identify plugin-related arguments.
-        cli_args (Sequence[str]): A sequence of command-line arguments to parse.
-
-    Returns:
-        dict: A dictionary containing the parsed plugin arguments where:
-            - Keys are the argument names (with '--' prefix removed if present)
-            - Values are the corresponding argument values
-
-    Example:
-        >>> args = ['--env.discover_packages_path=my_package',
-        ...         '--other_arg=value']
-        >>> parse_plugin_args('discover_packages_path', args)
-        {'env.discover_packages_path': 'my_package'}
-    """
-    plugin_args = {}
-    for arg in args:
-        if "=" in arg and plugin_arg_suffix in arg:
-            key, value = arg.split("=", 1)
-            # Remove leading '--' if present
-            if key.startswith("--"):
-                key = key[2:]
-            plugin_args[key] = value
-    return plugin_args
-
-
-class PluginLoadError(Exception):
-    """Raised when a plugin fails to load."""
-
-
-def load_plugin(plugin_path: str) -> None:
-    """Load and initialize a plugin from a given Python package path.
-
-    This function attempts to load a plugin by importing its package and any submodules.
-    Plugin registration is expected to happen during package initialization, i.e. when
-    the package is imported the gym environment should be registered and the config classes
-    registered with their parents using the `register_subclass` decorator.
-
-    Args:
-        plugin_path (str): The Python package path to the plugin (e.g. "mypackage.plugins.myplugin")
-
-    Raises:
-        PluginLoadError: If the plugin cannot be loaded due to import errors or if the package path is invalid.
-
-    Examples:
-        >>> load_plugin("external_plugin.core")       # Loads plugin from external package
-
-    Notes:
-        - The plugin package should handle its own registration during import
-        - All submodules in the plugin package will be imported
-        - Implementation follows the plugin discovery pattern from Python packaging guidelines
-
-    See Also:
-        https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/
-    """
-    try:
-        package_module = importlib.import_module(plugin_path, __package__)
-    except (ImportError, ModuleNotFoundError) as e:
-        raise PluginLoadError(
-            f"Failed to load plugin '{plugin_path}'. Verify the path and installation: {str(e)}"
-        ) from e
-
-    def iter_namespace(ns_pkg):
-        return pkgutil.iter_modules(ns_pkg.__path__, ns_pkg.__name__ + ".")
-
-    try:
-        for _finder, pkg_name, _ispkg in iter_namespace(package_module):
-            importlib.import_module(pkg_name)
-    except ImportError as e:
-        raise PluginLoadError(
-            f"Failed to load plugin '{plugin_path}'. Verify the path and installation: {str(e)}"
-        ) from e
-
-
 def get_path_arg(field_name: str, args: Sequence[str] | None = None) -> str | None:
     return parse_arg(f"{field_name}.{PATH_KEY}", args)
 
@@ -188,13 +105,10 @@ def filter_path_args(fields_to_filter: str | list[str], args: Sequence[str] | No
 
 def wrap(config_path: Path | None = None):
     """
-    HACK: Similar to draccus.wrap but does three additional things:
+    HACK: Similar to draccus.wrap but does two additional things:
         - Will remove '.path' arguments from CLI in order to process them later on.
         - If a 'config_path' is passed and the main config class has a 'from_pretrained' method, will
           initialize it from there to allow to fetch configs from the hub directly
-        - Will load plugins specified in the CLI arguments. These plugins will typically register
-            their own subclasses of config classes, so that draccus can find the right class to instantiate
-            from the CLI '.type' arguments
     """
 
     def wrapper_outer(fn):
@@ -207,14 +121,6 @@ def wrap(config_path: Path | None = None):
                 args = args[1:]
             else:
                 cli_args = sys.argv[1:]
-                plugin_args = parse_plugin_args(PLUGIN_DISCOVERY_SUFFIX, cli_args)
-                for plugin_cli_arg, plugin_path in plugin_args.items():
-                    try:
-                        load_plugin(plugin_path)
-                    except PluginLoadError as e:
-                        # add the relevant CLI arg to the error message
-                        raise PluginLoadError(f"{e}\nFailed plugin CLI Arg: {plugin_cli_arg}") from e
-                    cli_args = filter_arg(plugin_cli_arg, cli_args)
                 config_path_cli = parse_arg("config_path", cli_args)
                 if has_method(argtype, "__get_path_fields__"):
                     path_fields = argtype.__get_path_fields__()

lerobot/configs/train.py

@@ -22,12 +22,11 @@ from huggingface_hub import hf_hub_download
 from huggingface_hub.errors import HfHubHTTPError
 
 from lerobot.common import envs
-from lerobot.common.optim import OptimizerConfig
-from lerobot.common.optim.schedulers import LRSchedulerConfig
+from lerobot.common.optim import LRSchedulerConfig, OptimizerConfig
+from lerobot.common.policies import PreTrainedConfig
 from lerobot.common.utils.hub import HubMixin
 from lerobot.configs import parser
 from lerobot.configs.default import DatasetConfig, EvalConfig, WandBConfig
-from lerobot.configs.policies import PreTrainedConfig
 
 TRAIN_CONFIG_NAME = "train_config.json"
 

lerobot/scripts/control_robot.py

@@ -141,7 +141,7 @@ from pprint import pformat
 
 # from safetensors.torch import load_file, save_file
 from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.policies.factory import make_policy
+from lerobot.common.policies import make_policy
 from lerobot.common.robot_devices.control_configs import (
     CalibrateControlConfig,
     ControlPipelineConfig,

lerobot/scripts/eval.py

@@ -67,8 +67,7 @@ from tqdm import trange
 
 from lerobot.common.envs.factory import make_env
 from lerobot.common.envs.utils import preprocess_observation
-from lerobot.common.policies.factory import make_policy
-from lerobot.common.policies.pretrained import PreTrainedPolicy
+from lerobot.common.policies import PreTrainedPolicy, make_policy
 from lerobot.common.policies.utils import get_device_from_parameters
 from lerobot.common.utils.io_utils import write_video
 from lerobot.common.utils.random_utils import set_seed

lerobot/scripts/train.py

@@ -29,8 +29,7 @@ from lerobot.common.datasets.sampler import EpisodeAwareSampler
 from lerobot.common.datasets.utils import cycle
 from lerobot.common.envs.factory import make_env
 from lerobot.common.optim.factory import make_optimizer_and_scheduler
-from lerobot.common.policies.factory import make_policy
-from lerobot.common.policies.pretrained import PreTrainedPolicy
+from lerobot.common.policies import PreTrainedPolicy, make_policy
 from lerobot.common.policies.utils import get_device_from_parameters
 from lerobot.common.utils.logging_utils import AverageMeter, MetricsTracker
 from lerobot.common.utils.random_utils import set_seed

lerobot/scripts/visualize_dataset.py

@@ -265,25 +265,13 @@ def main():
         ),
     )
 
-    parser.add_argument(
-        "--tolerance-s",
-        type=float,
-        default=1e-4,
-        help=(
-            "Tolerance in seconds used to ensure data timestamps respect the dataset fps value"
-            "This is argument passed to the constructor of LeRobotDataset and maps to its tolerance_s constructor argument"
-            "If not given, defaults to 1e-4."
-        ),
-    )
-
     args = parser.parse_args()
     kwargs = vars(args)
     repo_id = kwargs.pop("repo_id")
     root = kwargs.pop("root")
-    tolerance_s = kwargs.pop("tolerance_s")
 
     logging.info("Loading dataset")
-    dataset = LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)
+    dataset = LeRobotDataset(repo_id, root=root)
 
     visualize_dataset(dataset, **vars(args))
 

lerobot/scripts/visualize_dataset_html.py

@@ -446,31 +446,15 @@ def main():
         help="Delete the output directory if it exists already.",
     )
 
-    parser.add_argument(
-        "--tolerance-s",
-        type=float,
-        default=1e-4,
-        help=(
-            "Tolerance in seconds used to ensure data timestamps respect the dataset fps value"
-            "This is argument passed to the constructor of LeRobotDataset and maps to its tolerance_s constructor argument"
-            "If not given, defaults to 1e-4."
-        ),
-    )
-
     args = parser.parse_args()
     kwargs = vars(args)
     repo_id = kwargs.pop("repo_id")
     load_from_hf_hub = kwargs.pop("load_from_hf_hub")
     root = kwargs.pop("root")
-    tolerance_s = kwargs.pop("tolerance_s")
 
     dataset = None
     if repo_id:
-        dataset = (
-            LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)
-            if not load_from_hf_hub
-            else get_dataset_info(repo_id)
-        )
+        dataset = LeRobotDataset(repo_id, root=root) if not load_from_hf_hub else get_dataset_info(repo_id)
 
     visualize_dataset_html(dataset, **vars(args))
 

pyproject.toml

@@ -56,6 +56,7 @@ dependencies = [
     "gymnasium==0.29.1",                                                 # TODO(rcadene, aliberts): Make gym 1.0.0 work
     "h5py>=3.10.0",
     "huggingface-hub[hf-transfer,cli]>=0.27.1 ; python_version < '4.0'",
+    "hydra-core>=1.3.2",
     "imageio[ffmpeg]>=2.34.0",
     "jsonlines>=4.0.0",
     "numba>=0.59.0",
@@ -69,8 +70,8 @@ dependencies = [
     "rerun-sdk>=0.21.0",
     "termcolor>=2.4.0",
     "torch>=2.2.1",
-    "torchcodec>=0.2.1",
     "torchvision>=0.21.0",
+    "transformers>=4.48.0",
     "wandb>=0.16.3",
     "zarr>=2.17.0",
 ]
@@ -84,7 +85,6 @@ dora = [
 dynamixel = ["dynamixel-sdk>=3.7.31", "pynput>=1.7.7"]
 feetech = ["feetech-servo-sdk>=1.0.0", "pynput>=1.7.7"]
 intelrealsense = ["pyrealsense2>=2.55.1.6486 ; sys_platform != 'darwin'"]
-pi0 = ["transformers>=4.48.0"]
 pusht = ["gym-pusht>=0.1.5 ; python_version < '4.0'"]
 stretch = [
     "hello-robot-stretch-body>=0.7.27 ; python_version < '4.0' and sys_platform == 'linux'",

tests/configs/test_plugin_loading.py

@@ -1,89 +0,0 @@
-import sys
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Generator
-
-import pytest
-
-from lerobot.common.envs.configs import EnvConfig
-from lerobot.configs.parser import PluginLoadError, load_plugin, parse_plugin_args, wrap
-
-
-def create_plugin_code(*, base_class: str = "EnvConfig", plugin_name: str = "test_env") -> str:
-    """Creates a dummy plugin module that implements its own EnvConfig subclass."""
-    return f"""
-from dataclasses import dataclass
-from lerobot.common.envs.configs import {base_class}
-
-@{base_class}.register_subclass("{plugin_name}")
-@dataclass
-class TestPluginConfig:
-    value: int = 42
-    """
-
-
-@pytest.fixture
-def plugin_dir(tmp_path: Path) -> Generator[Path, None, None]:
-    """Creates a temporary plugin package structure."""
-    plugin_pkg = tmp_path / "test_plugin"
-    plugin_pkg.mkdir()
-    (plugin_pkg / "__init__.py").touch()
-
-    with open(plugin_pkg / "my_plugin.py", "w") as f:
-        f.write(create_plugin_code())
-
-    # Add tmp_path to Python path so we can import from it
-    sys.path.insert(0, str(tmp_path))
-    yield plugin_pkg
-    sys.path.pop(0)
-
-
-def test_parse_plugin_args():
-    cli_args = [
-        "--env.type=test",
-        "--model.discover_packages_path=some.package",
-        "--env.discover_packages_path=other.package",
-    ]
-    plugin_args = parse_plugin_args("discover_packages_path", cli_args)
-    assert plugin_args == {
-        "model.discover_packages_path": "some.package",
-        "env.discover_packages_path": "other.package",
-    }
-
-
-def test_load_plugin_success(plugin_dir: Path):
-    # Import should work and register the plugin with the real EnvConfig
-    load_plugin("test_plugin")
-
-    assert "test_env" in EnvConfig.get_known_choices()
-    plugin_cls = EnvConfig.get_choice_class("test_env")
-    plugin_instance = plugin_cls()
-    assert plugin_instance.value == 42
-
-
-def test_load_plugin_failure():
-    with pytest.raises(PluginLoadError) as exc_info:
-        load_plugin("nonexistent_plugin")
-    assert "Failed to load plugin 'nonexistent_plugin'" in str(exc_info.value)
-
-
-def test_wrap_with_plugin(plugin_dir: Path):
-    @dataclass
-    class Config:
-        env: EnvConfig
-
-    @wrap()
-    def dummy_func(cfg: Config):
-        return cfg
-
-    # Test loading plugin via CLI args
-    sys.argv = [
-        "dummy_script.py",
-        "--env.discover_packages_path=test_plugin",
-        "--env.type=test_env",
-    ]
-
-    cfg = dummy_func()
-    assert isinstance(cfg, Config)
-    assert isinstance(cfg.env, EnvConfig.get_choice_class("test_env"))
-    assert cfg.env.value == 42

tests/mock_cv2.py

@@ -15,11 +15,6 @@ from functools import cache
 
 import numpy as np
 
-CAP_V4L2 = 200
-CAP_DSHOW = 700
-CAP_AVFOUNDATION = 1200
-CAP_ANY = -1
-
 CAP_PROP_FPS = 5
 CAP_PROP_FRAME_WIDTH = 3
 CAP_PROP_FRAME_HEIGHT = 4

tests/scripts/save_policy_to_safetensors.py

@@ -21,10 +21,11 @@ from safetensors.torch import save_file
 
 from lerobot.common.datasets.factory import make_dataset
 from lerobot.common.optim.factory import make_optimizer_and_scheduler
-from lerobot.common.policies.factory import make_policy, make_policy_config
+from lerobot.common.policies import make_policy
 from lerobot.common.utils.random_utils import set_seed
 from lerobot.configs.default import DatasetConfig
 from lerobot.configs.train import TrainPipelineConfig
+from tests.utils import make_policy_config
 
 
 def get_policy_stats(ds_repo_id: str, policy_name: str, policy_kwargs: dict):

tests/test_cameras.py

@@ -85,8 +85,8 @@ def test_camera(request, camera_type, mock):
     camera.connect()
     assert camera.is_connected
     assert camera.fps is not None
-    assert camera.capture_width is not None
-    assert camera.capture_height is not None
+    assert camera.width is not None
+    assert camera.height is not None
 
     # Test connecting twice raises an error
     with pytest.raises(RobotDeviceAlreadyConnectedError):
@@ -204,49 +204,3 @@ def test_save_images_from_cameras(tmp_path, request, camera_type, mock):
 
     # Small `record_time_s` to speedup unit tests
     save_images_from_cameras(tmp_path, record_time_s=0.02, mock=mock)
-
-
-@pytest.mark.parametrize("camera_type, mock", TEST_CAMERA_TYPES)
-@require_camera
-def test_camera_rotation(request, camera_type, mock):
-    config_kwargs = {"camera_type": camera_type, "mock": mock, "width": 640, "height": 480, "fps": 30}
-
-    # No rotation.
-    camera = make_camera(**config_kwargs, rotation=None)
-    camera.connect()
-    assert camera.capture_width == 640
-    assert camera.capture_height == 480
-    assert camera.width == 640
-    assert camera.height == 480
-    no_rot_img = camera.read()
-    h, w, c = no_rot_img.shape
-    assert h == 480 and w == 640 and c == 3
-    camera.disconnect()
-
-    # Rotation = 90 (clockwise).
-    camera = make_camera(**config_kwargs, rotation=90)
-    camera.connect()
-    # With a 90° rotation, we expect the metadata dimensions to be swapped.
-    assert camera.capture_width == 640
-    assert camera.capture_height == 480
-    assert camera.width == 480
-    assert camera.height == 640
-    import cv2
-
-    assert camera.rotation == cv2.ROTATE_90_CLOCKWISE
-    rot_img = camera.read()
-    h, w, c = rot_img.shape
-    assert h == 640 and w == 480 and c == 3
-    camera.disconnect()
-
-    # Rotation = 180.
-    camera = make_camera(**config_kwargs, rotation=None)
-    camera.connect()
-    assert camera.capture_width == 640
-    assert camera.capture_height == 480
-    assert camera.width == 640
-    assert camera.height == 480
-    no_rot_img = camera.read()
-    h, w, c = no_rot_img.shape
-    assert h == 480 and w == 640 and c == 3
-    camera.disconnect()

tests/test_control_robot.py

@@ -41,15 +41,14 @@ from unittest.mock import patch
 
 import pytest
 
+from lerobot.common.policies import PreTrainedConfig, make_policy
 from lerobot.common.policies.act.configuration_act import ACTConfig
-from lerobot.common.policies.factory import make_policy
 from lerobot.common.robot_devices.control_configs import (
     CalibrateControlConfig,
     RecordControlConfig,
     ReplayControlConfig,
     TeleoperateControlConfig,
 )
-from lerobot.configs.policies import PreTrainedConfig
 from lerobot.scripts.control_robot import calibrate, record, replay, teleoperate
 from tests.test_robots import make_robot
 from tests.utils import TEST_ROBOT_TYPES, mock_calibration_dir, require_robot

tests/test_datasets.py

@@ -40,12 +40,11 @@ from lerobot.common.datasets.utils import (
     unflatten_dict,
 )
 from lerobot.common.envs.factory import make_env_config
-from lerobot.common.policies.factory import make_policy_config
 from lerobot.common.robot_devices.robots.utils import make_robot
 from lerobot.configs.default import DatasetConfig
 from lerobot.configs.train import TrainPipelineConfig
 from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_REPO_ID
-from tests.utils import require_x86_64_kernel
+from tests.utils import make_policy_config, require_x86_64_kernel
 
 
 @pytest.fixture

tests/test_policies.py

@@ -28,11 +28,11 @@ from lerobot.common.datasets.utils import cycle, dataset_to_policy_features
 from lerobot.common.envs.factory import make_env, make_env_config
 from lerobot.common.envs.utils import preprocess_observation
 from lerobot.common.optim.factory import make_optimizer_and_scheduler
+from lerobot.common.policies.act import ACTConfig
 from lerobot.common.policies.act.modeling_act import ACTTemporalEnsembler
 from lerobot.common.policies.factory import (
     get_policy_class,
     make_policy,
-    make_policy_config,
 )
 from lerobot.common.policies.normalize import Normalize, Unnormalize
 from lerobot.common.policies.pretrained import PreTrainedPolicy
@@ -41,7 +41,7 @@ from lerobot.configs.default import DatasetConfig
 from lerobot.configs.train import TrainPipelineConfig
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from tests.scripts.save_policy_to_safetensors import get_policy_stats
-from tests.utils import DEVICE, require_cpu, require_env, require_x86_64_kernel
+from tests.utils import DEVICE, make_policy_config, require_cpu, require_env, require_x86_64_kernel
 
 
 @pytest.fixture
@@ -208,11 +208,10 @@ def test_act_backbone_lr():
     """
     Test that the ACT policy can be instantiated with a different learning rate for the backbone.
     """
-
     cfg = TrainPipelineConfig(
         # TODO(rcadene, aliberts): remove dataset download
         dataset=DatasetConfig(repo_id="lerobot/aloha_sim_insertion_scripted", episodes=[0]),
-        policy=make_policy_config("act", optimizer_lr=0.01, optimizer_lr_backbone=0.001),
+        policy=ACTConfig(optimizer_lr=0.01, optimizer_lr_backbone=0.001),
     )
     cfg.validate()  # Needed for auto-setting some parameters
 
@@ -368,7 +367,7 @@ def test_normalize(insert_temporal_dim):
         # was changed to true. For some reason, tests would pass locally, but not in CI. So here we override
         # to test with `policy.use_mpc=false`.
         ("lerobot/xarm_lift_medium", "tdmpc", {"use_mpc": False}, "use_policy"),
-        # ("lerobot/xarm_lift_medium", "tdmpc", {"use_mpc": True}, "use_mpc"),
+        ("lerobot/xarm_lift_medium", "tdmpc", {"use_mpc": True}, "use_mpc"),
         # TODO(rcadene): the diffusion model was normalizing the image in mean=0.5 std=0.5 which is a hack supposed to
         # to normalize the image at all. In our current codebase we dont normalize at all. But there is still a minor difference
         # that fails the test. However, by testing to normalize the image with 0.5 0.5 in the current codebase, the test pass.

tests/utils.py

@@ -23,6 +23,7 @@ import pytest
 import torch
 
 from lerobot import available_cameras, available_motors, available_robots
+from lerobot.common.policies import PreTrainedConfig
 from lerobot.common.robot_devices.cameras.utils import Camera
 from lerobot.common.robot_devices.cameras.utils import make_camera as make_camera_device
 from lerobot.common.robot_devices.motors.utils import MotorsBus
@@ -329,3 +330,8 @@ def make_motors_bus(motor_type: str, **kwargs) -> MotorsBus:
 
     else:
         raise ValueError(f"The motor type '{motor_type}' is not valid.")
+
+
+def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
+    policy_cfg_cls = PreTrainedConfig.get_choice_class(policy_type)
+    return policy_cfg_cls(**kwargs)