feat(dataset-tools): add dataset utilities and example script

- Introduced dataset tools for LeRobotDataset, including functions for deleting episodes, splitting datasets, adding/removing features, and merging datasets.
- Added an example script demonstrating the usage of these utilities.
- Implemented comprehensive tests for all new functionalities to ensure reliability and correctness.

.github/PULL_REQUEST_TEMPLATE.md

@@ -30,7 +30,7 @@ pytest -sx tests/test_stuff.py::test_something
 ```
 
 ```bash
-python -m lerobot.scripts.train --some.option=true
+lerobot-train --some.option=true
 ```
 
 ## SECTION TO REMOVE BEFORE SUBMITTING YOUR PR

.github/workflows/nightly.yml

@@ -29,8 +29,8 @@ on:
 env:
   UV_VERSION: "0.8.0"
   PYTHON_VERSION: "3.10"
-  DOCKER_IMAGE_NAME_CPU: huggingface/lerobot-gpu:latest
-  DOCKER_IMAGE_NAME_GPU: huggingface/lerobot-cpu:latest
+  DOCKER_IMAGE_NAME_CPU: huggingface/lerobot-cpu:latest
+  DOCKER_IMAGE_NAME_GPU: huggingface/lerobot-gpu:latest
 
 # Ensures that only the latest commit is built, canceling older runs.
 concurrency:

Makefile

@@ -44,7 +44,7 @@ test-end-to-end:
 	${MAKE} DEVICE=$(DEVICE) test-smolvla-ete-eval
 
 test-act-ete-train:
-	python -m lerobot.scripts.train \
+	lerobot-train \
 		--policy.type=act \
 		--policy.dim_model=64 \
 		--policy.n_action_steps=20 \
@@ -68,12 +68,12 @@ test-act-ete-train:
 		--output_dir=tests/outputs/act/
 
 test-act-ete-train-resume:
-	python -m lerobot.scripts.train \
+	lerobot-train \
 		--config_path=tests/outputs/act/checkpoints/000002/pretrained_model/train_config.json \
 		--resume=true
 
 test-act-ete-eval:
-	python -m lerobot.scripts.eval \
+	lerobot-eval \
 		--policy.path=tests/outputs/act/checkpoints/000004/pretrained_model \
 		--policy.device=$(DEVICE) \
 		--env.type=aloha \
@@ -82,7 +82,7 @@ test-act-ete-eval:
 		--eval.batch_size=1
 
 test-diffusion-ete-train:
-	python -m lerobot.scripts.train \
+	lerobot-train \
 		--policy.type=diffusion \
 		--policy.down_dims='[64,128,256]' \
 		--policy.diffusion_step_embed_dim=32 \
@@ -106,7 +106,7 @@ test-diffusion-ete-train:
 		--output_dir=tests/outputs/diffusion/
 
 test-diffusion-ete-eval:
-	python -m lerobot.scripts.eval \
+	lerobot-eval \
 		--policy.path=tests/outputs/diffusion/checkpoints/000002/pretrained_model \
 		--policy.device=$(DEVICE) \
 		--env.type=pusht \
@@ -115,7 +115,7 @@ test-diffusion-ete-eval:
 		--eval.batch_size=1
 
 test-tdmpc-ete-train:
-	python -m lerobot.scripts.train \
+	lerobot-train \
 		--policy.type=tdmpc \
 		--policy.device=$(DEVICE) \
 		--policy.push_to_hub=false \
@@ -137,7 +137,7 @@ test-tdmpc-ete-train:
 		--output_dir=tests/outputs/tdmpc/
 
 test-tdmpc-ete-eval:
-	python -m lerobot.scripts.eval \
+	lerobot-eval \
 		--policy.path=tests/outputs/tdmpc/checkpoints/000002/pretrained_model \
 		--policy.device=$(DEVICE) \
 		--env.type=xarm \
@@ -148,7 +148,7 @@ test-tdmpc-ete-eval:
 
 
 test-smolvla-ete-train:
-	python -m lerobot.scripts.train \
+	lerobot-train \
 		--policy.type=smolvla \
 		--policy.n_action_steps=20 \
 		--policy.chunk_size=20 \
@@ -171,7 +171,7 @@ test-smolvla-ete-train:
 		--output_dir=tests/outputs/smolvla/
 
 test-smolvla-ete-eval:
-	python -m lerobot.scripts.eval \
+	lerobot-eval \
 		--policy.path=tests/outputs/smolvla/checkpoints/000004/pretrained_model \
 		--policy.device=$(DEVICE) \
 		--env.type=aloha \

README.md

@@ -6,7 +6,7 @@
 
 <div align="center">
 
-[![Tests](https://github.com/huggingface/lerobot/actions/workflows/nightly.yml/badge.svg?branch=main)](https://github.com/huggingface/lerobot/actions/workflows/nighty.yml?query=branch%3Amain)
+[![Tests](https://github.com/huggingface/lerobot/actions/workflows/nightly.yml/badge.svg?branch=main)](https://github.com/huggingface/lerobot/actions/workflows/nightly.yml?query=branch%3Amain)
 [![Python versions](https://img.shields.io/pypi/pyversions/lerobot)](https://www.python.org/downloads/)
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/huggingface/lerobot/blob/main/LICENSE)
 [![Status](https://img.shields.io/pypi/status/lerobot)](https://pypi.org/project/lerobot/)
@@ -233,7 +233,7 @@ Under the hood, the `LeRobotDataset` format makes use of several ways to seriali
 
 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):
@@ -246,20 +246,30 @@ dataset attributes:
   │  ├ timestamp (float32): timestamp in the episode
   │  ├ next.done (bool): indicates the end of an episode ; True for the last frame in each episode
   │  └ index (int64): general index in the whole dataset
-  ├ episode_data_index: contains 2 tensors with the start and end indices of each episode
-  │  ├ from (1D int64 tensor): first frame index for each episode — shape (num episodes,) starts with 0
-  │  └ to: (1D int64 tensor): last frame index for each episode — shape (num episodes,)
-  ├ stats: a dictionary of statistics (max, mean, min, std) for each feature in the dataset, for instance
-  │  ├ observation.images.cam_high: {'max': tensor with same number of dimensions (e.g. `(c, 1, 1)` for images, `(c,)` for states), etc.}
-  │  ...
-  ├ info: a dictionary of metadata on the dataset
-  │  ├ codebase_version (str): this is to keep track of the codebase version the dataset was created with
-  │  ├ fps (float): frame per second the dataset is recorded/synchronized to
-  │  ├ video (bool): indicates if frames are encoded in mp4 video files to save space or stored as png files
-  │  └ encoding (dict): if video, this documents the main options that were used with ffmpeg to encode the videos
-  ├ videos_dir (Path): where the mp4 videos or png images are stored/accessed
-  └ camera_keys (list of string): the keys to access camera features in the item returned by the dataset (e.g. `["observation.images.cam_high", ...]`)
-```
+  ├ meta: a LeRobotDatasetMetadata object containing:
+  │  ├ info: a dictionary of metadata on the dataset
+  │  │  ├ codebase_version (str): this is to keep track of the codebase version the dataset was created with
+  │  │  ├ fps (int): frame per second the dataset is recorded/synchronized to
+  │  │  ├ features (dict): all features contained in the dataset with their shapes and types
+  │  │  ├ total_episodes (int): total number of episodes in the dataset
+  │  │  ├ total_frames (int): total number of frames in the dataset
+  │  │  ├ robot_type (str): robot type used for recording
+  │  │  ├ data_path (str): formattable string for the parquet files
+  │  │  └ video_path (str): formattable string for the video files (if using videos)
+  │  ├ episodes: a DataFrame containing episode metadata with columns:
+  │  │  ├ episode_index (int): index of the episode
+  │  │  ├ tasks (list): list of tasks for this episode
+  │  │  ├ length (int): number of frames in this episode
+  │  │  ├ dataset_from_index (int): start index of this episode in the dataset
+  │  │  └ dataset_to_index (int): end index of this episode in the dataset
+  │  ├ stats: a dictionary of statistics (max, mean, min, std) for each feature in the dataset, for instance
+  │  │  ├ observation.images.front_cam: {'max': tensor with same number of dimensions (e.g. `(c, 1, 1)` for images, `(c,)` for states), etc.}
+  │  │  └ ...
+  │  └ tasks: a DataFrame containing task information with task names as index and task_index as values
+  ├ 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:
 
@@ -276,22 +286,22 @@ Check out [example 2](https://github.com/huggingface/lerobot/blob/main/examples/
 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
-python -m lerobot.scripts.eval \
+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
-python -m lerobot.scripts.eval --policy.path={OUTPUT_DIR}/checkpoints/last/pretrained_model
+lerobot-eval --policy.path={OUTPUT_DIR}/checkpoints/last/pretrained_model
 ```
 
-See `python -m lerobot.scripts.eval --help` for more instructions.
+See `lerobot-eval --help` for more instructions.
 
 ### Train your own policy
 
@@ -303,7 +313,7 @@ A link to the wandb logs for the run will also show up in yellow in your termina
 
 \<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 `python -m lerobot.scripts.eval --help` for more instructions.
+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)
 
@@ -311,7 +321,7 @@ We provide some pretrained policies on our [hub page](https://huggingface.co/ler
 You can reproduce their training by loading the config from their run. Simply running:
 
 ```bash
-python -m lerobot.scripts.train --config_path=lerobot/diffusion_pusht
+lerobot-train --config_path=lerobot/diffusion_pusht
 ```
 
 reproduces SOTA results for Diffusion Policy on the PushT task.

benchmarks/video/run_video_benchmark.py

@@ -108,7 +108,8 @@ def save_decoded_frames(
 
 
 def save_first_episode(imgs_dir: Path, dataset: LeRobotDataset) -> None:
-    ep_num_images = dataset.episode_data_index["to"][0].item()
+    episode_index = 0
+    ep_num_images = dataset.meta.episodes["length"][episode_index]
     if imgs_dir.exists() and len(list(imgs_dir.glob("frame_*.png"))) == ep_num_images:
         return
 
@@ -265,7 +266,8 @@ def benchmark_encoding_decoding(
             overwrite=True,
         )
 
-    ep_num_images = dataset.episode_data_index["to"][0].item()
+    episode_index = 0
+    ep_num_images = dataset.meta.episodes["length"][episode_index]
     width, height = tuple(dataset[0][dataset.meta.camera_keys[0]].shape[-2:])
     num_pixels = width * height
     video_size_bytes = video_path.stat().st_size

docs/source/_toctree.yml

@@ -19,21 +19,16 @@
     title: Train RL in Simulation
   - local: async
     title: Use Async Inference
+  - local: porting_datasets_v3
+    title: Porting Large Datasets
   title: "Tutorials"
 - sections:
   - local: smolvla
     title: Finetune SmolVLA
   title: "Policies"
-
-- sections:
-  - local: introduction_processors
-    title: Introduction to Robot Processors
-  - local: implement_your_own_processor
-    title: Implement your own processor
-  - local: processors_robots_teleop
-    title: Processors for Robots and Teleoperators
-  title: "Robot Processors"
 - sections:
+  - local: hope_jr
+    title: Hope Jr
   - local: so101
     title: SO-101
   - local: so100
@@ -42,13 +37,7 @@
     title: Koch v1.1
   - local: lekiwi
     title: LeKiwi
-  - local: hope_jr
-    title: Hope Jr
   title: "Robots"
-- sections:
-  - local: phone_teleop
-    title: Phone
-  title: "Teleoperators"
 - sections:
   - local: notebooks
     title: Notebooks

docs/source/cameras.mdx

@@ -9,7 +9,7 @@ To instantiate a camera, you need a camera identifier. This identifier might cha
 To find the camera indices of the cameras plugged into your system, run the following script:
 
 ```bash
-python -m lerobot.find_cameras opencv # or realsense for Intel Realsense cameras
+lerobot-find-cameras opencv # or realsense for Intel Realsense cameras
 ```
 
 The output will look something like this if you have two cameras connected:

docs/source/hilserl.mdx

@@ -412,7 +412,7 @@ Example configuration for training the [reward classifier](https://huggingface.c
 To train the classifier, use the `train.py` script with your configuration:
 
 ```bash
-python -m lerobot.scripts.train --config_path path/to/reward_classifier_train_config.json
+lerobot-train --config_path path/to/reward_classifier_train_config.json
 ```
 
 **Deploying and Testing the Model**
@@ -458,7 +458,7 @@ The reward classifier will automatically provide rewards based on the visual inp
 3. **Train the classifier**:
 
    ```bash
-   python -m lerobot.scripts.train --config_path src/lerobot/configs/reward_classifier_train_config.json
+   lerobot-train --config_path src/lerobot/configs/reward_classifier_train_config.json
    ```
 
 4. **Test the classifier**:

docs/source/hope_jr.mdx

@@ -19,7 +19,7 @@ pip install -e ".[hopejr]"
 Before starting calibration and operation, you need to identify the USB ports for each HopeJR component. Run this script to find the USB ports for the arm, hand, glove, and exoskeleton:
 
 ```bash
-python -m lerobot.find_port
+lerobot-find-port
 ```
 
 This will display the available USB ports and their associated devices. Make note of the port paths (e.g., `/dev/tty.usbmodem58760433331`, `/dev/tty.usbmodem11301`) as you'll need to specify them in the `--robot.port` and `--teleop.port` parameters when recording data, replaying episodes, or running teleoperation scripts.
@@ -31,7 +31,7 @@ Before performing teleoperation, HopeJR's limbs need to be calibrated. Calibrati
 ### 1.1 Calibrate Robot Hand
 
 ```bash
-python -m lerobot.calibrate \
+lerobot-calibrate \
     --robot.type=hope_jr_hand \
     --robot.port=/dev/tty.usbmodem58760432281 \
     --robot.id=blue \
@@ -81,7 +81,7 @@ Once you have set the appropriate boundaries for all joints, click "Save" to sav
 ### 1.2 Calibrate Teleoperator Glove
 
 ```bash
-python -m lerobot.calibrate \
+lerobot-calibrate \
     --teleop.type=homunculus_glove \
     --teleop.port=/dev/tty.usbmodem11201 \
     --teleop.id=red \
@@ -120,7 +120,7 @@ Once calibration is complete, the system will save the calibration to `/Users/yo
 ### 1.3 Calibrate Robot Arm
 
 ```bash
-python -m lerobot.calibrate \
+lerobot-calibrate \
     --robot.type=hope_jr_arm \
     --robot.port=/dev/tty.usbserial-1110 \
     --robot.id=white
@@ -146,7 +146,7 @@ Use the calibration interface to set the range boundaries for each joint. Move e
 ### 1.4 Calibrate Teleoperator Exoskeleton
 
 ```bash
-python -m lerobot.calibrate \
+lerobot-calibrate \
     --teleop.type=homunculus_arm \
     --teleop.port=/dev/tty.usbmodem11201 \
     --teleop.id=black
@@ -178,7 +178,7 @@ Due to global variable conflicts in the Feetech middleware, teleoperation for ar
 ### Hand
 
 ```bash
-python -m lerobot.teleoperate \
+lerobot-teleoperate \
     --robot.type=hope_jr_hand \
     --robot.port=/dev/tty.usbmodem58760432281 \
     --robot.id=blue \
@@ -194,7 +194,7 @@ python -m lerobot.teleoperate \
 ### Arm
 
 ```bash
-python -m lerobot.teleoperate \
+lerobot-teleoperate \
     --robot.type=hope_jr_arm \
     --robot.port=/dev/tty.usbserial-1110 \
     --robot.id=white \
@@ -214,7 +214,7 @@ Record, Replay and Train with Hope-JR is still experimental.
 This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
 
 ```bash
-python -m lerobot.record \
+lerobot-record \
     --robot.type=hope_jr_hand \
     --robot.port=/dev/tty.usbmodem58760432281 \
     --robot.id=right \
@@ -236,7 +236,7 @@ python -m lerobot.record \
 ### Replay
 
 ```bash
-python -m lerobot.replay \
+lerobot-replay \
     --robot.type=hope_jr_hand \
     --robot.port=/dev/tty.usbmodem58760432281 \
     --robot.id=right \
@@ -248,7 +248,7 @@ python -m lerobot.replay \
 ### Train
 
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
   --dataset.repo_id=nepyope/hand_record_test_with_video_data \
   --policy.type=act \
   --output_dir=outputs/train/hopejr_hand \
@@ -263,7 +263,7 @@ python -m lerobot.scripts.train \
 This training run can be viewed as an example [here](https://wandb.ai/tino/lerobot/runs/rp0k8zvw?nw=nwusertino).
 
 ```bash
-python -m lerobot.record \
+lerobot-record \
   --robot.type=hope_jr_hand \
   --robot.port=/dev/tty.usbmodem58760432281 \
   --robot.id=right \

docs/source/il_robots.mdx

@@ -45,7 +45,7 @@ Note that the `id` associated with a robot is used to store the calibration file
 <hfoptions id="teleoperate_so101">
 <hfoption id="Command">
 ```bash
-python -m lerobot.teleoperate \
+lerobot-teleoperate \
     --robot.type=so101_follower \
     --robot.port=/dev/tty.usbmodem58760431541 \
     --robot.id=my_awesome_follower_arm \
@@ -101,7 +101,7 @@ With `rerun`, you can teleoperate again while simultaneously visualizing the cam
 <hfoptions id="teleoperate_koch_camera">
 <hfoption id="Command">
 ```bash
-python -m lerobot.teleoperate \
+lerobot-teleoperate \
     --robot.type=koch_follower \
     --robot.port=/dev/tty.usbmodem58760431541 \
     --robot.id=my_awesome_follower_arm \
@@ -174,7 +174,7 @@ Now you can record a dataset. To record 5 episodes and upload your dataset to th
 <hfoptions id="record">
 <hfoption id="Command">
 ```bash
-python -m lerobot.record \
+lerobot-record \
     --robot.type=so101_follower \
     --robot.port=/dev/tty.usbmodem585A0076841 \
     --robot.id=my_awesome_follower_arm \
@@ -376,7 +376,7 @@ You can replay the first episode on your robot with either the command below or
 <hfoptions id="replay">
 <hfoption id="Command">
 ```bash
-python -m lerobot.replay \
+lerobot-replay \
     --robot.type=so101_follower \
     --robot.port=/dev/tty.usbmodem58760431541 \
     --robot.id=my_awesome_follower_arm \
@@ -428,10 +428,10 @@ Your robot should replicate movements similar to those you recorded. For example
 
 ## Train a policy
 
-To train a policy to control your robot, use the [`python -m lerobot.scripts.train`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
+To train a policy to control your robot, use the [`lerobot-train`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
 
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
   --dataset.repo_id=${HF_USER}/so101_test \
   --policy.type=act \
   --output_dir=outputs/train/act_so101_test \
@@ -453,7 +453,7 @@ Training should take several hours. You will find checkpoints in `outputs/train/
 To resume training from a checkpoint, below is an example command to resume from `last` checkpoint of the `act_so101_test` policy:
 
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
   --config_path=outputs/train/act_so101_test/checkpoints/last/pretrained_model/train_config.json \
   --resume=true
 ```
@@ -490,7 +490,7 @@ You can use the `record` script from [`lerobot/record.py`](https://github.com/hu
 <hfoptions id="eval">
 <hfoption id="Command">
 ```bash
-python -m lerobot.record  \
+lerobot-record  \
   --robot.type=so100_follower \
   --robot.port=/dev/ttyACM1 \
   --robot.cameras="{ up: {type: opencv, index_or_path: /dev/video10, width: 640, height: 480, fps: 30}, side: {type: intelrealsense, serial_number_or_name: 233522074606, width: 640, height: 480, fps: 30}}" \
@@ -519,14 +519,11 @@ 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.record import record_loop
-from lerobot.policies.factory import make_processor
 
 NUM_EPISODES = 5
 FPS = 30
 EPISODE_TIME_SEC = 60
 TASK_DESCRIPTION = "My task description"
-HF_MODEL_ID = "<hf_username>/<model_repo_id>"
-HF_DATASET_ID = "<hf_username>/<eval_dataset_repo_id>"
 
 # Create the robot configuration
 camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
@@ -538,7 +535,7 @@ robot_config = SO100FollowerConfig(
 robot = SO100Follower(robot_config)
 
 # Initialize the policy
-policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
+policy = ACTPolicy.from_pretrained("<hf_username>/<my_policy_repo_id>")
 
 # Configure the dataset features
 action_features = hw_to_dataset_features(robot.action_features, "action")
@@ -547,7 +544,7 @@ dataset_features = {**action_features, **obs_features}
 
 # Create the dataset
 dataset = LeRobotDataset.create(
-    repo_id=HF_DATASET_ID,
+    repo_id="<hf_username>/eval_<dataset_repo_id>",
     fps=FPS,
     features=dataset_features,
     robot_type=robot.name,
@@ -562,12 +559,6 @@ _init_rerun(session_name="recording")
 # Connect the robot
 robot.connect()
 
-preprocessor, postprocessor = make_processor(
-    policy_cfg=policy,
-    pretrained_path=HF_MODEL_ID,
-    dataset_stats=dataset.meta.stats,
-)
-
 for episode_idx in range(NUM_EPISODES):
     log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
 
@@ -577,8 +568,6 @@ for episode_idx in range(NUM_EPISODES):
         events=events,
         fps=FPS,
         policy=policy,
-        preprocessor=preprocessor,
-        postprocessor=postprocessor,
         dataset=dataset,
         control_time_s=EPISODE_TIME_SEC,
         single_task=TASK_DESCRIPTION,

docs/source/il_sim.mdx

@@ -96,10 +96,10 @@ If you uploaded your dataset to the hub you can [visualize your dataset online](
 
 ## Train a policy
 
-To train a policy to control your robot, use the [`python -m lerobot.scripts.train`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
+To train a policy to control your robot, use the [`lerobot-train`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
 
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
   --dataset.repo_id=${HF_USER}/il_gym \
   --policy.type=act \
   --output_dir=outputs/train/il_sim_test \

docs/source/implement_your_own_processor.mdx

@@ -1,323 +0,0 @@
-# Implement your own Robot Processor
-
-In this tutorial, you'll learn how to implement your own Robot Processor.
-It begins by exploring the need for a custom processor, then uses the Normalization processors as the running example to explain how to implement, configure, and serialize a processor. Finally, it lists all helper processors that ship with LeRobot.
-
-## Why would you need a custom processor?
-
-In most cases, when reading raw data from a sensor like the camera and robot motor encoders,
-you will need to process this data to transform it into a format that is compatible to use with the policies in LeRobot.
-For example, raw images are encoded with `uint8` and the values are in the range `[0, 255]`.
-To use these images with the policies, you will need to cast them to `float32` and normalize them to the range `[0, 1]`.
-
-For example, in LeRobot's `VanillaObservationProcessor`, raw images come from the environment as numpy arrays with `uint8` values in range `[0, 255]` and in channel-last format `(H, W, C)`. The processor transforms them into PyTorch tensors with `float32` values in range `[0, 1]` and channel-first format `(C, H, W)`:
-
-```python
-# Input: numpy array with shape (480, 640, 3) and dtype uint8
-raw_image = env_observation["pixels"]  # Values in [0, 255]
-
-# After processing: torch tensor with shape (1, 3, 480, 640) and dtype float32
-processed_image = processor(transition)["observation"]["observation.image"]  # Values in [0, 1]
-```
-
-On the other hand, when a model returns a certain action to be executed on the robot, it is often that one has to post-process this action to make it compatible to run on the robot.
-For example, the model might return joint positions values that range from `[-1, 1]` and one would need to scale them to the ranges of the minimum and maximum joint angle positions of the robot.
-
-In LeRobot, this normalization workflow is handled by the `NormalizerProcessor` (for inputs) and the `UnnormalizerProcessor` (for outputs). These processors are heavily used by policies (e.g., Pi0, SmolVLA) and integrate tightly with the `RobotProcessor`'s `get_config`, `state_dict`, and `load_state_dict` APIs.
-
-For instance, `UnnormalizerProcessor` converts model outputs in `[-1, 1]` back to actual robot joint ranges:
-
-```python
-# Input: model action with normalized values in [-1, 1]
-normalized_action = torch.tensor([-0.5, 0.8, -1.0, 0.2])  # Model output
-
-# After post-processing: real joint positions in robot's native ranges
-# Example: joints range from [-180.0, 180.0]
-real_action = unnormalizer(transition)["action"]
-# real action after post-processing: [ -90.,  144., -180.,   36.]
-```
-
-The unnormalizer uses the dataset statistics to convert back:
-
-```python
-# For MIN_MAX normalization: action = (normalized + 1) * (max - min) / 2 + min
-real_action = (normalized_action + 1) * (max_val - min_val) / 2 + min_val
-```
-
-All these situations point us towards the need for a mechanism to preprocess the data before being passed to the policies and then post-process the action that are returned to be executed on the robot.
-
-To that end, LeRobot provides a pipeline mechanism to implement a sequence of processing steps for the input data and the output action.
-
-## How to implement your own processor?
-
-We'll use the `NormalizerProcessor` as a concrete running example because it is central to most policies and demonstrates configuration and state serialization cleanly.
-
-Prepare the sequence of processing steps necessary for your problem. A processor step is a class that implements the following methods:
-
-- `__call__`: implements the processing step for the input transition.
-- `get_config`: gets the configuration of the processor step.
-- `state_dict`: gets the state of the processor step.
-- `load_state_dict`: loads the state of the processor step.
-- `reset`: resets the state of the processor step.
-- `feature_contract`: displays the modification to the feature space during the processor step.
-
-### Implement the `__call__` method
-
-The `__call__` method is the core of your processor step. It takes an `EnvTransition` and returns a modified `EnvTransition`. Here's how the `NormalizerProcessor` conceptually works (simplified):
-
-```python
-from dataclasses import dataclass
-import torch
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
-from lerobot.processor.pipeline import EnvTransition, TransitionKey
-
-@dataclass
-class NormalizerProcessor:
-    features: dict[str, PolicyFeature]
-    norm_map: dict[FeatureType, NormalizationMode]
-    stats: dict[str, dict[str, torch.Tensor]]
-    eps: float = 1e-8
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        normalized_info = {}
-
-        obs = transition.get(TransitionKey.OBSERVATION)
-        act = transition.get(TransitionKey.ACTION)
-
-        new_obs = self._normalize_observation(obs, normalized_info)
-        new_act = self._normalize_action(act, normalized_info)
-
-        new_transition = transition.copy()
-        new_transition[TransitionKey.OBSERVATION] = new_obs
-        new_transition[TransitionKey.ACTION] = new_act
-
-        # Record what was normalized into complementary_data
-        if normalized_info:
-            comp = new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}
-            comp = dict(comp)
-            comp["normalized_keys"] = normalized_info
-            new_transition[TransitionKey.COMPLEMENTARY_DATA] = comp
-
-        return new_transition
-```
-
-See the full implementation in `src/lerobot/processor/normalize_processor.py` for details on mean/std and min/max modes and key selection.
-
-**Key principles:**
-
-- Always check if required data exists before processing
-- Return unchanged transition if no processing is needed
-- Use `transition.copy()` to avoid side effects
-- Only modify the specific keys your processor handles
-
-**Tip**: For observation-only processors, you can inherit from `ObservationProcessor` to avoid writing `__call__` boilerplate. The normalizer is mixed (observations and actions), so it implements `__call__` directly.
-
-### Configuration and State Management
-
-Processors support serialization through three methods that separate configuration from tensor state. This is especially important for normalization processors, which carry dataset statistics (tensors) in their state, and hyperparameters in their config:
-
-```python
-from dataclasses import dataclass, field
-from typing import Any
-import torch
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
-
-@dataclass
-class NormalizerProcessor:
-    features: dict[str, PolicyFeature]
-    norm_map: dict[FeatureType, NormalizationMode]
-    eps: float = 1e-8
-    _tensor_stats: dict[str, dict[str, torch.Tensor]] = field(default_factory=dict, init=False, repr=False)
-
-    def get_config(self) -> dict[str, Any]:
-        """JSON-serializable configuration (no tensors)."""
-        return {
-            "eps": self.eps,
-            "features": {k: {"type": v.type.value, "shape": v.shape} for k, v in self.features.items()},
-            "norm_map": {ft.value: nm.value for ft, nm in self.norm_map.items()},
-        }
-
-    def state_dict(self) -> dict[str, torch.Tensor]:
-        """Tensor state only (e.g., dataset statistics)."""
-        flat: dict[str, torch.Tensor] = {}
-        for key, sub in self._tensor_stats.items():
-            for stat_name, tensor in sub.items():
-                flat[f"{key}.{stat_name}"] = tensor
-        return flat
-
-    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
-        """Restore tensor state at runtime."""
-        self._tensor_stats.clear()
-        for flat_key, tensor in state.items():
-            key, stat_name = flat_key.rsplit(".", 1)
-            self._tensor_stats.setdefault(key, {})[stat_name] = tensor
-```
-
-**Usage:**
-
-```python
-# Save (e.g., inside a policy)
-config = processor.get_config()
-tensors = processor.state_dict()
-
-# Restore (e.g., loading a pretrained policy)
-new_processor = NormalizerProcessor(**config)
-new_processor.load_state_dict(tensors)
-```
-
-### Transform features
-
-The `transform_features` method defines how your processor transforms feature names and shapes. This is crucial for policy configuration and debugging.
-
-Normalization typically preserves the feature keys and shapes, so `NormalizerProcessor.transform_features` returns the input features unchanged. When your processor renames or reshapes, implement this method to reflect the mapping for downstream components. For example, a simple rename processor:
-
-```python
-def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-    # Simple renaming
-    if "pixels" in features:
-        features["observation.image"] = features.pop("pixels")
-
-    # Pattern-based renaming
-    for key in list(features.keys()):
-        if key.startswith("env_state."):
-            suffix = key[len("env_state."):]
-            features[f"observation.{suffix}"] = features.pop(key)
-
-    return features
-```
-
-**Key principles:**
-
-- Use `features.pop(old_key)` to remove and get the old feature
-- Use `features[new_key] = old_feature` to add the renamed feature
-- Always return the modified features dictionary
-- Document transformations clearly in the docstring
-
-### Example of usage from the codebase
-
-`transform_features` is used by `RobotProcessor` to derive the dataset/policy feature contract from an initial feature set by applying each step's transformation. You can see concrete examples in the codebase:
-
-- Phone teleoperation record pipeline (`examples/phone_so100_record.py`): processors like `ForwardKinematicsJointsToEE`, `GripperVelocityToJoint`, and `EEBoundsAndSafety` implement `transform_features` to declare which action/observation keys should be materialized in the dataset.
-- SO100 follower kinematics (`src/lerobot/robots/so100_follower/robot_kinematic_processor.py`): each processor's `transform_features` method adds or refines feature keys such as `observation.state.ee.{x,y,z,wx,wy,wz}` or `action.gripper.pos`.
-- Rename and tokenizer processors (`src/lerobot/processor/rename_processor.py`, `src/lerobot/processor/tokenizer_processor.py`): demonstrate key renaming and adding language token features to the contract.
-
-In practice, you will often aggregate features by running `RobotProcessor.transform_features(...)` with your initial features to compute the final contract before recording or training.
-
-## Helper Classes
-
-LeRobot provides pre-built processor classes for common transformations. Below is a comprehensive list of registered processors in the codebase.
-
-### Core processors (observations, actions, normalization)
-
-- **`VanillaObservationProcessor`** (`observation_processor`): Images and state processing to LeRobot format.
-- **`NormalizerProcessor`** (`normalizer_processor`): Normalize observations/actions (mean/std or min/max to [-1, 1]).
-- **`UnnormalizerProcessor`** (`unnormalizer_processor`): Inverse of the normalizer for model outputs.
-- **`DeviceProcessor`** (`device_processor`): Move tensors to a specific device (CPU/GPU) and optional float dtype.
-- **`ToBatchProcessor`** (`to_batch_processor`): Add batch dimension to observations/actions when missing.
-- **`RenameProcessor`** (`rename_processor`): Rename observation keys using a mapping dictionary.
-- **`TokenizerProcessor`** (`tokenizer_processor`): Tokenize language tasks into `observation.language.*` tensors.
-
-### Teleoperation mapping processors
-
-- **`MapDeltaActionToRobotAction`** (`map_delta_action_to_robot_action`): Map teleop deltas (e.g., gamepad) to `action.target_*` fields.
-- **`MapPhoneActionToRobotAction`** (`map_phone_action_to_robot_action`): Map calibrated phone pose/buttons to `action.target_*` and gripper.
-
-### Robot kinematics processors (SO100 follower example)
-
-- **`EEReferenceAndDelta`** (`ee_reference_and_delta`): Compute desired EE pose from target deltas and current pose.
-- **`EEBoundsAndSafety`** (`ee_bounds_and_safety`): Clip EE pose to bounds and check for jumps.
-- **`InverseKinematicsEEToJoints`** (`inverse_kinematics_ee_to_joints`): Convert EE pose to joint targets via IK.
-- **`GripperVelocityToJoint`** (`gripper_velocity_to_joint`): Convert gripper velocity input to joint position command.
-- **`ForwardKinematicsJointsToEE`** (`forward_kinematics_joints_to_ee`): Compute EE pose features from joint positions via FK.
-- **`AddRobotObservationAsComplimentaryData`** (`add_robot_observation`): Read robot observation and insert `raw_joint_positions` into complementary data.
-
-### Policy-specific utility processors
-
-- **`Pi0NewLineProcessor`** (`pi0_new_line_processor`): Ensure text tasks end with a newline (Pi0 tokenizer compatibility).
-- **`SmolVLANewLineProcessor`** (`smolvla_new_line_processor`): Ensure text tasks end with a newline (SmolVLA tokenizer compatibility).
-
-### Usage Example
-
-```python
-from lerobot.processor import NormalizerProcessor, DeviceProcessor, RobotProcessor, ToBatchProcessor
-
-# Create a processing pipeline (typical policy preprocessor)
-steps = [
-    NormalizerProcessor(features=features, norm_map=norm_map, stats=stats),
-    ToBatchProcessor(),
-    DeviceProcessor(device="cuda"),
-]
-
-# Use in RobotProcessor
-processor = RobotProcessor(steps=steps)
-processed_transition = processor(raw_transition)
-```
-
-### Using overrides
-
-You can override step parameters at load-time using `overrides`. This is handy for non-serializable objects or site-specific settings. It works both in policy factories and with `RobotProcessor.from_pretrained(...)`.
-
-Example: during policy evaluation on the robot, override the device and rename map.
-Use this to run a policy trained on CUDA on a CPU-only robot, or to remap camera keys when the robot uses different names than the dataset.
-
-```437:445:src/lerobot/record.py
-preprocessor, postprocessor = make_processor(
-    policy_cfg=cfg.policy,
-    pretrained_path=cfg.policy.pretrained_path,
-    dataset_stats=rename_stats(dataset.meta.stats, cfg.dataset.rename_map),
-    preprocessor_overrides={
-        "device_processor": {"device": cfg.policy.device},
-        "rename_processor": {"rename_map": cfg.dataset.rename_map},
-    },
-)
-```
-
-Direct usage with `from_pretrained`:
-
-```python
-from lerobot.processor import RobotProcessor
-
-processor = RobotProcessor.from_pretrained(
-    "username/my-processor",
-    overrides={
-        "device_processor": {"device": "cuda:0"},  # registry name for registered steps
-        "CustomStep": {"param": 42},               # class name for non-registered steps
-    },
-)
-```
-
-## Best Practices
-
-- **Keep processors atomic** - One transformation per processor for reusability and debugging
-- **Use dataclasses** - Clean initialization with `@dataclass`
-- **Always register processors** - Use `@ProcessorStepRegistry.register("name")` for discoverability
-- **Check for None** - Always validate required data exists before processing
-- **Use copy() for safety** - Avoid side effects with `transition.copy()`
-- **Separate config and state** - JSON-serializable config vs tensor state_dict
-- **Use base classes** - Inherit from `ObservationProcessor` for observation-only processing
-
-```python
-@ProcessorStepRegistry.register("my_processor")
-@dataclass
-class MyProcessor(ObservationProcessor):
-    threshold: float = 0.5
-
-    def observation(self, observation):
-        if observation is None:
-            return observation
-        # Your processing logic here
-        return processed_observation
-```
-
-## Conclusion
-
-You now have all the tools to implement custom processors in LeRobot! The key steps are:
-
-1. **Define your processor** as a dataclass with the required methods (`__call__`, `get_config`, `state_dict`, `load_state_dict`, `reset`, `feature_contract`)
-2. **Register it** using `@ProcessorStepRegistry.register("name")` for discoverability
-3. **Integrate it** into a `RobotProcessor` pipeline with other processing steps
-4. **Use base classes** like `ObservationProcessor` when possible to reduce boilerplate
-
-The processor system is designed to be modular and composable, allowing you to build complex data processing pipelines from simple, focused components. Whether you're preprocessing sensor data for training or post-processing model outputs for robot execution, custom processors give you the flexibility to handle any data transformation your robotics application requires. Policies like Pi0 and SmolVLA use the same normalization processors described above, so your understanding here will transfer directly when wiring policy preprocessors and postprocessors.
-
-Start simple, test thoroughly, and leverage the existing helper classes to build robust data processing pipelines for your robot learning workflows.

docs/source/introduction_processors.mdx

@@ -1,991 +0,0 @@
-# Introduction to Processors
-
-In robotics, there's a fundamental mismatch between the data that robots and humans produce and what machine learning models expect. This creates several translation challenges:
-
-**Raw Robot Data → Model Input:**
-
-- Robots output raw sensor data (camera images, joint positions, force readings) that need normalization, batching, and device placement before models can process them
-- Language instructions from humans ("pick up the red cube") must be tokenized into numerical representations
-- Different robots use different coordinate systems and units that need standardization
-
-**Model Output → Robot Commands:**
-
-- Models might output end-effector positions, but robots need joint-space commands
-- Teleoperators (like gamepads) produce relative movements (delta positions), but robots expect absolute commands
-- Model predictions are often normalized and need to be converted back to real-world scales
-
-**Cross-Domain Translation:**
-
-- Training data from one robot setup needs adaptation for deployment on different hardware
-- Models trained with specific camera configurations must work with new camera arrangements
-- Datasets with different naming conventions need harmonization
-
-**That's where processors come in.** They serve as the universal translators that bridge these gaps, ensuring seamless data flow from sensors to models to actuators.
-
-Processors are the data transformation backbone of LeRobot. They handle all the preprocessing and postprocessing steps needed to convert raw environment data into model-ready inputs and vice versa. This guide will walk you through everything you need to know about processors - from basic concepts to advanced usage patterns.
-
-## 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:
-
-- **Normalized**: Scaled to appropriate ranges for neural network processing
-- **Batched**: Organized with proper dimensions for batch processing
-- **Tokenized**: Text converted to numerical representations
-- **Device-placed**: Moved to the right hardware (CPU/GPU)
-- **Type-converted**: Cast to appropriate data types
-
-Processors handle these transformations through composable, reusable steps that can be chained together into pipelines. Think of them as a modular assembly line where each station performs a specific transformation on your data.
-
-## Core Concepts
-
-### EnvTransition: The Universal Data Container
-
-The `EnvTransition` is the fundamental data structure that flows through all processors. It's a typed dictionary that represents a complete robot-environment interaction:
-
-```python
-from lerobot.processor.pipeline import TransitionKey, EnvTransition
-
-# Example transition from a robot collecting data
-transition: EnvTransition = {
-    TransitionKey.OBSERVATION: {
-        "observation.images.camera0": camera0_image_tensor,  # Shape: (H, W, C)
-        "observation.images.camera1": camera1_image_tensor,  # Shape: (H, W, C)
-        "observation.state": joint_positions_tensor,         # Shape: (7,) for 7-DOF arm
-        "observation.environment_state": env_state_tensor    # Shape: (3,) for object position
-    },
-    TransitionKey.ACTION: action_tensor,                    # Shape: (7,) for joint velocities
-    TransitionKey.REWARD: 0.0,                             # Scalar reward signal
-    TransitionKey.DONE: False,                             # Episode termination flag
-    TransitionKey.TRUNCATED: False,                        # Episode truncation flag
-    TransitionKey.INFO: {"success": False},                # Additional metadata
-    TransitionKey.COMPLEMENTARY_DATA: {
-        "task": "pick up the red cube",                    # Language instruction
-        "task_index": 0,                                   # Task identifier
-        "index": 42                                        # Frame index
-    }
-}
-```
-
-Each key in the transition has a specific purpose:
-
-- **OBSERVATION**: All sensor data (images, states, proprioception)
-- **ACTION**: The action to execute or that was executed
-- **REWARD**: Reinforcement learning signal
-- **DONE/TRUNCATED**: Episode boundary indicators
-- **INFO**: Arbitrary metadata
-- **COMPLEMENTARY_DATA**: Task descriptions, indices, padding flags, inter-step data (e.g., you need to compute the velocities and then use this velocity to clip the action)
-
-### ProcessorStep: The Building Block Interface
-
-A `ProcessorStep` is a single transformation unit that processes transitions. It's a protocol (interface) that any processor step must implement:
-
-```python
-from lerobot.processor.pipeline import ProcessorStep, EnvTransition
-from lerobot.configs.types import PolicyFeature
-from typing import Any
-import torch
-
-class MyProcessorStep:
-    """Example processor step interface - all methods must be implemented."""
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        """Transform the transition - this is the main processing logic."""
-        raise NotImplementedError
-
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        """Declare how this step transforms feature shapes/types."""
-        raise NotImplementedError
-
-    def get_config(self) -> dict[str, Any]:
-        """Return JSON-serializable configuration for saving/loading."""
-        raise NotImplementedError
-
-    def state_dict(self) -> dict[str, torch.Tensor]:
-        """Return any learnable parameters (tensors only)."""
-        raise NotImplementedError
-
-    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
-        """Load learnable parameters from saved state."""
-        raise NotImplementedError
-
-    def reset(self) -> None:
-        """Reset any internal state between episodes."""
-        raise NotImplementedError
-```
-
-### RobotProcessor: The Pipeline Orchestrator
-
-The `RobotProcessor` chains multiple `ProcessorStep` instances together, executing them sequentially. It provides automatic format conversion to handle both batch dictionaries (from datasets) and EnvTransition dictionaries:
-
-```python
-from lerobot.processor.pipeline import RobotProcessor, _default_batch_to_transition, _default_transition_to_batch
-
-# Create a processing pipeline
-processor = RobotProcessor(
-    steps=[
-        step1,  # First transformation
-        step2,  # Second transformation
-        step3   # Third transformation
-    ],
-    name="my_preprocessing_pipeline",
-
-    # Optional: Custom converters for input/output formats
-    to_transition=_default_batch_to_transition,  # How to convert batch dict → EnvTransition
-    to_output=_default_transition_to_batch       # How to convert EnvTransition → output format
-)
-
-# The processor automatically handles different input formats:
-# 1. If input is a batch dict (from dataset), converts to EnvTransition
-# 2. Passes through each step sequentially
-# 3. Converts back to original format (or custom output format)
-
-# Example with batch dict input (common in training)
-batch_dict = {"observation.state": tensor, "action": tensor}
-output = processor(batch_dict)  # Automatically converted to/from EnvTransition
-
-# Example with EnvTransition input (common in inference)
-transition = {TransitionKey.OBSERVATION: {...}, TransitionKey.ACTION: ...}
-output = processor(transition)  # Stays as EnvTransition throughout
-```
-
-The `to_transition` and `to_output` converters enable seamless integration with existing codebases.
-By default, they handle the standard LeRobot batch format, but you can customize them for different data structures.
-
-### Additional Converter Functions
-
-LeRobot provides several specialized converter functions for common robotics scenarios:
-
-```python
-from lerobot.processor.converters import (
-    to_transition_teleop_action,
-    to_transition_robot_observation,
-    to_output_robot_action,
-    to_dataset_frame
-)
-```
-
-**`to_transition_teleop_action`** - Converts teleoperation device actions to EnvTransitions:
-
-```python
-# Use case: Phone, gamepad, or other teleop device control
-phone_action = {"x": 0.1, "y": -0.2, "gripper": 0.8}
-transition = to_transition_teleop_action(phone_action)
-# Creates: {ACTION: {"action.x": 0.1, "action.y": -0.2, "action.gripper": 0.8}, ...}
-```
-
-**`to_transition_robot_observation`** - Converts robot sensor data to EnvTransitions:
-
-```python
-# Use case: Live robot observation during inference
-robot_obs = {
-    "joint_1": 0.5, "joint_2": -0.3,  # joint positions
-    "camera_0": image_array            # camera images
-}
-transition = to_transition_robot_observation(robot_obs)
-# Creates: {OBSERVATION: {"observation.state.joint_1": 0.5, "observation.images.camera_0": image, ...}}
-```
-
-**`to_output_robot_action`** - Extracts robot-executable actions from EnvTransitions:
-
-```python
-# Use case: Converting model outputs back to robot commands
-model_transition = {ACTION: {"action.joint_1": 0.2, "action.joint_2": 0.1}}
-robot_action = to_output_robot_action(model_transition)
-# Returns: {"joint_1": 0.2, "joint_2": 0.1} - ready for robot.send_action()
-```
-
-**`to_dataset_frame`** - Converts transitions to dataset-compatible format:
-
-```python
-# Use case: Saving processed data or creating training batches
-features = {
-    "action": {"names": ["joint_1", "joint_2"]},
-    "observation.state": {"names": ["joint_1", "joint_2"]},
-    "observation.images.camera0": {...}
-}
-batch = to_dataset_frame(transition, features)
-# Returns: {"action": [0.2, 0.1], "observation.state": [0.5, -0.3], ...}
-```
-
-These converters are particularly useful when integrating with real robots, as shown in the examples:
-
-```python
-# Example from phone_so100_teleop.py - Real robot teleoperation
-phone_to_robot_ee_pose = RobotProcessor(
-    steps=[...],
-    to_transition=to_transition_teleop_action,  # Phone → EnvTransition
-    to_output=lambda tr: tr                     # Keep as EnvTransition
-)
-
-# Example from phone_so100_eval.py - Robot action execution
-robot_ee_to_joints = RobotProcessor(
-    steps=[...],
-    to_transition=lambda tr: tr,                # Already EnvTransition
-    to_output=to_output_robot_action           # EnvTransition → Robot action
-)
-
-# Example from phone_so100_record.py - Dataset recording
-robot_joints_to_ee_pose = RobotProcessor(
-    steps=[...],
-    to_transition=to_transition_robot_observation,  # Robot obs → EnvTransition
-    to_output=lambda tr: tr                         # Keep as EnvTransition for dataset
-)
-```
-
-### Data Format Conversion
-
-Different data sources have different formats, but processors need a unified `EnvTransition` structure internally.
-The default converters handle LeRobot datasets, but you can customize them:
-
-```python
-# Default: LeRobot batch format
-lerobot_batch = {
-    "observation.state": torch.tensor(...),
-    "action": torch.tensor(...),
-    "next.reward": torch.tensor(...),
-    "task": ["pick cube", ...]
-}
-# → Converts to EnvTransition → Processes → Converts back
-
-# Custom: Live robot data
-robot_data = {
-    "cameras": {"wrist_cam": np.array(...)},
-    "joint_positions": np.array(...),
-    "gripper_state": 0.5
-}
-
-def robot_to_transition(data: dict) -> EnvTransition:
-    return {
-        TransitionKey.OBSERVATION: {
-            "observation.images.wrist": torch.from_numpy(data["cameras"]["wrist_cam"]),
-            "observation.state": torch.from_numpy(data["joint_positions"])
-        },
-        TransitionKey.ACTION: None,
-        # ... other fields with defaults
-    }
-
-# Use custom converter
-processor = RobotProcessor(
-    steps=[...],
-    to_transition=robot_to_transition,
-    to_output=lambda transition: transition  # Keep as EnvTransition
-)
-```
-
-**When to customize:** Live robot data, Gymnasium environments, legacy datasets, or any non-LeRobot format.
-
-## Common Processor Steps
-
-LeRobot provides a rich set of pre-built processor steps for common transformations.
-Let's explore each in detail:
-
-### Data Normalization
-
-Normalization is crucial for neural network training and inference.
-The `NormalizerProcessor` handles both mean-std normalization and min-max scaling:
-
-```python
-from lerobot.processor.normalize_processor import NormalizerProcessor, UnnormalizerProcessor
-from lerobot.configs.types import PolicyFeature, FeatureType, NormalizationMode
-
-# Define what features exist in your data
-features = {
-    "observation.images.camera0": PolicyFeature(
-        type=FeatureType.IMAGE,
-        shape=(224, 224, 3)
-    ),
-    "observation.state": PolicyFeature(
-        type=FeatureType.STATE,
-        shape=(7,)
-    ),
-    "action": PolicyFeature(
-        type=FeatureType.ACTION,
-        shape=(7,)
-    )
-}
-
-# Define normalization strategy per feature type
-norm_map = {
-    FeatureType.IMAGE: NormalizationMode.MEAN_STD,    # Images: (x - mean) / std
-    FeatureType.STATE: NormalizationMode.MIN_MAX,     # States: scale to [-1, 1]
-    FeatureType.ACTION: NormalizationMode.MIN_MAX     # Actions: scale to [-1, 1]
-}
-
-# Create normalizer with dataset statistics
-normalizer = NormalizerProcessor(
-    features=features,
-    norm_map=norm_map,
-    stats=dataset.meta.stats,  # Contains mean, std, min, max per feature
-    normalize_keys={"observation.state", "action"}  # Optional: only normalize specific keys
-)
-
-# For postprocessing: inverse transformation
-unnormalizer = UnnormalizerProcessor(
-    features=features,
-    norm_map=norm_map,
-    stats=dataset.meta.stats
-)
-
-# The normalizer automatically:
-# - Detects which normalization to apply based on feature type
-# - Handles device placement of statistics tensors
-# - Skips keys not in stats or not in normalize_keys
-# - Adds metadata about what was normalized
-```
-
-### Device Management
-
-The `DeviceProcessor` ensures tensors are on the right device with the right dtype:
-
-```python
-from lerobot.processor.device_processor import DeviceProcessor
-
-# Basic GPU placement
-gpu_processor = DeviceProcessor(device="cuda:0")
-
-# Advanced: GPU with half-precision for inference
-efficient_processor = DeviceProcessor(
-    device="cuda:0",
-    float_dtype="float16"  # Convert float32 -> float16 for memory efficiency
-)
-
-# The processor:
-# - Moves all tensors to specified device
-# - Preserves non-tensor data unchanged
-# - Optionally converts float dtypes while preserving int/bool types
-# - Uses non_blocking transfers for CUDA devices
-# - Handles nested structures (observations, complementary_data)
-
-# Supported float dtypes:
-# "float16" / "half": 16-bit floating point
-# "float32" / "float": 32-bit floating point (default)
-# "float64" / "double": 64-bit floating point
-# "bfloat16": Brain floating point (better for training)
-```
-
-### Batch Processing
-
-Models expect batched inputs, but robot interactions often produce unbatched data:
-
-```python
-from lerobot.processor.batch_processor import ToBatchProcessor
-
-batch_processor = ToBatchProcessor()
-
-# Automatically adds batch dimensions where needed:
-# State: (7,) -> (1, 7)
-# Image: (224, 224, 3) -> (1, 224, 224, 3)
-# Action: (4,) -> (1, 4)
-# Task: "pick_cube" -> ["pick_cube"]
-# Already batched: (1, 7) -> (1, 7) [unchanged]
-
-# The processor intelligently:
-# - Detects tensor dimensionality
-# - Adds batch dim to 1D states/actions
-# - Adds batch dim to 3D images
-# - Wraps string tasks in lists
-# - Preserves already-batched data
-
-# Example usage in inference:
-single_observation = robot.get_observation()  # Unbatched
-batched_input = batch_processor({"observation": single_observation})
-model_output = model(batched_input)  # Model expects batch dim
-```
-
-### Text Tokenization
-
-For language-conditioned policies, text instructions must be tokenized:
-
-```python
-from lerobot.processor.tokenizer_processor import TokenizerProcessor
-from transformers import AutoTokenizer
-
-# Option 1: Auto-load tokenizer by name
-tokenizer_proc = TokenizerProcessor(
-    tokenizer_name="google/paligemma-3b-pt-224",
-    max_length=128,
-    task_key="task",           # Where to find text in complementary_data
-    padding="max_length",       # Pad to max_length
-    padding_side="right",
-    truncation=True            # Truncate if longer than max_length
-)
-
-# Option 2: Provide custom tokenizer
-custom_tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
-custom_proc = TokenizerProcessor(
-    tokenizer=custom_tokenizer,
-    max_length=256,
-    padding_side="left"        # For autoregressive models
-)
-
-# The processor:
-# - Extracts task text from complementary_data
-# - Tokenizes using HuggingFace tokenizer
-# - Adds tokens and attention_mask to observations
-# - Handles both single strings and lists of strings
-# - Preserves original task in complementary_data
-
-# Output structure:
-# observation["observation.language.tokens"] = tensor([101, 2032, ...])
-# observation["observation.language.attention_mask"] = tensor([1, 1, 0, ...])
-```
-
-### Key Renaming
-
-Different datasets and models may use different naming conventions.
-The `RenameProcessor` solves this mismatch:
-
-**Why is this useful?**
-
-- When loading a model trained on a different dataset with different key names
-- When using foundation models that expect specific key naming conventions
-- When standardizing datasets from different sources
-- When adapting legacy code to new naming standards
-
-```python
-from lerobot.processor.rename_processor import RenameProcessor
-
-# Example 1: Dataset uses "top"/"wrist", model expects "camera0"/"camera1"
-rename_proc = RenameProcessor(
-    rename_map={
-        "observation.images.top": "observation.images.camera0",
-        "observation.images.wrist": "observation.images.camera1",
-    }
-)
-
-# Example 2: Foundation model compatibility
-# Your dataset: "observation.state", Foundation model: "proprio"
-foundation_rename = RenameProcessor(
-    rename_map={
-        "observation.state": "proprio",
-        "observation.images.main": "rgb",
-    }
-)
-
-# Example 3: Standardizing multiple datasets
-standardize_rename = RenameProcessor(
-    rename_map={
-        # Different robots might use different names
-        "observation.joint_positions": "observation.state",
-        "observation.gripper_state": "observation.end_effector",
-        "observation.arm_camera": "observation.images.wrist",
-    }
-)
-```
-
-## Building Complete Pipelines
-
-Let's build a real-world preprocessing and postprocessing pipeline for a vision-based
-manipulation policy:
-
-```python
-# Consolidated imports
-from lerobot.processor import (
-    RobotProcessor,
-    NormalizerProcessor,
-    UnnormalizerProcessor,
-    DeviceProcessor,
-    ToBatchProcessor,
-    TokenizerProcessor,
-    RenameProcessor
-)
-
-# Step 1: Define the preprocessing pipeline
-preprocessor = RobotProcessor(
-    steps=[
-        # 1. Standardize naming from dataset
-        RenameProcessor(
-            rename_map={
-                "observation.images.top": "observation.images.camera0",
-                "observation.images.wrist": "observation.images.camera1"
-            }
-        ),
-
-        # 2. Add batch dimensions for model
-        ToBatchProcessor(),
-
-        # 3. Tokenize language instructions if present
-        TokenizerProcessor(
-            tokenizer_name="google/paligemma-3b-pt-224",
-            max_length=64,
-            task_key="task"
-        ),
-
-        # 4. Normalize numerical data
-        NormalizerProcessor(
-            features=policy_features,
-            norm_map={
-                FeatureType.IMAGE: NormalizationMode.MEAN_STD,
-                FeatureType.STATE: NormalizationMode.MIN_MAX,
-                FeatureType.ACTION: NormalizationMode.MIN_MAX
-            },
-            stats=dataset.meta.stats
-        ),
-
-        # 5. Move to GPU and convert to half precision
-        DeviceProcessor(
-            device="cuda:0",
-            float_dtype="float16"
-        )
-    ],
-    name="robot_preprocessor"
-)
-
-# Step 2: Define the postprocessing pipeline
-postprocessor = RobotProcessor(
-    steps=[
-        # 1. Move back to CPU for robot hardware
-        DeviceProcessor(device="cpu"),
-
-        # 2. Denormalize actions to original scale
-        UnnormalizerProcessor(
-            features=policy_features,
-            norm_map={
-                FeatureType.ACTION: NormalizationMode.MIN_MAX
-            },
-            stats=dataset.meta.stats
-        )
-    ],
-    name="robot_postprocessor"
-)
-```
-
-## Using Processors in Practice
-
-### Training Loop Integration
-
-Here's how processors integrate into a training loop using the policy's forward method:
-
-```python
-from torch.utils.data import DataLoader
-
-# Create dataset and dataloader
-dataset = LeRobotDataset(repo_id="your_dataset")
-dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
-
-# Initialize model and processors
-model = YourPolicy.from_pretrained("your_model")
-preprocessor = RobotProcessor.from_pretrained(
-    "your_model",
-    config_filename="robot_preprocessor.json"
-)
-
-# Training loop
-for epoch in range(num_epochs):
-    for batch in dataloader:
-        # Preprocess batch
-        processed_batch = preprocessor(batch)
-
-        # Forward pass - returns loss and optional metrics
-        loss, metrics = model.forward(processed_batch)
-
-        # Backward pass
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-
-        # Log metrics if available
-        if metrics:
-            wandb.log(metrics)
-```
-
-### Inference Pipeline
-
-For deployment, processors ensure consistent data handling with real robots:
-
-```python
-# Load model and processors
-policy = YourPolicy.from_pretrained("path/to/model")
-preprocessor = RobotProcessor.from_pretrained(
-    "path/to/model",
-    config_filename="robot_preprocessor.json"
-)
-postprocessor = RobotProcessor.from_pretrained(
-    "path/to/model",
-    config_filename="robot_postprocessor.json"
-)
-
-# Connect to robot
-robot = make_robot_from_config(robot_config)
-robot.connect()
-
-# Inference loop
-policy.eval()
-# Reset the policy and processors
-policy.reset()
-preprocessor.reset()
-postprocessor.reset()
-
-with torch.no_grad():
-    while not done:
-        # Get observation from robot
-        observation = robot.get_observation()
-
-        # Build dataset-compatible frame
-        observation_frame = build_dataset_frame(
-            dataset.features,
-            observation,
-            prefix="observation"
-        )
-
-        # Add task instruction to complementary data
-        observation_frame["task"] = "pick up the red cube"
-
-        # Preprocess for model
-        model_input = preprocessor(observation_frame)
-
-        # Run policy
-        raw_action = policy.select_action(model_input)
-
-        # Postprocess action
-        action_transition = {TransitionKey.ACTION: raw_action}
-        processed = postprocessor(action_transition)
-        action = processed[TransitionKey.ACTION]
-
-        # Convert to robot action format
-        robot_action = {
-            key: action[i].item()
-            for i, key in enumerate(robot.action_features)
-        }
-
-        # Execute on robot
-        robot.send_action(robot_action)
-```
-
-## Saving and Loading Processors
-
-Processors can be persisted and shared just like models, making them portable across different
-environments and ensuring reproducibility:
-
-### Local Save/Load
-
-```python
-# Save processor configuration and state
-preprocessor.save_pretrained(
-    "./my_robot_processor",
-    config_filename="preprocessor.json"  # Optional custom name
-)
-
-# The save creates:
-# my_robot_processor/
-# ├── preprocessor.json                    # Configuration
-# ├── preprocessor_step_0_normalizer.safetensors  # Step 0 state (stats)
-# └── preprocessor_step_1_device.safetensors     # Step 1 state (if any)
-
-# Load processor
-loaded = RobotProcessor.from_pretrained(
-    "./my_robot_processor",
-    config_filename="preprocessor.json"
-)
-```
-
-### HuggingFace Hub Integration
-
-The HuggingFace Hub provides a centralized place to share and version your processors.
-This is particularly useful for sharing preprocessing configurations with models,
-ensuring that anyone who downloads your model can reproduce your exact preprocessing pipeline.
-It also enables versioning and collaboration on preprocessing strategies.
-
-```python
-# Save to HuggingFace Hub
-preprocessor.save_pretrained("username/my-robot-policy")
-
-# Load from Hub with automatic download
-hub_processor = RobotProcessor.from_pretrained(
-    "username/my-robot-policy",
-    config_filename="robot_preprocessor.json",
-    revision="main",  # Optional: specific revision
-    cache_dir="./cache"  # Optional: local cache directory
-)
-
-# The Hub integration provides:
-# - Automatic versioning with git
-# - Public or private sharing
-# - Download caching for efficiency
-# - Integration with model repositories
-```
-
-### Loading with Overrides
-
-Sometimes you need to modify loaded processors for new environments or datasets.
-The override mechanism allows you to update specific processor configurations without modifying
-the saved files:
-
-```python
-# Load processor with configuration overrides
-processor = RobotProcessor.from_pretrained(
-    "./saved_processor",
-    overrides={
-        # Change device for different hardware
-        "device_processor": {"device": "cuda:1"},
-
-        # Update statistics for new dataset
-        "normalizer_processor": {"stats": new_dataset.meta.stats},
-
-        # Provide non-serializable objects (like tokenizers)
-        "tokenizer_processor": {"tokenizer": custom_tokenizer}
-    }
-)
-
-# Common override scenarios:
-# 1. Adapting to different hardware (GPU availability)
-# 2. Fine-tuning on new datasets with different statistics
-# 3. Providing runtime dependencies that can't be serialized
-# 4. Testing variations without creating new saved configs
-```
-
-## Creating Custom Processor Steps
-
-Build your own processor steps for specialized transformations.
-The key is implementing the required interface:
-
-### Basic Custom Step with Registration
-
-The registration mechanism allows your custom processors to be saved and loaded by name rather
-than by module path.
-This makes them more portable and easier to share:
-
-```python
-from dataclasses import dataclass
-from lerobot.processor.pipeline import ProcessorStepRegistry, ObservationProcessor
-
-# The @register decorator adds your processor to the global registry
-# Use a unique name, preferably namespaced to avoid conflicts
-@dataclass
-@ProcessorStepRegistry.register("my_company/gaussian_noise")
-class GaussianNoiseProcessor(ObservationProcessor):
-    """Add Gaussian noise to observations for robustness training."""
-
-    noise_std: float = 0.01
-    training_only: bool = True
-    is_training: bool = True
-
-    def observation(self, observation):
-        """Add noise to observation tensors."""
-        if not self.is_training and self.training_only:
-            return observation
-
-        noisy_obs = {}
-        for key, value in observation.items():
-            if isinstance(value, torch.Tensor) and "image" not in key:
-                # Add noise to non-image observations
-                noise = torch.randn_like(value) * self.noise_std
-                noisy_obs[key] = value + noise
-            else:
-                noisy_obs[key] = value
-
-        return noisy_obs
-
-    def get_config(self):
-        return {
-            "noise_std": self.noise_std,
-            "training_only": self.training_only,
-            "is_training": self.is_training
-        }
-
-# Why register?
-# 1. Enables saving by name: config saves "my_company/gaussian_noise" instead of full module path
-# 2. More portable: Others can use your processor without your exact module structure
-# 3. Version-safe: Module refactoring won't break saved configs
-# 4. Cleaner configs: JSON shows readable names instead of long import paths
-```
-
-### Using Base Classes for Common Patterns
-
-LeRobot provides base classes like `ObservationProcessor`, `ActionProcessor`, etc., that handle
-the boilerplate of extracting and reinserting specific components:
-
-```python
-from lerobot.processor import ActionProcessor
-
-@dataclass
-@ProcessorStepRegistry.register("my_company/action_clipper")
-class ActionClipProcessor(ActionProcessor):
-    """Clip actions to safe ranges."""
-
-    min_value: float = -1.0
-    max_value: float = 1.0
-
-    def action(self, action):
-        """Process only the action component."""
-        # No need to handle transition dict - base class does it
-        return torch.clamp(action, self.min_value, self.max_value)
-
-    def get_config(self):
-        return {"min_value": self.min_value, "max_value": self.max_value}
-```
-
-For more advanced processor patterns including stateful processors, see [Implement Your Own Processor](implement_your_own_processor.mdx).
-
-## Advanced Features
-
-### Debugging with Hooks
-
-Processors support hooks for monitoring and debugging without modifying the pipeline code:
-
-```python
-# Define monitoring hooks
-def log_shapes(step_idx: int, transition: EnvTransition):
-    """Log tensor shapes after each step."""
-    obs = transition.get(TransitionKey.OBSERVATION)
-    if obs:
-        print(f"Step {step_idx} shapes:")
-        for key, value in obs.items():
-            if isinstance(value, torch.Tensor):
-                print(f"  {key}: {value.shape}")
-
-def check_nans(step_idx: int, transition: EnvTransition):
-    """Check for NaN values."""
-    obs = transition.get(TransitionKey.OBSERVATION)
-    if obs:
-        for key, value in obs.items():
-            if isinstance(value, torch.Tensor) and torch.isnan(value).any():
-                print(f"Warning: NaN detected in {key} at step {step_idx}")
-
-# Register hooks
-processor.register_after_step_hook(log_shapes)
-processor.register_after_step_hook(check_nans)
-
-# Process data - hooks will be called after each step
-output = processor(input_data)
-
-# Remove hooks when done debugging
-processor.unregister_after_step_hook(log_shapes)
-processor.unregister_after_step_hook(check_nans)
-```
-
-### Step-by-Step Inspection
-
-Use `step_through()` for detailed debugging of the transformation pipeline:
-
-```python
-# Inspect data at each transformation stage
-for i, intermediate in enumerate(processor.step_through(data)):
-    print(f"\n=== After step {i} ===")
-
-    # Check observation shapes
-    obs = intermediate.get(TransitionKey.OBSERVATION)
-    if obs:
-        for key, value in obs.items():
-            if isinstance(value, torch.Tensor):
-                print(f"{key}: shape={value.shape}, "
-                      f"dtype={value.dtype}, "
-                      f"device={value.device}, "
-                      f"range=[{value.min():.3f}, {value.max():.3f}]")
-
-    # Check action if present
-    action = intermediate.get(TransitionKey.ACTION)
-    if action is not None and isinstance(action, torch.Tensor):
-        print(f"action: shape={action.shape}, range=[{action.min():.3f}, {action.max():.3f}]")
-```
-
-### Pipeline Slicing
-
-Extract subsets of a pipeline for testing or creating variations:
-
-```python
-# Get specific steps
-first_three_steps = processor[:3]  # Returns new RobotProcessor
-middle_step = processor[2]         # Returns single ProcessorStep
-
-# Test individual steps
-test_input = {...}
-step_output = processor[0](test_input)  # Test first step only
-
-# Create variations
-variant_processor = RobotProcessor(
-    steps=processor.steps[:-1] + [new_final_step],
-    name="variant"
-)
-```
-
-## Best Practices and Tips
-
-### 1. Order Matters
-
-The sequence of processors is crucial. Follow this general order:
-
-```python
-# Preprocessing: Raw → Model-ready
-1. Rename (standardize keys)
-2. Batch (add dimensions)
-3. Tokenize (text → tokens)
-4. Normalize (scale values)
-5. Device (move to GPU)
-
-# Postprocessing: Model → Robot-ready
-1. Device (move to CPU)
-2. Unnormalize (restore scale)
-3. Unbatch (remove dimensions if needed)
-```
-
-### 2. Registration Best Practices
-
-```python
-# Always register custom steps for better portability
-@ProcessorStepRegistry.register("my_company/special_processor")
-class SpecialProcessor:
-    ...
-
-# Use namespaced names to avoid conflicts
-# Good: "my_company/augmentation"
-# Bad: "augmentation" (too generic)
-
-# Check registered processors
-print(ProcessorStepRegistry.list())  # See all registered processors
-```
-
-### 3. Common Pitfalls and Solutions
-
-**Tensor Device Mismatch:**
-
-```python
-# Problem: RuntimeError: Expected all tensors on same device
-# Solution: Ensure DeviceProcessor is in pipeline
-preprocessor = RobotProcessor(
-    steps=[
-        NormalizerProcessor(...),
-        DeviceProcessor(device="cuda")  # Add this
-    ]
-)
-```
-
-**Missing Statistics:**
-
-```python
-# Problem: NormalizerProcessor has no stats
-# Solution 1: Compute stats from dataset
-from lerobot.datasets.compute_stats import compute_stats
-stats = compute_stats(dataset)
-
-# Solution 2: Load with overrides
-processor = RobotProcessor.from_pretrained(
-    "model_path",
-    overrides={"normalizer_processor": {"stats": dataset.meta.stats}}
-)
-```
-
-## Next Steps
-
-Now that you understand processors, explore these topics:
-
-- [**Implement Your Own Processor**](implement_your_own_processor.mdx) - Deep dive into creating custom processors with advanced features like stateful processing
-- [**Policy Documentation**](policies.mdx) - Learn how different policies use processors
-- [**Dataset Documentation**](datasets.mdx) - Understand the data format that processors transform
-- [**Training Guide**](training.mdx) - See processors in action during model training
-- [**Evaluation Guide**](evaluation.mdx) - Learn about processor usage during policy evaluation
-
-## Summary
-
-Processors are the unsung heroes of robotics pipelines, handling the critical transformations between raw sensor data and model-ready tensors. By understanding and effectively using processors, you can:
-
-- Build robust, reusable data pipelines
-- Share preprocessing configurations across projects
-- Debug data transformations systematically
-- Ensure consistency between training and deployment
-- Create custom transformations for specialized tasks
-
-Remember: good preprocessing is often the difference between a model that works in theory
-and one that works in practice!
-The modular pipeline approach ensures your transformations are testable, reproducible,
-and portable across different robots and environments.

docs/source/koch.mdx

@@ -31,7 +31,7 @@ pip install -e ".[dynamixel]"
 To find the port for each bus servo adapter, run this script:
 
 ```bash
-python -m lerobot.find_port
+lerobot-find-port
 ```
 
 <hfoptions id="example">
@@ -98,7 +98,7 @@ For a visual reference on how to set the motor ids please refer to [this video](
 <hfoption id="Command">
 
 ```bash
-python -m lerobot.setup_motors \
+lerobot-setup-motors \
     --robot.type=koch_follower \
     --robot.port=/dev/tty.usbmodem575E0031751 # <- paste here the port found at previous step
 ```
@@ -174,7 +174,7 @@ Do the same steps for the leader arm but modify the command or script accordingl
 <hfoption id="Command">
 
 ```bash
-python -m lerobot.setup_motors \
+lerobot-setup-motors \
     --teleop.type=koch_leader \
     --teleop.port=/dev/tty.usbmodem575E0031751 \  # <- paste here the port found at previous step
 ```
@@ -211,7 +211,7 @@ Run the following command or API example to calibrate the follower arm:
 <hfoption id="Command">
 
 ```bash
-python -m lerobot.calibrate \
+lerobot-calibrate \
     --robot.type=koch_follower \
     --robot.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --robot.id=my_awesome_follower_arm  # <- Give the robot a unique name
@@ -249,7 +249,7 @@ Do the same steps to calibrate the leader arm, run the following command or API
 <hfoption id="Command">
 
 ```bash
-python -m lerobot.calibrate \
+lerobot-calibrate \
     --teleop.type=koch_leader \
     --teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --teleop.id=my_awesome_leader_arm  # <- Give the robot a unique name

docs/source/lekiwi.mdx

@@ -60,7 +60,7 @@ First, we will assemble the two SO100/SO101 arms. One to attach to the mobile ba
 To find the port for each bus servo adapter, run this script:
 
 ```bash
-python -m lerobot.find_port
+lerobot-find-port
 ```
 
 <hfoptions id="example">
@@ -116,7 +116,7 @@ The instructions for configuring the motors can be found in the SO101 [docs](./s
 You can run this command to setup motors for LeKiwi. It will first setup the motors for arm (id 6..1) and then setup motors for wheels (9,8,7)
 
 ```bash
-python -m lerobot.setup_motors \
+lerobot-setup-motors \
     --robot.type=lekiwi \
     --robot.port=/dev/tty.usbmodem58760431551 # <- paste here the port found at previous step
 ```
@@ -174,7 +174,7 @@ The calibration process is very important because it allows a neural network tra
 Make sure the arm is connected to the Raspberry Pi and run this script or API example (on the Raspberry Pi via SSH) to launch calibration of the follower arm:
 
 ```bash
-python -m lerobot.calibrate \
+lerobot-calibrate \
     --robot.type=lekiwi \
     --robot.id=my_awesome_kiwi # <- Give the robot a unique name
 ```
@@ -193,7 +193,7 @@ Then, to calibrate the leader arm (which is attached to the laptop/pc). Run the
 <hfoption id="Command">
 
 ```bash
-python -m lerobot.calibrate \
+lerobot-calibrate \
     --teleop.type=so100_leader \
     --teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --teleop.id=my_awesome_leader_arm # <- Give the robot a unique name

docs/source/phone_teleop.mdx

@@ -1,195 +0,0 @@
-# Phone
-
-Use your phone (iOS or Android) to control your robot.
-
-**In this guide you'll learn:**
-
-- How to connect an iOS/Android phone
-- How phone pose is mapped to robot end‑effector (EE) targets
-- How to tweak safety limits, gripper control, and IK settings
-
-To use phone to control your robot, install the relevant dependencies with:
-
-```bash
-pip install lerobot[phone]
-```
-
-## Get started
-
-### Supported platforms
-
-- iOS: Uses the HEBI Mobile I/O app (ARKit pose + buttons). Download the app first, open it and the examples will discover it on your network and stream the phone pose and inputs.
-- Android: Uses the `teleop` package (WebXR). When you start the Python process, it prints a local URL. Open the link on your phone, tap Start, then use Move to stream pose.
-
-Links:
-
-- Android WebXR library: [`teleop` on PyPI](https://pypi.org/project/teleop/)
-- iOS app: [HEBI Mobile I/O](https://docs.hebi.us/tools.html#mobile-io)
-
-### 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.
-- 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.
-- Gripper control:
-  - 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.
-
-### 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.
-
-Teleoperation example:
-
-```36:43:examples/phone_so100_teleop.py
-from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
-
-teleop_config = PhoneConfig(phone_os=PhoneOS.IOS)  # or PhoneOS.ANDROID
-teleop_device = Phone(teleop_config)
-```
-
-### Step 2: Connect and calibrate
-
-When `Phone(teleop_config)` is created and `connect()` is called, calibration is prompted automatically. Hold the phone in the orientation described above, then:
-
-- iOS: press and hold `B1` to capture the reference pose.
-- Android: press `Move` button on the WebXR page to capture the reference pose.
-
-Why calibrate? We capture the current pose so subsequent poses are expressed in a robot aligned frame. When you again press the button to enable control, the position is recaptured to avoid drift when your phone is repositioned while it was disabled.
-
-### Step 3: Run an example
-
-Run on of the examples scripts to teleoperate, record a dataset, replay a dataset or evaluate a policy.
-
-All scripts assume you configured your robot (e.g., SO-100 follower) and set the correct serial port.
-
-- 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_so100_teleop.py
-  ```
-
-- Run this example to record a dataset, which saves absolute end effector observations and actions:
-
-  ```bash
-  python examples/phone_so100_record.py
-  ```
-
-- Run this example to replay recorded episodes:
-
-  ```bash
-  python examples/phone_so100_replay.py
-  ```
-
-- Run this example to evaluate a pretrained policy:
-
-  ```bash
-  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.
-
-  ```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.
-
-  ```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.
-
-  ```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()),
-  )
-  ```
-
-- 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.
-
-  ```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,
-      max_ee_twist_step_rad=0.50,
-  )
-  ```
-
-- 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.
-
-  ```78:81:examples/phone_so100_teleop.py
-  GripperVelocityToJoint(
-      motor_names=list(robot.bus.motors.keys()),
-      speed_factor=20.0,
-  )
-  ```
-
-#### Different IK initial guesses
-
-We use different IK initial guesses in the kinematic steps. As initial guess either the current measured joints or the previous IK solution is used.
-
-- Closed loop (used in record/eval): sets `initial_guess_current_joints=True` so IK starts from the measured joints each frame.
-
-  ```71:76:examples/phone_so100_eval.py
-  InverseKinematicsEEToJoints(
-      kinematics=kinematics_solver,
-      motor_names=list(robot.bus.motors.keys()),
-      initial_guess_current_joints=True,  # closed loop
-  )
-  ```
-
-- 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.
-
-  ```80:86:examples/phone_so100_replay.py
-  InverseKinematicsEEToJoints(
-      kinematics=kinematics_solver,
-      motor_names=list(robot.bus.motors.keys()),
-      initial_guess_current_joints=False,  # open loop
-  )
-  ```
-
-### 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.
-- GripperVelocityToJoint: integrates a velocity‑like gripper input into an absolute gripper position using the current measured state.
-- ForwardKinematicsJointsToEE: computes `observation.state.ee.*` from observed joints for logging and training on EE state.
-
-### Troubleshooting
-
-- iOS not discovered: ensure HEBI Mobile I/O is open and your laptop/phone are on the same network.
-- Android URL not reachable: check local you used `https` instead of `http`, use the exact IP printed by the script and allow your browser to enter and ignore the certificate issue.
-- Motion feels inverted: adjust the sign flips in `MapPhoneActionToRobotAction` or swap axes to match your setup.

docs/source/porting_datasets_v3.mdx

@@ -0,0 +1,321 @@
+# Porting Large Datasets to LeRobot Dataset v3.0
+
+This tutorial explains how to port large-scale robotic datasets to the LeRobot Dataset v3.0 format. We'll use the **DROID 1.0.1** dataset as our primary example, which demonstrates handling multi-terabyte datasets with thousands of shards across SLURM clusters.
+
+## File Organization: v2.1 vs v3.0
+
+Dataset v3.0 fundamentally changes how data is organized and stored:
+
+**v2.1 Structure (Episode-based)**:
+
+```
+dataset/
+├── data/chunk-000/episode_000000.parquet
+├── data/chunk-000/episode_000001.parquet
+├── videos/chunk-000/camera/episode_000000.mp4
+└── meta/episodes.jsonl
+```
+
+**v3.0 Structure (File-based)**:
+
+```
+dataset/
+├── data/chunk-000/file-000.parquet        # Multiple episodes per file
+├── videos/camera/chunk-000/file-000.mp4   # Consolidated video chunks
+└── meta/episodes/chunk-000/file-000.parquet  # Structured metadata
+```
+
+This transition from individual episode files to file-based chunks dramatically improves performance and reduces storage overhead.
+
+## What's New in Dataset v3.0
+
+Dataset v3.0 introduces significant improvements for handling large datasets:
+
+### 🏗️ **Enhanced File Organization**
+
+- **File-based structure**: Episodes are now grouped into chunked files rather than individual episode files
+- **Configurable file sizes**: for data and video files
+- **Improved storage efficiency**: Better compression and reduced overhead
+
+### 📊 **Modern Metadata Management**
+
+- **Parquet-based metadata**: Replaced JSON Lines with efficient parquet format
+- **Structured episode access**: Direct pandas DataFrame access via `dataset.meta.episodes`
+- **Per-episode statistics**: Enhanced statistics tracking at episode level
+
+### 🚀 **Performance Enhancements**
+
+- **Memory-mapped access**: Improved RAM usage through PyArrow memory mapping
+- **Faster loading**: Significantly reduced dataset initialization time
+- **Better scalability**: Designed for datasets with millions of episodes
+
+## Prerequisites
+
+Before porting large datasets, ensure you have:
+
+- **LeRobot installed** with v3.0 support. Follow our [Installation Guide](./installation).
+- **Sufficient storage**: Raw datasets can be very large (e.g., DROID requires 2TB)
+- **Cluster access** (recommended for large datasets): SLURM or similar job scheduler
+- **Dataset-specific dependencies**: For DROID, you'll need TensorFlow Dataset utilities
+
+## Understanding the DROID Dataset
+
+[DROID 1.0.1](https://droid-dataset.github.io/droid/the-droid-dataset) is an excellent example of a large-scale robotic dataset:
+
+- **Size**: 1.7TB (RLDS format), 8.7TB (raw data)
+- **Structure**: 2048 pre-defined TensorFlow dataset shards
+- **Content**: 76,000+ robot manipulation trajectories from Franka Emika Panda robots
+- **Scope**: Real-world manipulation tasks across multiple environments and objects
+- **Format**: Originally in TensorFlow Records/RLDS format, requiring conversion to LeRobot format
+- **Hosting**: Google Cloud Storage with public access via `gsutil`
+
+The dataset contains diverse manipulation demonstrations with:
+
+- Multiple camera views (wrist camera, exterior cameras)
+- Natural language task descriptions
+- Robot proprioceptive state and actions
+- Success/failure annotations
+
+### DROID Features Schema
+
+```python
+DROID_FEATURES = {
+    # Episode markers
+    "is_first": {"dtype": "bool", "shape": (1,)},
+    "is_last": {"dtype": "bool", "shape": (1,)},
+    "is_terminal": {"dtype": "bool", "shape": (1,)},
+
+    # Language instructions
+    "language_instruction": {"dtype": "string", "shape": (1,)},
+    "language_instruction_2": {"dtype": "string", "shape": (1,)},
+    "language_instruction_3": {"dtype": "string", "shape": (1,)},
+
+    # Robot state
+    "observation.state.gripper_position": {"dtype": "float32", "shape": (1,)},
+    "observation.state.cartesian_position": {"dtype": "float32", "shape": (6,)},
+    "observation.state.joint_position": {"dtype": "float32", "shape": (7,)},
+
+    # Camera observations
+    "observation.images.wrist_left": {"dtype": "image"},
+    "observation.images.exterior_1_left": {"dtype": "image"},
+    "observation.images.exterior_2_left": {"dtype": "image"},
+
+    # Actions
+    "action.gripper_position": {"dtype": "float32", "shape": (1,)},
+    "action.cartesian_position": {"dtype": "float32", "shape": (6,)},
+    "action.joint_position": {"dtype": "float32", "shape": (7,)},
+
+    # Standard LeRobot format
+    "observation.state": {"dtype": "float32", "shape": (8,)},  # joints + gripper
+    "action": {"dtype": "float32", "shape": (8,)},  # joints + gripper
+}
+```
+
+## Approach 1: Single Computer Porting
+
+### Step 1: Install Dependencies
+
+For DROID specifically:
+
+```bash
+pip install tensorflow
+pip install tensorflow_datasets
+```
+
+For other datasets, install the appropriate readers for your source format.
+
+### Step 2: Download Raw Data
+
+Download DROID from Google Cloud Storage using `gsutil`:
+
+```bash
+# Install Google Cloud SDK if not already installed
+# https://cloud.google.com/sdk/docs/install
+
+# Download the full RLDS dataset (1.7TB)
+gsutil -m cp -r gs://gresearch/robotics/droid/1.0.1 /your/data/
+
+# Or download just the 100-episode sample (2GB) for testing
+gsutil -m cp -r gs://gresearch/robotics/droid_100 /your/data/
+```
+
+> [!WARNING]
+> Large datasets require substantial time and storage:
+>
+> - **Full DROID (1.7TB)**: Several days to download depending on bandwidth
+> - **Processing time**: 7+ days for local porting of full dataset
+> - **Upload time**: 3+ days to push to Hugging Face Hub
+> - **Local storage**: ~400GB for processed LeRobot format
+
+### Step 3: Port the Dataset
+
+```bash
+python examples/port_datasets/droid_rlds/port_droid.py \
+    --raw-dir /your/data/droid/1.0.1 \
+    --repo-id your_id/droid_1.0.1 \
+    --push-to-hub
+```
+
+### Development and Testing
+
+For development, you can port a single shard:
+
+```bash
+python examples/port_datasets/droid_rlds/port_droid.py \
+    --raw-dir /your/data/droid/1.0.1 \
+    --repo-id your_id/droid_1.0.1_test \
+    --num-shards 2048 \
+    --shard-index 0
+```
+
+This approach works for smaller datasets or testing, but large datasets require cluster computing.
+
+## Approach 2: SLURM Cluster Porting (Recommended)
+
+For large datasets like DROID, parallel processing across multiple nodes dramatically reduces processing time.
+
+### Step 1: Install Cluster Dependencies
+
+```bash
+pip install datatrove  # Hugging Face's distributed processing library
+```
+
+### Step 2: Configure Your SLURM Environment
+
+Find your partition information:
+
+```bash
+sinfo --format="%R"  # List available partitions
+sinfo -N -p your_partition -h -o "%N cpus=%c mem=%m"  # Check resources
+```
+
+Choose a **CPU partition** - no GPU needed for dataset porting.
+
+### Step 3: Launch Parallel Porting Jobs
+
+```bash
+python examples/port_datasets/droid_rlds/slurm_port_shards.py \
+    --raw-dir /your/data/droid/1.0.1 \
+    --repo-id your_id/droid_1.0.1 \
+    --logs-dir /your/logs \
+    --job-name port_droid \
+    --partition your_partition \
+    --workers 2048 \
+    --cpus-per-task 8 \
+    --mem-per-cpu 1950M
+```
+
+#### Parameter Guidelines
+
+- **`--workers`**: Number of parallel jobs (max 2048 for DROID's shard count)
+- **`--cpus-per-task`**: 8 CPUs recommended for frame encoding parallelization
+- **`--mem-per-cpu`**: ~16GB total RAM (8×1950M) for loading raw frames
+
+> [!TIP]
+> Start with fewer workers (e.g., 100) to test your cluster configuration before launching thousands of jobs.
+
+### Step 4: Monitor Progress
+
+Check running jobs:
+
+```bash
+squeue -u $USER
+```
+
+Monitor overall progress:
+
+```bash
+jobs_status /your/logs
+```
+
+Inspect individual job logs:
+
+```bash
+less /your/logs/port_droid/slurm_jobs/JOB_ID_WORKER_ID.out
+```
+
+Debug failed jobs:
+
+```bash
+failed_logs /your/logs/port_droid
+```
+
+### Step 5: Aggregate Shards
+
+Once all porting jobs complete:
+
+```bash
+python examples/port_datasets/droid_rlds/slurm_aggregate_shards.py \
+    --repo-id your_id/droid_1.0.1 \
+    --logs-dir /your/logs \
+    --job-name aggr_droid \
+    --partition your_partition \
+    --workers 2048 \
+    --cpus-per-task 8 \
+    --mem-per-cpu 1950M
+```
+
+### Step 6: Upload to Hub
+
+```bash
+python examples/port_datasets/droid_rlds/slurm_upload.py \
+    --repo-id your_id/droid_1.0.1 \
+    --logs-dir /your/logs \
+    --job-name upload_droid \
+    --partition your_partition \
+    --workers 50 \
+    --cpus-per-task 4 \
+    --mem-per-cpu 1950M
+```
+
+> [!NOTE]
+> Upload uses fewer workers (50) since it's network-bound rather than compute-bound.
+
+## Dataset v3.0 File Structure
+
+Your completed dataset will have this modern structure:
+
+```
+dataset/
+├── meta/
+│   ├── episodes/
+│   │   └── chunk-000/
+│   │       └── file-000.parquet    # Episode metadata
+│   ├── tasks.parquet               # Task definitions
+│   ├── stats.json                  # Aggregated statistics
+│   └── info.json                   # Dataset information
+├── data/
+│   └── chunk-000/
+│       └── file-000.parquet        # Consolidated episode data
+└── videos/
+    └── camera_key/
+        └── chunk-000/
+            └── file-000.mp4        # Consolidated video files
+```
+
+This replaces the old episode-per-file structure with efficient, optimally-sized chunks.
+
+## Migrating from Dataset v2.1
+
+If you have existing datasets in v2.1 format, use the migration tool:
+
+```bash
+python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
+    --repo-id your_id/existing_dataset
+```
+
+This automatically:
+
+- Converts file structure to v3.0 format
+- Migrates metadata from JSON Lines to parquet
+- Aggregates statistics and creates per-episode stats
+- Updates version information
+
+## Performance Benefits
+
+Dataset v3.0 provides significant improvements for large datasets:
+
+- **Faster loading**: 3-5x reduction in initialization time
+- **Memory efficiency**: Better RAM usage through memory mapping
+- **Scalable processing**: Handles millions of episodes efficiently
+- **Storage optimization**: Reduced file count and improved compression

docs/source/processors_robots_teleop.mdx

@@ -1,148 +0,0 @@
-# Processors for Robots and Teleoperators
-
-This guide shows how to build and modify processing pipelines that connect teleoperators (e.g., phone) to robots and datasets. Pipelines standardize conversions between different action/observation spaces so you can swap teleops and robots without rewriting glue code.
-
-We use the Phone to SO‑100 follower examples for concreteness, but the same patterns apply to other robots.
-
-**What you'll learn**
-
-- Absolute vs. relative EE control: What each means, trade‑offs, and how to choose for your task.
-- Three-pipeline pattern: How to map teleop actions → dataset actions → robot commands, and robot observations → dataset observations.
-- Adapters (`to_transition` / `to_output`): How these convert raw dicts to `EnvTransition` and back to reduce boilerplate.
-- Dataset feature contracts: How steps declare features via `transform_features(...)`, and how to aggregate/merge them for recording.
-- Choosing a representation: When to store joints, absolute EE poses, or relative EE deltas—and how that affects training.
-- Pipeline customization guidance: How to swap robots/URDFs safely and tune bounds, step sizes, and options like IK initialization.
-
-### Absolute vs relative EE control
-
-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.
-
-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
-
-We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
-Each of these pipelines handle different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
-
-1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets)
-2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
-3. Pipeline 3: Robot observation space → dataset observation space (joints → EE pose)
-
-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 = 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 = 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 = 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 dicts and the pipeline’s `EnvTransition` format.
-In the phone to SO-100 follower examples we use the following adapters:
-
-- `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.*`)
-
-See `src/lerobot/processor/converters.py` for more details.
-
-## Dataset feature contracts
-
-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:
-
-```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
-```
-
-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 follower examples:
-
-```121:145:examples/phone_so100_record.py
-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`).
-- `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
-
-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, etc.
-
-Pick what you want to use for your policy action and observation space and configure/modify the pipelines and steps accordingly.
-
-### Different robots
-
-- 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.
-
-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

@@ -54,7 +54,7 @@ If you don't have a gpu device, you can train using our notebook on [![Google Co
 Pass your dataset to the training script using `--dataset.repo_id`. If you want to test your installation, run the following command where we use one of the datasets we collected for the [SmolVLA Paper](https://huggingface.co/papers/2506.01844).
 
 ```bash
-cd lerobot && python -m lerobot.scripts.train \
+cd lerobot && lerobot-train \
   --policy.path=lerobot/smolvla_base \
   --dataset.repo_id=${HF_USER}/mydataset \
   --batch_size=64 \
@@ -73,7 +73,7 @@ cd lerobot && python -m lerobot.scripts.train \
 Fine-tuning is an art. For a complete overview of the options for finetuning, run
 
 ```bash
-python -m lerobot.scripts.train --help
+lerobot-train --help
 ```
 
 <p align="center">
@@ -97,7 +97,7 @@ Similarly for when recording an episode, it is recommended that you are logged i
 Once you are logged in, you can run inference in your setup by doing:
 
 ```bash
-python -m lerobot.record \
+lerobot-record \
   --robot.type=so101_follower \
   --robot.port=/dev/ttyACM0 \ # <- Use your port
   --robot.id=my_blue_follower_arm \ # <- Use your robot id

docs/source/so100.mdx

@@ -26,7 +26,7 @@ Unlike the SO-101, the motor connectors are not easily accessible once the arm i
 To find the port for each bus servo adapter, run this script:
 
 ```bash
-python -m lerobot.find_port
+lerobot-find-port
 ```
 
 <hfoptions id="example">
@@ -93,7 +93,7 @@ For a visual reference on how to set the motor ids please refer to [this video](
 <hfoption id="Command">
 
 ```bash
-python -m lerobot.setup_motors \
+lerobot-setup-motors \
     --robot.type=so100_follower \
     --robot.port=/dev/tty.usbmodem585A0076841  # <- paste here the port found at previous step
 ```
@@ -168,7 +168,7 @@ Do the same steps for the leader arm.
 <hfoptions id="setup_motors">
 <hfoption id="Command">
 ```bash
-python -m lerobot.setup_motors \
+lerobot-setup-motors \
     --teleop.type=so100_leader \
     --teleop.port=/dev/tty.usbmodem575E0031751  # <- paste here the port found at previous step
 ```
@@ -568,7 +568,7 @@ Run the following command or API example to calibrate the follower arm:
 <hfoption id="Command">
 
 ```bash
-python -m lerobot.calibrate \
+lerobot-calibrate \
     --robot.type=so100_follower \
     --robot.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --robot.id=my_awesome_follower_arm # <- Give the robot a unique name
@@ -606,7 +606,7 @@ Do the same steps to calibrate the leader arm, run the following command or API
 <hfoption id="Command">
 
 ```bash
-python -m lerobot.calibrate \
+lerobot-calibrate \
     --teleop.type=so100_leader \
     --teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --teleop.id=my_awesome_leader_arm # <- Give the robot a unique name

docs/source/so101.mdx

@@ -162,7 +162,7 @@ It is advisable to install one 3-pin cable in the motor after placing them befor
 To find the port for each bus servo adapter, connect MotorBus to your computer via USB and power. Run the following script and disconnect the MotorBus when prompted:
 
 ```bash
-python -m lerobot.find_port
+lerobot-find-port
 ```
 
 <hfoptions id="example">
@@ -240,7 +240,7 @@ Connect the usb cable from your computer and the power supply to the follower ar
 <hfoption id="Command">
 
 ```bash
-python -m lerobot.setup_motors \
+lerobot-setup-motors \
     --robot.type=so101_follower \
     --robot.port=/dev/tty.usbmodem585A0076841  # <- paste here the port found at previous step
 ```
@@ -316,7 +316,7 @@ Do the same steps for the leader arm.
 <hfoption id="Command">
 
 ```bash
-python -m lerobot.setup_motors \
+lerobot-setup-motors \
     --teleop.type=so101_leader \
     --teleop.port=/dev/tty.usbmodem575E0031751  # <- paste here the port found at previous step
 ```
@@ -353,7 +353,7 @@ Run the following command or API example to calibrate the follower arm:
 <hfoption id="Command">
 
 ```bash
-python -m lerobot.calibrate \
+lerobot-calibrate \
     --robot.type=so101_follower \
     --robot.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --robot.id=my_awesome_follower_arm # <- Give the robot a unique name
@@ -402,7 +402,7 @@ Do the same steps to calibrate the leader arm, run the following command or API
 <hfoption id="Command">
 
 ```bash
-python -m lerobot.calibrate \
+lerobot-calibrate \
     --teleop.type=so101_leader \
     --teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --teleop.id=my_awesome_leader_arm # <- Give the robot a unique name

examples/1_load_lerobot_dataset.py

@@ -92,11 +92,11 @@ print(dataset.hf_dataset)
 # LeRobot datasets also subclasses PyTorch datasets so you can do everything you know and love from working
 # with the latter, like iterating through the dataset.
 # The __getitem__ iterates over the frames of the dataset. Since our datasets are also structured by
-# episodes, you can access the frame indices of any episode using the episode_data_index. Here, we access
+# episodes, you can access the frame indices of any episode using dataset.meta.episodes. Here, we access
 # frame indices associated to the first episode:
 episode_index = 0
-from_idx = dataset.episode_data_index["from"][episode_index].item()
-to_idx = dataset.episode_data_index["to"][episode_index].item()
+from_idx = dataset.meta.episodes["dataset_from_index"][episode_index]
+to_idx = dataset.meta.episodes["dataset_to_index"][episode_index]
 
 # Then we grab all the image frames from the first camera:
 camera_key = dataset.meta.camera_keys[0]

examples/4_train_policy_with_script.md

@@ -62,7 +62,7 @@ By default, every field takes its default value specified in the dataclass. If a
 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
-python -m lerobot.scripts.train \
+lerobot-train \
     --dataset.repo_id=lerobot/pusht \
     --policy.type=diffusion \
     --env.type=pusht
@@ -77,7 +77,7 @@ Let's break this down:
 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
-python -m lerobot.scripts.train \
+lerobot-train \
     --policy.type=act \
     --dataset.repo_id=lerobot/aloha_sim_insertion_human \
     --env.type=aloha \
@@ -90,7 +90,7 @@ We now want to train a different policy for aloha on another task. We'll change
 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
-python -m lerobot.scripts.train \
+lerobot-train \
     --policy.type=act \
     --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
     --env.type=aloha \
@@ -127,7 +127,7 @@ Now, let's assume that we want to reproduce the run just above. That run has pro
 We can then simply load the config values from this file using:
 
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
     --config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
     --output_dir=outputs/train/act_aloha_transfer_2
 ```
@@ -137,7 +137,7 @@ python -m lerobot.scripts.train \
 Similarly to Hydra, we can still override some parameters in the CLI if we want to, e.g.:
 
 ```bash
-python -m lerobot.scripts.train \
+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
@@ -148,7 +148,7 @@ python -m lerobot.scripts.train \
 `--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
-python -m lerobot.scripts.train --config_path=lerobot/diffusion_pusht
+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)
@@ -160,7 +160,7 @@ Being able to resume a training run is important in case it crashed or aborted f
 Let's reuse the command from the previous run and add a few more options:
 
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
     --policy.type=act \
     --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
     --env.type=aloha \
@@ -179,7 +179,7 @@ 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
-python -m lerobot.scripts.train \
+lerobot-train \
     --config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
     --resume=true
 ```
@@ -190,7 +190,7 @@ Another reason for which you might want to resume a run is simply to extend trai
 You could double the number of steps of the previous run with:
 
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
     --config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
     --resume=true \
     --steps=200000
@@ -224,7 +224,7 @@ In addition to the features currently in Draccus, we've added a special `.path`
 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
-python -m lerobot.scripts.train \
+lerobot-train \
     --policy.path=lerobot/act_aloha_sim_transfer_cube_human \
     --dataset.repo_id=lerobot/aloha_sim_insertion_human \
     --env.type=aloha \
@@ -270,7 +270,7 @@ We'll summarize here the main use cases to remember from this tutorial.
 #### Train a policy from scratch – CLI
 
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
     --policy.type=act \  # <- select 'act' policy
     --env.type=pusht \  # <- select 'pusht' environment
     --dataset.repo_id=lerobot/pusht  # <- train on this dataset
@@ -279,7 +279,7 @@ python -m lerobot.scripts.train \
 #### Train a policy from scratch - config file + CLI
 
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
     --config_path=path/to/pretrained_model \  # <- can also be a repo_id
     --policy.n_action_steps=80  # <- you may still override values
 ```
@@ -287,7 +287,7 @@ python -m lerobot.scripts.train \
 #### Resume/continue a training run
 
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
     --config_path=checkpoint/pretrained_model/ \
     --resume=true \
     --steps=200000  # <- you can change some training parameters
@@ -296,7 +296,7 @@ python -m lerobot.scripts.train \
 #### Fine-tuning
 
 ```bash
-python -m lerobot.scripts.train \
+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 \

examples/backward_compatibility/replay.py

@@ -18,7 +18,7 @@ Replays the actions of an episode from a dataset on a robot.
 Example:
 
 ```shell
-python -m lerobot.replay \
+lerobot-replay \
     --robot.type=so100_follower \
     --robot.port=/dev/tty.usbmodem58760431541 \
     --robot.id=black \

examples/lekiwi/evaluate.py

@@ -1,7 +1,6 @@
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.policies.act.modeling_act import ACTPolicy
-from lerobot.policies.factory import make_processor
 from lerobot.record import record_loop
 from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
 from lerobot.utils.control_utils import init_keyboard_listener
@@ -12,14 +11,12 @@ NUM_EPISODES = 2
 FPS = 30
 EPISODE_TIME_SEC = 60
 TASK_DESCRIPTION = "My task description"
-HF_MODEL_ID = "<hf_username>/<model_repo_id>"
-HF_DATASET_ID = "<hf_username>/<eval_dataset_repo_id>"
 
 # Create the robot and teleoperator configurations
 robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
 robot = LeKiwiClient(robot_config)
 
-policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
+policy = ACTPolicy.from_pretrained("<hf_username>/<policy_repo_id>")
 
 # Configure the dataset features
 action_features = hw_to_dataset_features(robot.action_features, "action")
@@ -28,7 +25,7 @@ dataset_features = {**action_features, **obs_features}
 
 # Create the dataset
 dataset = LeRobotDataset.create(
-    repo_id=HF_DATASET_ID,
+    repo_id="<hf_username>/<eval_dataset_repo_id>",
     fps=FPS,
     features=dataset_features,
     robot_type=robot.name,
@@ -46,12 +43,6 @@ listener, events = init_keyboard_listener()
 if not robot.is_connected:
     raise ValueError("Robot is not connected!")
 
-preprocessor, postprocessor = make_processor(
-    policy_cfg=policy,
-    pretrained_path=HF_MODEL_ID,
-    dataset_stats=dataset.meta.stats,
-)
-
 recorded_episodes = 0
 while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
     log_say(f"Running inference, recording eval episode {recorded_episodes} of {NUM_EPISODES}")
@@ -62,8 +53,6 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
         events=events,
         fps=FPS,
         policy=policy,
-        preprocessor=preprocessor,
-        postprocessor=postprocessor,
         dataset=dataset,
         control_time_s=EPISODE_TIME_SEC,
         single_task=TASK_DESCRIPTION,

examples/lekiwi/teleoperate.py

@@ -38,7 +38,7 @@ while True:
     keyboard_keys = keyboard.get_action()
     base_action = robot._from_keyboard_to_base_action(keyboard_keys)
 
-    log_rerun_data(observation=observation, action={**arm_action, **base_action})
+    log_rerun_data(observation, {**arm_action, **base_action})
 
     action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
 

examples/phone_so100_eval.py

@@ -1,158 +0,0 @@
-# !/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features
-from lerobot.datasets.utils import merge_features
-from lerobot.model.kinematics import RobotKinematics
-from lerobot.policies.act.modeling_act import ACTPolicy
-from lerobot.policies.factory import make_processor
-from lerobot.processor.converters import (
-    to_output_robot_action,
-    to_transition_robot_observation,
-)
-from lerobot.processor.pipeline import RobotProcessor
-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.utils.control_utils import init_keyboard_listener
-from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import _init_rerun
-
-NUM_EPISODES = 5
-FPS = 30
-EPISODE_TIME_SEC = 60
-TASK_DESCRIPTION = "My task description"
-HF_MODEL_ID = "<hf_username>/<model_repo_id>"
-HF_DATASET_ID = "<hf_username>/<dataset_repo_id>"
-
-# Initialize the robot with degrees
-camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem58760434471",
-    id="my_awesome_follower_arm",
-    cameras=camera_config,
-    use_degrees=True,
-)
-
-# Initialize the robot
-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="./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 pose action to joint action
-robot_ee_to_joints = RobotProcessor(
-    steps=[
-        AddRobotObservationAsComplimentaryData(robot=robot),
-        InverseKinematicsEEToJoints(
-            kinematics=kinematics_solver,
-            motor_names=list(robot.bus.motors.keys()),
-            initial_guess_current_joints=True,
-        ),
-    ],
-    to_transition=lambda tr: tr,
-    to_output=to_output_robot_action,
-)
-
-# Build pipeline to convert joint observation to ee pose observation
-robot_joints_to_ee_pose = RobotProcessor(
-    steps=[
-        ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
-    ],
-    to_transition=to_transition_robot_observation,
-    to_output=lambda tr: tr,
-)
-
-# Build dataset action and gripper features
-action_ee_and_gripper = aggregate_pipeline_dataset_features(
-    pipeline=robot_ee_to_joints,
-    initial_features={},
-    use_videos=True,
-    patterns=["action.ee", "action.gripper.pos", "observation.state.gripper.pos"],
-)  # Get all ee action features + gripper pos action features
-
-# Build dataset observation features
-obs_ee = aggregate_pipeline_dataset_features(
-    pipeline=robot_joints_to_ee_pose,
-    initial_features=robot.observation_features,
-    use_videos=True,
-    patterns=["observation.state.ee"],
-)  # Get all ee observation features
-
-dataset_features = merge_features(obs_ee, action_ee_and_gripper)
-
-print("All dataset features: ", dataset_features)
-
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id=HF_DATASET_ID,
-    fps=FPS,
-    features=dataset_features,
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
-
-# Initialize the keyboard listener and rerun visualization
-_, events = init_keyboard_listener()
-_init_rerun(session_name="recording_phone")
-
-# Connect the robot and teleoperator
-robot.connect()
-
-episode_idx = 0
-
-policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
-preprocessor, postprocessor = make_processor(
-    policy_cfg=policy,
-    pretrained_path=HF_MODEL_ID,
-    dataset_stats=dataset.meta.stats,
-)
-
-for episode_idx in range(NUM_EPISODES):
-    log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
-
-    record_loop(
-        robot=robot,
-        events=events,
-        fps=FPS,
-        policy=policy,
-        preprocessor=preprocessor,
-        postprocessor=postprocessor,
-        dataset=dataset,
-        control_time_s=EPISODE_TIME_SEC,
-        single_task=TASK_DESCRIPTION,
-        display_data=True,
-        robot_action_processor=robot_ee_to_joints,
-        robot_observation_processor=robot_joints_to_ee_pose,
-    )
-    dataset.save_episode()
-
-# Clean up
-log_say("Stop recording")
-robot.disconnect()
-dataset.push_to_hub()

examples/phone_so100_record.py

@@ -1,215 +0,0 @@
-# !/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features
-from lerobot.datasets.utils import merge_features
-from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor.converters import (
-    to_output_robot_action,
-    to_transition_robot_observation,
-    to_transition_teleop_action,
-)
-from lerobot.processor.pipeline import RobotProcessor
-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,
-    GripperVelocityToJoint,
-    InverseKinematicsEEToJoints,
-)
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
-from lerobot.teleoperators.phone.phone import Phone
-from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
-from lerobot.utils.control_utils import init_keyboard_listener
-from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import _init_rerun
-
-NUM_EPISODES = 10
-FPS = 30
-EPISODE_TIME_SEC = 60
-RESET_TIME_SEC = 30
-TASK_DESCRIPTION = "My task description"
-HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
-
-# Initialize the robot and teleoperator
-camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem58760434471",
-    id="my_awesome_follower_arm",
-    cameras=camera_config,
-    use_degrees=True,
-)
-teleop_config = PhoneConfig(phone_os=PhoneOS.IOS)  # or PhoneOS.ANDROID
-
-# Initialize the robot and teleoperator
-robot = SO100Follower(robot_config)
-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="./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
-phone_to_robot_ee_pose = RobotProcessor(
-    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.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,
-        ),
-    ],
-    to_transition=to_transition_teleop_action,
-    to_output=lambda tr: tr,
-)
-
-# Build pipeline to convert ee pose action to joint action
-robot_ee_to_joints = RobotProcessor(
-    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,
-)
-
-# Build pipeline to convert joint observation to ee pose observation
-robot_joints_to_ee_pose = RobotProcessor(
-    steps=[
-        ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
-    ],
-    to_transition=to_transition_robot_observation,
-    to_output=lambda tr: tr,
-)
-
-# Build dataset ee action features
-action_ee = aggregate_pipeline_dataset_features(
-    pipeline=phone_to_robot_ee_pose,
-    initial_features=phone.action_features,
-    use_videos=True,
-    patterns=["action.ee"],
-)
-
-# Get gripper pos action features
-gripper = aggregate_pipeline_dataset_features(
-    pipeline=robot_ee_to_joints,
-    initial_features={},
-    use_videos=True,
-    patterns=["action.gripper.pos", "observation.state.gripper.pos"],
-)
-
-# Build dataset ee observation features
-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)
-
-print("All dataset features: ", dataset_features)
-
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id=HF_REPO_ID,
-    fps=FPS,
-    features=dataset_features,
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
-
-# Initialize the keyboard listener and rerun visualization
-_, events = init_keyboard_listener()
-_init_rerun(session_name="recording_phone")
-
-# Connect the robot and teleoperator
-robot.connect()
-phone.connect()
-
-episode_idx = 0
-while episode_idx < NUM_EPISODES and not events["stop_recording"]:
-    log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
-
-    record_loop(
-        robot=robot,
-        events=events,
-        fps=FPS,
-        teleop=phone,
-        dataset=dataset,
-        control_time_s=EPISODE_TIME_SEC,
-        single_task=TASK_DESCRIPTION,
-        display_data=True,
-        teleop_action_processor=phone_to_robot_ee_pose,
-        robot_action_processor=robot_ee_to_joints,
-        robot_observation_processor=robot_joints_to_ee_pose,
-    )
-
-    # Reset the environment if not stopping or re-recording
-    if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
-        log_say("Reset the environment")
-        record_loop(
-            robot=robot,
-            events=events,
-            fps=FPS,
-            teleop=phone,
-            control_time_s=RESET_TIME_SEC,
-            single_task=TASK_DESCRIPTION,
-            display_data=True,
-            teleop_action_processor=phone_to_robot_ee_pose,
-            robot_action_processor=robot_ee_to_joints,
-            robot_observation_processor=robot_joints_to_ee_pose,
-        )
-
-    if events["rerecord_episode"]:
-        log_say("Re-recording episode")
-        events["rerecord_episode"] = False
-        events["exit_early"] = False
-        dataset.clear_episode_buffer()
-        continue
-
-    dataset.save_episode()
-    episode_idx += 1
-
-# Clean up
-log_say("Stop recording")
-robot.disconnect()
-phone.disconnect()
-dataset.push_to_hub()

examples/phone_so100_replay.py

@@ -1,106 +0,0 @@
-# !/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-import time
-
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor.converters import to_output_robot_action
-from lerobot.processor.pipeline import RobotProcessor
-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.robot_utils import busy_wait
-from lerobot.utils.utils import log_say
-
-EPISODE_IDX = 0
-HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
-
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", use_degrees=True
-)
-robot = SO100Follower(robot_config)
-robot.connect()
-
-dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
-actions = dataset.hf_dataset.select_columns("action")
-
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-kinematics_solver = RobotKinematics(
-    urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
-
-
-# This method converts the action from the dataset to a transition for pipeline
-def action_to_transition(action: dict):
-    act = {}
-
-    # EE pose
-    for k in ("ee.x", "ee.y", "ee.z", "ee.wx", "ee.wy", "ee.wz"):
-        if k in action:
-            act[f"action.{k}"] = float(action[k])
-
-    # Gripper: your dataset has absolute position
-    if "gripper.pos" in action:
-        act["action.gripper.pos"] = float(action["gripper.pos"])
-
-    return {
-        "observation": None,
-        "action": act,
-        "reward": None,
-        "done": False,
-        "truncated": False,
-        "info": {},
-        "complementary_data": {},
-    }
-
-
-# Build pipeline to convert ee pose action to joint action
-robot_ee_to_joints = RobotProcessor(
-    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=action_to_transition,
-    to_output=to_output_robot_action,
-)
-
-robot_ee_to_joints.reset()
-
-log_say(f"Replaying episode {EPISODE_IDX}")
-for idx in range(dataset.num_frames):
-    t0 = time.perf_counter()
-
-    ee_action = {
-        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
-    }
-
-    joint_action = robot_ee_to_joints(ee_action)
-    action_sent = robot.send_action(joint_action)
-
-    busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))
-
-robot.disconnect()

examples/phone_so100_teleop.py

@@ -1,109 +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 specif
-
-import time
-
-from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotProcessor
-from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action
-from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
-from lerobot.robots.so100_follower.robot_kinematic_processor import (
-    AddRobotObservationAsComplimentaryData,
-    EEBoundsAndSafety,
-    EEReferenceAndDelta,
-    GripperVelocityToJoint,
-    InverseKinematicsEEToJoints,
-)
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
-from lerobot.teleoperators.phone.phone import Phone
-from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
-
-# Initialize the robot and teleoperator
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", use_degrees=True
-)
-teleop_config = PhoneConfig(phone_os=PhoneOS.IOS)  # or PhoneOS.ANDROID
-
-# Initialize the robot and teleoperator
-robot = SO100Follower(robot_config)
-teleop_device = Phone(teleop_config)
-
-# 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="./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
-phone_to_robot_ee_pose = RobotProcessor(
-    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.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
-            max_ee_step_m=0.10,
-            max_ee_twist_step_rad=0.50,
-        ),
-    ],
-    to_transition=to_transition_teleop_action,
-    to_output=lambda tr: tr,
-)
-
-# Build pipeline to convert ee pose action to joint action
-robot_ee_to_joints = RobotProcessor(
-    steps=[
-        InverseKinematicsEEToJoints(
-            kinematics=kinematics_solver,
-            motor_names=list(robot.bus.motors.keys()),
-        ),
-        GripperVelocityToJoint(
-            motor_names=list(robot.bus.motors.keys()),
-            speed_factor=20.0,
-        ),
-    ],
-    to_transition=lambda tr: tr,
-    to_output=to_output_robot_action,
-)
-
-robot.connect()
-teleop_device.connect()
-
-print("Starting teleop loop. Move your phone to teleoperate the robot.")
-while True:
-    phone_obs = teleop_device.get_action()
-    if not phone_obs:
-        time.sleep(0.01)
-        continue
-
-    # Get teleop observation
-    phone_obs = teleop_device.get_action()
-
-    # Phone to EE pose transition
-    ee_transition = phone_to_robot_ee_pose(phone_obs)
-
-    # EE pose to Joints transition
-    joint_action = robot_ee_to_joints(ee_transition)
-
-    if joint_action:
-        robot.send_action(joint_action)
-
-    time.sleep(0.01)

examples/port_datasets/agibot_hdf5/port_agibot.py

@@ -0,0 +1,503 @@
+import json
+import logging
+import shutil
+import time
+from pathlib import Path
+
+import h5py
+import numpy as np
+import pandas as pd
+
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import (
+    DEFAULT_CHUNK_SIZE,
+    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
+    DEFAULT_VIDEO_PATH,
+    EPISODES_DIR,
+    get_video_duration_in_s,
+    get_video_size_in_mb,
+    update_chunk_file_indices,
+    write_info,
+)
+from lerobot.datasets.video_utils import concat_video_files
+from lerobot.utils.utils import get_elapsed_time_in_days_hours_minutes_seconds
+
+AGIBOT_FPS = 30
+AGIBOT_ROBOT_TYPE = "AgiBot_A2D"
+AGIBOT_FEATURES = {
+    # gripper open range in mm (0 for pull open, 1 for full close)
+    "observation.state.effector.position": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["left_gripper", "right_gripper"],
+        },
+    },
+    # flange xyz in meters
+    "observation.state.end.position": {
+        "dtype": "float32",
+        "shape": (6,),
+        "names": {
+            "axes": ["left_x", "left_y", "left_z", "right_x", "right_y", "right_z"],
+        },
+    },
+    # flange quaternion with xyzw
+    "observation.state.end.orientation": {
+        "dtype": "float32",
+        "shape": (8,),
+        "names": {
+            "axes": ["left_x", "left_y", "left_z", "left_w", "right_x", "right_y", "right_z", "right_w"],
+        },
+    },
+    # in radians
+    "observation.state.head.position": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["yaw", "pitch"],
+        },
+    },
+    # in motor steps
+    "observation.state.joint.current_value": {
+        "dtype": "float32",
+        "shape": (14,),
+        "names": {
+            "axes": [f"left_joint_{i}" for i in range(7)] + [f"right_joint_{i}" for i in range(7)],
+        },
+    },
+    # same as current_value but in radians
+    "observation.state.joint.position": {
+        "dtype": "float32",
+        "shape": (14,),
+        "names": {
+            "axes": [f"left_joint_{i}" for i in range(7)] + [f"right_joint_{i}" for i in range(7)],
+        },
+    },
+    # pitch in radians, lift in meters
+    "observation.state.waist.position": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["pitch", "lift"],
+        },
+    },
+    # concatenation of head.position, joint.position, effector.position, waist.position
+    "observation.state": {
+        "dtype": "float32",
+        "shape": (20,),
+        "names": {
+            "axes": ["head_yaw", "head_pitch"]
+            + [f"left_joint_{i}" for i in range(7)]
+            + ["left_gripper"]
+            + [f"right_joint_{i}" for i in range(7)]
+            + ["right_gripper"]
+            + ["waist_pitch", "waist_lift"],
+        },
+    },
+    # gripper open range in mm (0 for pull open, 1 for full close)
+    "action.effector.position": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["left_gripper", "right_gripper"],
+        },
+    },
+    # flange xyz in meters
+    "action.end.position": {
+        "dtype": "float32",
+        "shape": (6,),
+        "names": {
+            "axes": ["left_x", "left_y", "left_z", "right_x", "right_y", "right_z"],
+        },
+    },
+    # flange quaternion with xyzw
+    "action.end.orientation": {
+        "dtype": "float32",
+        "shape": (8,),
+        "names": {
+            "axes": ["left_x", "left_y", "left_z", "left_w", "right_x", "right_y", "right_z", "right_w"],
+        },
+    },
+    # in radians
+    "action.head.position": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["yaw", "pitch"],
+        },
+    },
+    # goal joint position in radians
+    "action.joint.position": {
+        "dtype": "float32",
+        "shape": (14,),
+        "names": {
+            "axes": [f"left_joint_{i}" for i in range(7)] + [f"right_joint_{i}" for i in range(7)],
+        },
+    },
+    "action.robot.velocity": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["velocity_x", "yaw_rate"],
+        },
+    },
+    # pitch in radians, lift in meters
+    "action.waist.position": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["pitch", "lift"],
+        },
+    },
+    # concatenation of head.position, joint.position, effector.position, waist.position, robot.velocity
+    "action": {
+        "dtype": "float32",
+        "shape": (22,),
+        "names": {
+            "axes": ["head_yaw", "head_pitch"]
+            + [f"left_joint_{i}" for i in range(7)]
+            + ["left_gripper"]
+            + [f"right_joint_{i}" for i in range(7)]
+            + ["right_gripper"]
+            + ["waist_pitch", "waist_lift"]
+            + ["velocity_x", "yaw_rate"],
+        },
+    },
+    # episode level annotation
+    "init_scene_text": {
+        "dtype": "string",
+        "shape": (1,),
+        "names": None,
+    },
+    # frame level annotation
+    "action_text": {
+        "dtype": "string",
+        "shape": (1,),
+        "names": None,
+    },
+    # frame level annotation
+    "skill": {
+        "dtype": "string",
+        "shape": (1,),
+        "names": None,
+    },
+}
+
+AGIBOT_IMAGES_FEATURES = {
+    "observation.images.top_head": {
+        "dtype": "video",
+        "shape": (480, 640, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.hand_left": {
+        "dtype": "video",
+        "shape": (480, 640, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.hand_right": {
+        "dtype": "video",
+        "shape": (480, 640, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.head_center_fisheye": {
+        "dtype": "video",
+        "shape": (748, 960, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.head_left_fisheye": {
+        "dtype": "video",
+        "shape": (748, 960, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.head_right_fisheye": {
+        "dtype": "video",
+        "shape": (748, 960, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.back_left_fisheye": {
+        "dtype": "video",
+        "shape": (748, 960, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.back_right_fisheye": {
+        "dtype": "video",
+        "shape": (748, 960, 3),
+        "names": ["height", "width", "channel"],
+    },
+}
+
+
+def load_info_per_task(raw_dir):
+    info_per_task = {}
+    task_info_dir = raw_dir / "task_info"
+    for path in task_info_dir.glob("task_*.json"):
+        task_index = int(path.name.replace("task_", "").replace(".json", ""))
+        with open(path) as f:
+            task_info = json.load(f)
+
+        task_info = {ep["episode_id"]: ep for ep in task_info}
+        info_per_task[task_index] = task_info
+
+    return info_per_task
+
+
+def create_frame_idx_to_frames_label_idx(ep_info):
+    frame_idx_to_frames_label_idx = {}
+    for label_idx, frames_label in enumerate(ep_info["label_info"]["action_config"]):
+        for frame_idx in range(frames_label["start_frame"], frames_label["end_frame"]):
+            frame_idx_to_frames_label_idx[frame_idx] = label_idx
+    return frame_idx_to_frames_label_idx
+
+
+def generate_lerobot_frames(raw_dir: Path, task_index: int, episode_index: int):
+    r"""/!\ The frames dont contain observation.cameras.*"""
+    info_per_task = load_info_per_task(raw_dir)
+    ep_info = info_per_task[task_index][episode_index]
+    frame_idx_to_frames_label_idx = create_frame_idx_to_frames_label_idx(ep_info)
+
+    # Empty features are commented out.
+    keys_mapping = {
+        # STATE
+        # "observation.state.effector.force": "state/effector/force",
+        "observation.state.effector.position": "state/effector/position",
+        # "observation.state.end.angular": "state/end/angular",
+        "observation.state.end.position": "state/end/position",
+        "observation.state.end.orientation": "state/end/orientation",
+        # "observation.state.end.velocity": "state/end/velocity",
+        # "observation.state.end.wrench": "state/end/wrench",
+        # "observation.state.head.effort": "state/head/effort",
+        "observation.state.head.position": "state/head/position",
+        # "observation.state.head.velocity": "state/head/velocity",
+        "observation.state.joint.current_value": "state/joint/current_value",
+        # "observation.state.joint.effort": "state/joint/effort",
+        "observation.state.joint.position": "state/joint/position",
+        # "observation.state.joint.velocity": "state/joint/velocity",
+        # "observation.state.robot.orientation": "state/robot/orientation",
+        # "observation.state.robot.orientation_drift": "state/robot/orientation_drift",
+        # "observation.state.robot.position": "state/robot/position",
+        # "observation.state.robot.position_drift": "state/robot/position_drift",
+        # "observation.state.waist.effort": "state/waist/effort",
+        "observation.state.waist.position": "state/waist/position",
+        # "observation.state.waist.velocity": "state/waist/velocity",
+        # ----- ACTION (index are also commented out) -----
+        # "action.effector.index": "action/effector/index",
+        "action.effector.position": "action/effector/position",
+        # "action.effector.force": "action/effector/force",
+        # "action.end.index": "action/end/index",
+        "action.end.position": "action/end/position",
+        "action.end.orientation": "action/end/orientation",
+        # "action.head.index": "action/head/index",
+        "action.head.position": "action/head/position",
+        # "action.joint.index": "action/joint/index",
+        "action.joint.position": "action/joint/position",
+        # "action.joint.effort": "action/joint/effort",
+        # "action.joint.velocity": "action/joint/velocity",
+        # "action.robot.index": "action/robot/index",
+        # "action.robot.position": "action/robot/position",
+        # "action.robot.orientation": "action/robot/orientation",
+        # "action.robot.angular": "action/robot/angular",
+        "action.robot.velocity": "action/robot/velocity",
+        # "action.waist.index": "action/waist/index",
+        "action.waist.position": "action/waist/position",
+    }
+
+    h5_path = raw_dir / f"proprio_stats/{task_index}/{episode_index}/proprio_stats.h5"
+    with h5py.File(h5_path) as h5:
+        num_frames = len(h5["state/joint/position"])
+
+        for h5_key in keys_mapping.values():
+            col_num_frames = h5[h5_key].shape[0]
+            if col_num_frames != num_frames:
+                raise ValueError(
+                    f"HDF5 column '{h5_key}' is expected to have {num_frames} but has {col_num_frames}' frames instead."
+                )
+
+        for i in range(num_frames):
+            # Create frame
+            f = {new_key: h5[h5_key][i] for new_key, h5_key in keys_mapping.items()}
+
+            for key in f:
+                f[key] = np.array(f[key]).astype(np.float32)
+
+            f["observation.state.end.position"] = f["observation.state.end.position"].reshape(6)
+            f["observation.state.end.orientation"] = f["observation.state.end.orientation"].reshape(8)
+            f["observation.state"] = np.concatenate(
+                [
+                    f["observation.state.head.position"],
+                    f["observation.state.joint.position"][:7],  # left
+                    f["observation.state.effector.position"][[0]],  # left
+                    f["observation.state.joint.position"][7:],  # right
+                    f["observation.state.effector.position"][[1]],  # right
+                    f["observation.state.waist.position"],
+                ]
+            )
+
+            f["action.end.position"] = f["action.end.position"].reshape(6)
+            f["action.end.orientation"] = f["action.end.orientation"].reshape(8)
+            f["action"] = np.concatenate(
+                [
+                    f["action.head.position"],
+                    f["action.joint.position"][:7],  # left
+                    f["action.effector.position"][[0]],  # left
+                    f["action.joint.position"][7:],  # right
+                    f["action.effector.position"][[1]],  # right
+                    f["action.waist.position"],
+                    f["action.robot.velocity"],
+                ]
+            )
+
+            # episode level annotation
+            f["task"] = ep_info["task_name"]
+            f["init_scene_text"] = ep_info["init_scene_text"]
+
+            # frame level annotation
+            if i in frame_idx_to_frames_label_idx:
+                frames_label_idx = frame_idx_to_frames_label_idx[i]
+                frames_label = ep_info["label_info"]["action_config"][frames_label_idx]
+                f["action_text"] = frames_label["action_text"]
+                f["skill"] = frames_label["skill"]
+            else:
+                f["action_text"] = ""
+                f["skill"] = ""
+
+            yield f
+
+
+def update_meta_data(
+    df,
+    ep_to_meta,
+):
+    def _update(row):
+        ep_idx = row["episode_index"]
+        for key, meta in ep_to_meta[ep_idx].items():
+            row[f"videos/{key}/chunk_index"] = meta["chunk_index"]
+            row[f"videos/{key}/file_index"] = meta["file_index"]
+            row[f"videos/{key}/from_timestamp"] = meta["from_timestamp"]
+            row[f"videos/{key}/to_timestamp"] = meta["to_timestamp"]
+        return row
+
+    return df.apply(_update, axis=1)
+
+
+def move_videos_to_lerobot_directory(lerobot_dataset, raw_dir, task_index, episode_names):
+    keys_mapping = {
+        "observation.images.top_head": "head_color",
+        "observation.images.hand_left": "hand_left_color",
+        "observation.images.hand_right": "hand_right_color",
+        "observation.images.head_center_fisheye": "head_center_fisheye_color",
+        "observation.images.head_left_fisheye": "head_left_fisheye_color",
+        "observation.images.head_right_fisheye": "head_right_fisheye_color",
+        "observation.images.back_left_fisheye": "back_left_fisheye_color",
+        "observation.images.back_right_fisheye": "back_right_fisheye_color",
+    }
+
+    # sanity check
+    for key in keys_mapping:
+        if key not in lerobot_dataset.meta.info["features"]:
+            raise ValueError(f"Key '{key}' not found in features.")
+
+    video_keys = keys_mapping.keys()
+    chunk_idx = dict.fromkeys(video_keys, 0)
+    file_idx = dict.fromkeys(video_keys, 0)
+    latest_duration_in_s = dict.fromkeys(video_keys, 0)
+    ep_to_meta = {}
+    for ep_idx, ep_name in enumerate(episode_names):
+        for key in video_keys:
+            raw_videos_dir = raw_dir / f"observations/{task_index}/{ep_name}/videos"
+            old_key = keys_mapping[key]
+            ep_path = raw_videos_dir / f"{old_key}.mp4"
+            ep_duration_in_s = get_video_duration_in_s(ep_path)
+
+            aggr_path = lerobot_dataset.root / DEFAULT_VIDEO_PATH.format(
+                video_key=key,
+                chunk_index=chunk_idx[key],
+                file_index=file_idx[key],
+            )
+            if not aggr_path.exists():
+                # First video
+                aggr_path.parent.mkdir(parents=True, exist_ok=True)
+                shutil.copy(str(ep_path), str(aggr_path))
+            else:
+                size_in_mb = get_video_size_in_mb(ep_path)
+                aggr_size_in_mb = get_video_size_in_mb(aggr_path)
+
+                if aggr_size_in_mb + size_in_mb >= DEFAULT_VIDEO_FILE_SIZE_IN_MB:
+                    # Size limit is reached, prepare new parquet file
+                    chunk_idx[key], file_idx[key] = update_chunk_file_indices(
+                        chunk_idx[key], file_idx[key], DEFAULT_CHUNK_SIZE
+                    )
+                    aggr_path = lerobot_dataset.root / DEFAULT_VIDEO_PATH.format(
+                        video_key=key,
+                        chunk_index=chunk_idx[key],
+                        file_index=file_idx[key],
+                    )
+                    aggr_path.parent.mkdir(parents=True, exist_ok=True)
+                    shutil.copy(str(ep_path), str(aggr_path))
+                    latest_duration_in_s[key] = 0
+                else:
+                    # Update the existing parquet file with new rows
+                    concat_video_files(
+                        [aggr_path, ep_path],
+                        lerobot_dataset.root,
+                        key,
+                        chunk_idx[key],
+                        file_idx[key],
+                    )
+
+            if ep_idx not in ep_to_meta:
+                ep_to_meta[ep_idx] = {}
+            ep_to_meta[ep_idx][key] = {
+                "chunk_index": chunk_idx[key],
+                "file_index": file_idx[key],
+                "from_timestamp": latest_duration_in_s[key],
+                "to_timestamp": latest_duration_in_s[key] + ep_duration_in_s,
+            }
+            latest_duration_in_s[key] += ep_duration_in_s
+
+    # Update episodes meta data
+    for meta_path in (lerobot_dataset.root / EPISODES_DIR).glob("chunk-*/file-*.parquet"):
+        df = pd.read_parquet(meta_path)
+        df = update_meta_data(df, ep_to_meta)
+        df.to_parquet(meta_path)
+
+
+def port_agibot(
+    raw_dir: Path, repo_id: str, task_index: int, episode_indices: list[int], push_to_hub: bool = False
+):
+    lerobot_dataset = LeRobotDataset.create(
+        repo_id=repo_id,
+        robot_type=AGIBOT_ROBOT_TYPE,
+        fps=AGIBOT_FPS,
+        features=AGIBOT_FEATURES,
+    )
+
+    start_time = time.time()
+    num_episodes = len(episode_indices)
+    logging.info(f"Number of episodes {num_episodes}")
+
+    for i, episode_index in enumerate(episode_indices):
+        elapsed_time = time.time() - start_time
+        d, h, m, s = get_elapsed_time_in_days_hours_minutes_seconds(elapsed_time)
+
+        logging.info(
+            f"{i} / {num_episodes} episodes processed (after {d} days, {h} hours, {m} minutes, {s:.3f} seconds)"
+        )
+
+        for frame in generate_lerobot_frames(raw_dir, task_index, episode_index):
+            lerobot_dataset.add_frame(frame)
+
+        lerobot_dataset.save_episode()
+        logging.info("Save_episode")
+
+    # Videos have already been encoded with the proper format, so we rely on hacks
+    # HACK: Add extra images features
+    lerobot_dataset.meta.info["features"].update(AGIBOT_IMAGES_FEATURES)
+    write_info(lerobot_dataset.meta.info, lerobot_dataset.meta.root)
+    move_videos_to_lerobot_directory(lerobot_dataset, raw_dir, task_index, episode_indices)
+
+    if push_to_hub:
+        lerobot_dataset.push_to_hub(
+            # Add agibot tag, since it belongs to the agibot collection of datasets
+            tags=["agibot"],
+            private=False,
+        )

examples/port_datasets/agibot_hdf5/slurm_port_shards.py

@@ -0,0 +1,198 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import logging
+import tarfile
+from pathlib import Path
+
+from datatrove.executor import LocalPipelineExecutor
+from datatrove.executor.slurm import SlurmPipelineExecutor
+from datatrove.pipeline.base import PipelineStep
+from port_datasets.agibot_hdf5.download import (
+    RAW_REPO_ID,
+    download_meta_data,
+    get_observations_files,
+)
+
+
+class PortAgiBotShards(PipelineStep):
+    def __init__(
+        self,
+        raw_dir: Path | str,
+        repo_id: str = None,
+    ):
+        super().__init__()
+        self.raw_dir = Path(raw_dir)
+        self.repo_id = repo_id
+
+    def run(self, data=None, rank: int = 0, world_size: int = 1):
+        import shutil
+
+        from datasets.utils.tqdm import disable_progress_bars
+        from port_datasets.agibot_hdf5.download import (
+            RAW_REPO_ID,
+            download,
+            get_observations_files,
+            no_depth,
+        )
+        from port_datasets.agibot_hdf5.port_agibot import port_agibot
+        from port_datasets.droid_rlds.port_droid import validate_dataset
+
+        from lerobot.constants import HF_LEROBOT_HOME
+        from lerobot.utils.utils import init_logging
+
+        init_logging()
+        disable_progress_bars()
+
+        shard_repo_id = f"{self.repo_id}_world_{world_size}_rank_{rank}"
+
+        dataset_dir = HF_LEROBOT_HOME / shard_repo_id
+        if dataset_dir.exists():
+            shutil.rmtree(dataset_dir)
+
+        obs_files, _ = get_observations_files(self.raw_dir, RAW_REPO_ID)
+        obs_file = obs_files[rank]
+
+        # Download subset
+        download(self.raw_dir, allow_patterns=obs_file)
+
+        tar_path = self.raw_dir / obs_file
+        with tarfile.open(tar_path, "r") as tar:
+            extracted_files = tar.getnames()
+
+        task_index = int(tar_path.parent.name)
+        episode_names = [int(p) for p in extracted_files if "/" not in p]
+
+        # Untar if needed
+        if not all((tar_path.parent / f"{ep_name}").exists() for ep_name in episode_names):
+            logging.info(f"Untar-ing {tar_path}...")
+            with tarfile.open(tar_path, "r") as tar:
+                tar.extractall(path=tar_path.parent, filter=no_depth)  # nosec B202
+
+        port_agibot(self.raw_dir, shard_repo_id, task_index, episode_names, push_to_hub=False)
+
+        for ep_name in episode_names:
+            shutil.rmtree(str(tar_path.parent / f"{ep_name}"))
+
+        tar_path.unlink()
+
+        validate_dataset(shard_repo_id)
+
+
+def make_port_executor(
+    raw_dir, repo_id, job_name, logs_dir, workers, partition, cpus_per_task, mem_per_cpu, slurm=True
+):
+    download_meta_data(raw_dir)
+    obs_files, _ = get_observations_files(raw_dir, RAW_REPO_ID)
+    num_shards = len(obs_files)
+
+    kwargs = {
+        "pipeline": [
+            PortAgiBotShards(raw_dir, repo_id),
+        ],
+        "logging_dir": str(logs_dir / job_name),
+    }
+
+    if slurm:
+        kwargs.update(
+            {
+                "job_name": job_name,
+                "tasks": num_shards,
+                "workers": workers,
+                "time": "08:00:00",
+                "partition": partition,
+                "cpus_per_task": cpus_per_task,
+                "sbatch_args": {"mem-per-cpu": mem_per_cpu},
+            }
+        )
+        executor = SlurmPipelineExecutor(**kwargs)
+    else:
+        kwargs.update(
+            {
+                "tasks": num_shards,
+                "workers": 1,
+            }
+        )
+        executor = LocalPipelineExecutor(**kwargs)
+
+    return executor
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--raw-dir",
+        type=Path,
+        required=True,
+        help="Directory containing input raw datasets (e.g. `path/to/dataset` or `path/to/dataset/version).",
+    )
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        help="Repositery identifier on Hugging Face: a community or a user name `/` the name of the dataset, required when push-to-hub is True.",
+    )
+    parser.add_argument(
+        "--logs-dir",
+        type=Path,
+        help="Path to logs directory for `datatrove`.",
+    )
+    parser.add_argument(
+        "--job-name",
+        type=str,
+        default="port_droid",
+        help="Job name used in slurm, and name of the directory created inside the provided logs directory.",
+    )
+    parser.add_argument(
+        "--slurm",
+        type=int,
+        default=1,
+        help="Launch over slurm. Use `--slurm 0` to launch sequentially (useful to debug).",
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=2048,
+        help="Number of slurm workers. It should be less than the maximum number of shards.",
+    )
+    parser.add_argument(
+        "--partition",
+        type=str,
+        help="Slurm partition. Ideally a CPU partition. No need for GPU partition.",
+    )
+    parser.add_argument(
+        "--cpus-per-task",
+        type=int,
+        default=8,
+        help="Number of cpus that each slurm worker will use.",
+    )
+    parser.add_argument(
+        "--mem-per-cpu",
+        type=str,
+        default="1950M",
+        help="Memory per cpu that each worker will use.",
+    )
+
+    args = parser.parse_args()
+    kwargs = vars(args)
+    kwargs["slurm"] = kwargs.pop("slurm") == 1
+    port_executor = make_port_executor(**kwargs)
+    port_executor.run()
+
+
+if __name__ == "__main__":
+    main()

examples/port_datasets/droid_rlds/display_error_files.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import json
+from pathlib import Path
+
+
+def find_missing_workers(completions_dir, world_size):
+    """Find workers that are not completed and returns their indices."""
+    full = list(range(world_size))
+
+    completed = []
+    for path in completions_dir.glob("*"):
+        if path.name in [".", ".."]:
+            continue
+        index = path.name.lstrip("0")
+        index = 0 if index == "" else int(index)
+        completed.append(index)
+
+    missing_workers = set(full) - set(completed)
+    return missing_workers
+
+
+def find_output_files(slurm_dir, worker_indices):
+    """Find output files associated to worker indices, and return tuples
+    of (worker index, output file path)
+    """
+    out_files = []
+    for path in slurm_dir.glob("*.out"):
+        _, worker_id = path.name.replace(".out", "").split("_")
+        worker_id = int(worker_id)
+        if worker_id in worker_indices:
+            out_files.append((worker_id, path))
+    return out_files
+
+
+def display_error_files(logs_dir, job_name):
+    executor_path = Path(logs_dir) / job_name / "executor.json"
+    completions_dir = Path(logs_dir) / job_name / "completions"
+
+    with open(executor_path) as f:
+        executor = json.load(f)
+
+    missing_workers = find_missing_workers(completions_dir, executor["world_size"])
+
+    for missing in sorted(missing_workers)[::-1]:
+        print(missing)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--logs-dir",
+        type=str,
+        help="Path to logs directory for `datatrove`.",
+    )
+    parser.add_argument(
+        "--job-name",
+        type=str,
+        default="port_droid",
+        help="Job name used in slurm, and name of the directory created inside the provided logs directory.",
+    )
+
+    args = parser.parse_args()
+
+    display_error_files(**vars(args))
+
+
+if __name__ == "__main__":
+    main()

examples/port_datasets/droid_rlds/port_droid.py

@@ -0,0 +1,430 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import logging
+import time
+from pathlib import Path
+
+import numpy as np
+import tensorflow_datasets as tfds
+
+from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.utils.utils import get_elapsed_time_in_days_hours_minutes_seconds
+
+DROID_SHARDS = 2048
+DROID_FPS = 15
+DROID_ROBOT_TYPE = "Franka"
+
+# Dataset schema slightly adapted from: https://droid-dataset.github.io/droid/the-droid-dataset.html#-dataset-schema
+DROID_FEATURES = {
+    # true on first step of the episode
+    "is_first": {
+        "dtype": "bool",
+        "shape": (1,),
+        "names": None,
+    },
+    # true on last step of the episode
+    "is_last": {
+        "dtype": "bool",
+        "shape": (1,),
+        "names": None,
+    },
+    # true on last step of the episode if it is a terminal step, True for demos
+    "is_terminal": {
+        "dtype": "bool",
+        "shape": (1,),
+        "names": None,
+    },
+    # language_instruction is also stored as "task" to follow LeRobot standard
+    "language_instruction": {
+        "dtype": "string",
+        "shape": (1,),
+        "names": None,
+    },
+    "language_instruction_2": {
+        "dtype": "string",
+        "shape": (1,),
+        "names": None,
+    },
+    "language_instruction_3": {
+        "dtype": "string",
+        "shape": (1,),
+        "names": None,
+    },
+    "observation.state.gripper_position": {
+        "dtype": "float32",
+        "shape": (1,),
+        "names": {
+            "axes": ["gripper"],
+        },
+    },
+    "observation.state.cartesian_position": {
+        "dtype": "float32",
+        "shape": (6,),
+        "names": {
+            "axes": ["x", "y", "z", "roll", "pitch", "yaw"],
+        },
+    },
+    "observation.state.joint_position": {
+        "dtype": "float32",
+        "shape": (7,),
+        "names": {
+            "axes": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6"],
+        },
+    },
+    # Add this new feature to follow LeRobot standard of using joint position + gripper
+    "observation.state": {
+        "dtype": "float32",
+        "shape": (8,),
+        "names": {
+            "axes": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6", "gripper"],
+        },
+    },
+    # Initially called wrist_image_left
+    "observation.images.wrist_left": {
+        "dtype": "video",
+        "shape": (180, 320, 3),
+        "names": [
+            "height",
+            "width",
+            "channels",
+        ],
+    },
+    # Initially called exterior_image_1_left
+    "observation.images.exterior_1_left": {
+        "dtype": "video",
+        "shape": (180, 320, 3),
+        "names": [
+            "height",
+            "width",
+            "channels",
+        ],
+    },
+    # Initially called exterior_image_2_left
+    "observation.images.exterior_2_left": {
+        "dtype": "video",
+        "shape": (180, 320, 3),
+        "names": [
+            "height",
+            "width",
+            "channels",
+        ],
+    },
+    "action.gripper_position": {
+        "dtype": "float32",
+        "shape": (1,),
+        "names": {
+            "axes": ["gripper"],
+        },
+    },
+    "action.gripper_velocity": {
+        "dtype": "float32",
+        "shape": (1,),
+        "names": {
+            "axes": ["gripper"],
+        },
+    },
+    "action.cartesian_position": {
+        "dtype": "float32",
+        "shape": (6,),
+        "names": {
+            "axes": ["x", "y", "z", "roll", "pitch", "yaw"],
+        },
+    },
+    "action.cartesian_velocity": {
+        "dtype": "float32",
+        "shape": (6,),
+        "names": {
+            "axes": ["x", "y", "z", "roll", "pitch", "yaw"],
+        },
+    },
+    "action.joint_position": {
+        "dtype": "float32",
+        "shape": (7,),
+        "names": {
+            "axes": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6"],
+        },
+    },
+    "action.joint_velocity": {
+        "dtype": "float32",
+        "shape": (7,),
+        "names": {
+            "axes": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6"],
+        },
+    },
+    # This feature was called "action" in RLDS dataset and consists of [6x joint velocities, 1x gripper position]
+    "action.original": {
+        "dtype": "float32",
+        "shape": (7,),
+        "names": {
+            "axes": ["x", "y", "z", "roll", "pitch", "yaw", "gripper"],
+        },
+    },
+    # Add this new feature to follow LeRobot standard of using joint position + gripper
+    "action": {
+        "dtype": "float32",
+        "shape": (8,),
+        "names": {
+            "axes": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6", "gripper"],
+        },
+    },
+    "discount": {
+        "dtype": "float32",
+        "shape": (1,),
+        "names": None,
+    },
+    "reward": {
+        "dtype": "float32",
+        "shape": (1,),
+        "names": None,
+    },
+    # Meta data that are the same for all frames in the episode
+    "task_category": {
+        "dtype": "string",
+        "shape": (1,),
+        "names": None,
+    },
+    "building": {
+        "dtype": "string",
+        "shape": (1,),
+        "names": None,
+    },
+    "collector_id": {
+        "dtype": "string",
+        "shape": (1,),
+        "names": None,
+    },
+    "date": {
+        "dtype": "string",
+        "shape": (1,),
+        "names": None,
+    },
+    "camera_extrinsics.wrist_left": {
+        "dtype": "float32",
+        "shape": (6,),
+        "names": {
+            "axes": ["x", "y", "z", "roll", "pitch", "yaw"],
+        },
+    },
+    "camera_extrinsics.exterior_1_left": {
+        "dtype": "float32",
+        "shape": (6,),
+        "names": {
+            "axes": ["x", "y", "z", "roll", "pitch", "yaw"],
+        },
+    },
+    "camera_extrinsics.exterior_2_left": {
+        "dtype": "float32",
+        "shape": (6,),
+        "names": {
+            "axes": ["x", "y", "z", "roll", "pitch", "yaw"],
+        },
+    },
+    "is_episode_successful": {
+        "dtype": "bool",
+        "shape": (1,),
+        "names": None,
+    },
+}
+
+
+def is_episode_successful(tf_episode_metadata):
+    # Adapted from: https://github.com/droid-dataset/droid_policy_learning/blob/dd1020eb20d981f90b5ff07dc80d80d5c0cb108b/robomimic/utils/rlds_utils.py#L8
+    return "/success/" in tf_episode_metadata["file_path"].numpy().decode()
+
+
+def generate_lerobot_frames(tf_episode):
+    m = tf_episode["episode_metadata"]
+    frame_meta = {
+        "task_category": m["building"].numpy().decode(),
+        "building": m["building"].numpy().decode(),
+        "collector_id": m["collector_id"].numpy().decode(),
+        "date": m["date"].numpy().decode(),
+        "camera_extrinsics.wrist_left": m["extrinsics_wrist_cam"].numpy(),
+        "camera_extrinsics.exterior_1_left": m["extrinsics_exterior_cam_1"].numpy(),
+        "camera_extrinsics.exterior_2_left": m["extrinsics_exterior_cam_2"].numpy(),
+        "is_episode_successful": np.array([is_episode_successful(m)]),
+    }
+    for f in tf_episode["steps"]:
+        # Dataset schema slightly adapted from: https://droid-dataset.github.io/droid/the-droid-dataset.html#-dataset-schema
+        frame = {
+            "is_first": np.array([f["is_first"].numpy()]),
+            "is_last": np.array([f["is_last"].numpy()]),
+            "is_terminal": np.array([f["is_terminal"].numpy()]),
+            "language_instruction": f["language_instruction"].numpy().decode(),
+            "language_instruction_2": f["language_instruction_2"].numpy().decode(),
+            "language_instruction_3": f["language_instruction_3"].numpy().decode(),
+            "observation.state.gripper_position": f["observation"]["gripper_position"].numpy(),
+            "observation.state.cartesian_position": f["observation"]["cartesian_position"].numpy(),
+            "observation.state.joint_position": f["observation"]["joint_position"].numpy(),
+            "observation.images.wrist_left": f["observation"]["wrist_image_left"].numpy(),
+            "observation.images.exterior_1_left": f["observation"]["exterior_image_1_left"].numpy(),
+            "observation.images.exterior_2_left": f["observation"]["exterior_image_2_left"].numpy(),
+            "action.gripper_position": f["action_dict"]["gripper_position"].numpy(),
+            "action.gripper_velocity": f["action_dict"]["gripper_velocity"].numpy(),
+            "action.cartesian_position": f["action_dict"]["cartesian_position"].numpy(),
+            "action.cartesian_velocity": f["action_dict"]["cartesian_velocity"].numpy(),
+            "action.joint_position": f["action_dict"]["joint_position"].numpy(),
+            "action.joint_velocity": f["action_dict"]["joint_velocity"].numpy(),
+            "discount": np.array([f["discount"].numpy()]),
+            "reward": np.array([f["reward"].numpy()]),
+            "action.original": f["action"].numpy(),
+        }
+
+        # language_instruction is also stored as "task" to follow LeRobot standard
+        frame["task"] = frame["language_instruction"]
+
+        # Add this new feature to follow LeRobot standard of using joint position + gripper
+        frame["observation.state"] = np.concatenate(
+            [frame["observation.state.joint_position"], frame["observation.state.gripper_position"]]
+        )
+        frame["action"] = np.concatenate([frame["action.joint_position"], frame["action.gripper_position"]])
+
+        # Meta data that are the same for all frames in the episode
+        frame.update(frame_meta)
+
+        # Cast fp64 to fp32
+        for key in frame:
+            if isinstance(frame[key], np.ndarray) and frame[key].dtype == np.float64:
+                frame[key] = frame[key].astype(np.float32)
+
+        yield frame
+
+
+def port_droid(
+    raw_dir: Path,
+    repo_id: str,
+    push_to_hub: bool = False,
+    num_shards: int | None = None,
+    shard_index: int | None = None,
+):
+    dataset_name = raw_dir.parent.name
+    version = raw_dir.name
+    data_dir = raw_dir.parent.parent
+
+    builder = tfds.builder(f"{dataset_name}/{version}", data_dir=data_dir, version="")
+
+    if num_shards is not None:
+        tfds_num_shards = builder.info.splits["train"].num_shards
+        if tfds_num_shards != DROID_SHARDS:
+            raise ValueError(
+                f"Number of shards of Droid dataset is expected to be {DROID_SHARDS} but is {tfds_num_shards}."
+            )
+        if num_shards != tfds_num_shards:
+            raise ValueError(
+                f"We only shard over the fixed number of shards provided by tensorflow dataset ({tfds_num_shards}), but {num_shards} shards provided instead."
+            )
+        if shard_index >= tfds_num_shards:
+            raise ValueError(
+                f"Shard index is greater than the num of shards ({shard_index} >= {num_shards})."
+            )
+
+        raw_dataset = builder.as_dataset(split=f"train[{shard_index}shard]")
+    else:
+        raw_dataset = builder.as_dataset(split="train")
+
+    lerobot_dataset = LeRobotDataset.create(
+        repo_id=repo_id,
+        robot_type=DROID_ROBOT_TYPE,
+        fps=DROID_FPS,
+        features=DROID_FEATURES,
+    )
+
+    start_time = time.time()
+    num_episodes = raw_dataset.cardinality().numpy().item()
+    logging.info(f"Number of episodes {num_episodes}")
+
+    for episode_index, episode in enumerate(raw_dataset):
+        elapsed_time = time.time() - start_time
+        d, h, m, s = get_elapsed_time_in_days_hours_minutes_seconds(elapsed_time)
+
+        logging.info(
+            f"{episode_index} / {num_episodes} episodes processed (after {d} days, {h} hours, {m} minutes, {s:.3f} seconds)"
+        )
+
+        for frame in generate_lerobot_frames(episode):
+            lerobot_dataset.add_frame(frame)
+
+        lerobot_dataset.save_episode()
+        logging.info("Save_episode")
+
+    if push_to_hub:
+        lerobot_dataset.push_to_hub(
+            # Add openx tag, since it belongs to the openx collection of datasets
+            tags=["openx"],
+            private=False,
+        )
+
+
+def validate_dataset(repo_id):
+    """Sanity check that ensure meta data can be loaded and all files are present."""
+    meta = LeRobotDatasetMetadata(repo_id)
+
+    if meta.total_episodes == 0:
+        raise ValueError("Number of episodes is 0.")
+
+    for ep_idx in range(meta.total_episodes):
+        data_path = meta.root / meta.get_data_file_path(ep_idx)
+
+        if not data_path.exists():
+            raise ValueError(f"Parquet file is missing in: {data_path}")
+
+        for vid_key in meta.video_keys:
+            vid_path = meta.root / meta.get_video_file_path(ep_idx, vid_key)
+            if not vid_path.exists():
+                raise ValueError(f"Video file is missing in: {vid_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--raw-dir",
+        type=Path,
+        required=True,
+        help="Directory containing input raw datasets (e.g. `path/to/dataset` or `path/to/dataset/version).",
+    )
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        help="Repositery identifier on Hugging Face: a community or a user name `/` the name of the dataset, required when push-to-hub is True",
+    )
+    parser.add_argument(
+        "--push-to-hub",
+        action="store_true",
+        help="Upload to hub.",
+    )
+    parser.add_argument(
+        "--num-shards",
+        type=int,
+        default=None,
+        help="Number of shards. Can be either None to load the full dataset, or 2048 to load one of the 2048 tensorflow dataset files.",
+    )
+    parser.add_argument(
+        "--shard-index",
+        type=int,
+        default=None,
+        help="Index of the shard. Can be either None to load the full dataset, or in [0,2047] to load one of the 2048 tensorflow dataset files.",
+    )
+
+    args = parser.parse_args()
+
+    port_droid(**vars(args))
+
+
+if __name__ == "__main__":
+    main()

examples/port_datasets/droid_rlds/slurm_aggregate_shards.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import logging
+from pathlib import Path
+
+from datatrove.executor import LocalPipelineExecutor
+from datatrove.executor.slurm import SlurmPipelineExecutor
+from datatrove.pipeline.base import PipelineStep
+from port_datasets.droid_rlds.port_droid import DROID_SHARDS
+
+from lerobot.datasets.aggregate import aggregate_datasets
+from lerobot.utils.utils import init_logging
+
+
+class AggregateDatasets(PipelineStep):
+    def __init__(
+        self,
+        repo_ids: list[str],
+        aggregated_repo_id: str,
+    ):
+        super().__init__()
+        self.repo_ids = repo_ids
+        self.aggr_repo_id = aggregated_repo_id
+
+    def run(self, data=None, rank: int = 0, world_size: int = 1):
+        init_logging()
+
+        # Since aggregate_datasets already handles parallel processing internally,
+        # we only need one worker to run the entire aggregation
+        if rank == 0:
+            logging.info(f"Starting aggregation of {len(self.repo_ids)} datasets into {self.aggr_repo_id}")
+            aggregate_datasets(self.repo_ids, self.aggr_repo_id)
+            logging.info("Aggregation complete!")
+        else:
+            logging.info(f"Worker {rank} skipping - only worker 0 performs aggregation")
+
+
+def make_aggregate_executor(
+    repo_ids, repo_id, job_name, logs_dir, workers, partition, cpus_per_task, mem_per_cpu, slurm=True
+):
+    kwargs = {
+        "pipeline": [
+            AggregateDatasets(repo_ids, repo_id),
+        ],
+        "logging_dir": str(logs_dir / job_name),
+    }
+
+    if slurm:
+        # For aggregation, we only need 1 task since aggregate_datasets handles everything
+        kwargs.update(
+            {
+                "job_name": job_name,
+                "tasks": 1,  # Only need 1 task for aggregation
+                "workers": 1,  # Only need 1 worker
+                "time": "08:00:00",
+                "partition": partition,
+                "cpus_per_task": cpus_per_task,
+                "sbatch_args": {"mem-per-cpu": mem_per_cpu},
+            }
+        )
+        executor = SlurmPipelineExecutor(**kwargs)
+    else:
+        kwargs.update(
+            {
+                "tasks": 1,
+                "workers": 1,
+            }
+        )
+        executor = LocalPipelineExecutor(**kwargs)
+
+    return executor
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        help="Repository identifier on Hugging Face: a community or a user name `/` the name of the dataset, required when push-to-hub is True.",
+    )
+    parser.add_argument(
+        "--logs-dir",
+        type=Path,
+        help="Path to logs directory for `datatrove`.",
+    )
+    parser.add_argument(
+        "--job-name",
+        type=str,
+        default="aggr_droid",
+        help="Job name used in slurm, and name of the directory created inside the provided logs directory.",
+    )
+    parser.add_argument(
+        "--slurm",
+        type=int,
+        default=1,
+        help="Launch over slurm. Use `--slurm 0` to launch sequentially (useful to debug).",
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=1,  # Changed default to 1 since aggregation doesn't need multiple workers
+        help="Number of slurm workers. For aggregation, this should be 1.",
+    )
+    parser.add_argument(
+        "--partition",
+        type=str,
+        help="Slurm partition. Ideally a CPU partition. No need for GPU partition.",
+    )
+    parser.add_argument(
+        "--cpus-per-task",
+        type=int,
+        default=8,
+        help="Number of cpus that each slurm worker will use.",
+    )
+    parser.add_argument(
+        "--mem-per-cpu",
+        type=str,
+        default="1950M",
+        help="Memory per cpu that each worker will use.",
+    )
+
+    args = parser.parse_args()
+    kwargs = vars(args)
+    kwargs["slurm"] = kwargs.pop("slurm") == 1
+
+    repo_ids = [f"{args.repo_id}_world_{DROID_SHARDS}_rank_{rank}" for rank in range(DROID_SHARDS)]
+    aggregate_executor = make_aggregate_executor(repo_ids, **kwargs)
+    aggregate_executor.run()
+
+
+if __name__ == "__main__":
+    main()

examples/port_datasets/droid_rlds/slurm_port_shards.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+from pathlib import Path
+
+from datatrove.executor import LocalPipelineExecutor
+from datatrove.executor.slurm import SlurmPipelineExecutor
+from datatrove.pipeline.base import PipelineStep
+from port_datasets.droid_rlds.port_droid import DROID_SHARDS
+
+
+class PortDroidShards(PipelineStep):
+    def __init__(
+        self,
+        raw_dir: Path | str,
+        repo_id: str = None,
+    ):
+        super().__init__()
+        self.raw_dir = Path(raw_dir)
+        self.repo_id = repo_id
+
+    def run(self, data=None, rank: int = 0, world_size: int = 1):
+        from datasets.utils.tqdm import disable_progress_bars
+        from port_datasets.droid_rlds.port_droid import port_droid, validate_dataset
+
+        from lerobot.utils.utils import init_logging
+
+        init_logging()
+        disable_progress_bars()
+
+        shard_repo_id = f"{self.repo_id}_world_{world_size}_rank_{rank}"
+
+        try:
+            validate_dataset(shard_repo_id)
+            return
+        except Exception:
+            pass  # nosec B110 - Dataset doesn't exist yet, continue with porting
+
+        port_droid(
+            self.raw_dir,
+            shard_repo_id,
+            push_to_hub=False,
+            num_shards=world_size,
+            shard_index=rank,
+        )
+
+        validate_dataset(shard_repo_id)
+
+
+def make_port_executor(
+    raw_dir, repo_id, job_name, logs_dir, workers, partition, cpus_per_task, mem_per_cpu, slurm=True
+):
+    kwargs = {
+        "pipeline": [
+            PortDroidShards(raw_dir, repo_id),
+        ],
+        "logging_dir": str(logs_dir / job_name),
+    }
+
+    if slurm:
+        kwargs.update(
+            {
+                "job_name": job_name,
+                "tasks": DROID_SHARDS,
+                "workers": workers,
+                "time": "08:00:00",
+                "partition": partition,
+                "cpus_per_task": cpus_per_task,
+                "sbatch_args": {"mem-per-cpu": mem_per_cpu},
+            }
+        )
+        executor = SlurmPipelineExecutor(**kwargs)
+    else:
+        kwargs.update(
+            {
+                "tasks": 1,
+                "workers": 1,
+            }
+        )
+        executor = LocalPipelineExecutor(**kwargs)
+
+    return executor
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--raw-dir",
+        type=Path,
+        required=True,
+        help="Directory containing input raw datasets (e.g. `path/to/dataset` or `path/to/dataset/version).",
+    )
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        help="Repositery identifier on Hugging Face: a community or a user name `/` the name of the dataset, required when push-to-hub is True.",
+    )
+    parser.add_argument(
+        "--logs-dir",
+        type=Path,
+        help="Path to logs directory for `datatrove`.",
+    )
+    parser.add_argument(
+        "--job-name",
+        type=str,
+        default="port_droid",
+        help="Job name used in slurm, and name of the directory created inside the provided logs directory.",
+    )
+    parser.add_argument(
+        "--slurm",
+        type=int,
+        default=1,
+        help="Launch over slurm. Use `--slurm 0` to launch sequentially (useful to debug).",
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=2048,
+        help="Number of slurm workers. It should be less than the maximum number of shards.",
+    )
+    parser.add_argument(
+        "--partition",
+        type=str,
+        help="Slurm partition. Ideally a CPU partition. No need for GPU partition.",
+    )
+    parser.add_argument(
+        "--cpus-per-task",
+        type=int,
+        default=8,
+        help="Number of cpus that each slurm worker will use.",
+    )
+    parser.add_argument(
+        "--mem-per-cpu",
+        type=str,
+        default="1950M",
+        help="Memory per cpu that each worker will use.",
+    )
+
+    args = parser.parse_args()
+    kwargs = vars(args)
+    kwargs["slurm"] = kwargs.pop("slurm") == 1
+    port_executor = make_port_executor(**kwargs)
+    port_executor.run()
+
+
+if __name__ == "__main__":
+    main()

examples/port_datasets/droid_rlds/slurm_upload.py

@@ -0,0 +1,281 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import logging
+import os
+from pathlib import Path
+
+from datatrove.executor import LocalPipelineExecutor
+from datatrove.executor.slurm import SlurmPipelineExecutor
+from datatrove.pipeline.base import PipelineStep
+from huggingface_hub import HfApi
+from huggingface_hub.constants import REPOCARD_NAME
+from port_datasets.droid_rlds.port_droid import DROID_SHARDS
+
+from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDatasetMetadata
+from lerobot.datasets.utils import create_lerobot_dataset_card
+from lerobot.utils.utils import init_logging
+
+
+class UploadDataset(PipelineStep):
+    def __init__(
+        self,
+        repo_id: str,
+        branch: str | None = None,
+        revision: str | None = None,
+        tags: list | None = None,
+        license: str | None = "apache-2.0",
+        private: bool = False,
+        distant_repo_id: str | None = None,
+        **card_kwargs,
+    ):
+        super().__init__()
+        self.repo_id = repo_id
+        self.distant_repo_id = self.repo_id if distant_repo_id is None else distant_repo_id
+        self.branch = branch
+        self.tags = tags
+        self.license = license
+        self.private = private
+        self.card_kwargs = card_kwargs
+        self.revision = revision if revision else CODEBASE_VERSION
+
+        if os.environ.get("HF_HUB_ENABLE_HF_TRANSFER", "0") != "1":
+            logging.warning(
+                'HF_HUB_ENABLE_HF_TRANSFER is not set to "1". Install hf_transfer and set the env '
+                "variable for faster uploads:\npip install hf-transfer\nexport HF_HUB_ENABLE_HF_TRANSFER=1"
+            )
+
+        self.create_repo()
+
+    def create_repo(self):
+        logging.info(f"Loading meta data from {self.repo_id}...")
+        meta = LeRobotDatasetMetadata(self.repo_id)
+
+        logging.info(f"Creating repo {self.distant_repo_id}...")
+        hub_api = HfApi()
+        hub_api.create_repo(
+            repo_id=self.distant_repo_id,
+            private=self.private,
+            repo_type="dataset",
+            exist_ok=True,
+        )
+        if self.branch:
+            hub_api.create_branch(
+                repo_id=self.distant_repo_id,
+                branch=self.branch,
+                revision=self.revision,
+                repo_type="dataset",
+                exist_ok=True,
+            )
+
+        if not hub_api.file_exists(
+            self.distant_repo_id, REPOCARD_NAME, repo_type="dataset", revision=self.branch
+        ):
+            card = create_lerobot_dataset_card(
+                tags=self.tags, dataset_info=meta.info, license=self.license, **self.card_kwargs
+            )
+            card.push_to_hub(repo_id=self.distant_repo_id, repo_type="dataset", revision=self.branch)
+
+            hub_api.create_tag(self.distant_repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
+
+        def list_files_recursively(directory):
+            base_path = Path(directory)
+            return [str(file.relative_to(base_path)) for file in base_path.rglob("*") if file.is_file()]
+
+        logging.info(f"Listing all local files from {self.repo_id}...")
+        self.file_paths = list_files_recursively(meta.root)
+        self.file_paths = sorted(self.file_paths)
+
+    def create_chunks(self, lst, n):
+        from itertools import islice
+
+        it = iter(lst)
+        return [list(islice(it, size)) for size in [len(lst) // n + (i < len(lst) % n) for i in range(n)]]
+
+    def create_commits(self, additions):
+        import logging
+        import math
+        import random
+        import time
+
+        from huggingface_hub import create_commit
+        from huggingface_hub.utils import HfHubHTTPError
+
+        FILES_BETWEEN_COMMITS = 10  # noqa: N806
+        BASE_DELAY = 0.1  # noqa: N806
+        MAX_RETRIES = 12  # noqa: N806
+
+        # Split the files into smaller chunks for faster commit
+        # and avoiding "A commit has happened since" error
+        num_chunks = math.ceil(len(additions) / FILES_BETWEEN_COMMITS)
+        chunks = self.create_chunks(additions, num_chunks)
+
+        for chunk in chunks:
+            retries = 0
+            while True:
+                try:
+                    create_commit(
+                        self.distant_repo_id,
+                        repo_type="dataset",
+                        operations=chunk,
+                        commit_message=f"DataTrove upload ({len(chunk)} files)",
+                        revision=self.branch,
+                    )
+                    # TODO: every 100 chunks super_squach_commits()
+                    logging.info("create_commit completed!")
+                    break
+                except HfHubHTTPError as e:
+                    if "A commit has happened since" in e.server_message:
+                        if retries >= MAX_RETRIES:
+                            logging.error(f"Failed to create commit after {MAX_RETRIES=}. Giving up.")
+                            raise e
+                        logging.info("Commit creation race condition issue. Waiting...")
+                        time.sleep(BASE_DELAY * 2**retries + random.uniform(0, 2))
+                        retries += 1
+                    else:
+                        raise e
+
+    def run(self, data=None, rank: int = 0, world_size: int = 1):
+        import logging
+
+        from datasets.utils.tqdm import disable_progress_bars
+        from huggingface_hub import CommitOperationAdd, preupload_lfs_files
+
+        from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
+        from lerobot.utils.utils import init_logging
+
+        init_logging()
+        disable_progress_bars()
+
+        chunks = self.create_chunks(self.file_paths, world_size)
+        file_paths = chunks[rank]
+
+        if len(file_paths) == 0:
+            raise ValueError(file_paths)
+
+        logging.info("Pre-uploading LFS files...")
+        for i, path in enumerate(file_paths):
+            logging.info(f"{i}: {path}")
+
+        meta = LeRobotDatasetMetadata(self.repo_id)
+        additions = [
+            CommitOperationAdd(path_in_repo=path, path_or_fileobj=meta.root / path) for path in file_paths
+        ]
+        preupload_lfs_files(
+            repo_id=self.distant_repo_id, repo_type="dataset", additions=additions, revision=self.branch
+        )
+
+        logging.info("Creating commits...")
+        self.create_commits(additions)
+        logging.info("Done!")
+
+
+def make_upload_executor(
+    repo_id, job_name, logs_dir, workers, partition, cpus_per_task, mem_per_cpu, slurm=True
+):
+    kwargs = {
+        "pipeline": [
+            UploadDataset(repo_id),
+        ],
+        "logging_dir": str(logs_dir / job_name),
+    }
+
+    if slurm:
+        kwargs.update(
+            {
+                "job_name": job_name,
+                "tasks": DROID_SHARDS,
+                "workers": workers,
+                "time": "08:00:00",
+                "partition": partition,
+                "cpus_per_task": cpus_per_task,
+                "sbatch_args": {"mem-per-cpu": mem_per_cpu},
+            }
+        )
+        executor = SlurmPipelineExecutor(**kwargs)
+    else:
+        kwargs.update(
+            {
+                "tasks": DROID_SHARDS,
+                "workers": 1,
+            }
+        )
+        executor = LocalPipelineExecutor(**kwargs)
+
+    return executor
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        help="Repositery identifier on Hugging Face: a community or a user name `/` the name of the dataset, required when push-to-hub is True.",
+    )
+    parser.add_argument(
+        "--logs-dir",
+        type=Path,
+        help="Path to logs directory for `datatrove`.",
+    )
+    parser.add_argument(
+        "--job-name",
+        type=str,
+        default="upload_droid",
+        help="Job name used in slurm, and name of the directory created inside the provided logs directory.",
+    )
+    parser.add_argument(
+        "--slurm",
+        type=int,
+        default=1,
+        help="Launch over slurm. Use `--slurm 0` to launch sequentially (useful to debug).",
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=50,
+        help="Number of slurm workers. It should be less than the maximum number of shards.",
+    )
+    parser.add_argument(
+        "--partition",
+        type=str,
+        help="Slurm partition. Ideally a CPU partition. No need for GPU partition.",
+    )
+    parser.add_argument(
+        "--cpus-per-task",
+        type=int,
+        default=8,
+        help="Number of cpus that each slurm worker will use.",
+    )
+    parser.add_argument(
+        "--mem-per-cpu",
+        type=str,
+        default="1950M",
+        help="Memory per cpu that each worker will use.",
+    )
+
+    init_logging()
+
+    args = parser.parse_args()
+    kwargs = vars(args)
+    kwargs["slurm"] = kwargs.pop("slurm") == 1
+    upload_executor = make_upload_executor(**kwargs)
+    upload_executor.run()
+
+
+if __name__ == "__main__":
+    main()

examples/use_dataset_tools.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python
+"""
+Example script demonstrating dataset tools utilities.
+
+This script shows how to:
+1. Delete episodes from a dataset
+2. Split a dataset into train/val sets
+3. Add/remove features
+4. Merge datasets
+
+Usage:
+    python examples/use_dataset_tools.py
+"""
+
+import numpy as np
+
+from lerobot.datasets.dataset_tools import (
+    add_feature,
+    delete_episodes,
+    merge_datasets,
+    remove_feature,
+    split_dataset,
+)
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+
+def main():
+    # Load an existing dataset (replace with your dataset)
+    dataset = LeRobotDataset("lerobot/pusht")
+
+    print(f"Original dataset: {dataset.meta.total_episodes} episodes, {dataset.meta.total_frames} frames")
+    print(f"Features: {list(dataset.meta.features.keys())}")
+
+    # Example 1: Delete episodes
+    print("\n1. Deleting episodes 0 and 2...")
+    filtered_dataset = delete_episodes(dataset, episode_indices=[0, 2], repo_id="pusht_filtered")
+    print(f"Filtered dataset: {filtered_dataset.meta.total_episodes} episodes")
+
+    # Example 2: Split dataset
+    print("\n2. Splitting dataset into train/val...")
+    splits = split_dataset(
+        dataset,
+        splits={"train": 0.8, "val": 0.2},
+    )
+    print(f"Train split: {splits['train'].meta.total_episodes} episodes")
+    print(f"Val split: {splits['val'].meta.total_episodes} episodes")
+
+    # Example 3: Add a feature
+    print("\n3. Adding a reward feature...")
+
+    # Method 1: Pre-computed values
+    reward_values = np.random.randn(dataset.meta.total_frames).astype(np.float32)
+    dataset_with_reward = add_feature(
+        dataset,
+        feature_name="reward",
+        feature_values=reward_values,
+        feature_info={
+            "dtype": "float32",
+            "shape": (1,),
+            "names": None,
+        },
+        repo_id="pusht_with_reward",
+    )
+
+    # Method 2: Using a callable
+    def compute_success(frame_dict, episode_idx, frame_idx):
+        # Example: mark last 10 frames of each episode as successful
+        episode_length = 10  # You'd get this from episode metadata
+        return float(frame_idx >= episode_length - 10)
+
+    dataset_with_success = add_feature(
+        dataset_with_reward,
+        feature_name="success",
+        feature_values=compute_success,
+        feature_info={
+            "dtype": "float32",
+            "shape": (1,),
+            "names": None,
+        },
+        repo_id="pusht_with_reward_and_success",
+    )
+
+    print(f"New features: {list(dataset_with_success.meta.features.keys())}")
+
+    # Example 4: Remove features
+    print("\n4. Removing the success feature...")
+    dataset_cleaned = remove_feature(dataset_with_success, feature_names="success", repo_id="pusht_cleaned")
+    print(f"Features after removal: {list(dataset_cleaned.meta.features.keys())}")
+
+    # Example 5: Merge datasets
+    print("\n5. Merging train and val splits back together...")
+    merged = merge_datasets([splits["train"], splits["val"]], output_repo_id="pusht_merged")
+    print(f"Merged dataset: {merged.meta.total_episodes} episodes")
+
+    # Example 6: Complex workflow
+    print("\n6. Complex workflow example...")
+
+    # Remove a camera if dataset has multiple
+    if len(dataset.meta.camera_keys) > 1:
+        camera_to_remove = dataset.meta.camera_keys[0]
+        print(f"Removing camera: {camera_to_remove}")
+        dataset_no_cam = remove_feature(
+            dataset, feature_names=camera_to_remove, repo_id="pusht_no_first_camera"
+        )
+        print(f"Remaining cameras: {dataset_no_cam.meta.camera_keys}")
+
+    print("\nDone! Check ~/.cache/huggingface/lerobot/ for the created datasets.")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -73,7 +73,6 @@ dependencies = [
     "pynput>=1.7.7",
     "pyserial>=3.5",
     "wandb>=0.20.0",
-    "scipy>=1.15.2",
 
     "torch>=2.2.1,<2.8.0", # TODO: Bumb dependency
     "torchcodec>=0.2.1,<0.6.0; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')", # TODO: Bumb dependency
@@ -96,7 +95,7 @@ dependencies = [
 # Common
 pygame-dep = ["pygame>=2.5.1"]
 placo-dep = ["placo>=0.9.6"]
-transformers-dep = ["transformers<=4.52.0"]
+transformers-dep = ["transformers>=4.50.3,<4.52.0"] # TODO: Bumb dependency
 grpcio-dep = ["grpcio==1.73.1", "protobuf==6.31.0"]
 
 # Motors
@@ -112,7 +111,6 @@ intelrealsense = [
     "pyrealsense2>=2.55.1.6486 ; sys_platform != 'darwin'",
     "pyrealsense2-macosx>=2.54 ; sys_platform == 'darwin'",
 ]
-phone = ["hebi-py>=2.8.0", "teleop>=0.1.0"]
 # stretch = [
 #     "hello-robot-stretch-body>=0.7.27 ; sys_platform == 'linux'",
 #     "pyrender @ git+https://github.com/mmatl/pyrender.git ; sys_platform == 'linux'",
@@ -154,8 +152,7 @@ all = [
     "lerobot[video_benchmark]",
     "lerobot[aloha]",
     "lerobot[pusht]",
-    "lerobot[xarm]",
-    "lerobot[phone]",
+    "lerobot[xarm]"
 ]
 
 [project.scripts]

scripts/convert_videos_to_images.py

@@ -1,74 +0,0 @@
-#!/usr/bin/env python
-
-"""
-Convert video dataset to image dataset for faster training.
-This pre-extracts all frames from MP4 files to PNG images.
-"""
-
-import argparse
-from pathlib import Path
-import logging
-import shutil
-
-def convert_dataset_videos_to_images(repo_id: str, root: str | None = None):
-    """Convert all videos in a LeRobot dataset to individual image files."""
-    from lerobot.datasets.lerobot_dataset import LeRobotDataset
-    from lerobot.datasets.video_utils import decode_video_frames
-    import torch
-    
-    # Load dataset
-    dataset = LeRobotDataset(repo_id, root=root, download_videos=True)
-    
-    total_frames_processed = 0
-    
-    for ep_idx in range(dataset.meta.total_episodes):
-        logging.info(f"Processing episode {ep_idx}/{dataset.meta.total_episodes}")
-        
-        for vid_key in dataset.meta.video_keys:
-            video_path = dataset.root / dataset.meta.get_video_file_path(ep_idx, vid_key)
-            
-            if not video_path.exists():
-                logging.warning(f"Video not found: {video_path}")
-                continue
-                
-            # Create image directory  
-            img_dir = dataset.root / f"images/chunk-{dataset.meta.get_episode_chunk(ep_idx)}/{vid_key}"
-            img_dir.mkdir(parents=True, exist_ok=True)
-            
-            # Decode all frames from video
-            # Get episode length to decode all frames
-            ep_length = dataset.meta.episodes[ep_idx]["length"]
-            timestamps = [i / dataset.fps for i in range(ep_length)]
-            
-            try:
-                frames = decode_video_frames(video_path, timestamps, dataset.tolerance_s, dataset.video_backend)
-                
-                # Save each frame as PNG
-                for i, frame in enumerate(frames.squeeze(0)):
-                    img_path = img_dir / f"episode_{ep_idx:06d}_{i:06d}.png"
-                    # Convert tensor to PIL and save
-                    import torchvision.transforms as T
-                    to_pil = T.ToPILImage()
-                    pil_frame = to_pil(frame)
-                    pil_frame.save(img_path)
-                    
-                total_frames_processed += len(frames.squeeze(0))
-                logging.info(f"  Extracted {len(frames.squeeze(0))} frames to {img_dir}")
-                
-            except Exception as e:
-                logging.error(f"Failed to process {video_path}: {e}")
-                continue
-    
-    logging.info(f"Conversion complete! Processed {total_frames_processed} total frames")
-    logging.info(f"You can now use download_videos=False to use the extracted images")
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Convert LeRobot video dataset to images")
-    parser.add_argument("repo_id", help="Dataset repo ID (e.g., 'kenmacken/record-test-2')")
-    parser.add_argument("--root", help="Local root directory", default=None)
-    
-    args = parser.parse_args()
-    
-    logging.basicConfig(level=logging.INFO)
-    
-    convert_dataset_videos_to_images(args.repo_id, args.root)

src/lerobot/calibrate.py

@@ -18,7 +18,7 @@ Helper to recalibrate your device (robot or teleoperator).
 Example:
 
 ```shell
-python -m lerobot.calibrate \
+lerobot-calibrate \
     --teleop.type=so100_leader \
     --teleop.port=/dev/tty.usbmodem58760431551 \
     --teleop.id=blue

src/lerobot/cameras/opencv/camera_opencv.py

@@ -60,7 +60,7 @@ class OpenCVCamera(Camera):
     or port changes, especially on Linux. Use the provided utility script to find
     available camera indices or paths:
     ```bash
-    python -m lerobot.find_cameras opencv
+    lerobot-find-cameras opencv
     ```
 
     The camera's default settings (FPS, resolution, color mode) are used unless
@@ -165,8 +165,7 @@ class OpenCVCamera(Camera):
             self.videocapture.release()
             self.videocapture = None
             raise ConnectionError(
-                f"Failed to open {self}."
-                f"Run `python -m lerobot.find_cameras opencv` to find available cameras."
+                f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
             )
 
         self._configure_capture_settings()

src/lerobot/cameras/realsense/camera_realsense.py

@@ -51,7 +51,7 @@ class RealSenseCamera(Camera):
 
     Use the provided utility script to find available camera indices and default profiles:
     ```bash
-    python -m lerobot.find_cameras realsense
+    lerobot-find-cameras realsense
     ```
 
     A `RealSenseCamera` instance requires a configuration object specifying the
@@ -176,8 +176,7 @@ class RealSenseCamera(Camera):
             self.rs_profile = None
             self.rs_pipeline = None
             raise ConnectionError(
-                f"Failed to open {self}."
-                "Run `python -m lerobot.find_cameras realsense` to find available cameras."
+                f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
             ) from e
 
         self._configure_capture_settings()

src/lerobot/configs/default.py

@@ -33,8 +33,6 @@ class DatasetConfig:
     # Root directory where the dataset will be stored (e.g. 'dataset/path').
     root: str | None = None
     episodes: list[int] | None = None
-    # Percentage of dataset to use (0-100). If set, overrides episodes parameter.
-    percentage: float | None = None
     image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
     revision: str | None = None
     use_imagenet_stats: bool = True

src/lerobot/configs/policies.py

@@ -26,7 +26,7 @@ from huggingface_hub import hf_hub_download
 from huggingface_hub.constants import CONFIG_NAME
 from huggingface_hub.errors import HfHubHTTPError
 
-from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.constants import ACTION, OBS_STATE
 from lerobot.optim.optimizers import OptimizerConfig
 from lerobot.optim.schedulers import LRSchedulerConfig
@@ -53,6 +53,7 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
     """
 
     n_obs_steps: int = 1
+    normalization_mapping: dict[str, NormalizationMode] = field(default_factory=dict)
 
     input_features: dict[str, PolicyFeature] = field(default_factory=dict)
     output_features: dict[str, PolicyFeature] = field(default_factory=dict)

src/lerobot/configs/types.py

@@ -24,7 +24,6 @@ class FeatureType(str, Enum):
     ENV = "ENV"
     ACTION = "ACTION"
     REWARD = "REWARD"
-    LANGUAGE = "LANGUAGE"
 
 
 class NormalizationMode(str, Enum):

src/lerobot/constants.py

@@ -21,7 +21,6 @@ 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"
 

src/lerobot/datasets/aggregate.py

@@ -0,0 +1,505 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import shutil
+from pathlib import Path
+
+import pandas as pd
+import tqdm
+
+from lerobot.datasets.compute_stats import aggregate_stats
+from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
+from lerobot.datasets.utils import (
+    DEFAULT_CHUNK_SIZE,
+    DEFAULT_DATA_FILE_SIZE_IN_MB,
+    DEFAULT_DATA_PATH,
+    DEFAULT_EPISODES_PATH,
+    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
+    DEFAULT_VIDEO_PATH,
+    get_parquet_file_size_in_mb,
+    get_video_size_in_mb,
+    to_parquet_with_hf_images,
+    update_chunk_file_indices,
+    write_info,
+    write_stats,
+    write_tasks,
+)
+from lerobot.datasets.video_utils import concat_video_files
+
+
+def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
+    """Validates that all dataset metadata have consistent properties.
+
+    Ensures all datasets have the same fps, robot_type, and features to guarantee
+    compatibility when aggregating them into a single dataset.
+
+    Args:
+        all_metadata: List of LeRobotDatasetMetadata objects to validate.
+
+    Returns:
+        tuple: A tuple containing (fps, robot_type, features) from the first metadata.
+
+    Raises:
+        ValueError: If any metadata has different fps, robot_type, or features
+                   than the first metadata in the list.
+    """
+
+    fps = all_metadata[0].fps
+    robot_type = all_metadata[0].robot_type
+    features = all_metadata[0].features
+
+    for meta in tqdm.tqdm(all_metadata, desc="Validate all meta data"):
+        if fps != meta.fps:
+            raise ValueError(f"Same fps is expected, but got fps={meta.fps} instead of {fps}.")
+        if robot_type != meta.robot_type:
+            raise ValueError(
+                f"Same robot_type is expected, but got robot_type={meta.robot_type} instead of {robot_type}."
+            )
+        if features != meta.features:
+            raise ValueError(
+                f"Same features is expected, but got features={meta.features} instead of {features}."
+            )
+
+    return fps, robot_type, features
+
+
+def update_data_df(df, src_meta, dst_meta):
+    """Updates a data DataFrame with new indices and task mappings for aggregation.
+
+    Adjusts episode indices, frame indices, and task indices to account for
+    previously aggregated data in the destination dataset.
+
+    Args:
+        df: DataFrame containing the data to be updated.
+        src_meta: Source dataset metadata.
+        dst_meta: Destination dataset metadata.
+
+    Returns:
+        pd.DataFrame: Updated DataFrame with adjusted indices.
+    """
+
+    def _update(row):
+        row["episode_index"] = row["episode_index"] + dst_meta.info["total_episodes"]
+        row["index"] = row["index"] + dst_meta.info["total_frames"]
+        task = src_meta.tasks.iloc[row["task_index"]].name
+        row["task_index"] = dst_meta.tasks.loc[task].task_index.item()
+        return row
+
+    return df.apply(_update, axis=1)
+
+
+def update_meta_data(
+    df,
+    dst_meta,
+    meta_idx,
+    data_idx,
+    videos_idx,
+):
+    """Updates metadata DataFrame with new chunk, file, and timestamp indices.
+
+    Adjusts all indices and timestamps to account for previously aggregated
+    data and videos in the destination dataset.
+
+    Args:
+        df: DataFrame containing the metadata to be updated.
+        dst_meta: Destination dataset metadata.
+        meta_idx: Dictionary containing current metadata chunk and file indices.
+        data_idx: Dictionary containing current data chunk and file indices.
+        videos_idx: Dictionary containing current video indices and timestamps.
+
+    Returns:
+        pd.DataFrame: Updated DataFrame with adjusted indices and timestamps.
+    """
+
+    def _update(row):
+        row["meta/episodes/chunk_index"] = row["meta/episodes/chunk_index"] + meta_idx["chunk"]
+        row["meta/episodes/file_index"] = row["meta/episodes/file_index"] + meta_idx["file"]
+        row["data/chunk_index"] = row["data/chunk_index"] + data_idx["chunk"]
+        row["data/file_index"] = row["data/file_index"] + data_idx["file"]
+        for key, video_idx in videos_idx.items():
+            row[f"videos/{key}/chunk_index"] = row[f"videos/{key}/chunk_index"] + video_idx["chunk"]
+            row[f"videos/{key}/file_index"] = row[f"videos/{key}/file_index"] + video_idx["file"]
+            row[f"videos/{key}/from_timestamp"] = (
+                row[f"videos/{key}/from_timestamp"] + video_idx["latest_duration"]
+            )
+            row[f"videos/{key}/to_timestamp"] = (
+                row[f"videos/{key}/to_timestamp"] + video_idx["latest_duration"]
+            )
+
+        row["dataset_from_index"] = row["dataset_from_index"] + dst_meta.info["total_frames"]
+        row["dataset_to_index"] = row["dataset_to_index"] + dst_meta.info["total_frames"]
+        row["episode_index"] = row["episode_index"] + dst_meta.info["total_episodes"]
+        return row
+
+    return df.apply(_update, axis=1)
+
+
+def aggregate_datasets(
+    repo_ids: list[str],
+    aggr_repo_id: str,
+    roots: list[Path] = None,
+    aggr_root: Path = None,
+    data_files_size_in_mb: float = None,
+    video_files_size_in_mb: float = None,
+    chunk_size: int = None,
+):
+    """Aggregates multiple LeRobot datasets into a single unified dataset.
+
+    This is the main function that orchestrates the aggregation process by:
+    1. Loading and validating all source dataset metadata
+    2. Creating a new destination dataset with unified tasks
+    3. Aggregating videos, data, and metadata from all source datasets
+    4. Finalizing the aggregated dataset with proper statistics
+
+    Args:
+        repo_ids: List of repository IDs for the datasets to aggregate.
+        aggr_repo_id: Repository ID for the aggregated output dataset.
+        roots: Optional list of root paths for the source datasets.
+        aggr_root: Optional root path for the aggregated dataset.
+        data_files_size_in_mb: Maximum size for data files in MB (defaults to DEFAULT_DATA_FILE_SIZE_IN_MB)
+        video_files_size_in_mb: Maximum size for video files in MB (defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB)
+        chunk_size: Maximum number of files per chunk (defaults to DEFAULT_CHUNK_SIZE)
+    """
+    logging.info("Start aggregate_datasets")
+
+    if data_files_size_in_mb is None:
+        data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
+    if video_files_size_in_mb is None:
+        video_files_size_in_mb = DEFAULT_VIDEO_FILE_SIZE_IN_MB
+    if chunk_size is None:
+        chunk_size = DEFAULT_CHUNK_SIZE
+
+    all_metadata = (
+        [LeRobotDatasetMetadata(repo_id) for repo_id in repo_ids]
+        if roots is None
+        else [
+            LeRobotDatasetMetadata(repo_id, root=root) for repo_id, root in zip(repo_ids, roots, strict=False)
+        ]
+    )
+    fps, robot_type, features = validate_all_metadata(all_metadata)
+    video_keys = [key for key in features if features[key]["dtype"] == "video"]
+
+    dst_meta = LeRobotDatasetMetadata.create(
+        repo_id=aggr_repo_id,
+        fps=fps,
+        robot_type=robot_type,
+        features=features,
+        root=aggr_root,
+    )
+
+    logging.info("Find all tasks")
+    unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
+    dst_meta.tasks = pd.DataFrame({"task_index": range(len(unique_tasks))}, index=unique_tasks)
+
+    meta_idx = {"chunk": 0, "file": 0}
+    data_idx = {"chunk": 0, "file": 0}
+    videos_idx = {
+        key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
+    }
+
+    dst_meta.episodes = {}
+
+    for src_meta in tqdm.tqdm(all_metadata, desc="Copy data and videos"):
+        videos_idx = aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size)
+        data_idx = aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size)
+
+        meta_idx = aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx)
+
+        dst_meta.info["total_episodes"] += src_meta.total_episodes
+        dst_meta.info["total_frames"] += src_meta.total_frames
+
+    finalize_aggregation(dst_meta, all_metadata)
+    logging.info("Aggregation complete.")
+
+
+def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size):
+    """Aggregates video chunks from a source dataset into the destination dataset.
+
+    Handles video file concatenation and rotation based on file size limits.
+    Creates new video files when size limits are exceeded.
+
+    Args:
+        src_meta: Source dataset metadata.
+        dst_meta: Destination dataset metadata.
+        videos_idx: Dictionary tracking video chunk and file indices.
+        video_files_size_in_mb: Maximum size for video files in MB (defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB)
+        chunk_size: Maximum number of files per chunk (defaults to DEFAULT_CHUNK_SIZE)
+
+    Returns:
+        dict: Updated videos_idx with current chunk and file indices.
+    """
+    for key, video_idx in videos_idx.items():
+        unique_chunk_file_pairs = {
+            (chunk, file)
+            for chunk, file in zip(
+                src_meta.episodes[f"videos/{key}/chunk_index"],
+                src_meta.episodes[f"videos/{key}/file_index"],
+                strict=False,
+            )
+        }
+        unique_chunk_file_pairs = sorted(unique_chunk_file_pairs)
+
+        chunk_idx = video_idx["chunk"]
+        file_idx = video_idx["file"]
+
+        for src_chunk_idx, src_file_idx in unique_chunk_file_pairs:
+            src_path = src_meta.root / DEFAULT_VIDEO_PATH.format(
+                video_key=key,
+                chunk_index=src_chunk_idx,
+                file_index=src_file_idx,
+            )
+
+            dst_path = dst_meta.root / DEFAULT_VIDEO_PATH.format(
+                video_key=key,
+                chunk_index=chunk_idx,
+                file_index=file_idx,
+            )
+
+            # If a new file is created, we don't want to increment the latest_duration
+            update_latest_duration = False
+
+            if not dst_path.exists():
+                # First write to this destination file
+                dst_path.parent.mkdir(parents=True, exist_ok=True)
+                shutil.copy(str(src_path), str(dst_path))
+                continue  # not accumulating further, already copied the file in place
+
+            # Check file sizes before appending
+            src_size = get_video_size_in_mb(src_path)
+            dst_size = get_video_size_in_mb(dst_path)
+
+            if dst_size + src_size >= video_files_size_in_mb:
+                # Rotate to a new chunk/file
+                chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, chunk_size)
+                dst_path = dst_meta.root / DEFAULT_VIDEO_PATH.format(
+                    video_key=key,
+                    chunk_index=chunk_idx,
+                    file_index=file_idx,
+                )
+                dst_path.parent.mkdir(parents=True, exist_ok=True)
+                shutil.copy(str(src_path), str(dst_path))
+            else:
+                # Get the timestamps shift for this video
+                timestamps_shift_s = dst_meta.info["total_frames"] / dst_meta.info["fps"]
+
+                # Append to existing video file
+                concat_video_files(
+                    [dst_path, src_path],
+                    dst_meta.root,
+                    key,
+                    chunk_idx,
+                    file_idx,
+                )
+                # Update the latest_duration when appending (shifts timestamps!)
+                update_latest_duration = not update_latest_duration
+
+        # Update the videos_idx with the final chunk and file indices for this key
+        videos_idx[key]["chunk"] = chunk_idx
+        videos_idx[key]["file"] = file_idx
+
+        if update_latest_duration:
+            videos_idx[key]["latest_duration"] += timestamps_shift_s
+
+    return videos_idx
+
+
+def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size):
+    """Aggregates data chunks from a source dataset into the destination dataset.
+
+    Reads source data files, updates indices to match the aggregated dataset,
+    and writes them to the destination with proper file rotation.
+
+    Args:
+        src_meta: Source dataset metadata.
+        dst_meta: Destination dataset metadata.
+        data_idx: Dictionary tracking data chunk and file indices.
+
+    Returns:
+        dict: Updated data_idx with current chunk and file indices.
+    """
+    unique_chunk_file_ids = {
+        (c, f)
+        for c, f in zip(
+            src_meta.episodes["data/chunk_index"], src_meta.episodes["data/file_index"], strict=False
+        )
+    }
+
+    unique_chunk_file_ids = sorted(unique_chunk_file_ids)
+
+    for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
+        src_path = src_meta.root / DEFAULT_DATA_PATH.format(
+            chunk_index=src_chunk_idx, file_index=src_file_idx
+        )
+        df = pd.read_parquet(src_path)
+        df = update_data_df(df, src_meta, dst_meta)
+
+        data_idx = append_or_create_parquet_file(
+            df,
+            src_path,
+            data_idx,
+            data_files_size_in_mb,
+            chunk_size,
+            DEFAULT_DATA_PATH,
+            contains_images=len(dst_meta.image_keys) > 0,
+            aggr_root=dst_meta.root,
+        )
+
+    return data_idx
+
+
+def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
+    """Aggregates metadata from a source dataset into the destination dataset.
+
+    Reads source metadata files, updates all indices and timestamps,
+    and writes them to the destination with proper file rotation.
+
+    Args:
+        src_meta: Source dataset metadata.
+        dst_meta: Destination dataset metadata.
+        meta_idx: Dictionary tracking metadata chunk and file indices.
+        data_idx: Dictionary tracking data chunk and file indices.
+        videos_idx: Dictionary tracking video indices and timestamps.
+
+    Returns:
+        dict: Updated meta_idx with current chunk and file indices.
+    """
+    chunk_file_ids = {
+        (c, f)
+        for c, f in zip(
+            src_meta.episodes["meta/episodes/chunk_index"],
+            src_meta.episodes["meta/episodes/file_index"],
+            strict=False,
+        )
+    }
+
+    chunk_file_ids = sorted(chunk_file_ids)
+    for chunk_idx, file_idx in chunk_file_ids:
+        src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
+        df = pd.read_parquet(src_path)
+        df = update_meta_data(
+            df,
+            dst_meta,
+            meta_idx,
+            data_idx,
+            videos_idx,
+        )
+
+        for k in videos_idx:
+            videos_idx[k]["latest_duration"] += videos_idx[k]["episode_duration"]
+
+        meta_idx = append_or_create_parquet_file(
+            df,
+            src_path,
+            meta_idx,
+            DEFAULT_DATA_FILE_SIZE_IN_MB,
+            DEFAULT_CHUNK_SIZE,
+            DEFAULT_EPISODES_PATH,
+            contains_images=False,
+            aggr_root=dst_meta.root,
+        )
+
+    return meta_idx
+
+
+def append_or_create_parquet_file(
+    df: pd.DataFrame,
+    src_path: Path,
+    idx: dict[str, int],
+    max_mb: float,
+    chunk_size: int,
+    default_path: str,
+    contains_images: bool = False,
+    aggr_root: Path = None,
+):
+    """Appends data to an existing parquet file or creates a new one based on size constraints.
+
+    Manages file rotation when size limits are exceeded to prevent individual files
+    from becoming too large. Handles both regular parquet files and those containing images.
+
+    Args:
+        df: DataFrame to write to the parquet file.
+        src_path: Path to the source file (used for size estimation).
+        idx: Dictionary containing current 'chunk' and 'file' indices.
+        max_mb: Maximum allowed file size in MB before rotation.
+        chunk_size: Maximum number of files per chunk before incrementing chunk index.
+        default_path: Format string for generating file paths.
+        contains_images: Whether the data contains images requiring special handling.
+        aggr_root: Root path for the aggregated dataset.
+
+    Returns:
+        dict: Updated index dictionary with current chunk and file indices.
+    """
+    dst_path = aggr_root / default_path.format(chunk_index=idx["chunk"], file_index=idx["file"])
+
+    if not dst_path.exists():
+        dst_path.parent.mkdir(parents=True, exist_ok=True)
+        if contains_images:
+            to_parquet_with_hf_images(df, dst_path)
+        else:
+            df.to_parquet(dst_path)
+        return idx
+
+    src_size = get_parquet_file_size_in_mb(src_path)
+    dst_size = get_parquet_file_size_in_mb(dst_path)
+
+    if dst_size + src_size >= max_mb:
+        idx["chunk"], idx["file"] = update_chunk_file_indices(idx["chunk"], idx["file"], chunk_size)
+        new_path = aggr_root / default_path.format(chunk_index=idx["chunk"], file_index=idx["file"])
+        new_path.parent.mkdir(parents=True, exist_ok=True)
+        final_df = df
+        target_path = new_path
+    else:
+        existing_df = pd.read_parquet(dst_path)
+        final_df = pd.concat([existing_df, df], ignore_index=True)
+        target_path = dst_path
+
+    if contains_images:
+        to_parquet_with_hf_images(final_df, target_path)
+    else:
+        final_df.to_parquet(target_path)
+
+    return idx
+
+
+def finalize_aggregation(aggr_meta, all_metadata):
+    """Finalizes the dataset aggregation by writing summary files and statistics.
+
+    Writes the tasks file, info file with total counts and splits, and
+    aggregated statistics from all source datasets.
+
+    Args:
+        aggr_meta: Aggregated dataset metadata.
+        all_metadata: List of all source dataset metadata objects.
+    """
+    logging.info("write tasks")
+    write_tasks(aggr_meta.tasks, aggr_meta.root)
+
+    logging.info("write info")
+    aggr_meta.info.update(
+        {
+            "total_tasks": len(aggr_meta.tasks),
+            "total_episodes": sum(m.total_episodes for m in all_metadata),
+            "total_frames": sum(m.total_frames for m in all_metadata),
+            "splits": {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"},
+        }
+    )
+    write_info(aggr_meta.info, aggr_meta.root)
+
+    logging.info("write stats")
+    aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
+    write_stats(aggr_meta.stats, aggr_meta.root)

src/lerobot/datasets/backward_compatibility.py

@@ -47,6 +47,18 @@ If you encounter a problem, contact LeRobot maintainers on [Discord](https://dis
 or open an [issue on GitHub](https://github.com/huggingface/lerobot/issues/new/choose).
 """
 
+V30_MESSAGE = """
+The dataset you requested ({repo_id}) is in {version} format.
+While current version of LeRobot is backward-compatible with it, the version of your dataset still uses global
+stats instead of per-episode stats. Update your dataset stats to the new format using this command:
+```
+python -m lerobot.datasets.v30.convert_dataset_v21_to_v30 --repo-id={repo_id}
+```
+
+If you encounter a problem, contact LeRobot maintainers on [Discord](https://discord.com/invite/s3KuuzsPFb)
+or open an [issue on GitHub](https://github.com/huggingface/lerobot/issues/new/choose).
+"""
+
 FUTURE_MESSAGE = """
 The dataset you requested ({repo_id}) is only available in {version} format.
 As we cannot ensure forward compatibility with it, please update your current version of lerobot.
@@ -58,7 +70,14 @@ class CompatibilityError(Exception): ...
 
 class BackwardCompatibilityError(CompatibilityError):
     def __init__(self, repo_id: str, version: packaging.version.Version):
-        message = V2_MESSAGE.format(repo_id=repo_id, version=version)
+        if version.major == 3:
+            message = V30_MESSAGE.format(repo_id=repo_id, version=version)
+        elif version.major == 2:
+            message = V2_MESSAGE.format(repo_id=repo_id, version=version)
+        else:
+            raise NotImplementedError(
+                "Contact the maintainer on [Discord](https://discord.com/invite/s3KuuzsPFb)."
+            )
         super().__init__(message)
 
 

src/lerobot/datasets/dataset_tools.py

@@ -0,0 +1,761 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Dataset tools utilities for LeRobotDataset.
+
+This module provides utilities for:
+- Deleting episodes from datasets
+- Splitting datasets into multiple smaller datasets
+- Adding/removing features from datasets
+- Merging datasets (wrapper around aggregate functionality)
+"""
+
+import logging
+import shutil
+from collections.abc import Callable
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+from tqdm import tqdm
+
+from lerobot.constants import HF_LEROBOT_HOME
+from lerobot.datasets.aggregate import aggregate_datasets
+from lerobot.datasets.compute_stats import aggregate_stats
+from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.datasets.utils import (
+    DEFAULT_CHUNK_SIZE,
+    DEFAULT_DATA_FILE_SIZE_IN_MB,
+    DEFAULT_DATA_PATH,
+    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
+    DEFAULT_VIDEO_PATH,
+    get_parquet_file_size_in_mb,
+    get_video_size_in_mb,
+    to_parquet_with_hf_images,
+    update_chunk_file_indices,
+    write_info,
+    write_stats,
+    write_tasks,
+)
+
+
+def delete_episodes(
+    dataset: LeRobotDataset,
+    episode_indices: list[int],
+    output_dir: str | Path | None = None,
+    repo_id: str | None = None,
+) -> LeRobotDataset:
+    """Delete episodes from a LeRobotDataset and create a new dataset.
+
+    Args:
+        dataset: The source LeRobotDataset.
+        episode_indices: List of episode indices to delete.
+        output_dir: Directory to save the new dataset. If None, uses default location.
+        repo_id: Repository ID for the new dataset. If None, appends "_filtered" to original.
+
+    Returns:
+        LeRobotDataset: New dataset with episodes removed.
+    """
+    if not episode_indices:
+        raise ValueError("No episodes to delete")
+
+    # Validate episode indices
+    valid_indices = set(range(dataset.meta.total_episodes))
+    invalid = set(episode_indices) - valid_indices
+    if invalid:
+        raise ValueError(f"Invalid episode indices: {invalid}")
+
+    logging.info(f"Deleting {len(episode_indices)} episodes from dataset")
+
+    # Create new dataset metadata
+    if repo_id is None:
+        repo_id = f"{dataset.repo_id}_filtered"
+    if output_dir is None:
+        output_dir = HF_LEROBOT_HOME / repo_id
+    else:
+        output_dir = Path(output_dir)
+
+    # Get episodes to keep
+    episodes_to_keep = [i for i in range(dataset.meta.total_episodes) if i not in episode_indices]
+    if not episodes_to_keep:
+        raise ValueError("Cannot delete all episodes from dataset")
+
+    # Create new dataset
+    new_meta = LeRobotDatasetMetadata.create(
+        repo_id=repo_id,
+        fps=dataset.meta.fps,
+        features=dataset.meta.features,
+        robot_type=dataset.meta.robot_type,
+        root=output_dir,
+        use_videos=len(dataset.meta.video_keys) > 0,
+    )
+
+    # Process episodes
+    episode_mapping = {}  # old_idx -> new_idx
+    new_episode_idx = 0
+
+    for old_idx in tqdm(episodes_to_keep, desc="Processing episodes"):
+        episode_mapping[old_idx] = new_episode_idx
+        new_episode_idx += 1
+
+    # Copy data files and update indices
+    _copy_and_reindex_data(dataset, new_meta, episode_mapping)
+
+    # Copy video files if present
+    if dataset.meta.video_keys:
+        _copy_and_reindex_videos(dataset, new_meta, episode_mapping)
+
+    # Create new dataset instance
+    new_dataset = LeRobotDataset(
+        repo_id=repo_id,
+        root=output_dir,
+        image_transforms=dataset.image_transforms,
+        delta_timestamps=dataset.delta_timestamps,
+        tolerance_s=dataset.tolerance_s,
+    )
+
+    logging.info(f"Created new dataset with {len(episodes_to_keep)} episodes")
+    return new_dataset
+
+
+def split_dataset(
+    dataset: LeRobotDataset,
+    splits: dict[str, list[int]] | dict[str, float],
+    output_dir: str | Path | None = None,
+) -> dict[str, LeRobotDataset]:
+    """Split a LeRobotDataset into multiple smaller datasets.
+
+    Args:
+        dataset: The source LeRobotDataset to split.
+        splits: Either a dict mapping split names to episode indices, or a dict mapping
+                split names to fractions (must sum to <= 1.0).
+        output_dir: Base directory for output datasets. If None, uses default location.
+
+    Returns:
+        dict[str, LeRobotDataset]: Dictionary mapping split names to new datasets.
+
+    Examples:
+        # Split by specific episodes
+        splits = {"train": [0, 1, 2], "val": [3, 4]}
+        datasets = split_dataset(dataset, splits)
+
+        # Split by fractions
+        splits = {"train": 0.8, "val": 0.2}
+        datasets = split_dataset(dataset, splits)
+    """
+    if not splits:
+        raise ValueError("No splits provided")
+
+    # Convert fractions to episode indices if needed
+    if all(isinstance(v, float) for v in splits.values()):
+        splits = _fractions_to_episode_indices(dataset.meta.total_episodes, splits)
+
+    # Validate episodes
+    all_episodes = set()
+    for split_name, episodes in splits.items():
+        if not episodes:
+            raise ValueError(f"Split '{split_name}' has no episodes")
+        episode_set = set(episodes)
+        if episode_set & all_episodes:
+            raise ValueError("Episodes cannot appear in multiple splits")
+        all_episodes.update(episode_set)
+
+    # Validate all episodes are valid
+    valid_indices = set(range(dataset.meta.total_episodes))
+    invalid = all_episodes - valid_indices
+    if invalid:
+        raise ValueError(f"Invalid episode indices: {invalid}")
+
+    if output_dir is None:
+        output_dir = HF_LEROBOT_HOME
+    else:
+        output_dir = Path(output_dir)
+
+    result_datasets = {}
+
+    for split_name, episodes in splits.items():
+        logging.info(f"Creating split '{split_name}' with {len(episodes)} episodes")
+
+        # Create repo_id for split
+        split_repo_id = f"{dataset.repo_id}_{split_name}"
+        split_output_dir = output_dir / split_repo_id
+
+        # Create episode mapping
+        episode_mapping = {old_idx: new_idx for new_idx, old_idx in enumerate(sorted(episodes))}
+
+        # Create new dataset metadata
+        new_meta = LeRobotDatasetMetadata.create(
+            repo_id=split_repo_id,
+            fps=dataset.meta.fps,
+            features=dataset.meta.features,
+            robot_type=dataset.meta.robot_type,
+            root=split_output_dir,
+            use_videos=len(dataset.meta.video_keys) > 0,
+        )
+
+        # Copy data and videos
+        _copy_and_reindex_data(dataset, new_meta, episode_mapping)
+        if dataset.meta.video_keys:
+            _copy_and_reindex_videos(dataset, new_meta, episode_mapping)
+
+        # Create new dataset instance
+        new_dataset = LeRobotDataset(
+            repo_id=split_repo_id,
+            root=split_output_dir,
+            image_transforms=dataset.image_transforms,
+            delta_timestamps=dataset.delta_timestamps,
+            tolerance_s=dataset.tolerance_s,
+        )
+
+        result_datasets[split_name] = new_dataset
+
+    return result_datasets
+
+
+def merge_datasets(
+    datasets: list[LeRobotDataset],
+    output_repo_id: str,
+    output_dir: str | Path | None = None,
+) -> LeRobotDataset:
+    """Merge multiple LeRobotDatasets into a single dataset.
+
+    This is a wrapper around the aggregate_datasets functionality with a cleaner API.
+
+    Args:
+        datasets: List of LeRobotDatasets to merge.
+        output_repo_id: Repository ID for the merged dataset.
+        output_dir: Directory to save the merged dataset. If None, uses default location.
+
+    Returns:
+        LeRobotDataset: The merged dataset.
+    """
+    if not datasets:
+        raise ValueError("No datasets to merge")
+
+    if output_dir is None:
+        output_dir = HF_LEROBOT_HOME / output_repo_id
+    else:
+        output_dir = Path(output_dir)
+
+    # Extract repo_ids and roots
+    repo_ids = [ds.repo_id for ds in datasets]
+    roots = [ds.root for ds in datasets]
+
+    # Call aggregate_datasets
+    aggregate_datasets(
+        repo_ids=repo_ids,
+        aggr_repo_id=output_repo_id,
+        roots=roots,
+        aggr_root=output_dir,
+    )
+
+    # Create and return the merged dataset
+    merged_dataset = LeRobotDataset(
+        repo_id=output_repo_id,
+        root=output_dir,
+        image_transforms=datasets[0].image_transforms,
+        delta_timestamps=datasets[0].delta_timestamps,
+        tolerance_s=datasets[0].tolerance_s,
+    )
+
+    return merged_dataset
+
+
+def add_feature(
+    dataset: LeRobotDataset,
+    feature_name: str,
+    feature_values: np.ndarray | torch.Tensor | Callable,
+    feature_info: dict,
+    output_dir: str | Path | None = None,
+    repo_id: str | None = None,
+) -> LeRobotDataset:
+    """Add a new feature to a LeRobotDataset.
+
+    Args:
+        dataset: The source LeRobotDataset.
+        feature_name: Name of the new feature.
+        feature_values: Either:
+            - Array/tensor of shape (num_frames, ...) with values for each frame
+            - Callable that takes (frame_dict, episode_index, frame_index) and returns feature value
+        feature_info: Dictionary with feature metadata (dtype, shape, names).
+        output_dir: Directory to save the new dataset. If None, uses default location.
+        repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
+
+    Returns:
+        LeRobotDataset: New dataset with the added feature.
+    """
+    if feature_name in dataset.meta.features:
+        raise ValueError(f"Feature '{feature_name}' already exists in dataset")
+
+    if repo_id is None:
+        repo_id = f"{dataset.repo_id}_modified"
+    if output_dir is None:
+        output_dir = HF_LEROBOT_HOME / repo_id
+    else:
+        output_dir = Path(output_dir)
+
+    # Validate feature_info
+    required_keys = {"dtype", "shape"}
+    if not required_keys.issubset(feature_info.keys()):
+        raise ValueError(f"feature_info must contain keys: {required_keys}")
+
+    # Create new features dict
+    new_features = dataset.meta.features.copy()
+    new_features[feature_name] = feature_info
+
+    # Create new dataset metadata
+    new_meta = LeRobotDatasetMetadata.create(
+        repo_id=repo_id,
+        fps=dataset.meta.fps,
+        features=new_features,
+        robot_type=dataset.meta.robot_type,
+        root=output_dir,
+        use_videos=len(dataset.meta.video_keys) > 0,
+    )
+
+    # Process data with new feature
+    _copy_data_with_feature_changes(
+        dataset=dataset,
+        new_meta=new_meta,
+        add_features={feature_name: (feature_values, feature_info)},
+    )
+
+    # Copy videos if present
+    if dataset.meta.video_keys:
+        _copy_videos(dataset, new_meta)
+
+    # Create new dataset instance
+    new_dataset = LeRobotDataset(
+        repo_id=repo_id,
+        root=output_dir,
+        image_transforms=dataset.image_transforms,
+        delta_timestamps=dataset.delta_timestamps,
+        tolerance_s=dataset.tolerance_s,
+    )
+
+    return new_dataset
+
+
+def remove_feature(
+    dataset: LeRobotDataset,
+    feature_names: str | list[str],
+    output_dir: str | Path | None = None,
+    repo_id: str | None = None,
+) -> LeRobotDataset:
+    """Remove features from a LeRobotDataset.
+
+    Args:
+        dataset: The source LeRobotDataset.
+        feature_names: Name(s) of features to remove. Can be a single string or list.
+        output_dir: Directory to save the new dataset. If None, uses default location.
+        repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
+
+    Returns:
+        LeRobotDataset: New dataset with features removed.
+    """
+    if isinstance(feature_names, str):
+        feature_names = [feature_names]
+
+    # Validate features exist
+    for name in feature_names:
+        if name not in dataset.meta.features:
+            raise ValueError(f"Feature '{name}' not found in dataset")
+
+    # Check if trying to remove required features
+    required_features = {"timestamp", "frame_index", "episode_index", "index", "task_index"}
+    if any(name in required_features for name in feature_names):
+        raise ValueError(f"Cannot remove required features: {required_features}")
+
+    if repo_id is None:
+        repo_id = f"{dataset.repo_id}_modified"
+    if output_dir is None:
+        output_dir = HF_LEROBOT_HOME / repo_id
+    else:
+        output_dir = Path(output_dir)
+
+    # Create new features dict
+    new_features = {k: v for k, v in dataset.meta.features.items() if k not in feature_names}
+
+    # Check if removing video features
+    video_keys_to_remove = [name for name in feature_names if name in dataset.meta.video_keys]
+
+    # Check if videos will remain after removal
+    remaining_video_keys = [k for k in dataset.meta.video_keys if k not in video_keys_to_remove]
+
+    # Create new dataset metadata
+    new_meta = LeRobotDatasetMetadata.create(
+        repo_id=repo_id,
+        fps=dataset.meta.fps,
+        features=new_features,
+        robot_type=dataset.meta.robot_type,
+        root=output_dir,
+        use_videos=len(remaining_video_keys) > 0,
+    )
+
+    # Process data with removed features
+    _copy_data_with_feature_changes(
+        dataset=dataset,
+        new_meta=new_meta,
+        remove_features=feature_names,
+    )
+
+    # Copy videos (excluding removed ones)
+    if new_meta.video_keys:
+        _copy_videos(dataset, new_meta, exclude_keys=video_keys_to_remove)
+
+    # Create new dataset instance
+    new_dataset = LeRobotDataset(
+        repo_id=repo_id,
+        root=output_dir,
+        image_transforms=dataset.image_transforms,
+        delta_timestamps=dataset.delta_timestamps,
+        tolerance_s=dataset.tolerance_s,
+    )
+
+    return new_dataset
+
+
+# Helper functions
+
+
+def _fractions_to_episode_indices(
+    total_episodes: int,
+    splits: dict[str, float],
+) -> dict[str, list[int]]:
+    """Convert split fractions to episode indices."""
+    if sum(splits.values()) > 1.0:
+        raise ValueError("Split fractions must sum to <= 1.0")
+
+    indices = list(range(total_episodes))
+    result = {}
+    start_idx = 0
+
+    for split_name, fraction in splits.items():
+        num_episodes = int(total_episodes * fraction)
+        end_idx = start_idx + num_episodes
+        if split_name == list(splits.keys())[-1]:  # Last split gets remaining episodes
+            end_idx = total_episodes
+        result[split_name] = indices[start_idx:end_idx]
+        start_idx = end_idx
+
+    return result
+
+
+def _copy_and_reindex_data(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+    episode_mapping: dict[int, int],
+) -> None:
+    """Copy data files and reindex episodes."""
+    # Get unique data files from episodes to keep
+    file_paths = set()
+    for old_idx in episode_mapping:
+        file_paths.add(src_dataset.meta.get_data_file_path(old_idx))
+
+    # Track global index
+    global_index = 0
+    chunk_idx, file_idx = 0, 0
+
+    # Process each data file
+    for src_path in tqdm(sorted(file_paths), desc="Processing data files"):
+        df = pd.read_parquet(src_dataset.root / src_path)
+
+        # Filter to keep only mapped episodes
+        mask = df["episode_index"].isin(episode_mapping.keys())
+        df = df[mask].copy()
+
+        if len(df) == 0:
+            continue
+
+        # Update episode indices
+        df["episode_index"] = df["episode_index"].map(episode_mapping)
+
+        # Update global index to be continuous
+        df["index"] = range(global_index, global_index + len(df))
+        global_index += len(df)
+
+        # Update task indices if needed
+        if dst_meta.tasks is None:
+            # Get unique tasks from filtered data
+            task_indices = df["task_index"].unique()
+            tasks = [src_dataset.meta.tasks.iloc[idx].name for idx in task_indices]
+            dst_meta.save_episode_tasks(list(set(tasks)))
+
+        # Remap task indices
+        task_mapping = {}
+        for old_task_idx in df["task_index"].unique():
+            task_name = src_dataset.meta.tasks.iloc[old_task_idx].name
+            new_task_idx = dst_meta.get_task_index(task_name)
+            task_mapping[old_task_idx] = new_task_idx
+        df["task_index"] = df["task_index"].map(task_mapping)
+
+        # Save processed data
+        chunk_idx, file_idx = _save_data_chunk(df, dst_meta, chunk_idx, file_idx)
+
+    # Process episodes metadata
+    _copy_and_reindex_episodes_metadata(src_dataset, dst_meta, episode_mapping)
+
+
+def _copy_and_reindex_videos(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+    episode_mapping: dict[int, int],
+) -> None:
+    """Copy video files and update metadata."""
+    for video_key in src_dataset.meta.video_keys:
+        video_files = set()
+        for old_idx in episode_mapping:
+            video_files.add(src_dataset.meta.get_video_file_path(old_idx, video_key))
+
+        chunk_idx, file_idx = 0, 0
+
+        for src_path in tqdm(sorted(video_files), desc=f"Processing {video_key} videos"):
+            dst_path = dst_meta.root / DEFAULT_VIDEO_PATH.format(
+                video_key=video_key,
+                chunk_index=chunk_idx,
+                file_index=file_idx,
+            )
+            dst_path.parent.mkdir(parents=True, exist_ok=True)
+
+            # For simplicity, copy entire video files
+            # In production, you might want to extract only relevant segments
+            shutil.copy(src_dataset.root / src_path, dst_path)
+
+            # Update indices for next file
+            file_size = get_video_size_in_mb(dst_path)
+            if file_size >= DEFAULT_VIDEO_FILE_SIZE_IN_MB * 0.9:  # 90% threshold
+                chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
+
+
+def _copy_and_reindex_episodes_metadata(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+    episode_mapping: dict[int, int],
+) -> None:
+    """Copy and reindex episodes metadata."""
+    all_stats = []
+    frame_offset = 0
+
+    for old_idx, new_idx in tqdm(
+        sorted(episode_mapping.items(), key=lambda x: x[1]), desc="Processing episodes metadata"
+    ):
+        # Get episode from source
+        src_episode = src_dataset.meta.episodes[old_idx]
+
+        # Create episode dict
+        episode_dict = {
+            "episode_index": new_idx,
+            "tasks": src_episode["tasks"],  # Already a list of task names
+            "length": src_episode["length"],
+        }
+
+        # Copy other metadata
+        episode_metadata = {
+            "data/chunk_index": 0,  # Will be recalculated when saving
+            "data/file_index": 0,  # Will be recalculated when saving
+            "dataset_from_index": frame_offset,
+            "dataset_to_index": frame_offset + src_episode["length"],
+        }
+
+        # Update frame offset for next episode
+        frame_offset += src_episode["length"]
+
+        # Copy stats metadata
+        for key in src_episode.keys():
+            if key.startswith("stats/"):
+                episode_dict[key] = src_episode[key]
+
+        # Add episode metadata
+        stats_dict = {
+            key.replace("stats/", ""): value
+            for key, value in episode_dict.items()
+            if key.startswith("stats/")
+        }
+        all_stats.append(stats_dict)
+
+        # Calculate stats from dict
+        episode_stats = {}
+        for key in dst_meta.features:
+            if key in stats_dict:
+                episode_stats[key] = stats_dict[key]
+
+        dst_meta.save_episode(
+            new_idx, episode_dict["length"], episode_dict["tasks"], episode_stats, episode_metadata
+        )
+
+    # Aggregate all stats
+    if all_stats:
+        aggregated_stats = aggregate_stats(all_stats)
+        write_stats(aggregated_stats, dst_meta.root)
+
+
+def _save_data_chunk(
+    df: pd.DataFrame,
+    meta: LeRobotDatasetMetadata,
+    chunk_idx: int = 0,
+    file_idx: int = 0,
+) -> tuple[int, int]:
+    """Save a data chunk and return updated indices."""
+    path = meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    if len(meta.image_keys) > 0:
+        to_parquet_with_hf_images(df, path)
+    else:
+        df.to_parquet(path)
+
+    # Check if we need to rotate files
+    file_size = get_parquet_file_size_in_mb(path)
+    if file_size >= DEFAULT_DATA_FILE_SIZE_IN_MB * 0.9:  # 90% threshold
+        chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
+
+    return chunk_idx, file_idx
+
+
+def _copy_data_with_feature_changes(
+    dataset: LeRobotDataset,
+    new_meta: LeRobotDatasetMetadata,
+    add_features: dict[str, tuple] | None = None,
+    remove_features: list[str] | None = None,
+) -> None:
+    """Copy data while adding or removing features."""
+    # Get all unique data files
+    file_paths = set()
+    for ep_idx in range(dataset.meta.total_episodes):
+        file_paths.add(dataset.meta.get_data_file_path(ep_idx))
+
+    frame_idx = 0
+
+    # Process each data file
+    for src_path in tqdm(sorted(file_paths), desc="Processing data files"):
+        df = pd.read_parquet(dataset.root / src_path)
+
+        # Remove features
+        if remove_features:
+            df = df.drop(columns=remove_features, errors="ignore")
+
+        # Add features
+        if add_features:
+            for feature_name, (values, _) in add_features.items():
+                if callable(values):
+                    # Compute values for each frame
+                    feature_values = []
+                    for _, row in df.iterrows():
+                        ep_idx = row["episode_index"]
+                        frame_in_ep = row["frame_index"]
+                        value = values(row.to_dict(), ep_idx, frame_in_ep)
+                        # Convert numpy arrays to scalars for single-element arrays
+                        if isinstance(value, np.ndarray) and value.size == 1:
+                            value = value.item()
+                        feature_values.append(value)
+                    df[feature_name] = feature_values
+                else:
+                    # Use provided values
+                    end_idx = frame_idx + len(df)
+                    # Convert to list to ensure proper shape handling
+                    feature_slice = values[frame_idx:end_idx]
+                    if len(feature_slice.shape) > 1 and feature_slice.shape[1] == 1:
+                        # Flatten single-element arrays to scalars for pandas
+                        df[feature_name] = feature_slice.flatten()
+                    else:
+                        df[feature_name] = feature_slice
+                    frame_idx = end_idx
+
+        # Save chunk
+        _save_data_chunk(df, new_meta)
+
+    # Copy episodes metadata and update stats
+    _copy_episodes_metadata_and_stats(dataset, new_meta)
+
+
+def _copy_videos(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+    exclude_keys: list[str] | None = None,
+) -> None:
+    """Copy video files, optionally excluding certain keys."""
+    if exclude_keys is None:
+        exclude_keys = []
+
+    for video_key in src_dataset.meta.video_keys:
+        if video_key in exclude_keys:
+            continue
+
+        # Get all video files for this key
+        video_files = set()
+        for ep_idx in range(src_dataset.meta.total_episodes):
+            video_files.add(src_dataset.meta.get_video_file_path(ep_idx, video_key))
+
+        # Copy video files
+        for src_path in tqdm(sorted(video_files), desc=f"Copying {video_key} videos"):
+            # Maintain same structure
+            rel_path = src_path.relative_to(src_dataset.root)
+            dst_path = dst_meta.root / rel_path
+            dst_path.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy(src_dataset.root / src_path, dst_path)
+
+
+def _copy_episodes_metadata_and_stats(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+) -> None:
+    """Copy episodes metadata and recalculate stats."""
+    # Copy tasks
+    if src_dataset.meta.tasks is not None:
+        write_tasks(src_dataset.meta.tasks, dst_meta.root)
+        dst_meta.tasks = src_dataset.meta.tasks.copy()
+
+    # Copy episodes metadata files
+    episodes_dir = src_dataset.root / "meta/episodes"
+    dst_episodes_dir = dst_meta.root / "meta/episodes"
+    if episodes_dir.exists():
+        shutil.copytree(episodes_dir, dst_episodes_dir, dirs_exist_ok=True)
+
+    # Update info
+    dst_meta.info.update(
+        {
+            "total_episodes": src_dataset.meta.total_episodes,
+            "total_frames": src_dataset.meta.total_frames,
+            "total_tasks": src_dataset.meta.total_tasks,
+            "splits": src_dataset.meta.info.get("splits", {"train": f"0:{src_dataset.meta.total_episodes}"}),
+        }
+    )
+
+    # Update video info if needed
+    if dst_meta.video_keys and src_dataset.meta.video_keys:
+        for key in dst_meta.video_keys:
+            if key in src_dataset.meta.features:
+                dst_meta.info["features"][key]["info"] = src_dataset.meta.info["features"][key].get(
+                    "info", {}
+                )
+
+    write_info(dst_meta.info, dst_meta.root)
+
+    # Recalculate stats if features changed
+    if set(dst_meta.features.keys()) != set(src_dataset.meta.features.keys()):
+        # Need to recalculate stats
+        logging.info("Recalculating dataset statistics...")
+        # This is a simplified version - in production you'd want to properly recalculate
+        if src_dataset.meta.stats:
+            new_stats = {}
+            for key in dst_meta.features:
+                if key in src_dataset.meta.stats:
+                    new_stats[key] = src_dataset.meta.stats[key]
+            write_stats(new_stats, dst_meta.root)
+    else:
+        # Copy existing stats
+        if src_dataset.meta.stats:
+            write_stats(src_dataset.meta.stats, dst_meta.root)

src/lerobot/datasets/factory.py

@@ -13,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import logging
 from pprint import pformat
 
 import torch
@@ -86,24 +87,10 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
             cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
         )
         delta_timestamps = resolve_delta_timestamps(cfg.policy, ds_meta)
-
-        # Handle percentage parameter
-        episodes = cfg.dataset.episodes
-        if cfg.dataset.percentage is not None:
-            # Calculate episodes based on percentage
-            total_episodes = ds_meta.total_episodes
-            num_episodes_to_use = max(1, int(total_episodes * cfg.dataset.percentage / 100))
-            episodes = list(range(num_episodes_to_use))
-            import logging
-
-            logging.info(
-                f"Using {cfg.dataset.percentage}% of dataset: {num_episodes_to_use}/{total_episodes} episodes"
-            )
-
         dataset = LeRobotDataset(
             cfg.dataset.repo_id,
             root=cfg.dataset.root,
-            episodes=episodes,
+            episodes=cfg.dataset.episodes,
             delta_timestamps=delta_timestamps,
             image_transforms=image_transforms,
             revision=cfg.dataset.revision,

src/lerobot/datasets/lerobot_dataset.py

@@ -16,16 +16,18 @@
 import contextlib
 import logging
 import shutil
+import tempfile
 from collections.abc import Callable
 from pathlib import Path
 
 import datasets
 import numpy as np
 import packaging.version
+import pandas as pd
 import PIL.Image
 import torch
 import torch.utils
-from datasets import concatenate_datasets, load_dataset
+from datasets import Dataset
 from huggingface_hub import HfApi, snapshot_download
 from huggingface_hub.constants import REPOCARD_NAME
 from huggingface_hub.errors import RevisionNotFoundError
@@ -34,46 +36,51 @@ 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 (
+    DEFAULT_EPISODES_PATH,
     DEFAULT_FEATURES,
     DEFAULT_IMAGE_PATH,
     INFO_PATH,
-    TASKS_PATH,
     _validate_feature_names,
-    append_jsonlines,
-    backward_compatible_episodes_stats,
     check_delta_timestamps,
-    check_timestamps_sync,
     check_version_compatibility,
     create_empty_dataset_info,
     create_lerobot_dataset_card,
     embed_images,
+    flatten_dict,
     get_delta_indices,
-    get_episode_data_index,
+    get_hf_dataset_size_in_mb,
     get_hf_features_from_features,
+    get_parquet_file_size_in_mb,
+    get_parquet_num_frames,
     get_safe_version,
+    get_video_duration_in_s,
+    get_video_size_in_mb,
     hf_transform_to_torch,
     is_valid_version,
     load_episodes,
-    load_episodes_stats,
     load_info,
+    load_nested_dataset,
     load_stats,
     load_tasks,
+    to_parquet_with_hf_images,
+    update_chunk_file_indices,
     validate_episode_buffer,
     validate_frame,
-    write_episode,
-    write_episode_stats,
     write_info,
     write_json,
+    write_stats,
+    write_tasks,
 )
 from lerobot.datasets.video_utils import (
     VideoFrame,
+    concat_video_files,
     decode_video_frames,
     encode_video_frames,
     get_safe_default_codec,
     get_video_info,
 )
 
-CODEBASE_VERSION = "v2.1"
+CODEBASE_VERSION = "v3.0"
 
 
 class LeRobotDatasetMetadata:
@@ -97,20 +104,18 @@ class LeRobotDatasetMetadata:
                 self.revision = get_safe_version(self.repo_id, self.revision)
 
             (self.root / "meta").mkdir(exist_ok=True, parents=True)
+            # TODO(rcadene): instead of downloading all episodes metadata files,
+            # download only the ones associated to the requested episodes. This would
+            # require adding `episodes: list[int]` as argument.
             self.pull_from_repo(allow_patterns="meta/")
             self.load_metadata()
 
     def load_metadata(self):
         self.info = load_info(self.root)
         check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
-        self.tasks, self.task_to_task_index = load_tasks(self.root)
+        self.tasks = load_tasks(self.root)
         self.episodes = load_episodes(self.root)
-        if self._version < packaging.version.parse("v2.1"):
-            self.stats = load_stats(self.root)
-            self.episodes_stats = backward_compatible_episodes_stats(self.stats, self.episodes)
-        else:
-            self.episodes_stats = load_episodes_stats(self.root)
-            self.stats = aggregate_stats(list(self.episodes_stats.values()))
+        self.stats = load_stats(self.root)
 
     def pull_from_repo(
         self,
@@ -132,18 +137,19 @@ class LeRobotDatasetMetadata:
         return packaging.version.parse(self.info["codebase_version"])
 
     def get_data_file_path(self, ep_index: int) -> Path:
-        ep_chunk = self.get_episode_chunk(ep_index)
-        fpath = self.data_path.format(episode_chunk=ep_chunk, episode_index=ep_index)
+        ep = self.episodes[ep_index]
+        chunk_idx = ep["data/chunk_index"]
+        file_idx = ep["data/file_index"]
+        fpath = self.data_path.format(chunk_index=chunk_idx, file_index=file_idx)
         return Path(fpath)
 
     def get_video_file_path(self, ep_index: int, vid_key: str) -> Path:
-        ep_chunk = self.get_episode_chunk(ep_index)
-        fpath = self.video_path.format(episode_chunk=ep_chunk, video_key=vid_key, episode_index=ep_index)
+        ep = self.episodes[ep_index]
+        chunk_idx = ep[f"videos/{vid_key}/chunk_index"]
+        file_idx = ep[f"videos/{vid_key}/file_index"]
+        fpath = self.video_path.format(video_key=vid_key, chunk_index=chunk_idx, file_index=file_idx)
         return Path(fpath)
 
-    def get_episode_chunk(self, ep_index: int) -> int:
-        return ep_index // self.chunks_size
-
     @property
     def data_path(self) -> str:
         """Formattable string for the parquet files."""
@@ -210,39 +216,108 @@ class LeRobotDatasetMetadata:
         return self.info["total_tasks"]
 
     @property
-    def total_chunks(self) -> int:
-        """Total number of chunks (groups of episodes)."""
-        return self.info["total_chunks"]
+    def chunks_size(self) -> int:
+        """Max number of files per chunk."""
+        return self.info["chunks_size"]
 
     @property
-    def chunks_size(self) -> int:
-        """Max number of episodes per chunk."""
-        return self.info["chunks_size"]
+    def data_files_size_in_mb(self) -> int:
+        """Max size of data file in mega bytes."""
+        return self.info["data_files_size_in_mb"]
+
+    @property
+    def video_files_size_in_mb(self) -> int:
+        """Max size of video file in mega bytes."""
+        return self.info["video_files_size_in_mb"]
 
     def get_task_index(self, task: str) -> int | None:
         """
         Given a task in natural language, returns its task_index if the task already exists in the dataset,
         otherwise return None.
         """
-        return self.task_to_task_index.get(task, None)
+        if task in self.tasks.index:
+            return int(self.tasks.loc[task].task_index)
+        else:
+            return None
 
-    def add_task(self, task: str):
+    def save_episode_tasks(self, tasks: list[str]):
+        if len(set(tasks)) != len(tasks):
+            raise ValueError(f"Tasks are not unique: {tasks}")
+
+        if self.tasks is None:
+            new_tasks = tasks
+            task_indices = range(len(tasks))
+            self.tasks = pd.DataFrame({"task_index": task_indices}, index=tasks)
+        else:
+            new_tasks = [task for task in tasks if task not in self.tasks.index]
+            new_task_indices = range(len(self.tasks), len(self.tasks) + len(new_tasks))
+            for task_idx, task in zip(new_task_indices, new_tasks, strict=False):
+                self.tasks.loc[task] = task_idx
+
+        if len(new_tasks) > 0:
+            # Update on disk
+            write_tasks(self.tasks, self.root)
+
+    def _save_episode_metadata(self, episode_dict: dict) -> None:
+        """Save episode metadata to a parquet file and update the Hugging Face dataset of episodes metadata.
+
+        This function processes episodes metadata from a dictionary, converts it into a Hugging Face dataset,
+        and saves it as a parquet file. It handles both the creation of new parquet files and the
+        updating of existing ones based on size constraints. After saving the metadata, it reloads
+        the Hugging Face dataset to ensure it is up-to-date.
+
+        Notes: We both need to update parquet files and HF dataset:
+        - `pandas` loads parquet file in RAM
+        - `datasets` relies on a memory mapping from pyarrow (no RAM). It either converts parquet files to a pyarrow cache on disk,
+          or loads directly from pyarrow cache.
         """
-        Given a task in natural language, add it to the dictionary of tasks.
-        """
-        if task in self.task_to_task_index:
-            raise ValueError(f"The task '{task}' already exists and can't be added twice.")
+        # Convert buffer into HF Dataset
+        episode_dict = {key: [value] for key, value in episode_dict.items()}
+        ep_dataset = Dataset.from_dict(episode_dict)
+        ep_size_in_mb = get_hf_dataset_size_in_mb(ep_dataset)
+        df = pd.DataFrame(ep_dataset)
+        num_frames = episode_dict["length"][0]
 
-        task_index = self.info["total_tasks"]
-        self.task_to_task_index[task] = task_index
-        self.tasks[task_index] = task
-        self.info["total_tasks"] += 1
+        if self.episodes is None:
+            # Initialize indices and frame count for a new dataset made of the first episode data
+            chunk_idx, file_idx = 0, 0
+            df["meta/episodes/chunk_index"] = [chunk_idx]
+            df["meta/episodes/file_index"] = [file_idx]
+            df["dataset_from_index"] = [0]
+            df["dataset_to_index"] = [num_frames]
+        else:
+            # Retrieve information from the latest parquet file
+            latest_ep = self.episodes[-1]
+            chunk_idx = latest_ep["meta/episodes/chunk_index"]
+            file_idx = latest_ep["meta/episodes/file_index"]
 
-        task_dict = {
-            "task_index": task_index,
-            "task": task,
-        }
-        append_jsonlines(task_dict, self.root / TASKS_PATH)
+            latest_path = self.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
+            latest_size_in_mb = get_parquet_file_size_in_mb(latest_path)
+
+            if latest_size_in_mb + ep_size_in_mb >= self.data_files_size_in_mb:
+                # Size limit is reached, prepare new parquet file
+                chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, self.chunks_size)
+
+            # Update the existing pandas dataframe with new row
+            df["meta/episodes/chunk_index"] = [chunk_idx]
+            df["meta/episodes/file_index"] = [file_idx]
+            df["dataset_from_index"] = [latest_ep["dataset_to_index"]]
+            df["dataset_to_index"] = [latest_ep["dataset_to_index"] + num_frames]
+
+            if latest_size_in_mb + ep_size_in_mb < self.data_files_size_in_mb:
+                # Size limit wasnt reached, concatenate latest dataframe with new one
+                latest_df = pd.read_parquet(latest_path)
+                df = pd.concat([latest_df, df], ignore_index=True)
+
+        # Write the resulting dataframe from RAM to disk
+        path = self.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        df.to_parquet(path, index=False)
+
+        # Update the Hugging Face dataset by reloading it.
+        # This process should be fast because only the latest Parquet file has been modified.
+        # Therefore, only this file needs to be converted to PyArrow; the rest is loaded from the PyArrow memory-mapped cache.
+        self.episodes = load_episodes(self.root)
 
     def save_episode(
         self,
@@ -250,30 +325,28 @@ class LeRobotDatasetMetadata:
         episode_length: int,
         episode_tasks: list[str],
         episode_stats: dict[str, dict],
+        episode_metadata: dict,
     ) -> None:
-        self.info["total_episodes"] += 1
-        self.info["total_frames"] += episode_length
-
-        chunk = self.get_episode_chunk(episode_index)
-        if chunk >= self.total_chunks:
-            self.info["total_chunks"] += 1
-
-        self.info["splits"] = {"train": f"0:{self.info['total_episodes']}"}
-        self.info["total_videos"] += len(self.video_keys)
-
-        write_info(self.info, self.root)
-
         episode_dict = {
             "episode_index": episode_index,
             "tasks": episode_tasks,
             "length": episode_length,
         }
-        self.episodes[episode_index] = episode_dict
-        write_episode(episode_dict, self.root)
+        episode_dict.update(episode_metadata)
+        episode_dict.update(flatten_dict({"stats": episode_stats}))
+        self._save_episode_metadata(episode_dict)
 
-        self.episodes_stats[episode_index] = episode_stats
-        self.stats = aggregate_stats([self.stats, episode_stats]) if self.stats else episode_stats
-        write_episode_stats(episode_index, episode_stats, self.root)
+        # Update info
+        self.info["total_episodes"] += 1
+        self.info["total_frames"] += episode_length
+        self.info["total_tasks"] = len(self.tasks)
+        self.info["splits"] = {"train": f"0:{self.info['total_episodes']}"}
+        if len(self.video_keys) > 0:
+            self.update_video_info()
+        write_info(self.info, self.root)
+
+        self.stats = aggregate_stats([self.stats, episode_stats]) if self.stats is not None else episode_stats
+        write_stats(self.stats, self.root)
 
     def update_video_info(self) -> None:
         """
@@ -313,12 +386,12 @@ class LeRobotDatasetMetadata:
 
         obj.root.mkdir(parents=True, exist_ok=False)
 
-        # TODO(aliberts, rcadene): implement sanity check for features
         features = {**features, **DEFAULT_FEATURES}
         _validate_feature_names(features)
 
-        obj.tasks, obj.task_to_task_index = {}, {}
-        obj.episodes_stats, obj.stats, obj.episodes = {}, {}, {}
+        obj.tasks = None
+        obj.episodes = None
+        obj.stats = None
         obj.info = create_empty_dataset_info(CODEBASE_VERSION, fps, features, use_videos, robot_type)
         if len(obj.video_keys) > 0 and not use_videos:
             raise ValueError()
@@ -340,7 +413,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         force_cache_sync: bool = False,
         download_videos: bool = True,
         video_backend: str | None = None,
-        batch_encoding_size: int = 1,
     ):
         """
         2 modes are available for instantiating this class, depending on 2 different use cases:
@@ -354,9 +426,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
             - On the Hugging Face Hub at the address https://huggingface.co/datasets/{repo_id} and not on
               your local disk in the 'root' folder. Instantiating this class with this 'repo_id' will download
               the dataset from that address and load it, pending your dataset is compliant with
-              codebase_version v2.0. If your dataset has been created before this new format, you will be
-              prompted to convert it using our conversion script from v1.6 to v2.0, which you can find at
-              lerobot/datasets/v2/convert_dataset_v1_to_v2.py.
+              codebase_version v3.0. If your dataset has been created before this new format, you will be
+              prompted to convert it using our conversion script from v2.1 to v3.0, which you can find at
+              lerobot/datasets/v30/convert_dataset_v21_to_v30.py.
 
 
         2. Your dataset doesn't already exists (either on local disk or on the Hub): you can create an empty
@@ -377,38 +449,47 @@ class LeRobotDataset(torch.utils.data.Dataset):
         .
         ├── data
         │   ├── chunk-000
-        │   │   ├── episode_000000.parquet
-        │   │   ├── episode_000001.parquet
-        │   │   ├── episode_000002.parquet
+        │   │   ├── file-000.parquet
+        │   │   ├── file-001.parquet
         │   │   └── ...
         │   ├── chunk-001
-        │   │   ├── episode_001000.parquet
-        │   │   ├── episode_001001.parquet
-        │   │   ├── episode_001002.parquet
+        │   │   ├── file-000.parquet
+        │   │   ├── file-001.parquet
         │   │   └── ...
         │   └── ...
         ├── meta
-        │   ├── episodes.jsonl
+        │   ├── episodes
+        │   │   ├── chunk-000
+        │   │   │   ├── file-000.parquet
+        │   │   │   ├── file-001.parquet
+        │   │   │   └── ...
+        │   │   ├── chunk-001
+        │   │   │   └── ...
+        │   │   └── ...
         │   ├── info.json
         │   ├── stats.json
-        │   └── tasks.jsonl
+        │   └── tasks.parquet
         └── videos
-            ├── chunk-000
-            │   ├── observation.images.laptop
-            │   │   ├── episode_000000.mp4
-            │   │   ├── episode_000001.mp4
-            │   │   ├── episode_000002.mp4
+            ├── observation.images.laptop
+            │   ├── chunk-000
+            │   │   ├── file-000.mp4
+            │   │   ├── file-001.mp4
             │   │   └── ...
-            │   ├── observation.images.phone
-            │   │   ├── episode_000000.mp4
-            │   │   ├── episode_000001.mp4
-            │   │   ├── episode_000002.mp4
+            │   ├── chunk-001
             │   │   └── ...
-            ├── chunk-001
+            │   └── ...
+            ├── observation.images.phone
+            │   ├── chunk-000
+            │   │   ├── file-000.mp4
+            │   │   ├── file-001.mp4
+            │   │   └── ...
+            │   ├── chunk-001
+            │   │   └── ...
+            │   └── ...
             └── ...
 
-        Note that this file-based structure is designed to be as versatile as possible. The files are split by
-        episodes which allows a more granular control over which episodes one wants to use and download. The
+        Note that this file-based structure is designed to be as versatile as possible. Multiple episodes are
+        consolidated into chunked files which improves storage efficiency and loading performance. The
         structure of the dataset is entirely described in the info.json file, which can be easily downloaded
         or viewed directly on the hub before downloading any actual data. The type of files used are very
         simple and do not need complex tools to be read, it only uses .parquet, .json and .mp4 files (and .md
@@ -442,8 +523,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
                 True.
             video_backend (str | None, optional): Video backend to use for decoding videos. Defaults to torchcodec when available int the platform; otherwise, defaults to 'pyav'.
                 You can also use the 'pyav' decoder used by Torchvision, which used to be the default option, or 'video_reader' which is another decoder of Torchvision.
-            batch_encoding_size (int, optional): Number of episodes to accumulate before batch encoding videos.
-                Set to 1 for immediate encoding (default), or higher for batched encoding. Defaults to 1.
         """
         super().__init__()
         self.repo_id = repo_id
@@ -455,8 +534,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         self.revision = revision if revision else CODEBASE_VERSION
         self.video_backend = video_backend if video_backend else get_safe_default_codec()
         self.delta_indices = None
-        self.batch_encoding_size = batch_encoding_size
-        self.episodes_since_last_encoding = 0
 
         # Unused attributes
         self.image_writer = None
@@ -468,29 +545,20 @@ class LeRobotDataset(torch.utils.data.Dataset):
         self.meta = LeRobotDatasetMetadata(
             self.repo_id, self.root, self.revision, force_cache_sync=force_cache_sync
         )
-        if self.episodes is not None and self.meta._version >= packaging.version.parse("v2.1"):
-            episodes_stats = [self.meta.episodes_stats[ep_idx] for ep_idx in self.episodes]
-            self.stats = aggregate_stats(episodes_stats)
 
         # Load actual data
         try:
             if force_cache_sync:
                 raise FileNotFoundError
-            assert all((self.root / fpath).is_file() for fpath in self.get_episodes_file_paths())
             self.hf_dataset = self.load_hf_dataset()
+            # Check if cached dataset contains all requested episodes
+            if not self._check_cached_episodes_sufficient():
+                raise FileNotFoundError("Cached dataset doesn't contain all requested episodes")
         except (AssertionError, FileNotFoundError, NotADirectoryError):
             self.revision = get_safe_version(self.repo_id, self.revision)
-            self.download_episodes(download_videos)
+            self.download(download_videos)
             self.hf_dataset = self.load_hf_dataset()
 
-        self.episode_data_index = get_episode_data_index(self.meta.episodes, self.episodes)
-
-        # Check timestamps
-        timestamps = torch.stack(self.hf_dataset["timestamp"]).numpy()
-        episode_indices = torch.stack(self.hf_dataset["episode_index"]).numpy()
-        ep_data_index_np = {k: t.numpy() for k, t in self.episode_data_index.items()}
-        check_timestamps_sync(timestamps, episode_indices, ep_data_index_np, self.fps, self.tolerance_s)
-
         # Setup delta_indices
         if self.delta_timestamps is not None:
             check_delta_timestamps(self.delta_timestamps, self.fps, self.tolerance_s)
@@ -566,7 +634,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
             ignore_patterns=ignore_patterns,
         )
 
-    def download_episodes(self, download_videos: bool = True) -> None:
+    def download(self, download_videos: bool = True) -> None:
         """Downloads the dataset from the given 'repo_id' at the provided version. If 'episodes' is given, this
         will only download those episodes (selected by their episode_index). If 'episodes' is None, the whole
         dataset will be downloaded. Thanks to the behavior of snapshot_download, if the files are already present
@@ -574,11 +642,10 @@ class LeRobotDataset(torch.utils.data.Dataset):
         """
         # TODO(rcadene, aliberts): implement faster transfer
         # https://huggingface.co/docs/huggingface_hub/en/guides/download#faster-downloads
-        files = None
         ignore_patterns = None if download_videos else "videos/"
+        files = None
         if self.episodes is not None:
             files = self.get_episodes_file_paths()
-
         self.pull_from_repo(allow_patterns=files, ignore_patterns=ignore_patterns)
 
     def get_episodes_file_paths(self) -> list[Path]:
@@ -591,28 +658,40 @@ class LeRobotDataset(torch.utils.data.Dataset):
                 for ep_idx in episodes
             ]
             fpaths += video_files
-
+        # episodes are stored in the same files, so we return unique paths only
+        fpaths = list(set(fpaths))
         return fpaths
 
     def load_hf_dataset(self) -> datasets.Dataset:
         """hf_dataset contains all the observations, states, actions, rewards, etc."""
-        if self.episodes is None:
-            path = str(self.root / "data")
-            hf_dataset = load_dataset("parquet", data_dir=path, split="train")
-        else:
-            files = [str(self.root / self.meta.get_data_file_path(ep_idx)) for ep_idx in self.episodes]
-            hf_dataset = load_dataset("parquet", data_files=files, split="train")
-
-        # TODO(aliberts): hf_dataset.set_format("torch")
+        features = get_hf_features_from_features(self.features)
+        hf_dataset = load_nested_dataset(self.root / "data", features=features)
         hf_dataset.set_transform(hf_transform_to_torch)
         return hf_dataset
 
+    def _check_cached_episodes_sufficient(self) -> bool:
+        """Check if the cached dataset contains all requested episodes."""
+        if self.hf_dataset is None or len(self.hf_dataset) == 0:
+            return False
+
+        # Get available episode indices from cached dataset
+        available_episodes = set(self.hf_dataset["episode_index"])
+
+        # Determine requested episodes
+        if self.episodes is None:
+            # Requesting all episodes - check if we have all episodes from metadata
+            requested_episodes = set(range(self.meta.total_episodes))
+        else:
+            # Requesting specific episodes
+            requested_episodes = set(self.episodes)
+
+        # Check if all requested episodes are available in cached data
+        return requested_episodes.issubset(available_episodes)
+
     def create_hf_dataset(self) -> datasets.Dataset:
         features = get_hf_features_from_features(self.features)
         ft_dict = {col: [] for col in features}
         hf_dataset = datasets.Dataset.from_dict(ft_dict, features=features, split="train")
-
-        # TODO(aliberts): hf_dataset.set_format("torch")
         hf_dataset.set_transform(hf_transform_to_torch)
         return hf_dataset
 
@@ -644,15 +723,16 @@ class LeRobotDataset(torch.utils.data.Dataset):
             return get_hf_features_from_features(self.features)
 
     def _get_query_indices(self, idx: int, ep_idx: int) -> tuple[dict[str, list[int | bool]]]:
-        ep_start = self.episode_data_index["from"][ep_idx]
-        ep_end = self.episode_data_index["to"][ep_idx]
+        ep = self.meta.episodes[ep_idx]
+        ep_start = ep["dataset_from_index"]
+        ep_end = ep["dataset_to_index"]
         query_indices = {
-            key: [max(ep_start.item(), min(ep_end.item() - 1, idx + delta)) for delta in delta_idx]
+            key: [max(ep_start, min(ep_end - 1, idx + delta)) for delta in delta_idx]
             for key, delta_idx in self.delta_indices.items()
         }
         padding = {  # Pad values outside of current episode range
             f"{key}_is_pad": torch.BoolTensor(
-                [(idx + delta < ep_start.item()) | (idx + delta >= ep_end.item()) for delta in delta_idx]
+                [(idx + delta < ep_start) | (idx + delta >= ep_end) for delta in delta_idx]
             )
             for key, delta_idx in self.delta_indices.items()
         }
@@ -666,7 +746,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         query_timestamps = {}
         for key in self.meta.video_keys:
             if query_indices is not None and key in query_indices:
-                timestamps = self.hf_dataset.select(query_indices[key])["timestamp"]
+                timestamps = self.hf_dataset[query_indices[key]]["timestamp"]
                 query_timestamps[key] = torch.stack(timestamps).tolist()
             else:
                 query_timestamps[key] = [current_ts]
@@ -675,7 +755,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
     def _query_hf_dataset(self, query_indices: dict[str, list[int]]) -> dict:
         return {
-            key: torch.stack(self.hf_dataset.select(q_idx)[key])
+            key: torch.stack(self.hf_dataset[q_idx][key])
             for key, q_idx in query_indices.items()
             if key not in self.meta.video_keys
         }
@@ -686,10 +766,17 @@ class LeRobotDataset(torch.utils.data.Dataset):
         Segmentation Fault. This probably happens because a memory reference to the video loader is created in
         the main process and a subprocess fails to access it.
         """
+        ep = self.meta.episodes[ep_idx]
         item = {}
         for vid_key, query_ts in query_timestamps.items():
+            # Episodes are stored sequentially on a single mp4 to reduce the number of files.
+            # Thus we load the start timestamp of the episode on this mp4 and,
+            # shift the query timestamp accordingly.
+            from_timestamp = ep[f"videos/{vid_key}/from_timestamp"]
+            shifted_query_ts = [from_timestamp + ts for ts in query_ts]
+
             video_path = self.root / self.meta.get_video_file_path(ep_idx, vid_key)
-            frames = decode_video_frames(video_path, query_ts, self.tolerance_s, self.video_backend)
+            frames = decode_video_frames(video_path, shifted_query_ts, self.tolerance_s, self.video_backend)
             item[vid_key] = frames.squeeze(0)
 
         return item
@@ -727,8 +814,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
         # Add task as a string
         task_idx = item["task_index"].item()
-        item["task"] = self.meta.tasks[task_idx]
-
+        item["task"] = self.meta.tasks.iloc[task_idx].name
         return item
 
     def __repr__(self):
@@ -758,6 +844,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
         )
         return self.root / fpath
 
+    def _get_image_file_dir(self, episode_index: int, image_key: str) -> Path:
+        return self._get_image_file_path(episode_index, image_key, frame_index=0).parent
+
     def _save_image(self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path) -> None:
         if self.image_writer is None:
             if isinstance(image, torch.Tensor):
@@ -766,7 +855,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         else:
             self.image_writer.save_image(image=image, fpath=fpath)
 
-    def add_frame(self, frame: dict, task: str, timestamp: float | None = None) -> None:
+    def add_frame(self, frame: dict) -> None:
         """
         This function only adds the frame to the episode_buffer. Apart from images — which are written in a
         temporary directory — nothing is written to disk. To save those frames, the 'save_episode()' method
@@ -784,11 +873,10 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
         # Automatically add frame_index and timestamp to episode buffer
         frame_index = self.episode_buffer["size"]
-        if timestamp is None:
-            timestamp = frame_index / self.fps
+        timestamp = frame.pop("timestamp") if "timestamp" in frame else frame_index / self.fps
         self.episode_buffer["frame_index"].append(frame_index)
         self.episode_buffer["timestamp"].append(timestamp)
-        self.episode_buffer["task"].append(task)
+        self.episode_buffer["task"].append(frame.pop("task"))  # Remove task from frame after processing
 
         # Add frame features to episode_buffer
         for key in frame:
@@ -814,17 +902,12 @@ class LeRobotDataset(torch.utils.data.Dataset):
         """
         This will save to disk the current episode in self.episode_buffer.
 
-        Video encoding is handled automatically based on batch_encoding_size:
-        - If batch_encoding_size == 1: Videos are encoded immediately after each episode
-        - If batch_encoding_size > 1: Videos are encoded in batches.
-
         Args:
             episode_data (dict | None, optional): Dict containing the episode data to save. If None, this will
                 save the current episode in self.episode_buffer, which is filled with 'add_frame'. Defaults to
                 None.
         """
-        if not episode_data:
-            episode_buffer = self.episode_buffer
+        episode_buffer = episode_data if episode_data is not None else self.episode_buffer
 
         validate_episode_buffer(episode_buffer, self.meta.total_episodes, self.features)
 
@@ -837,11 +920,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
         episode_buffer["index"] = np.arange(self.meta.total_frames, self.meta.total_frames + episode_length)
         episode_buffer["episode_index"] = np.full((episode_length,), episode_index)
 
-        # Add new tasks to the tasks dictionary
-        for task in episode_tasks:
-            task_index = self.meta.get_task_index(task)
-            if task_index is None:
-                self.meta.add_task(task)
+        # Update tasks and task indices with new tasks if any
+        self.meta.save_episode_tasks(episode_tasks)
 
         # Given tasks in natural language, find their corresponding task indices
         episode_buffer["task_index"] = np.array([self.meta.get_task_index(task) for task in tasks])
@@ -853,72 +933,142 @@ class LeRobotDataset(torch.utils.data.Dataset):
                 continue
             episode_buffer[key] = np.stack(episode_buffer[key])
 
+        # Wait for image writer to end, so that episode stats over images can be computed
         self._wait_image_writer()
-        self._save_episode_table(episode_buffer, episode_index)
         ep_stats = compute_episode_stats(episode_buffer, self.features)
 
-        has_video_keys = len(self.meta.video_keys) > 0
-        use_batched_encoding = self.batch_encoding_size > 1
+        ep_metadata = self._save_episode_data(episode_buffer)
+        for video_key in self.meta.video_keys:
+            ep_metadata.update(self._save_episode_video(video_key, episode_index))
 
-        if has_video_keys and not use_batched_encoding:
-            self.encode_episode_videos(episode_index)
+        # `meta.save_episode` need to be executed after encoding the videos
+        self.meta.save_episode(episode_index, episode_length, episode_tasks, ep_stats, ep_metadata)
 
-        # `meta.save_episode` should be executed after encoding the videos
-        self.meta.save_episode(episode_index, episode_length, episode_tasks, ep_stats)
+        if not episode_data:
+            # Reset episode buffer and clean up temporary images
+            self.clear_episode_buffer()
 
-        # Check if we should trigger batch encoding
-        if has_video_keys and use_batched_encoding:
-            self.episodes_since_last_encoding += 1
-            if self.episodes_since_last_encoding == self.batch_encoding_size:
-                start_ep = self.num_episodes - self.batch_encoding_size
-                end_ep = self.num_episodes
-                logging.info(
-                    f"Batch encoding {self.batch_encoding_size} videos for episodes {start_ep} to {end_ep - 1}"
-                )
-                self.batch_encode_videos(start_ep, end_ep)
-                self.episodes_since_last_encoding = 0
+    def _save_episode_data(self, episode_buffer: dict) -> dict:
+        """Save episode data to a parquet file and update the Hugging Face dataset of frames data.
 
-        # Episode data index and timestamp checking
-        ep_data_index = get_episode_data_index(self.meta.episodes, [episode_index])
-        ep_data_index_np = {k: t.numpy() for k, t in ep_data_index.items()}
-        check_timestamps_sync(
-            episode_buffer["timestamp"],
-            episode_buffer["episode_index"],
-            ep_data_index_np,
-            self.fps,
-            self.tolerance_s,
-        )
+        This function processes episodes data from a buffer, converts it into a Hugging Face dataset,
+        and saves it as a parquet file. It handles both the creation of new parquet files and the
+        updating of existing ones based on size constraints. After saving the data, it reloads
+        the Hugging Face dataset to ensure it is up-to-date.
 
-        # Verify that we have one parquet file per episode and the number of video files matches the number of encoded episodes
-        parquet_files = list(self.root.rglob("*.parquet"))
-        assert len(parquet_files) == self.num_episodes
-        video_files = list(self.root.rglob("*.mp4"))
-        assert len(video_files) == (self.num_episodes - self.episodes_since_last_encoding) * len(
-            self.meta.video_keys
-        )
-
-        if not episode_data:  # Reset the buffer
-            self.episode_buffer = self.create_episode_buffer()
-
-    def _save_episode_table(self, episode_buffer: dict, episode_index: int) -> None:
-        episode_dict = {key: episode_buffer[key] for key in self.hf_features}
-        ep_dataset = datasets.Dataset.from_dict(episode_dict, features=self.hf_features, split="train")
+        Notes: We both need to update parquet files and HF dataset:
+        - `pandas` loads parquet file in RAM
+        - `datasets` relies on a memory mapping from pyarrow (no RAM). It either converts parquet files to a pyarrow cache on disk,
+          or loads directly from pyarrow cache.
+        """
+        # Convert buffer into HF Dataset
+        ep_dict = {key: episode_buffer[key] for key in self.hf_features}
+        ep_dataset = datasets.Dataset.from_dict(ep_dict, features=self.hf_features, split="train")
         ep_dataset = embed_images(ep_dataset)
-        self.hf_dataset = concatenate_datasets([self.hf_dataset, ep_dataset])
-        self.hf_dataset.set_transform(hf_transform_to_torch)
-        ep_data_path = self.root / self.meta.get_data_file_path(ep_index=episode_index)
-        ep_data_path.parent.mkdir(parents=True, exist_ok=True)
-        ep_dataset.to_parquet(ep_data_path)
+        ep_size_in_mb = get_hf_dataset_size_in_mb(ep_dataset)
+        ep_num_frames = len(ep_dataset)
+        df = pd.DataFrame(ep_dataset)
+
+        if self.meta.episodes is None:
+            # Initialize indices and frame count for a new dataset made of the first episode data
+            chunk_idx, file_idx = 0, 0
+            latest_num_frames = 0
+        else:
+            # Retrieve information from the latest parquet file
+            latest_ep = self.meta.episodes[-1]
+            chunk_idx = latest_ep["data/chunk_index"]
+            file_idx = latest_ep["data/file_index"]
+
+            latest_path = self.root / self.meta.data_path.format(chunk_index=chunk_idx, file_index=file_idx)
+            latest_size_in_mb = get_parquet_file_size_in_mb(latest_path)
+            latest_num_frames = get_parquet_num_frames(latest_path)
+
+            # Determine if a new parquet file is needed
+            if latest_size_in_mb + ep_size_in_mb >= self.meta.data_files_size_in_mb:
+                # Size limit is reached, prepare new parquet file
+                chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, self.meta.chunks_size)
+                latest_num_frames = 0
+            else:
+                # Update the existing parquet file with new rows
+                latest_df = pd.read_parquet(latest_path)
+                df = pd.concat([latest_df, df], ignore_index=True)
+
+        # Write the resulting dataframe from RAM to disk
+        path = self.root / self.meta.data_path.format(chunk_index=chunk_idx, file_index=file_idx)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        if len(self.meta.image_keys) > 0:
+            to_parquet_with_hf_images(df, path)
+        else:
+            df.to_parquet(path)
+
+        # Update the Hugging Face dataset by reloading it.
+        # This process should be fast because only the latest Parquet file has been modified.
+        # Therefore, only this file needs to be converted to PyArrow; the rest is loaded from the PyArrow memory-mapped cache.
+        self.hf_dataset = self.load_hf_dataset()
+
+        metadata = {
+            "data/chunk_index": chunk_idx,
+            "data/file_index": file_idx,
+            "dataset_from_index": latest_num_frames,
+            "dataset_to_index": latest_num_frames + ep_num_frames,
+        }
+        return metadata
+
+    def _save_episode_video(self, video_key: str, episode_index: int):
+        # Encode episode frames into a temporary video
+        ep_path = self._encode_temporary_episode_video(video_key, episode_index)
+        ep_size_in_mb = get_video_size_in_mb(ep_path)
+        ep_duration_in_s = get_video_duration_in_s(ep_path)
+
+        if self.meta.episodes is None:
+            # Initialize indices for a new dataset made of the first episode data
+            chunk_idx, file_idx = 0, 0
+            latest_duration_in_s = 0
+            new_path = self.root / self.meta.video_path.format(
+                video_key=video_key, chunk_index=chunk_idx, file_index=file_idx
+            )
+            new_path.parent.mkdir(parents=True, exist_ok=True)
+            shutil.move(str(ep_path), str(new_path))
+        else:
+            # Retrieve information from the latest video file
+            latest_ep = self.meta.episodes[-1]
+            chunk_idx = latest_ep[f"videos/{video_key}/chunk_index"]
+            file_idx = latest_ep[f"videos/{video_key}/file_index"]
+
+            latest_path = self.root / self.meta.video_path.format(
+                video_key=video_key, chunk_index=chunk_idx, file_index=file_idx
+            )
+            latest_size_in_mb = get_video_size_in_mb(latest_path)
+            latest_duration_in_s = get_video_duration_in_s(latest_path)
+
+            if latest_size_in_mb + ep_size_in_mb >= self.meta.video_files_size_in_mb:
+                # Move temporary episode video to a new video file in the dataset
+                chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, self.meta.chunks_size)
+                new_path = self.root / self.meta.video_path.format(
+                    video_key=video_key, chunk_index=chunk_idx, file_index=file_idx
+                )
+                new_path.parent.mkdir(parents=True, exist_ok=True)
+                shutil.move(str(ep_path), str(new_path))
+            else:
+                # Update latest video file
+                concat_video_files([latest_path, ep_path], self.root, video_key, chunk_idx, file_idx)
+
+        # Remove temporary directory
+        shutil.rmtree(str(ep_path.parent))
+
+        metadata = {
+            "episode_index": episode_index,
+            f"videos/{video_key}/chunk_index": chunk_idx,
+            f"videos/{video_key}/file_index": file_idx,
+            f"videos/{video_key}/from_timestamp": latest_duration_in_s,
+            f"videos/{video_key}/to_timestamp": latest_duration_in_s + ep_duration_in_s,
+        }
+        return metadata
 
     def clear_episode_buffer(self) -> None:
-        episode_index = self.episode_buffer["episode_index"]
-
-        # Clean up image files for the current episode buffer
         if self.image_writer is not None:
             for cam_key in self.meta.camera_keys:
-                img_dir = self._get_image_file_path(
-                    episode_index=episode_index, image_key=cam_key, frame_index=0
-                ).parent
+                img_dir = self.root / "images" / cam_key
                 if img_dir.is_dir():
                     shutil.rmtree(img_dir)
 
@@ -939,7 +1089,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
     def stop_image_writer(self) -> None:
         """
         Whenever wrapping this dataset inside a parallelized DataLoader, this needs to be called first to
-        remove the image_writer in order for the LeRobotDataset object to be picklable and parallelized.
+        remove the image_writer in order for the LeRobotDataset object to be pickleable and parallelized.
         """
         if self.image_writer is not None:
             self.image_writer.stop()
@@ -950,55 +1100,16 @@ class LeRobotDataset(torch.utils.data.Dataset):
         if self.image_writer is not None:
             self.image_writer.wait_until_done()
 
-    def encode_episode_videos(self, episode_index: int) -> None:
+    def _encode_temporary_episode_video(self, video_key: str, episode_index: int) -> dict:
         """
         Use ffmpeg to convert frames stored as png into mp4 videos.
         Note: `encode_video_frames` is a blocking call. Making it asynchronous shouldn't speedup encoding,
         since video encoding with ffmpeg is already using multithreading.
-
-        This method handles video encoding steps:
-        - Video encoding via ffmpeg
-        - Video info updating in metadata
-        - Raw image cleanup
-
-        Args:
-            episode_index (int): Index of the episode to encode.
         """
-        for key in self.meta.video_keys:
-            video_path = self.root / self.meta.get_video_file_path(episode_index, key)
-            if video_path.is_file():
-                # Skip if video is already encoded. Could be the case when resuming data recording.
-                continue
-            img_dir = self._get_image_file_path(
-                episode_index=episode_index, image_key=key, frame_index=0
-            ).parent
-            encode_video_frames(img_dir, video_path, self.fps, overwrite=True)
-            shutil.rmtree(img_dir)
-
-        # Update video info (only needed when first episode is encoded since it reads from episode 0)
-        if len(self.meta.video_keys) > 0 and episode_index == 0:
-            self.meta.update_video_info()
-            write_info(self.meta.info, self.meta.root)  # ensure video info always written properly
-
-    def batch_encode_videos(self, start_episode: int = 0, end_episode: int | None = None) -> None:
-        """
-        Batch encode videos for multiple episodes.
-
-        Args:
-            start_episode: Starting episode index (inclusive)
-            end_episode: Ending episode index (exclusive). If None, encodes all episodes from start_episode
-        """
-        if end_episode is None:
-            end_episode = self.meta.total_episodes
-
-        logging.info(f"Starting batch video encoding for episodes {start_episode} to {end_episode - 1}")
-
-        # Encode all episodes with cleanup enabled for individual episodes
-        for ep_idx in range(start_episode, end_episode):
-            logging.info(f"Encoding videos for episode {ep_idx}")
-            self.encode_episode_videos(ep_idx)
-
-        logging.info("Batch video encoding completed")
+        temp_path = Path(tempfile.mkdtemp(dir=self.root)) / f"{video_key}_{episode_index:03d}.mp4"
+        img_dir = self._get_image_file_dir(episode_index, video_key)
+        encode_video_frames(img_dir, temp_path, self.fps, overwrite=True)
+        return temp_path
 
     @classmethod
     def create(
@@ -1013,7 +1124,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         image_writer_processes: int = 0,
         image_writer_threads: int = 0,
         video_backend: str | None = None,
-        batch_encoding_size: int = 1,
     ) -> "LeRobotDataset":
         """Create a LeRobot Dataset from scratch in order to record data."""
         obj = cls.__new__(cls)
@@ -1030,8 +1140,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         obj.revision = None
         obj.tolerance_s = tolerance_s
         obj.image_writer = None
-        obj.batch_encoding_size = batch_encoding_size
-        obj.episodes_since_last_encoding = 0
 
         if image_writer_processes or image_writer_threads:
             obj.start_image_writer(image_writer_processes, image_writer_threads)
@@ -1044,7 +1152,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         obj.image_transforms = None
         obj.delta_timestamps = None
         obj.delta_indices = None
-        obj.episode_data_index = None
         obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec()
         return obj
 

src/lerobot/datasets/online_buffer.py

@@ -337,13 +337,11 @@ def compute_sampler_weights(
     if len(offline_dataset) > 0:
         offline_data_mask_indices = []
         for start_index, end_index in zip(
-            offline_dataset.episode_data_index["from"],
-            offline_dataset.episode_data_index["to"],
+            offline_dataset.meta.episodes["dataset_from_index"],
+            offline_dataset.meta.episodes["dataset_to_index"],
             strict=True,
         ):
-            offline_data_mask_indices.extend(
-                range(start_index.item(), end_index.item() - offline_drop_n_last_frames)
-            )
+            offline_data_mask_indices.extend(range(start_index, end_index - offline_drop_n_last_frames))
         offline_data_mask = torch.zeros(len(offline_dataset), dtype=torch.bool)
         offline_data_mask[torch.tensor(offline_data_mask_indices)] = True
         weights.append(

src/lerobot/datasets/pipeline_features.py

@@ -1,94 +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.
-
-from collections.abc import Sequence
-from typing import Any
-
-from lerobot.datasets.utils import hw_to_dataset_features
-from lerobot.processor.pipeline import RobotProcessor
-
-
-def aggregate_pipeline_dataset_features(
-    pipeline: RobotProcessor,
-    initial_features: dict[str, Any],
-    *,
-    use_videos: bool = True,
-    patterns: Sequence[str] | None = None,
-) -> dict[str, dict]:
-    """
-    Aggregates the pipeline's features and returns a features dict ready for the dataset,
-    filtered to only those keys matching any of the given patterns (for action/state only).
-
-    - `initial_features`: raw camera specs, e.g. {"front": (h,w,c), ...}
-    - `use_videos`: whether to treat image features as video streams
-    - `patterns`: regexes to filter action & state features; images are included
-                  whenever use_videos=True, regardless of patterns.
-    """
-    import re
-
-    # Gather everything the pipeline features specifies, seeded with hardware cams:
-    all_features = pipeline.transform_features(initial_features)
-
-    # Helper to decide which action/state keys survive the `patterns` filter:
-    def keep(key: str) -> bool:
-        if patterns is None:
-            return True
-        return any(re.search(pat, key) for pat in patterns)
-
-    # Start with hardware dict, injecting initial cameras if videos are ON:
-    hw: dict[str, dict[str, Any]] = {}
-    if use_videos:
-        cams = {
-            name: shape
-            for name, shape in initial_features.items()
-            if isinstance(shape, tuple) and len(shape) == 3
-        }
-        if cams:
-            hw["observation"] = dict(cams)
-
-    # Go over every feature from the pipeline and merge:
-    for full_key, ty in all_features.items():
-        if full_key.startswith("action."):
-            # action.<feat>
-            if not keep(full_key):
-                continue
-            name = full_key[len("action.") :]
-            hw.setdefault("action", {})[name] = ty
-
-        elif full_key.startswith("observation.state."):
-            # observation.state.<feat>
-            if not keep(full_key):
-                continue
-            name = full_key[len("observation.state.") :]
-            hw.setdefault("observation", {})[name] = ty
-
-        elif full_key.startswith("observation.images."):
-            # observation.images.<cam>
-            # images obey ONLY the use_videos flag, not patterns
-            if not use_videos:
-                continue
-            name = full_key[len("observation.images.") :]
-            hw.setdefault("observation", {})[name] = ty
-
-        else:
-            # anything else (e.g. policy-only features) is ignored here
-            continue
-
-    out: dict[str, dict] = {}
-    if "action" in hw:
-        out.update(hw_to_dataset_features(hw["action"], "action", use_videos))
-    if "observation" in hw:
-        out.update(hw_to_dataset_features(hw["observation"], "observation", use_videos))
-
-    return out

src/lerobot/datasets/sampler.py

@@ -21,7 +21,8 @@ import torch
 class EpisodeAwareSampler:
     def __init__(
         self,
-        episode_data_index: dict,
+        dataset_from_indices: list[int],
+        dataset_to_indices: list[int],
         episode_indices_to_use: list | None = None,
         drop_n_first_frames: int = 0,
         drop_n_last_frames: int = 0,
@@ -30,7 +31,8 @@ class EpisodeAwareSampler:
         """Sampler that optionally incorporates episode boundary information.
 
         Args:
-            episode_data_index: Dictionary with keys 'from' and 'to' containing the start and end indices of each episode.
+            dataset_from_indices: List of indices containing the start of each episode in the dataset.
+            dataset_to_indices: List of indices containing the end of each episode in the dataset.
             episode_indices_to_use: List of episode indices to use. If None, all episodes are used.
                                     Assumes that episodes are indexed from 0 to N-1.
             drop_n_first_frames: Number of frames to drop from the start of each episode.
@@ -39,12 +41,10 @@ class EpisodeAwareSampler:
         """
         indices = []
         for episode_idx, (start_index, end_index) in enumerate(
-            zip(episode_data_index["from"], episode_data_index["to"], strict=True)
+            zip(dataset_from_indices, dataset_to_indices, strict=True)
         ):
             if episode_indices_to_use is None or episode_idx in episode_indices_to_use:
-                indices.extend(
-                    range(start_index.item() + drop_n_first_frames, end_index.item() - drop_n_last_frames)
-                )
+                indices.extend(range(start_index + drop_n_first_frames, end_index - drop_n_last_frames))
 
         self.indices = indices
         self.shuffle = shuffle

src/lerobot/datasets/utils.py

@@ -17,43 +17,59 @@ import contextlib
 import importlib.resources
 import json
 import logging
+import subprocess
 from collections.abc import Iterator
-from itertools import accumulate
 from pathlib import Path
 from pprint import pformat
 from types import SimpleNamespace
 from typing import Any
 
 import datasets
-import jsonlines
 import numpy as np
 import packaging.version
+import pandas
+import pandas as pd
+import pyarrow.parquet as pq
 import torch
+from datasets import Dataset, concatenate_datasets
 from datasets.table import embed_table_storage
 from huggingface_hub import DatasetCard, DatasetCardData, HfApi
 from huggingface_hub.errors import RevisionNotFoundError
 from PIL import Image as PILImage
 from torchvision import transforms
 
-from lerobot.configs.types import DictLike, FeatureType, PolicyFeature
+from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.datasets.backward_compatibility import (
     V21_MESSAGE,
     BackwardCompatibilityError,
     ForwardCompatibilityError,
 )
+from lerobot.robots import Robot
 from lerobot.utils.utils import is_valid_numpy_dtype_string
 
-DEFAULT_CHUNK_SIZE = 1000  # Max number of episodes per chunk
+DEFAULT_CHUNK_SIZE = 1000  # Max number of files per chunk
+DEFAULT_DATA_FILE_SIZE_IN_MB = 100  # Max size per file
+DEFAULT_VIDEO_FILE_SIZE_IN_MB = 500  # Max size per file
 
 INFO_PATH = "meta/info.json"
-EPISODES_PATH = "meta/episodes.jsonl"
 STATS_PATH = "meta/stats.json"
-EPISODES_STATS_PATH = "meta/episodes_stats.jsonl"
-TASKS_PATH = "meta/tasks.jsonl"
 
-DEFAULT_VIDEO_PATH = "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4"
-DEFAULT_PARQUET_PATH = "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet"
-DEFAULT_IMAGE_PATH = "images/{image_key}/episode_{episode_index:06d}/frame_{frame_index:06d}.png"
+EPISODES_DIR = "meta/episodes"
+DATA_DIR = "data"
+VIDEO_DIR = "videos"
+
+CHUNK_FILE_PATTERN = "chunk-{chunk_index:03d}/file-{file_index:03d}"
+DEFAULT_TASKS_PATH = "meta/tasks.parquet"
+DEFAULT_EPISODES_PATH = EPISODES_DIR + "/" + CHUNK_FILE_PATTERN + ".parquet"
+DEFAULT_DATA_PATH = DATA_DIR + "/" + CHUNK_FILE_PATTERN + ".parquet"
+DEFAULT_VIDEO_PATH = VIDEO_DIR + "/{video_key}/" + CHUNK_FILE_PATTERN + ".mp4"
+DEFAULT_IMAGE_PATH = "images/{image_key}/episode-{episode_index:06d}/frame-{frame_index:06d}.png"
+
+LEGACY_EPISODES_PATH = "meta/episodes.jsonl"
+LEGACY_EPISODES_STATS_PATH = "meta/episodes_stats.jsonl"
+LEGACY_TASKS_PATH = "meta/tasks.jsonl"
+LEGACY_DEFAULT_VIDEO_PATH = "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4"
+LEGACY_DEFAULT_PARQUET_PATH = "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet"
 
 DATASET_CARD_TEMPLATE = """
 ---
@@ -74,6 +90,79 @@ DEFAULT_FEATURES = {
 }
 
 
+def get_parquet_file_size_in_mb(parquet_path: str | Path) -> float:
+    metadata = pq.read_metadata(parquet_path)
+    total_uncompressed_size = 0
+    for row_group in range(metadata.num_row_groups):
+        rg_metadata = metadata.row_group(row_group)
+        for column in range(rg_metadata.num_columns):
+            col_metadata = rg_metadata.column(column)
+            total_uncompressed_size += col_metadata.total_uncompressed_size
+    return total_uncompressed_size / (1024**2)
+
+
+def get_hf_dataset_size_in_mb(hf_ds: Dataset) -> int:
+    return hf_ds.data.nbytes // (1024**2)
+
+
+def update_chunk_file_indices(chunk_idx: int, file_idx: int, chunks_size: int) -> tuple[int, int]:
+    if file_idx == chunks_size - 1:
+        file_idx = 0
+        chunk_idx += 1
+    else:
+        file_idx += 1
+    return chunk_idx, file_idx
+
+
+def load_nested_dataset(pq_dir: Path, features: datasets.Features | None = None) -> Dataset:
+    """Find parquet files in provided directory {pq_dir}/chunk-xxx/file-xxx.parquet
+    Convert parquet files to pyarrow memory mapped in a cache folder for efficient RAM usage
+    Concatenate all pyarrow references to return HF Dataset format
+
+    Args:
+        pq_dir: Directory containing parquet files
+        features: Optional features schema to ensure consistent loading of complex types like images
+    """
+    paths = sorted(pq_dir.glob("*/*.parquet"))
+    if len(paths) == 0:
+        raise FileNotFoundError(f"Provided directory does not contain any parquet file: {pq_dir}")
+
+    # TODO(rcadene): set num_proc to accelerate conversion to pyarrow
+    datasets = [Dataset.from_parquet(str(path), features=features) for path in paths]
+    return concatenate_datasets(datasets)
+
+
+def get_parquet_num_frames(parquet_path: str | Path) -> int:
+    metadata = pq.read_metadata(parquet_path)
+    return metadata.num_rows
+
+
+def get_video_size_in_mb(mp4_path: Path) -> float:
+    file_size_bytes = mp4_path.stat().st_size
+    file_size_mb = file_size_bytes / (1024**2)
+    return file_size_mb
+
+
+def get_video_duration_in_s(mp4_file: Path) -> float:
+    # TODO(rcadene): move to video_utils.py
+    command = [
+        "ffprobe",
+        "-v",
+        "error",
+        "-show_entries",
+        "format=duration",
+        "-of",
+        "default=noprint_wrappers=1:nokey=1",
+        str(mp4_file),
+    ]
+    result = subprocess.run(
+        command,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+    )
+    return float(result.stdout)
+
+
 def flatten_dict(d: dict, parent_key: str = "", sep: str = "/") -> dict:
     """Flatten a nested dictionary structure by collapsing nested keys into one key with a separator.
 
@@ -82,6 +171,7 @@ def flatten_dict(d: dict, parent_key: str = "", sep: str = "/") -> dict:
     >>> dct = {"a": {"b": 1, "c": {"d": 2}}, "e": 3}`
     >>> print(flatten_dict(dct))
     {"a/b": 1, "a/c/d": 2, "e": 3}
+    ```
     """
     items = []
     for k, v in d.items():
@@ -106,23 +196,13 @@ def unflatten_dict(d: dict, sep: str = "/") -> dict:
     return outdict
 
 
-def get_nested_item(obj: DictLike, flattened_key: str, sep: str = "/") -> Any:
-    split_keys = flattened_key.split(sep)
-    getter = obj[split_keys[0]]
-    if len(split_keys) == 1:
-        return getter
-
-    for key in split_keys[1:]:
-        getter = getter[key]
-
-    return getter
-
-
 def serialize_dict(stats: dict[str, torch.Tensor | np.ndarray | dict]) -> dict:
     serialized_dict = {}
     for key, value in flatten_dict(stats).items():
         if isinstance(value, (torch.Tensor, np.ndarray)):
             serialized_dict[key] = value.tolist()
+        elif isinstance(value, list) and isinstance(value[0], (int, float, list)):
+            serialized_dict[key] = value
         elif isinstance(value, np.generic):
             serialized_dict[key] = value.item()
         elif isinstance(value, (int, float)):
@@ -152,24 +232,7 @@ def write_json(data: dict, fpath: Path) -> None:
         json.dump(data, f, indent=4, ensure_ascii=False)
 
 
-def load_jsonlines(fpath: Path) -> list[Any]:
-    with jsonlines.open(fpath, "r") as reader:
-        return list(reader)
-
-
-def write_jsonlines(data: dict, fpath: Path) -> None:
-    fpath.parent.mkdir(exist_ok=True, parents=True)
-    with jsonlines.open(fpath, "w") as writer:
-        writer.write_all(data)
-
-
-def append_jsonlines(data: dict, fpath: Path) -> None:
-    fpath.parent.mkdir(exist_ok=True, parents=True)
-    with jsonlines.open(fpath, "a") as writer:
-        writer.write(data)
-
-
-def write_info(info: dict, local_dir: Path):
+def write_info(info: dict, local_dir: Path) -> None:
     write_json(info, local_dir / INFO_PATH)
 
 
@@ -180,65 +243,55 @@ def load_info(local_dir: Path) -> dict:
     return info
 
 
-def write_stats(stats: dict, local_dir: Path):
+def write_stats(stats: dict, local_dir: Path) -> None:
     serialized_stats = serialize_dict(stats)
     write_json(serialized_stats, local_dir / STATS_PATH)
 
 
-def cast_stats_to_numpy(stats) -> dict[str, dict[str, np.ndarray]]:
+def cast_stats_to_numpy(stats: dict) -> dict[str, dict[str, np.ndarray]]:
     stats = {key: np.array(value) for key, value in flatten_dict(stats).items()}
     return unflatten_dict(stats)
 
 
-def load_stats(local_dir: Path) -> dict[str, dict[str, np.ndarray]]:
+def load_stats(local_dir: Path) -> dict[str, dict[str, np.ndarray]] | None:
     if not (local_dir / STATS_PATH).exists():
         return None
     stats = load_json(local_dir / STATS_PATH)
     return cast_stats_to_numpy(stats)
 
 
-def write_task(task_index: int, task: dict, local_dir: Path):
-    task_dict = {
-        "task_index": task_index,
-        "task": task,
-    }
-    append_jsonlines(task_dict, local_dir / TASKS_PATH)
+def write_tasks(tasks: pandas.DataFrame, local_dir: Path) -> None:
+    path = local_dir / DEFAULT_TASKS_PATH
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tasks.to_parquet(path)
 
 
-def load_tasks(local_dir: Path) -> tuple[dict, dict]:
-    tasks = load_jsonlines(local_dir / TASKS_PATH)
-    tasks = {item["task_index"]: item["task"] for item in sorted(tasks, key=lambda x: x["task_index"])}
-    task_to_task_index = {task: task_index for task_index, task in tasks.items()}
-    return tasks, task_to_task_index
+def load_tasks(local_dir: Path) -> pandas.DataFrame:
+    tasks = pd.read_parquet(local_dir / DEFAULT_TASKS_PATH)
+    return tasks
 
 
-def write_episode(episode: dict, local_dir: Path):
-    append_jsonlines(episode, local_dir / EPISODES_PATH)
+def write_episodes(episodes: Dataset, local_dir: Path) -> None:
+    if get_hf_dataset_size_in_mb(episodes) > DEFAULT_DATA_FILE_SIZE_IN_MB:
+        raise NotImplementedError("Contact a maintainer.")
+
+    fpath = local_dir / DEFAULT_EPISODES_PATH.format(chunk_index=0, file_index=0)
+    fpath.parent.mkdir(parents=True, exist_ok=True)
+    episodes.to_parquet(fpath)
 
 
-def load_episodes(local_dir: Path) -> dict:
-    episodes = load_jsonlines(local_dir / EPISODES_PATH)
-    return {item["episode_index"]: item for item in sorted(episodes, key=lambda x: x["episode_index"])}
-
-
-def write_episode_stats(episode_index: int, episode_stats: dict, local_dir: Path):
-    # We wrap episode_stats in a dictionary since `episode_stats["episode_index"]`
-    # is a dictionary of stats and not an integer.
-    episode_stats = {"episode_index": episode_index, "stats": serialize_dict(episode_stats)}
-    append_jsonlines(episode_stats, local_dir / EPISODES_STATS_PATH)
-
-
-def load_episodes_stats(local_dir: Path) -> dict:
-    episodes_stats = load_jsonlines(local_dir / EPISODES_STATS_PATH)
-    return {
-        item["episode_index"]: cast_stats_to_numpy(item["stats"])
-        for item in sorted(episodes_stats, key=lambda x: x["episode_index"])
-    }
+def load_episodes(local_dir: Path) -> datasets.Dataset:
+    episodes = load_nested_dataset(local_dir / EPISODES_DIR)
+    # Select episode features/columns containing references to episode data and videos
+    # (e.g. tasks, dataset_from_index, dataset_to_index, data/chunk_index, data/file_index, etc.)
+    # This is to speedup access to these data, instead of having to load episode stats.
+    episodes = episodes.select_columns([key for key in episodes.features if not key.startswith("stats/")])
+    return episodes
 
 
 def backward_compatible_episodes_stats(
     stats: dict[str, dict[str, np.ndarray]], episodes: list[int]
-) -> dict[str, dict[str, np.ndarray]]:
+) -> dict[int, dict[str, dict[str, np.ndarray]]]:
     return dict.fromkeys(episodes, stats)
 
 
@@ -254,7 +307,7 @@ def load_image_as_numpy(
     return img_array
 
 
-def hf_transform_to_torch(items_dict: dict[torch.Tensor | None]):
+def hf_transform_to_torch(items_dict: dict[str, list[Any]]) -> dict[str, list[torch.Tensor | str]]:
     """Get a transform function that convert items from Hugging Face dataset (pyarrow)
     to torch tensors. Importantly, images are converted from PIL, which corresponds to
     a channel last representation (h w c) of uint8 type, to a torch image representation
@@ -439,6 +492,17 @@ def build_dataset_frame(
     return frame
 
 
+def get_features_from_robot(robot: Robot, use_videos: bool = True) -> dict:
+    # TODO(rcadene): add fps for each feature
+    camera_ft = {}
+    if robot.cameras:
+        camera_ft = {
+            key: {"dtype": "video" if use_videos else "image", **ft}
+            for key, ft in robot.camera_features.items()
+        }
+    return {**robot.motor_features, **camera_ft, **DEFAULT_FEATURES}
+
+
 def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFeature]:
     # TODO(aliberts): Implement "type" in dataset features and simplify this
     policy_features = {}
@@ -470,50 +534,6 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea
     return policy_features
 
 
-def merge_features(*dicts: dict) -> dict:
-    """
-    Merge LeRobot grouped feature dicts.
-
-    - For 1D numeric specs (dtype not image/video/string) with "names": we merge the names and recompute the shape.
-    - For others (observation.images.*), last one wins (if they are identical).
-    """
-    out: dict = {}
-    for d in dicts:
-        for key, value in d.items():
-            if not isinstance(value, dict):
-                out[key] = value
-                continue
-
-            dtype = value.get("dtype")
-            shape = value.get("shape")
-            is_vector = (
-                dtype not in ("image", "video", "string")
-                and isinstance(shape, tuple)
-                and len(shape) == 1
-                and "names" in value
-            )
-
-            if is_vector:
-                # Initialize or retrieve the accumulating dict for this feature key
-                target = out.setdefault(key, {"dtype": dtype, "names": [], "shape": (0,)})
-                # Ensure consistent data types across merged entries
-                if "dtype" in target and dtype != target["dtype"]:
-                    raise ValueError(f"dtype mismatch for '{key}': {target['dtype']} vs {dtype}")
-
-                # Merge feature names: append only new ones to preserve order without duplicates
-                seen = set(target["names"])
-                for n in value["names"]:
-                    if n not in seen:
-                        target["names"].append(n)
-                        seen.add(n)
-                # Recompute the shape to reflect the updated number of features
-                target["shape"] = (len(target["names"]),)
-            else:
-                # For images/videos and non-1D entries: override with the latest definition
-                out[key] = value
-    return out
-
-
 def create_empty_dataset_info(
     codebase_version: str,
     fps: int,
@@ -527,104 +547,17 @@ def create_empty_dataset_info(
         "total_episodes": 0,
         "total_frames": 0,
         "total_tasks": 0,
-        "total_videos": 0,
-        "total_chunks": 0,
         "chunks_size": DEFAULT_CHUNK_SIZE,
+        "data_files_size_in_mb": DEFAULT_DATA_FILE_SIZE_IN_MB,
+        "video_files_size_in_mb": DEFAULT_VIDEO_FILE_SIZE_IN_MB,
         "fps": fps,
         "splits": {},
-        "data_path": DEFAULT_PARQUET_PATH,
+        "data_path": DEFAULT_DATA_PATH,
         "video_path": DEFAULT_VIDEO_PATH if use_videos else None,
         "features": features,
     }
 
 
-def get_episode_data_index(
-    episode_dicts: dict[dict], episodes: list[int] | None = None
-) -> dict[str, torch.Tensor]:
-    episode_lengths = {ep_idx: ep_dict["length"] for ep_idx, ep_dict in episode_dicts.items()}
-    if episodes is not None:
-        episode_lengths = {ep_idx: episode_lengths[ep_idx] for ep_idx in episodes}
-
-    cumulative_lengths = list(accumulate(episode_lengths.values()))
-    return {
-        "from": torch.LongTensor([0] + cumulative_lengths[:-1]),
-        "to": torch.LongTensor(cumulative_lengths),
-    }
-
-
-def check_timestamps_sync(
-    timestamps: np.ndarray,
-    episode_indices: np.ndarray,
-    episode_data_index: dict[str, np.ndarray],
-    fps: int,
-    tolerance_s: float,
-    raise_value_error: bool = True,
-) -> bool:
-    """
-    This check is to make sure that each timestamp is separated from the next by (1/fps) +/- tolerance
-    to account for possible numerical error.
-
-    Args:
-        timestamps (np.ndarray): Array of timestamps in seconds.
-        episode_indices (np.ndarray): Array indicating the episode index for each timestamp.
-        episode_data_index (dict[str, np.ndarray]): A dictionary that includes 'to',
-            which identifies indices for the end of each episode.
-        fps (int): Frames per second. Used to check the expected difference between consecutive timestamps.
-        tolerance_s (float): Allowed deviation from the expected (1/fps) difference.
-        raise_value_error (bool): Whether to raise a ValueError if the check fails.
-
-    Returns:
-        bool: True if all checked timestamp differences lie within tolerance, False otherwise.
-
-    Raises:
-        ValueError: If the check fails and `raise_value_error` is True.
-    """
-    if timestamps.shape != episode_indices.shape:
-        raise ValueError(
-            "timestamps and episode_indices should have the same shape. "
-            f"Found {timestamps.shape=} and {episode_indices.shape=}."
-        )
-
-    # Consecutive differences
-    diffs = np.diff(timestamps)
-    within_tolerance = np.abs(diffs - (1.0 / fps)) <= tolerance_s
-
-    # Mask to ignore differences at the boundaries between episodes
-    mask = np.ones(len(diffs), dtype=bool)
-    ignored_diffs = episode_data_index["to"][:-1] - 1  # indices at the end of each episode
-    mask[ignored_diffs] = False
-    filtered_within_tolerance = within_tolerance[mask]
-
-    # Check if all remaining diffs are within tolerance
-    if not np.all(filtered_within_tolerance):
-        # Track original indices before masking
-        original_indices = np.arange(len(diffs))
-        filtered_indices = original_indices[mask]
-        outside_tolerance_filtered_indices = np.nonzero(~filtered_within_tolerance)[0]
-        outside_tolerance_indices = filtered_indices[outside_tolerance_filtered_indices]
-
-        outside_tolerances = []
-        for idx in outside_tolerance_indices:
-            entry = {
-                "timestamps": [timestamps[idx], timestamps[idx + 1]],
-                "diff": diffs[idx],
-                "episode_index": episode_indices[idx].item()
-                if hasattr(episode_indices[idx], "item")
-                else episode_indices[idx],
-            }
-            outside_tolerances.append(entry)
-
-        if raise_value_error:
-            raise ValueError(
-                f"""One or several timestamps unexpectedly violate the tolerance inside episode range.
-                This might be due to synchronization issues during data collection.
-                \n{pformat(outside_tolerances)}"""
-            )
-        return False
-
-    return True
-
-
 def check_delta_timestamps(
     delta_timestamps: dict[str, list[float]], fps: int, tolerance_s: float, raise_value_error: bool = True
 ) -> bool:
@@ -663,7 +596,7 @@ def get_delta_indices(delta_timestamps: dict[str, list[float]], fps: int) -> dic
     return delta_indices
 
 
-def cycle(iterable):
+def cycle(iterable: Any) -> Iterator[Any]:
     """The equivalent of itertools.cycle, but safe for Pytorch dataloaders.
 
     See https://github.com/pytorch/pytorch/issues/23900 for information on why itertools.cycle is not safe.
@@ -676,7 +609,7 @@ def cycle(iterable):
             iterator = iter(iterable)
 
 
-def create_branch(repo_id, *, branch: str, repo_type: str | None = None) -> None:
+def create_branch(repo_id: str, *, branch: str, repo_type: str | None = None) -> None:
     """Create a branch on a existing Hugging Face repo. Delete the branch if it already
     exists before creating it.
     """
@@ -697,7 +630,7 @@ def create_lerobot_dataset_card(
     **kwargs,
 ) -> DatasetCard:
     """
-    Keyword arguments will be used to replace values in src/lerobot/datasets/card_template.md.
+    Keyword arguments will be used to replace values in ./lerobot/datasets/card_template.md.
     Note: If specified, license must be one of https://huggingface.co/docs/hub/repositories-licenses.
     """
     card_tags = ["LeRobot"]
@@ -784,21 +717,28 @@ class IterableNamespace(SimpleNamespace):
         return vars(self).keys()
 
 
-def validate_frame(frame: dict, features: dict):
+def validate_frame(frame: dict, features: dict) -> None:
     expected_features = set(features) - set(DEFAULT_FEATURES)
     actual_features = set(frame)
 
-    error_message = validate_features_presence(actual_features, expected_features)
+    # task is a special required field that's not part of regular features
+    if "task" not in actual_features:
+        raise ValueError("Feature mismatch in `frame` dictionary:\nMissing features: {'task'}\n")
 
-    common_features = actual_features & expected_features
-    for name in common_features - {"task"}:
+    # Remove task from actual_features for regular feature validation
+    actual_features_for_validation = actual_features - {"task"}
+
+    error_message = validate_features_presence(actual_features_for_validation, expected_features)
+
+    common_features = actual_features_for_validation & expected_features
+    for name in common_features:
         error_message += validate_feature_dtype_and_shape(name, features[name], frame[name])
 
     if error_message:
         raise ValueError(error_message)
 
 
-def validate_features_presence(actual_features: set[str], expected_features: set[str]):
+def validate_features_presence(actual_features: set[str], expected_features: set[str]) -> str:
     error_message = ""
     missing_features = expected_features - actual_features
     extra_features = actual_features - expected_features
@@ -813,7 +753,9 @@ def validate_features_presence(actual_features: set[str], expected_features: set
     return error_message
 
 
-def validate_feature_dtype_and_shape(name: str, feature: dict, value: np.ndarray | PILImage.Image | str):
+def validate_feature_dtype_and_shape(
+    name: str, feature: dict, value: np.ndarray | PILImage.Image | str
+) -> str:
     expected_dtype = feature["dtype"]
     expected_shape = feature["shape"]
     if is_valid_numpy_dtype_string(expected_dtype):
@@ -828,7 +770,7 @@ def validate_feature_dtype_and_shape(name: str, feature: dict, value: np.ndarray
 
 def validate_feature_numpy_array(
     name: str, expected_dtype: str, expected_shape: list[int], value: np.ndarray
-):
+) -> str:
     error_message = ""
     if isinstance(value, np.ndarray):
         actual_dtype = value.dtype
@@ -845,7 +787,9 @@ def validate_feature_numpy_array(
     return error_message
 
 
-def validate_feature_image_or_video(name: str, expected_shape: list[str], value: np.ndarray | PILImage.Image):
+def validate_feature_image_or_video(
+    name: str, expected_shape: list[str], value: np.ndarray | PILImage.Image
+) -> str:
     # Note: The check of pixels range ([0,1] for float and [0,255] for uint8) is done by the image writer threads.
     error_message = ""
     if isinstance(value, np.ndarray):
@@ -861,13 +805,13 @@ def validate_feature_image_or_video(name: str, expected_shape: list[str], value:
     return error_message
 
 
-def validate_feature_string(name: str, value: str):
+def validate_feature_string(name: str, value: str) -> str:
     if not isinstance(value, str):
         return f"The feature '{name}' is expected to be of type 'str', but type '{type(value)}' provided instead.\n"
     return ""
 
 
-def validate_episode_buffer(episode_buffer: dict, total_episodes: int, features: dict):
+def validate_episode_buffer(episode_buffer: dict, total_episodes: int, features: dict) -> None:
     if "size" not in episode_buffer:
         raise ValueError("size key not found in episode_buffer")
 
@@ -891,3 +835,11 @@ def validate_episode_buffer(episode_buffer: dict, total_episodes: int, features:
             f"In episode_buffer not in features: {buffer_keys - set(features)}"
             f"In features not in episode_buffer: {set(features) - buffer_keys}"
         )
+
+
+def to_parquet_with_hf_images(df: pandas.DataFrame, path: Path) -> None:
+    """This function correctly writes to parquet a panda DataFrame that contains images encoded by HF dataset.
+    This way, it can be loaded by HF dataset and correctly formatted images are returned.
+    """
+    # TODO(qlhoest): replace this weird synthax by `df.to_parquet(path)` only
+    datasets.Dataset.from_dict(df.to_dict(orient="list")).to_parquet(path)

src/lerobot/datasets/v2/convert_dataset_v1_to_v2.py

@@ -121,12 +121,12 @@ from safetensors.torch import load_file
 
 from lerobot.datasets.utils import (
     DEFAULT_CHUNK_SIZE,
-    DEFAULT_PARQUET_PATH,
-    DEFAULT_VIDEO_PATH,
-    EPISODES_PATH,
     INFO_PATH,
+    LEGACY_DEFAULT_PARQUET_PATH,
+    LEGACY_DEFAULT_VIDEO_PATH,
+    LEGACY_EPISODES_PATH,
+    LEGACY_TASKS_PATH,
     STATS_PATH,
-    TASKS_PATH,
     create_branch,
     create_lerobot_dataset_card,
     flatten_dict,
@@ -290,12 +290,12 @@ def split_parquet_by_episodes(
     for ep_chunk in range(total_chunks):
         ep_chunk_start = DEFAULT_CHUNK_SIZE * ep_chunk
         ep_chunk_end = min(DEFAULT_CHUNK_SIZE * (ep_chunk + 1), total_episodes)
-        chunk_dir = "/".join(DEFAULT_PARQUET_PATH.split("/")[:-1]).format(episode_chunk=ep_chunk)
+        chunk_dir = "/".join(LEGACY_DEFAULT_PARQUET_PATH.split("/")[:-1]).format(episode_chunk=ep_chunk)
         (output_dir / chunk_dir).mkdir(parents=True, exist_ok=True)
         for ep_idx in range(ep_chunk_start, ep_chunk_end):
             ep_table = table.filter(pc.equal(table["episode_index"], ep_idx))
             episode_lengths.insert(ep_idx, len(ep_table))
-            output_file = output_dir / DEFAULT_PARQUET_PATH.format(
+            output_file = output_dir / LEGACY_DEFAULT_PARQUET_PATH.format(
                 episode_chunk=ep_chunk, episode_index=ep_idx
             )
             pq.write_table(ep_table, output_file)
@@ -344,13 +344,13 @@ def move_videos(
         ep_chunk_start = DEFAULT_CHUNK_SIZE * ep_chunk
         ep_chunk_end = min(DEFAULT_CHUNK_SIZE * (ep_chunk + 1), total_episodes)
         for vid_key in video_keys:
-            chunk_dir = "/".join(DEFAULT_VIDEO_PATH.split("/")[:-1]).format(
+            chunk_dir = "/".join(LEGACY_DEFAULT_VIDEO_PATH.split("/")[:-1]).format(
                 episode_chunk=ep_chunk, video_key=vid_key
             )
             (work_dir / chunk_dir).mkdir(parents=True, exist_ok=True)
 
             for ep_idx in range(ep_chunk_start, ep_chunk_end):
-                target_path = DEFAULT_VIDEO_PATH.format(
+                target_path = LEGACY_DEFAULT_VIDEO_PATH.format(
                     episode_chunk=ep_chunk, video_key=vid_key, episode_index=ep_idx
                 )
                 video_file = V1_VIDEO_FILE.format(video_key=vid_key, episode_index=ep_idx)
@@ -418,7 +418,7 @@ def _get_lfs_untracked_videos(work_dir: Path, video_files: list[str]) -> list[st
 def get_videos_info(repo_id: str, local_dir: Path, video_keys: list[str], branch: str) -> dict:
     # Assumes first episode
     video_files = [
-        DEFAULT_VIDEO_PATH.format(episode_chunk=0, video_key=vid_key, episode_index=0)
+        LEGACY_DEFAULT_VIDEO_PATH.format(episode_chunk=0, video_key=vid_key, episode_index=0)
         for vid_key in video_keys
     ]
     hub_api = HfApi()
@@ -495,7 +495,7 @@ def convert_dataset(
 
     assert set(tasks) == {task for ep_tasks in tasks_by_episodes.values() for task in ep_tasks}
     tasks = [{"task_index": task_idx, "task": task} for task_idx, task in enumerate(tasks)]
-    write_jsonlines(tasks, v20_dir / TASKS_PATH)
+    write_jsonlines(tasks, v20_dir / LEGACY_TASKS_PATH)
     features["task_index"] = {
         "dtype": "int64",
         "shape": (1,),
@@ -545,7 +545,7 @@ def convert_dataset(
         {"episode_index": ep_idx, "tasks": tasks_by_episodes[ep_idx], "length": episode_lengths[ep_idx]}
         for ep_idx in episode_indices
     ]
-    write_jsonlines(episodes, v20_dir / EPISODES_PATH)
+    write_jsonlines(episodes, v20_dir / LEGACY_EPISODES_PATH)
 
     # Assemble metadata v2.0
     metadata_v2_0 = {
@@ -559,8 +559,8 @@ def convert_dataset(
         "chunks_size": DEFAULT_CHUNK_SIZE,
         "fps": metadata_v1["fps"],
         "splits": {"train": f"0:{total_episodes}"},
-        "data_path": DEFAULT_PARQUET_PATH,
-        "video_path": DEFAULT_VIDEO_PATH if video_keys else None,
+        "data_path": LEGACY_DEFAULT_PARQUET_PATH,
+        "video_path": LEGACY_DEFAULT_VIDEO_PATH if video_keys else None,
         "features": features,
     }
     write_json(metadata_v2_0, v20_dir / INFO_PATH)

src/lerobot/datasets/v21/convert_dataset_v20_to_v21.py

@@ -13,24 +13,20 @@
 # limitations under the License.
 
 """
-This script converts a LeRobot dataset already pushed to the Hub from codebase version 2.0 to 2.1.
-It downloads metadata from a SOURCE dataset repo, computes/validates per-episode stats, updates
-the codebase version in `info.json`, and uploads the result to a DESTINATION dataset repo.
-It will:
+This script will help you convert any LeRobot dataset already pushed to the hub from codebase version 2.0 to
+2.1. It will:
 
 - Generate per-episodes stats and writes them in `episodes_stats.jsonl`
 - Check consistency between these new stats and the old ones.
 - Remove the deprecated `stats.json`.
 - Update codebase_version in `info.json`.
-- Push this new version to the destination repo/branch and tag it with the current codebase version.
+- Push this new version to the hub on the 'main' branch and tags it with "v2.1".
 
 Usage:
 
 ```bash
 python -m lerobot.datasets.v21.convert_dataset_v20_to_v21 \
-    --source-repo-id=namespace/source_dataset \
-    --dest-repo-id=namespace/destination_dataset \
-    --branch=main
+    --repo-id=aliberts/koch_tutorial
 ```
 
 """
@@ -41,7 +37,7 @@ import logging
 from huggingface_hub import HfApi
 
 from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
-from lerobot.datasets.utils import EPISODES_STATS_PATH, STATS_PATH, load_stats, write_info
+from lerobot.datasets.utils import STATS_PATH, load_stats, write_info
 from lerobot.datasets.v21.convert_stats import check_aggregate_stats, convert_stats
 
 V20 = "v2.0"
@@ -58,67 +54,45 @@ class SuppressWarnings:
 
 
 def convert_dataset(
-    source_repo_id: str,
-    dest_repo_id: str,
+    repo_id: str,
     branch: str | None = None,
     num_workers: int = 4,
 ):
-    # Download metadata from the source repo at v2.0
     with SuppressWarnings():
-        dataset = LeRobotDataset(source_repo_id, revision=V20, force_cache_sync=True)
+        dataset = LeRobotDataset(repo_id, revision=V20, force_cache_sync=True)
 
-    # Ensure we recompute fresh episodes stats
-    if (dataset.root / EPISODES_STATS_PATH).is_file():
-        (dataset.root / EPISODES_STATS_PATH).unlink()
-
-    # Compute and validate stats
     convert_stats(dataset, num_workers=num_workers)
     ref_stats = load_stats(dataset.root)
     check_aggregate_stats(dataset, ref_stats)
 
-    # Update codebase version in info.json
     dataset.meta.info["codebase_version"] = CODEBASE_VERSION
     write_info(dataset.meta.info, dataset.root)
 
-    # Remove deprecated stats.json locally so it won't be uploaded
-    if (dataset.root / STATS_PATH).is_file():
+    dataset.push_to_hub(branch=branch, tag_version=False, allow_patterns="meta/")
+
+    # delete old stats.json file
+    if (dataset.root / STATS_PATH).is_file:
         (dataset.root / STATS_PATH).unlink()
 
-    # Push only meta/ to destination repo
     hub_api = HfApi()
-    hub_api.create_repo(repo_id=dest_repo_id, private=False, repo_type="dataset", exist_ok=True)
-    if branch:
-        hub_api.create_branch(repo_id=dest_repo_id, branch=branch, repo_type="dataset", exist_ok=True)
+    if hub_api.file_exists(
+        repo_id=dataset.repo_id, filename=STATS_PATH, revision=branch, repo_type="dataset"
+    ):
+        hub_api.delete_file(
+            path_in_repo=STATS_PATH, repo_id=dataset.repo_id, revision=branch, repo_type="dataset"
+        )
 
-    hub_api.upload_folder(
-        repo_id=dest_repo_id,
-        folder_path=str(dataset.root),
-        repo_type="dataset",
-        revision=branch,
-        allow_patterns="meta/",
-    )
-
-    # Ensure old stats.json is deleted on destination
-    if hub_api.file_exists(repo_id=dest_repo_id, filename=STATS_PATH, revision=branch, repo_type="dataset"):
-        hub_api.delete_file(path_in_repo=STATS_PATH, repo_id=dest_repo_id, revision=branch, repo_type="dataset")
-
-    # Tag destination with current codebase version
-    hub_api.create_tag(dest_repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
+    hub_api.create_tag(repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "--source-repo-id",
+        "--repo-id",
         type=str,
         required=True,
-        help="Source dataset repo id to download from (must be v2.0).",
-    )
-    parser.add_argument(
-        "--dest-repo-id",
-        type=str,
-        required=True,
-        help="Destination dataset repo id to upload the converted metadata to.",
+        help="Repository identifier on Hugging Face: a community or a user name `/` the name of the dataset "
+        "(e.g. `lerobot/pusht`, `cadene/aloha_sim_insertion_human`).",
     )
     parser.add_argument(
         "--branch",

src/lerobot/datasets/v21/convert_stats.py

@@ -13,13 +13,28 @@
 # limitations under the License.
 
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
 
+import jsonlines
 import numpy as np
 from tqdm import tqdm
 
 from lerobot.datasets.compute_stats import aggregate_stats, get_feature_stats, sample_indices
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.utils import write_episode_stats
+from lerobot.datasets.utils import LEGACY_EPISODES_STATS_PATH, serialize_dict
+
+
+def append_jsonlines(data: dict, fpath: Path) -> None:
+    fpath.parent.mkdir(exist_ok=True, parents=True)
+    with jsonlines.open(fpath, "a") as writer:
+        writer.write(data)
+
+
+def legacy_write_episode_stats(episode_index: int, episode_stats: dict, local_dir: Path):
+    # We wrap episode_stats in a dictionary since `episode_stats["episode_index"]`
+    # is a dictionary of stats and not an integer.
+    episode_stats = {"episode_index": episode_index, "stats": serialize_dict(episode_stats)}
+    append_jsonlines(episode_stats, local_dir / LEGACY_EPISODES_STATS_PATH)
 
 
 def sample_episode_video_frames(dataset: LeRobotDataset, episode_index: int, ft_key: str) -> np.ndarray:
@@ -72,7 +87,7 @@ def convert_stats(dataset: LeRobotDataset, num_workers: int = 0):
             convert_episode_stats(dataset, ep_idx)
 
     for ep_idx in tqdm(range(total_episodes)):
-        write_episode_stats(ep_idx, dataset.meta.episodes_stats[ep_idx], dataset.root)
+        legacy_write_episode_stats(ep_idx, dataset.meta.episodes_stats[ep_idx], dataset.root)
 
 
 def check_aggregate_stats(

src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py

@@ -0,0 +1,480 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script will help you convert any LeRobot dataset already pushed to the hub from codebase version 2.1 to
+3.0. It will:
+
+- Generate per-episodes stats and writes them in `episodes_stats.jsonl`
+- Check consistency between these new stats and the old ones.
+- Remove the deprecated `stats.json`.
+- Update codebase_version in `info.json`.
+- Push this new version to the hub on the 'main' branch and tags it with "v3.0".
+
+Usage:
+
+```bash
+python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
+    --repo-id=lerobot/pusht
+```
+
+"""
+
+import argparse
+import shutil
+from pathlib import Path
+from typing import Any
+
+import jsonlines
+import pandas as pd
+import pyarrow as pa
+import tqdm
+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 (
+    DEFAULT_CHUNK_SIZE,
+    DEFAULT_DATA_FILE_SIZE_IN_MB,
+    DEFAULT_DATA_PATH,
+    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
+    DEFAULT_VIDEO_PATH,
+    LEGACY_EPISODES_PATH,
+    LEGACY_EPISODES_STATS_PATH,
+    LEGACY_TASKS_PATH,
+    cast_stats_to_numpy,
+    flatten_dict,
+    get_parquet_file_size_in_mb,
+    get_parquet_num_frames,
+    get_video_duration_in_s,
+    get_video_size_in_mb,
+    load_info,
+    update_chunk_file_indices,
+    write_episodes,
+    write_info,
+    write_stats,
+    write_tasks,
+)
+from lerobot.datasets.video_utils import concat_video_files
+
+V21 = "v2.1"
+
+
+"""
+-------------------------
+OLD
+data/chunk-000/episode_000000.parquet
+
+NEW
+data/chunk-000/file_000.parquet
+-------------------------
+OLD
+videos/chunk-000/CAMERA/episode_000000.mp4
+
+NEW
+videos/chunk-000/file_000.mp4
+-------------------------
+OLD
+episodes.jsonl
+{"episode_index": 1, "tasks": ["Put the blue block in the green bowl"], "length": 266}
+
+NEW
+meta/episodes/chunk-000/episodes_000.parquet
+episode_index | video_chunk_index | video_file_index | data_chunk_index | data_file_index | tasks | length
+-------------------------
+OLD
+tasks.jsonl
+{"task_index": 1, "task": "Put the blue block in the green bowl"}
+
+NEW
+meta/tasks/chunk-000/file_000.parquet
+task_index | task
+-------------------------
+OLD
+episodes_stats.jsonl
+
+NEW
+meta/episodes_stats/chunk-000/file_000.parquet
+episode_index | mean | std | min | max
+-------------------------
+UPDATE
+meta/info.json
+-------------------------
+"""
+
+
+def load_jsonlines(fpath: Path) -> list[Any]:
+    with jsonlines.open(fpath, "r") as reader:
+        return list(reader)
+
+
+def legacy_load_episodes(local_dir: Path) -> dict:
+    episodes = load_jsonlines(local_dir / LEGACY_EPISODES_PATH)
+    return {item["episode_index"]: item for item in sorted(episodes, key=lambda x: x["episode_index"])}
+
+
+def legacy_load_episodes_stats(local_dir: Path) -> dict:
+    episodes_stats = load_jsonlines(local_dir / LEGACY_EPISODES_STATS_PATH)
+    return {
+        item["episode_index"]: cast_stats_to_numpy(item["stats"])
+        for item in sorted(episodes_stats, key=lambda x: x["episode_index"])
+    }
+
+
+def legacy_load_tasks(local_dir: Path) -> tuple[dict, dict]:
+    tasks = load_jsonlines(local_dir / LEGACY_TASKS_PATH)
+    tasks = {item["task_index"]: item["task"] for item in sorted(tasks, key=lambda x: x["task_index"])}
+    task_to_task_index = {task: task_index for task_index, task in tasks.items()}
+    return tasks, task_to_task_index
+
+
+def convert_tasks(root, new_root):
+    tasks, _ = legacy_load_tasks(root)
+    task_indices = tasks.keys()
+    task_strings = tasks.values()
+    df_tasks = pd.DataFrame({"task_index": task_indices}, index=task_strings)
+    write_tasks(df_tasks, new_root)
+
+
+def concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys):
+    # TODO(rcadene): to save RAM use Dataset.from_parquet(file) and concatenate_datasets
+    dataframes = [pd.read_parquet(file) for file in paths_to_cat]
+    # Concatenate all DataFrames along rows
+    concatenated_df = pd.concat(dataframes, ignore_index=True)
+
+    path = new_root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    if len(image_keys) > 0:
+        schema = pa.Schema.from_pandas(concatenated_df)
+        features = Features.from_arrow_schema(schema)
+        for key in image_keys:
+            features[key] = Image()
+        schema = features.arrow_schema
+    else:
+        schema = None
+
+    concatenated_df.to_parquet(path, index=False, schema=schema)
+
+
+def convert_data(root, new_root):
+    data_dir = root / "data"
+    ep_paths = sorted(data_dir.glob("*/*.parquet"))
+
+    image_keys = get_image_keys(root)
+
+    ep_idx = 0
+    chunk_idx = 0
+    file_idx = 0
+    size_in_mb = 0
+    num_frames = 0
+    paths_to_cat = []
+    episodes_metadata = []
+    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 = {
+            "episode_index": ep_idx,
+            "data/chunk_index": chunk_idx,
+            "data/file_index": file_idx,
+            "dataset_from_index": num_frames,
+            "dataset_to_index": num_frames + ep_num_frames,
+        }
+        size_in_mb += ep_size_in_mb
+        num_frames += ep_num_frames
+        episodes_metadata.append(ep_metadata)
+        ep_idx += 1
+
+        if size_in_mb < DEFAULT_DATA_FILE_SIZE_IN_MB:
+            paths_to_cat.append(ep_path)
+            continue
+
+        concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
+
+        # Reset for the next file
+        size_in_mb = ep_size_in_mb
+        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)
+
+    # Write remaining data if any
+    if paths_to_cat:
+        concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
+
+    return episodes_metadata
+
+
+def get_video_keys(root):
+    info = load_info(root)
+    features = info["features"]
+    video_keys = [key for key, ft in features.items() if ft["dtype"] == "video"]
+    return video_keys
+
+
+def get_image_keys(root):
+    info = load_info(root)
+    features = info["features"]
+    image_keys = [key for key, ft in features.items() if ft["dtype"] == "image"]
+    return image_keys
+
+
+def convert_videos(root: Path, new_root: Path):
+    video_keys = get_video_keys(root)
+    if len(video_keys) == 0:
+        return None
+
+    video_keys = sorted(video_keys)
+
+    eps_metadata_per_cam = []
+    for camera in video_keys:
+        eps_metadata = convert_videos_of_camera(root, new_root, camera)
+        eps_metadata_per_cam.append(eps_metadata)
+
+    num_eps_per_cam = [len(eps_cam_map) for eps_cam_map in eps_metadata_per_cam]
+    if len(set(num_eps_per_cam)) != 1:
+        raise ValueError(f"All cams dont have same number of episodes ({num_eps_per_cam}).")
+
+    episods_metadata = []
+    num_cameras = len(video_keys)
+    num_episodes = num_eps_per_cam[0]
+    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]
+        if len(set(ep_ids)) != 1:
+            raise ValueError(f"All episode indices need to match ({ep_ids}).")
+
+        ep_dict = {}
+        for cam_idx in range(num_cameras):
+            ep_dict.update(eps_metadata_per_cam[cam_idx][ep_idx])
+        episods_metadata.append(ep_dict)
+
+    return episods_metadata
+
+
+def convert_videos_of_camera(root: Path, new_root: Path, video_key):
+    # Access old paths to mp4
+    videos_dir = root / "videos"
+    ep_paths = sorted(videos_dir.glob(f"*/{video_key}/*.mp4"))
+
+    ep_idx = 0
+    chunk_idx = 0
+    file_idx = 0
+    size_in_mb = 0
+    duration_in_s = 0.0
+    paths_to_cat = []
+    episodes_metadata = []
+    for ep_path in tqdm.tqdm(ep_paths, desc=f"convert videos of {video_key}"):
+        ep_size_in_mb = get_video_size_in_mb(ep_path)
+        ep_duration_in_s = get_video_duration_in_s(ep_path)
+
+        # Check if adding this episode would exceed the limit
+        if size_in_mb + ep_size_in_mb >= DEFAULT_VIDEO_FILE_SIZE_IN_MB and len(paths_to_cat) > 0:
+            # Size limit would be exceeded, save current accumulation WITHOUT this episode
+            concat_video_files(paths_to_cat, new_root, video_key, chunk_idx, file_idx)
+
+            # Update episodes metadata for the file we just saved
+            for i, _ in enumerate(paths_to_cat):
+                past_ep_idx = ep_idx - len(paths_to_cat) + i
+                episodes_metadata[past_ep_idx][f"videos/{video_key}/chunk_index"] = chunk_idx
+                episodes_metadata[past_ep_idx][f"videos/{video_key}/file_index"] = file_idx
+
+            # Move to next file and start fresh with current episode
+            chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
+            size_in_mb = 0
+            duration_in_s = 0.0
+            paths_to_cat = []
+
+        # Add current episode metadata
+        ep_metadata = {
+            "episode_index": ep_idx,
+            f"videos/{video_key}/chunk_index": chunk_idx,  # Will be updated when file is saved
+            f"videos/{video_key}/file_index": file_idx,  # Will be updated when file is saved
+            f"videos/{video_key}/from_timestamp": duration_in_s,
+            f"videos/{video_key}/to_timestamp": duration_in_s + ep_duration_in_s,
+        }
+        episodes_metadata.append(ep_metadata)
+
+        # Add current episode to accumulation
+        paths_to_cat.append(ep_path)
+        size_in_mb += ep_size_in_mb
+        duration_in_s += ep_duration_in_s
+        ep_idx += 1
+
+    # Write remaining videos if any
+    if paths_to_cat:
+        concat_video_files(paths_to_cat, new_root, video_key, chunk_idx, file_idx)
+
+        # Update episodes metadata for the final file
+        for i, _ in enumerate(paths_to_cat):
+            past_ep_idx = ep_idx - len(paths_to_cat) + i
+            episodes_metadata[past_ep_idx][f"videos/{video_key}/chunk_index"] = chunk_idx
+            episodes_metadata[past_ep_idx][f"videos/{video_key}/file_index"] = file_idx
+
+    return episodes_metadata
+
+
+def generate_episode_metadata_dict(
+    episodes_legacy_metadata, episodes_metadata, episodes_stats, episodes_videos=None
+):
+    num_episodes = len(episodes_metadata)
+    episodes_legacy_metadata_vals = list(episodes_legacy_metadata.values())
+    episodes_stats_vals = list(episodes_stats.values())
+    episodes_stats_keys = list(episodes_stats.keys())
+
+    for i in range(num_episodes):
+        ep_legacy_metadata = episodes_legacy_metadata_vals[i]
+        ep_metadata = episodes_metadata[i]
+        ep_stats = episodes_stats_vals[i]
+
+        ep_ids_set = {
+            ep_legacy_metadata["episode_index"],
+            ep_metadata["episode_index"],
+            episodes_stats_keys[i],
+        }
+
+        if episodes_videos is None:
+            ep_video = {}
+        else:
+            ep_video = episodes_videos[i]
+            ep_ids_set.add(ep_video["episode_index"])
+
+        if len(ep_ids_set) != 1:
+            raise ValueError(f"Number of episodes is not the same ({ep_ids_set}).")
+
+        ep_dict = {**ep_metadata, **ep_video, **ep_legacy_metadata, **flatten_dict({"stats": ep_stats})}
+        ep_dict["meta/episodes/chunk_index"] = 0
+        ep_dict["meta/episodes/file_index"] = 0
+        yield ep_dict
+
+
+def convert_episodes_metadata(root, new_root, episodes_metadata, episodes_video_metadata=None):
+    episodes_legacy_metadata = legacy_load_episodes(root)
+    episodes_stats = legacy_load_episodes_stats(root)
+
+    num_eps_set = {len(episodes_legacy_metadata), len(episodes_metadata)}
+    if episodes_video_metadata is not None:
+        num_eps_set.add(len(episodes_video_metadata))
+
+    if len(num_eps_set) != 1:
+        raise ValueError(f"Number of episodes is not the same ({num_eps_set}).")
+
+    ds_episodes = Dataset.from_generator(
+        lambda: generate_episode_metadata_dict(
+            episodes_legacy_metadata, episodes_metadata, episodes_stats, episodes_video_metadata
+        )
+    )
+    write_episodes(ds_episodes, new_root)
+
+    stats = aggregate_stats(list(episodes_stats.values()))
+    write_stats(stats, new_root)
+
+
+def convert_info(root, new_root):
+    info = load_info(root)
+    info["codebase_version"] = "v3.0"
+    del info["total_chunks"]
+    del info["total_videos"]
+    info["data_files_size_in_mb"] = DEFAULT_DATA_FILE_SIZE_IN_MB
+    info["video_files_size_in_mb"] = DEFAULT_VIDEO_FILE_SIZE_IN_MB
+    info["data_path"] = DEFAULT_DATA_PATH
+    info["video_path"] = DEFAULT_VIDEO_PATH
+    info["fps"] = float(info["fps"])
+    for key in info["features"]:
+        if info["features"][key]["dtype"] == "video":
+            # already has fps in video_info
+            continue
+        info["features"][key]["fps"] = info["fps"]
+    write_info(info, new_root)
+
+
+def convert_dataset(
+    repo_id: str,
+    branch: str | None = None,
+    num_workers: int = 4,
+):
+    root = HF_LEROBOT_HOME / repo_id
+    old_root = HF_LEROBOT_HOME / f"{repo_id}_old"
+    new_root = HF_LEROBOT_HOME / f"{repo_id}_v30"
+
+    if old_root.is_dir() and root.is_dir():
+        shutil.rmtree(str(root))
+        shutil.move(str(old_root), str(root))
+
+    if new_root.is_dir():
+        shutil.rmtree(new_root)
+
+    snapshot_download(
+        repo_id,
+        repo_type="dataset",
+        revision=V21,
+        local_dir=root,
+    )
+
+    convert_info(root, new_root)
+    convert_tasks(root, new_root)
+    episodes_metadata = convert_data(root, new_root)
+    episodes_videos_metadata = convert_videos(root, new_root)
+    convert_episodes_metadata(root, new_root, episodes_metadata, episodes_videos_metadata)
+
+    shutil.move(str(root), str(old_root))
+    shutil.move(str(new_root), str(root))
+
+    hub_api = HfApi()
+    try:
+        hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
+    except HTTPError as e:
+        print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
+        pass
+    hub_api.delete_files(
+        delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"],
+        repo_id=repo_id,
+        revision=branch,
+        repo_type="dataset",
+    )
+    hub_api.create_tag(repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
+
+    LeRobotDataset(repo_id).push_to_hub()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        required=True,
+        help="Repository identifier on Hugging Face: a community or a user name `/` the name of the dataset "
+        "(e.g. `lerobot/pusht`, `cadene/aloha_sim_insertion_human`).",
+    )
+    parser.add_argument(
+        "--branch",
+        type=str,
+        default=None,
+        help="Repo branch to push your dataset. Defaults to the main branch.",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=4,
+        help="Number of workers for parallelizing stats compute. Defaults to 4.",
+    )
+
+    args = parser.parse_args()
+    convert_dataset(**vars(args))

src/lerobot/datasets/video_utils.py

@@ -13,22 +13,26 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import glob
 import importlib
+import json
 import logging
 import shutil
+import subprocess
+import tempfile
 import warnings
+from collections import OrderedDict
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, ClassVar
 
-import av
 import pyarrow as pa
 import torch
 import torchvision
 from datasets.features.features import register_feature
 from PIL import Image
 
+from lerobot.datasets.utils import DEFAULT_VIDEO_PATH
+
 
 def get_safe_default_codec():
     if importlib.util.find_spec("torchcodec"):
@@ -102,7 +106,7 @@ def decode_video_frames_torchvision(
     keyframes_only = False
     torchvision.set_video_backend(backend)
     if backend == "pyav":
-        keyframes_only = True  # pyav doesn't support accurate seek
+        keyframes_only = True  # pyav doesnt support accuracte seek
 
     # set a video stream reader
     # TODO(rcadene): also load audio stream at the same time
@@ -155,6 +159,7 @@ def decode_video_frames_torchvision(
     )
 
     # get closest frames to the query timestamps
+    # TODO(rcadene): remove torch.stack
     closest_frames = torch.stack([loaded_frames[idx] for idx in argmin_])
     closest_ts = loaded_ts[argmin_]
 
@@ -252,83 +257,104 @@ def encode_video_frames(
     g: int | None = 2,
     crf: int | None = 30,
     fast_decode: int = 0,
-    log_level: int | None = av.logging.ERROR,
+    log_level: str | None = "quiet",
     overwrite: bool = False,
 ) -> None:
     """More info on ffmpeg arguments tuning on `benchmark/video/README.md`"""
-    # Check encoder availability
-    if vcodec not in ["h264", "hevc", "libsvtav1"]:
-        raise ValueError(f"Unsupported video codec: {vcodec}. Supported codecs are: h264, hevc, libsvtav1.")
-
     video_path = Path(video_path)
     imgs_dir = Path(imgs_dir)
+    video_path.parent.mkdir(parents=True, exist_ok=True)
 
-    video_path.parent.mkdir(parents=True, exist_ok=overwrite)
-
-    # Encoders/pixel formats incompatibility check
-    if (vcodec == "libsvtav1" or vcodec == "hevc") and pix_fmt == "yuv444p":
-        logging.warning(
-            f"Incompatible pixel format 'yuv444p' for codec {vcodec}, auto-selecting format 'yuv420p'"
-        )
-        pix_fmt = "yuv420p"
-
-    # Get input frames
-    template = "frame_" + ("[0-9]" * 6) + ".png"
-    input_list = sorted(
-        glob.glob(str(imgs_dir / template)), key=lambda x: int(x.split("_")[-1].split(".")[0])
+    ffmpeg_args = OrderedDict(
+        [
+            ("-f", "image2"),
+            ("-r", str(fps)),
+            ("-i", str(imgs_dir / "frame-%06d.png")),
+            ("-vcodec", vcodec),
+            ("-pix_fmt", pix_fmt),
+        ]
     )
 
-    # Define video output frame size (assuming all input frames are the same size)
-    if len(input_list) == 0:
-        raise FileNotFoundError(f"No images found in {imgs_dir}.")
-    dummy_image = Image.open(input_list[0])
-    width, height = dummy_image.size
-
-    # Define video codec options
-    video_options = {}
-
     if g is not None:
-        video_options["g"] = str(g)
+        ffmpeg_args["-g"] = str(g)
 
     if crf is not None:
-        video_options["crf"] = str(crf)
+        ffmpeg_args["-crf"] = str(crf)
 
     if fast_decode:
-        key = "svtav1-params" if vcodec == "libsvtav1" else "tune"
+        key = "-svtav1-params" if vcodec == "libsvtav1" else "-tune"
         value = f"fast-decode={fast_decode}" if vcodec == "libsvtav1" else "fastdecode"
-        video_options[key] = value
+        ffmpeg_args[key] = value
 
-    # Set logging level
     if log_level is not None:
-        # "While less efficient, it is generally preferable to modify logging with Python’s logging"
-        logging.getLogger("libav").setLevel(log_level)
+        ffmpeg_args["-loglevel"] = str(log_level)
 
-    # Create and open output file (overwrite by default)
-    with av.open(str(video_path), "w") as output:
-        output_stream = output.add_stream(vcodec, fps, options=video_options)
-        output_stream.pix_fmt = pix_fmt
-        output_stream.width = width
-        output_stream.height = height
+    ffmpeg_args = [item for pair in ffmpeg_args.items() for item in pair]
+    if overwrite:
+        ffmpeg_args.append("-y")
 
-        # Loop through input frames and encode them
-        for input_data in input_list:
-            input_image = Image.open(input_data).convert("RGB")
-            input_frame = av.VideoFrame.from_image(input_image)
-            packet = output_stream.encode(input_frame)
-            if packet:
-                output.mux(packet)
-
-        # Flush the encoder
-        packet = output_stream.encode()
-        if packet:
-            output.mux(packet)
-
-    # Reset logging level
-    if log_level is not None:
-        av.logging.restore_default_callback()
+    ffmpeg_cmd = ["ffmpeg"] + ffmpeg_args + [str(video_path)]
+    # redirect stdin to subprocess.DEVNULL to prevent reading random keyboard inputs from terminal
+    subprocess.run(ffmpeg_cmd, check=True, stdin=subprocess.DEVNULL)
 
     if not video_path.exists():
-        raise OSError(f"Video encoding did not work. File not found: {video_path}.")
+        raise OSError(
+            f"Video encoding did not work. File not found: {video_path}. "
+            f"Try running the command manually to debug: `{''.join(ffmpeg_cmd)}`"
+        )
+
+
+def concat_video_files(paths_to_cat: list[Path], root: Path, video_key: str, chunk_idx: int, file_idx: int):
+    """
+    Concatenate multiple video files into a single video file using ffmpeg.
+
+    This function takes a list of video file paths and concatenates them into a single
+    output video file. It uses ffmpeg's concat demuxer with stream copy mode for fast
+    concatenation without re-encoding.
+
+    Args:
+        paths_to_cat: List of video file paths to concatenate, in order.
+        root: Root directory where temporary files and output will be created.
+        video_key: Video key identifier (e.g., camera name) used in output path.
+        chunk_idx: Chunk index for organizing output files.
+        file_idx: File index within the chunk.
+
+    Note:
+        - Creates a temporary directory for intermediate files that is cleaned up after use.
+        - Uses ffmpeg's concat demuxer which requires all input videos to have the same
+          codec, resolution, and frame rate for proper concatenation.
+        - Output path follows the DEFAULT_VIDEO_PATH pattern with video_key, chunk_idx,
+          and file_idx parameters.
+    """
+
+    tmp_dir = Path(tempfile.mkdtemp(dir=root))
+    path_concat_video_files = tmp_dir / "concat_video_files.txt"
+    with open(path_concat_video_files, "w") as f:
+        for ep_path in paths_to_cat:
+            f.write(f"file '{str(ep_path)}'\n")
+
+    path_tmp_output = tmp_dir / "tmp_output.mp4"
+    command = [
+        "ffmpeg",
+        "-y",
+        "-f",
+        "concat",
+        "-safe",
+        "0",
+        "-i",
+        str(path_concat_video_files),
+        "-c",
+        "copy",
+        str(path_tmp_output),
+    ]
+    subprocess.run(command, check=True)
+
+    output_path = root / DEFAULT_VIDEO_PATH.format(
+        video_key=video_key, chunk_index=chunk_idx, file_index=file_idx
+    )
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    shutil.move(str(path_tmp_output), str(output_path))
+    shutil.rmtree(str(tmp_dir))
 
 
 @dataclass
@@ -364,68 +390,78 @@ with warnings.catch_warnings():
 
 
 def get_audio_info(video_path: Path | str) -> dict:
-    # Set logging level
-    logging.getLogger("libav").setLevel(av.logging.ERROR)
+    ffprobe_audio_cmd = [
+        "ffprobe",
+        "-v",
+        "error",
+        "-select_streams",
+        "a:0",
+        "-show_entries",
+        "stream=channels,codec_name,bit_rate,sample_rate,bit_depth,channel_layout,duration",
+        "-of",
+        "json",
+        str(video_path),
+    ]
+    result = subprocess.run(ffprobe_audio_cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        raise RuntimeError(f"Error running ffprobe: {result.stderr}")
 
-    # Getting audio stream information
-    audio_info = {}
-    with av.open(str(video_path), "r") as audio_file:
-        try:
-            audio_stream = audio_file.streams.audio[0]
-        except IndexError:
-            # Reset logging level
-            av.logging.restore_default_callback()
-            return {"has_audio": False}
+    info = json.loads(result.stdout)
+    audio_stream_info = info["streams"][0] if info.get("streams") else None
+    if audio_stream_info is None:
+        return {"has_audio": False}
 
-        audio_info["audio.channels"] = audio_stream.channels
-        audio_info["audio.codec"] = audio_stream.codec.canonical_name
-        # In an ideal loseless case : bit depth x sample rate x channels = bit rate.
-        # In an actual compressed case, the bit rate is set according to the compression level : the lower the bit rate, the more compression is applied.
-        audio_info["audio.bit_rate"] = audio_stream.bit_rate
-        audio_info["audio.sample_rate"] = audio_stream.sample_rate  # Number of samples per second
-        # In an ideal loseless case : fixed number of bits per sample.
-        # In an actual compressed case : variable number of bits per sample (often reduced to match a given depth rate).
-        audio_info["audio.bit_depth"] = audio_stream.format.bits
-        audio_info["audio.channel_layout"] = audio_stream.layout.name
-        audio_info["has_audio"] = True
-
-    # Reset logging level
-    av.logging.restore_default_callback()
-
-    return audio_info
+    # Return the information, defaulting to None if no audio stream is present
+    return {
+        "has_audio": True,
+        "audio.channels": audio_stream_info.get("channels", None),
+        "audio.codec": audio_stream_info.get("codec_name", None),
+        "audio.bit_rate": int(audio_stream_info["bit_rate"]) if audio_stream_info.get("bit_rate") else None,
+        "audio.sample_rate": int(audio_stream_info["sample_rate"])
+        if audio_stream_info.get("sample_rate")
+        else None,
+        "audio.bit_depth": audio_stream_info.get("bit_depth", None),
+        "audio.channel_layout": audio_stream_info.get("channel_layout", None),
+    }
 
 
 def get_video_info(video_path: Path | str) -> dict:
-    # Set logging level
-    logging.getLogger("libav").setLevel(av.logging.ERROR)
+    ffprobe_video_cmd = [
+        "ffprobe",
+        "-v",
+        "error",
+        "-select_streams",
+        "v:0",
+        "-show_entries",
+        "stream=r_frame_rate,width,height,codec_name,nb_frames,duration,pix_fmt",
+        "-of",
+        "json",
+        str(video_path),
+    ]
+    result = subprocess.run(ffprobe_video_cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        raise RuntimeError(f"Error running ffprobe: {result.stderr}")
 
-    # Getting video stream information
-    video_info = {}
-    with av.open(str(video_path), "r") as video_file:
-        try:
-            video_stream = video_file.streams.video[0]
-        except IndexError:
-            # Reset logging level
-            av.logging.restore_default_callback()
-            return {}
+    info = json.loads(result.stdout)
+    video_stream_info = info["streams"][0]
 
-        video_info["video.height"] = video_stream.height
-        video_info["video.width"] = video_stream.width
-        video_info["video.codec"] = video_stream.codec.canonical_name
-        video_info["video.pix_fmt"] = video_stream.pix_fmt
-        video_info["video.is_depth_map"] = False
+    # Calculate fps from r_frame_rate
+    r_frame_rate = video_stream_info["r_frame_rate"]
+    num, denom = map(int, r_frame_rate.split("/"))
+    fps = num / denom
 
-        # Calculate fps from r_frame_rate
-        video_info["video.fps"] = int(video_stream.base_rate)
+    pixel_channels = get_video_pixel_channels(video_stream_info["pix_fmt"])
 
-        pixel_channels = get_video_pixel_channels(video_stream.pix_fmt)
-        video_info["video.channels"] = pixel_channels
-
-    # Reset logging level
-    av.logging.restore_default_callback()
-
-    # Adding audio stream information
-    video_info.update(**get_audio_info(video_path))
+    video_info = {
+        "video.fps": fps,
+        "video.height": video_stream_info["height"],
+        "video.width": video_stream_info["width"],
+        "video.channels": pixel_channels,
+        "video.codec": video_stream_info["codec_name"],
+        "video.pix_fmt": video_stream_info["pix_fmt"],
+        "video.is_depth_map": False,
+        **get_audio_info(video_path),
+    }
 
     return video_info
 
@@ -452,66 +488,3 @@ def get_image_pixel_channels(image: Image):
         return 4  # RGBA
     else:
         raise ValueError("Unknown format")
-
-
-class VideoEncodingManager:
-    """
-    Context manager that ensures proper video encoding and data cleanup even if exceptions occur.
-
-    This manager handles:
-    - Batch encoding for any remaining episodes when recording interrupted
-    - Cleaning up temporary image files from interrupted episodes
-    - Removing empty image directories
-
-    Args:
-        dataset: The LeRobotDataset instance
-    """
-
-    def __init__(self, dataset):
-        self.dataset = dataset
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        # Handle any remaining episodes that haven't been batch encoded
-        if self.dataset.episodes_since_last_encoding > 0:
-            if exc_type is not None:
-                logging.info("Exception occurred. Encoding remaining episodes before exit...")
-            else:
-                logging.info("Recording stopped. Encoding remaining episodes...")
-
-            start_ep = self.dataset.num_episodes - self.dataset.episodes_since_last_encoding
-            end_ep = self.dataset.num_episodes
-            logging.info(
-                f"Encoding remaining {self.dataset.episodes_since_last_encoding} episodes, "
-                f"from episode {start_ep} to {end_ep - 1}"
-            )
-            self.dataset.batch_encode_videos(start_ep, end_ep)
-
-        # Clean up episode images if recording was interrupted
-        if exc_type is not None:
-            interrupted_episode_index = self.dataset.num_episodes
-            for key in self.dataset.meta.video_keys:
-                img_dir = self.dataset._get_image_file_path(
-                    episode_index=interrupted_episode_index, image_key=key, frame_index=0
-                ).parent
-                if img_dir.exists():
-                    logging.debug(
-                        f"Cleaning up interrupted episode images for episode {interrupted_episode_index}, camera {key}"
-                    )
-                    shutil.rmtree(img_dir)
-
-        # Clean up any remaining images directory if it's empty
-        img_dir = self.dataset.root / "images"
-        # Check for any remaining PNG files
-        png_files = list(img_dir.rglob("*.png"))
-        if len(png_files) == 0:
-            # Only remove the images directory if no PNG files remain
-            if img_dir.exists():
-                shutil.rmtree(img_dir)
-                logging.debug("Cleaned up empty images directory")
-        else:
-            logging.debug(f"Images directory is not empty, containing {len(png_files)} PNG files")
-
-        return False  # Don't suppress the original exception

src/lerobot/find_cameras.py

@@ -20,7 +20,7 @@ Helper to find the camera devices available in your system.
 Example:
 
 ```shell
-python -m lerobot.find_cameras
+lerobot-find-cameras
 ```
 """
 

src/lerobot/find_port.py

@@ -18,7 +18,7 @@ Helper to find the USB port associated with your MotorsBus.
 Example:
 
 ```shell
-python -m lerobot.find_port
+lerobot-find-port
 ```
 """
 

src/lerobot/motors/motors_bus.py

@@ -222,7 +222,7 @@ class MotorsBus(abc.ABC):
     A MotorsBus subclass instance requires a port (e.g. `FeetechMotorsBus(port="/dev/tty.usbmodem575E0031751"`)).
     To find the port, you can run our utility script:
     ```bash
-    python -m lerobot.find_port.py
+    lerobot-find-port.py
     >>> Finding all available ports for the MotorsBus.
     >>> ["/dev/tty.usbmodem575E0032081", "/dev/tty.usbmodem575E0031751"]
     >>> Remove the usb cable from your MotorsBus and press Enter when done.
@@ -446,7 +446,7 @@ class MotorsBus(abc.ABC):
         except (FileNotFoundError, OSError, serial.SerialException) as e:
             raise ConnectionError(
                 f"\nCould not connect on port '{self.port}'. Make sure you are using the correct port."
-                "\nTry running `python -m lerobot.find_port`\n"
+                "\nTry running `lerobot-find-port`\n"
             ) from e
 
     @abc.abstractmethod

src/lerobot/policies/__init__.py

@@ -15,19 +15,6 @@
 from .act.configuration_act import ACTConfig as ACTConfig
 from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
 from .pi0.configuration_pi0 import PI0Config as PI0Config
-from .pi0.processor_pi0 import Pi0NewLineProcessor
-from .rlearn.configuration_rlearn import RLearNConfig as RLearNConfig
 from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
-from .smolvla.processor_smolvla import SmolVLANewLineProcessor
 from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
 from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
-
-__all__ = [
-    "ACTConfig",
-    "DiffusionConfig",
-    "PI0Config",
-    "SmolVLAConfig",
-    "TDMPCConfig",
-    "VQBeTConfig",
-    "RLearNConfig",
-]

src/lerobot/policies/act/modeling_act.py

@@ -35,6 +35,7 @@ from torchvision.ops.misc import FrozenBatchNorm2d
 
 from lerobot.constants import ACTION, OBS_IMAGES
 from lerobot.policies.act.configuration_act import ACTConfig
+from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pretrained import PreTrainedPolicy
 
 
@@ -50,16 +51,27 @@ class ACTPolicy(PreTrainedPolicy):
     def __init__(
         self,
         config: ACTConfig,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
             config: Policy configuration class instance or None, in which case the default instantiation of
                     the configuration class is used.
+            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
+                that they will be passed with a call to `load_state_dict` before the policy is used.
         """
         super().__init__(config)
         config.validate_features()
         self.config = config
 
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
         self.model = ACT(config)
 
         if config.temporal_ensemble_coeff is not None:
@@ -125,19 +137,23 @@ class ACTPolicy(PreTrainedPolicy):
         """Predict a chunk of actions given environment observations."""
         self.eval()
 
+        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGES] = [batch[key] for key in self.config.image_features]
 
         actions = self.model(batch)[0]
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
         return actions
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """Run the batch through the model and compute the loss for training or validation."""
+        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGES] = [batch[key] for key in self.config.image_features]
 
+        batch = self.normalize_targets(batch)
         actions_hat, (mu_hat, log_sigma_x2_hat) = self.model(batch)
 
         l1_loss = (
@@ -287,7 +303,7 @@ class ACT(nn.Module):
                                 └───────────────────────┘
     """
 
-    def __init__(self, config: ACTConfig, dataset_stats=None):
+    def __init__(self, config: ACTConfig):
         # BERT style VAE encoder with input tokens [cls, robot_state, *action_sequence].
         # The cls token forms parameters of the latent's distribution (like this [*means, *log_variances]).
         super().__init__()

src/lerobot/policies/act/processor_act.py

@@ -1,50 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 Tony Z. Zhao and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-
-from lerobot.policies.act.configuration_act import ACTConfig
-from lerobot.processor import (
-    DeviceProcessor,
-    NormalizerProcessor,
-    RenameProcessor,
-    RobotProcessor,
-    ToBatchProcessor,
-    UnnormalizerProcessor,
-)
-
-
-def make_act_processor(
-    config: ACTConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
-) -> tuple[RobotProcessor, RobotProcessor]:
-    input_steps = [
-        RenameProcessor(rename_map={}),
-        NormalizerProcessor(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        ToBatchProcessor(),
-        DeviceProcessor(device=config.device),
-    ]
-    output_steps = [
-        DeviceProcessor(device="cpu"),
-        UnnormalizerProcessor(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
-        steps=output_steps, name="robot_postprocessor"
-    )

src/lerobot/policies/diffusion/modeling_diffusion.py

@@ -35,6 +35,7 @@ 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.normalize import Normalize, Unnormalize
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import (
     get_device_from_parameters,
@@ -56,6 +57,7 @@ class DiffusionPolicy(PreTrainedPolicy):
     def __init__(
         self,
         config: DiffusionConfig,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
@@ -68,6 +70,14 @@ class DiffusionPolicy(PreTrainedPolicy):
         config.validate_features()
         self.config = config
 
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
         # queues are populated during rollout of the policy, they contain the n latest observations and actions
         self._queues = None
 
@@ -96,6 +106,9 @@ class DiffusionPolicy(PreTrainedPolicy):
         batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
         actions = self.diffusion.generate_actions(batch)
 
+        # TODO(rcadene): make above methods return output dictionary?
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
+
         return actions
 
     @torch.no_grad()
@@ -124,6 +137,7 @@ class DiffusionPolicy(PreTrainedPolicy):
         if ACTION in batch:
             batch.pop(ACTION)
 
+        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
@@ -139,9 +153,11 @@ class DiffusionPolicy(PreTrainedPolicy):
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, None]:
         """Run the batch through the model and compute the loss for training or validation."""
+        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
+        batch = self.normalize_targets(batch)
         loss = self.diffusion.compute_loss(batch)
         # no output_dict so returning None
         return loss, None

src/lerobot/policies/diffusion/processor_diffusion.py

@@ -1,51 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 Columbia Artificial Intelligence, Robotics Lab,
-# and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-
-from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
-from lerobot.processor import (
-    DeviceProcessor,
-    NormalizerProcessor,
-    RenameProcessor,
-    RobotProcessor,
-    ToBatchProcessor,
-    UnnormalizerProcessor,
-)
-
-
-def make_diffusion_processor(
-    config: DiffusionConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
-) -> tuple[RobotProcessor, RobotProcessor]:
-    input_steps = [
-        RenameProcessor(rename_map={}),
-        NormalizerProcessor(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        ToBatchProcessor(),
-        DeviceProcessor(device=config.device),
-    ]
-    output_steps = [
-        DeviceProcessor(device="cpu"),
-        UnnormalizerProcessor(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
-        steps=output_steps, name="robot_postprocessor"
-    )

src/lerobot/policies/factory.py

@@ -14,14 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import logging
-from typing import Any, TypedDict, cast
 
-import torch
 from torch import nn
-from typing_extensions import Unpack
 
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType
@@ -34,16 +29,14 @@ from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.policies.rlearn.configuration_rlearn import RLearNConfig
 from lerobot.policies.sac.configuration_sac import SACConfig
 from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
-from lerobot.processor.pipeline import RobotProcessor
 
 
-def get_policy_class(name: str) -> type[PreTrainedPolicy]:
+def get_policy_class(name: str) -> PreTrainedPolicy:
     """Get the policy's class and config class given a name (matching the policy class' `name` attribute)."""
     if name == "tdmpc":
         from lerobot.policies.tdmpc.modeling_tdmpc import TDMPCPolicy
@@ -81,10 +74,6 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
         from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
 
         return SmolVLAPolicy
-    elif name == "rlearn":
-        from lerobot.policies.rlearn.modeling_rlearn import RLearNPolicy
-
-        return RLearNPolicy
     else:
         raise NotImplementedError(f"Policy with name {name} is not implemented.")
 
@@ -108,143 +97,14 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return SmolVLAConfig(**kwargs)
     elif policy_type == "reward_classifier":
         return RewardClassifierConfig(**kwargs)
-    elif policy_type == "rlearn":
-        return RLearNConfig(**kwargs)
     else:
         raise ValueError(f"Policy type '{policy_type}' is not available.")
 
 
-class ProcessorConfigKwargs(TypedDict, total=False):
-    """Keyword arguments for the processor config."""
-
-    preprocessor_config_filename: str | None
-    postprocessor_config_filename: str | None
-    preprocessor_overrides: dict[str, Any] | None
-    postprocessor_overrides: dict[str, Any] | None
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None
-
-
-def make_processor(
-    policy_cfg: PreTrainedConfig,
-    pretrained_path: str | None = None,
-    **kwargs: Unpack[ProcessorConfigKwargs],
-) -> tuple[RobotProcessor, RobotProcessor]:
-    """Make a processor instance for a given policy type.
-
-    This function creates the appropriate processor configuration based on the policy type.
-    Each policy type has its own processor with specific preprocessing steps.
-
-    Args:
-        policy_cfg: The config of the policy to create a processor for (e.g., "act", "diffusion", etc.)
-        pretrained_path: Optional path to load a pretrained processor from. If provided, loads
-            the processor from this path instead of creating a new one.
-        **kwargs: Additional keyword arguments passed to the processor creation.
-
-    Returns:
-        Tuple of (input_processor, output_processor) for the policy.
-
-    Raises:
-        NotImplementedError: If the policy type doesn't have a processor implemented.
-    """
-    if pretrained_path:
-        # Load a pretrained processor
-        # TODO(azouitine): Handle this case.
-        return (
-            RobotProcessor.from_pretrained(
-                pretrained_model_name_or_path=pretrained_path,
-                config_filename=kwargs.get("preprocessor_config_filename", "robot_preprocessor.json"),
-                overrides=kwargs.get("preprocessor_overrides", {}),
-            ),
-            RobotProcessor.from_pretrained(
-                pretrained_model_name_or_path=pretrained_path,
-                config_filename=kwargs.get("postprocessor_config_filename", "robot_postprocessor.json"),
-                overrides=kwargs.get("postprocessor_overrides", {}),
-            ),
-        )
-
-    # Create a new processor based on policy type
-    if policy_cfg.type == "tdmpc":
-        from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
-        from lerobot.policies.tdmpc.processor_tdmpc import make_tdmpc_processor
-
-        processors = make_tdmpc_processor(
-            config=cast(TDMPCConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
-        )
-
-    elif policy_cfg.type == "diffusion":
-        from lerobot.policies.diffusion.processor_diffusion import make_diffusion_processor
-
-        processors = make_diffusion_processor(
-            cast(DiffusionConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
-        )
-
-    elif policy_cfg.type == "act":
-        from lerobot.policies.act.processor_act import make_act_processor
-
-        processors = make_act_processor(
-            config=cast(ACTConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
-        )
-
-    elif policy_cfg.type == "vqbet":
-        from lerobot.policies.vqbet.processor_vqbet import make_vqbet_processor
-
-        processors = make_vqbet_processor(
-            config=cast(VQBeTConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
-        )
-
-    elif policy_cfg.type == "pi0":
-        from lerobot.policies.pi0.processor_pi0 import make_pi0_processor
-
-        processors = make_pi0_processor(
-            config=cast(PI0Config, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
-        )
-
-    elif policy_cfg.type == "pi0fast":
-        from lerobot.policies.pi0fast.processor_pi0fast import make_pi0fast_processor
-
-        processors = make_pi0fast_processor(
-            cast(PI0Config, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
-        )
-
-    elif policy_cfg.type == "sac":
-        from lerobot.policies.sac.processor_sac import make_sac_processor
-
-        processors = make_sac_processor(
-            cast(SACConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
-        )
-
-    elif policy_cfg.type == "reward_classifier":
-        from lerobot.policies.sac.reward_model.processor_classifier import make_classifier_processor
-
-        processors = make_classifier_processor(
-            cast(RewardClassifierConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
-        )
-
-    elif policy_cfg.type == "smolvla":
-        from lerobot.policies.smolvla.processor_smolvla import make_smolvla_processor
-
-        processors = make_smolvla_processor(
-            cast(SmolVLAConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
-        )
-
-    elif policy_cfg.type == "rlearn":
-        from lerobot.policies.rlearn.processor_rlearn import make_rlearn_processor
-
-        processors = make_rlearn_processor(
-            cast(RLearNConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
-        )
-
-    else:
-        raise NotImplementedError(f"Processor for policy type '{policy_cfg.type}' is not implemented.")
-
-    return processors
-
-
 def make_policy(
     cfg: PreTrainedConfig,
     ds_meta: LeRobotDatasetMetadata | None = None,
     env_cfg: EnvConfig | None = None,
-    episode_data_index: dict | None = None,
 ) -> PreTrainedPolicy:
     """Make an instance of a policy class.
 
@@ -287,6 +147,7 @@ def make_policy(
     kwargs = {}
     if ds_meta is not None:
         features = dataset_to_policy_features(ds_meta.features)
+        kwargs["dataset_stats"] = ds_meta.stats
     else:
         if not cfg.pretrained_path:
             logging.warning(
@@ -294,18 +155,12 @@ def make_policy(
                 "rather than a dataset. Normalization modules inside the policy will have infinite values "
                 "by default without stats from a dataset."
             )
-        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
 
-    # Pass episode_data_index for RLearN policy to calculate proper progress
-    if cfg.type == "rlearn" and episode_data_index is not None:
-        kwargs["episode_data_index"] = episode_data_index
-
     if cfg.pretrained_path:
         # Load a pretrained policy and override the config if needed (for example, if there are inference-time
         # hyperparameters that we want to vary).

src/lerobot/policies/pi0/modeling_pi0.py

@@ -30,7 +30,7 @@ pip install -e ".[pi0]"
 
 Example of finetuning the pi0 pretrained model (`pi0_base` in `openpi`):
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
 --policy.path=lerobot/pi0 \
 --dataset.repo_id=danaaubakirova/koch_test
 ```
@@ -38,7 +38,7 @@ python -m lerobot.scripts.train \
 Example of finetuning the pi0 neural network with PaliGemma and expert Gemma
 pretrained with VLM default parameters before pi0 finetuning:
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
 --policy.type=pi0 \
 --dataset.repo_id=danaaubakirova/koch_test
 ```
@@ -56,15 +56,18 @@ from collections import deque
 import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
+from transformers import AutoTokenizer
 
-from lerobot.constants import ACTION, OBS_LANGUAGE, OBS_STATE
+from lerobot.constants import ACTION, OBS_STATE
+from lerobot.policies.normalize import Normalize, Unnormalize
 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.utils import get_safe_dtype
+from lerobot.policies.utils import log_model_loading_keys
+from lerobot.utils.utils import get_safe_dtype, init_logging
 
 
 def create_sinusoidal_pos_embedding(
@@ -220,17 +223,28 @@ class PI0Policy(PreTrainedPolicy):
     def __init__(
         self,
         config: PI0Config,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
             config: Policy configuration class instance or None, in which case the default instantiation of
                     the configuration class is used.
+            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
+                that they will be passed with a call to `load_state_dict` before the policy is used.
         """
 
         super().__init__(config)
         config.validate_features()
         self.config = config
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
 
+        self.language_tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224")
         self.model = PI0FlowMatching(config)
 
         self.reset()
@@ -239,6 +253,99 @@ class PI0Policy(PreTrainedPolicy):
         """This should be called whenever the environment is reset."""
         self._action_queue = deque([], maxlen=self.config.n_action_steps)
 
+    @classmethod
+    def _transform_state_dict_keys(cls, state_dict: dict) -> dict:
+        """
+        Transform state dict keys to match expected model structure.
+
+        Transformations:
+        - model.paligemma_with_expert.paligemma.language_model.lm_head ->
+          model.paligemma_with_expert.paligemma.lm_head
+        - model.paligemma_with_expert.paligemma.language_model.model ->
+          model.paligemma_with_expert.paligemma.model.language_model
+        - model.paligemma_with_expert.paligemma.vision_tower ->
+          model.paligemma_with_expert.paligemma.model.vision_tower
+        - model.paligemma_with_expert.paligemma.multi_modal_projector ->
+          model.paligemma_with_expert.paligemma.model.multi_modal_projector
+
+        Also handles tied weights between lm_head.weight and
+        embed_tokens.weight.
+        """
+        import re
+
+        transformed_dict = {}
+
+        transformations = [
+            (
+                re.compile(r"\.paligemma_with_expert\.paligemma\.language_model\.lm_head"),
+                ".paligemma_with_expert.paligemma.lm_head",
+            ),
+            (
+                re.compile(r"\.paligemma_with_expert\.paligemma\.language_model\.model"),
+                ".paligemma_with_expert.paligemma.model.language_model",
+            ),
+            (
+                re.compile(r"\.paligemma_with_expert\.paligemma\.vision_tower"),
+                ".paligemma_with_expert.paligemma.model.vision_tower",
+            ),
+            (
+                re.compile(r"\.paligemma_with_expert\.paligemma\.multi_modal_projector"),
+                ".paligemma_with_expert.paligemma.model.multi_modal_projector",
+            ),
+        ]
+
+        for key, value in state_dict.items():
+            new_key = key
+            for pattern, replacement in transformations:
+                new_key = pattern.sub(replacement, new_key)
+            transformed_dict[new_key] = value
+
+        # Handle tied weights: lm_head.weight and embed_tokens.weight share memory
+        lm_head_key = None
+        embed_tokens_key = None
+
+        for key in transformed_dict:
+            if key.endswith(".paligemma_with_expert.paligemma.lm_head.weight"):
+                lm_head_key = key
+            elif key.endswith(".paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"):
+                embed_tokens_key = key
+            if lm_head_key and embed_tokens_key:
+                break
+
+        if lm_head_key and not embed_tokens_key:
+            embed_tokens_key = lm_head_key.replace(
+                ".lm_head.weight", ".model.language_model.embed_tokens.weight"
+            )
+            transformed_dict[embed_tokens_key] = transformed_dict[lm_head_key]
+        elif embed_tokens_key and not lm_head_key:
+            lm_head_key = embed_tokens_key.replace(
+                ".model.language_model.embed_tokens.weight", ".lm_head.weight"
+            )
+            transformed_dict[lm_head_key] = transformed_dict[embed_tokens_key]
+
+        return transformed_dict
+
+    @classmethod
+    def _load_as_safetensor(
+        cls, model: "PI0Policy", model_file: str, map_location: str, strict: bool
+    ) -> "PI0Policy":
+        """Override to apply key transformations before loading."""
+        from safetensors.torch import load_file
+
+        init_logging()
+        # Load the state dict from file safely
+        state_dict = load_file(model_file, device=map_location)
+
+        # Apply key transformations
+        transformed_state_dict = cls._transform_state_dict_keys(state_dict)
+
+        # Load the transformed state dict
+        msg = model.load_state_dict(transformed_state_dict, strict=strict)
+
+        # Log message
+        log_model_loading_keys(msg.missing_keys, msg.unexpected_keys)
+        return model
+
     def get_optim_params(self) -> dict:
         return self.parameters()
 
@@ -270,13 +377,14 @@ class PI0Policy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
 
+        batch = self.normalize_inputs(batch)
+
         # Action queue logic for n_action_steps > 1. When the action_queue is depleted, populate it by
         # querying the policy.
         if len(self._action_queue) == 0:
             images, img_masks = self.prepare_images(batch)
             state = self.prepare_state(batch)
-            lang_tokens = batch[f"{OBS_LANGUAGE}.tokens"]
-            lang_masks = batch[f"{OBS_LANGUAGE}.attention_mask"]
+            lang_tokens, lang_masks = self.prepare_language(batch)
 
             actions = self.model.sample_actions(
                 images, img_masks, lang_tokens, lang_masks, state, noise=noise
@@ -286,6 +394,8 @@ class PI0Policy(PreTrainedPolicy):
             original_action_dim = self.config.action_feature.shape[0]
             actions = actions[:, :, :original_action_dim]
 
+            actions = self.unnormalize_outputs({"action": actions})["action"]
+
             if self.config.adapt_to_pi_aloha:
                 actions = self._pi_aloha_encode_actions(actions)
 
@@ -300,10 +410,12 @@ class PI0Policy(PreTrainedPolicy):
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
             batch[ACTION] = self._pi_aloha_encode_actions_inv(batch[ACTION])
 
+        batch = self.normalize_inputs(batch)
+        batch = self.normalize_targets(batch)
+
         images, img_masks = self.prepare_images(batch)
         state = self.prepare_state(batch)
-        lang_tokens = batch[f"{OBS_LANGUAGE}.tokens"]
-        lang_masks = batch[f"{OBS_LANGUAGE}.attention_mask"]
+        lang_tokens, lang_masks = self.prepare_language(batch)
         actions = self.prepare_action(batch)
         actions_is_pad = batch.get("action_is_pad")
 
@@ -370,6 +482,26 @@ class PI0Policy(PreTrainedPolicy):
 
         return images, img_masks
 
+    def prepare_language(self, batch) -> tuple[Tensor, Tensor]:
+        """Tokenize the text input"""
+        device = batch[OBS_STATE].device
+        tasks = batch["task"]
+
+        # PaliGemma prompt has to end with a new line
+        tasks = [task if task.endswith("\n") else f"{task}\n" for task in tasks]
+
+        tokenized_prompt = self.language_tokenizer.__call__(
+            tasks,
+            padding="max_length",
+            padding_side="right",
+            max_length=self.config.tokenizer_max_length,
+            return_tensors="pt",
+        )
+        lang_tokens = tokenized_prompt["input_ids"].to(device=device)
+        lang_masks = tokenized_prompt["attention_mask"].to(device=device, dtype=torch.bool)
+
+        return lang_tokens, lang_masks
+
     def _pi_aloha_decode_state(self, state):
         # Flip the joints.
         for motor_idx in [1, 2, 8, 9]:
@@ -435,7 +567,7 @@ class PI0FlowMatching(nn.Module):
     └──────────────────────────────┘
     """
 
-    def __init__(self, config: PI0Config):
+    def __init__(self, config):
         super().__init__()
         self.config = config
 

src/lerobot/policies/pi0/processor_pi0.py

@@ -1,120 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 Physical Intelligence and 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.
-
-from typing import Any
-
-import torch
-
-from lerobot.configs.types import PolicyFeature
-from lerobot.policies.pi0.configuration_pi0 import PI0Config
-from lerobot.processor import (
-    DeviceProcessor,
-    NormalizerProcessor,
-    RobotProcessor,
-    ToBatchProcessor,
-    TokenizerProcessor,
-    UnnormalizerProcessor,
-)
-from lerobot.processor.pipeline import (
-    EnvTransition,
-    ProcessorStep,
-    ProcessorStepRegistry,
-    TransitionKey,
-)
-from lerobot.processor.rename_processor import RenameProcessor
-
-
-@ProcessorStepRegistry.register(name="pi0_new_line_processor")
-class Pi0NewLineProcessor(ProcessorStep):
-    """Add a new line to the end of the task if it doesn't have one.
-    This is required for the PaliGemma tokenizer.
-    """
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        # Check if complementary_data exists
-        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
-        if complementary_data is None or "task" not in complementary_data:
-            return transition
-
-        task = complementary_data["task"]
-        if task is None:
-            return transition
-
-        # Handle both string and list of strings
-        if isinstance(task, str):
-            # Single string: add newline if not present
-            if not task.endswith("\n"):
-                complementary_data["task"] = f"{task}\n"
-        elif isinstance(task, list) and all(isinstance(t, str) for t in task):
-            # List of strings: add newline to each if not present
-            complementary_data["task"] = [t if t.endswith("\n") else f"{t}\n" for t in task]
-        # If task is neither string nor list of strings, leave unchanged
-
-        return transition
-
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        """Add tokenized task features to the features."""
-        return features
-
-    def state_dict(self) -> dict[str, torch.Tensor]:
-        """Return state dictionary (empty for this processor)."""
-        return {}
-
-    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
-        """Load state dictionary (no-op for this processor)."""
-        pass
-
-    def reset(self) -> None:
-        """Reset processor state (no-op for this processor)."""
-        pass
-
-    def get_config(self) -> dict[str, Any]:
-        """Return configuration for serialization."""
-        return {}
-
-
-def make_pi0_processor(
-    config: PI0Config, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
-) -> tuple[RobotProcessor, RobotProcessor]:
-    # Add remaining processors
-    input_steps: list[ProcessorStep] = [
-        RenameProcessor(rename_map={}),  # To mimic the same processor as pretrained one
-        NormalizerProcessor(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        ToBatchProcessor(),
-        Pi0NewLineProcessor(),  # Add newlines before tokenization for PaliGemma
-        TokenizerProcessor(
-            tokenizer_name="google/paligemma-3b-pt-224",
-            max_length=config.tokenizer_max_length,
-            padding_side="right",
-            padding="max_length",
-        ),
-        DeviceProcessor(device=config.device),
-    ]
-
-    output_steps: list[ProcessorStep] = [
-        DeviceProcessor(device="cpu"),
-        UnnormalizerProcessor(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-
-    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
-        steps=output_steps, name="robot_postprocessor"
-    )

src/lerobot/policies/pi0fast/modeling_pi0fast.py

@@ -25,14 +25,14 @@ Disclaimer: It is not expected to perform as well as the original implementation
 
 Example of finetuning the pi0+FAST pretrained model (`pi0_fast_base` in `openpi`):
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
 --policy.path=lerobot/pi0fast_base \
 --dataset.repo_id=danaaubakirova/koch_test
 ```
 
 Example of training the pi0+FAST neural network with from scratch:
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
 --policy.type=pi0fast \
 --dataset.repo_id=danaaubakirova/koch_test
 ```
@@ -58,6 +58,7 @@ from transformers.cache_utils import HybridCache, StaticCache
 from transformers.models.auto import CONFIG_MAPPING
 
 from lerobot.constants import ACTION, OBS_STATE
+from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
 
@@ -145,6 +146,14 @@ class PI0FASTPolicy(PreTrainedPolicy):
         config.validate_features()
         self.config = config
 
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
         self.language_tokenizer = AutoProcessor.from_pretrained("google/paligemma-3b-pt-224")
         self.model = PI0FAST(config)
 
@@ -212,6 +221,8 @@ class PI0FASTPolicy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
 
+        batch = self.normalize_inputs(batch)
+
         # Action queue logic for n_action_steps > 1. When the action_queue is depleted, populate it by
         # querying the policy.
         if len(self._action_queue) == 0:
@@ -224,6 +235,8 @@ class PI0FASTPolicy(PreTrainedPolicy):
             ]  # self.config.max_action_dim  # self.config.action_feature.shape[0]
             actions = actions[:, :, :original_action_dim]
 
+            actions = self.unnormalize_outputs({"action": actions})["action"]
+
             if self.config.adapt_to_pi_aloha:
                 actions = self._pi_aloha_encode_actions(actions)
 
@@ -236,6 +249,8 @@ class PI0FASTPolicy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
             batch[ACTION] = self._pi_aloha_encode_actions_inv(batch[ACTION])
+        batch = self.normalize_inputs(batch)
+        batch = self.normalize_targets(batch)
         loss_dict = self.model.forward(batch)
         return loss_dict["loss"], loss_dict
 

src/lerobot/policies/pi0fast/processor_pi0fast.py

@@ -1,51 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 Physical Intelligence and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import torch
-
-from lerobot.policies.pi0.configuration_pi0 import PI0Config
-from lerobot.processor import (
-    DeviceProcessor,
-    NormalizerProcessor,
-    RenameProcessor,
-    RobotProcessor,
-    ToBatchProcessor,
-    UnnormalizerProcessor,
-)
-
-
-def make_pi0fast_processor(
-    config: PI0Config, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
-) -> tuple[RobotProcessor, RobotProcessor]:
-    input_steps = [
-        RenameProcessor(rename_map={}),  # To mimic the same processor as pretrained one
-        NormalizerProcessor(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        ToBatchProcessor(),
-        DeviceProcessor(device=config.device),
-    ]
-    output_steps = [
-        DeviceProcessor(device="cpu"),
-        UnnormalizerProcessor(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
-        steps=output_steps, name="robot_postprocessor"
-    )

src/lerobot/policies/rlearn/configuration_rlearn.py

@@ -1,128 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from dataclasses import dataclass, field
-
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import NormalizationMode
-
-
-@PreTrainedConfig.register_subclass("rlearn")
-@dataclass
-class RLearNConfig(PreTrainedConfig):
-    """Configuration for a video-language conditioned reward model (RLearN).
-
-    Inputs:
-      - Visual frames (one or multiple cameras). Optionally a short sequence.
-      - A language instruction/goal string.
-
-    Output:
-      - Per-timestep reward logits or a single-step reward logit.
-
-    Notes:
-      - This follows the ReWiND paper architecture. It uses frozen vision/text encoders
-        (DINOv3 for vision, SigLIP2 for language) and trains a
-        lightweight temporal aggregator + head.
-    """
-
-    # Encoders - Use SigLIP2 for both vision and text (shared checkpoint)
-    vision_model_name: str = "google/siglip2-base-patch16-224"
-    text_model_name: str = "google/siglip2-base-patch16-224"
-    freeze_backbones: bool = True
-
-    # Sequence length, amount of past frames including current one to use in the temporal model
-    max_seq_len: int = 16
-    # Temporal sampling stride
-    temporal_sampling_stride: int = 3 # Open x mostly has fps 10, and rewind has seq len 16, ours is 30fps so 30/10 = 3 stride lenght to have same timeframe!
-
-    # Model dimensions and transformer
-    dim_model: int = 512
-    num_layers: int = 4
-    num_heads: int = 8
-    ff_mult: int = 4  # Feed-forward multiplier, hidden = dim_model * ff_mult
-    dropout: float = 0.05
-
-    # --- reward head options ---
-    use_categorical_rewards: bool = False      # classification over bins
-    num_reward_bins: int = 25
-    reward_min_value: float = 0.0              # for HL-Gauss range
-    reward_max_value: float = 1.0
-    use_hl_gauss_loss: bool = True             # if False -> plain regression
-    hl_gauss_num_bins: int = 25                # histogram resolution
-
-    # Inference-time subsampling and regularization
-    inference_stride: int = 1 # inference_stride is an extra, second downsampling applied in forward after window sampling/rewind. Keep it at 1 to disable extra skipping
-    frame_dropout_p: float = 0.10
-
-    # Training
-    learning_rate: float = 5e-4
-    weight_decay: float = 0.01
-    head_lr_multiplier: float = 5.0
-    logit_eps: float = 1e-4
-    regularizer_warmup_steps: int = 500
-    
-    # Performance optimizations
-    use_amp: bool = False
-    compile_model: bool = True
-
-    # ReWiND augmentation
-    rewind_prob: float = 0.3 #0.8
-    rewind_last3_prob: float = 0.0 #0.3
-    mismatch_prob: float = 0.0 #0.2
-    
-    # Normalization presets
-    normalization_mapping: dict[str, NormalizationMode] = field(
-        default_factory=lambda: {
-            "VISUAL": NormalizationMode.MEAN_STD,
-        }
-    )
-
-    # Required path to episodes.jsonl for episode boundaries
-    episodes_jsonl_path: str | None = "meta/episodes.jsonl"
-
-    def validate_features(self) -> None:
-        # Require at least one image feature. Language is recommended but optional (can be blank).
-        if not self.image_features:
-            raise ValueError(
-                "You must provide at least one image feature for RLearN (e.g. 'observation.image')."
-            )
-
-    @property
-    def observation_delta_indices(self) -> list | None:
-        # Request a long enough context so in-window stride sampling can be >1.
-        # We ask for (max_seq_len * temporal_sampling_stride) frames ending at t=0.
-        # Example: max_seq_len=16, temporal_sampling_stride=3 → 48 deltas → ~46 frames available.
-        total_needed = self.max_seq_len * max(1, int(self.temporal_sampling_stride))
-        return list(range(1 - total_needed, 1))
-
-    @property
-    def action_delta_indices(self) -> list | None:
-        # Not an action chunking policy.
-        return None
-
-    @property
-    def reward_delta_indices(self) -> list | None:
-        # ReWiND generates progress labels on-the-fly, doesn't need reward data
-        return None
-
-    def get_optimizer_preset(self):  # type: ignore[override]
-        from lerobot.optim.optimizers import AdamWConfig
-
-        return AdamWConfig(lr=self.learning_rate, weight_decay=self.weight_decay)
-
-    def get_scheduler_preset(self):  # type: ignore[override]
-        # No scheduler by default.
-        return None

src/lerobot/policies/rlearn/eval_script.py

@@ -1,392 +0,0 @@
-#!/usr/bin/env python
-
-"""
-Standalone evaluation script for RLearN models.
-
-This script evaluates RLearN reward models on episodes from a dataset,
-generating comparison plots between ground truth rewards and model predictions.
-
-Usage:
-    python src/lerobot/policies/rlearn/eval_script.py --model MODEL_NAME --dataset DATASET_REPO --episodes N
-
-Example:
-    python src/lerobot/policies/rlearn/eval_script.py --model pepijn223/rlearn_18 --dataset pepijn223/phone_pipeline_pickup1 --episodes 2
-"""
-
-import argparse
-import os
-import sys
-from pathlib import Path
-
-# Add src to path for imports
-sys.path.append(str(Path(__file__).parent.parent.parent.parent))
-
-import warnings
-
-import matplotlib.pyplot as plt
-import numpy as np
-import torch
-from scipy.stats import spearmanr
-from tqdm import tqdm
-
-warnings.filterwarnings("ignore")
-
-# LeRobot imports
-from lerobot.constants import OBS_IMAGE, OBS_IMAGES, OBS_LANGUAGE
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.policies.rlearn.modeling_rlearn import RLearNPolicy
-
-
-def _to_chw_float01(img):
-    """Ensure CHW float in [0,1]."""
-    if isinstance(img, np.ndarray):
-        img = torch.from_numpy(img)
-    # HWC -> CHW if needed
-    if len(img.shape) == 3 and img.shape[-1] in (1, 3, 4):
-        img = img.permute(2, 0, 1)
-    if img.dtype == torch.uint8:
-        img = img.float() / 255.0
-    else:
-        img = img.float()
-    return torch.clamp(img, 0.0, 1.0)
-
-
-def _get_language(frame_data):
-    lang = None
-    if OBS_LANGUAGE in frame_data:
-        lang = frame_data[OBS_LANGUAGE]
-        if isinstance(lang, list) and len(lang) > 0:
-            lang = lang[0]
-    elif "task" in frame_data:
-        lang = frame_data["task"]
-    return lang if isinstance(lang, str) else "No language provided"
-
-
-def _get_ground_truth_reward(frame_data):
-    """Try common keys for ground-truth reward. Return None if unavailable."""
-    for key in ("reward", "rewards", "gt_reward", "progress"):
-        if key in frame_data:
-            r = frame_data[key]
-            # unwrap single-element lists/arrays
-            if isinstance(r, (list, np.ndarray)) and np.array(r).size == 1:
-                r = float(np.array(r).reshape(-1)[0])
-            try:
-                return float(r)
-            except Exception:
-                pass
-    return None
-
-
-def extract_episode_frames_and_gt(dataset, episode_idx):
-    """Load a full episode: frames (T, C, H, W), language (str), gt_rewards (np.ndarray or None)."""
-    ep_start = dataset.episode_data_index["from"][episode_idx].item()
-    ep_end = dataset.episode_data_index["to"][episode_idx].item()
-    T = ep_end - ep_start
-
-    frames = []
-    gt_rewards = []
-    language = None
-
-    for t in range(T):
-        item = dataset[ep_start + t]
-
-        # image(s)
-        if OBS_IMAGES in item:
-            img = item[OBS_IMAGES]
-        elif OBS_IMAGE in item:
-            img = item[OBS_IMAGE]
-        else:
-            # try to find an image-like key
-            img_keys = [k for k in item.keys() if "image" in k.lower()]
-            if not img_keys:
-                continue
-            img = item[img_keys[0]]
-
-        frames.append(_to_chw_float01(img))
-
-        # language once
-        if language is None:
-            language = _get_language(item)
-
-        # ground-truth reward (optional)
-        r = _get_ground_truth_reward(item)
-        gt_rewards.append(r)
-
-    if not frames:
-        return None, None, None
-
-    frames = torch.stack(frames)  # (T, C, H, W)
-
-    # If all GT entries are None, treat as missing
-    if all(r is None for r in gt_rewards):
-        gt_rewards = None
-    else:
-        # Replace None by forward filling
-        arr = np.array([np.nan if r is None else float(r) for r in gt_rewards], dtype=float)
-        # forward/back fill
-        if np.isnan(arr[0]):
-            first_valid = np.flatnonzero(~np.isnan(arr))
-            if len(first_valid) > 0:
-                arr[0] = arr[first_valid[0]]
-            else:
-                arr[0] = 0.0
-        for i in range(1, len(arr)):
-            if np.isnan(arr[i]):
-                arr[i] = arr[i - 1]
-        gt_rewards = arr
-
-    return frames, language or "No language provided", gt_rewards
-
-
-@torch.no_grad()
-def predict_rewards_sliding(model, frames, language, max_seq_len=16, batch_size=64, device="cuda", temporal_stride: int | None = None):
-    """
-    Sliding-window prediction: for each frame i, create a window [max(0, i-L+1) .. i],
-    left-pad by repeating the first frame to length L (<= 16), and take the prediction 
-    corresponding to the current frame's position in the window.
-    Returns np.ndarray of shape (T,).
-    """
-    T = frames.shape[0]
-    cfg = getattr(model, "config", object())
-    L = int(getattr(cfg, "max_seq_len", max_seq_len))
-    L = min(L, max_seq_len)  # hard-cap at 16
-    # Use the same temporal stride as training (skip s-1 frames, take 1)
-    if temporal_stride is None:
-        temporal_stride = int(getattr(cfg, "temporal_sampling_stride", 1))
-    temporal_stride = max(1, int(temporal_stride))
-
-    # Preprocessed tensor on device
-    frames = frames.to(device)
-
-    windows = []
-    frame_positions = []  # Track which temporal position each frame should use
-    left_pad_counts = []  # Number of left-pad (OOB) frames per window
-    
-    for i in range(T):
-        # Build indices with stride s: [..., i-3, i] etc., left-padded by clamping to 0
-        idxs = [i - (L - 1 - j) * temporal_stride for j in range(L)]
-        pad_needed = sum(1 for k in idxs if k < 0)
-        clamped = [0 if k < 0 else (T - 1 if k >= T else k) for k in idxs]
-        window = frames[clamped]  # (L, C, H, W)
-
-        # Use the last temporal position (current frame) for reading model output
-        frame_pos = L - 1
-
-        windows.append(window)
-        frame_positions.append(frame_pos)
-        left_pad_counts.append(pad_needed)
-
-    preds = np.zeros(T, dtype=float)
-
-    for s in range(0, T, batch_size):
-        e = min(s + batch_size, T)
-        batch_windows = torch.stack(windows[s:e])  # (B, L, C, H, W)
-        batch_positions = frame_positions[s:e]
-
-        batch = {OBS_IMAGES: batch_windows, OBS_LANGUAGE: [language] * (e - s)}  # expects (B, L, C, H, W)
-
-        # Model returns (B, L) predictions for each temporal position
-        values = model.predict_rewards(batch)  # torch.Tensor (B, L)
-
-        # Apply eval-time padding rule: predictions for left-padded (OOB) frames are zero
-        if values.dim() == 2 and len(left_pad_counts) >= (e - s):
-            for b_idx in range(e - s):
-                pad_n = left_pad_counts[s + b_idx]
-                if pad_n > 0:
-                    values[b_idx, :pad_n] = 0.0
-
-        # Debug output removed - issue was identified and fixed
-
-        if values.dim() == 2:
-            # Extract the prediction corresponding to each frame's position in its window
-            batch_preds = []
-            for b_idx, pos in enumerate(batch_positions):
-                batch_preds.append(values[b_idx, pos].item())
-            preds[s:e] = np.array(batch_preds)
-        else:
-            # Fallback: if model returns (B,), use as is
-            preds[s:e] = values.detach().float().cpu().numpy()
-
-    return preds
-
-
-def plot_episode_eval(episode_idx, gt, pred, language, save_path=None, show=False, title_prefix="RLearN Eval"):
-    """Plot GT vs Predicted over time. Saves PNG if save_path is provided."""
-    T = len(pred)
-    x = np.arange(T)
-
-    plt.figure(figsize=(14, 8))
-    plt.plot(x, pred, linewidth=2.5, marker="o", markersize=3, label="Predicted Reward", color="blue")
-
-    if gt is not None:
-        plt.plot(x, gt, linestyle="--", linewidth=2.5, label="Ground-Truth Reward", color="orange")
-        # Correlation between GT and Pred
-        corr, p = spearmanr(gt, pred)
-        corr_str = f"ρ(GT, Pred) = {0.0 if np.isnan(corr) else corr:.3f} (p={0.0 if np.isnan(p) else p:.3f})"
-    else:
-        expected = np.linspace(0, 1, T)
-        plt.plot(x, expected, linestyle="--", linewidth=2.5, label="Expected Progress (0→1)", color="orange")
-        corr, p = spearmanr(x, pred)
-        corr_str = f"VOC-S ρ(t, Pred) = {0.0 if np.isnan(corr) else corr:.3f} (p={0.0 if np.isnan(p) else p:.3f})"
-
-    plt.title(f"{title_prefix} — Episode {episode_idx}\n{language}\n{corr_str}", fontsize=14)
-    plt.xlabel("Frame Index", fontsize=12)
-    plt.ylabel("Reward / Progress", fontsize=12)
-    plt.legend(fontsize=11)
-    plt.grid(True, alpha=0.3)
-    plt.tight_layout()
-
-    if save_path is not None:
-        plt.savefig(save_path, dpi=200, bbox_inches="tight")
-        print(f"Saved eval image to: {save_path}")
-
-    if show:
-        plt.show()
-    else:
-        plt.close()
-
-
-def eval_episode_sliding(
-    episode_idx, dataset, model, save_dir=".", device="cuda", max_seq_len=16, batch_size=64, title_prefix="RLearN Eval"
-):
-    """End-to-end: load episode, predict with sliding 16-frame windows, and save PNG."""
-    frames, language, gt = extract_episode_frames_and_gt(dataset, episode_idx)
-    if frames is None:
-        print(f"[Episode {episode_idx}] No frames found.")
-        return None
-
-    model.eval()
-
-    pred = predict_rewards_sliding(
-        model=model, frames=frames, language=language, max_seq_len=max_seq_len, batch_size=batch_size, device=device
-    )
-
-    # Basic stats
-    print(f"Episode {episode_idx}: T={len(pred)}, pred∈[{pred.min():.3f},{pred.max():.3f}]")
-    if gt is not None:
-        print(f"GT available: gt∈[{np.nanmin(gt):.3f},{np.nanmax(gt):.3f}]")
-
-    save_path = f"{save_dir}/episode_{episode_idx:04d}_eval.png"
-    plot_episode_eval(
-        episode_idx=episode_idx, gt=gt, pred=pred, language=language, save_path=save_path, show=False, title_prefix=title_prefix
-    )
-    return save_path
-
-
-def main():
-    """Main evaluation script for RLearN models."""
-    # Parse command line arguments
-    parser = argparse.ArgumentParser(description="Evaluate RLearN model on episodes with GT vs Predicted rewards")
-    parser.add_argument("--model", type=str, required=True, help="Model name/path (e.g., pepijn223/rlearn_mse5)")
-    parser.add_argument("--dataset", type=str, required=True, help="Dataset repo (e.g., pepijn223/phone_pipeline_pickup1)")
-    parser.add_argument("--episodes", type=int, default=5, help="Number of episodes to evaluate")
-    parser.add_argument("--output", type=str, default="./eval_results", help="Output directory for images")
-    parser.add_argument(
-        "--device",
-        type=str,
-        default="cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu",
-        help="Device to use",
-    )
-    parser.add_argument("--batch_size", type=int, default=32, help="Batch size for sliding window evaluation")
-
-    args = parser.parse_args()
-
-    # Create output directory
-    output_dir = Path(args.output)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    print("🎯 RLearN Model Evaluation")
-    print("=" * 60)
-    print(f"Model: {args.model}")
-    print(f"Dataset: {args.dataset}")
-    print(f"Episodes: {args.episodes}")
-    print(f"Device: {args.device}")
-    print(f"Output: {output_dir}")
-    print("=" * 60)
-
-    try:
-        # Load dataset
-        print("📁 Loading dataset...")
-
-        dataset = LeRobotDataset(
-            repo_id=args.dataset,
-            episodes=list(range(min(args.episodes, 50))),  # Load enough episodes
-            download_videos=True,
-        )
-
-        print(f"✅ Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
-        print(f"   Features: {list(dataset.features.keys())}")
-        print(f"   FPS: {dataset.fps}")
-
-        # Load model
-        print("\n🤖 Loading model...")
-
-        model = RLearNPolicy.from_pretrained(args.model)
-        model = model.to(args.device)
-        model.eval()
-
-        print(f"✅ Model loaded on {args.device}")
-        print(f"   Parameters: {sum(p.numel() for p in model.parameters()):,}")
-        print(f"   Trainable: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
-        print(f"   Max sequence length: {model.config.max_seq_len}")
-
-        # Select episodes to evaluate
-        total_available = min(dataset.num_episodes, args.episodes)
-        episode_indices = list(range(total_available))
-
-        print(f"\n📊 Evaluating {len(episode_indices)} episodes...")
-        print("=" * 60)
-
-        # Run sliding window evaluation on each episode
-        saved_paths = []
-        for i, ep_idx in enumerate(episode_indices):
-            print(f"\n[{i+1}/{len(episode_indices)}] Processing Episode {ep_idx}")
-            print("-" * 40)
-
-            try:
-                save_path = eval_episode_sliding(
-                    episode_idx=ep_idx,
-                    dataset=dataset,
-                    model=model,
-                    save_dir=str(output_dir),
-                    device=args.device,
-                    batch_size=args.batch_size,
-                    title_prefix="RLearN Ground Truth vs Predicted",
-                )
-
-                if save_path:
-                    saved_paths.append(save_path)
-
-            except Exception as e:
-                print(f"❌ Error processing episode {ep_idx}: {e}")
-                import traceback
-
-                traceback.print_exc()
-                continue
-
-        # Summary
-        print("\n" + "=" * 60)
-        print("✅ EVALUATION COMPLETE")
-        print(f"📈 Generated {len(saved_paths)} evaluation plots")
-        print(f"📁 Results saved to: {output_dir}")
-        print("\nGenerated files:")
-        for path in saved_paths:
-            print(f"  • {path}")
-
-        if saved_paths:
-            print(f"\n💡 View the plots to compare ground truth vs predicted rewards!")
-            print(f"   Each plot shows the model's sliding 16-frame window predictions")
-            print(f"   against available ground truth rewards over the episode timeline.")
-
-        return 0
-
-    except Exception as e:
-        print(f"❌ Error during evaluation: {e}")
-        import traceback
-
-        traceback.print_exc()
-        return 1
-
-
-if __name__ == "__main__":
-    exit(main())

src/lerobot/policies/rlearn/processor_rlearn.py

@@ -1,128 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from dataclasses import dataclass
-from typing import Any
-
-from lerobot.configs.types import PolicyFeature
-from lerobot.constants import OBS_LANGUAGE
-from lerobot.policies.rlearn.configuration_rlearn import RLearNConfig
-from lerobot.processor import (
-    DeviceProcessor,
-    NormalizerProcessor,
-    RenameProcessor,
-    RobotProcessor,
-    ToBatchProcessor,
-    TokenizerProcessor,
-    UnnormalizerProcessor,
-)
-from lerobot.processor.pipeline import (
-    ComplementaryDataProcessor,
-    EnvTransition,
-    ProcessorStepRegistry,
-    TransitionKey,
-)
-
-
-def make_rlearn_processor(
-    config: RLearNConfig, dataset_stats: dict[str, dict[str, Any]] | None = None
-) -> tuple[RobotProcessor, RobotProcessor]:
-    """Build pre/post processors for RLearN.
-
-    Responsibilities moved out of the model:
-      - Normalize inputs (images) using dataset stats
-      - Ensure batching
-      - Map complementary_data.task to observation.language when available
-      - Tokenize language into observation.language.tokens / attention_mask
-      - Move to/from device
-    """
-
-    input_steps = [
-        # No renaming by default, but keep for future extensibility
-        RenameProcessor(rename_map={}),
-        # Move heavy normalization to GPU after transfer for better parallelism
-        ToBatchProcessor(),
-        RLearnLanguageFromTaskProcessor(),
-        # Use SigLIP2 for tokenizer to keep vocab aligned with text tower
-        TokenizerProcessor(
-            tokenizer_name=config.text_model_name,
-            max_length=64,
-            padding="max_length",
-            truncation=True,
-            padding_side="right",
-        ),
-        DeviceProcessor(device=config.device),
-        # Move normalization after GPU transfer to use GPU acceleration
-        NormalizerProcessor(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-    ]
-
-    output_steps = [
-        DeviceProcessor(device="cpu"),
-        UnnormalizerProcessor(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-
-    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
-        steps=output_steps, name="robot_postprocessor"
-    )
-
-
-@dataclass
-@ProcessorStepRegistry.register(name="rlearn_language_from_task")
-class RLearnLanguageFromTaskProcessor(ComplementaryDataProcessor):
-    """Copy complementary_data['task'] into observation['observation.language'] if present.
-
-    This ensures the model can consume a raw language string when tokenization is not used,
-    while TokenizerProcessor can still create tokenized fields.
-    """
-
-    task_key: str = "task"
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:  # type: ignore[override]
-        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
-        if not complementary_data or self.task_key not in complementary_data:
-            return transition
-
-        task = complementary_data.get(self.task_key)
-        if task is None:
-            return transition
-
-        # Normalize to list[str]
-        if isinstance(task, str):
-            task_list = [task]
-        elif isinstance(task, list) and all(isinstance(t, str) for t in task):
-            task_list = task
-        else:
-            return transition
-
-        observation = transition.get(TransitionKey.OBSERVATION) or {}
-        # Do not overwrite if user already provided observation.language
-        if OBS_LANGUAGE not in observation:
-            observation[OBS_LANGUAGE] = task_list
-            transition[TransitionKey.OBSERVATION] = observation
-        return transition
-
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:  # noqa: D401
-        # Adds nothing to features; only mirrors complementary_data.task into observation
-        return features
-
-    def get_config(self) -> dict[str, Any]:
-        return {"task_key": self.task_key}

src/lerobot/policies/rlearn/rlearn_plan.md

@@ -1,101 +0,0 @@
-## General Value/Reward Learning:
-
-I want to implement a general/universal vision and language value function or reward model for robotics/video tasks. Also called a video language conditioned reward model. Integrated with already existing LeRobot code if convenient, use the LeRobot Dataset for dataset and store the reward for a frame in the lerobot frame itself.
-
-Inspired by these papers:
-
-- ReWiND; https://arxiv.org/pdf/2505.10911 (Most applicable and main paper I want to implement ideas from) and code: https://github.com/lucidrains/rewind-reward-pytorch
-- LIV; https://arxiv.org/pdf/2306.00958 (Most applicable and 2nd main paper I want to implement ideas from) and code https://github.com/penn-pal-lab/LI
-- VLC: Video-Language Critic: Transferable Reward Functions for Language-Conditioned Robotics: https://arxiv.org/pdf/2405.19988 (Most applicable and 3rd paper I want to implement ideas from) and code: https://github.com/minttusofia/video_language_critic
-
-And these papers which are also relevant:
-
-- https://www.dyna.co/dyna-1/research (Main company I want to reproduce the eventual results from)
-- vip; https://arxiv.org/pdf/2210.00030
-- uvd; https://arxiv.org/pdf/2310.08581
-- vlm in context; https://arxiv.org/pdf/2411.04549
-- https://www.youtube.com/watch?v=JfZYtpEisoM
-
-Little less relevant but still similar papers:
-
-- Learning Generalizable Robotic Reward Functions from “In-The-Wild” Human Videos,
-- XIRL: Cross-embodiment Inverse Reinforcement Learning,
-- Video-Language Critic: Transferable Reward https://arxiv.org/pdf/2405.19988
-- Functions for Language-Conditioned Robotics,
-- LORel, Language-Driven Representation Learning for Robotics https://sites.google.com/view/robotlorel
-- RoboCLIP: One Demonstration is Enough to Learn Robot Policies https://arxiv.org/pdf/2310.07899
-- Points2Rewards: learn first key points and then uses the keypoints to learn general value function/policy https://semrob.github.io/docs/2025_rss_semrob.github.io_paper20.pdf
-- Language-Driven Representation Learning for Robotics: https://arxiv.org/pdf/2302.12766v1
-- R3M: A Universal Visual Representation for Robot Manipulation: https://arxiv.org/pdf/2203.12601v3
-
-Input should be the current image or whole video and the task goal specified in text/language. Output is current reward.
-Archiutecture:
-_ inputs: video o1:T (or current o1:t), language z;
-_ DINO v3 ViT-B/16 (86M params): https://huggingface.co/facebook/dinov3-vitb16-pretrain-lvd1689m for vision encoding
-\_ sentence-transformers/all-MiniLM-L12-v2: https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2 for text encoding \* Temporal module: small causal transformer ("cross-modal sequential aggregator"), with first-frame positional embedding (to avoid position cheating), frame-dropout, and stride sampling; outputs per-timestep logits.
-
-Loss: See this chatgpt thread: https://chatgpt.com/s/t_68999a50a0b081919abc365cdd205e01
-
-Past images: (for example a reward method go to 3rd floor, has to know what floor it was on and what pas actions it did, can we attend or encorperate images of decision from history in one way?) Maybe via this paper: Learning Long-Context Diffusion Policies via Past-Token Prediction
-
-Amount of frames needed for test/generalization: 1M frames? or ~20% of IPEC-COMMUNITY/bc_z_lerobot
-
-Eval:
-Implement something like voc score , or ROC rank order correlation between reward leanredna and ev reward from sim, or use something else to do additional evaluation
-
-Ideas:
-
-- Incorporate training on multiple horizons: as in label same dataset for longer horizons: make a sandwich (long), put cheese on bread (medium) and even smaller horizons: go down or close gripper (small)
-- Incorporate navigation goals “walk towards the kitchen”, make sure we fix CLIP contrastive learning issue of positional text misunderstanding where model doesnnt learn difference between "horse right of cow" and "horse left of cow" “Move right” potentially train with more other data or even actionable world models such as Genie 3 (https://deepmind.google/discover/blog/genie-3-a-new-frontier-for-world-models/)
-
-How to use a general reward model (use cases): - Train rl policy on it - Success detection - Do exploraion - Do task via planning and search to optimize reward - Filter out bad episodes in large datasets from imitation learning
-
-Potential Datasets: (start with dataset that is most clean for this and works best with chosen way of doing evals)
-_ Epic-Kitchens-100
-_ Something-Something v. 2 Dataset https://www.qualcomm.com/developer/software/something-something-v-2-dataset
-_ Ego4D (3000 hours)
-_ Open X-Embodiment (OXE)
-\_ Agi bot world: https://huggingface.co/datasets/agibot-world/AgiBotWorld-Alpha
-
-- GalexiAI dataset: https://opengalaxea.github.io/G0/
-  _ GTEA+ Gaze: https://cbs.ic.gatech.edu/fpv/
-  _ YouCook2 dataset
-  \_ HOWTO100M: https://www.di.ens.fr/willow/research/howto100m/
-- Genie generated dataset?
-
-### TODOs:
-
-- Implement first architecture [x]
-- Implement processors [x]
-- Choose right loss metric(s) [x]
-- Make dataset with script that generated the dataset (IPEC-COMMUNITY/bc_z_lerobot) ready in lerobot format (and be able to visualize in dataset visualizer)
-  - Annotate with ReWiND-style 0→1 progress rewards [x]
-  - Visualize to check [x]
-- Implement eval score or metric that is robust and can deal with generalization/is a good metric to try different architectures. And use it in an eval jupyter notebook with visalization of the live reward next to the video for part of the dataset: VOC score and score with correct and incorrect language captions [x]
-- Do first training [x]
-- Implement on-the-fly progress label generation (no need for pre-annotated rewards) [x]
-- Try different losses
-  - Only rewind loss [x]
-  - Exactly similar to: https://github.com/lucidrains/rewind-reward-pytorch/blob/main/rewind_reward_pytorch/rewind_reward.py#L11 [x]
-  - Try DINO v2 as encoder Base 86 M: with https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2 [x]
-  - Test rewind (evaluate) [x]
-- benchmark siglip 2 vs this implementation forward pass, debug speed [x]
-- use siglip 2 [x]
-- Fix evaluation bug !!! []
-- Fix sample episode padding bug !!! []
-- Overfit on one episode []
-- Cleanup code? [] + enable language loss
-- Convert python -m lerobot.datasets.v21.convert_dataset_v20_to_v21 --repo-id=IPEC-COMMUNITY/bc_z_lerobot and train on 1 percent
-- Then on 10 percent []
-- Ablation 16 sucessive frame vs 16 frame samples with stride 2 or 4 []
-- Add more artificial text to dataset generated by vlm (google gemini) []
-  - See google gemini vlm caption [] https://gemini.google.com/app/7e332ffaf32580f2
-  - Multiple captions per video, creat method to generate as much data as possible etc [] https://arxiv.org/abs/2508.13446, https://arxiv.org/pdf/2412.04453
-- Add other datasets from OXE metioned in rewind []
-- Extend evaluation []
-- Ablation for size vision encoder, language encoder, temporal head []
-- Ablation one mlp head per frame or single mlp head []
-- Add other datasets metnioned here []
-- How can we improve spatial aware learning? solve issue of Contrastive learning and position []
-
-

src/lerobot/policies/sac/modeling_sac.py

@@ -28,6 +28,7 @@ import torch.nn.functional as F  # noqa: N812
 from torch import Tensor
 from torch.distributions import MultivariateNormal, TanhTransform, Transform, TransformedDistribution
 
+from lerobot.policies.normalize import NormalizeBuffer
 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
@@ -44,6 +45,7 @@ class SACPolicy(
     def __init__(
         self,
         config: SACConfig | None = None,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         super().__init__(config)
         config.validate_features()
@@ -51,6 +53,7 @@ class SACPolicy(
 
         # Determine action dimension and initialize all components
         continuous_action_dim = config.output_features["action"].shape[0]
+        self._init_normalization(dataset_stats)
         self._init_encoders()
         self._init_critics(continuous_action_dim)
         self._init_actor(continuous_action_dim)
@@ -85,7 +88,8 @@ class SACPolicy(
 
         observations_features = None
         if self.shared_encoder and self.actor.encoder.has_images:
-            observations_features = self.actor.encoder.get_cached_image_features(batch)
+            # Cache and normalize image features
+            observations_features = self.actor.encoder.get_cached_image_features(batch, normalize=True)
 
         actions, _, _ = self.actor(batch, observations_features)
 
@@ -387,12 +391,28 @@ class SACPolicy(
         actor_loss = ((self.temperature * log_probs) - min_q_preds).mean()
         return actor_loss
 
+    def _init_normalization(self, dataset_stats):
+        """Initialize input/output normalization modules."""
+        self.normalize_inputs = nn.Identity()
+        self.normalize_targets = nn.Identity()
+        if self.config.dataset_stats is not None:
+            params = _convert_normalization_params_to_tensor(self.config.dataset_stats)
+            self.normalize_inputs = NormalizeBuffer(
+                self.config.input_features, self.config.normalization_mapping, params
+            )
+            stats = dataset_stats or params
+            self.normalize_targets = NormalizeBuffer(
+                self.config.output_features, self.config.normalization_mapping, stats
+            )
+
     def _init_encoders(self):
         """Initialize shared or separate encoders for actor and critic."""
         self.shared_encoder = self.config.shared_encoder
-        self.encoder_critic = SACObservationEncoder(self.config)
+        self.encoder_critic = SACObservationEncoder(self.config, self.normalize_inputs)
         self.encoder_actor = (
-            self.encoder_critic if self.shared_encoder else SACObservationEncoder(self.config)
+            self.encoder_critic
+            if self.shared_encoder
+            else SACObservationEncoder(self.config, self.normalize_inputs)
         )
 
     def _init_critics(self, continuous_action_dim):
@@ -404,7 +424,9 @@ class SACPolicy(
             )
             for _ in range(self.config.num_critics)
         ]
-        self.critic_ensemble = CriticEnsemble(encoder=self.encoder_critic, ensemble=heads)
+        self.critic_ensemble = CriticEnsemble(
+            encoder=self.encoder_critic, ensemble=heads, output_normalization=self.normalize_targets
+        )
         target_heads = [
             CriticHead(
                 input_dim=self.encoder_critic.output_dim + continuous_action_dim,
@@ -412,7 +434,9 @@ class SACPolicy(
             )
             for _ in range(self.config.num_critics)
         ]
-        self.critic_target = CriticEnsemble(encoder=self.encoder_critic, ensemble=target_heads)
+        self.critic_target = CriticEnsemble(
+            encoder=self.encoder_critic, ensemble=target_heads, output_normalization=self.normalize_targets
+        )
         self.critic_target.load_state_dict(self.critic_ensemble.state_dict())
 
         if self.config.use_torch_compile:
@@ -466,9 +490,10 @@ class SACPolicy(
 class SACObservationEncoder(nn.Module):
     """Encode image and/or state vector observations."""
 
-    def __init__(self, config: SACConfig) -> None:
+    def __init__(self, config: SACConfig, input_normalizer: nn.Module) -> None:
         super().__init__()
         self.config = config
+        self.input_normalization = input_normalizer
         self._init_image_layers()
         self._init_state_layers()
         self._compute_output_dim()
@@ -543,10 +568,11 @@ class SACObservationEncoder(nn.Module):
     def forward(
         self, obs: dict[str, Tensor], cache: dict[str, Tensor] | None = None, detach: bool = False
     ) -> Tensor:
+        obs = self.input_normalization(obs)
         parts = []
         if self.has_images:
             if cache is None:
-                cache = self.get_cached_image_features(obs)
+                cache = self.get_cached_image_features(obs, normalize=False)
             parts.append(self._encode_images(cache, detach))
         if self.has_env:
             parts.append(self.env_encoder(obs["observation.environment_state"]))
@@ -559,7 +585,7 @@ class SACObservationEncoder(nn.Module):
             "No parts to concatenate, you should have at least one image or environment state or state"
         )
 
-    def get_cached_image_features(self, obs: dict[str, Tensor]) -> dict[str, Tensor]:
+    def get_cached_image_features(self, obs: dict[str, Tensor], normalize: bool = False) -> dict[str, Tensor]:
         """Extract and optionally cache image features from observations.
 
         This function processes image observations through the vision encoder once and returns
@@ -571,17 +597,26 @@ class SACObservationEncoder(nn.Module):
         - The vision encoder forward pass is typically the main computational bottleneck during training and inference
         - Caching these features can provide 2-4x speedup in training and inference
 
+        Normalization behavior:
+        - When called from inside forward(): set normalize=False since inputs are already normalized
+        - When called from outside forward(): set normalize=True to ensure proper input normalization
+
         Usage patterns:
-        - Called in select_action()
+        - Called in select_action() with normalize=True
         - Called in learner.py's get_observation_features() to pre-compute features for all policy components
-        - Called internally by forward()
+        - Called internally by forward() with normalize=False
 
         Args:
             obs: Dictionary of observation tensors containing image keys
+            normalize: Whether to normalize observations before encoding
+                      Set to True when calling directly from outside the encoder's forward method
+                      Set to False when calling from within forward() where inputs are already normalized
 
         Returns:
             Dictionary mapping image keys to their corresponding encoded features
         """
+        if normalize:
+            obs = self.input_normalization(obs)
         batched = torch.cat([obs[k] for k in self.image_keys], dim=0)
         out = self.image_encoder(batched)
         chunks = torch.chunk(out, len(self.image_keys), dim=0)
@@ -712,6 +747,7 @@ class CriticEnsemble(nn.Module):
     Args:
         encoder (SACObservationEncoder): encoder for observations.
         ensemble (List[CriticHead]): list of critic heads.
+        output_normalization (nn.Module): normalization layer for actions.
         init_final (float | None): optional initializer scale for final layers.
 
     Forward returns a tensor of shape (num_critics, batch_size) containing Q-values.
@@ -721,11 +757,13 @@ class CriticEnsemble(nn.Module):
         self,
         encoder: SACObservationEncoder,
         ensemble: list[CriticHead],
+        output_normalization: nn.Module,
         init_final: float | None = None,
     ):
         super().__init__()
         self.encoder = encoder
         self.init_final = init_final
+        self.output_normalization = output_normalization
         self.critics = nn.ModuleList(ensemble)
 
     def forward(
@@ -737,6 +775,11 @@ class CriticEnsemble(nn.Module):
         device = get_device_from_parameters(self)
         # Move each tensor in observations to device
         observations = {k: v.to(device) for k, v in observations.items()}
+        # NOTE: We normalize actions it helps for sample efficiency
+        actions: dict[str, torch.tensor] = {"action": actions}
+        # NOTE: Normalization layer took dict in input and outputs a dict that why
+        actions = self.output_normalization(actions)["action"]
+        actions = actions.to(device)
 
         obs_enc = self.encoder(observations, cache=observation_features)
 

src/lerobot/policies/sac/processor_sac.py

@@ -1,52 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team.
-# All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import torch
-
-from lerobot.policies.sac.configuration_sac import SACConfig
-from lerobot.processor import (
-    DeviceProcessor,
-    NormalizerProcessor,
-    RenameProcessor,
-    RobotProcessor,
-    ToBatchProcessor,
-    UnnormalizerProcessor,
-)
-
-
-def make_sac_processor(
-    config: SACConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
-) -> tuple[RobotProcessor, RobotProcessor]:
-    input_steps = [
-        RenameProcessor(rename_map={}),
-        NormalizerProcessor(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        ToBatchProcessor(),
-        DeviceProcessor(device=config.device),
-    ]
-    output_steps = [
-        DeviceProcessor(device="cpu"),
-        UnnormalizerProcessor(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
-        steps=output_steps, name="robot_postprocessor"
-    )

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

@@ -20,6 +20,7 @@ import torch
 from torch import Tensor, nn
 
 from lerobot.constants import OBS_IMAGE, REWARD
+from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
 
@@ -107,12 +108,22 @@ class Classifier(PreTrainedPolicy):
     def __init__(
         self,
         config: RewardClassifierConfig,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         from transformers import AutoModel
 
         super().__init__(config)
         self.config = config
 
+        # Initialize normalization (standardized with the policy framework)
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
         # Set up encoder
         encoder = AutoModel.from_pretrained(self.config.model_name, trust_remote_code=True)
         # Extract vision model if we're given a multimodal model
@@ -236,6 +247,10 @@ class Classifier(PreTrainedPolicy):
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Tensor]]:
         """Standard forward pass for training compatible with train.py."""
+        # Normalize inputs if needed
+        batch = self.normalize_inputs(batch)
+        batch = self.normalize_targets(batch)
+
         # Extract images and labels
         images, labels = self.extract_images_and_labels(batch)
 

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

@@ -1,42 +0,0 @@
-# !/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-
-from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
-from lerobot.processor import (
-    DeviceProcessor,
-    IdentityProcessor,
-    NormalizerProcessor,
-    RobotProcessor,
-)
-
-
-def make_classifier_processor(
-    config: RewardClassifierConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
-) -> tuple[RobotProcessor, RobotProcessor]:
-    input_steps = [
-        NormalizerProcessor(
-            features=config.input_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-        NormalizerProcessor(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-        DeviceProcessor(device=config.device),
-    ]
-    output_steps = [DeviceProcessor(device="cpu"), IdentityProcessor()]
-    return RobotProcessor(steps=input_steps, name="classifier_preprocessor"), RobotProcessor(
-        steps=output_steps, name="classifier_postprocessor"
-    )

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -28,7 +28,7 @@ pip install -e ".[smolvla]"
 
 Example of finetuning the smolvla pretrained model (`smolvla_base`):
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
 --policy.path=lerobot/smolvla_base \
 --dataset.repo_id=danaaubakirova/svla_so100_task1_v3 \
 --batch_size=64 \
@@ -38,7 +38,7 @@ python -m lerobot.scripts.train \
 Example of finetuning a smolVLA. SmolVLA is composed of a pretrained VLM,
 and an action expert.
 ```bash
-python -m lerobot.scripts.train \
+lerobot-train \
 --policy.type=smolvla \
 --dataset.repo_id=danaaubakirova/svla_so100_task1_v3 \
 --batch_size=64 \
@@ -53,13 +53,21 @@ policy = SmolVLAPolicy.from_pretrained("lerobot/smolvla_base")
 """
 
 import math
+import os
+import re
 from collections import deque
 
+import safetensors
 import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
+from transformers import AutoProcessor
 
-from lerobot.constants import ACTION, OBS_LANGUAGE, OBS_STATE
+from lerobot.constants import ACTION, OBS_STATE
+from lerobot.policies.normalize import (
+    Normalize,
+    Unnormalize,
+)
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.smolvla.smolvlm_with_expert import SmolVLMWithExpertModel
@@ -68,6 +76,102 @@ from lerobot.policies.utils import (
 )
 from lerobot.utils.utils import get_safe_dtype
 
+# Matches ".soNNN", optionally followed by "-something", up to the "_buffer_" marker
+_VARIANT_RE = re.compile(r"\.so\d+(?:-[\w]+)?_buffer_")
+
+
+def canonicalise(k: str) -> str:
+    """
+    Remove dataset-variant markers like '.so100-blue_' or '.so100_' from a
+    normalisation-buffer key.
+    """
+    return _VARIANT_RE.sub(".buffer_", k)
+
+
+def standardise_state_dict(
+    checkpoint: dict[str, torch.Tensor], ref_keys: set[str], *, verbose: bool = True
+) -> tuple[dict[str, torch.Tensor], list[str]]:
+    """
+    • Re-keys `checkpoint ` so that every entry matches the *reference* key set.
+    • If several variant keys collapse to the same canonical name we keep the
+      first one and log the collision.
+    • Returns the new dict + a list of entries that could not be matched.
+    """
+    out, collisions, unmatched = {}, {}, []
+
+    for k, v in checkpoint.items():
+        canon = canonicalise(k)
+        if canon in ref_keys:
+            if canon in out:  # duplicate after collapsing
+                collisions.setdefault(canon, []).append(k)
+            else:
+                out[canon] = v
+        else:
+            unmatched.append(k)
+
+    if verbose:
+        for canon, variants in collisions.items():
+            print(f"[standardise_state_dict] '{canon}'  ←  {variants}")
+        if unmatched:
+            print(f"[standardise_state_dict] kept {len(unmatched)} unmatched keys")
+
+    out.update({k: checkpoint[k] for k in unmatched})
+    return out, unmatched
+
+
+def rename_checkpoint_keys(checkpoint: dict, rename_str: str):
+    """
+    Renames keys in a checkpoint dictionary based on the given rename string.
+
+    Args:
+        checkpoint (dict): The checkpoint dictionary.
+        rename_str (str): A string specifying key mappings in the format "old1//new1,old2//new2".
+
+    Returns:
+        dict: The modified checkpoint with renamed keys.
+    """
+
+    rename_dict = dict(pair.split("//") for pair in rename_str.split(","))
+
+    new_checkpoint = {}
+    for k, v in checkpoint.items():
+        for old_key, new_key in rename_dict.items():
+            if old_key in k:
+                k = k.replace(old_key, new_key)
+        new_checkpoint[k] = v
+    return new_checkpoint
+
+
+def load_smolvla(
+    model: torch.nn.Module,
+    filename: str | os.PathLike,
+    *,
+    device: str = "cpu",
+    checkpoint_keys_mapping: str = "",
+) -> torch.nn.Module:
+    state_dict = safetensors.torch.load_file(filename, device=device)
+
+    # Optional user-supplied renames (e.g. "model._orig_mod.//model.")
+    if checkpoint_keys_mapping and "//" in checkpoint_keys_mapping:
+        state_dict = rename_checkpoint_keys(state_dict, checkpoint_keys_mapping)
+
+    state_dict, _ = standardise_state_dict(state_dict, set(model.state_dict().keys()))
+
+    # HACK(aliberts): to not overwrite normalization parameters as they should come from the dataset
+    norm_keys = ("normalize_inputs", "normalize_targets", "unnormalize_outputs")
+    state_dict = {k: v for k, v in state_dict.items() if not k.startswith(norm_keys)}
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+
+    if not all(key.startswith(norm_keys) for key in missing) or unexpected:
+        raise RuntimeError(
+            "SmolVLA %d missing / %d unexpected keys",
+            len(missing),
+            len(unexpected),
+        )
+
+    return model
+
 
 def create_sinusoidal_pos_embedding(
     time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu"
@@ -222,17 +326,28 @@ class SmolVLAPolicy(PreTrainedPolicy):
     def __init__(
         self,
         config: SmolVLAConfig,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
             config: Policy configuration class instance or None, in which case the default instantiation of
                     the configuration class is used.
+            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
+                that they will be passed with a call to `load_state_dict` before the policy is used.
         """
 
         super().__init__(config)
         config.validate_features()
         self.config = config
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
 
+        self.language_tokenizer = AutoProcessor.from_pretrained(self.config.vlm_model_name).tokenizer
         self.model = VLAFlowMatching(config)
         self.reset()
 
@@ -242,6 +357,23 @@ class SmolVLAPolicy(PreTrainedPolicy):
             ACTION: deque(maxlen=self.config.n_action_steps),
         }
 
+    # HACK(aliberts, danaaubakirova): we overwrite this classmethod here to fix smolVLA-specific issues
+    @classmethod
+    def _load_as_safetensor(
+        cls,
+        model: "SmolVLAPolicy",
+        model_file: str,
+        map_location: str,
+        strict: bool,
+    ):
+        safetensors.torch.load_model(model, model_file, strict=strict, device=map_location)
+        return load_smolvla(
+            model,
+            model_file,
+            device=map_location,
+            checkpoint_keys_mapping="model._orig_mod.//model.",
+        )
+
     def get_optim_params(self) -> dict:
         return self.parameters()
 
@@ -257,8 +389,7 @@ class SmolVLAPolicy(PreTrainedPolicy):
 
         images, img_masks = self.prepare_images(batch)
         state = self.prepare_state(batch)
-        lang_tokens = batch[f"{OBS_LANGUAGE}.tokens"]
-        lang_masks = batch[f"{OBS_LANGUAGE}.attention_mask"]
+        lang_tokens, lang_masks = self.prepare_language(batch)
 
         actions = self.model.sample_actions(images, img_masks, lang_tokens, lang_masks, state, noise=noise)
 
@@ -266,6 +397,8 @@ class SmolVLAPolicy(PreTrainedPolicy):
         original_action_dim = self.config.action_feature.shape[0]
         actions = actions[:, :, :original_action_dim]
 
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
+
         if self.config.adapt_to_pi_aloha:
             actions = self._pi_aloha_encode_actions(actions)
 
@@ -275,6 +408,8 @@ class SmolVLAPolicy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
 
+        batch = self.normalize_inputs(batch)
+
         return batch
 
     @torch.no_grad()
@@ -315,11 +450,11 @@ class SmolVLAPolicy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
             batch[ACTION] = self._pi_aloha_encode_actions_inv(batch[ACTION])
-
+        batch = self.normalize_inputs(batch)
+        batch = self.normalize_targets(batch)
         images, img_masks = self.prepare_images(batch)
         state = self.prepare_state(batch)
-        lang_tokens = batch[f"{OBS_LANGUAGE}.tokens"]
-        lang_masks = batch[f"{OBS_LANGUAGE}.attention_mask"]
+        lang_tokens, lang_masks = self.prepare_language(batch)
         actions = self.prepare_action(batch)
         actions_is_pad = batch.get("actions_id_pad")
         loss_dict = {}
@@ -383,6 +518,30 @@ class SmolVLAPolicy(PreTrainedPolicy):
             img_masks.append(mask)
         return images, img_masks
 
+    def prepare_language(self, batch) -> tuple[Tensor, Tensor]:
+        """Tokenize the text input"""
+        device = batch[OBS_STATE].device
+        tasks = batch["task"]
+        if isinstance(tasks, str):
+            tasks = [tasks]
+
+        if len(tasks) == 1:
+            tasks = [tasks[0] for _ in range(batch[OBS_STATE].shape[0])]
+
+        tasks = [task if task.endswith("\n") else f"{task}\n" for task in tasks]
+
+        tokenized_prompt = self.language_tokenizer.__call__(
+            tasks,
+            padding=self.config.pad_language_to,
+            padding_side="right",
+            max_length=self.config.tokenizer_max_length,
+            return_tensors="pt",
+        )
+        lang_tokens = tokenized_prompt["input_ids"].to(device=device)
+        lang_masks = tokenized_prompt["attention_mask"].to(device=device, dtype=torch.bool)
+
+        return lang_tokens, lang_masks
+
     def _pi_aloha_decode_state(self, state):
         # Flip the joints.
         for motor_idx in [1, 2, 8, 9]:

src/lerobot/policies/smolvla/processor_smolvla.py

@@ -1,109 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 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.
-from typing import Any
-
-import torch
-
-from lerobot.configs.types import PolicyFeature
-from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
-from lerobot.processor import (
-    DeviceProcessor,
-    NormalizerProcessor,
-    RenameProcessor,
-    RobotProcessor,
-    ToBatchProcessor,
-    TokenizerProcessor,
-    UnnormalizerProcessor,
-)
-from lerobot.processor.pipeline import EnvTransition, ProcessorStep, ProcessorStepRegistry, TransitionKey
-
-
-def make_smolvla_processor(
-    config: SmolVLAConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
-) -> tuple[RobotProcessor, RobotProcessor]:
-    input_steps = [
-        RenameProcessor(rename_map={}),  # To mimic the same processor as pretrained one
-        NormalizerProcessor(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        ToBatchProcessor(),
-        SmolVLANewLineProcessor(),
-        TokenizerProcessor(
-            tokenizer_name=config.vlm_model_name,
-            padding=config.pad_language_to,
-            padding_side="right",
-            max_length=config.tokenizer_max_length,
-        ),
-        DeviceProcessor(device=config.device),
-    ]
-    output_steps = [
-        DeviceProcessor(device="cpu"),
-        UnnormalizerProcessor(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
-        steps=output_steps, name="robot_postprocessor"
-    )
-
-
-@ProcessorStepRegistry.register(name="smolvla_new_line_processor")
-class SmolVLANewLineProcessor(ProcessorStep):
-    """Add a new line to the end of the task if it doesn't have one."""
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        # Check if complementary_data exists
-        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
-        if complementary_data is None or "task" not in complementary_data:
-            return transition
-
-        task = complementary_data["task"]
-        if task is None:
-            return transition
-
-        # Handle both string and list of strings
-        if isinstance(task, str):
-            # Single string: add newline if not present
-            if not task.endswith("\n"):
-                complementary_data["task"] = f"{task}\n"
-        elif isinstance(task, list) and all(isinstance(t, str) for t in task):
-            # List of strings: add newline to each if not present
-            complementary_data["task"] = [t if t.endswith("\n") else f"{t}\n" for t in task]
-        # If task is neither string nor list of strings, leave unchanged
-
-        return transition
-
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        """Adds nothing to the features."""
-        return features
-
-    def state_dict(self) -> dict[str, torch.Tensor]:
-        """Return state dictionary (empty for this processor)."""
-        return {}
-
-    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
-        """Load state dictionary (no-op for this processor)."""
-        pass
-
-    def reset(self) -> None:
-        """Reset processor state (no-op for this processor)."""
-        pass
-
-    def get_config(self) -> dict[str, Any]:
-        """Return configuration for serialization."""
-        return {}

src/lerobot/policies/tdmpc/modeling_tdmpc.py

@@ -36,6 +36,7 @@ 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.normalize import Normalize, Unnormalize
 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
@@ -62,19 +63,26 @@ class TDMPCPolicy(PreTrainedPolicy):
     config_class = TDMPCConfig
     name = "tdmpc"
 
-    def __init__(
-        self,
-        config: TDMPCConfig,
-    ):
+    def __init__(self, config: TDMPCConfig, dataset_stats: dict[str, dict[str, Tensor]] | None = None):
         """
         Args:
             config: Policy configuration class instance or None, in which case the default instantiation of
                 the configuration class is used.
+            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
+                that they will be passed with a call to `load_state_dict` before the policy is used.
         """
         super().__init__(config)
         config.validate_features()
         self.config = config
 
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
         self.model = TDMPCTOLD(config)
         self.model_target = deepcopy(self.model)
         for param in self.model_target.parameters():
@@ -129,6 +137,7 @@ class TDMPCPolicy(PreTrainedPolicy):
 
         actions = torch.clamp(actions, -1, +1)
 
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
         return actions
 
     @torch.no_grad()
@@ -138,12 +147,11 @@ class TDMPCPolicy(PreTrainedPolicy):
         if ACTION in batch:
             batch.pop(ACTION)
 
+        batch = self.normalize_inputs(batch)
+
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGE] = batch[next(iter(self.config.image_features))]
-        # NOTE: for offline evaluation, we have action in the batch, so we need to pop it out
-        if ACTION in batch:
-            batch.pop(ACTION)
 
         self._queues = populate_queues(self._queues, batch)
 
@@ -312,9 +320,11 @@ class TDMPCPolicy(PreTrainedPolicy):
         """
         device = get_device_from_parameters(self)
 
+        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGE] = batch[next(iter(self.config.image_features))]
+        batch = self.normalize_targets(batch)
 
         info = {}
 

src/lerobot/policies/tdmpc/processor_tdmpc.py

@@ -1,51 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 Nicklas Hansen, Xiaolong Wang, Hao Su,
-# and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-
-from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
-from lerobot.processor import (
-    DeviceProcessor,
-    NormalizerProcessor,
-    RenameProcessor,
-    RobotProcessor,
-    ToBatchProcessor,
-    UnnormalizerProcessor,
-)
-
-
-def make_tdmpc_processor(
-    config: TDMPCConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
-) -> tuple[RobotProcessor, RobotProcessor]:
-    input_steps = [
-        RenameProcessor(rename_map={}),
-        NormalizerProcessor(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        ToBatchProcessor(),
-        DeviceProcessor(device=config.device),
-    ]
-    output_steps = [
-        DeviceProcessor(device="cpu"),
-        UnnormalizerProcessor(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
-        steps=output_steps, name="robot_postprocessor"
-    )

src/lerobot/policies/vqbet/modeling_vqbet.py

@@ -28,6 +28,7 @@ import torchvision
 from torch import Tensor, nn
 
 from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE
+from lerobot.policies.normalize import Normalize, Unnormalize
 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
@@ -47,6 +48,7 @@ class VQBeTPolicy(PreTrainedPolicy):
     def __init__(
         self,
         config: VQBeTConfig | None = None,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
@@ -59,6 +61,14 @@ class VQBeTPolicy(PreTrainedPolicy):
         config.validate_features()
         self.config = config
 
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
         self.vqbet = VQBeTModel(config)
 
         self.reset()
@@ -118,6 +128,7 @@ class VQBeTPolicy(PreTrainedPolicy):
     def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
         batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
         actions = self.vqbet(batch, rollout=True)[:, : self.config.action_chunk_size]
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
         return actions
 
     @torch.no_grad()
@@ -131,12 +142,10 @@ class VQBeTPolicy(PreTrainedPolicy):
         # NOTE: for offline evaluation, we have action in the batch, so we need to pop it out
         if ACTION in batch:
             batch.pop(ACTION)
+        batch = self.normalize_inputs(batch)
         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["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)
 
         self._queues = populate_queues(self._queues, batch)
 
@@ -156,8 +165,10 @@ class VQBeTPolicy(PreTrainedPolicy):
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """Run the batch through the model and compute the loss for training or validation."""
+        batch = self.normalize_inputs(batch)
         batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
         batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
+        batch = self.normalize_targets(batch)
         # VQ-BeT discretizes action using VQ-VAE before training BeT (please refer to section 3.2 in the VQ-BeT paper https://huggingface.co/papers/2403.03181)
         if not self.vqbet.action_head.vqvae_model.discretized.item():
             # loss: total loss of training RVQ

src/lerobot/policies/vqbet/processor_vqbet.py

@@ -1,52 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 Seungjae Lee and Yibin Wang and Haritheja Etukuru
-# and H. Jin Kim and Nur Muhammad Mahi Shafiullah and Lerrel Pinto
-# and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-
-from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
-from lerobot.processor import (
-    DeviceProcessor,
-    NormalizerProcessor,
-    RenameProcessor,
-    RobotProcessor,
-    ToBatchProcessor,
-    UnnormalizerProcessor,
-)
-
-
-def make_vqbet_processor(
-    config: VQBeTConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
-) -> tuple[RobotProcessor, RobotProcessor]:
-    input_steps = [
-        RenameProcessor(rename_map={}),  # Let the possibility to the user to rename the keys
-        NormalizerProcessor(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        ToBatchProcessor(),
-        DeviceProcessor(device=config.device),
-    ]
-    output_steps = [
-        DeviceProcessor(device="cpu"),
-        UnnormalizerProcessor(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
-        steps=output_steps, name="robot_postprocessor"
-    )

src/lerobot/processor/__init__.py

@@ -14,9 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .batch_processor import ToBatchProcessor
 from .device_processor import DeviceProcessor
-from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor, hotswap_stats
+from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor
 from .observation_processor import VanillaObservationProcessor
 from .pipeline import (
     ActionProcessor,
@@ -33,7 +32,6 @@ from .pipeline import (
     TruncatedProcessor,
 )
 from .rename_processor import RenameProcessor
-from .tokenizer_processor import TokenizerProcessor
 
 __all__ = [
     "ActionProcessor",
@@ -44,15 +42,12 @@ __all__ = [
     "InfoProcessor",
     "NormalizerProcessor",
     "UnnormalizerProcessor",
-    "hotswap_stats",
     "ObservationProcessor",
     "ProcessorStep",
     "ProcessorStepRegistry",
     "RenameProcessor",
     "RewardProcessor",
     "RobotProcessor",
-    "ToBatchProcessor",
-    "TokenizerProcessor",
     "TransitionKey",
     "TruncatedProcessor",
     "VanillaObservationProcessor",

src/lerobot/processor/batch_processor.py

@@ -1,139 +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.
-from dataclasses import dataclass
-from typing import Any
-
-import torch
-from torch import Tensor
-
-from lerobot.configs.types import PolicyFeature
-from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
-from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
-
-
-@dataclass
-@ProcessorStepRegistry.register(name="to_batch_processor")
-class ToBatchProcessor:
-    """Processor that adds batch dimensions to observations and actions when needed.
-
-    This processor ensures that observations and actions have proper batch dimensions for model processing:
-
-    - For state observations (observation.state, observation.environment_state):
-      Adds batch dimension (unsqueeze at dim=0) if tensor is 1-dimensional
-
-    - For image observations (observation.image, observation.images.*):
-      Adds batch dimension (unsqueeze at dim=0) if tensor is 3-dimensional (H, W, C)
-
-    - For actions:
-      Adds batch dimension (unsqueeze at dim=0) if tensor is 1-dimensional
-
-    - For task field in complementary data:
-      Wraps string task in a list to add batch dimension
-      (task must be a string or list of strings)
-
-    This is useful when processing single transitions that need to be batched for
-    model inference or when converting from unbatched environment outputs to
-    batched model inputs.
-
-    The processor only modifies tensors that need batching and leaves already
-    batched tensors unchanged.
-
-    Example:
-        ```python
-        # State: (7,) -> (1, 7)
-        # Image: (224, 224, 3) -> (1, 224, 224, 3)
-        # Action: (4,) -> (1, 4)
-        # Task: "pick_cube" -> ["pick_cube"]
-        # Already batched: (1, 7) -> (1, 7) [unchanged]
-        ```
-    """
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        self._process_observation(transition)
-        self._process_action(transition)
-        self._process_complementary_data(transition)
-        return transition
-
-    def _process_observation(self, transition: EnvTransition) -> None:
-        """Process observation component in-place, adding batch dimensions where needed."""
-        observation = transition.get(TransitionKey.OBSERVATION)
-        if observation is None:
-            return
-
-        # Process state observations - add batch dim if 1D
-        for state_key in [OBS_STATE, OBS_ENV_STATE]:
-            if state_key in observation:
-                state_value = observation[state_key]
-                if isinstance(state_value, Tensor) and state_value.dim() == 1:
-                    observation[state_key] = state_value.unsqueeze(0)
-
-        # Process single image observation - add batch dim if 3D
-        if OBS_IMAGE in observation:
-            image_value = observation[OBS_IMAGE]
-            if isinstance(image_value, Tensor) and image_value.dim() == 3:
-                observation[OBS_IMAGE] = image_value.unsqueeze(0)
-
-        # Process multiple image observations - add batch dim if 3D
-        for key, value in observation.items():
-            if key.startswith(f"{OBS_IMAGES}.") and isinstance(value, Tensor) and value.dim() == 3:
-                observation[key] = value.unsqueeze(0)
-
-    def _process_action(self, transition: EnvTransition) -> None:
-        """Process action component in-place, adding batch dimension if needed."""
-        action = transition.get(TransitionKey.ACTION)
-        if action is not None and isinstance(action, Tensor) and action.dim() == 1:
-            transition[TransitionKey.ACTION] = action.unsqueeze(0)
-
-    def _process_complementary_data(self, transition: EnvTransition) -> None:
-        """Process complementary data in-place, handling task field batching."""
-        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
-        if complementary_data is None:
-            return
-
-        # Process task field - wrap string in list to add batch dimension
-        if "task" in complementary_data:
-            task_value = complementary_data["task"]
-            if isinstance(task_value, str):
-                complementary_data["task"] = [task_value]
-
-        # Process index field - add batch dim if 0D
-        if "index" in complementary_data:
-            index_value = complementary_data["index"]
-            if isinstance(index_value, Tensor) and index_value.dim() == 0:
-                complementary_data["index"] = index_value.unsqueeze(0)
-
-        # Process task_index field - add batch dim if 0D
-        if "task_index" in complementary_data:
-            task_index_value = complementary_data["task_index"]
-            if isinstance(task_index_value, Tensor) and task_index_value.dim() == 0:
-                complementary_data["task_index"] = task_index_value.unsqueeze(0)
-
-    def get_config(self) -> dict[str, Any]:
-        """Return configuration for serialization."""
-        return {}
-
-    def state_dict(self) -> dict[str, torch.Tensor]:
-        """Return state dictionary (empty for this processor)."""
-        return {}
-
-    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
-        """Load state dictionary (no-op for this processor)."""
-        pass
-
-    def reset(self) -> None:
-        """Reset processor state (no-op for this processor)."""
-        pass
-
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        return features

src/lerobot/processor/converters.py

@@ -1,225 +0,0 @@
-# !/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-from collections.abc import Iterable, Sequence
-from copy import deepcopy
-from typing import Any
-
-import numpy as np
-import torch
-from scipy.spatial.transform import Rotation
-
-from .pipeline import EnvTransition, TransitionKey
-
-
-def _to_tensor(x: torch.Tensor | np.ndarray | Sequence[int | float]):
-    if isinstance(x, torch.Tensor):
-        return x
-    if isinstance(x, np.ndarray):
-        # Keep images (uint8 HWC) and python objects as-is
-        if x.dtype == np.uint8 or x.dtype == np.object_:
-            return x
-        # Scalars/arrays to float32 tensor
-        return torch.as_tensor(x, dtype=torch.float32)
-    # Anything else to float32 tensor
-    return torch.as_tensor(x, dtype=torch.float32)
-
-
-def _from_tensor(x: Any):
-    if isinstance(x, torch.Tensor):
-        return x.item() if x.numel() == 1 else x.detach().cpu().numpy()
-    return x
-
-
-def _is_image(arr: Any) -> bool:
-    return isinstance(arr, np.ndarray) and arr.dtype == np.uint8 and arr.ndim == 3
-
-
-def _split_obs_to_state_and_images(obs: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
-    state, images = {}, {}
-    for k, v in obs.items():
-        if _is_image(v):
-            images[k] = v
-        else:
-            state[k] = v
-    return state, images
-
-
-def make_obs_act_transition(
-    *, obs: dict[str, Any] | None = None, act: dict[str, Any] | None = None
-) -> EnvTransition:
-    return {
-        TransitionKey.OBSERVATION: {} if obs is None else obs,
-        TransitionKey.ACTION: {} if act is None else act,
-        TransitionKey.INFO: {},
-        TransitionKey.COMPLEMENTARY_DATA: {},
-        TransitionKey.REWARD: None,
-        TransitionKey.DONE: None,
-        TransitionKey.TRUNCATED: None,
-    }
-
-
-def to_transition_teleop_action(action: dict[str, Any]) -> EnvTransition:
-    """
-    Convert a raw teleop action dict into an EnvTransition under the ACTION TransitionKey.
-    """
-    act_dict: dict[str, Any] = {}
-    for k, v in action.items():
-        # Check if the value is a type that should not be converted to a tensor.
-        if isinstance(v, (Rotation, dict)):
-            act_dict[f"action.{k}"] = v
-            continue
-
-        arr = np.array(v) if np.isscalar(v) else v
-        act_dict[f"action.{k}"] = _to_tensor(arr)
-
-    return make_obs_act_transition(act=act_dict)
-
-
-# TODO(Adil, Pepijn): Overtime we can maybe add these converters to pipeline.py itself
-def to_transition_robot_observation(observation: dict[str, Any]) -> EnvTransition:
-    """
-    Convert a raw robot observation dict into an EnvTransition under the OBSERVATION TransitionKey.
-    """
-    state, images = _split_obs_to_state_and_images(observation)
-
-    obs_dict: dict[str, Any] = {}
-    for k, v in state.items():
-        arr = np.array(v) if np.isscalar(v) else v
-        obs_dict[f"observation.state.{k}"] = _to_tensor(arr)
-
-    for cam, img in images.items():
-        obs_dict[f"observation.images.{cam}"] = img
-
-    return make_obs_act_transition(obs=obs_dict)
-
-
-def to_output_robot_action(transition: EnvTransition) -> dict[str, Any]:
-    """
-    Converts a EnvTransition under the ACTION TransitionKey to a dict with keys ending in '.pos' for raw robot actions.
-    """
-    out: dict[str, Any] = {}
-    action_dict = transition.get(TransitionKey.ACTION) or {}
-
-    for k, v in action_dict.items():
-        if isinstance(k, str) and k.startswith("action.") and k.endswith((".pos", ".vel")):
-            out_key = k[len("action.") :]  # Strip the 'action.' prefix.
-            out[out_key] = float(v)
-
-    return out
-
-
-def to_dataset_frame(
-    transitions_or_transition: EnvTransition | Iterable[EnvTransition], features: dict[str, dict]
-) -> dict[str, any]:
-    """
-    Converts a single EnvTransition or an iterable of them into a flat,
-    dataset-friendly dictionary for training or evaluation, according to
-    the provided `features` spec.
-
-    Args:
-        transitions_or_transition: Either a single EnvTransition dict
-            or an iterable of them (which will be merged).
-        features (dict[str, dict]):
-            A feature specification dictionary:
-              - 'action': dict with 'names': list of action feature names
-              - 'observation.state': dict with 'names': list of state feature names
-              - keys starting with 'observation.images.' are passed through
-
-    Returns:
-        batch (dict[str, any]): Flat dictionary containing:
-          - numpy arrays for "observation.state" and "action"
-          - any image tensors defined in features
-          - next.{reward,done,truncated}
-          - info dict
-          - *_is_pad flags and task from complementary_data
-    """
-    action_names = features.get("action", {}).get("names", [])
-    obs_state_names = features.get("observation.state", {}).get("names", [])
-    image_keys = [k for k in features if k.startswith("observation.images.")]
-
-    def _merge(base: EnvTransition, other: EnvTransition) -> EnvTransition:
-        out = deepcopy(base)
-        for key in (
-            TransitionKey.OBSERVATION,
-            TransitionKey.ACTION,
-            TransitionKey.INFO,
-            TransitionKey.COMPLEMENTARY_DATA,
-        ):
-            if other.get(key):
-                out.setdefault(key, {}).update(deepcopy(other[key]))
-        for k in (TransitionKey.REWARD, TransitionKey.DONE, TransitionKey.TRUNCATED):
-            if k in other:
-                out[k] = other[k]
-        return out
-
-    def _ensure_transition(obj) -> EnvTransition:
-        # single transition
-        if isinstance(obj, dict) and any(isinstance(k, TransitionKey) for k in obj):
-            return obj
-        # iterable of transitions
-        if isinstance(obj, Iterable):
-            items = list(obj)
-            if not items:
-                return {}
-            acc = items[0]
-            for t in items[1:]:
-                acc = _merge(acc, t)
-            return acc
-        raise TypeError("Expected EnvTransition or iterable of them")
-
-    tr = _ensure_transition(transitions_or_transition)
-    obs = tr.get(TransitionKey.OBSERVATION, {}) or {}
-    act = tr.get(TransitionKey.ACTION, {}) or {}
-    batch: dict[str, any] = {}
-
-    # Images passthrough
-    for k in image_keys:
-        if k in obs:
-            batch[k] = obs[k]
-
-    # Observation.state vector
-    if obs_state_names:
-        vals = [_from_tensor(obs.get(f"observation.state.{n}", 0.0)) for n in obs_state_names]
-        batch["observation.state"] = np.asarray(vals, dtype=np.float32)
-
-    # Action vector
-    if action_names:
-        vals = [_from_tensor(act.get(f"action.{n}", 0.0)) for n in action_names]
-        batch["action"] = np.asarray(vals, dtype=np.float32)
-
-    # Next.* fields
-    if tr.get(TransitionKey.REWARD) is not None:
-        batch["next.reward"] = _from_tensor(tr[TransitionKey.REWARD])
-    if tr.get(TransitionKey.DONE) is not None:
-        batch["next.done"] = _from_tensor(tr[TransitionKey.DONE])
-    if tr.get(TransitionKey.TRUNCATED) is not None:
-        batch["next.truncated"] = _from_tensor(tr[TransitionKey.TRUNCATED])
-
-    # Complementary data flags and task
-    comp = tr.get(TransitionKey.COMPLEMENTARY_DATA) or {}
-    if comp:
-        # pad flags
-        for k, v in comp.items():
-            if k.endswith("_is_pad"):
-                batch[k] = v
-        # task label
-        if comp.get("task") is not None:
-            batch["task"] = comp["task"]
-
-    return batch

src/lerobot/processor/device_processor.py

@@ -19,63 +19,24 @@ from typing import Any
 import torch
 
 from lerobot.configs.types import PolicyFeature
-from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
+from lerobot.processor.pipeline import EnvTransition, TransitionKey
 from lerobot.utils.utils import get_safe_torch_device
 
 
-@ProcessorStepRegistry.register("device_processor")
 @dataclass
 class DeviceProcessor:
-    """Processes transitions by moving tensors to the specified device and optionally converting float dtypes.
+    """Processes transitions by moving tensors to the specified device.
 
     This processor ensures that all tensors in the transition are moved to the
-    specified device (CPU or GPU) before they are returned. It can also convert
-    floating-point tensors to a specified dtype while preserving non-float types
-    (int, long, bool, etc.).
+    specified device (CPU or GPU) before they are returned.
     """
 
-    device: str = "cpu"
-    float_dtype: str | None = None
-    _device: torch.device | None = None
+    device: torch.device = "cpu"
 
     def __post_init__(self):
-        self._device = get_safe_torch_device(self.device)
-        self.device = self._device.type
+        self.device = get_safe_torch_device(self.device)
         self.non_blocking = "cuda" in str(self.device)
 
-        # Validate and convert float_dtype string to torch dtype
-        if self.float_dtype is not None:
-            dtype_mapping = {
-                "float16": torch.float16,
-                "float32": torch.float32,
-                "float64": torch.float64,
-                "bfloat16": torch.bfloat16,
-                "half": torch.float16,
-                "float": torch.float32,
-                "double": torch.float64,
-            }
-
-            if self.float_dtype not in dtype_mapping:
-                available_dtypes = list(dtype_mapping.keys())
-                raise ValueError(
-                    f"Invalid float_dtype '{self.float_dtype}'. Available options: {available_dtypes}"
-                )
-
-            self._target_float_dtype = dtype_mapping[self.float_dtype]
-        else:
-            self._target_float_dtype = None
-
-    def _process_tensor(self, tensor: torch.Tensor) -> torch.Tensor:
-        """Process a tensor by moving to device and optionally converting float dtype."""
-        # Move to device first
-        tensor = tensor.to(self.device, non_blocking=self.non_blocking)
-
-        # Convert float dtype if specified and tensor is floating point
-        if self._target_float_dtype is not None and tensor.is_floating_point():
-            tensor = tensor.to(dtype=self._target_float_dtype)
-
-        return tensor
-
     def __call__(self, transition: EnvTransition) -> EnvTransition:
         # Create a copy of the transition
         new_transition = transition.copy()
@@ -84,7 +45,7 @@ class DeviceProcessor:
         observation = transition.get(TransitionKey.OBSERVATION)
         if observation is not None:
             new_observation = {
-                k: self._process_tensor(v) if isinstance(v, torch.Tensor) else v
+                k: v.to(self.device, non_blocking=self.non_blocking) if isinstance(v, torch.Tensor) else v
                 for k, v in observation.items()
             }
             new_transition[TransitionKey.OBSERVATION] = new_observation
@@ -92,54 +53,30 @@ class DeviceProcessor:
         # Process action tensor
         action = transition.get(TransitionKey.ACTION)
         if action is not None and isinstance(action, torch.Tensor):
-            new_transition[TransitionKey.ACTION] = self._process_tensor(action)
+            new_transition[TransitionKey.ACTION] = action.to(self.device, non_blocking=self.non_blocking)
 
         # Process reward tensor
         reward = transition.get(TransitionKey.REWARD)
         if reward is not None and isinstance(reward, torch.Tensor):
-            new_transition[TransitionKey.REWARD] = self._process_tensor(reward)
+            new_transition[TransitionKey.REWARD] = reward.to(self.device, non_blocking=self.non_blocking)
 
         # Process done tensor
         done = transition.get(TransitionKey.DONE)
         if done is not None and isinstance(done, torch.Tensor):
-            new_transition[TransitionKey.DONE] = self._process_tensor(done)
+            new_transition[TransitionKey.DONE] = done.to(self.device, non_blocking=self.non_blocking)
 
         # Process truncated tensor
         truncated = transition.get(TransitionKey.TRUNCATED)
         if truncated is not None and isinstance(truncated, torch.Tensor):
-            new_transition[TransitionKey.TRUNCATED] = self._process_tensor(truncated)
-
-        # Process complementary data tensors
-        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
-        if complementary_data is not None:
-            new_complementary_data = {}
-
-            # Process all items in complementary_data
-            for key, value in complementary_data.items():
-                if isinstance(value, torch.Tensor):
-                    new_complementary_data[key] = self._process_tensor(value)
-                else:
-                    new_complementary_data[key] = value
-
-            new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
+            new_transition[TransitionKey.TRUNCATED] = truncated.to(
+                self.device, non_blocking=self.non_blocking
+            )
 
         return new_transition
 
     def get_config(self) -> dict[str, Any]:
         """Return configuration for serialization."""
-        return {"device": self.device, "float_dtype": self.float_dtype}
+        return {"device": self.device}
 
-    def state_dict(self) -> dict[str, torch.Tensor]:
-        """Return state dictionary (empty for this processor)."""
-        return {}
-
-    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
-        """Load state dictionary (no-op for this processor)."""
-        pass
-
-    def reset(self) -> None:
-        """Reset processor state (no-op for this processor)."""
-        pass
-
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features

src/lerobot/processor/migrate_policy_normalization.py

@@ -1,502 +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.
-
-"""
-Generic script to migrate any policy model with normalization layers to the new pipeline-based system.
-
-This script:
-1. Loads an existing pretrained policy model
-2. Extracts normalization statistics from the model
-3. Creates both preprocessor and postprocessor:
-   - Preprocessor: normalizes both inputs (observations) and outputs (actions) for training
-   - Postprocessor: unnormalizes outputs (actions) for inference
-4. Removes normalization layers from the model state_dict
-5. Saves the new model and both processors
-
-Usage:
-    python src/lerobot/processor/migrate_policy_normalization.py \
-        --pretrained-path lerobot/act_aloha_sim_transfer_cube_human \
-        --policy-type act \
-        --push-to-hub
-"""
-
-import argparse
-import importlib
-import json
-import os
-from copy import deepcopy
-from pathlib import Path
-from typing import Any
-
-import torch
-from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file as load_safetensors
-
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
-from lerobot.processor.batch_processor import ToBatchProcessor
-from lerobot.processor.device_processor import DeviceProcessor
-from lerobot.processor.normalize_processor import NormalizerProcessor, UnnormalizerProcessor
-from lerobot.processor.pipeline import RobotProcessor
-from lerobot.processor.rename_processor import RenameProcessor
-
-# Policy type to class mapping
-POLICY_CLASSES = {
-    "act": "lerobot.policies.act.modeling_act.ACTPolicy",
-    "diffusion": "lerobot.policies.diffusion.modeling_diffusion.DiffusionPolicy",
-    "pi0": "lerobot.policies.pi0.modeling_pi0.PI0Policy",
-    "pi0fast": "lerobot.policies.pi0fast.modeling_pi0fast.PI0FASTPolicy",
-    "smolvla": "lerobot.policies.smolvla.modeling_smolvla.SmolVLAPolicy",
-    "tdmpc": "lerobot.policies.tdmpc.modeling_tdmpc.TDMPCPolicy",
-    "vqbet": "lerobot.policies.vqbet.modeling_vqbet.VQBeTPolicy",
-    "sac": "lerobot.policies.sac.modeling_sac.SACPolicy",
-    "classifier": "lerobot.policies.classifier.modeling_classifier.ClassifierPolicy",
-}
-
-
-def extract_normalization_stats(state_dict: dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]:
-    """Extract normalization statistics from model state_dict."""
-    stats = {}
-
-    # Define patterns to match and their prefixes to remove
-    normalization_patterns = [
-        "normalize_inputs.buffer_",
-        "unnormalize_outputs.buffer_",
-        "normalize_targets.buffer_",
-        "normalize.",  # Must come after normalize_* patterns
-        "unnormalize.",  # Must come after unnormalize_* patterns
-        "input_normalizer.",
-        "output_normalizer.",
-    ]
-
-    # Process each key in state_dict
-    for key, tensor in state_dict.items():
-        # Try each pattern
-        for pattern in normalization_patterns:
-            if key.startswith(pattern):
-                # Extract the remaining part after the pattern
-                remaining = key[len(pattern) :]
-                parts = remaining.split(".")
-
-                # Need at least feature name and stat type
-                if len(parts) >= 2:
-                    # Last part is the stat type (mean, std, min, max, etc.)
-                    stat_type = parts[-1]
-                    # Everything else is the feature name
-                    feature_name = ".".join(parts[:-1]).replace("_", ".")
-
-                    # Add to stats
-                    if feature_name not in stats:
-                        stats[feature_name] = {}
-                    stats[feature_name][stat_type] = tensor.clone()
-
-                # Only process the first matching pattern
-                break
-
-    return stats
-
-
-def detect_features_and_norm_modes(
-    config: dict[str, Any], stats: dict[str, dict[str, torch.Tensor]]
-) -> tuple[dict[str, PolicyFeature], dict[FeatureType, NormalizationMode]]:
-    """Detect features and normalization modes from config and stats."""
-    features = {}
-    norm_modes = {}
-
-    # First, check if there's a normalization_mapping in the config
-    if "normalization_mapping" in config:
-        print(f"Found normalization_mapping in config: {config['normalization_mapping']}")
-        # Extract normalization modes from config
-        for feature_name, mode_str in config["normalization_mapping"].items():
-            # Convert string to NormalizationMode enum
-            if mode_str == "mean_std":
-                mode = NormalizationMode.MEAN_STD
-            elif mode_str == "min_max":
-                mode = NormalizationMode.MIN_MAX
-            else:
-                print(f"Warning: Unknown normalization mode '{mode_str}' for feature '{feature_name}'")
-                continue
-
-            # Determine feature type from feature name
-            if "image" in feature_name or "visual" in feature_name:
-                feature_type = FeatureType.VISUAL
-            elif "state" in feature_name:
-                feature_type = FeatureType.STATE
-            elif "action" in feature_name:
-                feature_type = FeatureType.ACTION
-            else:
-                feature_type = FeatureType.STATE
-
-            norm_modes[feature_type] = mode
-
-    # Try to extract from config
-    if "features" in config:
-        for key, feature_config in config["features"].items():
-            shape = feature_config.get("shape", feature_config.get("dim"))
-            shape = (shape,) if isinstance(shape, int) else tuple(shape)
-
-            # Determine feature type
-            if "image" in key or "visual" in key:
-                feature_type = FeatureType.VISUAL
-            elif "state" in key:
-                feature_type = FeatureType.STATE
-            elif "action" in key:
-                feature_type = FeatureType.ACTION
-            else:
-                feature_type = FeatureType.STATE  # Default
-
-            features[key] = PolicyFeature(feature_type, shape)
-
-    # If no features in config, infer from stats
-    if not features:
-        for key, stat_dict in stats.items():
-            # Get shape from any stat tensor
-            tensor = next(iter(stat_dict.values()))
-            shape = tuple(tensor.shape)
-
-            # Determine feature type based on key
-            if "image" in key or "visual" in key or "pixels" in key:
-                feature_type = FeatureType.VISUAL
-            elif "state" in key or "joint" in key or "position" in key:
-                feature_type = FeatureType.STATE
-            elif "action" in key:
-                feature_type = FeatureType.ACTION
-            else:
-                feature_type = FeatureType.STATE
-
-            features[key] = PolicyFeature(feature_type, shape)
-
-    # If normalization modes weren't in config, determine based on available stats
-    if not norm_modes:
-        for key, stat_dict in stats.items():
-            if key in features:
-                if "mean" in stat_dict and "std" in stat_dict:
-                    feature_type = features[key].type
-                    if feature_type not in norm_modes:
-                        norm_modes[feature_type] = NormalizationMode.MEAN_STD
-                elif "min" in stat_dict and "max" in stat_dict:
-                    feature_type = features[key].type
-                    if feature_type not in norm_modes:
-                        norm_modes[feature_type] = NormalizationMode.MIN_MAX
-
-    # Default normalization modes if not detected
-    if FeatureType.VISUAL not in norm_modes:
-        norm_modes[FeatureType.VISUAL] = NormalizationMode.MEAN_STD
-    if FeatureType.STATE not in norm_modes:
-        norm_modes[FeatureType.STATE] = NormalizationMode.MIN_MAX
-    if FeatureType.ACTION not in norm_modes:
-        norm_modes[FeatureType.ACTION] = NormalizationMode.MEAN_STD
-
-    return features, norm_modes
-
-
-def remove_normalization_layers(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
-    """Remove normalization layers from state_dict."""
-    new_state_dict = {}
-
-    # Patterns to remove
-    remove_patterns = [
-        "normalize_inputs.",
-        "unnormalize_outputs.",
-        "normalize_targets.",  # Added pattern for target normalization
-        "normalize.",
-        "unnormalize.",
-        "input_normalizer.",
-        "output_normalizer.",
-        "normalizer.",
-    ]
-
-    for key, tensor in state_dict.items():
-        should_remove = any(pattern in key for pattern in remove_patterns)
-        if not should_remove:
-            new_state_dict[key] = tensor
-
-    return new_state_dict
-
-
-def convert_features_to_policy_features(features_dict: dict[str, dict]) -> dict[str, PolicyFeature]:
-    """Convert features from old format to PolicyFeature objects."""
-    converted_features = {}
-
-    for key, feature_dict in features_dict.items():
-        # Determine feature type based on key
-        if "image" in key or "visual" in key:
-            feature_type = FeatureType.VISUAL
-        elif "state" in key:
-            feature_type = FeatureType.STATE
-        elif "action" in key:
-            feature_type = FeatureType.ACTION
-        else:
-            feature_type = FeatureType.STATE
-
-        # Get shape from feature dict
-        shape = feature_dict.get("shape", feature_dict.get("dim"))
-        shape = (shape,) if isinstance(shape, int) else tuple(shape)
-
-        converted_features[key] = PolicyFeature(feature_type, shape)
-
-    return converted_features
-
-
-def load_model_from_hub(
-    repo_id: str, revision: str = None
-) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, Any]]:
-    """Load model state_dict and config from hub."""
-    # Download files
-    safetensors_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", revision=revision)
-
-    config_path = hf_hub_download(repo_id=repo_id, filename="config.json", revision=revision)
-    train_config_path = hf_hub_download(repo_id=repo_id, filename="train_config.json", revision=revision)
-
-    # Load state_dict
-    state_dict = load_safetensors(safetensors_path)
-
-    # Load config
-    with open(config_path) as f:
-        config = json.load(f)
-
-    with open(train_config_path) as f:
-        train_config = json.load(f)
-
-    return state_dict, config, train_config
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Migrate policy models with normalization layers to new pipeline system"
-    )
-    parser.add_argument(
-        "--pretrained-path",
-        type=str,
-        required=True,
-        help="Path to pretrained model (hub repo or local directory)",
-    )
-    parser.add_argument(
-        "--output-dir",
-        type=str,
-        default=None,
-        help="Output directory for migrated model (default: same as pretrained-path)",
-    )
-    parser.add_argument("--push-to-hub", action="store_true", help="Push migrated model to hub")
-    parser.add_argument(
-        "--hub-repo-id",
-        type=str,
-        default=None,
-        help="Hub repository ID for pushing (default: same as pretrained-path)",
-    )
-    parser.add_argument("--revision", type=str, default=None, help="Revision of the model to load")
-    parser.add_argument("--private", action="store_true", help="Make the hub repository private")
-
-    args = parser.parse_args()
-
-    # Load model and config
-    print(f"Loading model from {args.pretrained_path}...")
-    if os.path.isdir(args.pretrained_path):
-        # Local directory
-        state_dict = load_safetensors(os.path.join(args.pretrained_path, "model.safetensors"))
-        with open(os.path.join(args.pretrained_path, "config.json")) as f:
-            config = json.load(f)
-        with open(os.path.join(args.pretrained_path, "train_config.json")) as f:
-            train_config = json.load(f)
-    else:
-        # Hub repository
-        state_dict, config, train_config = load_model_from_hub(args.pretrained_path, args.revision)
-
-    # Extract normalization statistics
-    print("Extracting normalization statistics...")
-    stats = extract_normalization_stats(state_dict)
-
-    print(f"Found normalization statistics for: {list(stats.keys())}")
-
-    # Detect input features and normalization modes
-    print("Detecting features and normalization modes...")
-    features, norm_map = detect_features_and_norm_modes(config, stats)
-
-    print(f"Detected features: {list(features.keys())}")
-    print(f"Normalization modes: {norm_map}")
-
-    # Remove normalization layers from state_dict
-    print("Removing normalization layers from model...")
-    new_state_dict = remove_normalization_layers(state_dict)
-
-    removed_keys = set(state_dict.keys()) - set(new_state_dict.keys())
-    if removed_keys:
-        print(f"Removed {len(removed_keys)} normalization layer keys")
-
-    # Determine output path
-    if args.output_dir:
-        output_dir = Path(args.output_dir)
-    else:
-        if os.path.isdir(args.pretrained_path):
-            output_dir = Path(args.pretrained_path).parent / f"{Path(args.pretrained_path).name}_migrated"
-        else:
-            output_dir = Path(f"./{args.pretrained_path.replace('/', '_')}_migrated")
-
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    # Clean up config - remove normalization_mapping field
-    cleaned_config = dict(config)
-    if "normalization_mapping" in cleaned_config:
-        print("Removing 'normalization_mapping' field from config")
-        del cleaned_config["normalization_mapping"]
-    policy_type = deepcopy(cleaned_config["type"])
-
-    del cleaned_config["type"]
-
-    # Instantiate the policy model with cleaned config and load the cleaned state dict
-    print(f"Instantiating {policy_type} policy model...")
-    policy_class_path = POLICY_CLASSES[policy_type]
-    module_path, class_name = policy_class_path.rsplit(".", 1)
-
-    module = importlib.import_module(module_path)
-    policy_class = getattr(module, class_name)
-
-    # Create config class instance
-    config_module_path = module_path.replace("modeling", "configuration")
-    config_module = importlib.import_module(config_module_path)
-    # Handle special cases for config class names
-    config_class_names = {
-        "act": "ACTConfig",
-        "diffusion": "DiffusionConfig",
-        "pi0": "PI0Config",
-        "pi0fast": "PI0FASTConfig",
-        "smolvla": "SmolVLAConfig",
-        "tdmpc": "TDMPCConfig",
-        "vqbet": "VQBeTConfig",
-        "sac": "SACConfig",
-        "classifier": "ClassifierConfig",
-    }
-    config_class_name = config_class_names.get(policy_type, f"{policy_type.upper()}Config")
-    config_class = getattr(config_module, config_class_name)
-
-    # Convert input_features and output_features to PolicyFeature objects - these are mandatory
-    if "input_features" not in cleaned_config:
-        raise ValueError("Missing mandatory 'input_features' in config")
-    if "output_features" not in cleaned_config:
-        raise ValueError("Missing mandatory 'output_features' in config")
-
-    cleaned_config["input_features"] = convert_features_to_policy_features(cleaned_config["input_features"])
-    cleaned_config["output_features"] = convert_features_to_policy_features(cleaned_config["output_features"])
-
-    # Create config instance from cleaned config dict
-    policy_config = config_class(**cleaned_config)
-
-    # Create policy instance - some policies expect dataset_stats
-    policy = policy_class(policy_config)
-
-    # Load the cleaned state dict
-    policy.load_state_dict(new_state_dict, strict=True)
-    print("Successfully loaded cleaned state dict into policy model")
-
-    # Now create preprocessor and postprocessor with cleaned_config available
-    print("Creating preprocessor and postprocessor...")
-    # The pattern from existing processor factories:
-    # - Preprocessor has two NormalizerProcessors: one for input_features, one for output_features
-    # - Postprocessor has one UnnormalizerProcessor for output_features only
-
-    # Get features from cleaned_config (now they're PolicyFeature objects)
-    input_features = cleaned_config.get("input_features", {})
-    output_features = cleaned_config.get("output_features", {})
-
-    # Create preprocessor with two normalizers (following the pattern from processor factories)
-    preprocessor_steps = [
-        RenameProcessor(rename_map={}),
-        NormalizerProcessor(
-            features={**input_features, **output_features},
-            norm_map=norm_map,
-            stats=stats,
-        ),
-        ToBatchProcessor(),
-        DeviceProcessor(device=policy_config.device),
-    ]
-    preprocessor = RobotProcessor(steps=preprocessor_steps, name="robot_preprocessor")
-
-    # Create postprocessor with unnormalizer for outputs only
-    postprocessor_steps = [
-        DeviceProcessor(device="cpu"),
-        UnnormalizerProcessor(features=output_features, norm_map=norm_map, stats=stats),
-    ]
-    postprocessor = RobotProcessor(steps=postprocessor_steps, name="robot_postprocessor")
-
-    # Determine hub repo ID if pushing to hub
-    if args.push_to_hub:
-        if args.hub_repo_id:
-            hub_repo_id = args.hub_repo_id
-        else:
-            if not os.path.isdir(args.pretrained_path):
-                # Use same repo with "_migrated" suffix
-                hub_repo_id = f"{args.pretrained_path}_migrated"
-            else:
-                raise ValueError("--hub-repo-id must be specified when pushing local model to hub")
-    else:
-        hub_repo_id = None
-
-    # Save preprocessor and postprocessor to root directory
-    print(f"Saving preprocessor to {output_dir}...")
-    preprocessor.save_pretrained(output_dir)
-    if args.push_to_hub:
-        preprocessor.push_to_hub(repo_id=hub_repo_id, private=args.private)
-
-    print(f"Saving postprocessor to {output_dir}...")
-    postprocessor.save_pretrained(output_dir)
-    if args.push_to_hub:
-        postprocessor.push_to_hub(repo_id=hub_repo_id, private=args.private)
-
-    # Save model using the policy's save_pretrained method
-    print(f"Saving model to {output_dir}...")
-    policy.save_pretrained(
-        output_dir, push_to_hub=args.push_to_hub, repo_id=hub_repo_id, private=args.private
-    )
-
-    # Generate and save model card
-    print("Generating model card...")
-    # Get metadata from original config
-    dataset_repo_id = train_config.get("repo_id", "unknown")
-    license = config.get("license", "apache-2.0")
-
-    tags = config.get("tags", ["robotics", "lerobot", policy_type]) or ["robotics", "lerobot", policy_type]
-    tags = set(tags).union({"robotics", "lerobot", policy_type})
-    tags = list(tags)
-
-    # Generate model card
-    card = policy.generate_model_card(
-        dataset_repo_id=dataset_repo_id, model_type=policy_type, license=license, tags=tags
-    )
-
-    # Save model card locally
-    card.save(str(output_dir / "README.md"))
-    print(f"Model card saved to {output_dir / 'README.md'}")
-    # Push model card to hub if requested
-    if args.push_to_hub:
-        from huggingface_hub import HfApi
-
-        api = HfApi()
-        api.upload_file(
-            path_or_fileobj=str(output_dir / "README.md"),
-            path_in_repo="README.md",
-            repo_id=hub_repo_id,
-            repo_type="model",
-            commit_message="Add model card for migrated model",
-        )
-        print("Model card pushed to hub")
-
-    print("\nMigration complete!")
-    print(f"Migrated model saved to: {output_dir}")
-    if args.push_to_hub:
-        print(f"Successfully pushed to https://huggingface.co/{hub_repo_id}")
-
-
-if __name__ == "__main__":
-    main()

src/lerobot/processor/normalize_processor.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from collections.abc import Mapping
-from copy import deepcopy
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -11,7 +10,7 @@ from torch import Tensor
 
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, RobotProcessor, TransitionKey
+from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
 
 
 def _convert_stats_to_tensors(stats: dict[str, dict[str, Any]]) -> dict[str, dict[str, Tensor]]:
@@ -116,7 +115,7 @@ class NormalizerProcessor:
         if self.normalize_keys is not None and not isinstance(self.normalize_keys, set):
             self.normalize_keys = set(self.normalize_keys)
 
-    def _normalize_obs(self, observation, normalized_info):
+    def _normalize_obs(self, observation):
         if observation is None:
             return None
 
@@ -129,20 +128,7 @@ class NormalizerProcessor:
 
         processed = dict(observation)
         for key in keys_to_norm:
-            if key not in processed or key not in self.features:
-                continue
-
-            # Check the normalization mode for this feature type
-            feature = self.features[key]
-            norm_mode = self.norm_map.get(feature.type, NormalizationMode.IDENTITY)
-
-            # Skip normalization if mode is IDENTITY
-            if norm_mode is NormalizationMode.IDENTITY:
-                normalized_info[key] = "IDENTITY"
-                continue
-
-            # Skip if no stats available for this key
-            if key not in self._tensor_stats:
+            if key not in processed or key not in self._tensor_stats:
                 continue
 
             orig_val = processed[key]
@@ -153,35 +139,16 @@ class NormalizerProcessor:
             )
             stats = {k: v.to(tensor.device) for k, v in self._tensor_stats[key].items()}
 
-            if norm_mode is NormalizationMode.MEAN_STD:
-                if "mean" in stats and "std" in stats:
-                    mean, std = stats["mean"], stats["std"]
-                    processed[key] = (tensor - mean) / (std + self.eps)
-                    normalized_info[key] = "MEAN_STD"
-            elif norm_mode is NormalizationMode.MIN_MAX:
-                if "min" in stats and "max" in stats:
-                    min_val, max_val = stats["min"], stats["max"]
-                    processed[key] = 2 * (tensor - min_val) / (max_val - min_val + self.eps) - 1
-                    normalized_info[key] = "MIN_MAX"
-            else:
-                raise ValueError(f"Unsupported normalization mode: {norm_mode}")
-
+            if "mean" in stats and "std" in stats:
+                mean, std = stats["mean"], stats["std"]
+                processed[key] = (tensor - mean) / (std + self.eps)
+            elif "min" in stats and "max" in stats:
+                min_val, max_val = stats["min"], stats["max"]
+                processed[key] = 2 * (tensor - min_val) / (max_val - min_val + self.eps) - 1
         return processed
 
-    def _normalize_action(self, action, normalized_info):
-        if action is None:
-            return action
-
-        # Check the normalization mode for actions
-        norm_mode = self.norm_map.get(FeatureType.ACTION, NormalizationMode.IDENTITY)
-
-        # Skip normalization if mode is IDENTITY
-        if norm_mode is NormalizationMode.IDENTITY:
-            normalized_info["action"] = "IDENTITY"
-            return action
-
-        # Skip if no stats available for actions
-        if "action" not in self._tensor_stats:
+    def _normalize_action(self, action):
+        if action is None or "action" not in self._tensor_stats:
             return action
 
         tensor = (
@@ -190,42 +157,22 @@ class NormalizerProcessor:
             else torch.as_tensor(action, dtype=torch.float32)
         )
         stats = {k: v.to(tensor.device) for k, v in self._tensor_stats["action"].items()}
-
-        if norm_mode is NormalizationMode.MEAN_STD:
-            if "mean" in stats and "std" in stats:
-                mean, std = stats["mean"], stats["std"]
-                normalized_info["action"] = "MEAN_STD"
-                return (tensor - mean) / (std + self.eps)
-        elif norm_mode is NormalizationMode.MIN_MAX:
-            if "min" in stats and "max" in stats:
-                min_val, max_val = stats["min"], stats["max"]
-                normalized_info["action"] = "MIN_MAX"
-                return 2 * (tensor - min_val) / (max_val - min_val + self.eps) - 1
-        else:
-            raise ValueError(f"Unsupported normalization mode: {norm_mode}")
-
-        # If we reach here, the required stats for the normalization mode are not available
-        raise ValueError(f"Action stats must contain appropriate values for {norm_mode} normalization")
+        if "mean" in stats and "std" in stats:
+            mean, std = stats["mean"], stats["std"]
+            return (tensor - mean) / (std + self.eps)
+        if "min" in stats and "max" in stats:
+            min_val, max_val = stats["min"], stats["max"]
+            return 2 * (tensor - min_val) / (max_val - min_val + self.eps) - 1
+        raise ValueError("Action stats must contain either ('mean','std') or ('min','max')")
 
     def __call__(self, transition: EnvTransition) -> EnvTransition:
-        # Track what was normalized
-        normalized_info = {}
-
-        observation = self._normalize_obs(transition.get(TransitionKey.OBSERVATION), normalized_info)
-        action = self._normalize_action(transition.get(TransitionKey.ACTION), normalized_info)
+        observation = self._normalize_obs(transition.get(TransitionKey.OBSERVATION))
+        action = self._normalize_action(transition.get(TransitionKey.ACTION))
 
         # Create a new transition with normalized values
         new_transition = transition.copy()
         new_transition[TransitionKey.OBSERVATION] = observation
         new_transition[TransitionKey.ACTION] = action
-
-        # Add normalization info to complementary data
-        if normalized_info:
-            comp_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
-            comp_data = {} if comp_data is None else dict(comp_data)
-            comp_data["normalized_keys"] = normalized_info
-            new_transition[TransitionKey.COMPLEMENTARY_DATA] = comp_data
-
         return new_transition
 
     def get_config(self) -> dict[str, Any]:
@@ -257,7 +204,7 @@ class NormalizerProcessor:
     def reset(self):
         pass
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features
 
 
@@ -306,28 +253,14 @@ class UnnormalizerProcessor:
         self.stats = self.stats or {}
         self._tensor_stats = _convert_stats_to_tensors(self.stats)
 
-    def _unnormalize_obs(self, observation, unnormalized_info):
+    def _unnormalize_obs(self, observation):
         if observation is None:
             return None
         keys = [k for k, ft in self.features.items() if ft.type is not FeatureType.ACTION]
         processed = dict(observation)
         for key in keys:
-            if key not in processed or key not in self.features:
+            if key not in processed or key not in self._tensor_stats:
                 continue
-
-            # Check the normalization mode for this feature type
-            feature = self.features[key]
-            norm_mode = self.norm_map.get(feature.type, NormalizationMode.IDENTITY)
-
-            # Skip unnormalization if mode is IDENTITY
-            if norm_mode is NormalizationMode.IDENTITY:
-                unnormalized_info[key] = "IDENTITY"
-                continue
-
-            # Skip if no stats available for this key
-            if key not in self._tensor_stats:
-                continue
-
             orig_val = processed[key]
             tensor = (
                 orig_val.to(dtype=torch.float32)
@@ -335,80 +268,39 @@ class UnnormalizerProcessor:
                 else torch.as_tensor(orig_val, dtype=torch.float32)
             )
             stats = {k: v.to(tensor.device) for k, v in self._tensor_stats[key].items()}
-
-            if norm_mode is NormalizationMode.MEAN_STD:
-                if "mean" in stats and "std" in stats:
-                    mean, std = stats["mean"], stats["std"]
-                    processed[key] = tensor * std + mean
-                    unnormalized_info[key] = "MEAN_STD"
-            elif norm_mode is NormalizationMode.MIN_MAX:
-                if "min" in stats and "max" in stats:
-                    min_val, max_val = stats["min"], stats["max"]
-                    processed[key] = (tensor + 1) / 2 * (max_val - min_val) + min_val
-                    unnormalized_info[key] = "MIN_MAX"
-            else:
-                raise ValueError(f"Unsupported normalization mode: {norm_mode}")
-
+            if "mean" in stats and "std" in stats:
+                mean, std = stats["mean"], stats["std"]
+                processed[key] = tensor * std + mean
+            elif "min" in stats and "max" in stats:
+                min_val, max_val = stats["min"], stats["max"]
+                processed[key] = (tensor + 1) / 2 * (max_val - min_val) + min_val
         return processed
 
-    def _unnormalize_action(self, action, unnormalized_info):
-        if action is None:
+    def _unnormalize_action(self, action):
+        if action is None or "action" not in self._tensor_stats:
             return action
-
-        # Check the normalization mode for actions
-        norm_mode = self.norm_map.get(FeatureType.ACTION, NormalizationMode.IDENTITY)
-
-        # Skip unnormalization if mode is IDENTITY
-        if norm_mode is NormalizationMode.IDENTITY:
-            unnormalized_info["action"] = "IDENTITY"
-            return action
-
-        # Skip if no stats available for actions
-        if "action" not in self._tensor_stats:
-            return action
-
         tensor = (
             action.to(dtype=torch.float32)
             if isinstance(action, torch.Tensor)
             else torch.as_tensor(action, dtype=torch.float32)
         )
         stats = {k: v.to(tensor.device) for k, v in self._tensor_stats["action"].items()}
-
-        if norm_mode is NormalizationMode.MEAN_STD:
-            if "mean" in stats and "std" in stats:
-                mean, std = stats["mean"], stats["std"]
-                unnormalized_info["action"] = "MEAN_STD"
-                return tensor * std + mean
-        elif norm_mode is NormalizationMode.MIN_MAX:
-            if "min" in stats and "max" in stats:
-                min_val, max_val = stats["min"], stats["max"]
-                unnormalized_info["action"] = "MIN_MAX"
-                return (tensor + 1) / 2 * (max_val - min_val) + min_val
-        else:
-            raise ValueError(f"Unsupported normalization mode: {norm_mode}")
-
-        # If we reach here, the required stats for the normalization mode are not available
-        raise ValueError(f"Action stats must contain appropriate values for {norm_mode} normalization")
+        if "mean" in stats and "std" in stats:
+            mean, std = stats["mean"], stats["std"]
+            return tensor * std + mean
+        if "min" in stats and "max" in stats:
+            min_val, max_val = stats["min"], stats["max"]
+            return (tensor + 1) / 2 * (max_val - min_val) + min_val
+        raise ValueError("Action stats must contain either ('mean','std') or ('min','max')")
 
     def __call__(self, transition: EnvTransition) -> EnvTransition:
-        # Track what was unnormalized
-        unnormalized_info = {}
-
-        observation = self._unnormalize_obs(transition.get(TransitionKey.OBSERVATION), unnormalized_info)
-        action = self._unnormalize_action(transition.get(TransitionKey.ACTION), unnormalized_info)
+        observation = self._unnormalize_obs(transition.get(TransitionKey.OBSERVATION))
+        action = self._unnormalize_action(transition.get(TransitionKey.ACTION))
 
         # Create a new transition with unnormalized values
         new_transition = transition.copy()
         new_transition[TransitionKey.OBSERVATION] = observation
         new_transition[TransitionKey.ACTION] = action
-
-        # Add unnormalization info to complementary data
-        if unnormalized_info:
-            comp_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
-            comp_data = {} if comp_data is None else dict(comp_data)
-            comp_data["unnormalized_keys"] = unnormalized_info
-            new_transition[TransitionKey.COMPLEMENTARY_DATA] = comp_data
-
         return new_transition
 
     def get_config(self) -> dict[str, Any]:
@@ -435,41 +327,5 @@ class UnnormalizerProcessor:
     def reset(self):
         pass
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features
-
-
-def hotswap_stats(robot_processor: RobotProcessor, stats: dict[str, dict[str, Any]]) -> RobotProcessor:
-    robot_processor = deepcopy(robot_processor)
-    for step in robot_processor.steps:
-        if isinstance(step, NormalizerProcessor) or isinstance(step, UnnormalizerProcessor):
-            step: NormalizerProcessor | UnnormalizerProcessor
-            step.stats = stats
-            step._tensor_stats = _convert_stats_to_tensors(stats)
-    return robot_processor
-
-
-def rename_stats(stats: dict[str, dict[str, Any]], rename_map: dict[str, str]) -> dict[str, dict[str, Any]]:
-    """Rename keys in the stats dictionary according to the provided mapping.
-
-    Args:
-        stats: The statistics dictionary with structure {feature_key: {stat_name: value}}
-        rename_map: Dictionary mapping old key names to new key names
-
-    Returns:
-        A new stats dictionary with renamed keys
-
-    Example:
-        >>> stats = {"observation.state": {"mean": 0.0, "std": 1.0}, "action": {"mean": 0.5, "std": 0.5}}
-        >>> rename_map = {"observation.state": "observation.robot_state"}
-        >>> new_stats = rename_stats(stats, rename_map)
-        >>> # new_stats will have "observation.robot_state" instead of "observation.state"
-    """
-    renamed_stats = {}
-
-    for old_key, sub_stats in stats.items():
-        # Use the new key if it exists in the rename map, otherwise keep the old key
-        new_key = rename_map.get(old_key, old_key)
-        renamed_stats[new_key] = deepcopy(sub_stats)
-
-    return renamed_stats

src/lerobot/processor/observation_processor.py

@@ -106,8 +106,9 @@ class VanillaObservationProcessor(ObservationProcessor):
     def observation(self, observation):
         return self._process_observation(observation)
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         """Transforms feature keys to a standardized contract.
+
         This method handles several renaming patterns:
         - Exact matches (e.g., 'pixels' -> 'OBS_IMAGE').
         - Prefixed exact matches (e.g., 'observation.pixels' -> 'OBS_IMAGE').

src/lerobot/processor/pipeline.py

@@ -23,7 +23,7 @@ from copy import deepcopy
 from dataclasses import dataclass, field
 from enum import Enum
 from pathlib import Path
-from typing import Any, Protocol, TypedDict, runtime_checkable
+from typing import Any, Protocol, TypedDict
 
 import torch
 from huggingface_hub import ModelHubMixin, hf_hub_download
@@ -132,7 +132,6 @@ class ProcessorStepRegistry:
         cls._registry.clear()
 
 
-@runtime_checkable
 class ProcessorStep(Protocol):
     """Structural typing interface for a single processor step.
 
@@ -146,6 +145,7 @@ class ProcessorStep(Protocol):
 
     **Required**:
         - ``__call__(transition: EnvTransition) -> EnvTransition``
+        - ``feature_contract(features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]``
 
     Optional helper protocol:
     * ``get_config() -> dict[str, Any]`` – User-defined JSON-serializable
@@ -158,8 +158,6 @@ class ProcessorStep(Protocol):
     * ``load_state_dict(state)`` – Inverse of ``state_dict``. Receives a dict
       containing torch tensors only.
     * ``reset()`` – Clear internal buffers at episode boundaries.
-    * ``transform_features(features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]``
-    If present, this method will be called to aggregate the dataset features of all steps.
 
     Example separation:
     - get_config(): {"name": "my_step", "learning_rate": 0.01, "window_size": 10}
@@ -176,7 +174,7 @@ class ProcessorStep(Protocol):
 
     def reset(self) -> None: ...
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]: ...
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]: ...
 
 
 def _default_batch_to_transition(batch: dict[str, Any]) -> EnvTransition:  # noqa: D401
@@ -203,16 +201,10 @@ def _default_batch_to_transition(batch: dict[str, Any]) -> EnvTransition:  # noq
     observation_keys = {k: v for k, v in batch.items() if k.startswith("observation.")}
     observation = observation_keys if observation_keys else None
 
-    # Extract padding, task, index, and task_index keys for complementary data
+    # Extract padding and task keys for complementary data
     pad_keys = {k: v for k, v in batch.items() if "_is_pad" in k}
     task_key = {"task": batch["task"]} if "task" in batch else {}
-    index_key = {"index": batch["index"]} if "index" in batch else {}
-    task_index_key = {"task_index": batch["task_index"]} if "task_index" in batch else {}
-    complementary_data = (
-        {**pad_keys, **task_key, **index_key, **task_index_key}
-        if pad_keys or task_key or index_key or task_index_key
-        else {}
-    )
+    complementary_data = {**pad_keys, **task_key} if pad_keys or task_key else {}
 
     transition: EnvTransition = {
         TransitionKey.OBSERVATION: observation,
@@ -239,7 +231,7 @@ def _default_transition_to_batch(transition: EnvTransition) -> dict[str, Any]:
         "info": transition.get(TransitionKey.INFO, {}),
     }
 
-    # Add padding, task, index, and task_index data from complementary_data
+    # Add padding and task data from complementary_data
     complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
     if complementary_data:
         pad_data = {k: v for k, v in complementary_data.items() if "_is_pad" in k}
@@ -248,12 +240,6 @@ def _default_transition_to_batch(transition: EnvTransition) -> dict[str, Any]:
         if "task" in complementary_data:
             batch["task"] = complementary_data["task"]
 
-        if "index" in complementary_data:
-            batch["index"] = complementary_data["index"]
-
-        if "task_index" in complementary_data:
-            batch["task_index"] = complementary_data["task_index"]
-
     # Handle observation - flatten dict to observation.* keys if it's a dict
     observation = transition.get(TransitionKey.OBSERVATION)
     if isinstance(observation, dict):
@@ -356,10 +342,7 @@ class RobotProcessor(ModelHubMixin):
                 hook(idx, current_transition)
 
         # Convert back to original format if needed
-        if called_with_batch or self.to_output is not _default_transition_to_batch:
-            return self.to_output(current_transition)
-        else:
-            return current_transition
+        return self.to_output(current_transition) if called_with_batch else current_transition
 
     def _prepare_transition(self, data: EnvTransition | dict[str, Any]) -> tuple[EnvTransition, bool]:
         """Prepare and validate transition data for processing.
@@ -592,9 +575,10 @@ class RobotProcessor(ModelHubMixin):
             if config_filename is None:
                 # Try common config names
                 common_names = [
-                    "robot_processor.json",
-                    "robot_preprocessor.json",
-                    "robot_postprocessor.json",
+                    "processor.json",
+                    "preprocessor.json",
+                    "postprocessor.json",
+                    "robotprocessor.json",
                 ]
                 config_path = None
                 for name in common_names:
@@ -824,15 +808,23 @@ class RobotProcessor(ModelHubMixin):
                     f"Step {i} ({type(step).__name__}) must define __call__(transition) -> EnvTransition"
                 )
 
-    def transform_features(self, initial_features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+            fc = getattr(step, "feature_contract", None)
+            if not callable(fc):
+                raise TypeError(
+                    f"Step {i} ({type(step).__name__}) must define feature_contract(features) -> dict[str, Any]"
+                )
+
+    def feature_contract(self, initial_features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         """
-        Apply ALL steps in order. Only if a step has a features method, it will be called.
-        We aggregate the dataset features of all steps.
+        Apply ALL steps in order. Each step must implement
+        feature_contract(features) and return a dict (full or incremental schema).
         """
         features: dict[str, PolicyFeature] = deepcopy(initial_features)
 
         for _, step in enumerate(self.steps):
-            out = step.transform_features(features)
+            out = step.feature_contract(features)
+            if not isinstance(out, dict):
+                raise TypeError(f"{step.__class__.__name__}.feature_contract must return dict[str, Any]")
             features = out
         return features
 
@@ -892,7 +884,7 @@ class ObservationProcessor:
     def reset(self) -> None:
         pass
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features
 
 
@@ -952,7 +944,7 @@ class ActionProcessor:
     def reset(self) -> None:
         pass
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features
 
 
@@ -1011,7 +1003,7 @@ class RewardProcessor:
     def reset(self) -> None:
         pass
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features
 
 
@@ -1075,7 +1067,7 @@ class DoneProcessor:
     def reset(self) -> None:
         pass
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features
 
 
@@ -1135,7 +1127,7 @@ class TruncatedProcessor:
     def reset(self) -> None:
         pass
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features
 
 
@@ -1200,7 +1192,7 @@ class InfoProcessor:
     def reset(self) -> None:
         pass
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features
 
 
@@ -1246,7 +1238,7 @@ class ComplementaryDataProcessor:
     def reset(self) -> None:
         pass
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features
 
 
@@ -1268,5 +1260,5 @@ class IdentityProcessor:
     def reset(self) -> None:
         pass
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features

src/lerobot/processor/rename_processor.py

@@ -43,7 +43,7 @@ class RenameProcessor(ObservationProcessor):
     def get_config(self) -> dict[str, Any]:
         return {"rename_map": self.rename_map}
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         """Transforms:
         - Each key in the observation that appears in `rename_map` is renamed to its value.
         - Keys not in `rename_map` remain unchanged.

src/lerobot/processor/tokenizer_processor.py

@@ -1,229 +0,0 @@
-"""
-Tokenizer processor for handling text tokenization in robot transitions.
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any
-
-import torch
-
-from lerobot.configs.types import FeatureType, PolicyFeature
-from lerobot.constants import OBS_LANGUAGE
-from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
-from lerobot.utils.import_utils import _transformers_available
-
-if TYPE_CHECKING or _transformers_available:
-    from transformers import AutoTokenizer
-else:
-    AutoTokenizer = None
-
-
-@dataclass
-@ProcessorStepRegistry.register(name="tokenizer_processor")
-class TokenizerProcessor:
-    """Tokenizes text tasks in complementary data using a huggingface tokenizer.
-
-    This processor handles tokenization of task strings found in the complementary_data
-    using a specified pretrained tokenizer from Hugging Face. It adds tokenized versions
-    to the observation data for model processing while preserving the original task string.
-
-    The processor supports both single strings and lists of strings as task inputs.
-
-    Args:
-        tokenizer_name: Name of the pretrained tokenizer to load from Hugging Face Hub
-            (e.g., "bert-base-uncased", "microsoft/DialoGPT-medium"). This will be used
-            with AutoTokenizer.from_pretrained(). If tokenizer is provided, this is ignored.
-        tokenizer: A tokenizer object (e.g., from transformers library) that implements
-            the __call__ method. If provided, tokenizer_name is ignored. This parameter
-            is not serialized and must be provided via overrides when loading.
-        max_length: Maximum sequence length for tokenization. Defaults to 512.
-        task_key: Key in complementary_data containing the task text. Defaults to "task".
-        padding: Padding strategy for tokenization. Defaults to "max_length".
-        truncation: Whether to truncate sequences longer than max_length. Defaults to True.
-
-    Examples:
-        Using tokenizer name (auto-loaded):
-        ```python
-        processor = TokenizerProcessor(tokenizer_name="bert-base-uncased", max_length=128)
-        ```
-
-        Using custom tokenizer object:
-        ```python
-        from transformers import AutoTokenizer
-
-        custom_tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-        processor = TokenizerProcessor(tokenizer=custom_tokenizer, max_length=128)
-        ```
-    """
-
-    tokenizer_name: str | None = None
-    tokenizer: Any | None = None  # Otherwise transformers is not available in the core dependencies
-    max_length: int = 512
-    task_key: str = "task"
-    padding_side: str = "right"
-    padding: str = "max_length"
-    truncation: bool = True
-
-    # Internal tokenizer instance (not serialized)
-    _tokenizer: Any = field(default=None, init=False, repr=False)
-
-    def __post_init__(self):
-        """Initialize the tokenizer from the provided tokenizer or tokenizer name."""
-        if not _transformers_available:
-            raise ImportError(
-                "The 'transformers' library is not installed. "
-                "Please install it with `pip install 'lerobot[transformers-dep]'` to use TokenizerProcessor."
-            )
-
-        if self.tokenizer is not None:
-            # Use provided tokenizer object directly
-            self._tokenizer = self.tokenizer
-        elif self.tokenizer_name is not None:
-            if AutoTokenizer is None:
-                raise ImportError("AutoTokenizer is not available")
-            self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
-        else:
-            raise ValueError(
-                "Either 'tokenizer' or 'tokenizer_name' must be provided. "
-                "Pass a tokenizer object directly or a tokenizer name to auto-load."
-            )
-
-    def get_task(self, transition: EnvTransition) -> list[str] | None:
-        """Extract and normalize task from complementary data.
-
-        Args:
-            transition: Input transition containing complementary_data.
-
-        Returns:
-            List of task strings if task is present, None otherwise.
-        """
-        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
-        if complementary_data is None:
-            return None
-
-        if self.task_key not in complementary_data:
-            return None
-
-        task = complementary_data[self.task_key]
-        if task is None:
-            return None
-
-        # Convert to list of strings
-        if isinstance(task, str):
-            return [task]
-        elif isinstance(task, list) and all(isinstance(t, str) for t in task):
-            return task
-
-        return None
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        """Process the transition by tokenizing the task text.
-
-        Args:
-            transition: Input transition containing complementary_data with task text.
-
-        Returns:
-            Modified transition with tokenized task added to observation.
-
-        Raises:
-            ValueError: If tokenizer initialization failed.
-        """
-        task = self.get_task(transition)
-        if task is None:
-            return transition
-
-        # Tokenize the task
-        tokenized_prompt = self._tokenize_text(task)
-
-        # Get or create observation dict
-        observation = transition.get(TransitionKey.OBSERVATION)
-        if observation is None:
-            observation = {}
-        else:
-            observation = dict(observation)  # Make a copy
-
-        # Add tokenized data to observation
-        input_ids = tokenized_prompt["input_ids"]
-        attention_mask = tokenized_prompt.get("attention_mask")
-        if attention_mask is None:
-            # Some tokenizers (e.g., SigLIP text) may not return attention_mask; default to ones
-            attention_mask = torch.ones_like(input_ids)
-        observation[f"{OBS_LANGUAGE}.tokens"] = input_ids
-        observation[f"{OBS_LANGUAGE}.attention_mask"] = attention_mask.to(dtype=torch.bool)
-
-        transition[TransitionKey.OBSERVATION.value] = observation  # type: ignore[misc]
-        return transition
-
-    def _tokenize_text(self, text: str | list[str]) -> dict[str, torch.Tensor]:
-        """Tokenize text using the configured tokenizer.
-
-        Args:
-            text: Text string or list of strings to tokenize.
-
-        Returns:
-            Dictionary containing tokenized output with keys like 'input_ids', 'attention_mask'.
-        """
-        return self._tokenizer(
-            text,
-            max_length=self.max_length,
-            truncation=self.truncation,
-            padding=self.padding,
-            padding_side=self.padding_side,
-            return_tensors="pt",
-        )
-
-    def get_config(self) -> dict[str, Any]:
-        """Return configuration for serialization.
-
-        Note: Only tokenizer_name is saved, not the tokenizer object itself.
-        When loading, provide the tokenizer via overrides if needed.
-        """
-        config = {
-            "max_length": self.max_length,
-            "task_key": self.task_key,
-            "padding_side": self.padding_side,
-            "padding": self.padding,
-            "truncation": self.truncation,
-        }
-
-        # Only include tokenizer_name if it was used (not when tokenizer object was provided)
-        if self.tokenizer_name is not None:
-            config["tokenizer_name"] = self.tokenizer_name
-
-        return config
-
-    def state_dict(self) -> dict[str, torch.Tensor]:
-        """Return state dictionary (empty for this processor)."""
-        return {}
-
-    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
-        """Load state dictionary (no-op for this processor)."""
-        pass
-
-    def reset(self) -> None:
-        """Reset processor state (no-op for this processor)."""
-        pass
-
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        """Add tokenized task features to the feature contract.
-
-        Args:
-            features: Input feature dictionary.
-
-        Returns:
-            Updated feature dictionary with tokenized task features added.
-        """
-        # Add features for tokenized output if they don't exist
-        # Standard tokenizer output includes tokens and attention_mask
-        tokens_key = f"{OBS_LANGUAGE}.tokens"
-        attention_mask_key = f"{OBS_LANGUAGE}.attention_mask"
-
-        if tokens_key not in features:
-            features[tokens_key] = PolicyFeature(type=FeatureType.LANGUAGE, shape=(self.max_length,))
-
-        if attention_mask_key not in features:
-            features[attention_mask_key] = PolicyFeature(type=FeatureType.LANGUAGE, shape=(self.max_length,))
-
-        return features

src/lerobot/record.py

@@ -18,7 +18,7 @@ Records a dataset. Actions for the robot can be either generated by teleoperatio
 Example:
 
 ```shell
-python -m lerobot.record \
+lerobot-record \
     --robot.type=so100_follower \
     --robot.port=/dev/tty.usbmodem58760431541 \
     --robot.cameras="{laptop: {type: opencv, camera_index: 0, width: 640, height: 480}}" \
@@ -36,7 +36,7 @@ python -m lerobot.record \
 
 Example recording with bimanual so100:
 ```shell
-python -m lerobot.record \
+lerobot-record \
   --robot.type=bi_so100_follower \
   --robot.left_arm_port=/dev/tty.usbmodem5A460851411 \
   --robot.right_arm_port=/dev/tty.usbmodem5A460812391 \
@@ -59,7 +59,7 @@ python -m lerobot.record \
 
 import logging
 import time
-from dataclasses import asdict, dataclass, field
+from dataclasses import asdict, dataclass
 from pathlib import Path
 from pprint import pformat
 
@@ -72,19 +72,9 @@ from lerobot.configs import parser
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.datasets.image_writer import safe_stop_image_writer
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.utils import hw_to_dataset_features
-from lerobot.datasets.video_utils import VideoEncodingManager
-from lerobot.policies.factory import make_policy, make_processor
+from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
+from lerobot.policies.factory import make_policy
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.processor import RobotProcessor
-from lerobot.processor.converters import (
-    to_dataset_frame,
-    to_output_robot_action,
-    to_transition_robot_observation,
-    to_transition_teleop_action,
-)
-from lerobot.processor.normalize_processor import rename_stats
-from lerobot.processor.pipeline import IdentityProcessor, TransitionKey
 from lerobot.robots import (  # noqa: F401
     Robot,
     RobotConfig,
@@ -158,8 +148,6 @@ class DatasetRecordConfig:
     # Number of episodes to record before batch encoding videos
     # Set to 1 for immediate encoding (default behavior), or higher for batched encoding
     video_encoding_batch_size: int = 1
-    # Rename map for the observation to override the image and state keys
-    rename_map: dict[str, str] = field(default_factory=dict)
 
     def __post_init__(self):
         if self.single_task is None:
@@ -198,36 +186,6 @@ class RecordConfig:
         return ["policy"]
 
 
-""" --------------- record_loop() data flow --------------------------
-       [ Robot ]
-           V
-     [ robot.get_observation() ] ---> raw_obs
-           V
-     [ robot_observation_processor ] ---> obs_transition
-           V
-     .-----( ACTION LOGIC )------------------.
-     V                                       V
-     [ From Teleoperator ]                   [ From Policy ]
-     |                                       |
-     |  [teleop.get_action] -> raw_action    |   [predict_action]
-     |          |                            |          |
-     |          V                            |          V
-     | [teleop_action_processor]             |          |
-     |          |                            |          |
-     '---> teleop_transition                 '---> policy_transition
-     |                                       |
-     '-------------------------.-------------'
-                               V
-                  [ robot_action_processor ] --> robot_action_to_send
-                               V
-                    [ robot.send_action() ] -- (Robot Executes)
-                               V
-        ( Transitions are merged & added to Dataset )
-                               V
-                  ( Rerun Log / Loop Wait )
-"""
-
-
 @safe_stop_image_writer
 def record_loop(
     robot: Robot,
@@ -236,30 +194,15 @@ def record_loop(
     dataset: LeRobotDataset | None = None,
     teleop: Teleoperator | list[Teleoperator] | None = None,
     policy: PreTrainedPolicy | None = None,
-    preprocessor: RobotProcessor | None = None,
-    postprocessor: RobotProcessor | None = None,
     control_time_s: int | None = None,
-    teleop_action_processor: RobotProcessor | None = None,  # runs after teleop
-    robot_action_processor: RobotProcessor | None = None,  # runs before robot
-    robot_observation_processor: RobotProcessor | None = None,  # runs after robot
     single_task: str | None = None,
     display_data: bool = False,
 ):
-    teleop_action_processor = teleop_action_processor or RobotProcessor(
-        steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr
-    )
-    robot_action_processor = robot_action_processor or RobotProcessor(
-        steps=[IdentityProcessor()], to_transition=lambda tr: tr, to_output=to_output_robot_action
-    )
-    robot_observation_processor = robot_observation_processor or RobotProcessor(
-        steps=[IdentityProcessor()], to_transition=to_transition_robot_observation, to_output=lambda tr: tr
-    )
-
     if dataset is not None and dataset.fps != fps:
         raise ValueError(f"The dataset fps should be equal to requested fps ({dataset.fps} != {fps}).")
 
     teleop_arm = teleop_keyboard = None
-    if isinstance(teleop, list):  # For LeKiwi
+    if isinstance(teleop, list):
         teleop_keyboard = next((t for t in teleop if isinstance(t, KeyboardTeleop)), None)
         teleop_arm = next(
             (
@@ -275,20 +218,9 @@ def record_loop(
                 "For multi-teleop, the list must contain exactly one KeyboardTeleop and one arm teleoperator. Currently only supported for LeKiwi robot."
             )
 
-    # Reset policy and processor if they are provided
-    if policy is not None and preprocessor is not None and postprocessor is not None:
+    # if policy is given it needs cleaning up
+    if policy is not None:
         policy.reset()
-        preprocessor.reset()
-        postprocessor.reset()
-
-    # Reset custom pipelines
-    teleop_action_processor.reset()
-    robot_action_processor.reset()
-    robot_observation_processor.reset()
-
-    policy_transition = None
-    teleop_transition = None
-    obs_transition = None
 
     timestamp = 0
     start_episode_t = time.perf_counter()
@@ -299,87 +231,51 @@ def record_loop(
             events["exit_early"] = False
             break
 
-        # Get robot observation
-        obs = robot.get_observation()
+        observation = robot.get_observation()
 
-        # Applies a pipeline to the raw robot observation, default is IdentityProcessor
-        obs_transition = robot_observation_processor(obs)
-
-        # Get action from either policy or teleop
-        if policy is not None and preprocessor is not None and postprocessor is not None:
-            if dataset is not None:
-                observation_frame = to_dataset_frame(
-                    obs_transition, dataset.features
-                )  # Convert the observation to the dataset format
+        if policy is not None or dataset is not None:
+            observation_frame = build_dataset_frame(dataset.features, observation, prefix="observation")
 
+        if policy is not None:
             action_values = predict_action(
-                observation=observation_frame,
-                policy=policy,
-                device=get_safe_torch_device(policy.config.device),
-                preprocessor=preprocessor,
-                postprocessor=postprocessor,
-                use_amp=policy.config.use_amp,
+                observation_frame,
+                policy,
+                get_safe_torch_device(policy.config.device),
+                policy.config.use_amp,
                 task=single_task,
                 robot_type=robot.robot_type,
             )
-
-            action_names = dataset.features["action"]["names"]
-            policy_action = {f"action.{name}": float(action_values[i]) for i, name in enumerate(action_names)}
-            policy_transition = {
-                TransitionKey.ACTION: policy_action,
-                TransitionKey.COMPLEMENTARY_DATA: {},
-            }
-
-        elif isinstance(teleop, Teleoperator):
-            act = teleop.get_action()
-
-            # Applies a pipeline to the raw teleop action, default is IdentityProcessor
-            teleop_transition = teleop_action_processor(act)
-
-        elif isinstance(teleop, list):
+            action = {key: action_values[i].item() for i, key in enumerate(robot.action_features)}
+        elif policy is None and isinstance(teleop, Teleoperator):
+            action = teleop.get_action()
+        elif policy is None and isinstance(teleop, list):
+            # TODO(pepijn, steven): clean the record loop for use of multiple robots (possibly with pipeline)
             arm_action = teleop_arm.get_action()
             arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
+
             keyboard_action = teleop_keyboard.get_action()
             base_action = robot._from_keyboard_to_base_action(keyboard_action)
-            act = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
-            teleop_transition = teleop_action_processor(act)
+
+            action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
         else:
             logging.info(
-                "No policy or teleoperator provided, skipping action generation. "
-                "This is likely to happen during environment reset."
+                "No policy or teleoperator provided, skipping action generation."
+                "This is likely to happen when resetting the environment without a teleop device."
+                "The robot won't be at its rest position at the start of the next episode."
             )
-            # Still continue to next loop to respect timing
+            continue
 
-        # Applies a pipeline to the action, default is IdentityProcessor
-        # IMPORTANT: action_pipeline.to_output must return a dict suitable for robot.send_action()
-        if policy_transition is not None:
-            robot_action_to_send = robot_action_processor(policy_transition)
-        else:
-            robot_action_to_send = robot_action_processor(teleop_transition)
-
-        # Send action to robot
         # Action can eventually be clipped using `max_relative_target`,
-        # so action actually sent is saved in the dataset. action = postprocessor.process(action)
-        # TODO(pepijn, adil): we should use a pipeline step to clip the action, so the sent action is the action that we input to the robot.
-        _ = robot.send_action(robot_action_to_send)
+        # so action actually sent is saved in the dataset.
+        sent_action = robot.send_action(action)
 
-        # Write to dataset
         if dataset is not None:
-            # If to_dataset_frame is provided, use it to merge the transitions.
-            merged = []
-            if obs_transition is not None:  # The observation from the robot
-                merged.append(obs_transition)
-            if teleop_transition is not None:  # The action from teleop
-                merged.append(teleop_transition)
-            if policy_transition is not None:  # The action from policy
-                merged.append(policy_transition)
-            frame = to_dataset_frame(
-                merged if len(merged) > 1 else merged[0], dataset.features
-            )  # Convert the observation to the dataset format
-            dataset.add_frame(frame, task=single_task)
+            action_frame = build_dataset_frame(dataset.features, sent_action, prefix="action")
+            frame = {**observation_frame, **action_frame, "task": single_task}
+            dataset.add_frame(frame)
 
         if display_data:
-            log_rerun_data([obs_transition, teleop_transition or policy_transition])
+            log_rerun_data(observation, action)
 
         dt_s = time.perf_counter() - start_loop_t
         busy_wait(1 / fps - dt_s)
@@ -405,7 +301,6 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
         dataset = LeRobotDataset(
             cfg.dataset.repo_id,
             root=cfg.dataset.root,
-            batch_encoding_size=cfg.dataset.video_encoding_batch_size,
         )
 
         if hasattr(robot, "cameras") and len(robot.cameras) > 0:
@@ -426,23 +321,10 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
             use_videos=cfg.dataset.video,
             image_writer_processes=cfg.dataset.num_image_writer_processes,
             image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras),
-            batch_encoding_size=cfg.dataset.video_encoding_batch_size,
         )
 
     # Load pretrained policy
     policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta)
-    preprocessor = None
-    postprocessor = None
-    if cfg.policy is not None:
-        preprocessor, postprocessor = make_processor(
-            policy_cfg=cfg.policy,
-            pretrained_path=cfg.policy.pretrained_path,
-            dataset_stats=rename_stats(dataset.meta.stats, cfg.dataset.rename_map),
-            preprocessor_overrides={
-                "device_processor": {"device": cfg.policy.device},
-                "rename_processor": {"rename_map": cfg.dataset.rename_map},
-            },
-        )
 
     robot.connect()
     if teleop is not None:
@@ -450,49 +332,46 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
 
     listener, events = init_keyboard_listener()
 
-    with VideoEncodingManager(dataset):
-        recorded_episodes = 0
-        while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]:
-            log_say(f"Recording episode {dataset.num_episodes}", cfg.play_sounds)
+    recorded_episodes = 0
+    while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]:
+        log_say(f"Recording episode {dataset.num_episodes}", cfg.play_sounds)
+        record_loop(
+            robot=robot,
+            events=events,
+            fps=cfg.dataset.fps,
+            teleop=teleop,
+            policy=policy,
+            dataset=dataset,
+            control_time_s=cfg.dataset.episode_time_s,
+            single_task=cfg.dataset.single_task,
+            display_data=cfg.display_data,
+        )
+
+        # Execute a few seconds without recording to give time to manually reset the environment
+        # Skip reset for the last episode to be recorded
+        if not events["stop_recording"] and (
+            (recorded_episodes < cfg.dataset.num_episodes - 1) or events["rerecord_episode"]
+        ):
+            log_say("Reset the environment", cfg.play_sounds)
             record_loop(
                 robot=robot,
                 events=events,
                 fps=cfg.dataset.fps,
                 teleop=teleop,
-                policy=policy,
-                preprocessor=preprocessor,
-                postprocessor=postprocessor,
-                dataset=dataset,
-                control_time_s=cfg.dataset.episode_time_s,
+                control_time_s=cfg.dataset.reset_time_s,
                 single_task=cfg.dataset.single_task,
                 display_data=cfg.display_data,
             )
 
-            # Execute a few seconds without recording to give time to manually reset the environment
-            # Skip reset for the last episode to be recorded
-            if not events["stop_recording"] and (
-                (recorded_episodes < cfg.dataset.num_episodes - 1) or events["rerecord_episode"]
-            ):
-                log_say("Reset the environment", cfg.play_sounds)
-                record_loop(
-                    robot=robot,
-                    events=events,
-                    fps=cfg.dataset.fps,
-                    teleop=teleop,
-                    control_time_s=cfg.dataset.reset_time_s,
-                    single_task=cfg.dataset.single_task,
-                    display_data=cfg.display_data,
-                )
+        if events["rerecord_episode"]:
+            log_say("Re-record episode", cfg.play_sounds)
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-            if events["rerecord_episode"]:
-                log_say("Re-record episode", cfg.play_sounds)
-                events["rerecord_episode"] = False
-                events["exit_early"] = False
-                dataset.clear_episode_buffer()
-                continue
-
-            dataset.save_episode()
-            recorded_episodes += 1
+        dataset.save_episode()
+        recorded_episodes += 1
 
     log_say("Stop recording", cfg.play_sounds, blocking=True)
 
@@ -510,5 +389,9 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
     return dataset
 
 
-if __name__ == "__main__":
+def main():
     record()
+
+
+if __name__ == "__main__":
+    main()