precommit

- Modified tests to assert that the processor raises appropriate exceptions when the task key is missing or has an invalid value in the complementary data.
- Ensured that the tests cover cases for None, integer, and mixed list task values, improving robustness against invalid inputs.

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -25,7 +25,7 @@ body:
     id: system-info
     attributes:
       label: System Info
-      description: Please share your LeRobot configuration by running `lerobot-info` (if installed) or `python -m lerobot.scripts.display_sys_info` (if not installed) and pasting the output below.
+      description: If needed, you can share your lerobot configuration with us by running `python -m lerobot.scripts.display_sys_info` and copy-pasting its outputs below
       render: Shell
       placeholder: lerobot version, OS, python version, numpy version, torch version, and lerobot's configuration
     validations:

.github/workflows/stale.yml

@@ -1,68 +0,0 @@
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# This workflow handles closing stale issues and PRs.
-name: Stale
-on:
-  # Allows running this workflow manually from the Actions tab
-  workflow_dispatch:
-
-  # Runs at 02:00
-  schedule:
-    - cron: "0 2 * * *"
-
-env:
-  CLOSE_ISSUE_MESSAGE: >
-    This issue was closed because it has been stalled for 14 days with no activity.
-    Feel free to reopen if is still relevant, or to ping a collaborator if you have any questions.
-  CLOSE_PR_MESSAGE: >
-    This PR was closed because it has been stalled for 14 days with no activity.
-    Feel free to reopen if is still relevant, or to ping a collaborator if you have any questions.
-  WARN_ISSUE_MESSAGE: >
-    This issue has been automatically marked as stale because it has not had
-    recent activity (6 months). It will be closed if no further activity occurs.
-    Thank you for your contributions.
-  WARN_PR_MESSAGE: >
-    This PR has been automatically marked as stale because it has not had
-    recent activity (6 months). It will be closed if no further activity occurs.
-    Thank you for your contributions.
-
-jobs:
-  # This job runs the actions/stale action to close stale issues and PRs.
-  stale:
-    name: Close Stale Issues and PRs
-    runs-on: ubuntu-latest
-    permissions:
-      actions: write
-      contents: write # only for delete-branch option
-      issues: write
-      pull-requests: write
-    steps:
-      - uses: actions/stale@v10
-        with:
-          repo-token: ${{ secrets.GITHUB_TOKEN }}
-          stale-issue-label: stale
-          stale-pr-label: stale
-          exempt-issue-labels: never-stale
-          exempt-pr-labels: never-stale
-          days-before-issue-stale: 180 # TODO(Steven): Will modify this to 90 after initial cleanup
-          days-before-issue-close: 14
-          days-before-pr-stale: 180
-          days-before-pr-close: 14
-          delete-branch: true
-          close-issue-message: ${{ env.CLOSE_ISSUE_MESSAGE }}
-          close-pr-message: ${{ env.CLOSE_PR_MESSAGE }}
-          stale-issue-message: ${{ env.WARN_ISSUE_MESSAGE }}
-          stale-pr-message: ${{ env.WARN_PR_MESSAGE }}
-          operations-per-run: 500

README.md

@@ -202,7 +202,7 @@ Check out [example 1](https://github.com/huggingface/lerobot/blob/main/examples/
 You can also locally visualize episodes from a dataset on the hub by executing our script from the command line:
 
 ```bash
-lerobot-dataset-viz \
+python -m lerobot.scripts.visualize_dataset \
     --repo-id lerobot/pusht \
     --episode-index 0
 ```
@@ -210,7 +210,7 @@ lerobot-dataset-viz \
 or from a dataset in a local folder with the `root` option and the `--local-files-only` (in the following case the dataset will be searched for in `./my_local_data_dir/lerobot/pusht`)
 
 ```bash
-lerobot-dataset-viz \
+python -m lerobot.scripts.visualize_dataset \
     --repo-id lerobot/pusht \
     --root ./my_local_data_dir \
     --local-files-only 1 \
@@ -221,19 +221,19 @@ It will open `rerun.io` and display the camera streams, robot states and actions
 
 https://github-production-user-asset-6210df.s3.amazonaws.com/4681518/328035972-fd46b787-b532-47e2-bb6f-fd536a55a7ed.mov?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240505%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240505T172924Z&X-Amz-Expires=300&X-Amz-Signature=d680b26c532eeaf80740f08af3320d22ad0b8a4e4da1bcc4f33142c15b509eda&X-Amz-SignedHeaders=host&actor_id=24889239&key_id=0&repo_id=748713144
 
-Our script can also visualize datasets stored on a distant server. See `lerobot-dataset-viz --help` for more instructions.
+Our script can also visualize datasets stored on a distant server. See `python -m lerobot.scripts.visualize_dataset --help` for more instructions.
 
 ### The `LeRobotDataset` format
 
 A dataset in `LeRobotDataset` format is very simple to use. It can be loaded from a repository on the Hugging Face hub or a local folder simply with e.g. `dataset = LeRobotDataset("lerobot/aloha_static_coffee")` and can be indexed into like any Hugging Face and PyTorch dataset. For instance `dataset[0]` will retrieve a single temporal frame from the dataset containing observation(s) and an action as PyTorch tensors ready to be fed to a model.
 
-A specificity of `LeRobotDataset` is that, rather than retrieving a single frame by its index, we can retrieve several frames based on their temporal relationship with the indexed frame, by setting `delta_timestamps` to a list of relative times with respect to the indexed frame. For example, with `delta_timestamps = {"observation.image": [-1, -0.5, -0.2, 0]}` one can retrieve, for a given index, 4 frames: 3 "previous" frames 1 second, 0.5 seconds, and 0.2 seconds before the indexed frame, and the indexed frame itself (corresponding to the 0 entry). See example [1_load_lerobot_dataset.py](https://github.com/huggingface/lerobot/blob/main/examples/dataset/load_lerobot_dataset.py) for more details on `delta_timestamps`.
+A specificity of `LeRobotDataset` is that, rather than retrieving a single frame by its index, we can retrieve several frames based on their temporal relationship with the indexed frame, by setting `delta_timestamps` to a list of relative times with respect to the indexed frame. For example, with `delta_timestamps = {"observation.image": [-1, -0.5, -0.2, 0]}` one can retrieve, for a given index, 4 frames: 3 "previous" frames 1 second, 0.5 seconds, and 0.2 seconds before the indexed frame, and the indexed frame itself (corresponding to the 0 entry). See example [1_load_lerobot_dataset.py](https://github.com/huggingface/lerobot/blob/main/examples/1_load_lerobot_dataset.py) for more details on `delta_timestamps`.
 
 Under the hood, the `LeRobotDataset` format makes use of several ways to serialize data which can be useful to understand if you plan to work more closely with this format. We tried to make a flexible yet simple dataset format that would cover most type of features and specificities present in reinforcement learning and robotics, in simulation and in real-world, with a focus on cameras and robot states but easily extended to other types of sensory inputs as long as they can be represented by a tensor.
 
 Here are the important details and internal structure organization of a typical `LeRobotDataset` instantiated with `dataset = LeRobotDataset("lerobot/aloha_static_coffee")`. The exact features will change from dataset to dataset but not the main aspects:
 
-```
+````
 dataset attributes:
   ├ hf_dataset: a Hugging Face dataset (backed by Arrow/parquet). Typical features example:
   │  ├ observation.images.cam_high (VideoFrame):
@@ -269,7 +269,7 @@ dataset attributes:
   ├ root (Path): local directory where the dataset is stored
   ├ image_transforms (Callable): optional image transformations to apply to visual modalities
   └ delta_timestamps (dict): optional delta timestamps for temporal queries
-```
+decoding videos (e.g., 'pyav', 'torchcodec')
 
 A `LeRobotDataset` is serialised using several widespread file formats for each of its parts, namely:
 
@@ -279,6 +279,42 @@ A `LeRobotDataset` is serialised using several widespread file formats for each
 
 Dataset can be uploaded/downloaded from the HuggingFace hub seamlessly. To work on a local dataset, you can specify its location with the `root` argument if it's not in the default `~/.cache/huggingface/lerobot` location.
 
+### Evaluate a pretrained policy
+
+Check out [example 2](https://github.com/huggingface/lerobot/blob/main/examples/2_evaluate_pretrained_policy.py) that illustrates how to download a pretrained policy from Hugging Face hub, and run an evaluation on its corresponding environment.
+
+We also provide a more capable script to parallelize the evaluation over multiple environments during the same rollout. Here is an example with a pretrained model hosted on [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht):
+
+```bash
+lerobot-eval \
+    --policy.path=lerobot/diffusion_pusht \
+    --env.type=pusht \
+    --eval.batch_size=10 \
+    --eval.n_episodes=10 \
+    --policy.use_amp=false \
+    --policy.device=cuda
+````
+
+Note: After training your own policy, you can re-evaluate the checkpoints with:
+
+```bash
+lerobot-eval --policy.path={OUTPUT_DIR}/checkpoints/last/pretrained_model
+```
+
+See `lerobot-eval --help` for more instructions.
+
+### Train your own policy
+
+Check out [example 3](https://github.com/huggingface/lerobot/blob/main/examples/3_train_policy.py) that illustrates how to train a model using our core library in python, and [example 4](https://github.com/huggingface/lerobot/blob/main/examples/4_train_policy_with_script.md) that shows how to use our training script from command line.
+
+To use wandb for logging training and evaluation curves, make sure you've run `wandb login` as a one-time setup step. Then, when running the training command above, enable WandB in the configuration by adding `--wandb.enable=true`.
+
+A link to the wandb logs for the run will also show up in yellow in your terminal. Here is an example of what they look like in your browser. Please also check [here](https://github.com/huggingface/lerobot/blob/main/examples/4_train_policy_with_script.md#typical-logs-and-metrics) for the explanation of some commonly used metrics in logs.
+
+\<img src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/wandb.png" alt="WandB logs example"\>
+
+Note: For efficiency, during training every checkpoint is evaluated on a low number of episodes. You may use `--eval.n_episodes=500` to evaluate on more episodes than the default. Or, after training, you may want to re-evaluate your best checkpoints on more episodes or change the evaluation settings. See `lerobot-eval --help` for more instructions.
+
 #### Reproduce state-of-the-art (SOTA)
 
 We provide some pretrained policies on our [hub page](https://huggingface.co/lerobot) that can achieve state-of-the-art performances.
@@ -337,7 +373,3 @@ If you want, you can cite this work with:
 ## Star History
 
 [![Star History Chart](https://api.star-history.com/svg?repos=huggingface/lerobot&type=Timeline)](https://star-history.com/#huggingface/lerobot&Timeline)
-
-```
-
-```

benchmarks/video/run_video_benchmark.py

@@ -35,13 +35,12 @@ 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,
     encode_video_frames,
 )
-from lerobot.utils.constants import OBS_IMAGE
+from lerobot.utils.benchmark import TimeBenchmark
 
 BASE_ENCODING = OrderedDict(
     [
@@ -118,7 +117,7 @@ def save_first_episode(imgs_dir: Path, dataset: LeRobotDataset) -> None:
     hf_dataset = dataset.hf_dataset.with_format(None)
 
     # We only save images from the first camera
-    img_keys = [key for key in hf_dataset.features if key.startswith(OBS_IMAGE)]
+    img_keys = [key for key in hf_dataset.features if key.startswith("observation.image")]
     imgs_dataset = hf_dataset.select_columns(img_keys[0])
 
     for i, item in enumerate(

docs/source/_toctree.yml

@@ -23,14 +23,14 @@
 - sections:
   - local: lerobot-dataset-v3
     title: Using LeRobotDataset
+  - local: libero
+    title: Using Libero
   - local: porting_datasets_v3
     title: Porting Large Datasets
   title: "Datasets"
 - sections:
   - local: smolvla
     title: Finetune SmolVLA
-  - local: libero
-    title: Using Libero
   title: "Policies"
 
 - sections:

docs/source/async.mdx

@@ -31,7 +31,7 @@ Then, spin up a policy server (in one terminal, or in a separate machine) specif
 You can spin up a policy server running:
 
 ```shell
-python src/lerobot/async_inference/policy_server.py \
+python src/lerobot/scripts/server/policy_server.py \
     --host=127.0.0.1 \
     --port=8080 \
 ```
@@ -39,7 +39,7 @@ python src/lerobot/async_inference/policy_server.py \
 This will start a policy server listening on `127.0.0.1:8080` (`localhost`, port 8080). At this stage, the policy server is empty, as all information related to which policy to run and with which parameters are specified during the first handshake with the client. Spin up a client with:
 
 ```shell
-python src/lerobot/async_inference/robot_client.py \
+python src/lerobot/scripts/server/robot_client.py \
     --server_address=127.0.0.1:8080 \ # SERVER: the host address and port of the policy server
     --robot.type=so100_follower \ # ROBOT: your robot type
     --robot.port=/dev/tty.usbmodem585A0076841 \ # ROBOT: your robot port
@@ -122,8 +122,8 @@ python -m lerobot.scripts.server.policy_server \
 
 <!-- prettier-ignore-start -->
 ```python
-from lerobot.async_inference.configs import PolicyServerConfig
-from lerobot.async_inference.policy_server import serve
+from lerobot.scripts.server.configs import PolicyServerConfig
+from lerobot.scripts.server.policy_server import serve
 
 config = PolicyServerConfig(
     host="localhost",
@@ -148,7 +148,7 @@ The `RobotClient` streams observations to the `PolicyServer`, and receives actio
 <hfoptions id="start_robot_client">
 <hfoption id="Command">
 ```bash
-python src/lerobot/async_inference/robot_client.py \
+python src/lerobot/scripts/server/robot_client.py \
     --server_address=127.0.0.1:8080 \ # SERVER: the host address and port of the policy server
     --robot.type=so100_follower \ # ROBOT: your robot type
     --robot.port=/dev/tty.usbmodem585A0076841 \ # ROBOT: your robot port
@@ -171,9 +171,9 @@ python src/lerobot/async_inference/robot_client.py \
 import threading
 from lerobot.robots.so100_follower import SO100FollowerConfig
 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.async_inference.configs import RobotClientConfig
-from lerobot.async_inference.robot_client import RobotClient
-from lerobot.async_inference.helpers import visualize_action_queue_size
+from lerobot.scripts.server.configs import RobotClientConfig
+from lerobot.scripts.server.robot_client import RobotClient
+from lerobot.scripts.server.helpers import visualize_action_queue_size
 
 # 1. Create the robot instance
 """Check out the cameras available in your setup by running `python lerobot/find_cameras.py`"""

docs/source/hilserl.mdx

@@ -62,7 +62,7 @@ pip install -e ".[hilserl]"
 
 ### Understanding Configuration
 
-The training process begins with proper configuration for the HILSerl environment. The main configuration class is `GymManipulatorConfig` in `lerobot/rl/gym_manipulator.py`, which contains nested `HILSerlRobotEnvConfig` and `DatasetConfig`. The configuration is organized into focused, nested sub-configs:
+The training process begins with proper configuration for the HILSerl environment. The main configuration class is `GymManipulatorConfig` in `lerobot/scripts/rl/gym_manipulator.py`, which contains nested `HILSerlRobotEnvConfig` and `DatasetConfig`. The configuration is organized into focused, nested sub-configs:
 
 <!-- prettier-ignore-start -->
 ```python
@@ -160,8 +160,9 @@ The action processor (`action_processor`) handles outgoing actions and human int
 
 1. **AddTeleopActionAsComplimentaryDataStep**: Captures teleoperator actions for logging
 2. **AddTeleopEventsAsInfoStep**: Records intervention events and episode control signals
-3. **InterventionActionProcessorStep**: Handles human interventions and episode termination
-4. **Inverse Kinematics Pipeline** (when enabled):
+3. **AddRobotObservationAsComplimentaryData**: Stores raw robot state for processing
+4. **InterventionActionProcessorStep**: Handles human interventions and episode termination
+5. **Inverse Kinematics Pipeline** (when enabled):
    - **MapDeltaActionToRobotActionStep**: Converts delta actions to robot action format
    - **EEReferenceAndDelta**: Computes end-effector reference and delta movements
    - **EEBoundsAndSafety**: Enforces workspace safety bounds
@@ -304,19 +305,19 @@ Before collecting demonstrations, you need to determine the appropriate operatio
 
 This helps simplify the problem of learning on the real robot in two ways: 1) by limiting the robot's operational space to a specific region that solves the task and avoids unnecessary or unsafe exploration, and 2) by allowing training in end-effector space rather than joint space. Empirically, learning in joint space for reinforcement learning in manipulation is often a harder problem - some tasks are nearly impossible to learn in joint space but become learnable when the action space is transformed to end-effector coordinates.
 
-**Using lerobot-find-joint-limits**
+**Using find_joint_limits.py**
 
 This script helps you find the safe operational bounds for your robot's end-effector. Given that you have a follower and leader arm, you can use the script to find the bounds for the follower arm that will be applied during training.
 Bounding the action space will reduce the redundant exploration of the agent and guarantees safety.
 
 ```bash
-lerobot-find-joint-limits \
-  --robot.type=so100_follower \
-  --robot.port=/dev/tty.usbmodem58760431541 \
-  --robot.id=black \
-  --teleop.type=so100_leader \
-  --teleop.port=/dev/tty.usbmodem58760431551 \
-  --teleop.id=blue
+python -m lerobot.scripts.find_joint_limits \
+    --robot.type=so100_follower \
+    --robot.port=/dev/tty.usbmodem58760431541 \
+    --robot.id=black \
+    --teleop.type=so100_leader \
+    --teleop.port=/dev/tty.usbmodem58760431551 \
+    --teleop.id=blue
 ```
 
 **Workflow**
@@ -346,7 +347,7 @@ With the bounds defined, you can safely collect demonstrations for training. Tra
 
 **Setting Up Record Mode**
 
-Create a configuration file for recording demonstrations (or edit an existing one like [env_config.json](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/rl/env_config.json)):
+Create a configuration file for recording demonstrations (or edit an existing one like [env_config_so100.json](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_so100.json)):
 
 1. Set `mode` to `"record"` at the root level
 2. Specify a unique `repo_id` for your dataset in the `dataset` section (e.g., "username/task_name")
@@ -518,7 +519,7 @@ During the online training, press `space` to take over the policy and `space` ag
 Start the recording process, an example of the config file can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_so100.json):
 
 ```bash
-python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/env_config_so100.json
+python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config_so100.json
 ```
 
 During recording:
@@ -549,7 +550,7 @@ Note: If you already know the crop parameters, you can skip this step and just s
 Use the `crop_dataset_roi.py` script to interactively select regions of interest in your camera images:
 
 ```bash
-python -m lerobot.rl.crop_dataset_roi --repo-id username/pick_lift_cube
+python -m lerobot.scripts.rl.crop_dataset_roi --repo-id username/pick_lift_cube
 ```
 
 1. For each camera view, the script will display the first frame
@@ -618,7 +619,7 @@ Before training, you need to collect a dataset with labeled examples. The `recor
 To collect a dataset, you need to modify some parameters in the environment configuration based on HILSerlRobotEnvConfig.
 
 ```bash
-python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/reward_classifier_train_config.json
+python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/reward_classifier_train_config.json
 ```
 
 **Key Parameters for Data Collection**
@@ -764,7 +765,7 @@ or set the argument in the json config file.
 Run `gym_manipulator.py` to test the model.
 
 ```bash
-python -m lerobot.rl.gym_manipulator --config_path path/to/env_config.json
+python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config.json
 ```
 
 The reward classifier will automatically provide rewards based on the visual input from the robot's cameras.
@@ -772,12 +773,12 @@ The reward classifier will automatically provide rewards based on the visual inp
 **Example Workflow for training the reward classifier**
 
 1. **Create the configuration files**:
-   Create the necessary json configuration files for the reward classifier and the environment. Check the examples [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/reward_classifier/config.json).
+   Create the necessary json configuration files for the reward classifier and the environment. Check the examples [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/tree/main).
 
 2. **Collect a dataset**:
 
    ```bash
-   python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
+   python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
    ```
 
 3. **Train the classifier**:
@@ -788,7 +789,7 @@ The reward classifier will automatically provide rewards based on the visual inp
 
 4. **Test the classifier**:
    ```bash
-   python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
+   python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
    ```
 
 ### Training with Actor-Learner
@@ -797,7 +798,7 @@ The LeRobot system uses a distributed actor-learner architecture for training. T
 
 **Configuration Setup**
 
-Create a training configuration file (example available [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/rl/train_config.json)). The training config is based on the main `TrainRLServerPipelineConfig` class in `lerobot/configs/train.py`.
+Create a training configuration file (example available [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/train_config_hilserl_so100.json)). The training config is based on the main `TrainRLServerPipelineConfig` class in `lerobot/configs/train.py`.
 
 1. Configure the policy settings (`type="sac"`, `device`, etc.)
 2. Set `dataset` to your cropped dataset
@@ -810,7 +811,7 @@ Create a training configuration file (example available [here](https://huggingfa
 First, start the learner server process:
 
 ```bash
-python -m lerobot.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
+python -m lerobot.scripts.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
 ```
 
 The learner:
@@ -825,7 +826,7 @@ The learner:
 In a separate terminal, start the actor process with the same configuration:
 
 ```bash
-python -m lerobot.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
+python -m lerobot.scripts.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
 ```
 
 The actor:

docs/source/hilserl_sim.mdx

@@ -26,7 +26,7 @@ pip install -e ".[hilserl]"
 
 ## Configuration
 
-To use `gym_hil` with LeRobot, you need to create a configuration file. An example is provided [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/rl/gym_hil/env_config.json). Key configuration sections include:
+To use `gym_hil` with LeRobot, you need to create a configuration file. An example is provided [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/gym_hil_env.json). Key configuration sections include:
 
 ### Environment Type and Task
 
@@ -91,7 +91,7 @@ Important parameters:
 To run the environment, set mode to null:
 
 ```bash
-python -m lerobot.rl.gym_manipulator --config_path path/to/gym_hil_env.json
+python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/gym_hil_env.json
 ```
 
 ### Recording a Dataset
@@ -118,21 +118,21 @@ To collect a dataset, set the mode to `record` whilst defining the repo_id and n
 ```
 
 ```bash
-python -m lerobot.rl.gym_manipulator --config_path path/to/gym_hil_env.json
+python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/gym_hil_env.json
 ```
 
 ### Training a Policy
 
-To train a policy, checkout the configuration example available [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/rl/gym_hil/train_config.json) and run the actor and learner servers:
+To train a policy, checkout the configuration example available [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/train_gym_hil_env.json) and run the actor and learner servers:
 
 ```bash
-python -m lerobot.rl.actor --config_path path/to/train_gym_hil_env.json
+python -m lerobot.scripts.rl.actor --config_path path/to/train_gym_hil_env.json
 ```
 
 In a different terminal, run the learner server:
 
 ```bash
-python -m lerobot.rl.learner --config_path path/to/train_gym_hil_env.json
+python -m lerobot.scripts.rl.learner --config_path path/to/train_gym_hil_env.json
 ```
 
 The simulation environment provides a safe and repeatable way to develop and test your Human-In-the-Loop reinforcement learning components before deploying to real robots.

docs/source/il_robots.mdx

@@ -200,7 +200,7 @@ from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderCo
 from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 from lerobot.record import record_loop
 
 NUM_EPISODES = 5
@@ -237,7 +237,7 @@ dataset = LeRobotDataset.create(
 
 # Initialize the keyboard listener and rerun visualization
 _, events = init_keyboard_listener()
-init_rerun(session_name="recording")
+_init_rerun(session_name="recording")
 
 # Connect the robot and teleoperator
 robot.connect()
@@ -517,7 +517,7 @@ from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerCon
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 from lerobot.record import record_loop
 from lerobot.policies.factory import make_processor
 
@@ -557,7 +557,7 @@ dataset = LeRobotDataset.create(
 
 # Initialize the keyboard listener and rerun visualization
 _, events = init_keyboard_listener()
-init_rerun(session_name="recording")
+_init_rerun(session_name="recording")
 
 # Connect the robot
 robot.connect()

docs/source/il_sim.mdx

@@ -22,7 +22,7 @@ pip install -e ".[hilserl]"
 
 ## Teleoperate and Record a Dataset
 
-To use `gym_hil` with LeRobot, you need to use a configuration file. An example config file can be found [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/sim_il/env_config.json).
+To use `gym_hil` with LeRobot, you need to use a configuration file. An example config file can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_gym_hil_il.json).
 
 To teleoperate and collect a dataset, we need to modify this config file. Here's an example configuration for imitation learning data collection:
 
@@ -61,14 +61,14 @@ Then we can run this command to start:
 <hfoption id="Linux">
 
 ```bash
-python -m lerobot.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
+python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
 ```
 
 </hfoption>
 <hfoption id="MacOS">
 
 ```bash
-mjpython -m lerobot.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
+mjpython -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
 ```
 
 </hfoption>
@@ -165,7 +165,7 @@ huggingface-cli upload ${HF_USER}/il_sim_test${CKPT} \
 
 ## Evaluate your policy in Sim
 
-To evaluate your policy we have to use a configuration file. An example can be found [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/sim_il/eval_config.json).
+To evaluate your policy we have to use a configuration file. An example can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/eval_config_gym_hil.json).
 
 Here's an example evaluation configuration:
 
@@ -198,14 +198,14 @@ Then you can run this command to visualize your trained policy
 <hfoption id="Linux">
 
 ```bash
-python -m lerobot.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
+python -m lerobot.scripts.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
 ```
 
 </hfoption>
 <hfoption id="MacOS">
 
 ```bash
-mjpython -m lerobot.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
+mjpython -m lerobot.scripts.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
 ```
 
 </hfoption>

docs/source/introduction_processors.mdx

@@ -14,7 +14,7 @@ Training data from one robot setup needs adaptation for deployment on different
 **That's where processors come in.** They serve as universal translators that bridge these gaps, ensuring seamless data flow from sensors to models to actuators.
 Processors handle all the preprocessing and postprocessing steps needed to convert raw environment data into model-ready inputs and vice versa.
 
-This means that your favorite policy can be used like this:
+Now your favorite policy can be used like this:
 
 ```python
 import torch
@@ -40,7 +40,7 @@ postprocessed_action = postprocessor(action)
 
 ## What are Processors?
 
-In robotics, data comes in many forms: images from cameras, joint positions from sensors, text instructions from users, and more. Each type of data requires specific transformations before a model can use it effectively. Models need this data to be:
+In robotics, data comes in many forms - images from cameras, joint positions from sensors, text instructions from users, and more. Each type of data requires specific transformations before a model can use it effectively. Models need this data to be:
 
 - **Normalized**: Scaled to appropriate ranges for neural network processing
 - **Batched**: Organized with proper dimensions for batch processing
@@ -181,8 +181,8 @@ training_preprocessor = PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
         RenameObservationsProcessorStep(rename_map={}),     # Standardize keys
         AddBatchDimensionProcessorStep(),                   # Add batch dims
         TokenizerProcessorStep(tokenizer_name="...", ...),  # Tokenize language
-        DeviceProcessorStep(device="cuda"),                 # Move to GPU first
-        NormalizerProcessorStep(features=..., stats=...),   # Normalize on GPU
+        DeviceProcessorStep(device="cuda"),                 # Move to GPU first ⚡
+        NormalizerProcessorStep(features=..., stats=...),   # Normalize on GPU ⚡
     ]
 )
 

docs/source/koch.mdx

@@ -277,7 +277,7 @@ leader.disconnect()
 </hfoption>
 </hfoptions>
 
-Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./il_robots)
+Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./getting_started_real_world_robot)
 
 > [!TIP]
 > If you have any questions or need help, please reach out on [Discord](https://discord.com/invite/s3KuuzsPFb).

docs/source/lekiwi.mdx

@@ -323,7 +323,7 @@ To replay an episode run the API example below, make sure to change `remote_ip`,
 python examples/lekiwi/replay.py
 ```
 
-Congrats 🎉, your robot is all set to learn a task on its own. Start training it by the training part of this tutorial: [Getting started with real-world robots](./il_robots)
+Congrats 🎉, your robot is all set to learn a task on its own. Start training it by the training part of this tutorial: [Getting started with real-world robots](./getting_started_real_world_robot)
 
 ## Evaluate your policy
 

docs/source/lerobot-dataset-v3.mdx

@@ -8,7 +8,6 @@ This docs will guide you to:
 - Record a dataset and push it to the Hub
 - Load datasets for training with `LeRobotDataset`
 - Stream datasets without downloading using `StreamingLeRobotDataset`
-- Apply image transforms for data augmentation during training
 - Migrate existing `v2.1` datasets to `v3.0`
 
 ## What’s new in `v3`
@@ -151,117 +150,6 @@ dataset = StreamingLeRobotDataset(repo_id)  # streams directly from the Hub
   </figure>
 </div>
 
-## Image transforms
-
-Image transforms are data augmentations applied to camera frames during training to improve model robustness and generalization. LeRobot supports various transforms including brightness, contrast, saturation, hue, and sharpness adjustments.
-
-### Using transforms during dataset creation/recording
-
-Currently, transforms are applied during **training time only**, not during recording. When you create or record a dataset, the raw images are stored without transforms. This allows you to experiment with different augmentations later without re-recording data.
-
-### Adding transforms to existing datasets (API)
-
-Use the `image_transforms` parameter when loading a dataset for training:
-
-```python
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.transforms import ImageTransforms, ImageTransformsConfig, ImageTransformConfig
-
-# Option 1: Use default transform configuration (disabled by default)
-transforms_config = ImageTransformsConfig(
-    enable=True,  # Enable transforms
-    max_num_transforms=3,  # Apply up to 3 transforms per frame
-    random_order=False,  # Apply in standard order
-)
-transforms = ImageTransforms(transforms_config)
-
-dataset = LeRobotDataset(
-    repo_id="your-username/your-dataset",
-    image_transforms=transforms
-)
-
-# Option 2: Create custom transform configuration
-custom_transforms_config = ImageTransformsConfig(
-    enable=True,
-    max_num_transforms=2,
-    random_order=True,
-    tfs={
-        "brightness": ImageTransformConfig(
-            weight=1.0,
-            type="ColorJitter",
-            kwargs={"brightness": (0.7, 1.3)}  # Adjust brightness range
-        ),
-        "contrast": ImageTransformConfig(
-            weight=2.0,  # Higher weight = more likely to be selected
-            type="ColorJitter",
-            kwargs={"contrast": (0.8, 1.2)}
-        ),
-        "sharpness": ImageTransformConfig(
-            weight=0.5,  # Lower weight = less likely to be selected
-            type="SharpnessJitter",
-            kwargs={"sharpness": (0.3, 2.0)}
-        ),
-    }
-)
-
-dataset = LeRobotDataset(
-    repo_id="your-username/your-dataset",
-    image_transforms=ImageTransforms(custom_transforms_config)
-)
-
-# Option 3: Use pure torchvision transforms
-from torchvision.transforms import v2
-
-torchvision_transforms = v2.Compose([
-    v2.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
-    v2.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),
-])
-
-dataset = LeRobotDataset(
-    repo_id="your-username/your-dataset",
-    image_transforms=torchvision_transforms
-)
-```
-
-### Available transform types
-
-LeRobot provides several transform types:
-
-- **`ColorJitter`**: Adjusts brightness, contrast, saturation, and hue
-- **`SharpnessJitter`**: Randomly adjusts image sharpness
-- **`Identity`**: No transformation (useful for testing)
-
-You can also use any `torchvision.transforms.v2` transform by passing it directly to the `image_transforms` parameter.
-
-### Configuration options
-
-- **`enable`**: Enable/disable transforms (default: `False`)
-- **`max_num_transforms`**: Maximum number of transforms applied per frame (default: `3`)
-- **`random_order`**: Apply transforms in random order vs. standard order (default: `False`)
-- **`weight`**: Sampling probability for each transform (higher = more likely, if sum of weights is not 1, they will be normalized)
-- **`kwargs`**: Transform-specific parameters (e.g., brightness range)
-
-### Visualizing transforms
-
-Use the visualization script to preview how transforms affect your data:
-
-```bash
-lerobot-imgtransform-viz \
-  --repo-id=your-username/your-dataset \
-  --output-dir=./transform_examples \
-  --n-examples=5
-```
-
-This saves example images showing the effect of each transform, helping you tune parameters.
-
-### Best practices
-
-- **Start conservative**: Begin with small ranges (e.g., brightness 0.9-1.1) and increase gradually
-- **Test first**: Use the visualization script to ensure transforms look reasonable
-- **Monitor training**: Strong augmentations can hurt performance if too aggressive
-- **Match your domain**: If your robot operates in varying lighting, use brightness/contrast transforms
-- **Combine wisely**: Using too many transforms simultaneously can make training unstable
-
 ## Migrate `v2.1` → `v3.0`
 
 A converter aggregates per‑episode files into larger shards and writes episode offsets/metadata. Convert your dataset using the instructions below.

docs/source/libero.mdx

@@ -33,7 +33,7 @@ To Install LIBERO, after following LeRobot official instructions, just do:
 Evaluate a policy on one LIBERO suite:
 
 ```bash
-lerobot-eval \
+python src/lerobot/scripts/eval.py \
   --policy.path="your-policy-id" \
   --env.type=libero \
   --env.task=libero_object \
@@ -52,7 +52,7 @@ lerobot-eval \
 Benchmark a policy across multiple suites at once:
 
 ```bash
-lerobot-eval \
+python src/lerobot/scripts/eval.py \
   --policy.path="your-policy-id" \
   --env.type=libero \
   --env.task=libero_object,libero_spatial \
@@ -103,11 +103,10 @@ For reference, here is the **original dataset** published by Physical Intelligen
 ### Example training command
 
 ```bash
-lerobot-train \
+python src/lerobot/scripts/train.py \
   --policy.type=smolvla \
   --policy.repo_id=${HF_USER}/libero-test \
-  --policy.load_vlm_weights=true \
-  --dataset.repo_id=HuggingFaceVLA/libero \
+  --dataset.repo_id=jadechoghari/smol-libero3 \
   --env.type=libero \
   --env.task=libero_10 \
   --output_dir=./outputs/ \

docs/source/phone_teleop.mdx

@@ -28,7 +28,7 @@ Links:
 
 ### Phone orientation and controls
 
-- Orientation: hold the phone with the screen facing up and the top edge pointing in the same direction as the robot gripper. This ensures calibration aligns the phone’s frame with the robot frame so motion feels natural, see the image below for reference.
+- Orientation: hold the phone with the screen facing up and the top edge pointing in the same direction as the robot gripper. This ensures calibration aligns the phone’s frame with the robot frame so motion feels natural.
 - Enable/disable:
   - iOS: Hold `B1` to enable teleoperation, release to stop. The first press captures a reference pose.
   - Android: Press and hold the `Move` button, release to stop. The first press captures a reference pose.
@@ -36,8 +36,6 @@ Links:
   - iOS: Analog input `A3` controls the gripper as velocity input.
   - Android: Buttons `A` and `B` act like increment/decrement (A opens, B closes). You can tune velocity in the `GripperVelocityToJoint` step.
 
-<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/phone_teleop.webp" alt="Phone teleop orientation" title="Phone teleop orientation" width="40%">
-
 ### Step 1: Choose the platform
 
 Modify the examples to use `PhoneOS.IOS` or `PhoneOS.ANDROID` in `PhoneConfig`. The API is identical across platforms, only the input source differs. All examples are under `examples/` and have `phone_so100_*.py` variants.
@@ -66,79 +64,80 @@ Run on of the examples scripts to teleoperate, record a dataset, replay a datase
 
 All scripts assume you configured your robot (e.g., SO-100 follower) and set the correct serial port.
 
-Additionally you need to **copy the urdf of the robot to the examples folder**. For the examples in this tutorial (Using SO100/SO101) it is highly recommended to use the urdf in the [SO-ARM100 repo](https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf)
+- Android: after starting the script, open the printed local URL on your phone, tap Start, then press and hold Move.
+- iOS: open HEBI Mobile I/O first; B1 enables motion. A3 controls the gripper.
+
+You can customize mapping or safety limits by editing the processor steps shown in the examples.
+
+You can also remap inputs (e.g., use a different analog input) or adapt the pipeline to other robots (e.g., LeKiwi) by modifying the input and kinematics steps. More about this in the [Processors for Robots and Teleoperators](./processors_robots_teleop.mdx) guide.
 
 - Run this example to teleoperate:
 
   ```bash
-  python examples/phone_to_so100/teleoperate.py
+  python examples/phone_so100_teleop.py
   ```
 
-After running the example:
-
-- Android: after starting the script, open the printed local URL on your phone, tap Start, then press and hold Move.
-- iOS: open HEBI Mobile I/O first; B1 enables motion. A3 controls the gripper.
-
-Additionally you can customize mapping or safety limits by editing the processor steps shown in the examples. You can also remap inputs (e.g., use a different analog input) or adapt the pipeline to other robots (e.g., LeKiwi) by modifying the input and kinematics steps. More about this in the [Processors for Robots and Teleoperators](./processors_robots_teleop.mdx) guide.
-
 - Run this example to record a dataset, which saves absolute end effector observations and actions:
 
   ```bash
-  python examples/phone_to_so100/record.py
+  python examples/phone_so100_record.py
   ```
 
 - Run this example to replay recorded episodes:
 
   ```bash
-  python examples/phone_to_so100/replay.py
+  python examples/phone_so100_replay.py
   ```
 
 - Run this example to evaluate a pretrained policy:
 
   ```bash
-  python examples/phone_to_so100/evaluate.py
+  python examples/phone_so100_eval.py
   ```
 
 ### Important pipeline steps and options
 
 - Kinematics are used in multiple steps. We use [Placo](https://github.com/Rhoban/placo) which is a wrapper around Pinocchio for handling our kinematics. We construct the kinematics object by passing the robot's URDF and target frame. We set `target_frame_name` to the gripper frame.
 
-  ```examples/phone_to_so100/teleoperate.py
-  kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
+  ```44:49:examples/phone_so100_teleop.py
+  RobotKinematics(
+      urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
+      target_frame_name="gripper_frame_link",
+      joint_names=list(robot.bus.motors.keys()),
   )
-
   ```
 
 - The `MapPhoneActionToRobotAction` step converts the calibrated phone pose and inputs into target deltas and gripper commands, below is shown what the step outputs.
 
-  ```src/lerobot/teleoperators/phone/phone_processor.py
-  action["enabled"] = enabled
-        action["target_x"] = -pos[1] if enabled else 0.0
-        action["target_y"] = pos[0] if enabled else 0.0
-        action["target_z"] = pos[2] if enabled else 0.0
-        action["target_wx"] = rotvec[1] if enabled else 0.0
-        action["target_wy"] = rotvec[0] if enabled else 0.0
-        action["target_wz"] = -rotvec[2] if enabled else 0.0
-        action["gripper_vel"] = gripper_vel  # Still send gripper action when disabled
+  ```72:83:src/lerobot/teleoperators/phone/phone_processor.py
+  # Map calibrated phone pose to robot targets (enabled gates the motion)
+  act.update(
+      {
+          "action.enabled": enabled,
+          "action.target_x": -pos[1] if enabled else 0.0,
+          "action.target_y":  pos[0] if enabled else 0.0,
+          "action.target_z":  pos[2] if enabled else 0.0,
+          "action.target_wx": rotvec[1] if enabled else 0.0,
+          "action.target_wy": rotvec[0] if enabled else 0.0,
+          "action.target_wz": -rotvec[2] if enabled else 0.0,
+          "action.gripper":   gripper,
+      }
+  )
   ```
 
 - The `EEReferenceAndDelta` step converts target deltas to an absolute desired EE pose, storing a reference on enable, the `end_effector_step_sizes` are the step sizes for the EE pose and can be modified to change the motion speed.
 
-  ```examples/phone_to_so100/teleoperate.py
+  ```56:65:examples/phone_so100_teleop.py
   EEReferenceAndDelta(
       kinematics=kinematics_solver,
       end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
       motor_names=list(robot.bus.motors.keys()),
-      use_latched_reference=True,
-  ),
+  )
   ```
 
 - The `EEBoundsAndSafety` step clamps EE motion to a workspace and checks for large ee step jumps to ensure safety. The `end_effector_bounds` are the bounds for the EE pose and can be modified to change the workspace. The `max_ee_step_m` and `max_ee_twist_step_rad` are the step limits for the EE pose and can be modified to change the safety limits.
 
-  ```examples/phone_to_so100/teleoperate.py
+  ```61:66:examples/phone_so100_teleop.py
   EEBoundsAndSafety(
       end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
       max_ee_step_m=0.10,
@@ -148,8 +147,11 @@ Additionally you can customize mapping or safety limits by editing the processor
 
 - The `GripperVelocityToJoint` step turns a velocity‑like gripper input into absolute gripper position using the current measured state. The `speed_factor` is the factor by which the velocity is multiplied.
 
-  ```examples/phone_to_so100/teleoperate.py
-  GripperVelocityToJoint(speed_factor=20.0)
+  ```78:81:examples/phone_so100_teleop.py
+  GripperVelocityToJoint(
+      motor_names=list(robot.bus.motors.keys()),
+      speed_factor=20.0,
+  )
   ```
 
 #### Different IK initial guesses
@@ -158,7 +160,7 @@ We use different IK initial guesses in the kinematic steps. As initial guess eit
 
 - Closed loop (used in record/eval): sets `initial_guess_current_joints=True` so IK starts from the measured joints each frame.
 
-  ```examples/phone_to_so100/record.py
+  ```71:76:examples/phone_so100_eval.py
   InverseKinematicsEEToJoints(
       kinematics=kinematics_solver,
       motor_names=list(robot.bus.motors.keys()),
@@ -168,7 +170,7 @@ We use different IK initial guesses in the kinematic steps. As initial guess eit
 
 - Open loop (used in replay): sets `initial_guess_current_joints=False` so IK continues from the previous IK solution rather than the measured state. This preserves action stability when we replay without feedback.
 
-  ```examples/phone_to_so100/replay.py
+  ```80:86:examples/phone_so100_replay.py
   InverseKinematicsEEToJoints(
       kinematics=kinematics_solver,
       motor_names=list(robot.bus.motors.keys()),
@@ -179,6 +181,7 @@ We use different IK initial guesses in the kinematic steps. As initial guess eit
 ### Pipeline steps explained
 
 - MapPhoneActionToRobotAction: converts calibrated phone pose and inputs into target deltas and a gripper command. Motion is gated by an enable signal (B1 on iOS, Move on Android).
+- AddRobotObservationAsComplimentaryData: reads current robot joints and inserts them under `complementary_data.raw_joint_positions` for FK/IK steps to use.
 - EEReferenceAndDelta: latches a reference EE pose on enable and combines it with target deltas to produce an absolute desired EE pose each frame. When disabled, it keeps sending the last commanded pose.
 - EEBoundsAndSafety: clamps the EE pose to a workspace and rate‑limits jumps for safety. Also declares `action.ee.*` features.
 - InverseKinematicsEEToJoints: turns an EE pose into joint positions with IK. `initial_guess_current_joints=True` is recommended for closed‑loop control; set `False` for open‑loop replay for stability.

docs/source/processors_robots_teleop.mdx

@@ -17,7 +17,7 @@ We use the Phone to SO‑100 follower examples for concreteness, but the same pa
 
 The examples in this guide use absolute end effector (EE) poses because they are easy to reason about. In practice, relative EE deltas or joint position are often preferred as learning features.
 
-With processors, you choose the learning features you want to use for your policy. This could be joints positions/velocities, absolute EE, or relative EE positions. You can also choose to store other features, such as joint torques, motor currents, etc.
+You can choose what you save and learn from the teleop and robot action spaces, joints, absolute EE, or relative EE by using/implementing the right steps (and `transform_features()`) in your pipelines.
 
 ## Three pipelines
 
@@ -31,102 +31,99 @@ Each of these pipelines handle different conversions between different action an
 Below is an example of the three pipelines that we use in the phone to SO-100 follower examples:
 
 ```69:90:examples/phone_so100_record.py
-phone_to_robot_ee_pose_processor = RobotProcessorPipeline[RobotAction, RobotAction]( # teleop -> dataset action
-    steps=[
-        MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
-        EEReferenceAndDelta(
-            kinematics=kinematics_solver, end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5}, motor_names=list(robot.bus.motors.keys()),
-        ),
-        EEBoundsAndSafety(
-            end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]}, max_ee_step_m=0.20, max_ee_twist_step_rad=0.50,
-        ),
-        GripperVelocityToJoint(),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
+phone_to_robot_ee_pose = RobotProcessor(  # teleop -> dataset action
+    steps=[MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
+           AddRobotObservationAsComplimentaryData(robot=robot),
+            EEReferenceAndDelta(kinematics=kinematics_solver,
+                               end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
+                               motor_names=list(robot.bus.motors.keys())),
+           EEBoundsAndSafety(end_effector_bounds={"min": [-1, -1, -1], "max": [1, 1, 1]},
+                             max_ee_step_m=0.20, max_ee_twist_step_rad=0.50)],
+    to_transition=to_transition_teleop_action,
+    to_output=lambda tr: tr,
 )
 
-robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction]( # dataset action -> robot
-    steps=[
-        InverseKinematicsEEToJoints(
-            kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()), initial_guess_current_joints=True,
-        ),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
+robot_ee_to_joints = RobotProcessor(      # dataset action -> robot
+    steps=[InverseKinematicsEEToJoints(kinematics=kinematics_solver,
+                                       motor_names=list(robot.bus.motors.keys()),
+                                       initial_guess_current_joints=True),
+           GripperVelocityToJoint(motor_names=list(robot.bus.motors.keys()), speed_factor=20.0)],
+    to_transition=lambda tr: tr,
+    to_output=to_output_robot_action,
 )
 
-robot_joints_to_ee_pose = RobotProcessorPipeline[RobotObservation, RobotObservation]( # robot obs -> dataset obs
-    steps=[
-        ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
-    ],
-    to_transition=observation_to_transition,
-    to_output=transition_to_observation,
+robot_joints_to_ee_pose = RobotProcessor( # robot obs -> dataset obs
+    steps=[ForwardKinematicsJointsToEE(kinematics=kinematics_solver,
+                                       motor_names=list(robot.bus.motors.keys()))],
+    to_transition=to_transition_robot_observation,
+    to_output=lambda tr: tr,
 )
 ```
 
 ## Why to_transition / to_output
 
 To convert from robot/teleoperator to pipeline and back, we use the `to_transition` and `to_output` pipeline adapters.
-They standardize conversions to reduce boilerplate code, and form the bridge between the robot and teleoperators raw dictionaries and the pipeline’s `EnvTransition` format.
+They standardize conversions to reduce boilerplate code, and form the bridge between the robot and teleoperators raw dicts and the pipeline’s `EnvTransition` format.
 In the phone to SO-100 follower examples we use the following adapters:
 
-- `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
-- `transition_to_robot_action`: transforms the pipeline transition to a robot action dict.
-- `observation_to_transition`: transforms the robot observation dict to a pipeline transition.
-- `transition_to_observation`: transforms the pipeline transition to a observation dict.
+- `to_transition_teleop_action`: transforms the teleop action dict to a pipeline transition (puts keys under `action.*`, converts scalars/arrays to tensors, keeps objects like `Rotation` intact)
+- `to_output_robot_action`: transforms the pipeline transition to a robot action dict (extracts keys ending with `.pos`/`.vel` and strips `action.` prefix)
+- `to_transition_robot_observation`: transforms the robot observation dict to a pipeline transition (splits state vs images; stores state under `observation.state.*` and images under `observation.images.*`)
 
-Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
+See `src/lerobot/processor/converters.py` for more details.
 
 ## Dataset feature contracts
 
-Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`.
+Dataset features are the keys saved in the dataset. Each step can declare what its dataset features are via `transform_features(...)`. We can then aggregate features per pipeline with `aggregate_pipeline_dataset_features()` and merge multiple groups with `merge_features(...)`.
 
 Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
 
-```src/lerobot/robots/so100_follower/robot_kinematic_processor.py
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        # We only use the ee pose in the dataset, so we don't need the joint positions
-        for n in self.motor_names:
-            features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
-        # We specify the dataset features of this step that we want to be stored in the dataset
-        for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
-            features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
-                type=FeatureType.STATE, shape=(1,)
-            )
-        return features
+```203:211:src/lerobot/robots/so100_follower/robot_kinematic_processor.py
+def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    # Because this is last step we specify the dataset features of this step that we want to be stored in the dataset
+    features["action.ee.x"] = float
+    features["action.ee.y"] = float
+    features["action.ee.z"] = float
+    features["action.ee.wx"] = float
+    features["action.ee.wy"] = float
+    features["action.ee.wz"] = float
+    return features
 ```
 
-Here we declare what PolicyFeatures we modify in this step, so we know what features we can expect when we run the processor. These features can then be aggregated and used to create the dataset features.
+Tip: declare features at the last step that produces them (e.g., `EEBoundsAndSafety` declares `action.ee.*`, `ForwardKinematicsJointsToEE` declares `observation.state.ee.*`).
 
-Below is an example of how we aggregate and merge features in the phone to SO-100 record example:
+Below is an example of how we aggregate and merge features in the phone to SO-100 follower examples:
 
 ```121:145:examples/phone_so100_record.py
-features=combine_feature_dicts(
-        # Run the feature contract of the pipelines
-        # This tells you how the features would look like after the pipeline steps
-        aggregate_pipeline_dataset_features(
-            pipeline=phone_to_robot_ee_pose_processor,
-            initial_features=create_initial_features(action=phone.action_features), # <- Action features we can expect, these come from our teleop device (phone) and action processor
-            use_videos=True,
-        ),
-        aggregate_pipeline_dataset_features(
-            pipeline=robot_joints_to_ee_pose,
-            initial_features=create_initial_features(observation=robot.observation_features), # <- Observation features we can expect, these come from our robot and observation processor
-            use_videos=True,
-            patterns=["observation.state.ee"], # <- Here you could optionally filter the features we want to store in the dataset, with a specific pattern
+action_ee = aggregate_pipeline_dataset_features(
+    pipeline=phone_to_robot_ee_pose,
+    initial_features=phone.action_features,
+    use_videos=True,
+    patterns=["action.ee"],
+)
 
-        ),
-    ),
+gripper = aggregate_pipeline_dataset_features(
+    pipeline=robot_ee_to_joints,
+    initial_features={},
+    use_videos=True,
+    patterns=["action.gripper.pos", "observation.state.gripper.pos"],
+)
+
+observation_ee = aggregate_pipeline_dataset_features(
+    pipeline=robot_joints_to_ee_pose,
+    initial_features=robot.observation_features,
+    use_videos=True,
+    patterns=["observation.state.ee"],
+)
+
+dataset_features = merge_features(action_ee, gripper, observation_ee)
 ```
 
 How it works:
 
-- `aggregate_pipeline_dataset_features(...)`: applies `transform_features` across the pipeline and filters by patterns (images included when `use_videos=True`, and state features included when `patterns` is specified).
-- `combine_feature_dicts(...)`: combine multiple feature dicts.
-- Recording with `record_loop(...)` uses `build_dataset_frame(...)` to build frames consistent with `dataset.features` before we call `add_frame(...)` to add the frame to the dataset.
+- `aggregate_pipeline_dataset_features(...)`: applies `transform_features` across the pipeline and filters by patterns (images included when `use_videos=True`).
+- `merge_features(...)`: combine multiple feature dicts.
+- Recording uses `to_dataset_frame(...)` to build frames consistent with `dataset.features` before we call `add_frame(...)` to add the frame to the dataset.
 
 ## Guidance when customizing robot pipelines
 
@@ -135,17 +132,17 @@ You can store any of the following features as your action/observation space:
 - Joint positions
 - Absolute EE poses
 - Relative EE deltas
-- Other features: joint velocity, torques, etc.
+- Other features: joint velocity, etc.
 
 Pick what you want to use for your policy action and observation space and configure/modify the pipelines and steps accordingly.
 
 ### Different robots
 
-- You can easily reuse pipelines, for example to use another robot with phone teleop, modify the examples and swap the robot `RobotKinematics` (URDF) and `motor_names` to use your own robot with Phone teleop. Additionally you should ensure `target_frame_name` points to your gripper/wrist.
+- Swap `RobotKinematics` URDF and `motor_names`. Ensure `target_frame_name` points to your gripper/wrist.
 
 ### Safety first
 
 - When changing pipelines, start with tight bounds, implement safety steps when working with real robots.
 - Its advised to start with simulation first and then move to real robots.
 
-Thats it! We hope this guide helps you get started with customizing your robot pipelines, If you run into any issues at any point, jump into our [Discord community](https://discord.com/invite/s3KuuzsPFb) for support.
+Hope this guide helps you get started with customizing your robot pipelines, If you run into any issues at any point, jump into our [Discord community](https://discord.com/invite/s3KuuzsPFb) for support.

docs/source/smolvla.mdx

@@ -29,7 +29,7 @@ SmolVLA is Hugging Face’s lightweight foundation model for robotics. Designed
 ## Collect a dataset
 
 SmolVLA is a base model, so fine-tuning on your own data is required for optimal performance in your setup.
-We recommend recording ~50 episodes of your task as a starting point. Follow our guide to get started: [Recording a Dataset](./il_robots)
+We recommend recording ~50 episodes of your task as a starting point. Follow our guide to get started: [Recording a Dataset](https://huggingface.co/docs/lerobot/getting_started_real_world_robot#record-a-dataset)
 
 <Tip>
 
@@ -93,7 +93,7 @@ lerobot-train --help
 
 ## Evaluate the finetuned model and run it in real-time
 
-Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
+Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./getting_started_real_world_robot#record-a-dataset).
 Once you are logged in, you can run inference in your setup by doing:
 
 ```bash

docs/source/so100.mdx

@@ -634,7 +634,7 @@ leader.disconnect()
 </hfoption>
 </hfoptions>
 
-Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./il_robots)
+Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./getting_started_real_world_robot)
 
 > [!TIP]
 > If you have any questions or need help, please reach out on [Discord](https://discord.com/invite/s3KuuzsPFb).

docs/source/so101.mdx

@@ -430,7 +430,7 @@ leader.disconnect()
 </hfoption>
 </hfoptions>
 
-Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./il_robots)
+Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./getting_started_real_world_robot)
 
 > [!TIP]
 > If you have any questions or need help, please reach out on [Discord](https://discord.com/invite/s3KuuzsPFb).

examples/1_load_lerobot_dataset.py

@@ -136,7 +136,7 @@ print(f"{dataset[0]['action'].shape=}\n")  # (64, c)
 # PyTorch datasets.
 dataloader = torch.utils.data.DataLoader(
     dataset,
-    num_workers=4,
+    num_workers=0,
     batch_size=32,
     shuffle=True,
 )

examples/2_evaluate_pretrained_policy.py

@@ -0,0 +1,139 @@
+# 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.
+
+"""
+This script demonstrates how to evaluate a pretrained policy from the HuggingFace Hub or from your local
+training outputs directory. In the latter case, you might want to run examples/3_train_policy.py first.
+
+It requires the installation of the 'gym_pusht' simulation environment. Install it by running:
+```bash
+pip install -e ".[pusht]"
+```
+"""
+
+from pathlib import Path
+
+import gym_pusht  # noqa: F401
+import gymnasium as gym
+import imageio
+import numpy
+import torch
+
+from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
+
+# Create a directory to store the video of the evaluation
+output_directory = Path("outputs/eval/example_pusht_diffusion")
+output_directory.mkdir(parents=True, exist_ok=True)
+
+# Select your device
+device = "cuda"
+
+# Provide the [hugging face repo id](https://huggingface.co/lerobot/diffusion_pusht):
+pretrained_policy_path = "lerobot/diffusion_pusht"
+# OR a path to a local outputs/train folder.
+# pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")
+
+policy = DiffusionPolicy.from_pretrained(pretrained_policy_path)
+
+# Initialize evaluation environment to render two observation types:
+# an image of the scene and state/position of the agent. The environment
+# also automatically stops running after 300 interactions/steps.
+env = gym.make(
+    "gym_pusht/PushT-v0",
+    obs_type="pixels_agent_pos",
+    max_episode_steps=300,
+)
+
+# We can verify that the shapes of the features expected by the policy match the ones from the observations
+# produced by the environment
+print(policy.config.input_features)
+print(env.observation_space)
+
+# Similarly, we can check that the actions produced by the policy will match the actions expected by the
+# environment
+print(policy.config.output_features)
+print(env.action_space)
+
+# Reset the policy and environments to prepare for rollout
+policy.reset()
+numpy_observation, info = env.reset(seed=42)
+
+# Prepare to collect every rewards and all the frames of the episode,
+# from initial state to final state.
+rewards = []
+frames = []
+
+# Render frame of the initial state
+frames.append(env.render())
+
+step = 0
+done = False
+while not done:
+    # Prepare observation for the policy running in Pytorch
+    state = torch.from_numpy(numpy_observation["agent_pos"])
+    image = torch.from_numpy(numpy_observation["pixels"])
+
+    # Convert to float32 with image from channel first in [0,255]
+    # to channel last in [0,1]
+    state = state.to(torch.float32)
+    image = image.to(torch.float32) / 255
+    image = image.permute(2, 0, 1)
+
+    # Send data tensors from CPU to GPU
+    state = state.to(device, non_blocking=True)
+    image = image.to(device, non_blocking=True)
+
+    # Add extra (empty) batch dimension, required to forward the policy
+    state = state.unsqueeze(0)
+    image = image.unsqueeze(0)
+
+    # Create the policy input dictionary
+    observation = {
+        "observation.state": state,
+        "observation.image": image,
+    }
+
+    # Predict the next action with respect to the current observation
+    with torch.inference_mode():
+        action = policy.select_action(observation)
+
+    # Prepare the action for the environment
+    numpy_action = action.squeeze(0).to("cpu").numpy()
+
+    # Step through the environment and receive a new observation
+    numpy_observation, reward, terminated, truncated, info = env.step(numpy_action)
+    print(f"{step=} {reward=} {terminated=}")
+
+    # Keep track of all the rewards and frames
+    rewards.append(reward)
+    frames.append(env.render())
+
+    # The rollout is considered done when the success state is reached (i.e. terminated is True),
+    # or the maximum number of iterations is reached (i.e. truncated is True)
+    done = terminated | truncated | done
+    step += 1
+
+if terminated:
+    print("Success!")
+else:
+    print("Failure!")
+
+# Get the speed of environment (i.e. its number of frames per second).
+fps = env.metadata["render_fps"]
+
+# Encode all frames into a mp4 video.
+video_path = output_directory / "rollout.mp4"
+imageio.mimsave(str(video_path), numpy.stack(frames), fps=fps)
+
+print(f"Video of the evaluation is available in '{video_path}'.")

examples/3_train_policy.py

@@ -12,7 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""This script demonstrates how to train Diffusion Policy on the PushT environment."""
+"""This script demonstrates how to train Diffusion Policy on the PushT environment.
+
+Once you have trained a model with this script, you can try to evaluate it on
+examples/2_evaluate_pretrained_policy.py
+"""
 
 from pathlib import Path
 
@@ -23,7 +27,6 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetad
 from lerobot.datasets.utils import dataset_to_policy_features
 from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
-from lerobot.policies.factory import make_pre_post_processors
 
 
 def main():
@@ -53,10 +56,9 @@ def main():
     cfg = DiffusionConfig(input_features=input_features, output_features=output_features)
 
     # We can now instantiate our policy with this config and the dataset stats.
-    policy = DiffusionPolicy(cfg)
+    policy = DiffusionPolicy(cfg, dataset_stats=dataset_metadata.stats)
     policy.train()
     policy.to(device)
-    preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
 
     # Another policy-dataset interaction is with the delta_timestamps. Each policy expects a given number frames
     # which can differ for inputs, outputs and rewards (if there are some).
@@ -97,7 +99,7 @@ def main():
     done = False
     while not done:
         for batch in dataloader:
-            batch = preprocessor(batch)
+            batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
             loss, _ = policy.forward(batch)
             loss.backward()
             optimizer.step()
@@ -112,8 +114,6 @@ def main():
 
     # Save a policy checkpoint.
     policy.save_pretrained(output_directory)
-    preprocessor.save_pretrained(output_directory)
-    postprocessor.save_pretrained(output_directory)
 
 
 if __name__ == "__main__":

examples/4_train_policy_with_script.md

@@ -0,0 +1,311 @@
+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 assumes 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.
+
+## The training script
+
+LeRobot offers a training script at [`lerobot/scripts/train.py`](../src/lerobot/scripts/train.py). At a high level it does the following:
+
+- Initialize/load a configuration for the following steps using.
+- Instantiates a dataset.
+- (Optional) Instantiates a simulation environment corresponding to that dataset.
+- Instantiates a policy.
+- Runs a standard training loop with forward pass, backward pass, optimization step, and occasional logging, evaluation (of the policy on the environment), and checkpointing.
+
+## Overview of the configuration system
+
+In the training script, the main function `train` expects a `TrainPipelineConfig` object:
+
+<!-- prettier-ignore-start -->
+```python
+# train.py
+@parser.wrap()
+def train(cfg: TrainPipelineConfig):
+```
+<!-- prettier-ignore-end -->
+
+You can inspect the `TrainPipelineConfig` defined in [`lerobot/configs/train.py`](../src/lerobot/configs/train.py) (which is heavily commented and meant to be a reference to understand any option)
+
+When running the script, inputs for the command line are parsed thanks to the `@parser.wrap()` decorator and an instance of this class is automatically generated. Under the hood, this is done with [Draccus](https://github.com/dlwh/draccus) which is a tool dedicated to this purpose. If you're familiar with Hydra, Draccus can similarly load configurations from config files (.json, .yaml) and also override their values through command line inputs. Unlike Hydra, these configurations are pre-defined in the code through dataclasses rather than being defined entirely in config files. This allows for more rigorous serialization/deserialization, typing, and to manipulate configuration as objects directly in the code and not as dictionaries or namespaces (which enables nice features in an IDE such as autocomplete, jump-to-def, etc.)
+
+Let's have a look at a simplified example. Amongst other attributes, the training config has the following attributes:
+
+<!-- prettier-ignore-start -->
+```python
+@dataclass
+class TrainPipelineConfig:
+    dataset: DatasetConfig
+    env: envs.EnvConfig | None = None
+    policy: PreTrainedConfig | None = None
+```
+<!-- prettier-ignore-end -->
+
+in which `DatasetConfig` for example is defined as such:
+
+<!-- prettier-ignore-start -->
+```python
+@dataclass
+class DatasetConfig:
+    repo_id: str
+    episodes: list[int] | None = None
+    video_backend: str = "pyav"
+```
+<!-- prettier-ignore-end -->
+
+This creates a hierarchical relationship where, for example assuming we have a `cfg` instance of `TrainPipelineConfig`, we can access the `repo_id` value with `cfg.dataset.repo_id`.
+From the command line, we can specify this value by using a very similar syntax `--dataset.repo_id=repo/id`.
+
+By default, every field takes its default value specified in the dataclass. If a field doesn't have a default value, it needs to be specified either from the command line or from a config file – which path is also given in the command line (more in this below). In the example above, the `dataset` field doesn't have a default value which means it must be specified.
+
+## Specifying values from the CLI
+
+Let's say that we want to train [Diffusion Policy](../src/lerobot/policies/diffusion) on the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using the [gym_pusht](https://github.com/huggingface/gym-pusht) environment for evaluation. The command to do so would look like this:
+
+```bash
+lerobot-train \
+    --dataset.repo_id=lerobot/pusht \
+    --policy.type=diffusion \
+    --env.type=pusht
+```
+
+Let's break this down:
+
+- To specify the dataset, we just need to specify its `repo_id` on the hub which is the only required argument in the `DatasetConfig`. The rest of the fields have default values and in this case we are fine with those so we can just add the option `--dataset.repo_id=lerobot/pusht`.
+- To specify the policy, we can just select diffusion policy using `--policy` appended with `.type`. Here, `.type` is a special argument which allows us to select config classes inheriting from `draccus.ChoiceRegistry` and that have been decorated with the `register_subclass()` method. To have a better explanation of this feature, have a look at this [Draccus demo](https://github.com/dlwh/draccus?tab=readme-ov-file#more-flexible-configuration-with-choice-types). In our code, we use this mechanism mainly to select policies, environments, robots, and some other components like optimizers. The policies available to select are located in [lerobot/policies](../src/lerobot/policies)
+- Similarly, we select the environment with `--env.type=pusht`. The different environment configs are available in [`lerobot/envs/configs.py`](../src/lerobot/envs/configs.py)
+
+Let's see another example. Let's say you've been training [ACT](../src/lerobot/policies/act) on [lerobot/aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human) using the [gym-aloha](https://github.com/huggingface/gym-aloha) environment for evaluation with:
+
+```bash
+lerobot-train \
+    --policy.type=act \
+    --dataset.repo_id=lerobot/aloha_sim_insertion_human \
+    --env.type=aloha \
+    --output_dir=outputs/train/act_aloha_insertion
+```
+
+> Notice we added `--output_dir` to explicitly tell where to write outputs from this run (checkpoints, training state, configs etc.). This is not mandatory and if you don't specify it, a default directory will be created from the current date and time, env.type and policy.type. This will typically look like `outputs/train/2025-01-24/16-10-05_aloha_act`.
+
+We now want to train a different policy for aloha on another task. We'll change the dataset and use [lerobot/aloha_sim_transfer_cube_human](https://huggingface.co/datasets/lerobot/aloha_sim_transfer_cube_human) instead. Of course, we also need to change the task of the environment as well to match this other task.
+Looking at the [`AlohaEnv`](../src/lerobot/envs/configs.py) config, the task is `"AlohaInsertion-v0"` by default, which corresponds to the task we trained on in the command above. The [gym-aloha](https://github.com/huggingface/gym-aloha?tab=readme-ov-file#description) environment also has the `AlohaTransferCube-v0` task which corresponds to this other task we want to train on. Putting this together, we can train this new policy on this different task using:
+
+```bash
+lerobot-train \
+    --policy.type=act \
+    --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
+    --env.type=aloha \
+    --env.task=AlohaTransferCube-v0 \
+    --output_dir=outputs/train/act_aloha_transfer
+```
+
+## Loading from a config file
+
+Now, let's assume that we want to reproduce the run just above. That run has produced a `train_config.json` file in its checkpoints, which serializes the `TrainPipelineConfig` instance it used:
+
+```json
+{
+    "dataset": {
+        "repo_id": "lerobot/aloha_sim_transfer_cube_human",
+        "episodes": null,
+        ...
+    },
+    "env": {
+        "type": "aloha",
+        "task": "AlohaTransferCube-v0",
+        "fps": 50,
+        ...
+    },
+    "policy": {
+        "type": "act",
+        "n_obs_steps": 1,
+        ...
+    },
+    ...
+}
+```
+
+We can then simply load the config values from this file using:
+
+```bash
+lerobot-train \
+    --config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
+    --output_dir=outputs/train/act_aloha_transfer_2
+```
+
+`--config_path` is also a special argument which allows to initialize the config from a local config file. It can point to a directory that contains `train_config.json` or to the config file itself directly.
+
+Similarly to Hydra, we can still override some parameters in the CLI if we want to, e.g.:
+
+```bash
+lerobot-train \
+    --config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
+    --output_dir=outputs/train/act_aloha_transfer_2
+    --policy.n_action_steps=80
+```
+
+> Note: While `--output_dir` is not required in general, in this case we need to specify it since it will otherwise take the value from the `train_config.json` (which is `outputs/train/act_aloha_transfer`). In order to prevent accidental deletion of previous run checkpoints, we raise an error if you're trying to write in an existing directory. This is not the case when resuming a run, which is what you'll learn next.
+
+`--config_path` can also accept the repo_id of a repo on the hub that contains a `train_config.json` file, e.g. running:
+
+```bash
+lerobot-train --config_path=lerobot/diffusion_pusht
+```
+
+will start a training run with the same configuration used for training [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht)
+
+## Resume training
+
+Being able to resume a training run is important in case it crashed or aborted for any reason. We'll demonstrate how to do that here.
+
+Let's reuse the command from the previous run and add a few more options:
+
+```bash
+lerobot-train \
+    --policy.type=act \
+    --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
+    --env.type=aloha \
+    --env.task=AlohaTransferCube-v0 \
+    --log_freq=25 \
+    --save_freq=100 \
+    --output_dir=outputs/train/run_resumption
+```
+
+Here we've taken care to set up the log frequency and checkpointing frequency to low numbers so we can showcase resumption. You should be able to see some logging and have a first checkpoint within 1 minute (depending on hardware). Wait for the first checkpoint to happen, you should see a line that looks like this in your terminal:
+
+```
+INFO 2025-01-24 16:10:56 ts/train.py:263 Checkpoint policy after step 100
+```
+
+Now let's simulate a crash by killing the process (hit `ctrl`+`c`). We can then simply resume this run from the last checkpoint available with:
+
+```bash
+lerobot-train \
+    --config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
+    --resume=true
+```
+
+You should see from the logging that your training picks up from where it left off.
+
+Another reason for which you might want to resume a run is simply to extend training and add more training steps. The number of training steps is set by the option `--steps`, which is 100 000 by default.
+You could double the number of steps of the previous run with:
+
+```bash
+lerobot-train \
+    --config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
+    --resume=true \
+    --steps=200000
+```
+
+## Outputs of a run
+
+In the output directory, there will be a folder called `checkpoints` with the following structure:
+
+```bash
+outputs/train/run_resumption/checkpoints
+├── 000100  # checkpoint_dir for training step 100
+│   ├── pretrained_model/
+│   │   ├── config.json  # policy config
+│   │   ├── model.safetensors  # policy weights
+│   │   └── train_config.json  # train config
+│   └── training_state/
+│       ├── optimizer_param_groups.json  #  optimizer param groups
+│       ├── optimizer_state.safetensors  # optimizer state
+│       ├── rng_state.safetensors  # rng states
+│       ├── scheduler_state.json  # scheduler state
+│       └── training_step.json  # training step
+├── 000200
+└── last -> 000200  # symlink to the last available checkpoint
+```
+
+## Fine-tuning a pre-trained policy
+
+In addition to the features currently in Draccus, we've added a special `.path` argument for the policy, which allows to load a policy as you would with `PreTrainedPolicy.from_pretrained()`. In that case, `path` can be a local directory that contains a checkpoint or a repo_id pointing to a pretrained policy on the hub.
+
+For example, we could fine-tune a [policy pre-trained on the aloha transfer task](https://huggingface.co/lerobot/act_aloha_sim_transfer_cube_human) on the aloha insertion task. We can achieve this with:
+
+```bash
+lerobot-train \
+    --policy.path=lerobot/act_aloha_sim_transfer_cube_human \
+    --dataset.repo_id=lerobot/aloha_sim_insertion_human \
+    --env.type=aloha \
+    --env.task=AlohaInsertion-v0
+```
+
+When doing so, keep in mind that the features of the fine-tuning dataset would have to match the input/output features of the pretrained policy.
+
+## Typical logs and metrics
+
+When you start the training process, you will first see your full configuration being printed in the terminal. You can check it to make sure that you configured your run correctly. The final configuration will also be saved with the checkpoint.
+
+After that, you will see training log like this one:
+
+```
+INFO 2024-08-14 13:35:12 ts/train.py:192 step:0 smpl:64 ep:1 epch:0.00 loss:1.112 grdn:15.387 lr:2.0e-07 updt_s:1.738 data_s:4.774
+```
+
+or evaluation log:
+
+```
+INFO 2024-08-14 13:38:45 ts/train.py:226 step:100 smpl:6K ep:52 epch:0.25 ∑rwrd:20.693 success:0.0% eval_s:120.266
+```
+
+These logs will also be saved in wandb if `wandb.enable` is set to `true`. Here are the meaning of some abbreviations:
+
+- `smpl`: number of samples seen during training.
+- `ep`: number of episodes seen during training. An episode contains multiple samples in a complete manipulation task.
+- `epch`: number of time all unique samples are seen (epoch).
+- `grdn`: gradient norm.
+- `∑rwrd`: compute the sum of rewards in every evaluation episode and then take an average of them.
+- `success`: average success rate of eval episodes. Reward and success are usually different except for the sparsing reward setting, where reward=1 only when the task is completed successfully.
+- `eval_s`: time to evaluate the policy in the environment, in second.
+- `updt_s`: time to update the network parameters, in second.
+- `data_s`: time to load a batch of data, in second.
+
+Some metrics are useful for initial performance profiling. For example, if you find the current GPU utilization is low via the `nvidia-smi` command and `data_s` sometimes is too high, you may need to modify batch size or number of dataloading workers to accelerate dataloading. We also recommend [pytorch profiler](https://github.com/huggingface/lerobot?tab=readme-ov-file#improve-your-code-with-profiling) for detailed performance probing.
+
+## In short
+
+We'll summarize here the main use cases to remember from this tutorial.
+
+#### Train a policy from scratch – CLI
+
+```bash
+lerobot-train \
+    --policy.type=act \  # <- select 'act' policy
+    --env.type=pusht \  # <- select 'pusht' environment
+    --dataset.repo_id=lerobot/pusht  # <- train on this dataset
+```
+
+#### Train a policy from scratch - config file + CLI
+
+```bash
+lerobot-train \
+    --config_path=path/to/pretrained_model \  # <- can also be a repo_id
+    --policy.n_action_steps=80  # <- you may still override values
+```
+
+#### Resume/continue a training run
+
+```bash
+lerobot-train \
+    --config_path=checkpoint/pretrained_model/ \
+    --resume=true \
+    --steps=200000  # <- you can change some training parameters
+```
+
+#### Fine-tuning
+
+```bash
+lerobot-train \
+    --policy.path=lerobot/act_aloha_sim_transfer_cube_human \  # <- can also be a local path to a checkpoint
+    --dataset.repo_id=lerobot/aloha_sim_insertion_human \
+    --env.type=aloha \
+    --env.task=AlohaInsertion-v0
+```
+
+---
+
+Now that you know the basics of how to train a policy, you might want to know how to apply this knowledge to actual robots, or how to record your own datasets and train policies on your specific task?
+If that's the case, head over to the next tutorial [`7_get_started_with_real_robot.md`](./7_get_started_with_real_robot.md).
+
+Or in the meantime, happy training! 🤗

examples/5_train_with_streaming.py

@@ -13,20 +13,23 @@
 # limitations under the License.
 
 """This script demonstrates how to train a Diffusion Policy on the PushT environment,
-using a dataset processed in streaming mode."""
+using a dataset processed in streaming mode.
+
+Once you have trained a model with this script, you can try to evaluate it on
+examples/2_evaluate_pretrained_policy.py
+"""
 
 from pathlib import Path
 
 import torch
 
 from lerobot.configs.types import FeatureType
+from lerobot.constants import ACTION
 from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
 from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
 from lerobot.datasets.utils import dataset_to_policy_features
 from lerobot.policies.act.configuration_act import ACTConfig
 from lerobot.policies.act.modeling_act import ACTPolicy
-from lerobot.policies.factory import make_pre_post_processors
-from lerobot.utils.constants import ACTION
 
 
 def main():
@@ -47,7 +50,9 @@ def main():
     training_steps = 10
     log_freq = 1
 
-    dataset_id = "lerobot/droid_1.0.1"  # 26M frames! Would require 4TB of disk space if installed locally (:
+    dataset_id = (
+        "aractingi/droid_1.0.1"  # 26M frames! Would require 4TB of disk space if installed locally (:
+    )
     dataset_metadata = LeRobotDatasetMetadata(dataset_id)
     features = dataset_to_policy_features(dataset_metadata.features)
     output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
@@ -55,10 +60,9 @@ def main():
 
     # We can now instantiate our policy with this config and the dataset stats.
     cfg = ACTConfig(input_features=input_features, output_features=output_features)
-    policy = ACTPolicy(cfg)
+    policy = ACTPolicy(cfg, dataset_stats=dataset_metadata.stats)
     policy.train()
     policy.to(device)
-    preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
 
     # Delta timestamps are used to (1) augment frames used during training and (2) supervise the policy.
     # Here, we use delta-timestamps to only provide ground truth actions for supervision
@@ -85,7 +89,13 @@ def main():
     done = False
     while not done:
         for batch in dataloader:
-            batch = preprocessor(batch)
+            batch = {
+                k: (v.type(torch.float32) if isinstance(v, torch.Tensor) and v.dtype != torch.bool else v)
+                for k, v in batch.items()
+            }
+            batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
+
+            # batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
             loss, _ = policy.forward(batch)
             loss.backward()
             optimizer.step()
@@ -100,8 +110,6 @@ def main():
 
     # Save a policy checkpoint.
     policy.save_pretrained(output_directory)
-    preprocessor.save_pretrained(output_directory)
-    postprocessor.save_pretrained(output_directory)
 
 
 if __name__ == "__main__":

examples/backward_compatibility/replay.py

@@ -44,7 +44,6 @@ from lerobot.robots import (  # noqa: F401
     so100_follower,
     so101_follower,
 )
-from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import (
     init_logging,
@@ -79,16 +78,16 @@ def replay(cfg: ReplayConfig):
 
     robot = make_robot_from_config(cfg.robot)
     dataset = LeRobotDataset(cfg.dataset.repo_id, root=cfg.dataset.root, episodes=[cfg.dataset.episode])
-    actions = dataset.hf_dataset.select_columns(ACTION)
+    actions = dataset.hf_dataset.select_columns("action")
     robot.connect()
 
     log_say("Replaying episode", cfg.play_sounds, blocking=True)
     for idx in range(dataset.num_frames):
         start_episode_t = time.perf_counter()
 
-        action_array = actions[idx][ACTION]
+        action_array = actions[idx]["action"]
         action = {}
-        for i, name in enumerate(dataset.features[ACTION]["names"]):
+        for i, name in enumerate(dataset.features["action"]["names"]):
             key = f"{name.removeprefix('main_')}.pos"
             action[key] = action_array[i].item()
 

examples/dataset/use_dataset_image_transforms.py

@@ -1,177 +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.
-
-"""
-This example demonstrates how to use image transforms with LeRobot datasets for data augmentation during training.
-
-Image transforms are applied to camera frames to improve model robustness and generalization. They are applied
-at training time only, not during dataset recording, allowing you to experiment with different augmentations
-without re-recording data.
-"""
-
-import torch
-from torchvision.transforms import v2
-from torchvision.transforms.functional import to_pil_image
-
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.transforms import ImageTransformConfig, ImageTransforms, ImageTransformsConfig
-
-
-def save_image(tensor, filename):
-    """Helper function to save a tensor as an image file."""
-    if tensor.dim() == 3:  # [C, H, W]
-        if tensor.max() > 1.0:
-            tensor = tensor / 255.0
-        tensor = torch.clamp(tensor, 0.0, 1.0)
-        pil_image = to_pil_image(tensor)
-        pil_image.save(filename)
-        print(f"Saved: {filename}")
-    else:
-        print(f"Skipped {filename}: unexpected tensor shape {tensor.shape}")
-
-
-def example_1_default_transforms():
-    """Example 1: Use default transform configuration and save original vs transformed images"""
-    print("\n Example 1: Default Transform Configuration with Image Saving")
-
-    repo_id = "pepijn223/record_main_0"  # Example dataset
-
-    try:
-        # Load dataset without transforms (original)
-        dataset_original = LeRobotDataset(repo_id=repo_id)
-
-        # Load dataset with transforms enabled
-        transforms_config = ImageTransformsConfig(
-            enable=True,  # Enable transforms (disabled by default)
-            max_num_transforms=2,  # Apply up to 2 transforms per frame
-            random_order=False,  # Apply in standard order
-        )
-        dataset_with_transforms = LeRobotDataset(
-            repo_id=repo_id, image_transforms=ImageTransforms(transforms_config)
-        )
-
-        # Save original and transformed images for comparison
-        if len(dataset_original) > 0:
-            frame_idx = 0  # Use first frame
-            original_sample = dataset_original[frame_idx]
-            transformed_sample = dataset_with_transforms[frame_idx]
-
-            print(f"Saving comparison images (frame {frame_idx}):")
-
-            for cam_key in dataset_original.meta.camera_keys:
-                if cam_key in original_sample and cam_key in transformed_sample:
-                    cam_name = cam_key.replace(".", "_").replace("/", "_")
-
-                    # Save original and transformed images
-                    save_image(original_sample[cam_key], f"{cam_name}_original.png")
-                    save_image(transformed_sample[cam_key], f"{cam_name}_transformed.png")
-
-    except Exception as e:
-        print(f"Could not load dataset '{repo_id}': {e}")
-
-
-def example_2_custom_transforms():
-    """Example 2: Create custom transform configuration and save examples"""
-    print("\n Example 2: Custom Transform Configuration")
-
-    repo_id = "pepijn223/record_main_0"  # Example dataset
-
-    try:
-        # Create custom transform configuration with strong effects
-        custom_transforms_config = ImageTransformsConfig(
-            enable=True,
-            max_num_transforms=2,  # Apply up to 2 transforms per frame
-            random_order=True,  # Apply transforms in random order
-            tfs={
-                "brightness": ImageTransformConfig(
-                    weight=1.0,
-                    type="ColorJitter",
-                    kwargs={"brightness": (0.5, 1.5)},  # Strong brightness range
-                ),
-                "contrast": ImageTransformConfig(
-                    weight=1.0,  # Higher weight = more likely to be selected
-                    type="ColorJitter",
-                    kwargs={"contrast": (0.6, 1.4)},  # Strong contrast
-                ),
-                "sharpness": ImageTransformConfig(
-                    weight=0.5,  # Lower weight = less likely to be selected
-                    type="SharpnessJitter",
-                    kwargs={"sharpness": (0.2, 2.0)},  # Strong sharpness variation
-                ),
-            },
-        )
-
-        dataset_with_custom_transforms = LeRobotDataset(
-            repo_id=repo_id, image_transforms=ImageTransforms(custom_transforms_config)
-        )
-
-        # Save examples with strong transforms
-        if len(dataset_with_custom_transforms) > 0:
-            sample = dataset_with_custom_transforms[0]
-            print("Saving custom transform examples:")
-
-            for cam_key in dataset_with_custom_transforms.meta.camera_keys:
-                if cam_key in sample:
-                    cam_name = cam_key.replace(".", "_").replace("/", "_")
-                    save_image(sample[cam_key], f"{cam_name}_custom_transforms.png")
-
-    except Exception as e:
-        print(f"Could not load dataset '{repo_id}': {e}")
-
-
-def example_3_torchvision_transforms():
-    """Example 3: Use pure torchvision transforms and save examples"""
-    print("\n Example 3: Pure Torchvision Transforms")
-
-    repo_id = "pepijn223/record_main_0"  # Example dataset
-
-    try:
-        # Create torchvision transform pipeline
-        torchvision_transforms = v2.Compose(
-            [
-                v2.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1),
-                v2.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),
-                v2.RandomRotation(degrees=10),  # Small rotation
-            ]
-        )
-
-        dataset_with_torchvision = LeRobotDataset(repo_id=repo_id, image_transforms=torchvision_transforms)
-
-        # Save examples with torchvision transforms
-        if len(dataset_with_torchvision) > 0:
-            sample = dataset_with_torchvision[0]
-            print("Saving torchvision transform examples:")
-
-            for cam_key in dataset_with_torchvision.meta.camera_keys:
-                if cam_key in sample:
-                    cam_name = cam_key.replace(".", "_").replace("/", "_")
-                    save_image(sample[cam_key], f"{cam_name}_torchvision.png")
-
-    except Exception as e:
-        print(f"Could not load dataset '{repo_id}': {e}")
-
-
-def main():
-    """Run all examples"""
-    print("LeRobot Dataset Image Transforms Examples")
-
-    example_1_default_transforms()
-    example_2_custom_transforms()
-    example_3_torchvision_transforms()
-
-
-if __name__ == "__main__":
-    main()

examples/lekiwi/evaluate.py

@@ -19,12 +19,11 @@ from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.policies.act.modeling_act import ACTPolicy
 from lerobot.policies.factory import make_pre_post_processors
 from lerobot.processor import make_default_processors
+from lerobot.record import record_loop
 from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
-from lerobot.scripts.lerobot_record import record_loop
-from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 
 NUM_EPISODES = 2
 FPS = 30
@@ -42,8 +41,8 @@ robot = LeKiwiClient(robot_config)
 policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
 
 # Configure the dataset features
-action_features = hw_to_dataset_features(robot.action_features, ACTION)
-obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
+action_features = hw_to_dataset_features(robot.action_features, "action")
+obs_features = hw_to_dataset_features(robot.observation_features, "observation")
 dataset_features = {**action_features, **obs_features}
 
 # Create the dataset
@@ -61,8 +60,6 @@ preprocessor, postprocessor = make_pre_post_processors(
     policy_cfg=policy,
     pretrained_path=HF_MODEL_ID,
     dataset_stats=dataset.meta.stats,
-    # The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
-    preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
 )
 
 # Connect the robot
@@ -74,7 +71,7 @@ teleop_action_processor, robot_action_processor, robot_observation_processor = m
 
 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="lekiwi_evaluate")
+_init_rerun(session_name="lekiwi_evaluate")
 
 if not robot.is_connected:
     raise ValueError("Robot is not connected!")

examples/lekiwi/record.py

@@ -17,17 +17,16 @@
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.processor import make_default_processors
+from lerobot.record import record_loop
 from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
 from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.keyboard import KeyboardTeleop, KeyboardTeleopConfig
 from lerobot.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
-from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 
-NUM_EPISODES = 2
+NUM_EPISODES = 3
 FPS = 30
 EPISODE_TIME_SEC = 30
 RESET_TIME_SEC = 10
@@ -48,8 +47,8 @@ keyboard = KeyboardTeleop(keyboard_config)
 teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
 
 # Configure the dataset features
-action_features = hw_to_dataset_features(robot.action_features, ACTION)
-obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
+action_features = hw_to_dataset_features(robot.action_features, "action")
+obs_features = hw_to_dataset_features(robot.observation_features, "observation")
 dataset_features = {**action_features, **obs_features}
 
 # Create the dataset
@@ -70,7 +69,7 @@ keyboard.connect()
 
 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="lekiwi_record")
+_init_rerun(session_name="lekiwi_record")
 
 if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
     raise ValueError("Robot or teleop is not connected!")

examples/lekiwi/replay.py

@@ -19,7 +19,6 @@ import time
 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.utils import log_say
 
@@ -33,9 +32,7 @@ robot = LeKiwiClient(robot_config)
 
 # Fetch the dataset to replay
 dataset = LeRobotDataset("<hf_username>/<dataset_repo_id>", episodes=[EPISODE_IDX])
-# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
-episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
-actions = episode_frames.select_columns(ACTION)
+actions = dataset.hf_dataset.select_columns("action")
 
 # Connect to the robot
 robot.connect()
@@ -45,12 +42,12 @@ if not robot.is_connected:
 
 print("Starting replay loop...")
 log_say(f"Replaying episode {EPISODE_IDX}")
-for idx in range(len(episode_frames)):
+for idx in range(dataset.num_frames):
     t0 = time.perf_counter()
 
     # Get recorded action from dataset
     action = {
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
     }
 
     # Send action to robot

examples/lekiwi/teleoperate.py

@@ -20,7 +20,7 @@ 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.visualization_utils import init_rerun, log_rerun_data
+from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data
 
 FPS = 30
 
@@ -41,7 +41,7 @@ leader_arm.connect()
 keyboard.connect()
 
 # Init rerun viewer
-init_rerun(session_name="lekiwi_teleop")
+_init_rerun(session_name="lekiwi_teleop")
 
 if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
     raise ValueError("Robot or teleop is not connected!")

examples/phone_to_so100/evaluate.py

@@ -30,20 +30,21 @@ from lerobot.processor import (
 )
 from lerobot.processor.converters import (
     observation_to_transition,
-    robot_action_observation_to_transition,
+    robot_action_to_transition,
     transition_to_observation,
     transition_to_robot_action,
 )
+from lerobot.record import record_loop
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
+    AddRobotObservationAsComplimentaryData,
     ForwardKinematicsJointsToEE,
     InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 
 NUM_EPISODES = 5
 FPS = 30
@@ -68,21 +69,22 @@ policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
 
 # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
 kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
+    urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
     target_frame_name="gripper_frame_link",
     joint_names=list(robot.bus.motors.keys()),
 )
 
 # Build pipeline to convert EE action to joints action
-robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
     steps=[
+        AddRobotObservationAsComplimentaryData(robot=robot),
         InverseKinematicsEEToJoints(
             kinematics=kinematics_solver,
             motor_names=list(robot.bus.motors.keys()),
             initial_guess_current_joints=True,
         ),
     ],
-    to_transition=robot_action_observation_to_transition,
+    to_transition=robot_action_to_transition,
     to_output=transition_to_robot_action,
 )
 
@@ -128,8 +130,6 @@ preprocessor, postprocessor = make_pre_post_processors(
     policy_cfg=policy,
     pretrained_path=HF_MODEL_ID,
     dataset_stats=dataset.meta.stats,
-    # The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
-    preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
 )
 
 # Connect the robot
@@ -137,7 +137,7 @@ robot.connect()
 
 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="phone_so100_evaluate")
+_init_rerun(session_name="phone_so100_evaluate")
 
 if not robot.is_connected:
     raise ValueError("Robot is not connected!")

examples/phone_to_so100/record.py

@@ -22,12 +22,14 @@ from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
 from lerobot.processor.converters import (
     observation_to_transition,
-    robot_action_observation_to_transition,
+    robot_action_to_transition,
     transition_to_observation,
     transition_to_robot_action,
 )
+from lerobot.record import record_loop
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
+    AddRobotObservationAsComplimentaryData,
     EEBoundsAndSafety,
     EEReferenceAndDelta,
     ForwardKinematicsJointsToEE,
@@ -35,15 +37,14 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.scripts.lerobot_record import record_loop
 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.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 
-NUM_EPISODES = 2
+NUM_EPISODES = 10
 FPS = 30
 EPISODE_TIME_SEC = 60
 RESET_TIME_SEC = 30
@@ -53,7 +54,7 @@ HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
 # Create the robot and teleoperator configurations
 camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
 robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411",
+    port="/dev/tty.usbmodem58760434471",
     id="my_awesome_follower_arm",
     cameras=camera_config,
     use_degrees=True,
@@ -66,34 +67,34 @@ phone = Phone(teleop_config)
 
 # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
 kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
+    urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
     target_frame_name="gripper_frame_link",
     joint_names=list(robot.bus.motors.keys()),
 )
 
 # 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[RobotAction, RobotAction](
     steps=[
         MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
+        AddRobotObservationAsComplimentaryData(robot=robot),
         EEReferenceAndDelta(
             kinematics=kinematics_solver,
             end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
             motor_names=list(robot.bus.motors.keys()),
-            use_latched_reference=True,
         ),
         EEBoundsAndSafety(
             end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
             max_ee_step_m=0.20,
             max_ee_twist_step_rad=0.50,
         ),
-        GripperVelocityToJoint(speed_factor=20.0),
+        GripperVelocityToJoint(),
     ],
-    to_transition=robot_action_observation_to_transition,
+    to_transition=robot_action_to_transition,
     to_output=transition_to_robot_action,
 )
 
 # Build pipeline to convert EE action to joints action
-robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
     steps=[
         InverseKinematicsEEToJoints(
             kinematics=kinematics_solver,
@@ -101,7 +102,7 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
             initial_guess_current_joints=True,
         ),
     ],
-    to_transition=robot_action_observation_to_transition,
+    to_transition=robot_action_to_transition,
     to_output=transition_to_robot_action,
 )
 
@@ -143,7 +144,7 @@ phone.connect()
 
 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="phone_so100_record")
+_init_rerun(session_name="phone_so100_record")
 
 if not robot.is_connected or not phone.is_connected:
     raise ValueError("Robot or teleop is not connected!")

examples/phone_to_so100/replay.py

@@ -18,17 +18,14 @@ import time
 
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
-from lerobot.processor.converters import (
-    robot_action_observation_to_transition,
-    transition_to_robot_action,
-)
+from lerobot.processor import RobotAction, RobotProcessorPipeline
+from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
+    AddRobotObservationAsComplimentaryData,
     InverseKinematicsEEToJoints,
 )
 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.utils import log_say
 
@@ -37,7 +34,7 @@ HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
 
 # Initialize the robot config
 robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
+    port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", use_degrees=True
 )
 
 # Initialize the robot
@@ -45,29 +42,28 @@ robot = SO100Follower(robot_config)
 
 # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
 kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
+    urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
     target_frame_name="gripper_frame_link",
     joint_names=list(robot.bus.motors.keys()),
 )
 
 # Build pipeline to convert EE action to joints action
-robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
     steps=[
+        AddRobotObservationAsComplimentaryData(robot=robot),
         InverseKinematicsEEToJoints(
             kinematics=kinematics_solver,
             motor_names=list(robot.bus.motors.keys()),
             initial_guess_current_joints=False,  # Because replay is open loop
         ),
     ],
-    to_transition=robot_action_observation_to_transition,
+    to_transition=robot_action_to_transition,
     to_output=transition_to_robot_action,
 )
 
 # Fetch the dataset to replay
 dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
-# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
-episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
-actions = episode_frames.select_columns(ACTION)
+actions = dataset.hf_dataset.select_columns("action")
 
 # Connect to the robot
 robot.connect()
@@ -77,19 +73,16 @@ if not robot.is_connected:
 
 print("Starting replay loop...")
 log_say(f"Replaying episode {EPISODE_IDX}")
-for idx in range(len(episode_frames)):
+for idx in range(dataset.num_frames):
     t0 = time.perf_counter()
 
     # Get recorded action from dataset
     ee_action = {
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
     }
 
-    # Get robot observation
-    robot_obs = robot.get_observation()
-
     # Dataset EE -> robot joints
-    joint_action = robot_ee_to_joints_processor((ee_action, robot_obs))
+    joint_action = robot_ee_to_joints_processor(ee_action)
 
     # Send action to robot
     _ = robot.send_action(joint_action)

examples/phone_to_so100/teleoperate.py

@@ -16,13 +16,11 @@
 import time
 
 from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
-from lerobot.processor.converters import (
-    robot_action_observation_to_transition,
-    transition_to_robot_action,
-)
+from lerobot.processor import RobotAction, RobotProcessorPipeline
+from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
+    AddRobotObservationAsComplimentaryData,
     EEBoundsAndSafety,
     EEReferenceAndDelta,
     GripperVelocityToJoint,
@@ -33,13 +31,13 @@ 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.visualization_utils import init_rerun, log_rerun_data
+from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data
 
 FPS = 30
 
 # Initialize the robot and teleoperator
 robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
+    port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", use_degrees=True
 )
 teleop_config = PhoneConfig(phone_os=PhoneOS.IOS)  # or PhoneOS.ANDROID
 
@@ -49,20 +47,20 @@ teleop_device = Phone(teleop_config)
 
 # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
 kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
+    urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
     target_frame_name="gripper_frame_link",
     joint_names=list(robot.bus.motors.keys()),
 )
 
 # Build pipeline to convert phone action to ee pose action to joint action
-phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+phone_to_robot_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
     steps=[
         MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
+        AddRobotObservationAsComplimentaryData(robot=robot),
         EEReferenceAndDelta(
             kinematics=kinematics_solver,
             end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
             motor_names=list(robot.bus.motors.keys()),
-            use_latched_reference=True,
         ),
         EEBoundsAndSafety(
             end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
@@ -75,10 +73,9 @@ phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, Robo
         InverseKinematicsEEToJoints(
             kinematics=kinematics_solver,
             motor_names=list(robot.bus.motors.keys()),
-            initial_guess_current_joints=True,
         ),
     ],
-    to_transition=robot_action_observation_to_transition,
+    to_transition=robot_action_to_transition,
     to_output=transition_to_robot_action,
 )
 
@@ -87,7 +84,7 @@ robot.connect()
 teleop_device.connect()
 
 # Init rerun viewer
-init_rerun(session_name="phone_so100_teleop")
+_init_rerun(session_name="phone_so100_teleop")
 
 if not robot.is_connected or not teleop_device.is_connected:
     raise ValueError("Robot or teleop is not connected!")
@@ -96,14 +93,11 @@ print("Starting teleop loop. Move your phone to teleoperate the robot...")
 while True:
     t0 = time.perf_counter()
 
-    # Get robot observation
-    robot_obs = robot.get_observation()
-
     # Get teleop action
     phone_obs = teleop_device.get_action()
 
     # Phone -> EE pose -> Joints transition
-    joint_action = phone_to_robot_joints_processor((phone_obs, robot_obs))
+    joint_action = phone_to_robot_joints_processor(phone_obs)
 
     # Send action to robot
     _ = robot.send_action(joint_action)

examples/so100_to_so100_EE/evaluate.py

@@ -30,20 +30,21 @@ from lerobot.processor import (
 )
 from lerobot.processor.converters import (
     observation_to_transition,
-    robot_action_observation_to_transition,
+    robot_action_to_transition,
     transition_to_observation,
     transition_to_robot_action,
 )
+from lerobot.record import record_loop
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
+    AddRobotObservationAsComplimentaryData,
     ForwardKinematicsJointsToEE,
     InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 
 NUM_EPISODES = 5
 FPS = 30
@@ -68,21 +69,22 @@ policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
 
 # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
 kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
+    urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
     target_frame_name="gripper_frame_link",
     joint_names=list(robot.bus.motors.keys()),
 )
 
 # Build pipeline to convert EE action to joints action
-robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
     steps=[
+        AddRobotObservationAsComplimentaryData(robot=robot),
         InverseKinematicsEEToJoints(
             kinematics=kinematics_solver,
             motor_names=list(robot.bus.motors.keys()),
             initial_guess_current_joints=True,
         ),
     ],
-    to_transition=robot_action_observation_to_transition,
+    to_transition=robot_action_to_transition,
     to_output=transition_to_robot_action,
 )
 
@@ -129,8 +131,6 @@ preprocessor, postprocessor = make_pre_post_processors(
     policy_cfg=policy,
     pretrained_path=HF_MODEL_ID,
     dataset_stats=dataset.meta.stats,
-    # The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
-    preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
 )
 
 # Connect the robot and teleoperator
@@ -138,7 +138,7 @@ robot.connect()
 
 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="so100_so100_evaluate")
+_init_rerun(session_name="so100_so100_evaluate")
 
 if not robot.is_connected:
     raise ValueError("Robot is not connected!")

examples/so100_to_so100_EE/record.py

@@ -23,25 +23,26 @@ from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
 from lerobot.processor.converters import (
     observation_to_transition,
-    robot_action_observation_to_transition,
+    robot_action_to_transition,
     transition_to_observation,
     transition_to_robot_action,
 )
+from lerobot.record import record_loop
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
+    AddRobotObservationAsComplimentaryData,
     EEBoundsAndSafety,
     ForwardKinematicsJointsToEE,
     InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig
 from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 
-NUM_EPISODES = 2
+NUM_EPISODES = 10
 FPS = 30
 EPISODE_TIME_SEC = 60
 RESET_TIME_SEC = 30
@@ -61,14 +62,14 @@ leader = SO100Leader(leader_config)
 
 # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
 follower_kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
+    urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
     target_frame_name="gripper_frame_link",
     joint_names=list(follower.bus.motors.keys()),
 )
 
 # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
 leader_kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
+    urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
     target_frame_name="gripper_frame_link",
     joint_names=list(leader.bus.motors.keys()),
 )
@@ -85,19 +86,20 @@ follower_joints_to_ee = RobotProcessorPipeline[RobotObservation, RobotObservatio
 )
 
 # Build pipeline to convert leader joints to EE action
-leader_joints_to_ee = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+leader_joints_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
     steps=[
         ForwardKinematicsJointsToEE(
             kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
         ),
     ],
-    to_transition=robot_action_observation_to_transition,
+    to_transition=robot_action_to_transition,
     to_output=transition_to_robot_action,
 )
 
 # Build pipeline to convert EE action to follower joints
-ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+ee_to_follower_joints = RobotProcessorPipeline[RobotAction, RobotAction](
     [
+        AddRobotObservationAsComplimentaryData(robot=follower),
         EEBoundsAndSafety(
             end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
             max_ee_step_m=0.10,
@@ -106,10 +108,9 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
         InverseKinematicsEEToJoints(
             kinematics=follower_kinematics_solver,
             motor_names=list(follower.bus.motors.keys()),
-            initial_guess_current_joints=True,
         ),
     ],
-    to_transition=robot_action_observation_to_transition,
+    to_transition=robot_action_to_transition,
     to_output=transition_to_robot_action,
 )
 
@@ -143,7 +144,7 @@ follower.connect()
 
 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="recording_phone")
+_init_rerun(session_name="recording_phone")
 
 if not leader.is_connected or not follower.is_connected:
     raise ValueError("Robot or teleop is not connected!")

examples/so100_to_so100_EE/replay.py

@@ -19,17 +19,14 @@ import time
 
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
-from lerobot.processor.converters import (
-    robot_action_observation_to_transition,
-    transition_to_robot_action,
-)
+from lerobot.processor import RobotAction, RobotProcessorPipeline
+from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
+    AddRobotObservationAsComplimentaryData,
     InverseKinematicsEEToJoints,
 )
 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.utils import log_say
 
@@ -46,29 +43,28 @@ robot = SO100Follower(robot_config)
 
 # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
 kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
+    urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
     target_frame_name="gripper_frame_link",
     joint_names=list(robot.bus.motors.keys()),
 )
 
 # Build pipeline to convert EE action to joints action
-robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
     steps=[
+        AddRobotObservationAsComplimentaryData(robot=robot),
         InverseKinematicsEEToJoints(
             kinematics=kinematics_solver,
             motor_names=list(robot.bus.motors.keys()),
             initial_guess_current_joints=False,  # Because replay is open loop
         ),
     ],
-    to_transition=robot_action_observation_to_transition,
+    to_transition=robot_action_to_transition,
     to_output=transition_to_robot_action,
 )
 
 # Fetch the dataset to replay
 dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
-# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
-episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
-actions = episode_frames.select_columns(ACTION)
+actions = dataset.hf_dataset.select_columns("action")
 
 # Connect to the robot
 robot.connect()
@@ -78,19 +74,16 @@ if not robot.is_connected:
 
 print("Starting replay loop...")
 log_say(f"Replaying episode {EPISODE_IDX}")
-for idx in range(len(episode_frames)):
+for idx in range(dataset.num_frames):
     t0 = time.perf_counter()
 
     # Get recorded action from dataset
     ee_action = {
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
     }
 
-    # Get robot observation
-    robot_obs = robot.get_observation()
-
     # Dataset EE -> robot joints
-    joint_action = robot_ee_to_joints_processor((ee_action, robot_obs))
+    joint_action = robot_ee_to_joints_processor(ee_action)
 
     # Send action to robot
     _ = robot.send_action(joint_action)

examples/so100_to_so100_EE/teleoperate.py

@@ -17,14 +17,14 @@
 import time
 
 from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
+from lerobot.processor import RobotAction, RobotProcessorPipeline
 from lerobot.processor.converters import (
-    robot_action_observation_to_transition,
     robot_action_to_transition,
     transition_to_robot_action,
 )
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
+    AddRobotObservationAsComplimentaryData,
     EEBoundsAndSafety,
     ForwardKinematicsJointsToEE,
     InverseKinematicsEEToJoints,
@@ -33,7 +33,7 @@ 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.visualization_utils import init_rerun, log_rerun_data
+from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data
 
 FPS = 30
 
@@ -49,14 +49,14 @@ leader = SO100Leader(leader_config)
 
 # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
 follower_kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
+    urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
     target_frame_name="gripper_frame_link",
     joint_names=list(follower.bus.motors.keys()),
 )
 
 # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
 leader_kinematics_solver = RobotKinematics(
-    urdf_path="./SO101/so101_new_calib.urdf",
+    urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
     target_frame_name="gripper_frame_link",
     joint_names=list(leader.bus.motors.keys()),
 )
@@ -73,8 +73,9 @@ leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
 )
 
 # build pipeline to convert EE action to robot joints
-ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+ee_to_follower_joints = RobotProcessorPipeline[RobotAction, RobotAction](
     [
+        AddRobotObservationAsComplimentaryData(robot=follower),
         EEBoundsAndSafety(
             end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
             max_ee_step_m=0.10,
@@ -83,10 +84,9 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
         InverseKinematicsEEToJoints(
             kinematics=follower_kinematics_solver,
             motor_names=list(follower.bus.motors.keys()),
-            initial_guess_current_joints=False,
         ),
     ],
-    to_transition=robot_action_observation_to_transition,
+    to_transition=robot_action_to_transition,
     to_output=transition_to_robot_action,
 )
 
@@ -95,15 +95,12 @@ follower.connect()
 leader.connect()
 
 # Init rerun viewer
-init_rerun(session_name="so100_so100_EE_teleop")
+_init_rerun(session_name="so100_so100_EE_teleop")
 
 print("Starting teleop loop...")
 while True:
     t0 = time.perf_counter()
 
-    # Get robot observation
-    robot_obs = follower.get_observation()
-
     # Get teleop observation
     leader_joints_obs = leader.get_action()
 
@@ -111,7 +108,7 @@ while True:
     leader_ee_act = leader_to_ee(leader_joints_obs)
 
     # teleop EE -> robot joints
-    follower_joints_act = ee_to_follower_joints((leader_ee_act, robot_obs))
+    follower_joints_act = ee_to_follower_joints(leader_ee_act)
 
     # Send action to robot
     _ = follower.send_action(follower_joints_act)

examples/test.sh

@@ -0,0 +1,58 @@
+#!/bin/bash
+
+# storage / caches
+RAID=/raid/jade
+export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
+export HF_HOME=$RAID/.cache/huggingface
+export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
+export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
+export WANDB_CACHE_DIR=$RAID/.cache/wandb
+export TMPDIR=$RAID/.cache/tmp
+mkdir -p $TMPDIR
+export WANDB_MODE=offline
+export HF_DATASETS_OFFLINE=1
+export HF_HUB_OFFLINE=1
+export TOKENIZERS_PARALLELISM=false
+export MUJOCO_GL=egl
+export CUDA_VISIBLE_DEVICES=2
+
+# CONFIGURATION
+POLICY_PATH="/raid/jade/logs/lerobot/lerobot_2_HuggingFaceVLA_libero_smolvla_lr1e-4bs32steps100000/checkpoints/100000/pretrained_model"
+POLICY_PATH="/raid/jade/models/smolvla_pipe"
+TASK=libero_spatial
+ENV_TYPE="libero"
+BATCH_SIZE=1
+N_EPISODES=1
+# storage / caches
+RAID=/raid/jade
+N_ACTION_STEPS=1
+export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
+export HF_HOME=$RAID/.cache/huggingface
+export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
+export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
+export WANDB_CACHE_DIR=$RAID/.cache/wandb
+export TMPDIR=$RAID/.cache/tmp
+mkdir -p $TMPDIR
+export WANDB_MODE=offline
+# export HF_DATASETS_OFFLINE=1
+# export HF_HUB_OFFLINE=1
+export TOKENIZERS_PARALLELISM=false
+export MUJOCO_GL=egl
+export MUJOCO_GL=egl
+unset HF_HUB_OFFLINE
+# RUN EVALUATION
+python src/lerobot/scripts/eval.py \
+    --policy.path="$POLICY_PATH" \
+    --env.type="$ENV_TYPE" \
+    --eval.batch_size="$BATCH_SIZE" \
+    --eval.n_episodes="$N_EPISODES" \
+    --env.task=$TASK \
+    --env.max_parallel_tasks=10 \
+# python examples/evaluate_libero.py \
+#     --policy_path "$POLICY_PATH" \
+#     --task_suite_name "$TASK" \
+#     --num_steps_wait 10 \
+#     --num_trials_per_task 10 \
+#     --video_out_path "data/libero/videos" \
+#     --device "cuda" \
+#     --seed 7

pyproject.toml

@@ -59,7 +59,7 @@ keywords = ["lerobot", "huggingface", "robotics",  "machine learning", "artifici
 dependencies = [
 
     # Hugging Face dependencies
-    "datasets>=4.0.0",
+    "datasets>=2.19.0,<=3.6.0", # TODO: Bumb dependency
     "diffusers>=0.27.2",
     "huggingface-hub[hf-transfer,cli]>=0.34.2",
 
@@ -121,7 +121,7 @@ phone = ["hebi-py>=2.8.0", "teleop>=0.1.0"]
 # Policies
 pi0 = ["lerobot[transformers-dep]"]
 smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14", "accelerate>=1.7.0", "safetensors>=0.4.3"]
-hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.11", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
+hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.9", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
 
 # Features
 async = ["lerobot[grpcio-dep]", "matplotlib>=3.10.3"]
@@ -135,9 +135,21 @@ video_benchmark = ["scikit-image>=0.23.2", "pandas>=2.2.2"]
 aloha = ["gym-aloha>=0.1.1"]
 pusht = ["gym-pusht>=0.1.5", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead
 xarm = ["gym-xarm>=0.1.1"]
-libero = ["lerobot[transformers-dep]", "libero @ git+https://github.com/huggingface/lerobot-libero.git@main#egg=libero"]
-
-
+libero = [
+    "hydra-core>=1.2,<1.4",
+    "easydict>=1.9",
+    "lerobot[transformers-dep]",
+    "robomimic==0.2.0",
+    "thop>=0.1.0.post2206102148",
+    "robosuite==1.4.0",
+    "bddl==1.0.1",
+    "matplotlib>=3.5.3",
+    "cloudpickle>=2.0.0",
+    "gym>=0.25,<0.27",
+    "future>=0.18.3",
+    "egl_probe @ git+https://github.com/jadechoghari/egl_probe.git#egg=egl_probe",
+    "libero @ git+https://github.com/jadechoghari/LIBERO.git@main#egg=libero",
+]
 # All
 all = [
     "lerobot[dynamixel]",
@@ -158,23 +170,19 @@ all = [
     "lerobot[pusht]",
     "lerobot[xarm]",
     "lerobot[phone]",
-    "lerobot[libero]",
+    "lerobot[libero]"
 ]
 
 [project.scripts]
-lerobot-calibrate="lerobot.scripts.lerobot_calibrate:main"
-lerobot-find-cameras="lerobot.scripts.lerobot_find_cameras:main"
-lerobot-find-port="lerobot.scripts.lerobot_find_port:main"
-lerobot-record="lerobot.scripts.lerobot_record:main"
-lerobot-replay="lerobot.scripts.lerobot_replay:main"
-lerobot-setup-motors="lerobot.scripts.lerobot_setup_motors:main"
-lerobot-teleoperate="lerobot.scripts.lerobot_teleoperate:main"
-lerobot-eval="lerobot.scripts.lerobot_eval:main"
-lerobot-train="lerobot.scripts.lerobot_train:main"
-lerobot-dataset-viz="lerobot.scripts.lerobot_dataset_viz:main"
-lerobot-info="lerobot.scripts.lerobot_info:main"
-lerobot-find-joint-limits="lerobot.scripts.lerobot_find_joint_limits:main"
-lerobot-imgtransform-viz="lerobot.scripts.lerobot_imgtransform_viz:main"
+lerobot-calibrate="lerobot.calibrate:main"
+lerobot-find-cameras="lerobot.find_cameras:main"
+lerobot-find-port="lerobot.find_port:main"
+lerobot-record="lerobot.record:main"
+lerobot-replay="lerobot.replay:main"
+lerobot-setup-motors="lerobot.setup_motors:main"
+lerobot-teleoperate="lerobot.teleoperate:main"
+lerobot-eval="lerobot.scripts.eval:main"
+lerobot-train="lerobot.scripts.train:main"
 
 # ---------------- Tool Configurations ----------------
 [tool.setuptools.packages.find]
@@ -267,83 +275,8 @@ default.extend-ignore-identifiers-re = [
 # color = true
 # paths = ["src/lerobot"]
 
-# TODO: Enable mypy gradually module by module across multiple PRs
-# Uncomment [tool.mypy] first, then uncomment individual module overrides as they get proper type annotations
-
 # [tool.mypy]
 # python_version = "3.10"
 # warn_return_any = true
 # warn_unused_configs = true
 # ignore_missing_imports = false
-# strict = true
-# disallow_untyped_defs = true
-# disallow_incomplete_defs = true
-# check_untyped_defs = true
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.utils.*"
-# # include = "src/lerobot/utils/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.configs.*"
-# # include = "src/lerobot/configs/**/*.py"
-
-# # Data processing modules
-# [[tool.mypy.overrides]]
-# module = "lerobot.processor.*"
-# # include = "src/lerobot/processor/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.datasets.*"
-# # include = "src/lerobot/datasets/**/*.py"
-
-# # Core machine learning modules
-# [[tool.mypy.overrides]]
-# module = "lerobot.optim.*"
-# # include = "src/lerobot/optim/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.model.*"
-# # include = "src/lerobot/model/**/*.py"
-
-# # Hardware interfaces
-# [[tool.mypy.overrides]]
-# module = "lerobot.cameras.*"
-# # include = "src/lerobot/cameras/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.motors.*"
-# # include = "src/lerobot/motors/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.robots.*"
-# # include = "src/lerobot/robots/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.teleoperators.*"
-# # include = "src/lerobot/teleoperators/**/*.py"
-
-# # Complex modules (enable these last)
-# [[tool.mypy.overrides]]
-# module = "lerobot.policies.*"
-# # include = "src/lerobot/policies/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.rl.*"
-# # include = "src/lerobot/rl/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.envs.*"
-# # include = "src/lerobot/envs/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.async_inference.*"
-# # include = "src/lerobot/async_inference/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.transport.*"
-# # include = "src/lerobot/transport/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.scripts.*"
-# # include = "src/lerobot/scripts/**/*.py"

src/lerobot/cameras/opencv/camera_opencv.py

@@ -31,7 +31,7 @@ if platform.system() == "Windows" and "OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"
 import cv2
 import numpy as np
 
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 from ..camera import Camera
 from ..utils import get_cv2_backend, get_cv2_rotation

src/lerobot/cameras/reachy2_camera/reachy2_camera.py

@@ -31,7 +31,7 @@ import numpy as np
 from reachy2_sdk.media.camera import CameraView
 from reachy2_sdk.media.camera_manager import CameraManager
 
-from lerobot.utils.errors import DeviceNotConnectedError
+from lerobot.errors import DeviceNotConnectedError
 
 from ..camera import Camera
 from .configuration_reachy2_camera import ColorMode, Reachy2CameraConfig

src/lerobot/cameras/realsense/camera_realsense.py

@@ -29,7 +29,7 @@ try:
 except Exception as e:
     logging.info(f"Could not import realsense: {e}")
 
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 from ..camera import Camera
 from ..configs import ColorMode

src/lerobot/configs/policies.py

@@ -27,9 +27,9 @@ from huggingface_hub.constants import CONFIG_NAME
 from huggingface_hub.errors import HfHubHTTPError
 
 from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.constants import ACTION, OBS_STATE
 from lerobot.optim.optimizers import OptimizerConfig
 from lerobot.optim.schedulers import LRSchedulerConfig
-from lerobot.utils.constants import ACTION, OBS_STATE
 from lerobot.utils.hub import HubMixin
 from lerobot.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
 
@@ -196,10 +196,11 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
             config = json.load(f)
 
         config.pop("type")
-        with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as f:
+        with tempfile.NamedTemporaryFile("w+") as f:
             json.dump(config, f)
             config_file = f.name
+            f.flush()
 
-        cli_overrides = policy_kwargs.pop("cli_overrides", [])
-        with draccus.config_type("json"):
-            return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)
+            cli_overrides = policy_kwargs.pop("cli_overrides", [])
+            with draccus.config_type("json"):
+                return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)

src/lerobot/constants.py

@@ -17,21 +17,19 @@ from pathlib import Path
 
 from huggingface_hub.constants import HF_HOME
 
-OBS_STR = "observation"
-OBS_PREFIX = OBS_STR + "."
-OBS_ENV_STATE = OBS_STR + ".environment_state"
-OBS_STATE = OBS_STR + ".state"
-OBS_IMAGE = OBS_STR + ".image"
-OBS_IMAGES = OBS_IMAGE + "s"
-OBS_LANGUAGE = OBS_STR + ".language"
-OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
-OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
-
+OBS_ENV_STATE = "observation.environment_state"
+OBS_STATE = "observation.state"
+OBS_IMAGE = "observation.image"
+OBS_IMAGES = "observation.images"
+OBS_LANGUAGE = "observation.language"
 ACTION = "action"
 REWARD = "next.reward"
 TRUNCATED = "next.truncated"
 DONE = "next.done"
 
+OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
+OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
+
 ROBOTS = "robots"
 ROBOT_TYPE = "robot_type"
 TELEOPERATORS = "teleoperators"

src/lerobot/datasets/factory.py

@@ -27,7 +27,6 @@ from lerobot.datasets.lerobot_dataset import (
 )
 from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
 from lerobot.datasets.transforms import ImageTransforms
-from lerobot.utils.constants import ACTION, OBS_PREFIX, REWARD
 
 IMAGENET_STATS = {
     "mean": [[[0.485]], [[0.456]], [[0.406]]],  # (c,1,1)
@@ -55,11 +54,11 @@ def resolve_delta_timestamps(
     """
     delta_timestamps = {}
     for key in ds_meta.features:
-        if key == REWARD and cfg.reward_delta_indices is not None:
+        if key == "next.reward" and cfg.reward_delta_indices is not None:
             delta_timestamps[key] = [i / ds_meta.fps for i in cfg.reward_delta_indices]
-        if key == ACTION and cfg.action_delta_indices is not None:
+        if key == "action" and cfg.action_delta_indices is not None:
             delta_timestamps[key] = [i / ds_meta.fps for i in cfg.action_delta_indices]
-        if key.startswith(OBS_PREFIX) and cfg.observation_delta_indices is not None:
+        if key.startswith("observation.") and cfg.observation_delta_indices is not None:
             delta_timestamps[key] = [i / ds_meta.fps for i in cfg.observation_delta_indices]
 
     if len(delta_timestamps) == 0:

src/lerobot/datasets/lerobot_dataset.py

@@ -31,6 +31,7 @@ import torch.utils
 from huggingface_hub import HfApi, snapshot_download
 from huggingface_hub.errors import RevisionNotFoundError
 
+from lerobot.constants import HF_LEROBOT_HOME
 from lerobot.datasets.compute_stats import aggregate_stats, compute_episode_stats
 from lerobot.datasets.image_writer import AsyncImageWriter, write_image
 from lerobot.datasets.utils import (
@@ -78,7 +79,6 @@ from lerobot.datasets.video_utils import (
     get_video_duration_in_s,
     get_video_info,
 )
-from lerobot.utils.constants import HF_LEROBOT_HOME
 
 CODEBASE_VERSION = "v3.0"
 

src/lerobot/datasets/pipeline_features.py

@@ -17,9 +17,9 @@ from collections.abc import Sequence
 from typing import Any
 
 from lerobot.configs.types import PipelineFeatureType
+from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE
 from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.processor import DataProcessorPipeline
-from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR
 
 
 def create_initial_features(
@@ -92,8 +92,8 @@ def aggregate_pipeline_dataset_features(
 
     # Intermediate storage for categorized and filtered features.
     processed_features: dict[str, dict[str, Any]] = {
-        ACTION: {},
-        OBS_STR: {},
+        "action": {},
+        "observation": {},
     }
     images_token = OBS_IMAGES.split(".")[-1]
 
@@ -125,15 +125,17 @@ def aggregate_pipeline_dataset_features(
             # 3. Add the feature to the appropriate group with a clean name.
             name = strip_prefix(key, PREFIXES_TO_STRIP)
             if is_action:
-                processed_features[ACTION][name] = value
+                processed_features["action"][name] = value
             else:
-                processed_features[OBS_STR][name] = value
+                processed_features["observation"][name] = value
 
     # Convert the processed features into the final dataset format.
     dataset_features = {}
-    if processed_features[ACTION]:
-        dataset_features.update(hw_to_dataset_features(processed_features[ACTION], ACTION, use_videos))
-    if processed_features[OBS_STR]:
-        dataset_features.update(hw_to_dataset_features(processed_features[OBS_STR], OBS_STR, use_videos))
+    if processed_features["action"]:
+        dataset_features.update(hw_to_dataset_features(processed_features["action"], ACTION, use_videos))
+    if processed_features["observation"]:
+        dataset_features.update(
+            hw_to_dataset_features(processed_features["observation"], "observation", use_videos)
+        )
 
     return dataset_features

src/lerobot/datasets/streaming_dataset.py

@@ -21,6 +21,7 @@ import numpy as np
 import torch
 from datasets import load_dataset
 
+from lerobot.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
 from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDatasetMetadata
 from lerobot.datasets.utils import (
     Backtrackable,
@@ -37,7 +38,6 @@ from lerobot.datasets.video_utils import (
     VideoDecoderCache,
     decode_video_frames_torchcodec,
 )
-from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
 
 
 class StreamingLeRobotDataset(torch.utils.data.IterableDataset):

src/lerobot/datasets/utils.py

@@ -43,7 +43,6 @@ from lerobot.datasets.backward_compatibility import (
     BackwardCompatibilityError,
     ForwardCompatibilityError,
 )
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_STR
 from lerobot.utils.utils import is_valid_numpy_dtype_string
 
 DEFAULT_CHUNK_SIZE = 1000  # Max number of files per chunk
@@ -646,14 +645,14 @@ def hw_to_dataset_features(
     }
     cam_fts = {key: shape for key, shape in hw_features.items() if isinstance(shape, tuple)}
 
-    if joint_fts and prefix == ACTION:
+    if joint_fts and prefix == "action":
         features[prefix] = {
             "dtype": "float32",
             "shape": (len(joint_fts),),
             "names": list(joint_fts),
         }
 
-    if joint_fts and prefix == OBS_STR:
+    if joint_fts and prefix == "observation":
         features[f"{prefix}.state"] = {
             "dtype": "float32",
             "shape": (len(joint_fts),),
@@ -729,11 +728,11 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea
             # Backward compatibility for "channel" which is an error introduced in LeRobotDataset v2.0 for ported datasets.
             if names[2] in ["channel", "channels"]:  # (h, w, c) -> (c, h, w)
                 shape = (shape[2], shape[0], shape[1])
-        elif key == OBS_ENV_STATE:
+        elif key == "observation.environment_state":
             type = FeatureType.ENV
-        elif key.startswith(OBS_STR):
+        elif key.startswith("observation"):
             type = FeatureType.STATE
-        elif key.startswith(ACTION):
+        elif key.startswith("action"):
             type = FeatureType.ACTION
         else:
             continue

src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py

@@ -34,7 +34,6 @@ python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
 """
 
 import argparse
-import logging
 import shutil
 from pathlib import Path
 from typing import Any
@@ -47,6 +46,7 @@ from datasets import Dataset, Features, Image
 from huggingface_hub import HfApi, snapshot_download
 from requests import HTTPError
 
+from lerobot.constants import HF_LEROBOT_HOME
 from lerobot.datasets.compute_stats import aggregate_stats
 from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
 from lerobot.datasets.utils import (
@@ -71,8 +71,6 @@ from lerobot.datasets.utils import (
     write_tasks,
 )
 from lerobot.datasets.video_utils import concatenate_video_files, get_video_duration_in_s
-from lerobot.utils.constants import HF_LEROBOT_HOME
-from lerobot.utils.utils import init_logging
 
 V21 = "v2.1"
 
@@ -146,7 +144,6 @@ def legacy_load_tasks(local_dir: Path) -> tuple[dict, dict]:
 
 
 def convert_tasks(root, new_root):
-    logging.info(f"Converting tasks from {root} to {new_root}")
     tasks, _ = legacy_load_tasks(root)
     task_indices = tasks.keys()
     task_strings = tasks.values()
@@ -188,10 +185,7 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
     num_frames = 0
     paths_to_cat = []
     episodes_metadata = []
-
-    logging.info(f"Converting data files from {len(ep_paths)} episodes")
-
-    for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"):
+    for ep_path in ep_paths:
         ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
         ep_num_frames = get_parquet_num_frames(ep_path)
         ep_metadata = {
@@ -215,6 +209,7 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
 
         # Reset for the next file
         size_in_mb = ep_size_in_mb
+        num_frames = ep_num_frames
         paths_to_cat = [ep_path]
 
         chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
@@ -241,8 +236,6 @@ def get_image_keys(root):
 
 
 def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
-    logging.info(f"Converting videos from {root} to {new_root}")
-
     video_keys = get_video_keys(root)
     if len(video_keys) == 0:
         return None
@@ -261,7 +254,7 @@ def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
     episods_metadata = []
     num_cameras = len(video_keys)
     num_episodes = num_eps_per_cam[0]
-    for ep_idx in tqdm.tqdm(range(num_episodes), desc="convert videos"):
+    for ep_idx in range(num_episodes):
         # Sanity check
         ep_ids = [eps_metadata_per_cam[cam_idx][ep_idx]["episode_index"] for cam_idx in range(num_cameras)]
         ep_ids += [ep_idx]
@@ -288,7 +281,6 @@ def convert_videos_of_camera(root: Path, new_root: Path, video_key: str, video_f
     duration_in_s = 0.0
     paths_to_cat = []
     episodes_metadata = []
-
     for ep_path in tqdm.tqdm(ep_paths, desc=f"convert videos of {video_key}"):
         ep_size_in_mb = get_video_size_in_mb(ep_path)
         ep_duration_in_s = get_video_duration_in_s(ep_path)
@@ -382,8 +374,6 @@ def generate_episode_metadata_dict(
 
 
 def convert_episodes_metadata(root, new_root, episodes_metadata, episodes_video_metadata=None):
-    logging.info(f"Converting episodes metadata from {root} to {new_root}")
-
     episodes_legacy_metadata = legacy_load_episodes(root)
     episodes_stats = legacy_load_episodes_stats(root)
 
@@ -414,8 +404,7 @@ def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb):
     info["video_files_size_in_mb"] = video_file_size_in_mb
     info["data_path"] = DEFAULT_DATA_PATH
     info["video_path"] = DEFAULT_VIDEO_PATH
-    info["fps"] = int(info["fps"])
-    logging.info(f"Converting info from {root} to {new_root}")
+    info["fps"] = float(info["fps"])
     for key in info["features"]:
         if info["features"][key]["dtype"] == "video":
             # already has fps in video_info
@@ -480,7 +469,6 @@ def convert_dataset(
 
 
 if __name__ == "__main__":
-    init_logging()
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--repo-id",

src/lerobot/envs/configs.py

@@ -19,9 +19,9 @@ from typing import Any
 import draccus
 
 from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 from lerobot.robots import RobotConfig
 from lerobot.teleoperators.config import TeleoperatorConfig
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 
 
 @dataclass
@@ -31,7 +31,6 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC):
     features: dict[str, PolicyFeature] = field(default_factory=dict)
     features_map: dict[str, str] = field(default_factory=dict)
     max_parallel_tasks: int = 1
-    disable_env_checker: bool = True
 
     @property
     def type(self) -> str:
@@ -53,12 +52,12 @@ class AlohaEnv(EnvConfig):
     render_mode: str = "rgb_array"
     features: dict[str, PolicyFeature] = field(
         default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(14,)),
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(14,)),
         }
     )
     features_map: dict[str, str] = field(
         default_factory=lambda: {
-            ACTION: ACTION,
+            "action": ACTION,
             "agent_pos": OBS_STATE,
             "top": f"{OBS_IMAGE}.top",
             "pixels/top": f"{OBS_IMAGES}.top",
@@ -93,13 +92,13 @@ class PushtEnv(EnvConfig):
     visualization_height: int = 384
     features: dict[str, PolicyFeature] = field(
         default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
             "agent_pos": PolicyFeature(type=FeatureType.STATE, shape=(2,)),
         }
     )
     features_map: dict[str, str] = field(
         default_factory=lambda: {
-            ACTION: ACTION,
+            "action": ACTION,
             "agent_pos": OBS_STATE,
             "environment_state": OBS_ENV_STATE,
             "pixels": OBS_IMAGE,
@@ -135,13 +134,13 @@ class XarmEnv(EnvConfig):
     visualization_height: int = 384
     features: dict[str, PolicyFeature] = field(
         default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(4,)),
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(4,)),
             "pixels": PolicyFeature(type=FeatureType.VISUAL, shape=(84, 84, 3)),
         }
     )
     features_map: dict[str, str] = field(
         default_factory=lambda: {
-            ACTION: ACTION,
+            "action": ACTION,
             "agent_pos": OBS_STATE,
             "pixels": OBS_IMAGE,
         }
@@ -163,71 +162,33 @@ class XarmEnv(EnvConfig):
 
 
 @dataclass
-class ImagePreprocessingConfig:
-    crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
-    resize_size: tuple[int, int] | None = None
+class VideoRecordConfig:
+    """Configuration for video recording in ManiSkill environments."""
+
+    enabled: bool = False
+    record_dir: str = "videos"
+    trajectory_name: str = "trajectory"
 
 
 @dataclass
-class RewardClassifierConfig:
-    """Configuration for reward classification."""
-
-    pretrained_path: str | None = None
-    success_threshold: float = 0.5
-    success_reward: float = 1.0
-
-
-@dataclass
-class InverseKinematicsConfig:
-    """Configuration for inverse kinematics processing."""
-
-    urdf_path: str | None = None
-    target_frame_name: str | None = None
-    end_effector_bounds: dict[str, list[float]] | None = None
-    end_effector_step_sizes: dict[str, float] | None = None
-
-
-@dataclass
-class ObservationConfig:
-    """Configuration for observation processing."""
+class EnvTransformConfig:
+    """Configuration for environment wrappers."""
 
+    # ee_action_space_params: EEActionSpaceConfig = field(default_factory=EEActionSpaceConfig)
+    control_mode: str = "gamepad"
+    display_cameras: bool = False
     add_joint_velocity_to_observation: bool = False
     add_current_to_observation: bool = False
     add_ee_pose_to_observation: bool = False
-    display_cameras: bool = False
-
-
-@dataclass
-class GripperConfig:
-    """Configuration for gripper control and penalties."""
-
-    use_gripper: bool = True
-    gripper_penalty: float = 0.0
-    gripper_penalty_in_reward: bool = False
-
-
-@dataclass
-class ResetConfig:
-    """Configuration for environment reset behavior."""
-
+    crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
+    resize_size: tuple[int, int] | None = None
+    control_time_s: float = 20.0
     fixed_reset_joint_positions: Any | None = None
     reset_time_s: float = 5.0
-    control_time_s: float = 20.0
-    terminate_on_success: bool = True
-
-
-@dataclass
-class HILSerlProcessorConfig:
-    """Configuration for environment processing pipeline."""
-
-    control_mode: str = "gamepad"
-    observation: ObservationConfig | None = None
-    image_preprocessing: ImagePreprocessingConfig | None = None
-    gripper: GripperConfig | None = None
-    reset: ResetConfig | None = None
-    inverse_kinematics: InverseKinematicsConfig | None = None
-    reward_classifier: RewardClassifierConfig | None = None
-    max_gripper_pos: float | None = 100.0
+    use_gripper: bool = True
+    gripper_quantization_threshold: float | None = 0.8
+    gripper_penalty: float = 0.0
+    gripper_penalty_in_reward: bool = False
 
 
 @EnvConfig.register_subclass(name="gym_manipulator")
@@ -237,15 +198,82 @@ class HILSerlRobotEnvConfig(EnvConfig):
 
     robot: RobotConfig | None = None
     teleop: TeleoperatorConfig | None = None
-    processor: HILSerlProcessorConfig = field(default_factory=HILSerlProcessorConfig)
-
+    wrapper: EnvTransformConfig | None = None
+    fps: int = 10
     name: str = "real_robot"
+    mode: str | None = None  # Either "record", "replay", None
+    repo_id: str | None = None
+    dataset_root: str | None = None
+    task: str | None = ""
+    num_episodes: int = 10  # only for record mode
+    episode: int = 0
+    device: str = "cuda"
+    push_to_hub: bool = True
+    pretrained_policy_name_or_path: str | None = None
+    reward_classifier_pretrained_path: str | None = None
+    # For the reward classifier, to record more positive examples after a success
+    number_of_steps_after_success: int = 0
 
     @property
     def gym_kwargs(self) -> dict:
         return {}
 
 
+@EnvConfig.register_subclass("hil")
+@dataclass
+class HILEnvConfig(EnvConfig):
+    """Configuration for the HIL environment."""
+
+    name: str = "PandaPickCube"
+    task: str | None = "PandaPickCubeKeyboard-v0"
+    use_viewer: bool = True
+    gripper_penalty: float = 0.0
+    use_gamepad: bool = True
+    state_dim: int = 18
+    action_dim: int = 4
+    fps: int = 100
+    episode_length: int = 100
+    video_record: VideoRecordConfig = field(default_factory=VideoRecordConfig)
+    features: dict[str, PolicyFeature] = field(
+        default_factory=lambda: {
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(4,)),
+            "observation.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128)),
+            "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(18,)),
+        }
+    )
+    features_map: dict[str, str] = field(
+        default_factory=lambda: {
+            "action": ACTION,
+            "observation.image": OBS_IMAGE,
+            "observation.state": OBS_STATE,
+        }
+    )
+    ################# args from hilserlrobotenv
+    reward_classifier_pretrained_path: str | None = None
+    robot_config: RobotConfig | None = None
+    teleop_config: TeleoperatorConfig | None = None
+    wrapper: EnvTransformConfig | None = None
+    mode: str | None = None  # Either "record", "replay", None
+    repo_id: str | None = None
+    dataset_root: str | None = None
+    num_episodes: int = 10  # only for record mode
+    episode: int = 0
+    device: str = "cuda"
+    push_to_hub: bool = True
+    pretrained_policy_name_or_path: str | None = None
+    # For the reward classifier, to record more positive examples after a success
+    number_of_steps_after_success: int = 0
+    ############################
+
+    @property
+    def gym_kwargs(self) -> dict:
+        return {
+            "use_viewer": self.use_viewer,
+            "use_gamepad": self.use_gamepad,
+            "gripper_penalty": self.gripper_penalty,
+        }
+
+
 @EnvConfig.register_subclass("libero")
 @dataclass
 class LiberoEnv(EnvConfig):
@@ -259,12 +287,12 @@ class LiberoEnv(EnvConfig):
     camera_name_mapping: dict[str, str] | None = (None,)
     features: dict[str, PolicyFeature] = field(
         default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
         }
     )
     features_map: dict[str, str] = field(
         default_factory=lambda: {
-            ACTION: ACTION,
+            "action": ACTION,
             "agent_pos": OBS_STATE,
             "pixels/agentview_image": f"{OBS_IMAGES}.image",
             "pixels/robot0_eye_in_hand_image": f"{OBS_IMAGES}.image2",

src/lerobot/envs/factory.py

@@ -17,7 +17,7 @@ import importlib
 
 import gymnasium as gym
 
-from lerobot.envs.configs import AlohaEnv, EnvConfig, LiberoEnv, PushtEnv, XarmEnv
+from lerobot.envs.configs import AlohaEnv, EnvConfig, HILEnvConfig, LiberoEnv, PushtEnv, XarmEnv
 
 
 def make_env_config(env_type: str, **kwargs) -> EnvConfig:
@@ -27,6 +27,8 @@ def make_env_config(env_type: str, **kwargs) -> EnvConfig:
         return PushtEnv(**kwargs)
     elif env_type == "xarm":
         return XarmEnv(**kwargs)
+    elif env_type == "hil":
+        return HILEnvConfig(**kwargs)
     elif env_type == "libero":
         return LiberoEnv(**kwargs)
     else:
@@ -76,13 +78,14 @@ def make_env(
     try:
         importlib.import_module(package_name)
     except ModuleNotFoundError as e:
-        print(f"{package_name} is not installed. Please install it with `pip install 'lerobot[{cfg.type}]'`")
-        raise e
+        raise ModuleNotFoundError(
+            f'{package_name} is not installed. Install with: pip install "lerobot[{cfg.type}]"'
+        ) from e
 
     gym_handle = f"{package_name}/{cfg.task}"
 
     def _make_one():
-        return gym.make(gym_handle, disable_env_checker=cfg.disable_env_checker, **(cfg.gym_kwargs or {}))
+        return gym.make(gym_handle, disable_env_checker=True, **(cfg.gym_kwargs or {}))
 
     vec = env_cls([_make_one for _ in range(n_envs)])
 

src/lerobot/envs/libero.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 from __future__ import annotations
 
+import math
 import os
 from collections import defaultdict
 from collections.abc import Callable, Iterable, Mapping, Sequence
@@ -28,7 +29,6 @@ import torch
 from gymnasium import spaces
 from libero.libero import benchmark, get_libero_path
 from libero.libero.envs import OffScreenRenderEnv
-from robosuite.utils.transform_utils import quat2axisangle
 
 
 def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
@@ -44,7 +44,7 @@ def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
     return cams
 
 
-def _get_suite(name: str) -> benchmark.Benchmark:
+def _get_suite(name: str) -> Any:
     """Instantiate a LIBERO suite by name with clear validation."""
     bench = benchmark.get_benchmark_dict()
     if name not in bench:
@@ -66,6 +66,33 @@ def _select_task_ids(total_tasks: int, task_ids: Iterable[int] | None) -> list[i
     return ids
 
 
+def quat2axisangle(quat: np.ndarray) -> np.ndarray:
+    """
+    Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55
+
+    Converts quaternion to axis-angle format.
+    Returns a unit vector direction scaled by its angle in radians.
+
+    Args:
+        quat (np.array): (x,y,z,w) vec4 float angles
+
+    Returns:
+        np.array: (ax,ay,az) axis-angle exponential coordinates
+    """
+    # clip quaternion
+    if quat[3] > 1.0:
+        quat[3] = 1.0
+    elif quat[3] < -1.0:
+        quat[3] = -1.0
+
+    den = np.sqrt(1.0 - quat[3] * quat[3])
+    if math.isclose(den, 0.0):
+        # This is (close to) a zero degree rotation, immediately return
+        return np.zeros(3)
+
+    return (quat[:3] * 2.0 * math.acos(quat[3])) / den
+
+
 def get_task_init_states(task_suite: Any, i: int) -> np.ndarray:
     init_states_path = (
         Path(get_libero_path("init_states"))
@@ -83,10 +110,6 @@ def get_libero_dummy_action():
 
 OBS_STATE_DIM = 8
 ACTION_DIM = 7
-AGENT_POS_LOW = -1000.0
-AGENT_POS_HIGH = 1000.0
-ACTION_LOW = -1.0
-ACTION_HIGH = 1.0
 TASK_SUITE_MAX_STEPS: dict[str, int] = {
     "libero_spatial": 280,  # longest training demo has 193 steps
     "libero_object": 280,  # longest training demo has 254 steps
@@ -125,8 +148,8 @@ class LiberoEnv(gym.Env):
         self.visualization_width = visualization_width
         self.visualization_height = visualization_height
         self.init_states = init_states
-        self.camera_name = _parse_camera_names(
-            camera_name
+        self.camera_name = camera_name.split(
+            ","
         )  # agentview_image (main) or robot0_eye_in_hand_image (wrist)
 
         # Map raw camera names to "image1" and "image2".
@@ -176,17 +199,15 @@ class LiberoEnv(gym.Env):
                 {
                     "pixels": spaces.Dict(images),
                     "agent_pos": spaces.Box(
-                        low=AGENT_POS_LOW,
-                        high=AGENT_POS_HIGH,
+                        low=-1000.0,
+                        high=1000.0,
                         shape=(OBS_STATE_DIM,),
                         dtype=np.float64,
                     ),
                 }
             )
 
-        self.action_space = spaces.Box(
-            low=ACTION_LOW, high=ACTION_HIGH, shape=(ACTION_DIM,), dtype=np.float32
-        )
+        self.action_space = spaces.Box(low=-1, high=1, shape=(ACTION_DIM,), dtype=np.float32)
 
     def render(self):
         raw_obs = self._env.env._get_observations()
@@ -291,6 +312,7 @@ def _make_env_fns(
     gym_kwargs: Mapping[str, Any],
 ) -> list[Callable[[], LiberoEnv]]:
     """Build n_envs factory callables for a single (suite, task_id)."""
+    joined_cams = ",".join(camera_names)  # keep backward-compat: downstream expects a string
 
     def _make_env(episode_index: int, **kwargs) -> LiberoEnv:
         local_kwargs = dict(kwargs)
@@ -298,7 +320,7 @@ def _make_env_fns(
             task_suite=suite,
             task_id=task_id,
             task_suite_name=suite_name,
-            camera_name=camera_names,
+            camera_name=joined_cams,
             init_states=init_states,
             episode_index=episode_index,
             **local_kwargs,

src/lerobot/envs/utils.py

@@ -26,7 +26,6 @@ from torch import Tensor
 
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.envs.configs import EnvConfig
-from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 from lerobot.utils.utils import get_channel_first_image_shape
 
 
@@ -42,9 +41,9 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
     return_observations = {}
     if "pixels" in observations:
         if isinstance(observations["pixels"], dict):
-            imgs = {f"{OBS_IMAGES}.{key}": img for key, img in observations["pixels"].items()}
+            imgs = {f"observation.images.{key}": img for key, img in observations["pixels"].items()}
         else:
-            imgs = {OBS_IMAGE: observations["pixels"]}
+            imgs = {"observation.image": observations["pixels"]}
 
         for imgkey, img in imgs.items():
             # TODO(aliberts, rcadene): use transforms.ToTensor()?
@@ -73,13 +72,13 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
         if env_state.dim() == 1:
             env_state = env_state.unsqueeze(0)
 
-        return_observations[OBS_ENV_STATE] = env_state
+        return_observations["observation.environment_state"] = env_state
 
     # TODO(rcadene): enable pixels only baseline with `obs_type="pixels"` in environment by removing
     agent_pos = torch.from_numpy(observations["agent_pos"]).float()
     if agent_pos.dim() == 1:
         agent_pos = agent_pos.unsqueeze(0)
-    return_observations[OBS_STATE] = agent_pos
+    return_observations["observation.state"] = agent_pos
 
     return return_observations
 
@@ -100,7 +99,6 @@ def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]:
 
         policy_key = env_cfg.features_map[key]
         policy_features[policy_key] = feature
-
     return policy_features
 
 

src/lerobot/find_cameras.py

@@ -24,7 +24,7 @@ lerobot-find-cameras
 ```
 """
 
-# NOTE(Steven): RealSense can also be identified/opened as OpenCV cameras. If you know the camera is a RealSense, use the `lerobot-find-cameras realsense` flag to avoid confusion.
+# NOTE(Steven): RealSense can also be identified/opened as OpenCV cameras. If you know the camera is a RealSense, use the `lerobot.find_cameras realsense` flag to avoid confusion.
 # NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
 
 import argparse

src/lerobot/motors/dynamixel/dynamixel.py

@@ -22,7 +22,7 @@ import logging
 from copy import deepcopy
 from enum import Enum
 
-from lerobot.motors.encoding_utils import decode_twos_complement, encode_twos_complement
+from lerobot.utils.encoding_utils import decode_twos_complement, encode_twos_complement
 
 from ..motors_bus import Motor, MotorCalibration, MotorsBus, NameOrID, Value, get_address
 from .tables import (

src/lerobot/motors/feetech/feetech.py

@@ -17,7 +17,7 @@ from copy import deepcopy
 from enum import Enum
 from pprint import pformat
 
-from lerobot.motors.encoding_utils import decode_sign_magnitude, encode_sign_magnitude
+from lerobot.utils.encoding_utils import decode_sign_magnitude, encode_sign_magnitude
 
 from ..motors_bus import Motor, MotorCalibration, MotorsBus, NameOrID, Value, get_address
 from .tables import (

src/lerobot/motors/motors_bus.py

@@ -32,7 +32,7 @@ import serial
 from deepdiff import DeepDiff
 from tqdm import tqdm
 
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.utils.utils import enter_pressed, move_cursor_up
 
 NameOrID: TypeAlias = str | int

src/lerobot/optim/optimizers.py

@@ -22,11 +22,11 @@ import draccus
 import torch
 from safetensors.torch import load_file, save_file
 
-from lerobot.datasets.utils import flatten_dict, unflatten_dict, write_json
-from lerobot.utils.constants import (
+from lerobot.constants import (
     OPTIMIZER_PARAM_GROUPS,
     OPTIMIZER_STATE,
 )
+from lerobot.datasets.utils import flatten_dict, unflatten_dict, write_json
 from lerobot.utils.io_utils import deserialize_json_into_object
 
 

src/lerobot/optim/schedulers.py

@@ -22,8 +22,8 @@ import draccus
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LambdaLR, LRScheduler
 
+from lerobot.constants import SCHEDULER_STATE
 from lerobot.datasets.utils import write_json
-from lerobot.utils.constants import SCHEDULER_STATE
 from lerobot.utils.io_utils import deserialize_json_into_object
 
 

src/lerobot/policies/act/modeling_act.py

@@ -33,9 +33,9 @@ from torch import Tensor, nn
 from torchvision.models._utils import IntermediateLayerGetter
 from torchvision.ops.misc import FrozenBatchNorm2d
 
+from lerobot.constants import ACTION, OBS_IMAGES
 from lerobot.policies.act.configuration_act import ACTConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
 
 
 class ACTPolicy(PreTrainedPolicy):
@@ -394,25 +394,25 @@ class ACT(nn.Module):
             latent dimension.
         """
         if self.config.use_vae and self.training:
-            assert ACTION in batch, (
+            assert "action" in batch, (
                 "actions must be provided when using the variational objective in training mode."
             )
 
-        if OBS_IMAGES in batch:
-            batch_size = batch[OBS_IMAGES][0].shape[0]
+        if "observation.images" in batch:
+            batch_size = batch["observation.images"][0].shape[0]
         else:
-            batch_size = batch[OBS_ENV_STATE].shape[0]
+            batch_size = batch["observation.environment_state"].shape[0]
 
         # Prepare the latent for input to the transformer encoder.
-        if self.config.use_vae and ACTION in batch and self.training:
+        if self.config.use_vae and "action" in batch and self.training:
             # Prepare the input to the VAE encoder: [cls, *joint_space_configuration, *action_sequence].
             cls_embed = einops.repeat(
                 self.vae_encoder_cls_embed.weight, "1 d -> b 1 d", b=batch_size
             )  # (B, 1, D)
             if self.config.robot_state_feature:
-                robot_state_embed = self.vae_encoder_robot_state_input_proj(batch[OBS_STATE])
+                robot_state_embed = self.vae_encoder_robot_state_input_proj(batch["observation.state"])
                 robot_state_embed = robot_state_embed.unsqueeze(1)  # (B, 1, D)
-            action_embed = self.vae_encoder_action_input_proj(batch[ACTION])  # (B, S, D)
+            action_embed = self.vae_encoder_action_input_proj(batch["action"])  # (B, S, D)
 
             if self.config.robot_state_feature:
                 vae_encoder_input = [cls_embed, robot_state_embed, action_embed]  # (B, S+2, D)
@@ -430,7 +430,7 @@ class ACT(nn.Module):
             cls_joint_is_pad = torch.full(
                 (batch_size, 2 if self.config.robot_state_feature else 1),
                 False,
-                device=batch[OBS_STATE].device,
+                device=batch["observation.state"].device,
             )
             key_padding_mask = torch.cat(
                 [cls_joint_is_pad, batch["action_is_pad"]], axis=1
@@ -454,7 +454,7 @@ class ACT(nn.Module):
             mu = log_sigma_x2 = None
             # TODO(rcadene, alexander-soare): remove call to `.to` to speedup forward ; precompute and use buffer
             latent_sample = torch.zeros([batch_size, self.config.latent_dim], dtype=torch.float32).to(
-                batch[OBS_STATE].device
+                batch["observation.state"].device
             )
 
         # Prepare transformer encoder inputs.
@@ -462,16 +462,18 @@ class ACT(nn.Module):
         encoder_in_pos_embed = list(self.encoder_1d_feature_pos_embed.weight.unsqueeze(1))
         # Robot state token.
         if self.config.robot_state_feature:
-            encoder_in_tokens.append(self.encoder_robot_state_input_proj(batch[OBS_STATE]))
+            encoder_in_tokens.append(self.encoder_robot_state_input_proj(batch["observation.state"]))
         # Environment state token.
         if self.config.env_state_feature:
-            encoder_in_tokens.append(self.encoder_env_state_input_proj(batch[OBS_ENV_STATE]))
+            encoder_in_tokens.append(
+                self.encoder_env_state_input_proj(batch["observation.environment_state"])
+            )
 
         if self.config.image_features:
             # For a list of images, the H and W may vary but H*W is constant.
             # NOTE: If modifying this section, verify on MPS devices that
             # gradients remain stable (no explosions or NaNs).
-            for img in batch[OBS_IMAGES]:
+            for img in batch["observation.images"]:
                 cam_features = self.backbone(img)["feature_map"]
                 cam_pos_embed = self.encoder_cam_feat_pos_embed(cam_features).to(dtype=cam_features.dtype)
                 cam_features = self.encoder_img_feat_input_proj(cam_features)

src/lerobot/policies/act/processor_act.py

@@ -17,6 +17,7 @@ from typing import Any
 
 import torch
 
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.act.configuration_act import ACTConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -28,7 +29,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_act_pre_post_processors(
@@ -65,10 +65,10 @@ def make_act_pre_post_processors(
         ),
     ]
     output_steps = [
+        DeviceProcessorStep(device="cpu"),
         UnnormalizerProcessorStep(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
-        DeviceProcessorStep(device="cpu"),
     ]
 
     return (

src/lerobot/policies/diffusion/modeling_diffusion.py

@@ -33,6 +33,7 @@ from diffusers.schedulers.scheduling_ddim import DDIMScheduler
 from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
 from torch import Tensor, nn
 
+from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
 from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import (
@@ -41,7 +42,6 @@ from lerobot.policies.utils import (
     get_output_shape,
     populate_queues,
 )
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
 
 
 class DiffusionPolicy(PreTrainedPolicy):
@@ -81,13 +81,13 @@ class DiffusionPolicy(PreTrainedPolicy):
     def reset(self):
         """Clear observation and action queues. Should be called on `env.reset()`"""
         self._queues = {
-            OBS_STATE: deque(maxlen=self.config.n_obs_steps),
-            ACTION: deque(maxlen=self.config.n_action_steps),
+            "observation.state": deque(maxlen=self.config.n_obs_steps),
+            "action": deque(maxlen=self.config.n_action_steps),
         }
         if self.config.image_features:
-            self._queues[OBS_IMAGES] = deque(maxlen=self.config.n_obs_steps)
+            self._queues["observation.images"] = deque(maxlen=self.config.n_obs_steps)
         if self.config.env_state_feature:
-            self._queues[OBS_ENV_STATE] = deque(maxlen=self.config.n_obs_steps)
+            self._queues["observation.environment_state"] = deque(maxlen=self.config.n_obs_steps)
 
     @torch.no_grad()
     def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
@@ -234,7 +234,7 @@ class DiffusionModel(nn.Module):
         if self.config.image_features:
             if self.config.use_separate_rgb_encoder_per_camera:
                 # Combine batch and sequence dims while rearranging to make the camera index dimension first.
-                images_per_camera = einops.rearrange(batch[OBS_IMAGES], "b s n ... -> n (b s) ...")
+                images_per_camera = einops.rearrange(batch["observation.images"], "b s n ... -> n (b s) ...")
                 img_features_list = torch.cat(
                     [
                         encoder(images)
@@ -249,7 +249,7 @@ class DiffusionModel(nn.Module):
             else:
                 # Combine batch, sequence, and "which camera" dims before passing to shared encoder.
                 img_features = self.rgb_encoder(
-                    einops.rearrange(batch[OBS_IMAGES], "b s n ... -> (b s n) ...")
+                    einops.rearrange(batch["observation.images"], "b s n ... -> (b s n) ...")
                 )
                 # Separate batch dim and sequence dim back out. The camera index dim gets absorbed into the
                 # feature dim (effectively concatenating the camera features).
@@ -275,7 +275,7 @@ class DiffusionModel(nn.Module):
             "observation.environment_state": (B, n_obs_steps, environment_dim)
         }
         """
-        batch_size, n_obs_steps = batch[OBS_STATE].shape[:2]
+        batch_size, n_obs_steps = batch["observation.state"].shape[:2]
         assert n_obs_steps == self.config.n_obs_steps
 
         # Encode image features and concatenate them all together along with the state vector.
@@ -306,10 +306,10 @@ class DiffusionModel(nn.Module):
         }
         """
         # Input validation.
-        assert set(batch).issuperset({OBS_STATE, ACTION, "action_is_pad"})
-        assert OBS_IMAGES in batch or OBS_ENV_STATE in batch
-        n_obs_steps = batch[OBS_STATE].shape[1]
-        horizon = batch[ACTION].shape[1]
+        assert set(batch).issuperset({"observation.state", "action", "action_is_pad"})
+        assert "observation.images" in batch or "observation.environment_state" in batch
+        n_obs_steps = batch["observation.state"].shape[1]
+        horizon = batch["action"].shape[1]
         assert horizon == self.config.horizon
         assert n_obs_steps == self.config.n_obs_steps
 
@@ -317,7 +317,7 @@ class DiffusionModel(nn.Module):
         global_cond = self._prepare_global_conditioning(batch)  # (B, global_cond_dim)
 
         # Forward diffusion.
-        trajectory = batch[ACTION]
+        trajectory = batch["action"]
         # Sample noise to add to the trajectory.
         eps = torch.randn(trajectory.shape, device=trajectory.device)
         # Sample a random noising timestep for each item in the batch.
@@ -338,7 +338,7 @@ class DiffusionModel(nn.Module):
         if self.config.prediction_type == "epsilon":
             target = eps
         elif self.config.prediction_type == "sample":
-            target = batch[ACTION]
+            target = batch["action"]
         else:
             raise ValueError(f"Unsupported prediction type {self.config.prediction_type}")
 

src/lerobot/policies/diffusion/processor_diffusion.py

@@ -18,6 +18,7 @@ from typing import Any
 
 import torch
 
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -29,7 +30,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_diffusion_pre_post_processors(
@@ -73,10 +73,10 @@ def make_diffusion_pre_post_processors(
         ),
     ]
     output_steps = [
+        DeviceProcessorStep(device="cpu"),
         UnnormalizerProcessorStep(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
-        DeviceProcessorStep(device="cpu"),
     ]
     return (
         PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](

src/lerobot/policies/factory.py

@@ -24,6 +24,7 @@ from typing_extensions import Unpack
 
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
 from lerobot.datasets.utils import dataset_to_policy_features
 from lerobot.envs.configs import EnvConfig
@@ -45,7 +46,6 @@ from lerobot.processor.converters import (
     transition_to_batch,
     transition_to_policy_action,
 )
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def get_policy_class(name: str) -> type[PreTrainedPolicy]:
@@ -359,7 +359,6 @@ def make_policy(
         if env_cfg is None:
             raise ValueError("env_cfg cannot be None when ds_meta is not provided")
         features = env_to_policy_features(env_cfg)
-
     cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
     cfg.input_features = {key: ft for key, ft in features.items() if key not in cfg.output_features}
     kwargs["config"] = cfg
@@ -372,10 +371,7 @@ def make_policy(
     else:
         # Make a fresh policy.
         policy = policy_cls(**kwargs)
-
     policy.to(cfg.device)
     assert isinstance(policy, torch.nn.Module)
-
     # policy = torch.compile(policy, mode="reduce-overhead")
-
     return policy

src/lerobot/policies/pi0/configuration_pi0.py

@@ -20,7 +20,6 @@ from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
     CosineDecayWithWarmupSchedulerConfig,
 )
-from lerobot.utils.constants import OBS_IMAGES
 
 
 @PreTrainedConfig.register_subclass("pi0")
@@ -114,7 +113,7 @@ class PI0Config(PreTrainedConfig):
         #     raise ValueError("You must provide at least one image or the environment state among the inputs.")
 
         for i in range(self.empty_cameras):
-            key = f"{OBS_IMAGES}.empty_camera_{i}"
+            key = f"observation.images.empty_camera_{i}"
             empty_camera = PolicyFeature(
                 type=FeatureType.VISUAL,
                 shape=(3, 480, 640),

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

@@ -21,7 +21,6 @@ import torch
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
 from lerobot.policies.factory import make_policy
-from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
 
 
 def display(tensor: torch.Tensor):
@@ -61,26 +60,26 @@ def main():
 
     # Override stats
     dataset_meta = LeRobotDatasetMetadata(dataset_repo_id)
-    dataset_meta.stats[OBS_STATE]["mean"] = torch.tensor(
+    dataset_meta.stats["observation.state"]["mean"] = torch.tensor(
         norm_stats["norm_stats"]["state"]["mean"][:num_motors], dtype=torch.float32
     )
-    dataset_meta.stats[OBS_STATE]["std"] = torch.tensor(
+    dataset_meta.stats["observation.state"]["std"] = torch.tensor(
         norm_stats["norm_stats"]["state"]["std"][:num_motors], dtype=torch.float32
     )
 
     # Create LeRobot batch from Jax
     batch = {}
     for cam_key, uint_chw_array in example["images"].items():
-        batch[f"{OBS_IMAGES}.{cam_key}"] = torch.from_numpy(uint_chw_array) / 255.0
-    batch[OBS_STATE] = torch.from_numpy(example["state"])
-    batch[ACTION] = torch.from_numpy(outputs["actions"])
+        batch[f"observation.images.{cam_key}"] = torch.from_numpy(uint_chw_array) / 255.0
+    batch["observation.state"] = torch.from_numpy(example["state"])
+    batch["action"] = torch.from_numpy(outputs["actions"])
     batch["task"] = example["prompt"]
 
     if model_name == "pi0_aloha_towel":
-        del batch[f"{OBS_IMAGES}.cam_low"]
+        del batch["observation.images.cam_low"]
     elif model_name == "pi0_aloha_sim":
-        batch[f"{OBS_IMAGES}.top"] = batch[f"{OBS_IMAGES}.cam_high"]
-        del batch[f"{OBS_IMAGES}.cam_high"]
+        batch["observation.images.top"] = batch["observation.images.cam_high"]
+        del batch["observation.images.cam_high"]
 
     # Batchify
     for key in batch:
@@ -117,7 +116,7 @@ def main():
         actions.append(action)
 
     actions = torch.stack(actions, dim=1)
-    pi_actions = batch[ACTION]
+    pi_actions = batch["action"]
     print("actions")
     display(actions)
     print()

src/lerobot/policies/pi0/modeling_pi0.py

@@ -57,13 +57,13 @@ import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
 
+from lerobot.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.policies.pi0.paligemma_with_expert import (
     PaliGemmaWithExpertConfig,
     PaliGemmaWithExpertModel,
 )
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.utils.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.utils.utils import get_safe_dtype
 
 

src/lerobot/policies/pi0/processor_pi0.py

@@ -19,6 +19,7 @@ from typing import Any
 import torch
 
 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -34,7 +35,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 @ProcessorStepRegistry.register(name="pi0_new_line_processor")
@@ -146,10 +146,10 @@ def make_pi0_pre_post_processors(
     ]
 
     output_steps: list[ProcessorStep] = [
+        DeviceProcessorStep(device="cpu"),
         UnnormalizerProcessorStep(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
-        DeviceProcessorStep(device="cpu"),
     ]
 
     return (

src/lerobot/policies/pi0fast/configuration_pi0fast.py

@@ -6,7 +6,6 @@ from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
     CosineDecayWithWarmupSchedulerConfig,
 )
-from lerobot.utils.constants import OBS_IMAGES
 
 
 @PreTrainedConfig.register_subclass("pi0fast")
@@ -100,7 +99,7 @@ class PI0FASTConfig(PreTrainedConfig):
 
     def validate_features(self) -> None:
         for i in range(self.empty_cameras):
-            key = f"{OBS_IMAGES}.empty_camera_{i}"
+            key = f"observation.images.empty_camera_{i}"
             empty_camera = PolicyFeature(
                 type=FeatureType.VISUAL,
                 shape=(3, 480, 640),

src/lerobot/policies/pi0fast/modeling_pi0fast.py

@@ -57,9 +57,9 @@ from transformers import AutoProcessor, AutoTokenizer, PaliGemmaForConditionalGe
 from transformers.cache_utils import HybridCache, StaticCache
 from transformers.models.auto import CONFIG_MAPPING
 
+from lerobot.constants import ACTION, OBS_STATE
 from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.utils.constants import ACTION, OBS_STATE
 
 PRECISION = {
     "float16": torch.float16,

src/lerobot/policies/pi0fast/processor_pi0fast.py

@@ -18,6 +18,7 @@ from typing import Any
 
 import torch
 
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -29,7 +30,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_pi0fast_pre_post_processors(
@@ -73,10 +73,10 @@ def make_pi0fast_pre_post_processors(
         ),
     ]
     output_steps = [
+        DeviceProcessorStep(device="cpu"),
         UnnormalizerProcessorStep(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
-        DeviceProcessorStep(device="cpu"),
     ]
     return (
         PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](

src/lerobot/policies/pretrained.py

@@ -246,9 +246,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
             base_model=base_model,
         )
 
-        template_card = (
-            files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text(encoding="utf-8")
-        )
+        template_card = files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text()
         card = ModelCard.from_template(card_data, template_str=template_card)
         card.validate()
         return card

src/lerobot/policies/sac/configuration_sac.py

@@ -19,8 +19,8 @@ from dataclasses import dataclass, field
 
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
+from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE
 from lerobot.optim.optimizers import MultiAdamConfig
-from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE
 
 
 def is_image_feature(key: str) -> bool:
@@ -225,7 +225,7 @@ class SACConfig(PreTrainedConfig):
                 "You must provide either 'observation.state' or an image observation (key starting with 'observation.image') in the input features"
             )
 
-        if ACTION not in self.output_features:
+        if "action" not in self.output_features:
             raise ValueError("You must provide 'action' in the output features")
 
     @property

src/lerobot/policies/sac/modeling_sac.py

@@ -31,7 +31,6 @@ from torch.distributions import MultivariateNormal, TanhTransform, Transform, Tr
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.configuration_sac import SACConfig, is_image_feature
 from lerobot.policies.utils import get_device_from_parameters
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_STATE
 
 DISCRETE_DIMENSION_INDEX = -1  # Gripper is always the last dimension
 
@@ -51,7 +50,7 @@ class SACPolicy(
         self.config = config
 
         # Determine action dimension and initialize all components
-        continuous_action_dim = config.output_features[ACTION].shape[0]
+        continuous_action_dim = config.output_features["action"].shape[0]
         self._init_encoders()
         self._init_critics(continuous_action_dim)
         self._init_actor(continuous_action_dim)
@@ -158,7 +157,7 @@ class SACPolicy(
             The computed loss tensor
         """
         # Extract common components from batch
-        actions: Tensor = batch[ACTION]
+        actions: Tensor = batch["action"]
         observations: dict[str, Tensor] = batch["state"]
         observation_features: Tensor = batch.get("observation_feature")
 
@@ -514,17 +513,17 @@ class SACObservationEncoder(nn.Module):
             )
 
     def _init_state_layers(self) -> None:
-        self.has_env = OBS_ENV_STATE in self.config.input_features
-        self.has_state = OBS_STATE in self.config.input_features
+        self.has_env = "observation.environment_state" in self.config.input_features
+        self.has_state = "observation.state" in self.config.input_features
         if self.has_env:
-            dim = self.config.input_features[OBS_ENV_STATE].shape[0]
+            dim = self.config.input_features["observation.environment_state"].shape[0]
             self.env_encoder = nn.Sequential(
                 nn.Linear(dim, self.config.latent_dim),
                 nn.LayerNorm(self.config.latent_dim),
                 nn.Tanh(),
             )
         if self.has_state:
-            dim = self.config.input_features[OBS_STATE].shape[0]
+            dim = self.config.input_features["observation.state"].shape[0]
             self.state_encoder = nn.Sequential(
                 nn.Linear(dim, self.config.latent_dim),
                 nn.LayerNorm(self.config.latent_dim),
@@ -550,9 +549,9 @@ class SACObservationEncoder(nn.Module):
                 cache = self.get_cached_image_features(obs)
             parts.append(self._encode_images(cache, detach))
         if self.has_env:
-            parts.append(self.env_encoder(obs[OBS_ENV_STATE]))
+            parts.append(self.env_encoder(obs["observation.environment_state"]))
         if self.has_state:
-            parts.append(self.state_encoder(obs[OBS_STATE]))
+            parts.append(self.state_encoder(obs["observation.state"]))
         if parts:
             return torch.cat(parts, dim=-1)
 

src/lerobot/policies/sac/processor_sac.py

@@ -19,6 +19,7 @@ from typing import Any
 
 import torch
 
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.sac.configuration_sac import SACConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -30,7 +31,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_sac_pre_post_processors(
@@ -73,10 +73,10 @@ def make_sac_pre_post_processors(
         ),
     ]
     output_steps = [
+        DeviceProcessorStep(device="cpu"),
         UnnormalizerProcessorStep(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
-        DeviceProcessorStep(device="cpu"),
     ]
     return (
         PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](

src/lerobot/policies/sac/reward_model/configuration_classifier.py

@@ -19,7 +19,6 @@ from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
 from lerobot.optim.optimizers import AdamWConfig, OptimizerConfig
 from lerobot.optim.schedulers import LRSchedulerConfig
-from lerobot.utils.constants import OBS_IMAGE
 
 
 @PreTrainedConfig.register_subclass(name="reward_classifier")
@@ -70,7 +69,7 @@ class RewardClassifierConfig(PreTrainedConfig):
 
     def validate_features(self) -> None:
         """Validate feature configurations."""
-        has_image = any(key.startswith(OBS_IMAGE) for key in self.input_features)
+        has_image = any(key.startswith("observation.image") for key in self.input_features)
         if not has_image:
             raise ValueError(
                 "You must provide an image observation (key starting with 'observation.image') in the input features"

src/lerobot/policies/sac/reward_model/modeling_classifier.py

@@ -19,9 +19,9 @@ import logging
 import torch
 from torch import Tensor, nn
 
+from lerobot.constants import OBS_IMAGE, REWARD
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
-from lerobot.utils.constants import OBS_IMAGE, REWARD
 
 
 class ClassifierOutput:
@@ -59,9 +59,7 @@ class SpatialLearnedEmbeddings(nn.Module):
         super().__init__()
         self.height = height
         self.width = width
-        self.channel = (
-            channel  # TODO(fracapuano): this gives issues with non-square images bc is hardcoded to 4
-        )
+        self.channel = channel
         self.num_features = num_features
 
         self.kernel = nn.Parameter(torch.empty(channel, height, width, num_features))

src/lerobot/policies/smolvla/configuration_smolvla.py

@@ -20,7 +20,6 @@ from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
     CosineDecayWithWarmupSchedulerConfig,
 )
-from lerobot.utils.constants import OBS_IMAGES
 
 
 @PreTrainedConfig.register_subclass("smolvla")
@@ -118,7 +117,7 @@ class SmolVLAConfig(PreTrainedConfig):
 
     def validate_features(self) -> None:
         for i in range(self.empty_cameras):
-            key = f"{OBS_IMAGES}.empty_camera_{i}"
+            key = f"observation.images.empty_camera_{i}"
             empty_camera = PolicyFeature(
                 type=FeatureType.VISUAL,
                 shape=(3, 480, 640),

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -59,13 +59,13 @@ import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
 
+from lerobot.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.smolvla.smolvlm_with_expert import SmolVLMWithExpertModel
 from lerobot.policies.utils import (
     populate_queues,
 )
-from lerobot.utils.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.utils.utils import get_safe_dtype
 
 

src/lerobot/policies/smolvla/processor_smolvla.py

@@ -19,6 +19,7 @@ from typing import Any
 import torch
 
 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -33,7 +34,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_smolvla_pre_post_processors(
@@ -84,10 +84,10 @@ def make_smolvla_pre_post_processors(
         ),
     ]
     output_steps = [
+        DeviceProcessorStep(device="cpu"),
         UnnormalizerProcessorStep(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
-        DeviceProcessorStep(device="cpu"),
     ]
     return (
         PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](

src/lerobot/policies/tdmpc/modeling_tdmpc.py

@@ -35,10 +35,10 @@ import torch.nn as nn
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor
 
+from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_STATE, REWARD
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_PREFIX, OBS_STATE, OBS_STR, REWARD
 
 
 class TDMPCPolicy(PreTrainedPolicy):
@@ -91,13 +91,13 @@ class TDMPCPolicy(PreTrainedPolicy):
         called on `env.reset()`
         """
         self._queues = {
-            OBS_STATE: deque(maxlen=1),
-            ACTION: deque(maxlen=max(self.config.n_action_steps, self.config.n_action_repeats)),
+            "observation.state": deque(maxlen=1),
+            "action": deque(maxlen=max(self.config.n_action_steps, self.config.n_action_repeats)),
         }
         if self.config.image_features:
-            self._queues[OBS_IMAGE] = deque(maxlen=1)
+            self._queues["observation.image"] = deque(maxlen=1)
         if self.config.env_state_feature:
-            self._queues[OBS_ENV_STATE] = deque(maxlen=1)
+            self._queues["observation.environment_state"] = deque(maxlen=1)
         # Previous mean obtained from the cross-entropy method (CEM) used during MPC. It is used to warm start
         # CEM for the next step.
         self._prev_mean: torch.Tensor | None = None
@@ -325,7 +325,7 @@ class TDMPCPolicy(PreTrainedPolicy):
 
         action = batch[ACTION]  # (t, b, action_dim)
         reward = batch[REWARD]  # (t, b)
-        observations = {k: v for k, v in batch.items() if k.startswith(OBS_PREFIX)}
+        observations = {k: v for k, v in batch.items() if k.startswith("observation.")}
 
         # Apply random image augmentations.
         if self.config.image_features and self.config.max_random_shift_ratio > 0:
@@ -387,10 +387,10 @@ class TDMPCPolicy(PreTrainedPolicy):
                 temporal_loss_coeffs
                 * F.mse_loss(z_preds[1:], z_targets, reduction="none").mean(dim=-1)
                 # `z_preds` depends on the current observation and the actions.
-                * ~batch[f"{OBS_STR}.state_is_pad"][0]
+                * ~batch["observation.state_is_pad"][0]
                 * ~batch["action_is_pad"]
                 # `z_targets` depends on the next observation.
-                * ~batch[f"{OBS_STR}.state_is_pad"][1:]
+                * ~batch["observation.state_is_pad"][1:]
             )
             .sum(0)
             .mean()
@@ -403,7 +403,7 @@ class TDMPCPolicy(PreTrainedPolicy):
                 * F.mse_loss(reward_preds, reward, reduction="none")
                 * ~batch["next.reward_is_pad"]
                 # `reward_preds` depends on the current observation and the actions.
-                * ~batch[f"{OBS_STR}.state_is_pad"][0]
+                * ~batch["observation.state_is_pad"][0]
                 * ~batch["action_is_pad"]
             )
             .sum(0)
@@ -419,11 +419,11 @@ class TDMPCPolicy(PreTrainedPolicy):
                     reduction="none",
                 ).sum(0)  # sum over ensemble
                 # `q_preds_ensemble` depends on the first observation and the actions.
-                * ~batch[f"{OBS_STR}.state_is_pad"][0]
+                * ~batch["observation.state_is_pad"][0]
                 * ~batch["action_is_pad"]
                 # q_targets depends on the reward and the next observations.
                 * ~batch["next.reward_is_pad"]
-                * ~batch[f"{OBS_STR}.state_is_pad"][1:]
+                * ~batch["observation.state_is_pad"][1:]
             )
             .sum(0)
             .mean()
@@ -441,7 +441,7 @@ class TDMPCPolicy(PreTrainedPolicy):
                 temporal_loss_coeffs
                 * raw_v_value_loss
                 # `v_targets` depends on the first observation and the actions, as does `v_preds`.
-                * ~batch[f"{OBS_STR}.state_is_pad"][0]
+                * ~batch["observation.state_is_pad"][0]
                 * ~batch["action_is_pad"]
             )
             .sum(0)
@@ -477,7 +477,7 @@ class TDMPCPolicy(PreTrainedPolicy):
             * mse
             * temporal_loss_coeffs
             # `action_preds` depends on the first observation and the actions.
-            * ~batch[f"{OBS_STR}.state_is_pad"][0]
+            * ~batch["observation.state_is_pad"][0]
             * ~batch["action_is_pad"]
         ).mean()
 

src/lerobot/policies/tdmpc/processor_tdmpc.py

@@ -18,6 +18,7 @@ from typing import Any
 
 import torch
 
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -29,7 +30,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_tdmpc_pre_post_processors(
@@ -71,10 +71,10 @@ def make_tdmpc_pre_post_processors(
         ),
     ]
     output_steps = [
+        DeviceProcessorStep(device="cpu"),
         UnnormalizerProcessorStep(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
-        DeviceProcessorStep(device="cpu"),
     ]
     return (
         PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](

src/lerobot/policies/vqbet/modeling_vqbet.py

@@ -27,11 +27,11 @@ import torch.nn.functional as F  # noqa: N812
 import torchvision
 from torch import Tensor, nn
 
+from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.policies.vqbet.vqbet_utils import GPT, ResidualVQ
-from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
 
 # ruff: noqa: N806
 
@@ -133,7 +133,7 @@ class VQBeTPolicy(PreTrainedPolicy):
             batch.pop(ACTION)
         batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
         # NOTE: It's important that this happens after stacking the images into a single key.
-        batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
+        batch["observation.images"] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
         # NOTE: for offline evaluation, we have action in the batch, so we need to pop it out
         if ACTION in batch:
             batch.pop(ACTION)
@@ -340,12 +340,14 @@ class VQBeTModel(nn.Module):
 
     def forward(self, batch: dict[str, Tensor], rollout: bool) -> tuple[dict, dict]:
         # Input validation.
-        assert set(batch).issuperset({OBS_STATE, OBS_IMAGES})
-        batch_size, n_obs_steps = batch[OBS_STATE].shape[:2]
+        assert set(batch).issuperset({"observation.state", "observation.images"})
+        batch_size, n_obs_steps = batch["observation.state"].shape[:2]
         assert n_obs_steps == self.config.n_obs_steps
 
         # Extract image feature (first combine batch and sequence dims).
-        img_features = self.rgb_encoder(einops.rearrange(batch[OBS_IMAGES], "b s n ... -> (b s n) ..."))
+        img_features = self.rgb_encoder(
+            einops.rearrange(batch["observation.images"], "b s n ... -> (b s n) ...")
+        )
         # Separate batch and sequence dims.
         img_features = einops.rearrange(
             img_features, "(b s n) ... -> b s n ...", b=batch_size, s=n_obs_steps, n=self.num_images
@@ -357,7 +359,9 @@ class VQBeTModel(nn.Module):
             img_features
         )  # (batch, obs_step, number of different cameras, projection dims)
         input_tokens = [rgb_tokens[:, :, i] for i in range(rgb_tokens.size(2))]
-        input_tokens.append(self.state_projector(batch[OBS_STATE]))  # (batch, obs_step, projection dims)
+        input_tokens.append(
+            self.state_projector(batch["observation.state"])
+        )  # (batch, obs_step, projection dims)
         input_tokens.append(einops.repeat(self.action_token, "1 1 d -> b n d", b=batch_size, n=n_obs_steps))
         # Interleave tokens by stacking and rearranging.
         input_tokens = torch.stack(input_tokens, dim=2)

src/lerobot/policies/vqbet/processor_vqbet.py

@@ -19,6 +19,7 @@ from typing import Any
 
 import torch
 
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -30,7 +31,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_vqbet_pre_post_processors(
@@ -72,10 +72,10 @@ def make_vqbet_pre_post_processors(
         ),
     ]
     output_steps = [
+        DeviceProcessorStep(device="cpu"),
         UnnormalizerProcessorStep(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
-        DeviceProcessorStep(device="cpu"),
     ]
     return (
         PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](

src/lerobot/processor/__init__.py

@@ -36,10 +36,7 @@ from .factory import (
     make_default_robot_observation_processor,
     make_default_teleop_action_processor,
 )
-from .gym_action_processor import (
-    Numpy2TorchActionProcessorStep,
-    Torch2NumpyActionProcessorStep,
-)
+from .gym_action_processor import Numpy2TorchActionProcessorStep, Torch2NumpyActionProcessorStep
 from .hil_processor import (
     AddTeleopActionAsComplimentaryDataStep,
     AddTeleopEventsAsInfoStep,
@@ -70,10 +67,6 @@ from .pipeline import (
     RobotProcessorPipeline,
     TruncatedProcessorStep,
 )
-from .policy_robot_bridge import (
-    PolicyActionToRobotActionProcessorStep,
-    RobotActionToPolicyActionProcessorStep,
-)
 from .rename_processor import RenameObservationsProcessorStep
 from .tokenizer_processor import TokenizerProcessorStep
 
@@ -123,8 +116,6 @@ __all__ = [
     "RobotProcessorPipeline",
     "TokenizerProcessorStep",
     "Torch2NumpyActionProcessorStep",
-    "RobotActionToPolicyActionProcessorStep",
-    "PolicyActionToRobotActionProcessorStep",
     "transition_to_batch",
     "TransitionKey",
     "TruncatedProcessorStep",

src/lerobot/processor/batch_processor.py

@@ -25,7 +25,7 @@ from dataclasses import dataclass, field
 from torch import Tensor
 
 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
-from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
+from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 
 from .core import EnvTransition, PolicyAction
 from .pipeline import (

src/lerobot/processor/converters.py

@@ -23,8 +23,6 @@ from typing import Any
 import numpy as np
 import torch
 
-from lerobot.utils.constants import ACTION, DONE, OBS_PREFIX, REWARD, TRUNCATED
-
 from .core import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
 
 
@@ -209,37 +207,14 @@ def create_transition(
     }
 
 
-def robot_action_observation_to_transition(
-    action_observation: tuple[RobotAction, RobotObservation],
-) -> EnvTransition:
-    """
-    Convert a raw robot action and observation dictionary into a standardized `EnvTransition`.
-
-    Args:
-        action: The raw action dictionary from a teleoperation device or controller.
-        observation: The raw observation dictionary from the environment.
-
-    Returns:
-        An `EnvTransition` containing the formatted observation.
-    """
-    if not isinstance(action_observation, tuple):
-        raise ValueError("action_observation should be a tuple type with an action and observation")
-
-    action, observation = action_observation
-
-    if action is not None and not isinstance(action, dict):
-        raise ValueError(f"Action should be a RobotAction type got {type(action)}")
-
-    if observation is not None and not isinstance(observation, dict):
-        raise ValueError(f"Observation should be a RobotObservation type got {type(observation)}")
-
-    return create_transition(action=action, observation=observation)
-
-
 def robot_action_to_transition(action: RobotAction) -> EnvTransition:
     """
     Convert a raw robot action dictionary into a standardized `EnvTransition`.
 
+    The keys in the action dictionary are prefixed with "action." and stored under
+    the `ACTION` key in the transition. Values are converted to tensors, except for
+    special types like `Rotation`.
+
     Args:
         action: The raw action dictionary from a teleoperation device or controller.
 
@@ -255,6 +230,10 @@ def observation_to_transition(observation: RobotObservation) -> EnvTransition:
     """
     Convert a raw robot observation dictionary into a standardized `EnvTransition`.
 
+    The observation is split into state and image components. State keys are prefixed
+    with "observation.state." and image keys with "observation.images.". The result is
+    stored under the `OBSERVATION` key in the transition.
+
     Args:
         observation: The raw observation dictionary from the environment.
 
@@ -344,20 +323,20 @@ def batch_to_transition(batch: dict[str, Any]) -> EnvTransition:
     if not isinstance(batch, dict):
         raise ValueError(f"EnvTransition must be a dictionary. Got {type(batch).__name__}")
 
-    action = batch.get(ACTION)
+    action = batch.get("action")
     if action is not None and not isinstance(action, PolicyAction):
         raise ValueError(f"Action should be a PolicyAction type got {type(action)}")
 
     # Extract observation and complementary data keys.
-    observation_keys = {k: v for k, v in batch.items() if k.startswith(OBS_PREFIX)}
+    observation_keys = {k: v for k, v in batch.items() if k.startswith("observation.")}
     complementary_data = _extract_complementary_data(batch)
 
     return create_transition(
         observation=observation_keys if observation_keys else None,
-        action=batch.get(ACTION),
-        reward=batch.get(REWARD, 0.0),
-        done=batch.get(DONE, False),
-        truncated=batch.get(TRUNCATED, False),
+        action=batch.get("action"),
+        reward=batch.get("next.reward", 0.0),
+        done=batch.get("next.done", False),
+        truncated=batch.get("next.truncated", False),
         info=batch.get("info", {}),
         complementary_data=complementary_data if complementary_data else None,
     )
@@ -379,10 +358,10 @@ def transition_to_batch(transition: EnvTransition) -> dict[str, Any]:
         raise ValueError(f"Transition should be a EnvTransition type (dict) got {type(transition)}")
 
     batch = {
-        ACTION: transition.get(TransitionKey.ACTION),
-        REWARD: transition.get(TransitionKey.REWARD, 0.0),
-        DONE: transition.get(TransitionKey.DONE, False),
-        TRUNCATED: transition.get(TransitionKey.TRUNCATED, False),
+        "action": transition.get(TransitionKey.ACTION),
+        "next.reward": transition.get(TransitionKey.REWARD, 0.0),
+        "next.done": transition.get(TransitionKey.DONE, False),
+        "next.truncated": transition.get(TransitionKey.TRUNCATED, False),
         "info": transition.get(TransitionKey.INFO, {}),
     }
 

src/lerobot/processor/delta_action_processor.py

@@ -48,12 +48,12 @@ class MapTensorToDeltaActionDictStep(ActionProcessorStep):
 
         # TODO (maractingi): add rotation
         delta_action = {
-            "delta_x": action[0].item(),
-            "delta_y": action[1].item(),
-            "delta_z": action[2].item(),
+            "delta_x": action[0],
+            "delta_y": action[1],
+            "delta_z": action[2],
         }
         if self.use_gripper:
-            delta_action["gripper"] = action[3].item()
+            delta_action["gripper"] = action[3]
         return delta_action
 
     def transform_features(
@@ -126,7 +126,7 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
             "target_wx": target_wx,
             "target_wy": target_wy,
             "target_wz": target_wz,
-            "gripper_vel": float(gripper),
+            "gripper": float(gripper),
         }
 
         return action

src/lerobot/processor/device_processor.py

@@ -105,10 +105,6 @@ class DeviceProcessorStep(ProcessorStep):
             # In both cases, use the configured device.
             target_device = self.tensor_device
 
-        # MPS workaround: Convert float64 to float32 since MPS doesn't support float64
-        if target_device.type == "mps" and tensor.dtype == torch.float64:
-            tensor = tensor.to(dtype=torch.float32)
-
         # Only move if necessary
         if tensor.device != target_device:
             tensor = tensor.to(target_device, non_blocking=self.non_blocking)

src/lerobot/processor/factory.py

@@ -16,7 +16,7 @@
 
 from .converters import (
     observation_to_transition,
-    robot_action_observation_to_transition,
+    robot_action_to_transition,
     transition_to_observation,
     transition_to_robot_action,
 )
@@ -24,23 +24,19 @@ from .core import RobotAction, RobotObservation
 from .pipeline import IdentityProcessorStep, RobotProcessorPipeline
 
 
-def make_default_teleop_action_processor() -> RobotProcessorPipeline[
-    tuple[RobotAction, RobotObservation], RobotAction
-]:
-    teleop_action_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+def make_default_teleop_action_processor() -> RobotProcessorPipeline[RobotAction, RobotAction]:
+    teleop_action_processor = RobotProcessorPipeline[RobotAction, RobotAction](
         steps=[IdentityProcessorStep()],
-        to_transition=robot_action_observation_to_transition,
+        to_transition=robot_action_to_transition,
         to_output=transition_to_robot_action,
     )
     return teleop_action_processor
 
 
-def make_default_robot_action_processor() -> RobotProcessorPipeline[
-    tuple[RobotAction, RobotObservation], RobotAction
-]:
-    robot_action_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+def make_default_robot_action_processor() -> RobotProcessorPipeline[RobotAction, RobotAction]:
+    robot_action_processor = RobotProcessorPipeline[RobotAction, RobotAction](
         steps=[IdentityProcessorStep()],
-        to_transition=robot_action_observation_to_transition,
+        to_transition=robot_action_to_transition,
         to_output=transition_to_robot_action,
     )
     return robot_action_processor

src/lerobot/processor/gym_action_processor.py

@@ -19,8 +19,8 @@ from dataclasses import dataclass
 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
 
 from .converters import to_tensor
-from .core import EnvAction, EnvTransition, PolicyAction
-from .pipeline import ActionProcessorStep, ProcessorStep, ProcessorStepRegistry
+from .core import EnvAction, PolicyAction
+from .pipeline import ActionProcessorStep, ProcessorStepRegistry
 
 
 @ProcessorStepRegistry.register("torch2numpy_action_processor")
@@ -69,27 +69,23 @@ class Torch2NumpyActionProcessorStep(ActionProcessorStep):
 
 @ProcessorStepRegistry.register("numpy2torch_action_processor")
 @dataclass
-class Numpy2TorchActionProcessorStep(ProcessorStep):
-    """Converts a NumPy array action to a PyTorch tensor when action is present."""
+class Numpy2TorchActionProcessorStep(ActionProcessorStep):
+    """
+    Converts a NumPy array action to a PyTorch tensor.
 
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        """Converts numpy action to torch tensor if action exists, otherwise passes through."""
-        from .core import TransitionKey
+    This step is useful for converting actions from environments or hardware,
+    which are often NumPy arrays, into PyTorch tensors that can be processed
+    by a policy or model.
+    """
 
-        self._current_transition = transition.copy()
-        new_transition = self._current_transition
-
-        action = new_transition.get(TransitionKey.ACTION)
-        if action is not None:
-            if not isinstance(action, EnvAction):
-                raise TypeError(
-                    f"Expected np.ndarray or None, got {type(action).__name__}. "
-                    "Use appropriate processor for non-tensor actions."
-                )
-            torch_action = to_tensor(action, dtype=None)  # Preserve original dtype
-            new_transition[TransitionKey.ACTION] = torch_action
-
-        return new_transition
+    def action(self, action: EnvAction) -> PolicyAction:
+        if not isinstance(action, EnvAction):
+            raise TypeError(
+                f"Expected np.ndarray or None, got {type(action).__name__}. "
+                "Use appropriate processor for non-tensor actions."
+            )
+        torch_action = to_tensor(action, dtype=None)  # Preserve original dtype
+        return torch_action
 
     def transform_features(
         self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]

src/lerobot/processor/hil_processor.py

@@ -340,16 +340,11 @@ class GripperPenaltyProcessorStep(ComplementaryDataProcessorStep):
         """
         action = self.transition.get(TransitionKey.ACTION)
 
-        raw_joint_positions = complementary_data.get("raw_joint_positions", None)
-        if raw_joint_positions is None:
-            return complementary_data
-
-        current_gripper_pos = raw_joint_positions.get(GRIPPER_KEY, None)
+        current_gripper_pos = complementary_data.get("raw_joint_positions", None).get(GRIPPER_KEY, None)
         if current_gripper_pos is None:
             return complementary_data
 
-        # Gripper action is a PolicyAction at this stage
-        gripper_action = action[-1].item()
+        gripper_action = action[f"{GRIPPER_KEY}.pos"]
         gripper_action_normalized = gripper_action / self.max_gripper_pos
 
         # Normalize gripper state and action

src/lerobot/processor/joint_observations_processor.py

@@ -20,12 +20,12 @@ from typing import Any
 import torch
 
 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.constants import OBS_STATE
 from lerobot.processor.pipeline import (
     ObservationProcessorStep,
     ProcessorStepRegistry,
 )
 from lerobot.robots import Robot
-from lerobot.utils.constants import OBS_STATE
 
 
 @dataclass

src/lerobot/processor/migrate_policy_normalization.py

@@ -59,7 +59,6 @@ from safetensors.torch import load_file as load_safetensors
 
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.policies.factory import get_policy_class, make_policy_config, make_pre_post_processors
-from lerobot.utils.constants import ACTION
 
 
 def extract_normalization_stats(state_dict: dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]:
@@ -197,7 +196,7 @@ def detect_features_and_norm_modes(
                 feature_type = FeatureType.VISUAL
             elif "state" in key:
                 feature_type = FeatureType.STATE
-            elif ACTION in key:
+            elif "action" in key:
                 feature_type = FeatureType.ACTION
             else:
                 feature_type = FeatureType.STATE  # Default
@@ -216,7 +215,7 @@ def detect_features_and_norm_modes(
                 feature_type = FeatureType.VISUAL
             elif "state" in key or "joint" in key or "position" in key:
                 feature_type = FeatureType.STATE
-            elif ACTION in key:
+            elif "action" in key:
                 feature_type = FeatureType.ACTION
             else:
                 feature_type = FeatureType.STATE
@@ -322,7 +321,7 @@ def convert_features_to_policy_features(features_dict: dict[str, dict]) -> dict[
             feature_type = FeatureType.VISUAL
         elif "state" in key:
             feature_type = FeatureType.STATE
-        elif ACTION in key:
+        elif "action" in key:
             feature_type = FeatureType.ACTION
         else:
             feature_type = FeatureType.STATE