Merge branch 'feat/add_pi' into feat/validate_pi_libero

for more information, see https://pre-commit.ci

README.md

@@ -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,30 +246,20 @@ 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
-  ├ 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')
+  ├ 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", ...]`)
+```
 
 A `LeRobotDataset` is serialised using several widespread file formats for each of its parts, namely:
 
@@ -293,7 +283,7 @@ lerobot-eval \
     --eval.n_episodes=10 \
     --policy.use_amp=false \
     --policy.device=cuda
-````
+```
 
 Note: After training your own policy, you can re-evaluate the checkpoints with:
 

benchmarks/video/run_video_benchmark.py

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

docker/Dockerfile.user

@@ -29,7 +29,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
 
 # Install system dependencies and uv (as root)
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    build-essential git curl libglib2.0-0 libegl1-mesa ffmpeg \
+    build-essential git curl libglib2.0-0 libegl1-mesa-dev ffmpeg \
     libusb-1.0-0-dev speech-dispatcher libgeos-dev portaudio19-dev \
     && curl -LsSf https://astral.sh/uv/install.sh | sh \
     && mv /root/.local/bin/uv /usr/local/bin/uv \

docs/source/_toctree.yml

@@ -19,8 +19,8 @@
     title: Train RL in Simulation
   - local: async
     title: Use Async Inference
-  - local: porting_datasets_v3
-    title: Porting Large Datasets
+  - local: libero
+    title: Using LIBERO
   title: "Tutorials"
 - sections:
   - local: smolvla
@@ -37,10 +37,14 @@
     title: Koch v1.1
   - local: lekiwi
     title: LeKiwi
+  - local: reachy2
+    title: Reachy 2
   title: "Robots"
 - sections:
   - local: notebooks
     title: Notebooks
+  - local: feetech
+    title: Updating Feetech Firmware
   title: "Resources"
 - sections:
   - local: contributing

docs/source/feetech.mdx

@@ -0,0 +1,71 @@
+# Feetech Motor Firmware Update
+
+This tutorial guides you through updating the firmware of Feetech motors using the official Feetech software.
+
+## Prerequisites
+
+- Windows computer (Feetech software is only available for Windows)
+- Feetech motor control board
+- USB cable to connect the control board to your computer
+- Feetech motors connected to the control board
+
+## Step 1: Download Feetech Software
+
+1. Visit the official Feetech software download page: [https://www.feetechrc.com/software.html](https://www.feetechrc.com/software.html)
+2. Download the latest version of the Feetech debugging software (FD)
+3. Install the software on your Windows computer
+
+## Step 2: Hardware Setup
+
+1. Connect your Feetech motors to the motor control board
+2. Connect the motor control board to your Windows computer via USB cable
+3. Ensure power is supplied to the motors
+
+## Step 3: Configure Connection
+
+1. Launch the Feetech debugging software
+2. Select the correct COM port from the port dropdown menu
+   - If unsure which port to use, check Windows Device Manager under "Ports (COM & LPT)"
+3. Set the appropriate baud rate (typically 1000000 for most Feetech motors)
+4. Click "Open" to establish communication with the control board
+
+## Step 4: Scan for Motors
+
+1. Once connected, click the "Search" button to detect all connected motors
+2. The software will automatically discover and list all motors on the bus
+3. Each motor will appear with its ID number
+
+## Step 5: Update Firmware
+
+For each motor you want to update:
+
+1. **Select the motor** from the list by clicking on it
+2. **Click on Upgrade tab**:
+3. **Click on Online button**:
+   - If an potential firmware update is found, it will be displayed in the box
+4. **Click on Upgrade button**:
+   - The update progress will be displayed
+
+## Step 6: Verify Update
+
+1. After the update completes, the software should automatically refresh the motor information
+2. Verify that the firmware version has been updated to the expected version
+
+## Important Notes
+
+⚠️ **Warning**: Do not disconnect power or USB during firmware updates, it will potentially brick the motor.
+
+## Bonus: Motor Debugging on Linux/macOS
+
+For debugging purposes only, you can use the open-source Feetech Debug Tool:
+
+- **Repository**: [FT_SCServo_Debug_Qt](https://github.com/CarolinePascal/FT_SCServo_Debug_Qt/tree/fix/port-search-timer)
+
+### Installation Instructions
+
+Follow the instructions in the repository to install the tool, for Ubuntu you can directly install it, for MacOS you need to build it from source.
+
+**Limitations:**
+
+- This tool is for debugging and parameter adjustment only
+- Firmware updates must still be done on Windows with official Feetech software

docs/source/libero.mdx

@@ -0,0 +1,230 @@
+# LIBERO
+
+**LIBERO** is a benchmark designed to study **lifelong robot learning**. The idea is that robots won’t just be pretrained once in a factory, they’ll need to keep learning and adapting with their human users over time. This ongoing adaptation is called **lifelong learning in decision making (LLDM)**, and it’s a key step toward building robots that become truly personalized helpers. The benchmark was first introduced in the [LIBERO paper](https://arxiv.org/abs/2306.03310) and the [original repository](https://github.com/Lifelong-Robot-Learning/LIBERO).
+
+To make progress on this challenge, LIBERO provides a set of standardized tasks that focus on **knowledge transfer**: how well a robot can apply what it has already learned to new situations. By evaluating on LIBERO, different algorithms can be compared fairly and researchers can build on each other’s work.
+
+LIBERO includes **five task suites**:
+
+- **LIBERO-Spatial (`libero_spatial`)** – tasks that require reasoning about spatial relations.
+- **LIBERO-Object (`libero_object`)** – tasks centered on manipulating different objects.
+- **LIBERO-Goal (`libero_goal`)** – goal-conditioned tasks where the robot must adapt to changing targets.
+- **LIBERO-90 (`libero_90`)** – 90 short-horizon tasks from the LIBERO-100 collection.
+- **LIBERO-Long (`libero_10`)** – 10 long-horizon tasks from the LIBERO-100 collection.
+
+Together, these suites cover **130 tasks**, ranging from simple object manipulations to complex multi-step scenarios. LIBERO is meant to grow over time, and to serve as a shared benchmark where the community can test and improve lifelong learning algorithms.
+
+![An overview of the LIBERO benchmark](https://libero-project.github.io/assets/img/libero/fig1.png)
+_Figure 1: An overview of the LIBERO benchmark._
+
+## Evaluating with LIBERO
+
+At **LeRobot**, we ported [LIBERO](https://github.com/Lifelong-Robot-Learning/LIBERO) into our framework and used it primarily to **benchmark [SmolVLA](https://huggingface.co/docs/lerobot/en/smolvla)**, our lightweight Vision-Language-Action model, comparing it against state-of-the-art VLA models such as Pi0, OpenVLA, Octo, and Diffusion Policy.
+
+LIBERO is now part of our **multi-eval supported simulation**, allowing you to benchmark your policies either on a **single suite of tasks** or across **multiple suites at once** with just a single flag.
+
+To install LIBERO, first follow the [LeRobot Installation Guide](https://huggingface.co/docs/lerobot/installation).
+Once LeRobot is installed, there are two options:
+
+1. **Install via pip** (recommended):
+
+   ```bash
+   pip install "lerobot[libero,smolvla]"
+   ```
+
+2. **Install from source**:
+   ```bash
+   git clone https://github.com/huggingface/lerobot.git
+   cd lerobot
+   pip install -e ".[libero,smolvla]"
+   ```
+
+### Single-suite evaluation
+
+Evaluate a policy on one LIBERO suite:
+
+```bash
+python src/lerobot/scripts/eval.py \
+  --policy.path="your-policy-id" \
+  --env.type=libero \
+  --env.task=libero_object \
+  --env.multitask_eval=False \
+  --eval.batch_size=2 \
+  --eval.n_episodes=3
+```
+
+- `--env.task` picks the suite (`libero_object`, `libero_spatial`, etc.).
+- `--eval.batch_size` controls how many environments run in parallel.
+- `--eval.n_episodes` sets how many episodes to run in total.
+
+---
+
+### Multi-suite evaluation
+
+Benchmark a policy across multiple suites at once:
+
+```bash
+python src/lerobot/scripts/eval.py \
+  --policy.path="your-policy-id" \
+  --env.type=libero \
+  --env.task=libero_object \
+  --env.multitask_eval=True \
+  --eval.batch_size=1 \
+  --eval.n_episodes=2
+```
+
+- Pass a comma-separated list to `--env.task` for multi-suite evaluation.
+- Set `-env.multitask_eval=True` to enable evaluation across all tasks in those suites.
+
+### Policy inputs and outputs
+
+When using LIBERO through LeRobot, policies interact with the environment via **observations** and **actions**:
+
+- **Observations**
+  - `observation.state` – proprioceptive features (agent state).
+  - `observation.images.image` – main camera view (`agentview_image`).
+  - `observation.images.image2` – wrist camera view (`robot0_eye_in_hand_image`).
+
+  ⚠️ **Note:** LeRobot enforces the `.images.*` prefix for any visual features. Make sure your dataset metadata keys match this convention when evaluating.
+
+  ## Input Features and Metadata Alignment
+
+  To train or evaluate a policy, you use `make_policy`, which builds a feature-naming dictionary for the observations the policy expects.
+  This mapping can come from:
+  - Dataset metadata
+  - The evaluation environment
+  - The policy path (if a pretrained repo ID is provided)
+
+  ### Common Issues
+
+  A common problem is when the keys in the dataset, environment, and policy config do not match. For example:
+  - `wrist_image` vs `observation.images.image2`
+  - `observation.image2` (as in SmolVLA) vs the `.images.*` prefix convention
+
+  Such mismatches will cause `KeyError`s. This may be due to assumptions in `make_policy` or missing error handling.
+
+  ***
+
+  ### How to Check Expected Features
+  - Open your policy config (`config.json`), e.g. [example here](https://huggingface.co/jadechoghari/smolvla-libero/blob/main/config.json).
+  - Or add a breakpoint in `train.py` and inspect:
+
+    ````python
+    print(policy.config.input_features)
+    To ensure you can just check what your policy expects as `input_features`:
+
+    - Open your policy config (`config.json`), e.g. [example here](https://huggingface.co/jadechoghari/smolvla-libero/blob/main/config.json).
+    - Or add a breakpoint in `train.py` and inspect:
+      ```python
+      print(policy.config.input_features)
+    Fixing KeyErrors (Preprocessing Example)
+    ````
+
+## Fixing KeyErrors (Preprocessing Example)
+
+If your dataset columns do not follow the expected naming, you can rename them in-place before training:
+
+````python
+import pyarrow.parquet as pq
+import shutil
+
+def rename_columns(parquet_path, rename_map):
+    table = pq.read_table(parquet_path)
+    schema = table.schema
+    new_names = [rename_map.get(name, name) for name in schema.names]
+    renamed_table = table.rename_columns(new_names)
+    backup_path = parquet_path + ".bak"
+    shutil.copy(parquet_path, backup_path)
+    pq.write_table(renamed_table, parquet_path)
+    print(f"patched {parquet_path}, backup at {backup_path}")
+
+# example mapping: align dataset keys to LeRobot convention
+rename_map = {
+    "image": "observation.images.image",
+    "wrist_image": "observation.images.image2",
+}
+
+rename_columns("episode_000001.parquet", rename_map)
+
+
+
+- **Actions**
+  - Continuous control values in a `Box(-1, 1, shape=(7,))` space.
+
+We also provide a notebook for quick testing:
+Training with LIBERO
+
+## Training with LIBERO
+
+When training on LIBERO tasks, make sure your dataset parquet and metadata keys follow the LeRobot convention.
+
+The environment expects:
+
+- `observation.state` → 8-dim agent state
+- `observation.images.image` → main camera (`agentview_image`)
+- `observation.images.image2` → wrist camera (`robot0_eye_in_hand_image`)
+
+⚠️ Cleaning the dataset upfront is **cleaner and more efficient** than remapping keys inside the code. We plan to provide a script to easily preprocess such data.
+To avoid potential mismatches and `KeyError`s, we provide a **preprocessed LIBERO dataset** that is fully compatible with the current LeRobot codebase and requires no additional manipulations.
+
+- 🔗 [Preprocessed LIBERO dataset (Hugging Face LeRobot org)](https://huggingface.co/datasets/HuggingFaceVLA/libero)
+- 🔗 [Original LIBERO dataset (physical-intelligence)](https://huggingface.co/datasets/physical-intelligence/libero)
+
+The preprocessed dataset follows LeRobot naming conventions (e.g., `.images.*` prefix for visual features) and aligns with policy configs out-of-the-box.
+The original dataset is acknowledged here as the primary source.
+---
+
+### Example training command
+
+```bash
+python src/lerobot/scripts/train.py \
+  --policy.type=smolvla \
+  --policy.repo_id=${HF_USER}/libero-test \
+  --dataset.repo_id=jadechoghari/smol-libero3 \
+  --env.type=libero \
+  --env.task=libero_10 \
+  --output_dir=./outputs/ \
+  --steps=100000 \
+  --batch_size=4 \
+  --env.multitask_eval=True \
+  --eval.batch_size=1 \
+  --eval.n_episodes=1 \
+  --eval_freq=1000 \
+````
+
+---
+
+### Note on rendering
+
+LeRobot uses MuJoCo for simulation. You need to set the rendering backend before training or evaluation:
+
+- `export MUJOCO_GL=egl` → for headless servers (e.g. HPC, cloud)
+
+---
+
+## Colab Note on Parallel Evaluation
+
+When running evaluation on Colab, you may encounter warnings such as:
+
+```
+UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
+```
+
+This happens because Colab’s rendering contexts are **not thread-safe**, and `ThreadPoolExecutor(max_workers=num_workers)` can trigger segfaults or leaked semaphore warnings.
+
+**Colab Note:**
+Parallel evaluation is not supported in Colab. To avoid these issues, run sequentially or disable multitask evaluation:
+
+Run sequentially:
+
+```bash
+--env.max_parallel_tasks=1
+```
+
+Or disable multitask evaluation:
+
+```bash
+--env.multitask_eval=False
+```
+
+If you want to take advantage of **parallel evaluation**, we recommend **not using Colab**. Instead, run locally or on a proper compute environment where multi-threaded rendering is easily supported.

docs/source/porting_datasets_v3.mdx

@@ -1,321 +0,0 @@
-# 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/reachy2.mdx

@@ -0,0 +1,288 @@
+# Reachy 2
+
+Reachy 2 is an open-source humanoid robot made by Pollen Robotics, specifically designed for the development of embodied AI and real-world applications.
+Check out [Pollen Robotics website](https://www.pollen-robotics.com/reachy/), or access [Reachy 2 documentation](https://docs.pollen-robotics.com/) for more information on the platform!
+
+## Teleoperate Reachy 2
+
+Currently, there are two ways to teleoperate Reachy 2:
+
+- Pollen Robotics’ VR teleoperation (not included in LeRobot).
+- Robot-to-robot teleoperation (use one Reachy 2 to control another).
+
+## Reachy 2 Simulation
+
+**(Linux only)** You can run Reachy 2 in simulation (Gazebo or MuJoCo) using the provided [Docker image](https://hub.docker.com/r/pollenrobotics/reachy2_core).
+
+1. Install [Docker Engine](https://docs.docker.com/engine/).
+2. Run (for MuJoCo):
+
+```
+docker run --rm -it \
+  --name reachy \
+  --privileged \
+  --network host \
+  --ipc host \
+  --device-cgroup-rule='c 189:* rwm' \
+  --group-add audio \
+  -e ROS_DOMAIN_ID="$ROS_DOMAIN_ID" \
+  -e DISPLAY="$DISPLAY" \
+  -e RCUTILS_CONSOLE_OUTPUT_FORMAT="[{severity}]: {message}" \
+  -e REACHY2_CORE_SERVICE_FAKE="${REACHY2_CORE_SERVICE_FAKE:-true}" \
+  -v /dev:/dev \
+  -v "$HOME/.reachy_config":/home/reachy/.reachy_config_override \
+  -v "$HOME/.reachy.log":/home/reachy/.ros/log \
+  -v /usr/lib/x86_64-linux-gnu:/opt/host-libs \
+  --entrypoint /package/launch.sh \
+  pollenrobotics/reachy2_core:1.7.5.9_deploy \
+  start_rviz:=true start_sdk_server:=true mujoco:=true
+```
+
+> If MuJoCo runs slowly (low simulation frequency), append `-e LD_LIBRARY_PATH="/opt/host-libs:$LD_LIBRARY_PATH" \` to the previous command to improve performance:
+>
+> ```
+> docker run --rm -it \
+>   --name reachy \
+>   --privileged \
+>   --network host \
+>   --ipc host \
+>   --device-cgroup-rule='c 189:* rwm' \
+>   --group-add audio \
+>   -e ROS_DOMAIN_ID="$ROS_DOMAIN_ID" \
+>   -e DISPLAY="$DISPLAY" \
+>   -e RCUTILS_CONSOLE_OUTPUT_FORMAT="[{severity}]: {message}" \
+>   -e REACHY2_CORE_SERVICE_FAKE="${REACHY2_CORE_SERVICE_FAKE:-true}" \
+>   -e LD_LIBRARY_PATH="/opt/host-libs:$LD_LIBRARY_PATH" \
+>   -v /dev:/dev \
+>   -v "$HOME/.reachy_config":/home/reachy/.reachy_config_override \
+>   -v "$HOME/.reachy.log":/home/reachy/.ros/log \
+>   -v /usr/lib/x86_64-linux-gnu:/opt/host-libs \
+>   --entrypoint /package/launch.sh \
+>   pollenrobotics/reachy2_core:1.7.5.9_deploy \
+>   start_rviz:=true start_sdk_server:=true mujoco:=true
+> ```
+
+## Setup
+
+### Prerequisites
+
+- On your robot, check the **service images** meet the minimum versions:
+  - **reachy2-core >= 1.7.5.2**
+  - **webrtc >= 2.0.1.1**
+
+Then, if you want to use VR teleoperation:
+
+- Install the [Reachy 2 teleoperation application](https://docs.pollen-robotics.com/teleoperation/teleoperation-introduction/discover-teleoperation/).
+  Use version **>=v1.2.0**
+
+We recommend using two computers: one for teleoperation (Windows required) and another for recording with LeRobot.
+
+### Install LeRobot
+
+Follow the [installation instructions](https://github.com/huggingface/lerobot#installation) to install LeRobot.
+
+Install LeRobot with Reachy 2 dependencies:
+
+```bash
+pip install -e ".[reachy2]"
+```
+
+### (Optional but recommended) Install pollen_data_acquisition_server
+
+How you manage Reachy 2 recording sessions is up to you, but the **easiest** way is to use this server so you can control sessions directly from the VR teleoperation app.
+
+> **Note:** Currently, only the VR teleoperation application works as a client for this server, so this step primarily targets teleoperation. You’re free to develop custom clients to manage sessions to your needs.
+
+In your LeRobot environment, install the server from source:
+
+```bash
+git clone https://github.com/pollen-robotics/pollen_data_acquisition_server.git
+cd pollen_data_acquisition_server
+pip install -e .
+```
+
+Find the [pollen_data_acquisition_server documentation here](https://github.com/pollen-robotics/pollen_data_acquisition_server).
+
+## Step 1: Recording
+
+### Get Reachy 2 IP address
+
+Before starting teleoperation and data recording, find the [robot's IP address](https://docs.pollen-robotics.com/getting-started/setup-reachy2/connect-reachy2/).
+We strongly recommend connecting all devices (PC and robot) via **Ethernet**.
+
+### Launch recording
+
+There are two ways to manage recording sessions when using the Reachy 2 VR teleoperation application:
+
+- **Using the data acquisition server (recommended for VR teleop)**: The VR app orchestrates sessions (via the server it tells LeRobot when to create datasets, start/stop episodes) while also controlling the robot’s motions.
+- **Using LeRobot’s record script**: LeRobot owns session control and decides when to start/stop episodes. If you also use the VR teleop app, it’s only for motion control.
+
+### Option 1: Using Pollen data acquisition server (recommended for VR teleop)
+
+Make sure you have installed pollen_data_acquisition_server, as explained in the Setup section.
+
+Launch the data acquisition server to be able to manage your session directly from the teleoperation application:
+
+```bash
+python -m pollen_data_acquisition_server.server
+```
+
+Then get into the teleoperation application and choose "Data acquisition session".
+You can finally setup your session by following the screens displayed.
+
+> Even without the VR app, you can use the `pollen_data_acquisition_server` with your own client implementation.
+
+### Option 2: Using lerobot.record
+
+Reachy 2 is fully supported by LeRobot’s recording features.
+If you choose this option but still want to use the VR teleoperation application, select "Standard session" in the app.
+
+**Example: start a recording without the mobile base:**
+First add reachy2 and reachy2_teleoperator to the imports of the record script. Then you can use the following command:
+
+```bash
+python -m lerobot.record \
+    --robot.type=reachy2 \
+    --robot.ip_address=192.168.0.200 \
+    --robot.id=r2-0000 \
+    --robot.use_external_commands=true \
+    --robot.with_mobile_base=false \
+    --teleop.type=reachy2_teleoperator \
+    --teleop.ip_address=192.168.0.200 \
+    --teleop.with_mobile_base=false \
+    --dataset.repo_id=pollen_robotics/record_test \
+    --dataset.single_task="Reachy 2 recording test" \
+    --dataset.num_episodes=1 \
+    --dataset.episode_time_s=5 \
+    --dataset.fps=15 \
+    --dataset.push_to_hub=true \
+    --dataset.private=true \
+    --display_data=true
+```
+
+#### Specific Options
+
+**Extended setup overview (all options included):**
+
+```bash
+python -m lerobot.record \
+    --robot.type=reachy2 \
+    --robot.ip_address=192.168.0.200 \
+    --robot.use_external_commands=true \
+    --robot.with_mobile_base=true \
+    --robot.with_l_arm=true \
+    --robot.with_r_arm=true \
+    --robot.with_neck=true \
+    --robot.with_antennas=true \
+    --robot.with_left_teleop_camera=true \
+    --robot.with_right_teleop_camera=true \
+    --robot.with_torso_camera=false \
+    --robot.disable_torque_on_disconnect=false \
+    --robot.max_relative_target=5.0 \
+    --teleop.type=reachy2_teleoperator \
+    --teleop.ip_address=192.168.0.200 \
+    --teleop.use_present_position=false \
+    --teleop.with_mobile_base=false \
+    --teleop.with_l_arm=true \
+    --teleop.with_r_arm=true \
+    --teleop.with_neck=true \
+    --teleop.with_antennas=true \
+    --dataset.repo_id=pollen_robotics/record_test \
+    --dataset.single_task="Reachy 2 recording test" \
+    --dataset.num_episodes=1 \
+    --dataset.episode_time_s=5 \
+    --dataset.fps=15 \
+    --dataset.push_to_hub=true \
+    --dataset.private=true \
+    --display_data=true
+```
+
+##### `--robot.use_external_commands`
+
+Determine whether LeRobot robot.send_action() sends commands to the robot.
+**Must** be set to false while using the VR teleoperation application, as the app already sends commands.
+
+##### `--teleop.use_present_position`
+
+Determine whether the teleoperator reads the goal or present position of the robot.
+Must be set to true if a compliant Reachy 2 is used to control another one.
+
+##### Use the relevant parts
+
+From our initial tests, recording **all** joints when only some are moving can reduce model quality with certain policies.
+To avoid this, you can exclude specific parts from recording and replay using:
+
+````
+--robot.with_<part>=false
+```,
+with `<part>` being one of : `mobile_base`, `l_arm`, `r_arm", `neck`, `antennas`.
+It determine whether the corresponding part is recorded in the observations. True if not set.
+
+By default, **all parts are recorded**.
+
+The same per-part mechanism is available in `reachy2_teleoperator` as well.
+
+````
+
+--teleop.with\_<part>
+
+```
+with `<part>` being one of : `mobile_base`, `l_arm`, `r_arm", `neck`, `antennas`.
+Determine whether the corresponding part is recorded in the actions. True if not set.
+
+> **Important:** In a given session, the **enabled parts must match** on both the robot and the teleoperator.
+For example, if the robot runs with `--robot.with_mobile_base=false`, the teleoperator must disable the same part `--teleoperator.with_mobile_base=false`.
+
+##### Use the relevant cameras
+
+You can do the same for **cameras**. By default, only the **teleoperation cameras** are recorded (both `left_teleop_camera` and `right_teleop_camera`). Enable or disable each camera with:
+
+```
+
+--robot.with_left_teleop_camera=<true|false>
+--robot.with_right_teleop_camera=<true|false>
+--robot.with_torso_camera=<true|false>
+
+````
+
+
+## Step 2: Replay
+
+Make sure the robot is configured with the same parts as the dataset:
+
+```bash
+python -m lerobot.replay \
+    --robot.type=reachy2 \
+    --robot.ip_address=192.168.0.200 \
+    --robot.use_external_commands=false \
+    --robot.with_mobile_base=false \
+    --dataset.repo_id=pollen_robotics/record_test \
+    --dataset.episode=0
+    --display_data=true
+````
+
+## Step 3: Train
+
+```bash
+python -m lerobot.scripts.train \
+  --dataset.repo_id=pollen_robotics/record_test \
+  --policy.type=act \
+  --output_dir=outputs/train/reachy2_test \
+  --job_name=reachy2 \
+  --policy.device=mps \
+  --wandb.enable=true \
+  --policy.repo_id=pollen_robotics/record_test_policy
+```
+
+## Step 4: Evaluate
+
+```bash
+python -m lerobot.record \
+  --robot.type=reachy2 \
+  --robot.ip_address=192.168.0.200 \
+  --display_data=false \
+  --dataset.repo_id=pollen_robotics/eval_record_test \
+  --dataset.single_task="Evaluate reachy2 policy" \
+  --dataset.num_episodes=10 \
+  --policy.path=outputs/train/reachy2_test/checkpoints/last/pretrained_model
+```

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 dataset.meta.episodes. Here, we access
+# episodes, you can access the frame indices of any episode using the episode_data_index. Here, we access
 # frame indices associated to the first episode:
 episode_index = 0
-from_idx = dataset.meta.episodes["dataset_from_index"][episode_index]
-to_idx = dataset.meta.episodes["dataset_to_index"][episode_index]
+from_idx = dataset.episode_data_index["from"][episode_index].item()
+to_idx = dataset.episode_data_index["to"][episode_index].item()
 
 # Then we grab all the image frames from the first camera:
 camera_key = dataset.meta.camera_keys[0]

examples/6_evaluate_libero.sh

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

examples/port_datasets/agibot_hdf5/port_agibot.py

@@ -1,503 +0,0 @@
-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

@@ -1,198 +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 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

@@ -1,85 +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 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

@@ -1,430 +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 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

@@ -1,148 +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 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

@@ -1,162 +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 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

@@ -1,281 +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 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

@@ -1,111 +0,0 @@
-#!/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()

push_pi0_to_hub.py

@@ -0,0 +1,193 @@
+#!/usr/bin/env python
+
+"""Script to create and push a PI0OpenPI model to HuggingFace hub with proper config format."""
+
+import tempfile
+from pathlib import Path
+
+import torch
+from huggingface_hub import HfApi, create_repo
+
+from lerobot.policies.pi0_openpi import PI0OpenPIConfig, PI0OpenPIPolicy
+
+
+def create_and_push_model(
+    repo_id: str,
+    private: bool = False,
+    token: str = None,
+):
+    """Create a PI0OpenPI model with proper config and push to HuggingFace hub.
+
+    Args:
+        repo_id: HuggingFace repository ID (e.g., "username/model-name")
+        private: Whether to create a private repository
+        token: HuggingFace API token (optional, will use cached token if not provided)
+    """
+    print("=" * 60)
+    print("PI0OpenPI Model Hub Upload")
+    print("=" * 60)
+
+    # Create configuration
+    print("\nCreating PI0OpenPI configuration...")
+    config = PI0OpenPIConfig(
+        # Model architecture
+        paligemma_variant="gemma_2b",
+        action_expert_variant="gemma_300m",
+        pi05=False,  # Use PI0 (not PI0.5)
+        dtype="float32",  # Use float32 for compatibility
+        # Input/output dimensions
+        action_dim=32,  # see openpi `Pi0Config`
+        state_dim=32,
+        chunk_size=50,
+        n_action_steps=50,
+        # Image inputs, see openpi `model.py, IMAGE_KEYS`
+        image_keys=(
+            "observation.images.base_0_rgb",
+            "observation.images.left_wrist_0_rgb",
+            "observation.images.right_wrist_0_rgb",
+        ),
+        # Training settings
+        gradient_checkpointing=False,
+        compile_model=False,
+        device=None,  # Auto-detect
+        # Tokenizer settings
+        tokenizer_max_length=48,  # see openpi `__post_init__`, use pi0=48 and pi05=200
+    )
+
+    print(f"  - Config type: {config.__class__.__name__}")
+    print(f"  - PaliGemma variant: {config.paligemma_variant}")
+    print(f"  - Action expert variant: {config.action_expert_variant}")
+    print(f"  - Action dim: {config.action_dim}")
+    print(f"  - State dim: {config.state_dim}")
+
+    # Create dummy dataset stats for normalization
+    print("\nCreating dataset statistics...")
+    dataset_stats = {
+        "observation.state": {
+            "mean": torch.zeros(config.state_dim),
+            "std": torch.ones(config.state_dim),
+            "min": torch.full((config.state_dim,), -5.0),
+            "max": torch.full((config.state_dim,), 5.0),
+        },
+        "action": {
+            "mean": torch.zeros(config.action_dim),
+            "std": torch.ones(config.action_dim),
+            "min": torch.full((config.action_dim,), -1.0),
+            "max": torch.full((config.action_dim,), 1.0),
+        },
+    }
+
+    # Add image stats
+    for key in config.image_keys:
+        dataset_stats[key] = {
+            "mean": torch.tensor([0.485, 0.456, 0.406]),  # TODO(pepijn): fix this, now its ImageNet mean
+            "std": torch.tensor([0.229, 0.224, 0.225]),  # TODO(pepijn): fix this, now its ImageNet std
+            "min": torch.tensor([0.0, 0.0, 0.0]),
+            "max": torch.tensor([1.0, 1.0, 1.0]),
+        }
+
+    # Create the policy
+    print("\nInitializing PI0OpenPI policy...")
+    print("  (This may take a moment as it loads the tokenizer and initializes the model)")
+    policy = PI0OpenPIPolicy(config, dataset_stats)
+
+    # Initialize with small random weights (optional - for testing)
+    # Note: In practice, you would load your trained weights here
+    print("\nInitializing model weights...")
+    for name, param in policy.named_parameters():
+        if "weight" in name:
+            if "norm" in name.lower() or "layernorm" in name.lower():
+                torch.nn.init.ones_(param)
+            elif len(param.shape) >= 2:
+                torch.nn.init.xavier_uniform_(param, gain=0.01)
+            else:
+                torch.nn.init.normal_(param, mean=0.0, std=0.01)
+        elif "bias" in name:
+            torch.nn.init.zeros_(param)
+
+    print(f"  - Total parameters: {sum(p.numel() for p in policy.parameters()):,}")
+    print(f"  - Trainable parameters: {sum(p.numel() for p in policy.parameters() if p.requires_grad):,}")
+
+    # Create temporary directory for saving
+    with tempfile.TemporaryDirectory() as tmpdir:
+        save_path = Path(tmpdir) / "model"
+        save_path.mkdir(exist_ok=True)
+
+        print(f"\nSaving model to temporary directory: {save_path}")
+
+        # Save the model using LeRobot's save_pretrained method
+        # This ensures the config is saved in the correct format
+        policy.save_pretrained(save_path)
+
+        # List saved files
+        saved_files = list(save_path.glob("*"))
+        print("\nSaved files:")
+        for file in saved_files:
+            size = file.stat().st_size
+            print(f"  - {file.name}: {size:,} bytes")
+
+        # Create or get repository
+        print(f"\nCreating/accessing repository: {repo_id}")
+        api = HfApi(token=token)
+
+        try:
+            # Create repo if it doesn't exist
+            create_repo(
+                repo_id,
+                private=private,
+                token=token,
+                exist_ok=True,
+            )
+            print(f"  ✓ Repository ready: https://huggingface.co/{repo_id}")
+        except Exception as e:
+            print(f"  ⚠️  Note: {e}")
+
+        # Upload to hub
+        print("\nUploading to HuggingFace hub...")
+        api.upload_folder(
+            folder_path=str(save_path),
+            repo_id=repo_id,
+            repo_type="model",
+            token=token,
+            commit_message="Upload PI0OpenPI model with proper LeRobot config format",
+        )
+
+        print(f"\n✓ Model successfully uploaded to: https://huggingface.co/{repo_id}")
+
+    print("\n" + "=" * 60)
+    print("✓ Process complete!")
+    print("=" * 60)
+
+    return policy
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Push PI0OpenPI model to HuggingFace hub")
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        default="test-user/pi0-openpi-test",
+        help="HuggingFace repository ID (e.g., 'username/model-name')",
+    )
+    parser.add_argument(
+        "--private",
+        action="store_true",
+        help="Create a private repository",
+    )
+    parser.add_argument(
+        "--token",
+        type=str,
+        default=None,
+        help="HuggingFace API token (optional, uses cached token if not provided)",
+    )
+
+    args = parser.parse_args()
+
+    # Run the upload
+    create_and_push_model(
+        repo_id=args.repo_id,
+        private=args.private,
+        token=args.token,
+    )

pyproject.toml

@@ -29,7 +29,7 @@ version = "0.3.4"
 description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch"
 readme = "README.md"
 license = { text = "Apache-2.0" }
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 authors = [
     { name = "Rémi Cadène", email = "re.cadene@gmail.com" },
     { name = "Simon Alibert", email = "alibert.sim@gmail.com" },
@@ -50,7 +50,7 @@ classifiers = [
     "Intended Audience :: Education",
     "Intended Audience :: Science/Research",
     "License :: OSI Approved :: Apache Software License",
-    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
     "Topic :: Software Development :: Build Tools",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
@@ -95,7 +95,7 @@ dependencies = [
 # Common
 pygame-dep = ["pygame>=2.5.1"]
 placo-dep = ["placo>=0.9.6"]
-transformers-dep = ["transformers>=4.50.3,<4.52.0"] # TODO: Bumb dependency
+transformers-dep = ["transformers==4.53.2"]
 grpcio-dep = ["grpcio==1.73.1", "protobuf==6.31.0"]
 
 # Motors
@@ -106,6 +106,7 @@ dynamixel = ["dynamixel-sdk>=3.7.31"]
 gamepad = ["lerobot[pygame-dep]", "hidapi>=0.14.0"]
 hopejr = ["lerobot[feetech]", "lerobot[pygame-dep]"]
 lekiwi = ["lerobot[feetech]", "pyzmq>=26.2.1"]
+reachy2 = ["reachy2_sdk>=1.0.14"]
 kinematics = ["lerobot[placo-dep]"]
 intelrealsense = [
     "pyrealsense2>=2.55.1.6486 ; sys_platform != 'darwin'",
@@ -134,13 +135,33 @@ video_benchmark = ["scikit-image>=0.23.2", "pandas>=2.2.2"]
 aloha = ["gym-aloha>=0.1.1"]
 pusht = ["gym-pusht>=0.1.5", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead
 xarm = ["gym-xarm>=0.1.1"]
-
+libero = [
+    "hydra-core>=1.2,<1.4",
+    "numpy",
+    "wandb",
+    "easydict",
+    "transformers",
+    "opencv-python",
+    "robomimic==0.2.0",
+    "einops",
+    "thop",
+    "robosuite==1.4.0",
+    "mujoco>=2.3.7,<3.0.0",
+    "bddl==1.0.1",
+    "matplotlib",
+    "cloudpickle",
+    "future",
+    "gym",
+    "egl_probe @ git+https://github.com/jadechoghari/egl_probe.git#egg=egl_probe",
+    "libero @ git+https://github.com/jadechoghari/LIBERO.git@main#egg=libero",
+]
 # All
 all = [
     "lerobot[dynamixel]",
     "lerobot[gamepad]",
     "lerobot[hopejr]",
     "lerobot[lekiwi]",
+    "lerobot[reachy2]",
     "lerobot[kinematics]",
     "lerobot[intelrealsense]",
     "lerobot[pi0]",
@@ -152,7 +173,8 @@ all = [
     "lerobot[video_benchmark]",
     "lerobot[aloha]",
     "lerobot[pusht]",
-    "lerobot[xarm]"
+    "lerobot[xarm]",
+    "lerobot[libero]"
 ]
 
 [project.scripts]
@@ -258,7 +280,7 @@ default.extend-ignore-identifiers-re = [
 # paths = ["src/lerobot"]
 
 # [tool.mypy]
-# python_version = "3.10"
+# python_version = "3.11"
 # warn_return_any = true
 # warn_unused_configs = true
 # ignore_missing_imports = false

src/lerobot/cameras/reachy2_camera/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+
+from .configuration_reachy2_camera import Reachy2CameraConfig
+from .reachy2_camera import Reachy2Camera

src/lerobot/cameras/reachy2_camera/configuration_reachy2_camera.py

@@ -0,0 +1,78 @@
+# 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.
+
+from dataclasses import dataclass
+
+from ..configs import CameraConfig, ColorMode
+
+
+@CameraConfig.register_subclass("reachy2_camera")
+@dataclass
+class Reachy2CameraConfig(CameraConfig):
+    """Configuration class for Reachy 2 camera devices.
+
+    This class provides configuration options for Reachy 2 cameras,
+    supporting both the teleop and depth cameras. It includes settings
+    for resolution, frame rate, color mode, and the selection of the cameras.
+
+    Example configurations:
+    ```python
+    # Basic configurations
+    Reachy2CameraConfig(
+        name="teleop",
+        image_type="left",
+        ip_address="192.168.0.200",  # IP address of the robot
+        fps=15,
+        width=640,
+        height=480,
+        color_mode=ColorMode.RGB,
+    )  # Left teleop camera, 640x480 @ 15FPS
+    ```
+
+    Attributes:
+        name: Name of the camera device. Can be "teleop" or "depth".
+        image_type: Type of image stream. For "teleop" camera, can be "left" or "right".
+                    For "depth" camera, can be "rgb" or "depth". (depth is not supported yet)
+        fps: Requested frames per second for the color stream.
+        width: Requested frame width in pixels for the color stream.
+        height: Requested frame height in pixels for the color stream.
+        color_mode: Color mode for image output (RGB or BGR). Defaults to RGB.
+        ip_address: IP address of the robot. Defaults to "localhost".
+        port: Port number for the camera server. Defaults to 50065.
+
+    Note:
+        - Only 3-channel color output (RGB/BGR) is currently supported.
+    """
+
+    name: str
+    image_type: str
+    color_mode: ColorMode = ColorMode.RGB
+    ip_address: str | None = "localhost"
+    port: int = 50065
+    # use_depth: bool = False
+
+    def __post_init__(self):
+        if self.name not in ["teleop", "depth"]:
+            raise ValueError(f"`name` is expected to be 'teleop' or 'depth', but {self.name} is provided.")
+        if (self.name == "teleop" and self.image_type not in ["left", "right"]) or (
+            self.name == "depth" and self.image_type not in ["rgb", "depth"]
+        ):
+            raise ValueError(
+                f"`image_type` is expected to be 'left' or 'right' for teleop camera, and 'rgb' or 'depth' for depth camera, but {self.image_type} is provided."
+            )
+
+        if self.color_mode not in ["rgb", "bgr"]:
+            raise ValueError(
+                f"`color_mode` is expected to be 'rgb' or 'bgr', but {self.color_mode} is provided."
+            )

src/lerobot/cameras/reachy2_camera/reachy2_camera.py

@@ -0,0 +1,288 @@
+# 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.
+
+"""
+Provides the Reachy2Camera class for capturing frames from Reachy 2 cameras using Reachy 2's CameraManager.
+"""
+
+import logging
+import os
+import platform
+import time
+from threading import Event, Lock, Thread
+from typing import Any
+
+# Fix MSMF hardware transform compatibility for Windows before importing cv2
+if platform.system() == "Windows" and "OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS" not in os.environ:
+    os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
+import cv2
+import numpy as np
+from reachy2_sdk.media.camera import CameraView
+from reachy2_sdk.media.camera_manager import CameraManager
+
+from lerobot.errors import DeviceNotConnectedError
+
+from ..camera import Camera
+from .configuration_reachy2_camera import ColorMode, Reachy2CameraConfig
+
+logger = logging.getLogger(__name__)
+
+
+class Reachy2Camera(Camera):
+    """
+    Manages Reachy 2 camera using Reachy 2 CameraManager.
+
+    This class provides a high-level interface to connect to, configure, and read
+    frames from Reachy 2 cameras. It supports both synchronous and asynchronous
+    frame reading.
+
+    An Reachy2Camera instance requires a camera name (e.g., "teleop") and an image
+    type (e.g., "left") to be specified in the configuration.
+
+    The camera's default settings (FPS, resolution, color mode) are used unless
+    overridden in the configuration.
+    """
+
+    def __init__(self, config: Reachy2CameraConfig):
+        """
+        Initializes the Reachy2Camera instance.
+
+        Args:
+            config: The configuration settings for the camera.
+        """
+        super().__init__(config)
+
+        self.config = config
+
+        self.fps = config.fps
+        self.color_mode = config.color_mode
+
+        self.cam_manager: CameraManager | None = None
+
+        self.thread: Thread | None = None
+        self.stop_event: Event | None = None
+        self.frame_lock: Lock = Lock()
+        self.latest_frame: np.ndarray | None = None
+        self.new_frame_event: Event = Event()
+
+    def __str__(self) -> str:
+        return f"{self.__class__.__name__}({self.config.name}, {self.config.image_type})"
+
+    @property
+    def is_connected(self) -> bool:
+        """Checks if the camera is currently connected and opened."""
+        if self.config.name == "teleop":
+            return self.cam_manager._grpc_connected and self.cam_manager.teleop if self.cam_manager else False
+        elif self.config.name == "depth":
+            return self.cam_manager._grpc_connected and self.cam_manager.depth if self.cam_manager else False
+        else:
+            raise ValueError(f"Invalid camera name '{self.config.name}'. Expected 'teleop' or 'depth'.")
+
+    def connect(self, warmup: bool = True):
+        """
+        Connects to the Reachy2 CameraManager as specified in the configuration.
+        """
+        self.cam_manager = CameraManager(host=self.config.ip_address, port=self.config.port)
+        self.cam_manager.initialize_cameras()
+
+        logger.info(f"{self} connected.")
+
+    @staticmethod
+    def find_cameras(ip_address: str = "localhost", port: int = 50065) -> list[dict[str, Any]]:
+        """
+        Detects available Reachy 2 cameras.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries,
+            where each dictionary contains 'name', 'stereo',
+            and the default profile properties (width, height, fps).
+        """
+        initialized_cameras = []
+        camera_manager = CameraManager(host=ip_address, port=port)
+
+        for camera in [camera_manager.teleop, camera_manager.depth]:
+            if camera is None:
+                continue
+
+            height, width, _, _, _, _, _ = camera.get_parameters()
+
+            camera_info = {
+                "name": camera._cam_info.name,
+                "stereo": camera._cam_info.stereo,
+                "default_profile": {
+                    "width": width,
+                    "height": height,
+                    "fps": 30,
+                },
+            }
+            initialized_cameras.append(camera_info)
+
+        camera_manager.disconnect()
+        return initialized_cameras
+
+    def read(self, color_mode: ColorMode | None = None) -> np.ndarray:
+        """
+        Reads a single frame synchronously from the camera.
+
+        This is a blocking call.
+
+        Args:
+            color_mode (Optional[ColorMode]): If specified, overrides the default
+                color mode (`self.color_mode`) for this read operation (e.g.,
+                request RGB even if default is BGR).
+
+        Returns:
+            np.ndarray: The captured frame as a NumPy array in the format
+                       (height, width, channels), using the specified or default
+                       color mode and applying any configured rotation.
+        """
+        if not self.is_connected:
+            raise DeviceNotConnectedError(f"{self} is not connected.")
+
+        start_time = time.perf_counter()
+
+        frame = None
+
+        if self.cam_manager is None:
+            raise DeviceNotConnectedError(f"{self} is not connected.")
+        else:
+            if self.config.name == "teleop" and hasattr(self.cam_manager, "teleop"):
+                if self.config.image_type == "left":
+                    frame = self.cam_manager.teleop.get_frame(CameraView.LEFT, size=(640, 480))[0]
+                elif self.config.image_type == "right":
+                    frame = self.cam_manager.teleop.get_frame(CameraView.RIGHT, size=(640, 480))[0]
+            elif self.config.name == "depth" and hasattr(self.cam_manager, "depth"):
+                if self.config.image_type == "depth":
+                    frame = self.cam_manager.depth.get_depth_frame()[0]
+                elif self.config.image_type == "rgb":
+                    frame = self.cam_manager.depth.get_frame(size=(640, 480))[0]
+
+            if frame is None:
+                return np.empty((0, 0, 3), dtype=np.uint8)
+
+            if self.config.color_mode == "rgb":
+                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+        read_duration_ms = (time.perf_counter() - start_time) * 1e3
+        logger.debug(f"{self} read took: {read_duration_ms:.1f}ms")
+
+        return frame
+
+    def _read_loop(self):
+        """
+        Internal loop run by the background thread for asynchronous reading.
+
+        On each iteration:
+        1. Reads a color frame
+        2. Stores result in latest_frame (thread-safe)
+        3. Sets new_frame_event to notify listeners
+
+        Stops on DeviceNotConnectedError, logs other errors and continues.
+        """
+        while not self.stop_event.is_set():
+            try:
+                color_image = self.read()
+
+                with self.frame_lock:
+                    self.latest_frame = color_image
+                self.new_frame_event.set()
+
+            except DeviceNotConnectedError:
+                break
+            except Exception as e:
+                logger.warning(f"Error reading frame in background thread for {self}: {e}")
+
+    def _start_read_thread(self) -> None:
+        """Starts or restarts the background read thread if it's not running."""
+        if self.thread is not None and self.thread.is_alive():
+            self.thread.join(timeout=0.1)
+        if self.stop_event is not None:
+            self.stop_event.set()
+
+        self.stop_event = Event()
+        self.thread = Thread(target=self._read_loop, args=(), name=f"{self}_read_loop")
+        self.thread.daemon = True
+        self.thread.start()
+
+    def _stop_read_thread(self) -> None:
+        """Signals the background read thread to stop and waits for it to join."""
+        if self.stop_event is not None:
+            self.stop_event.set()
+
+        if self.thread is not None and self.thread.is_alive():
+            self.thread.join(timeout=2.0)
+
+        self.thread = None
+        self.stop_event = None
+
+    def async_read(self, timeout_ms: float = 200) -> np.ndarray:
+        """
+        Reads the latest available frame asynchronously.
+
+        This method retrieves the most recent frame captured by the background
+        read thread. It does not block waiting for the camera hardware directly,
+        but may wait up to timeout_ms for the background thread to provide a frame.
+
+        Args:
+            timeout_ms (float): Maximum time in milliseconds to wait for a frame
+                to become available. Defaults to 200ms (0.2 seconds).
+
+        Returns:
+            np.ndarray: The latest captured frame as a NumPy array in the format
+                       (height, width, channels), processed according to configuration.
+
+        Raises:
+            DeviceNotConnectedError: If the camera is not connected.
+            TimeoutError: If no frame becomes available within the specified timeout.
+            RuntimeError: If an unexpected error occurs.
+        """
+        if not self.is_connected:
+            raise DeviceNotConnectedError(f"{self} is not connected.")
+
+        if self.thread is None or not self.thread.is_alive():
+            self._start_read_thread()
+
+        if not self.new_frame_event.wait(timeout=timeout_ms / 1000.0):
+            thread_alive = self.thread is not None and self.thread.is_alive()
+            raise TimeoutError(
+                f"Timed out waiting for frame from camera {self} after {timeout_ms} ms. "
+                f"Read thread alive: {thread_alive}."
+            )
+
+        with self.frame_lock:
+            frame = self.latest_frame
+            self.new_frame_event.clear()
+
+        if frame is None:
+            raise RuntimeError(f"Internal error: Event set but no frame available for {self}.")
+
+        return frame
+
+    def disconnect(self):
+        """
+        Stops the background read thread (if running).
+
+        Raises:
+            DeviceNotConnectedError: If the camera is already disconnected.
+        """
+        if not self.is_connected and self.thread is None:
+            raise DeviceNotConnectedError(f"{self} not connected.")
+
+        if self.thread is not None:
+            self._stop_read_thread()
+
+        if self.cam_manager is not None:
+            self.cam_manager.disconnect()
+
+        logger.info(f"{self} disconnected.")

src/lerobot/cameras/utils.py

@@ -37,8 +37,14 @@ def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> dict[s
             from .realsense.camera_realsense import RealSenseCamera
 
             cameras[key] = RealSenseCamera(cfg)
+
+        elif cfg.type == "reachy2_camera":
+            from .reachy2_camera.reachy2_camera import Reachy2Camera
+
+            cameras[key] = Reachy2Camera(cfg)
+
         else:
-            raise ValueError(f"The motor type '{cfg.type}' is not valid.")
+            raise ValueError(f"The camera type '{cfg.type}' is not valid.")
 
     return cameras
 

src/lerobot/configs/policies.py

@@ -72,9 +72,11 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
     tags: list[str] | None = None
     # Add tags to your policy on the hub.
     license: str | None = None
+    # Either the repo ID of a model hosted on the Hub or a path to a directory containing weights
+    # saved using `Policy.save_pretrained`. If not provided, the policy is initialized from scratch.
+    pretrained_path: str | None = None
 
     def __post_init__(self):
-        self.pretrained_path = None
         if not self.device or not is_torch_device_available(self.device):
             auto_device = auto_select_torch_device()
             logging.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")

src/lerobot/datasets/aggregate.py

@@ -1,505 +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 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,18 +47,6 @@ 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.
@@ -70,14 +58,7 @@ class CompatibilityError(Exception): ...
 
 class BackwardCompatibilityError(CompatibilityError):
     def __init__(self, repo_id: str, version: packaging.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)."
-            )
+        message = V2_MESSAGE.format(repo_id=repo_id, version=version)
         super().__init__(message)
 
 

src/lerobot/datasets/dataset_tools.py

@@ -1,761 +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.
-"""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/lerobot_dataset.py

@@ -16,18 +16,16 @@
 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 Dataset
+from datasets import concatenate_datasets, load_dataset
 from huggingface_hub import HfApi, snapshot_download
 from huggingface_hub.constants import REPOCARD_NAME
 from huggingface_hub.errors import RevisionNotFoundError
@@ -36,51 +34,46 @@ 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_hf_dataset_size_in_mb,
+    get_episode_data_index,
     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 = "v3.0"
+CODEBASE_VERSION = "v2.1"
 
 
 class LeRobotDatasetMetadata:
@@ -104,18 +97,20 @@ 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 = load_tasks(self.root)
+        self.tasks, self.task_to_task_index = load_tasks(self.root)
         self.episodes = load_episodes(self.root)
-        self.stats = load_stats(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()))
 
     def pull_from_repo(
         self,
@@ -137,19 +132,18 @@ class LeRobotDatasetMetadata:
         return packaging.version.parse(self.info["codebase_version"])
 
     def get_data_file_path(self, ep_index: int) -> Path:
-        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)
+        ep_chunk = self.get_episode_chunk(ep_index)
+        fpath = self.data_path.format(episode_chunk=ep_chunk, episode_index=ep_index)
         return Path(fpath)
 
     def get_video_file_path(self, ep_index: int, vid_key: str) -> Path:
-        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)
+        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)
         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."""
@@ -215,109 +209,40 @@ class LeRobotDatasetMetadata:
         """Total number of different tasks performed in this dataset."""
         return self.info["total_tasks"]
 
+    @property
+    def total_chunks(self) -> int:
+        """Total number of chunks (groups of episodes)."""
+        return self.info["total_chunks"]
+
     @property
     def chunks_size(self) -> int:
-        """Max number of files per chunk."""
+        """Max number of episodes per chunk."""
         return self.info["chunks_size"]
 
-    @property
-    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.
         """
-        if task in self.tasks.index:
-            return int(self.tasks.loc[task].task_index)
-        else:
-            return None
+        return self.task_to_task_index.get(task, None)
 
-    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.
+    def add_task(self, task: str):
         """
-        # 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]
+        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.")
 
-        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_index = self.info["total_tasks"]
+        self.task_to_task_index[task] = task_index
+        self.tasks[task_index] = task
+        self.info["total_tasks"] += 1
 
-            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)
+        task_dict = {
+            "task_index": task_index,
+            "task": task,
+        }
+        append_jsonlines(task_dict, self.root / TASKS_PATH)
 
     def save_episode(
         self,
@@ -325,28 +250,30 @@ 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,
         }
-        episode_dict.update(episode_metadata)
-        episode_dict.update(flatten_dict({"stats": episode_stats}))
-        self._save_episode_metadata(episode_dict)
+        self.episodes[episode_index] = episode_dict
+        write_episode(episode_dict, 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)
+        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)
 
     def update_video_info(self) -> None:
         """
@@ -386,12 +313,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 = None
-        obj.episodes = None
-        obj.stats = None
+        obj.tasks, obj.task_to_task_index = {}, {}
+        obj.episodes_stats, obj.stats, obj.episodes = {}, {}, {}
         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()
@@ -413,6 +340,7 @@ 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:
@@ -426,9 +354,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 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.
+              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.
 
 
         2. Your dataset doesn't already exists (either on local disk or on the Hub): you can create an empty
@@ -449,47 +377,38 @@ class LeRobotDataset(torch.utils.data.Dataset):
         .
         ├── data
         │   ├── chunk-000
-        │   │   ├── file-000.parquet
-        │   │   ├── file-001.parquet
+        │   │   ├── episode_000000.parquet
+        │   │   ├── episode_000001.parquet
+        │   │   ├── episode_000002.parquet
         │   │   └── ...
         │   ├── chunk-001
-        │   │   ├── file-000.parquet
-        │   │   ├── file-001.parquet
+        │   │   ├── episode_001000.parquet
+        │   │   ├── episode_001001.parquet
+        │   │   ├── episode_001002.parquet
         │   │   └── ...
         │   └── ...
         ├── meta
-        │   ├── episodes
-        │   │   ├── chunk-000
-        │   │   │   ├── file-000.parquet
-        │   │   │   ├── file-001.parquet
-        │   │   │   └── ...
-        │   │   ├── chunk-001
-        │   │   │   └── ...
-        │   │   └── ...
+        │   ├── episodes.jsonl
         │   ├── info.json
         │   ├── stats.json
-        │   └── tasks.parquet
+        │   └── tasks.jsonl
         └── videos
-            ├── observation.images.laptop
-            │   ├── chunk-000
-            │   │   ├── file-000.mp4
-            │   │   ├── file-001.mp4
+            ├── chunk-000
+            │   ├── observation.images.laptop
+            │   │   ├── episode_000000.mp4
+            │   │   ├── episode_000001.mp4
+            │   │   ├── episode_000002.mp4
             │   │   └── ...
-            │   ├── chunk-001
+            │   ├── observation.images.phone
+            │   │   ├── episode_000000.mp4
+            │   │   ├── episode_000001.mp4
+            │   │   ├── episode_000002.mp4
             │   │   └── ...
-            │   └── ...
-            ├── observation.images.phone
-            │   ├── chunk-000
-            │   │   ├── file-000.mp4
-            │   │   ├── file-001.mp4
-            │   │   └── ...
-            │   ├── chunk-001
-            │   │   └── ...
-            │   └── ...
+            ├── chunk-001
             └── ...
 
-        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
+        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
         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
@@ -523,6 +442,8 @@ 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
@@ -534,6 +455,8 @@ 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
@@ -545,20 +468,29 @@ 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(download_videos)
+            self.download_episodes(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)
@@ -634,7 +566,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
             ignore_patterns=ignore_patterns,
         )
 
-    def download(self, download_videos: bool = True) -> None:
+    def download_episodes(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
@@ -642,10 +574,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
         """
         # TODO(rcadene, aliberts): implement faster transfer
         # https://huggingface.co/docs/huggingface_hub/en/guides/download#faster-downloads
-        ignore_patterns = None if download_videos else "videos/"
         files = None
+        ignore_patterns = None if download_videos else "videos/"
         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]:
@@ -658,40 +591,28 @@ 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."""
-        features = get_hf_features_from_features(self.features)
-        hf_dataset = load_nested_dataset(self.root / "data", features=features)
+        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")
         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
 
@@ -723,16 +644,15 @@ 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 = self.meta.episodes[ep_idx]
-        ep_start = ep["dataset_from_index"]
-        ep_end = ep["dataset_to_index"]
+        ep_start = self.episode_data_index["from"][ep_idx]
+        ep_end = self.episode_data_index["to"][ep_idx]
         query_indices = {
-            key: [max(ep_start, min(ep_end - 1, idx + delta)) for delta in delta_idx]
+            key: [max(ep_start.item(), min(ep_end.item() - 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) | (idx + delta >= ep_end) for delta in delta_idx]
+                [(idx + delta < ep_start.item()) | (idx + delta >= ep_end.item()) for delta in delta_idx]
             )
             for key, delta_idx in self.delta_indices.items()
         }
@@ -746,7 +666,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[query_indices[key]]["timestamp"]
+                timestamps = self.hf_dataset.select(query_indices[key])["timestamp"]
                 query_timestamps[key] = torch.stack(timestamps).tolist()
             else:
                 query_timestamps[key] = [current_ts]
@@ -755,7 +675,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[q_idx][key])
+            key: torch.stack(self.hf_dataset.select(q_idx)[key])
             for key, q_idx in query_indices.items()
             if key not in self.meta.video_keys
         }
@@ -766,17 +686,10 @@ 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, shifted_query_ts, self.tolerance_s, self.video_backend)
+            frames = decode_video_frames(video_path, query_ts, self.tolerance_s, self.video_backend)
             item[vid_key] = frames.squeeze(0)
 
         return item
@@ -814,7 +727,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
         # Add task as a string
         task_idx = item["task_index"].item()
-        item["task"] = self.meta.tasks.iloc[task_idx].name
+        item["task"] = self.meta.tasks[task_idx]
+
         return item
 
     def __repr__(self):
@@ -844,9 +758,6 @@ 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):
@@ -855,7 +766,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         else:
             self.image_writer.save_image(image=image, fpath=fpath)
 
-    def add_frame(self, frame: dict) -> None:
+    def add_frame(self, frame: dict, task: str, timestamp: float | None = None) -> 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
@@ -873,10 +784,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
         # Automatically add frame_index and timestamp to episode buffer
         frame_index = self.episode_buffer["size"]
-        timestamp = frame.pop("timestamp") if "timestamp" in frame else frame_index / self.fps
+        if timestamp is None:
+            timestamp = frame_index / self.fps
         self.episode_buffer["frame_index"].append(frame_index)
         self.episode_buffer["timestamp"].append(timestamp)
-        self.episode_buffer["task"].append(frame.pop("task"))  # Remove task from frame after processing
+        self.episode_buffer["task"].append(task)
 
         # Add frame features to episode_buffer
         for key in frame:
@@ -902,12 +814,19 @@ 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.
         """
-        episode_buffer = episode_data if episode_data is not None else self.episode_buffer
+        if not episode_data:
+            episode_buffer = self.episode_buffer
+        else:
+            episode_buffer = episode_data
 
         validate_episode_buffer(episode_buffer, self.meta.total_episodes, self.features)
 
@@ -920,8 +839,11 @@ 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)
 
-        # Update tasks and task indices with new tasks if any
-        self.meta.save_episode_tasks(episode_tasks)
+        # 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)
 
         # 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])
@@ -933,142 +855,72 @@ 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)
 
-        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))
+        has_video_keys = len(self.meta.video_keys) > 0
+        use_batched_encoding = self.batch_encoding_size > 1
 
-        # `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)
+        if has_video_keys and not use_batched_encoding:
+            self.encode_episode_videos(episode_index)
 
-        if not episode_data:
-            # Reset episode buffer and clean up temporary images
-            self.clear_episode_buffer()
+        # `meta.save_episode` should be executed after encoding the videos
+        self.meta.save_episode(episode_index, episode_length, episode_tasks, ep_stats)
 
-    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.
-
-        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.
-
-        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)
-        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
+        # 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}"
                 )
-                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)
+                self.batch_encode_videos(start_ep, end_ep)
+                self.episodes_since_last_encoding = 0
 
-        # Remove temporary directory
-        shutil.rmtree(str(ep_path.parent))
+        # 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,
+        )
 
-        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
+        # 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")
+        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)
 
     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.root / "images" / cam_key
+                img_dir = self._get_image_file_path(
+                    episode_index=episode_index, image_key=cam_key, frame_index=0
+                ).parent
                 if img_dir.is_dir():
                     shutil.rmtree(img_dir)
 
@@ -1089,7 +941,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 pickleable and parallelized.
+        remove the image_writer in order for the LeRobotDataset object to be picklable and parallelized.
         """
         if self.image_writer is not None:
             self.image_writer.stop()
@@ -1100,16 +952,55 @@ class LeRobotDataset(torch.utils.data.Dataset):
         if self.image_writer is not None:
             self.image_writer.wait_until_done()
 
-    def _encode_temporary_episode_video(self, video_key: str, episode_index: int) -> dict:
+    def encode_episode_videos(self, episode_index: int) -> None:
         """
         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.
         """
-        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
+        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")
 
     @classmethod
     def create(
@@ -1124,6 +1015,7 @@ 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)
@@ -1140,6 +1032,8 @@ 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)
@@ -1152,6 +1046,7 @@ 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,11 +337,13 @@ def compute_sampler_weights(
     if len(offline_dataset) > 0:
         offline_data_mask_indices = []
         for start_index, end_index in zip(
-            offline_dataset.meta.episodes["dataset_from_index"],
-            offline_dataset.meta.episodes["dataset_to_index"],
+            offline_dataset.episode_data_index["from"],
+            offline_dataset.episode_data_index["to"],
             strict=True,
         ):
-            offline_data_mask_indices.extend(range(start_index, end_index - offline_drop_n_last_frames))
+            offline_data_mask_indices.extend(
+                range(start_index.item(), end_index.item() - 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/sampler.py

@@ -21,8 +21,7 @@ import torch
 class EpisodeAwareSampler:
     def __init__(
         self,
-        dataset_from_indices: list[int],
-        dataset_to_indices: list[int],
+        episode_data_index: dict,
         episode_indices_to_use: list | None = None,
         drop_n_first_frames: int = 0,
         drop_n_last_frames: int = 0,
@@ -31,8 +30,7 @@ class EpisodeAwareSampler:
         """Sampler that optionally incorporates episode boundary information.
 
         Args:
-            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_data_index: Dictionary with keys 'from' and 'to' containing the start and end indices of each episode.
             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.
@@ -41,10 +39,12 @@ class EpisodeAwareSampler:
         """
         indices = []
         for episode_idx, (start_index, end_index) in enumerate(
-            zip(dataset_from_indices, dataset_to_indices, strict=True)
+            zip(episode_data_index["from"], episode_data_index["to"], strict=True)
         ):
             if episode_indices_to_use is None or episode_idx in episode_indices_to_use:
-                indices.extend(range(start_index + drop_n_first_frames, end_index - drop_n_last_frames))
+                indices.extend(
+                    range(start_index.item() + drop_n_first_frames, end_index.item() - drop_n_last_frames)
+                )
 
         self.indices = indices
         self.shuffle = shuffle

src/lerobot/datasets/utils.py

@@ -17,59 +17,43 @@ 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 FeatureType, PolicyFeature
+from lerobot.configs.types import DictLike, 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 files per chunk
-DEFAULT_DATA_FILE_SIZE_IN_MB = 100  # Max size per file
-DEFAULT_VIDEO_FILE_SIZE_IN_MB = 500  # Max size per file
+DEFAULT_CHUNK_SIZE = 1000  # Max number of episodes per chunk
 
 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"
 
-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"
+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"
 
 DATASET_CARD_TEMPLATE = """
 ---
@@ -90,79 +74,6 @@ 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.
 
@@ -171,7 +82,6 @@ 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():
@@ -196,13 +106,23 @@ 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)):
@@ -232,7 +152,24 @@ def write_json(data: dict, fpath: Path) -> None:
         json.dump(data, f, indent=4, ensure_ascii=False)
 
 
-def write_info(info: dict, local_dir: Path) -> None:
+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):
     write_json(info, local_dir / INFO_PATH)
 
 
@@ -243,55 +180,65 @@ def load_info(local_dir: Path) -> dict:
     return info
 
 
-def write_stats(stats: dict, local_dir: Path) -> None:
+def write_stats(stats: dict, local_dir: Path):
     serialized_stats = serialize_dict(stats)
     write_json(serialized_stats, local_dir / STATS_PATH)
 
 
-def cast_stats_to_numpy(stats: dict) -> dict[str, dict[str, np.ndarray]]:
+def cast_stats_to_numpy(stats) -> 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]] | None:
+def load_stats(local_dir: Path) -> dict[str, dict[str, np.ndarray]]:
     if not (local_dir / STATS_PATH).exists():
         return None
     stats = load_json(local_dir / STATS_PATH)
     return cast_stats_to_numpy(stats)
 
 
-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 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 load_tasks(local_dir: Path) -> pandas.DataFrame:
-    tasks = pd.read_parquet(local_dir / DEFAULT_TASKS_PATH)
-    return tasks
+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 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 write_episode(episode: dict, local_dir: Path):
+    append_jsonlines(episode, local_dir / EPISODES_PATH)
 
 
-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 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 backward_compatible_episodes_stats(
     stats: dict[str, dict[str, np.ndarray]], episodes: list[int]
-) -> dict[int, dict[str, dict[str, np.ndarray]]]:
+) -> dict[str, dict[str, np.ndarray]]:
     return dict.fromkeys(episodes, stats)
 
 
@@ -307,7 +254,7 @@ def load_image_as_numpy(
     return img_array
 
 
-def hf_transform_to_torch(items_dict: dict[str, list[Any]]) -> dict[str, list[torch.Tensor | str]]:
+def hf_transform_to_torch(items_dict: dict[torch.Tensor | None]):
     """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
@@ -492,17 +439,6 @@ 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 = {}
@@ -547,17 +483,104 @@ 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_DATA_PATH,
+        "data_path": DEFAULT_PARQUET_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:
@@ -596,7 +619,7 @@ def get_delta_indices(delta_timestamps: dict[str, list[float]], fps: int) -> dic
     return delta_indices
 
 
-def cycle(iterable: Any) -> Iterator[Any]:
+def cycle(iterable):
     """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.
@@ -609,7 +632,7 @@ def cycle(iterable: Any) -> Iterator[Any]:
             iterator = iter(iterable)
 
 
-def create_branch(repo_id: str, *, branch: str, repo_type: str | None = None) -> None:
+def create_branch(repo_id, *, 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.
     """
@@ -630,7 +653,7 @@ def create_lerobot_dataset_card(
     **kwargs,
 ) -> DatasetCard:
     """
-    Keyword arguments will be used to replace values in ./lerobot/datasets/card_template.md.
+    Keyword arguments will be used to replace values in src/lerobot/datasets/card_template.md.
     Note: If specified, license must be one of https://huggingface.co/docs/hub/repositories-licenses.
     """
     card_tags = ["LeRobot"]
@@ -717,28 +740,21 @@ class IterableNamespace(SimpleNamespace):
         return vars(self).keys()
 
 
-def validate_frame(frame: dict, features: dict) -> None:
+def validate_frame(frame: dict, features: dict):
     expected_features = set(features) - set(DEFAULT_FEATURES)
     actual_features = set(frame)
 
-    # 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")
+    error_message = validate_features_presence(actual_features, expected_features)
 
-    # 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:
+    common_features = actual_features & expected_features
+    for name in common_features - {"task"}:
         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]) -> str:
+def validate_features_presence(actual_features: set[str], expected_features: set[str]):
     error_message = ""
     missing_features = expected_features - actual_features
     extra_features = actual_features - expected_features
@@ -753,9 +769,7 @@ 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
-) -> str:
+def validate_feature_dtype_and_shape(name: str, feature: dict, value: np.ndarray | PILImage.Image | str):
     expected_dtype = feature["dtype"]
     expected_shape = feature["shape"]
     if is_valid_numpy_dtype_string(expected_dtype):
@@ -770,7 +784,7 @@ def validate_feature_dtype_and_shape(
 
 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
@@ -787,9 +801,7 @@ 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
-) -> str:
+def validate_feature_image_or_video(name: str, expected_shape: list[str], value: np.ndarray | PILImage.Image):
     # 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):
@@ -805,13 +817,13 @@ def validate_feature_image_or_video(
     return error_message
 
 
-def validate_feature_string(name: str, value: str) -> str:
+def validate_feature_string(name: str, value: 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) -> None:
+def validate_episode_buffer(episode_buffer: dict, total_episodes: int, features: dict):
     if "size" not in episode_buffer:
         raise ValueError("size key not found in episode_buffer")
 
@@ -835,11 +847,3 @@ 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(LEGACY_DEFAULT_PARQUET_PATH.split("/")[:-1]).format(episode_chunk=ep_chunk)
+        chunk_dir = "/".join(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 / LEGACY_DEFAULT_PARQUET_PATH.format(
+            output_file = output_dir / 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(LEGACY_DEFAULT_VIDEO_PATH.split("/")[:-1]).format(
+            chunk_dir = "/".join(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 = LEGACY_DEFAULT_VIDEO_PATH.format(
+                target_path = 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 = [
-        LEGACY_DEFAULT_VIDEO_PATH.format(episode_chunk=0, video_key=vid_key, episode_index=0)
+        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 / LEGACY_TASKS_PATH)
+    write_jsonlines(tasks, v20_dir / 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 / LEGACY_EPISODES_PATH)
+    write_jsonlines(episodes, v20_dir / 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": LEGACY_DEFAULT_PARQUET_PATH,
-        "video_path": LEGACY_DEFAULT_VIDEO_PATH if video_keys else None,
+        "data_path": DEFAULT_PARQUET_PATH,
+        "video_path": 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

@@ -37,7 +37,7 @@ import logging
 from huggingface_hub import HfApi
 
 from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
-from lerobot.datasets.utils import STATS_PATH, load_stats, write_info
+from lerobot.datasets.utils import EPISODES_STATS_PATH, STATS_PATH, load_stats, write_info
 from lerobot.datasets.v21.convert_stats import check_aggregate_stats, convert_stats
 
 V20 = "v2.0"
@@ -61,6 +61,9 @@ def convert_dataset(
     with SuppressWarnings():
         dataset = LeRobotDataset(repo_id, revision=V20, force_cache_sync=True)
 
+    if (dataset.root / EPISODES_STATS_PATH).is_file():
+        (dataset.root / EPISODES_STATS_PATH).unlink()
+
     convert_stats(dataset, num_workers=num_workers)
     ref_stats = load_stats(dataset.root)
     check_aggregate_stats(dataset, ref_stats)

src/lerobot/datasets/v21/convert_stats.py

@@ -13,28 +13,13 @@
 # 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 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)
+from lerobot.datasets.utils import write_episode_stats
 
 
 def sample_episode_video_frames(dataset: LeRobotDataset, episode_index: int, ft_key: str) -> np.ndarray:
@@ -87,7 +72,7 @@ def convert_stats(dataset: LeRobotDataset, num_workers: int = 0):
             convert_episode_stats(dataset, ep_idx)
 
     for ep_idx in tqdm(range(total_episodes)):
-        legacy_write_episode_stats(ep_idx, dataset.meta.episodes_stats[ep_idx], dataset.root)
+        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

@@ -1,480 +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.
-
-"""
-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,26 +13,22 @@
 # 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"):
@@ -106,7 +102,7 @@ def decode_video_frames_torchvision(
     keyframes_only = False
     torchvision.set_video_backend(backend)
     if backend == "pyav":
-        keyframes_only = True  # pyav doesnt support accuracte seek
+        keyframes_only = True  # pyav doesn't support accurate seek
 
     # set a video stream reader
     # TODO(rcadene): also load audio stream at the same time
@@ -159,7 +155,6 @@ 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_]
 
@@ -257,104 +252,83 @@ def encode_video_frames(
     g: int | None = 2,
     crf: int | None = 30,
     fast_decode: int = 0,
-    log_level: str | None = "quiet",
+    log_level: int | None = av.logging.ERROR,
     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)
 
-    ffmpeg_args = OrderedDict(
-        [
-            ("-f", "image2"),
-            ("-r", str(fps)),
-            ("-i", str(imgs_dir / "frame-%06d.png")),
-            ("-vcodec", vcodec),
-            ("-pix_fmt", pix_fmt),
-        ]
+    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])
     )
 
+    # 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:
-        ffmpeg_args["-g"] = str(g)
+        video_options["g"] = str(g)
 
     if crf is not None:
-        ffmpeg_args["-crf"] = str(crf)
+        video_options["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"
-        ffmpeg_args[key] = value
+        video_options[key] = value
 
+    # Set logging level
     if log_level is not None:
-        ffmpeg_args["-loglevel"] = str(log_level)
+        # "While less efficient, it is generally preferable to modify logging with Python’s logging"
+        logging.getLogger("libav").setLevel(log_level)
 
-    ffmpeg_args = [item for pair in ffmpeg_args.items() for item in pair]
-    if overwrite:
-        ffmpeg_args.append("-y")
+    # 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_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)
+        # 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()
 
     if not video_path.exists():
-        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))
+        raise OSError(f"Video encoding did not work. File not found: {video_path}.")
 
 
 @dataclass
@@ -390,78 +364,68 @@ with warnings.catch_warnings():
 
 
 def get_audio_info(video_path: Path | str) -> dict:
-    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}")
+    # Set logging level
+    logging.getLogger("libav").setLevel(av.logging.ERROR)
 
-    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}
+    # 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}
 
-    # 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),
-    }
+        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
 
 
 def get_video_info(video_path: Path | str) -> dict:
-    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}")
+    # Set logging level
+    logging.getLogger("libav").setLevel(av.logging.ERROR)
 
-    info = json.loads(result.stdout)
-    video_stream_info = info["streams"][0]
+    # 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 {}
 
-    # 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
+        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
 
-    pixel_channels = get_video_pixel_channels(video_stream_info["pix_fmt"])
+        # Calculate fps from r_frame_rate
+        video_info["video.fps"] = int(video_stream.base_rate)
 
-    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),
-    }
+        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))
 
     return video_info
 
@@ -488,3 +452,66 @@ 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/envs/configs.py

@@ -30,6 +30,8 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC):
     fps: int = 30
     features: dict[str, PolicyFeature] = field(default_factory=dict)
     features_map: dict[str, str] = field(default_factory=dict)
+    multitask_eval: bool = False
+    max_parallel_tasks: int = 5
 
     @property
     def type(self) -> str:
@@ -271,3 +273,53 @@ class HILEnvConfig(EnvConfig):
             "use_gamepad": self.use_gamepad,
             "gripper_penalty": self.gripper_penalty,
         }
+
+
+@EnvConfig.register_subclass("libero")
+@dataclass
+class LiberoEnv(EnvConfig):
+    task: str = "libero_10"  # can also choose libero_spatial, libero_object, etc.
+    fps: int = 30
+    episode_length: int = 520
+    obs_type: str = "pixels_agent_pos"
+    render_mode: str = "rgb_array"
+    camera_name: str = "agentview_image,robot0_eye_in_hand_image"
+    init_states: bool = True
+    multitask_eval: bool = True
+    features: dict[str, PolicyFeature] = field(
+        default_factory=lambda: {
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
+        }
+    )
+    features_map: dict[str, str] = field(
+        default_factory=lambda: {
+            "action": ACTION,
+            "agent_pos": OBS_STATE,
+            "pixels/agentview_image": f"{OBS_IMAGES}.image",
+            "pixels/robot0_eye_in_hand_image": f"{OBS_IMAGES}.image2",
+        }
+    )
+
+    def __post_init__(self):
+        if self.obs_type == "pixels":
+            self.features["pixels/agentview_image"] = PolicyFeature(
+                type=FeatureType.VISUAL, shape=(360, 360, 3)
+            )
+            self.features["pixels/robot0_eye_in_hand_image"] = PolicyFeature(
+                type=FeatureType.VISUAL, shape=(360, 360, 3)
+            )
+        elif self.obs_type == "pixels_agent_pos":
+            self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(8,))
+            self.features["pixels/agentview_image"] = PolicyFeature(
+                type=FeatureType.VISUAL, shape=(360, 360, 3)
+            )
+            self.features["pixels/robot0_eye_in_hand_image"] = PolicyFeature(
+                type=FeatureType.VISUAL, shape=(360, 360, 3)
+            )
+
+    @property
+    def gym_kwargs(self) -> dict:
+        return {
+            "obs_type": self.obs_type,
+            "render_mode": self.render_mode,
+        }

src/lerobot/envs/factory.py

@@ -17,7 +17,7 @@ import importlib
 
 import gymnasium as gym
 
-from lerobot.envs.configs import AlohaEnv, EnvConfig, HILEnvConfig, PushtEnv, XarmEnv
+from lerobot.envs.configs import AlohaEnv, EnvConfig, HILEnvConfig, LiberoEnv, PushtEnv, XarmEnv
 
 
 def make_env_config(env_type: str, **kwargs) -> EnvConfig:
@@ -29,11 +29,15 @@ def make_env_config(env_type: str, **kwargs) -> EnvConfig:
         return XarmEnv(**kwargs)
     elif env_type == "hil":
         return HILEnvConfig(**kwargs)
+    elif env_type == "libero":
+        return LiberoEnv(**kwargs)
     else:
         raise ValueError(f"Policy type '{env_type}' is not available.")
 
 
-def make_env(cfg: EnvConfig, n_envs: int = 1, use_async_envs: bool = False) -> gym.vector.VectorEnv | None:
+def make_env(
+    cfg: EnvConfig, n_envs: int = 1, use_async_envs: bool = False
+) -> gym.vector.VectorEnv | dict[str, dict[int, gym.vector.VectorEnv]]:
     """Makes a gym vector environment according to the config.
 
     Args:
@@ -48,24 +52,43 @@ def make_env(cfg: EnvConfig, n_envs: int = 1, use_async_envs: bool = False) -> g
 
     Returns:
         gym.vector.VectorEnv: The parallelized gym.env instance.
+        dict[str, dict[int, gym.vector.VectorEnv]]: A mapping from task suite
+            names to indexed vectorized environments (when multitask eval is used).
+
     """
     if n_envs < 1:
-        raise ValueError("`n_envs must be at least 1")
+        raise ValueError("`n_envs` must be at least 1")
+
+    env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
+
+    if "libero" in cfg.type:
+        from lerobot.envs.libero import create_libero_envs
+
+        return create_libero_envs(
+            task=cfg.task,
+            n_envs=n_envs,
+            camera_name=cfg.camera_name,
+            init_states=cfg.init_states,
+            gym_kwargs=cfg.gym_kwargs,
+            env_cls=env_cls,
+            multitask_eval=cfg.multitask_eval,
+        )
 
     package_name = f"gym_{cfg.type}"
-
     try:
         importlib.import_module(package_name)
     except ModuleNotFoundError as e:
-        print(f"{package_name} is not installed. Please install it with `pip install 'lerobot[{cfg.type}]'`")
-        raise e
+        raise ModuleNotFoundError(
+            f'{package_name} is not installed. Install with: pip install "lerobot[{cfg.type}]"'
+        ) from e
 
     gym_handle = f"{package_name}/{cfg.task}"
 
-    # batched version of the env that returns an observation of shape (b, c)
-    env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
-    env = env_cls(
-        [lambda: gym.make(gym_handle, disable_env_checker=True, **cfg.gym_kwargs) for _ in range(n_envs)]
-    )
+    def _make_one():
+        return gym.make(gym_handle, disable_env_checker=True, **(cfg.gym_kwargs or {}))
 
-    return env
+    vec = env_cls([_make_one for _ in range(n_envs)])
+
+    # normalize to {suite: {task_id: vec_env}} for consistency
+    suite_name = cfg.type  # e.g., "pusht", "aloha"
+    return {suite_name: {0: vec}}

src/lerobot/envs/libero.py

@@ -0,0 +1,497 @@
+from __future__ import annotations
+
+import logging
+import math
+import os
+from collections import defaultdict
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from itertools import chain
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+import torch
+from gymnasium import spaces
+from libero.libero import benchmark, get_libero_path
+from libero.libero.envs import OffScreenRenderEnv
+
+logger = logging.getLogger(__name__)
+
+# ---- Helpers -----------------------------------------------------------------
+
+
+def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
+    """Normalize camera_name into a non-empty list of strings."""
+    if isinstance(camera_name, str):
+        cams = [c.strip() for c in camera_name.split(",") if c.strip()]
+    elif isinstance(camera_name, (list, tuple)):
+        cams = [str(c).strip() for c in camera_name if str(c).strip()]
+    else:
+        raise TypeError(f"camera_name must be str or sequence[str], got {type(camera_name).__name__}")
+    if not cams:
+        raise ValueError("camera_name resolved to an empty list.")
+    return cams
+
+
+def _get_suite(name: str):
+    """Instantiate a LIBERO suite by name with clear validation."""
+    bench = benchmark.get_benchmark_dict()
+    if name not in bench:
+        raise ValueError(f"Unknown LIBERO suite '{name}'. Available: {', '.join(sorted(bench.keys()))}")
+    suite = bench[name]()
+    if not getattr(suite, "tasks", None):
+        raise ValueError(f"Suite '{name}' has no tasks.")
+    return suite
+
+
+def _select_task_ids(total_tasks: int, task_ids: Iterable[int] | None) -> list[int]:
+    """Validate/normalize task ids. If None → all tasks."""
+    if task_ids is None:
+        return list(range(total_tasks))
+    ids = sorted({int(t) for t in task_ids})
+    for t in ids:
+        if t < 0 or t >= total_tasks:
+            raise ValueError(f"task_id {t} out of range [0, {total_tasks - 1}].")
+    return ids
+
+
+def _make_env_fns(
+    *,
+    suite,
+    suite_name: str,
+    task_id: int,
+    n_envs: int,
+    camera_names: list[str],
+    init_states: bool,
+    gym_kwargs: Mapping[str, Any],
+    LiberoEnv: type,  # injected to avoid forward ref issues if needed
+) -> list[Callable[[], LiberoEnv]]:
+    """Build n_envs factory callables for a single (suite, task_id)."""
+    joined_cams = ",".join(camera_names)  # keep backward-compat: downstream expects a string
+    fns: list[Callable[[], LiberoEnv]] = []
+    for i in range(n_envs):
+
+        def _mk(
+            i=i,
+            suite=suite,
+            task_id=task_id,
+            suite_name=suite_name,
+            joined_cams=joined_cams,
+            init_states=init_states,
+            gym_kwargs=dict(gym_kwargs),
+        ):
+            return LiberoEnv(
+                task_suite=suite,
+                task_id=task_id,
+                task_suite_name=suite_name,
+                camera_name=joined_cams,
+                init_states=init_states,
+                episode_index=i,
+                **gym_kwargs,
+            )
+
+        fns.append(_mk)
+    return fns
+
+
+# ---- Main API ----------------------------------------------------------------
+
+
+def create_libero_envs(
+    task: str,
+    n_envs: int,
+    gym_kwargs: dict[str, Any] | None = None,
+    camera_name: str | Sequence[str] = "agentview_image,robot0_eye_in_hand_image",
+    init_states: bool = True,
+    env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
+    multitask_eval: bool = True,  # kept for signature compatibility; return type is consistent regardless
+) -> dict[str, dict[int, Any]]:
+    """
+    Create vectorized LIBERO environments with a consistent return shape.
+
+    Returns:
+        dict[suite_name][task_id] -> vec_env (env_cls([...]) with exactly n_envs factories)
+    Notes:
+        - n_envs is the number of rollouts *per task* (episode_index = 0..n_envs-1).
+        - `task` can be a single suite or a comma-separated list of suites.
+        - You may pass `task_ids` (list[int]) inside `gym_kwargs` to restrict tasks per suite.
+    """
+    if env_cls is None or not callable(env_cls):
+        raise ValueError("env_cls must be a callable that wraps a list of environment factory callables.")
+    if not isinstance(n_envs, int) or n_envs <= 0:
+        raise ValueError(f"n_envs must be a positive int; got {n_envs}.")
+
+    gym_kwargs = dict(gym_kwargs or {})
+    task_ids_filter = gym_kwargs.pop("task_ids", None)  # optional: limit to specific tasks
+
+    # Avoid circular import/type issues: assume LiberoEnv is defined in this module
+    try:
+        LiberoEnv  # type: ignore[name-defined]
+    except NameError:
+        # If LiberoEnv is in the same file, this won't run. If it's elsewhere, import here.
+        exit()
+        # from .libero_env import LiberoEnv  # adjust if your class lives in another module
+
+    camera_names = _parse_camera_names(camera_name)
+    suite_names = [s.strip() for s in str(task).split(",") if s.strip()]
+    if not suite_names:
+        raise ValueError("`task` must contain at least one LIBERO suite name.")
+
+    logger.info(
+        "Creating LIBERO envs | suites=%s | n_envs(per task)=%d | init_states=%s | multitask_eval=%s",
+        suite_names,
+        n_envs,
+        init_states,
+        bool(multitask_eval),
+    )
+    if task_ids_filter is not None:
+        logger.info("Restricting to task_ids=%s", task_ids_filter)
+
+    out: dict[str, dict[int, Any]] = defaultdict(dict)
+
+    for suite_name in suite_names:
+        suite = _get_suite(suite_name)
+        total = len(suite.tasks)
+        selected = _select_task_ids(total, task_ids_filter)
+
+        if not selected:
+            raise ValueError(f"No tasks selected for suite '{suite_name}' (available: {total}).")
+
+        for tid in selected:
+            fns = _make_env_fns(
+                suite=suite,
+                suite_name=suite_name,
+                task_id=tid,
+                n_envs=n_envs,
+                camera_names=camera_names,
+                init_states=init_states,
+                gym_kwargs=gym_kwargs,
+                LiberoEnv=LiberoEnv,
+            )
+            out[suite_name][tid] = env_cls(fns)
+            logger.debug("Built vec env | suite=%s | task_id=%d | n_envs=%d", suite_name, tid, n_envs)
+
+    # return plain dicts for predictability
+    return {suite: dict(task_map) for suite, task_map in out.items()}
+
+
+def quat2axisangle(quat):
+    """
+    Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55
+
+    Converts quaternion to axis-angle format.
+    Returns a unit vector direction scaled by its angle in radians.
+
+    Args:
+        quat (np.array): (x,y,z,w) vec4 float angles
+
+    Returns:
+        np.array: (ax,ay,az) axis-angle exponential coordinates
+    """
+    # clip quaternion
+    if quat[3] > 1.0:
+        quat[3] = 1.0
+    elif quat[3] < -1.0:
+        quat[3] = -1.0
+
+    den = np.sqrt(1.0 - quat[3] * quat[3])
+    if math.isclose(den, 0.0):
+        # This is (close to) a zero degree rotation, immediately return
+        return np.zeros(3)
+
+    return (quat[:3] * 2.0 * math.acos(quat[3])) / den
+
+
+def get_task_init_states(task_suite, i):
+    init_states_path = os.path.join(
+        get_libero_path("init_states"),
+        task_suite.tasks[i].problem_folder,
+        task_suite.tasks[i].init_states_file,
+    )
+    init_states = torch.load(init_states_path, weights_only=False)  # nosec B614
+    return init_states
+
+
+def get_libero_dummy_action():
+    """Get dummy/no-op action, used to roll out the simulation while the robot does nothing."""
+    return [0, 0, 0, 0, 0, 0, -1]
+
+
+OBS_STATE_DIM = 8
+ACTION_DIM = 7
+
+
+class LiberoEnv(gym.Env):
+    metadata = {"render_modes": ["rgb_array"], "render_fps": 80}
+
+    def __init__(
+        self,
+        task_suite,
+        task_id,
+        task_suite_name,
+        camera_name="agentview_image,robot0_eye_in_hand_image",
+        obs_type="pixels",
+        render_mode="rgb_array",
+        observation_width=256,
+        observation_height=256,
+        visualization_width=640,
+        visualization_height=480,
+        init_states=True,
+        episode_index=0,
+    ):
+        super().__init__()
+        self.task_id = task_id
+        self.obs_type = obs_type
+        self.render_mode = render_mode
+        self.observation_width = observation_width
+        self.observation_height = observation_height
+        self.visualization_width = visualization_width
+        self.visualization_height = visualization_height
+        self.init_states = init_states
+        self.camera_name = camera_name.split(
+            ","
+        )  # agentview_image (main) or robot0_eye_in_hand_image (wrist)
+
+        # Map raw camera names to "image1" and "image2".
+        # The preprocessing step `preprocess_observation` will then prefix these with `.images.*`,
+        # following the LeRobot convention (e.g., `observation.images.image`, `observation.images.image2`).
+        # This ensures the policy consistently receives observations in the
+        # expected format regardless of the original camera naming.
+        self.camera_name_mapping = {
+            "agentview_image": "image",
+            "robot0_eye_in_hand_image": "image2",
+        }
+
+        self.num_steps_wait = (
+            10  # Do nothing for the first few timesteps to wait for the simulator drops objects
+        )
+        self.episode_index = episode_index
+
+        self._env = self._make_envs_task(task_suite, self.task_id)
+        TASK_SUITE_MAX_STEPS: dict[str, int] = {
+            "libero_spatial": 220,  # longest training demo has 193 steps
+            "libero_object": 280,  # longest training demo has 254 steps
+            "libero_goal": 300,  # longest training demo has 270 steps
+            "libero_10": 520,  # longest training demo has 505 steps
+            "libero_90": 400,  # longest training demo has 373 steps
+        }
+        default_steps = 500
+        self._max_episode_steps = TASK_SUITE_MAX_STEPS.get(task_suite_name, default_steps)
+
+        images = {}
+        for cam in self.camera_name:
+            images[self.camera_name_mapping[cam]] = spaces.Box(
+                low=0,
+                high=255,
+                shape=(self.observation_height, self.observation_width, 3),
+                dtype=np.uint8,
+            )
+
+        if self.obs_type == "state":
+            raise NotImplementedError(
+                "The 'state' observation type is not supported in LiberoEnv. "
+                "Please switch to an image-based obs_type (e.g. 'pixels', 'pixels_agent_pos')."
+            )
+
+        elif self.obs_type == "pixels":
+            self.observation_space = spaces.Dict(
+                {
+                    "pixels": spaces.Dict(images),
+                }
+            )
+        elif self.obs_type == "pixels_agent_pos":
+            self.observation_space = spaces.Dict(
+                {
+                    "pixels": spaces.Dict(images),
+                    "agent_pos": spaces.Box(
+                        low=-1000.0,
+                        high=1000.0,
+                        shape=(OBS_STATE_DIM,),
+                        dtype=np.float64,
+                    ),
+                }
+            )
+
+        self.action_space = spaces.Box(low=-1, high=1, shape=(ACTION_DIM,), dtype=np.float32)
+
+    def render(self):
+        raw_obs = self._env.env._get_observations()
+        image = self._format_raw_obs(raw_obs)["pixels"]["image"]
+        return image
+
+    def _make_envs_task(self, task_suite, task_id: int = 0):
+        task = task_suite.get_task(task_id)
+        self.task = task.name
+        self.task_description = task.language
+        task_bddl_file = os.path.join(get_libero_path("bddl_files"), task.problem_folder, task.bddl_file)
+
+        env_args = {
+            "bddl_file_name": task_bddl_file,
+            "camera_heights": self.observation_height,
+            "camera_widths": self.observation_width,
+        }
+        env = OffScreenRenderEnv(**env_args)
+        env.reset()
+        if self.init_states:
+            init_states = get_task_init_states(
+                task_suite, task_id
+            )  # for benchmarking purpose, we fix the a set of initial states FIXME(mshukor): should be in the reset()?
+            init_state_id = self.episode_index  # episode index
+            env.set_init_state(init_states[init_state_id])
+
+        return env
+
+    def _format_raw_obs(self, raw_obs):
+        images = {}
+        for camera_name in self.camera_name:
+            image = raw_obs[camera_name]
+            image = image[::-1, ::-1]  # rotate 180 degrees
+            images[self.camera_name_mapping[camera_name]] = image
+        state = np.concatenate(
+            (
+                raw_obs["robot0_eef_pos"],
+                quat2axisangle(raw_obs["robot0_eef_quat"]),
+                raw_obs["robot0_gripper_qpos"],
+            )
+        )
+        agent_pos = state
+        if self.obs_type == "state":
+            raise NotImplementedError(
+                "The 'state' observation type is not supported in LiberoEnv. "
+                "Please switch to an image-based obs_type (e.g. 'pixels', 'pixels_agent_pos')."
+            )
+        elif self.obs_type == "pixels":
+            obs = {"pixels": images.copy()}
+        elif self.obs_type == "pixels_agent_pos":
+            obs = {
+                "pixels": images.copy(),
+                "agent_pos": agent_pos,
+            }
+        return obs
+
+    def reset(self, seed=None, **kwargs):
+        super().reset(seed=seed)
+
+        self._env.seed(seed)
+        raw_obs = self._env.reset()
+        # Do nothing for the first few timesteps to wait for the simulator drops objects
+        for _ in range(self.num_steps_wait):
+            raw_obs, _, _, _ = self._env.step(get_libero_dummy_action())
+        observation = self._format_raw_obs(raw_obs)
+        info = {"is_success": False}
+        return observation, info
+
+    def step(self, action):
+        if action.ndim != 1:
+            raise ValueError(
+                f"Expected action to be 1-D (shape (action_dim,)), "
+                f"but got shape {action.shape} with ndim={action.ndim}"
+            )
+        raw_obs, reward, done, info = self._env.step(action)
+
+        is_success = self._env.check_success()
+        terminated = done or is_success
+        info["is_success"] = done  # is_success
+
+        observation = self._format_raw_obs(raw_obs)
+        if done:
+            self.reset()
+            print(self.task, self.task_id, done, is_success)
+        truncated = False
+        return observation, reward, terminated, truncated, info
+
+    def close(self):
+        self._env.close()
+
+
+def create_libero_envs1(
+    task: str,
+    n_envs: int,
+    gym_kwargs: dict[str, Any] = None,
+    camera_name: str = "agentview_image,robot0_eye_in_hand_image",
+    init_states: bool = True,
+    env_cls: Callable = None,
+    multitask_eval: bool = True,
+) -> dict[str, dict[str, Any]]:
+    """
+    Here n_envs is per task and equal to the number of rollouts.
+    Returns:
+        dict[str, dict[str, list[LiberoEnv]]]: keys are task_suite and values are list of LiberoEnv envs.
+    """
+    print("num envs", n_envs)
+    print("multitask_eval", multitask_eval)
+    print("gym_kwargs", gym_kwargs)
+    if gym_kwargs is None:
+        gym_kwargs = {}
+
+    if not multitask_eval:
+        benchmark_dict = benchmark.get_benchmark_dict()
+        task_suite = benchmark_dict[task]()  # can also choose libero_spatial, libero_object, libero_10 etc.
+        tasks_id = list(range(len(task_suite.tasks)))
+        episode_indices = [0 for i in range(len(tasks_id))]
+        if len(tasks_id) == 1:
+            tasks_id = [tasks_id[0] for _ in range(n_envs)]
+            episode_indices = list(range(n_envs))
+        elif len(tasks_id) < n_envs and n_envs % len(tasks_id) == 0:
+            n_repeat = n_envs // len(tasks_id)
+            print("n_repeat", n_repeat)
+            episode_indices = []
+            for _ in range(len(tasks_id)):
+                episode_indices.extend(list(range(n_repeat)))
+            tasks_id = list(chain.from_iterable([[item] * n_repeat for item in tasks_id]))
+        elif n_envs < len(tasks_id):
+            tasks_id = tasks_id[:n_envs]
+            episode_indices = list(range(n_envs))[:n_envs]
+            print(f"WARNING: n_envs < len(tasks_id), evaluating only on {tasks_id}")
+        print(f"Creating Libero envs with task ids {tasks_id} from suite {task}")
+        assert n_envs == len(tasks_id), (
+            f"len(n_envs) and tasks_id should be the same, got {n_envs} and {len(tasks_id)}"
+        )
+        return env_cls(
+            [
+                lambda i=i: LiberoEnv(
+                    task_suite=task_suite,
+                    task_id=tasks_id[i],
+                    task_suite_name=task,
+                    camera_name=camera_name,
+                    init_states=init_states,
+                    episode_index=episode_indices[i],
+                    **gym_kwargs,
+                )
+                for i in range(n_envs)
+            ]
+        )
+    else:
+        envs = defaultdict(dict)
+        benchmark_dict = benchmark.get_benchmark_dict()
+        task = task.split(",")
+        for _task in task:
+            task_suite = benchmark_dict[
+                _task
+            ]()  # can also choose libero_spatial, libero_object, libero_10 etc.
+            tasks_ids = list(range(len(task_suite.tasks)))
+            for tasks_id in tasks_ids:
+                episode_indices = list(range(n_envs))
+                print(
+                    f"Creating Libero envs with task ids {tasks_id} from suite {_task}, episode_indices: {episode_indices}"
+                )
+                envs_list = [
+                    (
+                        lambda i=i,
+                        task_suite=task_suite,
+                        tasks_id=tasks_id,
+                        _task=_task,
+                        episode_indices=episode_indices: LiberoEnv(
+                            task_suite=task_suite,
+                            task_id=tasks_id,
+                            task_suite_name=_task,
+                            camera_name=camera_name,
+                            init_states=init_states,
+                            episode_index=episode_indices[i],
+                            **gym_kwargs,
+                        )
+                    )
+                    for i in range(n_envs)
+                ]
+                envs[_task][tasks_id] = env_cls(envs_list)
+        return envs

src/lerobot/envs/utils.py

@@ -134,3 +134,49 @@ def add_envs_task(env: gym.vector.VectorEnv, observation: dict[str, Any]) -> dic
         num_envs = observation[list(observation.keys())[0]].shape[0]
         observation["task"] = ["" for _ in range(num_envs)]
     return observation
+
+
+def _close_single_env(env: Any) -> None:
+    """Try to close a single env object if it exposes .close()."""
+    try:
+        close_fn = getattr(env, "close", None)
+        if callable(close_fn):
+            close_fn()
+    except Exception as exc:
+        # Best-effort close: log but don't raise
+        LOG.debug("Exception while closing env %s: %s", env, exc)
+
+
+def close_envs(env_or_collection: Any) -> None:
+    """
+    Close a single env or any nested structure of envs.
+
+    Accepts:
+      - a single env with .close()
+      - a Mapping of things (e.g. dict)
+      - a Sequence of things (list/tuple) but NOT str/bytes
+      - nested combinations of the above
+
+    This is intentionally permissive and best-effort: it will swallow exceptions
+    encountered while closing individual envs and continue.
+    """
+    # Guard: single object with close()
+    if hasattr(env_or_collection, "close") and not isinstance(env_or_collection, (Mapping, Sequence)):
+        _close_single_env(env_or_collection)
+        return
+
+    # Mapping (e.g., {suite: {task_id: vec_env}})
+    if isinstance(env_or_collection, Mapping):
+        for v in env_or_collection.values():
+            close_envs(v)
+        return
+
+    # Sequence (list/tuple) but skip str/bytes
+    if isinstance(env_or_collection, Sequence) and not isinstance(env_or_collection, (str, bytes)):
+        for v in env_or_collection:
+            close_envs(v)
+        return
+
+    # Fallback: try to close if possible
+    if hasattr(env_or_collection, "close"):
+        _close_single_env(env_or_collection)

src/lerobot/motors/dynamixel/tables.py

@@ -107,6 +107,8 @@ X_SERIES_ENCODINGS_TABLE = {
     "Goal_PWM": X_SERIES_CONTROL_TABLE["Goal_PWM"][1],
     "Goal_Current": X_SERIES_CONTROL_TABLE["Goal_Current"][1],
     "Goal_Velocity": X_SERIES_CONTROL_TABLE["Goal_Velocity"][1],
+    "Goal_Position": X_SERIES_CONTROL_TABLE["Goal_Position"][1],
+    "Present_Position": X_SERIES_CONTROL_TABLE["Present_Position"][1],
     "Present_PWM": X_SERIES_CONTROL_TABLE["Present_PWM"][1],
     "Present_Current": X_SERIES_CONTROL_TABLE["Present_Current"][1],
     "Present_Velocity": X_SERIES_CONTROL_TABLE["Present_Velocity"][1],

src/lerobot/policies/factory.py

@@ -27,7 +27,9 @@ from lerobot.envs.utils import env_to_policy_features
 from lerobot.policies.act.configuration_act import ACTConfig
 from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
+from lerobot.policies.pi0_openpi.configuration_pi0openpi import PI0OpenPIConfig
 from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
+from lerobot.policies.pi05_openpi.configuration_pi05openpi import PI05OpenPIConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.configuration_sac import SACConfig
 from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
@@ -62,6 +64,14 @@ def get_policy_class(name: str) -> PreTrainedPolicy:
         from lerobot.policies.pi0fast.modeling_pi0fast import PI0FASTPolicy
 
         return PI0FASTPolicy
+    elif name == "pi0_openpi":
+        from lerobot.policies.pi0_openpi.modeling_pi0openpi import PI0OpenPIPolicy
+
+        return PI0OpenPIPolicy
+    elif name == "pi05_openpi":
+        from lerobot.policies.pi05_openpi.modeling_pi05openpi import PI05OpenPIPolicy
+
+        return PI05OpenPIPolicy
     elif name == "sac":
         from lerobot.policies.sac.modeling_sac import SACPolicy
 
@@ -91,6 +101,10 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return PI0Config(**kwargs)
     elif policy_type == "pi0fast":
         return PI0FASTConfig(**kwargs)
+    elif policy_type == "pi0_openpi":
+        return PI0OpenPIConfig(**kwargs)
+    elif policy_type == "pi05_openpi":
+        return PI05OpenPIConfig(**kwargs)
     elif policy_type == "sac":
         return SACConfig(**kwargs)
     elif policy_type == "smolvla":
@@ -172,7 +186,5 @@ def make_policy(
 
     policy.to(cfg.device)
     assert isinstance(policy, nn.Module)
-
     # policy = torch.compile(policy, mode="reduce-overhead")
-
     return policy

src/lerobot/policies/pi05_openpi/README.md

@@ -0,0 +1,92 @@
+# π₀.₅ (pi05)
+
+This repository contains the Hugging Face port of **π₀.₅**, adapted from [OpenPI](https://github.com/Physical-Intelligence/openpi) by the Physical Intelligence.
+It is designed as a **Vision-Language-Action model with open-world generalization**.
+
+---
+
+### ⚠️ WARNING ⚠️
+
+This project requires **patching the Hugging Face `transformers` library**.
+
+1. Make sure you have the exact version installed:
+
+```bash
+   pip show transformers
+```
+
+It must be version **4.53.2**.
+
+2. Apply the custom patches by copying the modified files into your environment:
+
+   ```bash
+   cp -r ./src/lerobot/policies/pi0_openpi/transformers_replace/* \
+     $(python -c "import transformers, os; print(os.path.dirname(transformers.__file__))")
+   ```
+
+   These patches overwrite parts of `transformers` to:
+   - Support the **AdaRMS optimizer**,
+   - Correctly control the precision of activations,
+   - Allow the KV cache to be used without updates.
+
+**Important:**
+
+- This permanently modifies your `transformers` installation.
+- The changes survive reinstalls unless you explicitly remove the patched files or recreate the environment.
+
+To undo and restore a clean state:
+
+```bash
+pip uninstall transformers
+pip install transformers==4.53.2
+```
+
+---
+
+## Model Overview
+
+| Feature              | π₀                                                     | π₀.₅                                      |
+| -------------------- | ------------------------------------------------------ | ----------------------------------------- |
+| State Embedding      | Uses `state_proj` layer                                | No state embedding                        |
+| Time Conditioning    | Concatenates time with actions via `action_time_mlp_*` | Uses `time_mlp_*` for AdaRMS conditioning |
+| AdaRMS               | Not used                                               | Used in action expert                     |
+| Tokenizer Length     | 48 tokens                                              | 200 tokens                                |
+| Discrete State Input | False                                                  | True                                      |
+| Parameter Count      | Higher (includes state embedding)                      | Lower (no state embedding)                |
+
+---
+
+## Citation
+
+If you use this work, please cite both **OpenPI** and the π₀.₅ paper:
+
+```bibtex
+@misc{openpi2024,
+  author       = {Physical Intelligence Lab},
+  title        = {OpenPI: PyTorch Implementation of π0 and π0.5 Policies},
+  year         = {2024},
+  publisher    = {GitHub},
+  howpublished = {\url{https://github.com/Physical-Intelligence/openpi}},
+  license      = {Apache-2.0}
+}
+
+@misc{intelligence2025pi05visionlanguageactionmodelopenworld,
+  title        = {π₀.₅: a Vision-Language-Action Model with Open-World Generalization},
+  author       = {Physical Intelligence and Kevin Black and Noah Brown and James Darpinian and Karan Dhabalia and Danny Driess and Adnan Esmail and Michael Equi and Chelsea Finn and Niccolo Fusai and Manuel Y. Galliker and Dibya Ghosh and Lachy Groom and Karol Hausman and Brian Ichter and Szymon Jakubczak and Tim Jones and Liyiming Ke and Devin LeBlanc and Sergey Levine and Adrian Li-Bell and Mohith Mothukuri and Suraj Nair and Karl Pertsch and Allen Z. Ren and Lucy Xiaoyang Shi and Laura Smith and Jost Tobias Springenberg and Kyle Stachowicz and James Tanner and Quan Vuong and Homer Walke and Anna Walling and Haohuan Wang and Lili Yu and Ury Zhilinsky},
+  year         = {2025},
+  eprint       = {2504.16054},
+  archivePrefix= {arXiv},
+  primaryClass = {cs.LG},
+  url          = {https://arxiv.org/abs/2504.16054},
+}
+```
+
+---
+
+## License
+
+This port follows the **Apache 2.0 License**, consistent with the original [OpenPI repository](https://github.com/Physical-Intelligence/openpi).
+
+```
+
+```

src/lerobot/policies/pi05_openpi/__init__.py

@@ -0,0 +1,20 @@
+#!/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 .configuration_pi05openpi import PI05OpenPIConfig
+from .modeling_pi05openpi import PI05OpenPIPolicy
+
+__all__ = ["PI05OpenPIConfig", "PI05OpenPIPolicy"]

src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py

@@ -0,0 +1,137 @@
+#!/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 dataclasses import dataclass, field
+
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.types import NormalizationMode
+from lerobot.optim.optimizers import AdamWConfig
+from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
+
+
+@PreTrainedConfig.register_subclass("pi05_openpi")
+@dataclass
+class PI05OpenPIConfig(PreTrainedConfig):
+    # Model architecture
+    paligemma_variant: str = "gemma_2b"
+    action_expert_variant: str = "gemma_300m"
+    discrete_state_input: bool | None = (
+        True  # Whether to use discrete state input # see openpi `Pi0Config, __post_init__`
+    )
+    dtype: str = "float32"  # Options: "bfloat16", "float32"
+
+    # Input / output structure
+    n_obs_steps: int = 1
+    chunk_size: int = 50  # Number of action steps to predict, in openpi called "action_horizon"
+    n_action_steps: int = 50  # Number of action steps to execute
+
+    # Shorter state and action vectors will be padded to these dimensions
+    max_state_dim: int = 32  # State dimension (will be padded to 32)
+    max_action_dim: int = 32  # Action dimension (will be padded to 32)
+
+    # Flow matching parameters: see openpi `PI0Pytorch`
+    num_inference_steps: int = 10  # Number of denoising steps during inference
+    time_sampling_beta_alpha: float = 1.5  # Beta distribution alpha parameter for time sampling
+    time_sampling_beta_beta: float = 1.0  # Beta distribution beta parameter for time sampling
+    min_period: float = 4e-3  # Min period for sinusoidal positional encoding
+    max_period: float = 4.0  # Max period for sinusoidal positional encoding
+
+    # Image preprocessing
+    image_resolution: tuple[int, int] = (224, 224)  # see openpi `preprocessing_pytorch.py`
+
+    # Normalization
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.IDENTITY,  # Images are normalized to [-1, 1] in preprocessing
+            "STATE": NormalizationMode.MEAN_STD,
+            "ACTION": NormalizationMode.MEAN_STD,
+        }
+    )
+
+    # Training settings
+    gradient_checkpointing: bool = False  # Enable gradient checkpointing for memory optimization
+    compile_model: bool = False  # Whether to use torch.compile for model optimization
+    compile_mode: str = "max-autotune"  # Torch compile mode
+    device: str | None = None  # Device to use for the model (None = auto-detect)
+
+    # Optimizer settings: see openpi `AdamW` and
+    optimizer_lr: float = 2.5e-5  # see openpi `CosineDecaySchedule: peak_lr`
+    optimizer_betas: tuple[float, float] = (0.9, 0.95)
+    optimizer_eps: float = 1e-8
+    optimizer_weight_decay: float = 0.01
+    optimizer_grad_clip_norm: float = 1.0
+
+    # Scheduler settings: see openpi `CosineDecaySchedule`
+    scheduler_warmup_steps: int = 1_000
+    scheduler_decay_steps: int = 30_000
+    scheduler_decay_lr: float = 2.5e-6
+
+    tokenizer_max_length: int = 200  # see openpi `__post_init__`
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        # Validate configuration
+        if self.n_action_steps > self.chunk_size:
+            raise ValueError(
+                f"n_action_steps ({self.n_action_steps}) cannot be greater than chunk_size ({self.chunk_size})"
+            )
+
+        if self.paligemma_variant not in ["gemma_300m", "gemma_2b"]:
+            raise ValueError(f"Invalid paligemma_variant: {self.paligemma_variant}")
+
+        if self.action_expert_variant not in ["gemma_300m", "gemma_2b"]:
+            raise ValueError(f"Invalid action_expert_variant: {self.action_expert_variant}")
+
+        if self.dtype not in ["bfloat16", "float32"]:
+            raise ValueError(f"Invalid dtype: {self.dtype}")
+
+    def validate_features(self) -> None:
+        """Validate and set up input/output features."""
+        # Image features are now handled dynamically through dataset configuration
+        # No need to auto-add hardcoded image keys
+
+        # State and action features are also handled dynamically through dataset configuration
+        # The actual dimensions come from the feature shapes, max dimensions are used for padding only
+        pass
+
+    def get_optimizer_preset(self) -> AdamWConfig:
+        return AdamWConfig(
+            lr=self.optimizer_lr,
+            betas=self.optimizer_betas,
+            eps=self.optimizer_eps,
+            weight_decay=self.optimizer_weight_decay,
+            grad_clip_norm=self.optimizer_grad_clip_norm,
+        )
+
+    def get_scheduler_preset(self):
+        return CosineDecayWithWarmupSchedulerConfig(
+            peak_lr=self.optimizer_lr,
+            decay_lr=self.scheduler_decay_lr,
+            num_warmup_steps=self.scheduler_warmup_steps,
+            num_decay_steps=self.scheduler_decay_steps,
+        )
+
+    @property
+    def observation_delta_indices(self) -> None:
+        return None
+
+    @property
+    def action_delta_indices(self) -> list:
+        return list(range(self.chunk_size))
+
+    @property
+    def reward_delta_indices(self) -> None:
+        return None

src/lerobot/policies/pi05_openpi/transformers_replace/models/gemma/configuration_gemma.py

@@ -0,0 +1,173 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/gemma/modular_gemma.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_gemma.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2024 Google Inc. 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 ...configuration_utils import PretrainedConfig
+
+
+class GemmaConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`GemmaModel`]. It is used to instantiate an Gemma
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a similar configuration to that of the Gemma-7B.
+    e.g. [google/gemma-7b](https://huggingface.co/google/gemma-7b)
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+    Args:
+        vocab_size (`int`, *optional*, defaults to 256000):
+            Vocabulary size of the Gemma model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`GemmaModel`]
+        hidden_size (`int`, *optional*, defaults to 3072):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 24576):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 28):
+            Number of hidden layers in the Transformer decoder.
+        num_attention_heads (`int`, *optional*, defaults to 16):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        num_key_value_heads (`int`, *optional*, defaults to 16):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details, check out [this
+            paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
+            `num_attention_heads`.
+        head_dim (`int`, *optional*, defaults to 256):
+            The attention head dimension.
+        hidden_act (`str` or `function`, *optional*, defaults to `"gelu_pytorch_tanh"`):
+            The legacy activation function. It is overwritten by the `hidden_activation`.
+        hidden_activation (`str` or `function`, *optional*):
+            The non-linear activation function (function or string) in the decoder. Will default to `"gelu_pytorch_tanh"`
+            if not specified. `"gelu_pytorch_tanh"` uses an approximation of the `"gelu"` activation function.
+        max_position_embeddings (`int`, *optional*, defaults to 8192):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        pad_token_id (`int`, *optional*, defaults to 0):
+            Padding token id.
+        eos_token_id (`int`, *optional*, defaults to 1):
+            End of stream token id.
+        bos_token_id (`int`, *optional*, defaults to 2):
+            Beginning of stream token id.
+        tie_word_embeddings (`bool`, *optional*, defaults to `True`):
+            Whether to tie weight embeddings
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        use_adarms (`bool`, *optional*, defaults to `False`):
+            Whether to use ADARMS.
+        adarms_cond_dim (`int`, *optional*, defaults to `None`):
+            The dimension of the ADARMS condition.
+    ```python
+    >>> from transformers import GemmaModel, GemmaConfig
+    >>> # Initializing a Gemma gemma-7b style configuration
+    >>> configuration = GemmaConfig()
+    >>> # Initializing a model from the gemma-7b style configuration
+    >>> model = GemmaModel(configuration)
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "gemma"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=256000,
+        hidden_size=3072,
+        intermediate_size=24576,
+        num_hidden_layers=28,
+        num_attention_heads=16,
+        num_key_value_heads=16,
+        head_dim=256,
+        hidden_act="gelu_pytorch_tanh",
+        hidden_activation=None,
+        max_position_embeddings=8192,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=0,
+        eos_token_id=1,
+        bos_token_id=2,
+        tie_word_embeddings=True,
+        rope_theta=10000.0,
+        attention_bias=False,
+        attention_dropout=0.0,
+        use_adarms: bool = False,
+        adarms_cond_dim: int | None = None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.head_dim = head_dim
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.hidden_activation = hidden_activation
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.use_adarms = use_adarms
+        self.adarms_cond_dim = adarms_cond_dim
+
+        # Set default for adarms_cond_dim if use_adarms is True
+        if self.use_adarms and self.adarms_cond_dim is None:
+            self.adarms_cond_dim = self.hidden_size
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["GemmaConfig"]

src/lerobot/policies/pi05_openpi/transformers_replace/models/gemma/modeling_gemma.py

@@ -0,0 +1,895 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/gemma/modular_gemma.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_gemma.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2024 Google Inc. 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 Callable
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+    SequenceClassifierOutputWithPast,
+    TokenClassifierOutput,
+)
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import LossKwargs, auto_docstring, can_return_tuple, logging
+from .configuration_gemma import GemmaConfig
+
+logger = logging.get_logger(__name__)
+
+
+# Workaround for Python 3.10+ UnionType compatibility with transformers auto_docstring
+def safe_auto_docstring(func=None, **kwargs):
+    """Auto docstring decorator that handles Python 3.10+ UnionType gracefully."""
+
+    def decorator(f):
+        try:
+            return auto_docstring(f, **kwargs) if kwargs else auto_docstring(f)
+        except (AttributeError, TypeError):
+            # If auto_docstring fails due to UnionType, just return the function unchanged
+            return f
+
+    if func is None:
+        # Called with arguments, return the decorator
+        return decorator
+    else:
+        # Called without arguments, apply directly
+        return decorator(func)
+
+
+class GemmaRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6, cond_dim: int | None = None):
+        super().__init__()
+        self.eps = eps
+        self.dim = dim
+        self.cond_dim = cond_dim
+
+        # Dense layer for adaptive normalization (if cond_dim is provided)
+        if cond_dim is not None:
+            # self.dense = nn.Linear(cond_dim, dim * 3, bias=True, dtype=torch.bfloat16)
+            self.dense = nn.Linear(cond_dim, dim * 3, bias=True)
+            # Initialize with zeros (matches source implementation)
+            nn.init.zeros_(self.dense.weight)
+        else:
+            self.weight = nn.Parameter(torch.zeros(dim, dtype=torch.bfloat16))
+            self.dense = None
+
+    def _norm(self, x):
+        # Compute variance in float32 (like the source implementation)
+        var = torch.mean(torch.square(x.float()), dim=-1, keepdim=True)
+        # Compute normalization in float32
+        normed_inputs = x * torch.rsqrt(var + self.eps)
+        return normed_inputs
+
+    def forward(self, x, cond=None):
+        dtype = x.dtype  # original dtype, could be half-precision
+        normed_inputs = self._norm(x)
+
+        if cond is None or self.dense is None:
+            # regular RMSNorm
+            # scale by learned parameter in float32 (matches source implementation)
+            normed_inputs = normed_inputs * (1.0 + self.weight.float())
+            return normed_inputs.to(dtype), None  # return in original dtype with None gate
+
+        # adaptive RMSNorm (if cond is provided and dense layer exists)
+        if cond.shape[-1] != self.cond_dim:
+            raise ValueError(f"Expected cond dimension {self.cond_dim}, got {cond.shape[-1]}")
+
+        # self.dense.to(dtype=torch.bfloat16).to(dtype=torch.float32)
+        modulation = self.dense(cond)
+        # Reshape modulation to broadcast properly: [batch, 1, features] for [batch, seq, features]
+        if len(x.shape) == 3:  # [batch, seq, features]
+            modulation = modulation.unsqueeze(1)
+
+        scale, shift, gate = torch.chunk(modulation, 3, dim=-1)
+
+        # Apply adaptive normalization: use model weight dtype to ensure compatibility
+        # model_dtype = self.dense.weight.dtype  # Use the model's dtype (bfloat16)
+        # scale = scale.to(model_dtype)
+        # shift = shift.to(model_dtype)
+        # gate = gate.to(model_dtype)
+        # normed_inputs = normed_inputs.to(model_dtype)  # Convert normed_inputs to model dtype
+
+        normed_inputs = normed_inputs * (1 + scale.to(torch.float32)) + shift.to(torch.float32)
+
+        return normed_inputs.to(dtype), gate.to(dtype)
+
+    def extra_repr(self):
+        repr_str = f"{tuple(self.weight.shape)}, eps={self.eps}"
+        if self.dense is not None:
+            repr_str += f", adaptive=True, cond_dim={self.cond_dim}"
+        return repr_str
+
+
+class GemmaMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+class GemmaRotaryEmbedding(nn.Module):
+    def __init__(self, config: GemmaConfig, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = (
+            self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        )
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def _gated_residual(x, y, gate):
+    """
+    Applies gated residual connection with optional gate parameter.
+
+    Args:
+        x: Input tensor (residual)
+        y: Output tensor to be added
+        gate: Optional gate tensor to modulate the addition
+
+    Returns:
+        x + y if gate is None, otherwise x + y * gate
+    """
+    if x is None and y is None:
+        return None
+    if x is None or y is None:
+        return x if x is not None else y
+    if gate is None:
+        return x + y
+    return x + y * gate
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class GemmaAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: GemmaConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_value: Cache | None = None,
+        cache_position: torch.LongTensor | None = None,
+        use_cache: bool = False,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        # Use cache if provided
+        if past_key_value is not None:
+            if use_cache:
+                # sin and cos are specific to RoPE models; cache_position needed for the static cache
+                cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+                key_states, value_states = past_key_value.update(
+                    key_states, value_states, self.layer_idx, cache_kwargs
+                )
+            else:
+                key_states = torch.cat([past_key_value[self.layer_idx][0], key_states], dim=2)
+                value_states = torch.cat([past_key_value[self.layer_idx][1], value_states], dim=2)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class GemmaDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: GemmaConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = GemmaAttention(config=config, layer_idx=layer_idx)
+
+        self.mlp = GemmaMLP(config)
+        cond_dim = getattr(config, "adarms_cond_dim", None) if getattr(config, "use_adarms", False) else None
+        self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim)
+        self.post_attention_layernorm = GemmaRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_value: Cache | None = None,
+        output_attentions: bool | None = False,
+        use_cache: bool | None = False,
+        cache_position: torch.LongTensor | None = None,
+        position_embeddings: None
+        | (tuple[torch.Tensor, torch.Tensor]) = None,  # necessary, but kept here for BC
+        adarms_cond: torch.Tensor | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
+        residual = hidden_states
+        hidden_states, gate = self.input_layernorm(hidden_states, adarms_cond)
+
+        # Self Attention
+        hidden_states, self_attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = _gated_residual(residual, hidden_states, gate)
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states, gate = self.post_attention_layernorm(hidden_states, adarms_cond)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = _gated_residual(residual, hidden_states, gate)
+
+        outputs = (hidden_states,)
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        return outputs
+
+
+@safe_auto_docstring
+class GemmaPreTrainedModel(PreTrainedModel):
+    config_class = GemmaConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["GemmaDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_3 = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_cache_class = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, GemmaRMSNorm):
+            if hasattr(module, "weight"):
+                module.weight.data.fill_(1.0)
+
+
+@safe_auto_docstring
+class GemmaModel(GemmaPreTrainedModel):
+    def __init__(self, config: GemmaConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [GemmaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+
+        cond_dim = getattr(config, "adarms_cond_dim", None) if getattr(config, "use_adarms", False) else None
+        self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim)
+        self.rotary_emb = GemmaRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    @can_return_tuple
+    @safe_auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        adarms_cond: torch.Tensor | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> BaseModelOutputWithPast:
+        """
+        adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*):
+            Condition for ADARMS.
+        """
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+            )
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+        # embed positions
+        hidden_states = inputs_embeds
+        # Convert to bfloat16 if the first layer uses bfloat16
+        if len(self.layers) > 0 and self.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(torch.bfloat16)
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # normalized
+        # Gemma downcasts the below to float16, causing sqrt(3072)=55.4256 to become 55.5
+        # See https://github.com/huggingface/transformers/pull/29402
+        _normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype)
+        # hidden_states = hidden_states * normalizer
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                adarms_cond=adarms_cond,
+                **kwargs,
+            )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states, _ = self.norm(hidden_states, adarms_cond)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
+
+
+@safe_auto_docstring
+class GemmaForCausalLM(GemmaPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = GemmaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @can_return_tuple
+    @safe_auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        adarms_cond: torch.Tensor | None = None,
+        **kwargs: Unpack[KwargsForCausalLM],
+    ) -> CausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*):
+            Condition for ADARMS.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, GemmaForCausalLM
+
+        >>> model = GemmaForCausalLM.from_pretrained("google/gemma-7b")
+        >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b")
+
+        >>> prompt = "What is your favorite condiment?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "What is your favorite condiment?"
+        ```"""
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            cache_position=cache_position,
+            adarms_cond=adarms_cond,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@safe_auto_docstring(
+    custom_intro="""
+    The Gemma Model transformer with a sequence classification head on top (linear layer).
+
+    [`GemmaForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+    (e.g. GPT-2) do.
+
+    Since it does classification on the last token, it requires to know the position of the last token. If a
+    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+    each row of the batch).
+    """
+)
+class GemmaForSequenceClassification(GemmaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = GemmaModel(config)
+        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    @can_return_tuple
+    @safe_auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        adarms_cond: torch.Tensor | None = None,
+    ) -> SequenceClassifierOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+
+        adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*):
+            Condition for ADARMS.
+        """
+
+        transformer_outputs: BaseModelOutputWithPast = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            adarms_cond=adarms_cond,
+        )
+        hidden_states = transformer_outputs.last_hidden_state
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        if self.config.pad_token_id is None and batch_size != 1:
+            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+        if self.config.pad_token_id is None:
+            last_non_pad_token = -1
+        elif input_ids is not None:
+            # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
+            non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+            token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+            last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+        else:
+            last_non_pad_token = -1
+            logger.warning_once(
+                f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+                "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+            )
+
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config
+            )
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=pooled_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )
+
+
+@safe_auto_docstring
+class GemmaForTokenClassification(GemmaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = GemmaModel(config)
+        if getattr(config, "classifier_dropout", None) is not None:
+            classifier_dropout = config.classifier_dropout
+        elif getattr(config, "hidden_dropout", None) is not None:
+            classifier_dropout = config.hidden_dropout
+        else:
+            classifier_dropout = 0.1
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.score = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    @can_return_tuple
+    @safe_auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        adarms_cond: torch.Tensor | None = None,
+    ) -> TokenClassifierOutput:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+
+        adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*):
+            Condition for ADARMS.
+        """
+
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            adarms_cond=adarms_cond,
+        )
+        sequence_output = outputs.last_hidden_state
+        sequence_output = self.dropout(sequence_output)
+        logits = self.score(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.config)
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "GemmaModel",
+    "GemmaForCausalLM",
+    "GemmaForSequenceClassification",
+    "GemmaForTokenClassification",
+    "GemmaPreTrainedModel",
+]

src/lerobot/policies/pi05_openpi/transformers_replace/models/paligemma/modeling_paligemma.py

@@ -0,0 +1,666 @@
+# 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.
+"""PyTorch PaliGemmamodel."""
+
+from dataclasses import dataclass
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+
+from ...cache_utils import Cache, HybridCache, StaticCache
+from ...generation import GenerationMixin
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_outputs import BaseModelOutputWithPast
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import (
+    LossKwargs,
+    ModelOutput,
+    auto_docstring,
+    can_return_tuple,
+    is_torchdynamo_compiling,
+    logging,
+)
+from ..auto import AutoModel
+from .configuration_paligemma import PaliGemmaConfig
+
+logger = logging.get_logger(__name__)
+
+
+# Workaround for Python 3.10+ UnionType compatibility with transformers auto_docstring
+def safe_auto_docstring(func=None, **kwargs):
+    """Auto docstring decorator that handles Python 3.10+ UnionType gracefully."""
+
+    def decorator(f):
+        try:
+            return auto_docstring(f, **kwargs) if kwargs else auto_docstring(f)
+        except (AttributeError, TypeError):
+            # If auto_docstring fails due to UnionType, just return the function unchanged
+            return f
+
+    if func is None:
+        # Called with arguments, return the decorator
+        return decorator
+    else:
+        # Called without arguments, apply directly
+        return decorator(func)
+
+
+@dataclass
+@safe_auto_docstring(
+    custom_intro="""
+    Base class for Paligemma outputs, with hidden states and attentions.
+    """
+)
+class PaligemmaModelOutputWithPast(BaseModelOutputWithPast):
+    r"""
+    past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+        `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    image_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+    """
+
+    image_hidden_states: torch.FloatTensor | None = None
+
+
+@dataclass
+@safe_auto_docstring(
+    custom_intro="""
+    Base class for PaliGemma causal language model (or autoregressive) outputs.
+    """
+)
+class PaliGemmaCausalLMOutputWithPast(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+        Language modeling loss (for next-token prediction).
+    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.text_config.vocab_size)`):
+        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+    past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+        `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    image_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+        image_hidden_states of the model produced by the vision encoder after projecting last hidden state.
+    """
+
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    past_key_values: list[torch.FloatTensor] | Cache | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    attentions: tuple[torch.FloatTensor] | None = None
+    image_hidden_states: torch.FloatTensor | None = None
+
+
+class PaliGemmaMultiModalProjector(nn.Module):
+    def __init__(self, config: PaliGemmaConfig):
+        super().__init__()
+        self.linear = nn.Linear(
+            config.vision_config.hidden_size, config.vision_config.projection_dim, bias=True
+        )
+
+    def forward(self, image_features):
+        hidden_states = self.linear(image_features)
+
+        return hidden_states
+
+
+@safe_auto_docstring
+class PaliGemmaPreTrainedModel(PreTrainedModel):
+    config_class = PaliGemmaConfig
+    base_model_prefix = ""
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["PaliGemmaMultiModalProjector"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_cache_class = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module):
+        # important: this ported version of PaliGemmaisn't meant for training from scratch - only
+        # inference and fine-tuning
+        std = getattr(self.config, "initializer_range", self.config.get_text_config().initializer_range)
+
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+
+@safe_auto_docstring(
+    custom_intro="""
+    The Base Paligemma model which consists of a vision backbone and a language model without language modeling head.,
+    """
+)
+class PaliGemmaModel(PaliGemmaPreTrainedModel):
+    _checkpoint_conversion_mapping = {"language_model.model": "language_model"}
+    # we are filtering the logits/labels so we shouldn't divide the loss based on num_items_in_batch
+    accepts_loss_kwargs = False
+
+    def __init__(self, config: PaliGemmaConfig):
+        super().__init__(config)
+        self.vision_tower = AutoModel.from_config(config=config.vision_config)
+        self.multi_modal_projector = PaliGemmaMultiModalProjector(config)
+        self.vocab_size = config.text_config.vocab_size
+
+        language_model = AutoModel.from_config(config=config.text_config)
+        self.language_model = language_model
+
+        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
+        self.post_init()
+
+    # Copied from transformers.models.llava.modeling_llava.LlavaModel.get_input_embeddings with Llava->PaliGemma
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    # Copied from transformers.models.llava.modeling_llava.LlavaModel.set_input_embeddings with Llava->PaliGemma
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def set_decoder(self, decoder):
+        self.language_model = decoder
+
+    def get_decoder(self):
+        return self.language_model
+
+    def _update_causal_mask(
+        self,
+        attention_mask,
+        token_type_ids=None,
+        past_key_values=None,
+        cache_position=None,
+        input_tensor=None,
+        is_training: bool | None = None,
+    ):
+        if self.config.text_config._attn_implementation == "flash_attention_2":
+            if attention_mask is not None and 0.0 in attention_mask:
+                return attention_mask
+            return None
+        is_training = is_training if is_training is not None else self.training
+        using_static_cache = isinstance(past_key_values, StaticCache)
+        min_dtype = torch.finfo(self.dtype).min
+        if input_tensor is None:
+            input_tensor = attention_mask
+
+        inputs_lead_dim, sequence_length = input_tensor.shape[:2]
+        if using_static_cache:
+            target_length = past_key_values.get_max_cache_shape()
+        elif isinstance(past_key_values, HybridCache):
+            target_length = past_key_values.get_max_cache_shape()
+        else:
+            target_length = (
+                attention_mask.shape[-1]
+                if isinstance(attention_mask, torch.Tensor)
+                else cache_position[0] + sequence_length + 1
+            )
+
+        if attention_mask is not None and attention_mask.dim() == 4:
+            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+            return attention_mask
+
+        causal_mask = torch.full(
+            (sequence_length, target_length),
+            fill_value=min_dtype,
+            dtype=self.dtype,
+            device=cache_position.device,
+        )
+        # Causal diagonal mask only if training, otherwise attend to the whole prefix. Training-specific attn for prefix is handled below
+        if sequence_length != 1:
+            if is_training:
+                causal_mask = torch.triu(causal_mask, diagonal=1)
+            else:
+                causal_mask[:, :sequence_length] = 0.0
+
+        causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(
+            -1, 1
+        )
+        causal_mask = causal_mask[None, None, :, :].expand(inputs_lead_dim, 1, -1, -1)
+        if attention_mask is not None:
+            causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
+            mask_length = attention_mask.shape[-1]
+
+            # First unmask prefix tokens during training
+            if is_training:
+                if token_type_ids is None:
+                    raise ValueError("Token type ids must be provided during training")
+                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+                    token_type_ids[:, None, None, :].to(causal_mask.device) == 0, 0
+                )
+
+            # Then apply padding mask (will mask pad tokens)
+            padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+                causal_mask.device
+            )
+            padding_mask = padding_mask == 0
+            causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+                padding_mask, min_dtype
+            )
+
+        return causal_mask
+
+    def get_image_features(self, pixel_values: torch.FloatTensor):
+        """
+        Obtains image last hidden states from the vision tower and apply multimodal projection.
+
+        Args:
+            pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`)
+               The tensors corresponding to the input images.
+        Returns:
+            image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
+        """
+        image_outputs = self.vision_tower(pixel_values)
+        selected_image_feature = image_outputs.last_hidden_state
+        image_features = self.multi_modal_projector(selected_image_feature)
+        return image_features
+
+    @can_return_tuple
+    @safe_auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        pixel_values: torch.FloatTensor = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: list[torch.FloatTensor] | Cache | None = None,
+        token_type_ids: torch.LongTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple | PaligemmaModelOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.text_config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
+
+        >>> model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma2-3b-mix-224")
+        >>> processor = AutoProcessor.from_pretrained("google/paligemma2-3b-mix-224")
+
+        >>> prompt = "Where is the cat standing?"
+        >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+
+        >>> inputs = processor(images=image, text=prompt,  return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(**inputs,)
+        >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Where is the cat standing?\nsnow"
+        ```"""
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        is_training = token_type_ids is not None and labels is not None
+
+        # Replace image id with PAD if the image token if OOV, to avoid index-errors
+        if input_ids is not None and self.config.image_token_id >= self.vocab_size:
+            special_image_mask = input_ids == self.config.image_token_id
+            llm_input_ids = input_ids.clone()
+            llm_input_ids[special_image_mask] = 0
+        else:
+            llm_input_ids = input_ids
+
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(llm_input_ids)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0) + 1  # Paligemma positions are 1-indexed
+
+        # Merge text and images
+        if pixel_values is not None:
+            image_features = self.get_image_features(pixel_values)
+
+            if input_ids is None:
+                special_image_mask = inputs_embeds == self.get_input_embeddings()(
+                    torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+                )
+            else:
+                special_image_mask = (input_ids == self.config.image_token_id).unsqueeze(-1)
+                special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
+
+            if (
+                not is_torchdynamo_compiling()
+                and inputs_embeds[special_image_mask].numel() != image_features.numel()
+            ):
+                image_tokens_in_text = (special_image_mask).sum(dim=1).sum(dim=0)[0]
+                raise ValueError(
+                    f"Number of images does not match number of special image tokens in the input text. "
+                    f"Got {image_tokens_in_text} image tokens in the text but {image_features.shape[0] * image_features.shape[1]} "
+                    "tokens from image embeddings."
+                )
+            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+        causal_mask = self._update_causal_mask(
+            attention_mask, token_type_ids, past_key_values, cache_position, inputs_embeds, is_training
+        )
+        outputs = self.language_model(
+            attention_mask=causal_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        return PaligemmaModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=image_features if pixel_values is not None else None,
+        )
+
+
+class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
+
+
+@safe_auto_docstring(
+    custom_intro="""
+    The Base Paligemma model which consists of a vision backbone and a language model without language modeling head.,
+    """
+)
+class PaliGemmaForConditionalGeneration(PaliGemmaPreTrainedModel, GenerationMixin):
+    _checkpoint_conversion_mapping = {
+        "^language_model.model": "model.language_model",
+        "^vision_tower": "model.vision_tower",
+        "^multi_modal_projector": "model.multi_modal_projector",
+        "^language_model.lm_head": "lm_head",
+    }
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: PaliGemmaConfig):
+        super().__init__(config)
+        self.model = PaliGemmaModel(config)
+        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.model.get_decoder()
+
+    def get_image_features(self, pixel_values):
+        return self.model.get_image_features(pixel_values)
+
+    # Make modules available through conditional class for BC
+    @property
+    def language_model(self):
+        return self.model.language_model
+
+    @property
+    def vision_tower(self):
+        return self.model.vision_tower
+
+    @property
+    def multi_modal_projector(self):
+        return self.model.multi_modal_projector
+
+    @can_return_tuple
+    @safe_auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        pixel_values: torch.FloatTensor = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: list[torch.FloatTensor] | Cache | None = None,
+        token_type_ids: torch.LongTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[KwargsForCausalLM],
+    ) -> tuple | PaliGemmaCausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.text_config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
+
+        >>> model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma2-3b-mix-224")
+        >>> processor = AutoProcessor.from_pretrained("google/paligemma2-3b-mix-224")
+
+        >>> prompt = "Where is the cat standing?"
+        >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+
+        >>> inputs = processor(images=image, text=prompt,  return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(**inputs,)
+        >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Where is the cat standing?\nsnow"
+        ```"""
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            token_type_ids=token_type_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            labels=labels,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+            )
+
+        return PaliGemmaCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=outputs.image_hidden_states,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        pixel_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        use_cache=True,
+        logits_to_keep=None,
+        labels=None,
+        **kwargs,
+    ):
+        # Overwritten -- custom `position_ids` and `pixel_values` handling
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            cache_position=cache_position,
+            use_cache=use_cache,
+            logits_to_keep=logits_to_keep,
+            token_type_ids=token_type_ids,
+            **kwargs,
+        )
+
+        # position_ids in Paligemma are 1-indexed
+        if model_inputs.get("position_ids") is not None:
+            model_inputs["position_ids"] += 1
+        # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
+        # Otherwise we need pixel values to be passed to model. NOTE: use_cache=False needs pixel_values always
+        if cache_position[0] == 0:
+            model_inputs["pixel_values"] = pixel_values
+        is_training = token_type_ids is not None and labels is not None
+        if cache_position[0] == 0 and isinstance(past_key_values, HybridCache):
+            input_tensor = inputs_embeds if inputs_embeds is not None else input_ids
+            causal_mask = self.model._update_causal_mask(
+                attention_mask, token_type_ids, past_key_values, cache_position, input_tensor, is_training
+            )
+            model_inputs["attention_mask"] = causal_mask
+
+        return model_inputs
+
+    @staticmethod
+    # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position
+    def _prepare_4d_causal_attention_mask_with_cache_position(
+        attention_mask: torch.Tensor,
+        sequence_length: int,
+        target_length: int,
+        dtype: torch.dtype,
+        cache_position: torch.Tensor,
+        batch_size: int,
+        **kwargs,
+    ):
+        """
+        Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
+        `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
+
+        Args:
+            attention_mask (`torch.Tensor`):
+                A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape
+                `(batch_size, 1, query_length, key_value_length)`.
+            sequence_length (`int`):
+                The sequence length being processed.
+            target_length (`int`):
+                The target length: when generating with static cache, the mask should be as long as the static cache,
+                to account for the 0 padding, the part of the cache that is not filled yet.
+            dtype (`torch.dtype`):
+                The dtype to use for the 4D attention mask.
+            cache_position (`torch.Tensor`):
+                Indices depicting the position of the input sequence tokens in the sequence.
+            batch_size (`torch.Tensor`):
+                Batch size.
+        """
+        if attention_mask is not None and attention_mask.dim() == 4:
+            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+            causal_mask = attention_mask
+        else:
+            min_dtype = torch.finfo(dtype).min
+            causal_mask = torch.full(
+                (sequence_length, target_length),
+                fill_value=min_dtype,
+                dtype=dtype,
+                device=cache_position.device,
+            )
+            if sequence_length != 1:
+                causal_mask = torch.triu(causal_mask, diagonal=1)
+            causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(
+                -1, 1
+            )
+            causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
+            if attention_mask is not None:
+                causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
+                mask_length = attention_mask.shape[-1]
+                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+                    causal_mask.device
+                )
+                padding_mask = padding_mask == 0
+                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+                    padding_mask, min_dtype
+                )
+
+        return causal_mask
+
+
+__all__ = ["PaliGemmaForConditionalGeneration", "PaliGemmaPreTrainedModel", "PaliGemmaModel"]

src/lerobot/policies/pi05_openpi/transformers_replace/models/siglip/check.py

@@ -0,0 +1,5 @@
+import transformers
+
+
+def check_whether_transformers_replace_is_installed_correctly():
+    return transformers.__version__ == "4.53.2"

src/lerobot/policies/pi0_openpi/README.md

@@ -0,0 +1,92 @@
+# π₀ (pi0)
+
+This repository contains the Hugging Face port of **π₀**, adapted from [OpenPI](https://github.com/Physical-Intelligence/openpi) by the Physical Intelligence.
+It is designed as a **Vision-Language-Action flow model for general robot control**.
+
+---
+
+### ⚠️ WARNING ⚠️
+
+This project requires **patching the Hugging Face `transformers` library**.
+
+1. Make sure you have the exact version installed:
+
+```bash
+   pip show transformers
+```
+
+It must be version **4.53.2**.
+
+2. Apply the custom patches by copying the modified files into your environment:
+
+   ```bash
+   cp -r ./src/lerobot/policies/pi0_openpi/transformers_replace/* \
+     $(python -c "import transformers, os; print(os.path.dirname(transformers.__file__))")
+   ```
+
+   These patches overwrite parts of `transformers` to:
+   - Support the **AdaRMS optimizer**,
+   - Correctly control the precision of activations,
+   - Allow the KV cache to be used without updates.
+
+**Important:**
+
+- This permanently modifies your `transformers` installation.
+- The changes survive reinstalls unless you explicitly remove the patched files or recreate the environment.
+
+To undo and restore a clean state:
+
+```bash
+pip uninstall transformers
+pip install transformers==4.53.2
+```
+
+---
+
+## Model Overview
+
+| Feature              | π₀                                                     | π₀.₅                                      |
+| -------------------- | ------------------------------------------------------ | ----------------------------------------- |
+| State Embedding      | Uses `state_proj` layer                                | No state embedding                        |
+| Time Conditioning    | Concatenates time with actions via `action_time_mlp_*` | Uses `time_mlp_*` for AdaRMS conditioning |
+| AdaRMS               | Not used                                               | Used in action expert                     |
+| Tokenizer Length     | 48 tokens                                              | 200 tokens                                |
+| Discrete State Input | False                                                  | True                                      |
+| Parameter Count      | Higher (includes state embedding)                      | Lower (no state embedding)                |
+
+---
+
+## Citation
+
+If you use this work, please cite both **OpenPI** and the π₀ paper:
+
+```bibtex
+@misc{openpi2024,
+  author       = {Physical Intelligence Lab},
+  title        = {OpenPI: PyTorch Implementation of π0 and π0.5 Policies},
+  year         = {2024},
+  publisher    = {GitHub},
+  howpublished = {\url{https://github.com/Physical-Intelligence/openpi}},
+  license      = {Apache-2.0}
+}
+
+@misc{black2024pi0visionlanguageactionflowmodel,
+  title        = {π₀: A Vision-Language-Action Flow Model for General Robot Control},
+  author       = {Kevin Black and Noah Brown and Danny Driess and Adnan Esmail and Michael Equi and Chelsea Finn and Niccolo Fusai and Lachy Groom and Karol Hausman and Brian Ichter and Szymon Jakubczak and Tim Jones and Liyiming Ke and Sergey Levine and Adrian Li-Bell and Mohith Mothukuri and Suraj Nair and Karl Pertsch and Lucy Xiaoyang Shi and James Tanner and Quan Vuong and Anna Walling and Haohuan Wang and Ury Zhilinsky},
+  year         = {2024},
+  eprint       = {2410.24164},
+  archivePrefix= {arXiv},
+  primaryClass = {cs.LG},
+  url          = {https://arxiv.org/abs/2410.24164},
+}
+```
+
+---
+
+## License
+
+This port follows the **Apache 2.0 License**, consistent with the original [OpenPI repository](https://github.com/Physical-Intelligence/openpi).
+
+```
+
+```

src/lerobot/policies/pi0_openpi/__init__.py

@@ -0,0 +1,20 @@
+#!/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 .configuration_pi0openpi import PI0OpenPIConfig
+from .modeling_pi0openpi import PI0OpenPIPolicy
+
+__all__ = ["PI0OpenPIConfig", "PI0OpenPIPolicy"]

src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py

@@ -0,0 +1,134 @@
+#!/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 dataclasses import dataclass, field
+
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.types import NormalizationMode
+from lerobot.optim.optimizers import AdamWConfig
+from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
+
+
+@PreTrainedConfig.register_subclass("pi0_openpi")
+@dataclass
+class PI0OpenPIConfig(PreTrainedConfig):
+    # Model architecture
+    paligemma_variant: str = "gemma_2b"
+    action_expert_variant: str = "gemma_300m"
+    dtype: str = "float32"  # Options: "bfloat16", "float32"
+
+    # Input / output structure
+    n_obs_steps: int = 1
+    chunk_size: int = 50  # Number of action steps to predict, in openpi called "action_horizon"
+    n_action_steps: int = 50  # Number of action steps to execute
+
+    # Shorter state and action vectors will be padded to these dimensions
+    max_state_dim: int = 32  # State dimension (will be padded to 32)
+    max_action_dim: int = 32  # Action dimension (will be padded to 32)
+
+    # Flow matching parameters: see openpi `PI0Pytorch`
+    num_inference_steps: int = 10  # Number of denoising steps during inference
+    time_sampling_beta_alpha: float = 1.5  # Beta distribution alpha parameter for time sampling
+    time_sampling_beta_beta: float = 1.0  # Beta distribution beta parameter for time sampling
+    min_period: float = 4e-3  # Min period for sinusoidal positional encoding
+    max_period: float = 4.0  # Max period for sinusoidal positional encoding
+
+    # Image preprocessing
+    image_resolution: tuple[int, int] = (224, 224)  # see openpi `preprocessing_pytorch.py`
+
+    # Normalization
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.IDENTITY,  # Images are normalized to [-1, 1] in preprocessing
+            "STATE": NormalizationMode.MEAN_STD,
+            "ACTION": NormalizationMode.MEAN_STD,
+        }
+    )
+
+    # Training settings
+    gradient_checkpointing: bool = False  # Enable gradient checkpointing for memory optimization
+    compile_model: bool = False  # Whether to use torch.compile for model optimization
+    compile_mode: str = "max-autotune"  # Torch compile mode
+    device: str | None = None  # Device to use for the model (None = auto-detect)
+
+    # Optimizer settings: see openpi `AdamW` and
+    optimizer_lr: float = 2.5e-5  # see openpi `CosineDecaySchedule: peak_lr`
+    optimizer_betas: tuple[float, float] = (0.9, 0.95)
+    optimizer_eps: float = 1e-8
+    optimizer_weight_decay: float = 0.01
+    optimizer_grad_clip_norm: float = 1.0
+
+    # Scheduler settings: see openpi `CosineDecaySchedule`
+    scheduler_warmup_steps: int = 1_000
+    scheduler_decay_steps: int = 30_000
+    scheduler_decay_lr: float = 2.5e-6
+
+    tokenizer_max_length: int = 48  # pi0=48, see openpi `__post_init__`
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        # Validate configuration
+        if self.n_action_steps > self.chunk_size:
+            raise ValueError(
+                f"n_action_steps ({self.n_action_steps}) cannot be greater than chunk_size ({self.chunk_size})"
+            )
+
+        if self.paligemma_variant not in ["gemma_300m", "gemma_2b"]:
+            raise ValueError(f"Invalid paligemma_variant: {self.paligemma_variant}")
+
+        if self.action_expert_variant not in ["gemma_300m", "gemma_2b"]:
+            raise ValueError(f"Invalid action_expert_variant: {self.action_expert_variant}")
+
+        if self.dtype not in ["bfloat16", "float32"]:
+            raise ValueError(f"Invalid dtype: {self.dtype}")
+
+    def validate_features(self) -> None:
+        """Validate and set up input/output features."""
+        # Image features are now handled dynamically through dataset configuration
+        # No need to auto-add hardcoded image keys
+
+        # State and action features are also handled dynamically through dataset configuration
+        # The actual dimensions come from the feature shapes, max dimensions are used for padding only
+        pass
+
+    def get_optimizer_preset(self) -> AdamWConfig:
+        return AdamWConfig(
+            lr=self.optimizer_lr,
+            betas=self.optimizer_betas,
+            eps=self.optimizer_eps,
+            weight_decay=self.optimizer_weight_decay,
+            grad_clip_norm=self.optimizer_grad_clip_norm,
+        )
+
+    def get_scheduler_preset(self):
+        return CosineDecayWithWarmupSchedulerConfig(
+            peak_lr=self.optimizer_lr,
+            decay_lr=self.scheduler_decay_lr,
+            num_warmup_steps=self.scheduler_warmup_steps,
+            num_decay_steps=self.scheduler_decay_steps,
+        )
+
+    @property
+    def observation_delta_indices(self) -> None:
+        return None
+
+    @property
+    def action_delta_indices(self) -> list:
+        return list(range(self.chunk_size))
+
+    @property
+    def reward_delta_indices(self) -> None:
+        return None

src/lerobot/policies/pi0_openpi/transformers_replace/models/gemma/configuration_gemma.py

@@ -0,0 +1,173 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/gemma/modular_gemma.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_gemma.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2024 Google Inc. 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 ...configuration_utils import PretrainedConfig
+
+
+class GemmaConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`GemmaModel`]. It is used to instantiate an Gemma
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a similar configuration to that of the Gemma-7B.
+    e.g. [google/gemma-7b](https://huggingface.co/google/gemma-7b)
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+    Args:
+        vocab_size (`int`, *optional*, defaults to 256000):
+            Vocabulary size of the Gemma model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`GemmaModel`]
+        hidden_size (`int`, *optional*, defaults to 3072):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 24576):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 28):
+            Number of hidden layers in the Transformer decoder.
+        num_attention_heads (`int`, *optional*, defaults to 16):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        num_key_value_heads (`int`, *optional*, defaults to 16):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details, check out [this
+            paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
+            `num_attention_heads`.
+        head_dim (`int`, *optional*, defaults to 256):
+            The attention head dimension.
+        hidden_act (`str` or `function`, *optional*, defaults to `"gelu_pytorch_tanh"`):
+            The legacy activation function. It is overwritten by the `hidden_activation`.
+        hidden_activation (`str` or `function`, *optional*):
+            The non-linear activation function (function or string) in the decoder. Will default to `"gelu_pytorch_tanh"`
+            if not specified. `"gelu_pytorch_tanh"` uses an approximation of the `"gelu"` activation function.
+        max_position_embeddings (`int`, *optional*, defaults to 8192):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        pad_token_id (`int`, *optional*, defaults to 0):
+            Padding token id.
+        eos_token_id (`int`, *optional*, defaults to 1):
+            End of stream token id.
+        bos_token_id (`int`, *optional*, defaults to 2):
+            Beginning of stream token id.
+        tie_word_embeddings (`bool`, *optional*, defaults to `True`):
+            Whether to tie weight embeddings
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        use_adarms (`bool`, *optional*, defaults to `False`):
+            Whether to use ADARMS.
+        adarms_cond_dim (`int`, *optional*, defaults to `None`):
+            The dimension of the ADARMS condition.
+    ```python
+    >>> from transformers import GemmaModel, GemmaConfig
+    >>> # Initializing a Gemma gemma-7b style configuration
+    >>> configuration = GemmaConfig()
+    >>> # Initializing a model from the gemma-7b style configuration
+    >>> model = GemmaModel(configuration)
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "gemma"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=256000,
+        hidden_size=3072,
+        intermediate_size=24576,
+        num_hidden_layers=28,
+        num_attention_heads=16,
+        num_key_value_heads=16,
+        head_dim=256,
+        hidden_act="gelu_pytorch_tanh",
+        hidden_activation=None,
+        max_position_embeddings=8192,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=0,
+        eos_token_id=1,
+        bos_token_id=2,
+        tie_word_embeddings=True,
+        rope_theta=10000.0,
+        attention_bias=False,
+        attention_dropout=0.0,
+        use_adarms: bool = False,
+        adarms_cond_dim: int | None = None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.head_dim = head_dim
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.hidden_activation = hidden_activation
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.use_adarms = use_adarms
+        self.adarms_cond_dim = adarms_cond_dim
+
+        # Set default for adarms_cond_dim if use_adarms is True
+        if self.use_adarms and self.adarms_cond_dim is None:
+            self.adarms_cond_dim = self.hidden_size
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["GemmaConfig"]

src/lerobot/policies/pi0_openpi/transformers_replace/models/gemma/modeling_gemma.py

@@ -0,0 +1,895 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/gemma/modular_gemma.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_gemma.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2024 Google Inc. 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 Callable
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+    SequenceClassifierOutputWithPast,
+    TokenClassifierOutput,
+)
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import LossKwargs, auto_docstring, can_return_tuple, logging
+from .configuration_gemma import GemmaConfig
+
+logger = logging.get_logger(__name__)
+
+
+# Workaround for Python 3.10+ UnionType compatibility with transformers auto_docstring
+def safe_auto_docstring(func=None, **kwargs):
+    """Auto docstring decorator that handles Python 3.10+ UnionType gracefully."""
+
+    def decorator(f):
+        try:
+            return auto_docstring(f, **kwargs) if kwargs else auto_docstring(f)
+        except (AttributeError, TypeError):
+            # If auto_docstring fails due to UnionType, just return the function unchanged
+            return f
+
+    if func is None:
+        # Called with arguments, return the decorator
+        return decorator
+    else:
+        # Called without arguments, apply directly
+        return decorator(func)
+
+
+class GemmaRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6, cond_dim: int | None = None):
+        super().__init__()
+        self.eps = eps
+        self.dim = dim
+        self.cond_dim = cond_dim
+
+        # Dense layer for adaptive normalization (if cond_dim is provided)
+        if cond_dim is not None:
+            # self.dense = nn.Linear(cond_dim, dim * 3, bias=True, dtype=torch.bfloat16)
+            self.dense = nn.Linear(cond_dim, dim * 3, bias=True)
+            # Initialize with zeros (matches source implementation)
+            nn.init.zeros_(self.dense.weight)
+        else:
+            self.weight = nn.Parameter(torch.zeros(dim, dtype=torch.bfloat16))
+            self.dense = None
+
+    def _norm(self, x):
+        # Compute variance in float32 (like the source implementation)
+        var = torch.mean(torch.square(x.float()), dim=-1, keepdim=True)
+        # Compute normalization in float32
+        normed_inputs = x * torch.rsqrt(var + self.eps)
+        return normed_inputs
+
+    def forward(self, x, cond=None):
+        dtype = x.dtype  # original dtype, could be half-precision
+        normed_inputs = self._norm(x)
+
+        if cond is None or self.dense is None:
+            # regular RMSNorm
+            # scale by learned parameter in float32 (matches source implementation)
+            normed_inputs = normed_inputs * (1.0 + self.weight.float())
+            return normed_inputs.to(dtype), None  # return in original dtype with None gate
+
+        # adaptive RMSNorm (if cond is provided and dense layer exists)
+        if cond.shape[-1] != self.cond_dim:
+            raise ValueError(f"Expected cond dimension {self.cond_dim}, got {cond.shape[-1]}")
+
+        # self.dense.to(dtype=torch.bfloat16).to(dtype=torch.float32)
+        modulation = self.dense(cond)
+        # Reshape modulation to broadcast properly: [batch, 1, features] for [batch, seq, features]
+        if len(x.shape) == 3:  # [batch, seq, features]
+            modulation = modulation.unsqueeze(1)
+
+        scale, shift, gate = torch.chunk(modulation, 3, dim=-1)
+
+        # Apply adaptive normalization: use model weight dtype to ensure compatibility
+        # model_dtype = self.dense.weight.dtype  # Use the model's dtype (bfloat16)
+        # scale = scale.to(model_dtype)
+        # shift = shift.to(model_dtype)
+        # gate = gate.to(model_dtype)
+        # normed_inputs = normed_inputs.to(model_dtype)  # Convert normed_inputs to model dtype
+
+        normed_inputs = normed_inputs * (1 + scale.to(torch.float32)) + shift.to(torch.float32)
+
+        return normed_inputs.to(dtype), gate.to(dtype)
+
+    def extra_repr(self):
+        repr_str = f"{tuple(self.weight.shape)}, eps={self.eps}"
+        if self.dense is not None:
+            repr_str += f", adaptive=True, cond_dim={self.cond_dim}"
+        return repr_str
+
+
+class GemmaMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+class GemmaRotaryEmbedding(nn.Module):
+    def __init__(self, config: GemmaConfig, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = (
+            self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        )
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def _gated_residual(x, y, gate):
+    """
+    Applies gated residual connection with optional gate parameter.
+
+    Args:
+        x: Input tensor (residual)
+        y: Output tensor to be added
+        gate: Optional gate tensor to modulate the addition
+
+    Returns:
+        x + y if gate is None, otherwise x + y * gate
+    """
+    if x is None and y is None:
+        return None
+    if x is None or y is None:
+        return x if x is not None else y
+    if gate is None:
+        return x + y
+    return x + y * gate
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class GemmaAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: GemmaConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_value: Cache | None = None,
+        cache_position: torch.LongTensor | None = None,
+        use_cache: bool = False,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        # Use cache if provided
+        if past_key_value is not None:
+            if use_cache:
+                # sin and cos are specific to RoPE models; cache_position needed for the static cache
+                cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+                key_states, value_states = past_key_value.update(
+                    key_states, value_states, self.layer_idx, cache_kwargs
+                )
+            else:
+                key_states = torch.cat([past_key_value[self.layer_idx][0], key_states], dim=2)
+                value_states = torch.cat([past_key_value[self.layer_idx][1], value_states], dim=2)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class GemmaDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: GemmaConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = GemmaAttention(config=config, layer_idx=layer_idx)
+
+        self.mlp = GemmaMLP(config)
+        cond_dim = getattr(config, "adarms_cond_dim", None) if getattr(config, "use_adarms", False) else None
+        self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim)
+        self.post_attention_layernorm = GemmaRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_value: Cache | None = None,
+        output_attentions: bool | None = False,
+        use_cache: bool | None = False,
+        cache_position: torch.LongTensor | None = None,
+        position_embeddings: None
+        | (tuple[torch.Tensor, torch.Tensor]) = None,  # necessary, but kept here for BC
+        adarms_cond: torch.Tensor | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
+        residual = hidden_states
+        hidden_states, gate = self.input_layernorm(hidden_states, adarms_cond)
+
+        # Self Attention
+        hidden_states, self_attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = _gated_residual(residual, hidden_states, gate)
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states, gate = self.post_attention_layernorm(hidden_states, adarms_cond)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = _gated_residual(residual, hidden_states, gate)
+
+        outputs = (hidden_states,)
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        return outputs
+
+
+@safe_auto_docstring
+class GemmaPreTrainedModel(PreTrainedModel):
+    config_class = GemmaConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["GemmaDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_3 = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_cache_class = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, GemmaRMSNorm):
+            if hasattr(module, "weight"):
+                module.weight.data.fill_(1.0)
+
+
+@safe_auto_docstring
+class GemmaModel(GemmaPreTrainedModel):
+    def __init__(self, config: GemmaConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [GemmaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+
+        cond_dim = getattr(config, "adarms_cond_dim", None) if getattr(config, "use_adarms", False) else None
+        self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim)
+        self.rotary_emb = GemmaRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    @can_return_tuple
+    @safe_auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        adarms_cond: torch.Tensor | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> BaseModelOutputWithPast:
+        """
+        adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*):
+            Condition for ADARMS.
+        """
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+            )
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+        # embed positions
+        hidden_states = inputs_embeds
+        # Convert to bfloat16 if the first layer uses bfloat16
+        if len(self.layers) > 0 and self.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(torch.bfloat16)
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # normalized
+        # Gemma downcasts the below to float16, causing sqrt(3072)=55.4256 to become 55.5
+        # See https://github.com/huggingface/transformers/pull/29402
+        _normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype)
+        # hidden_states = hidden_states * normalizer
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                adarms_cond=adarms_cond,
+                **kwargs,
+            )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states, _ = self.norm(hidden_states, adarms_cond)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
+
+
+@safe_auto_docstring
+class GemmaForCausalLM(GemmaPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = GemmaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @can_return_tuple
+    @safe_auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        adarms_cond: torch.Tensor | None = None,
+        **kwargs: Unpack[KwargsForCausalLM],
+    ) -> CausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*):
+            Condition for ADARMS.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, GemmaForCausalLM
+
+        >>> model = GemmaForCausalLM.from_pretrained("google/gemma-7b")
+        >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b")
+
+        >>> prompt = "What is your favorite condiment?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "What is your favorite condiment?"
+        ```"""
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            cache_position=cache_position,
+            adarms_cond=adarms_cond,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@safe_auto_docstring(
+    custom_intro="""
+    The Gemma Model transformer with a sequence classification head on top (linear layer).
+
+    [`GemmaForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+    (e.g. GPT-2) do.
+
+    Since it does classification on the last token, it requires to know the position of the last token. If a
+    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+    each row of the batch).
+    """
+)
+class GemmaForSequenceClassification(GemmaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = GemmaModel(config)
+        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    @can_return_tuple
+    @safe_auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        adarms_cond: torch.Tensor | None = None,
+    ) -> SequenceClassifierOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+
+        adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*):
+            Condition for ADARMS.
+        """
+
+        transformer_outputs: BaseModelOutputWithPast = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            adarms_cond=adarms_cond,
+        )
+        hidden_states = transformer_outputs.last_hidden_state
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        if self.config.pad_token_id is None and batch_size != 1:
+            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+        if self.config.pad_token_id is None:
+            last_non_pad_token = -1
+        elif input_ids is not None:
+            # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
+            non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+            token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+            last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+        else:
+            last_non_pad_token = -1
+            logger.warning_once(
+                f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+                "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+            )
+
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config
+            )
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=pooled_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )
+
+
+@safe_auto_docstring
+class GemmaForTokenClassification(GemmaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = GemmaModel(config)
+        if getattr(config, "classifier_dropout", None) is not None:
+            classifier_dropout = config.classifier_dropout
+        elif getattr(config, "hidden_dropout", None) is not None:
+            classifier_dropout = config.hidden_dropout
+        else:
+            classifier_dropout = 0.1
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.score = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    @can_return_tuple
+    @safe_auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        adarms_cond: torch.Tensor | None = None,
+    ) -> TokenClassifierOutput:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+
+        adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*):
+            Condition for ADARMS.
+        """
+
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            adarms_cond=adarms_cond,
+        )
+        sequence_output = outputs.last_hidden_state
+        sequence_output = self.dropout(sequence_output)
+        logits = self.score(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.config)
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "GemmaModel",
+    "GemmaForCausalLM",
+    "GemmaForSequenceClassification",
+    "GemmaForTokenClassification",
+    "GemmaPreTrainedModel",
+]

src/lerobot/policies/pi0_openpi/transformers_replace/models/paligemma/modeling_paligemma.py

@@ -0,0 +1,666 @@
+# 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.
+"""PyTorch PaliGemmamodel."""
+
+from dataclasses import dataclass
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+
+from ...cache_utils import Cache, HybridCache, StaticCache
+from ...generation import GenerationMixin
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_outputs import BaseModelOutputWithPast
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import (
+    LossKwargs,
+    ModelOutput,
+    auto_docstring,
+    can_return_tuple,
+    is_torchdynamo_compiling,
+    logging,
+)
+from ..auto import AutoModel
+from .configuration_paligemma import PaliGemmaConfig
+
+logger = logging.get_logger(__name__)
+
+
+# Workaround for Python 3.10+ UnionType compatibility with transformers auto_docstring
+def safe_auto_docstring(func=None, **kwargs):
+    """Auto docstring decorator that handles Python 3.10+ UnionType gracefully."""
+
+    def decorator(f):
+        try:
+            return auto_docstring(f, **kwargs) if kwargs else auto_docstring(f)
+        except (AttributeError, TypeError):
+            # If auto_docstring fails due to UnionType, just return the function unchanged
+            return f
+
+    if func is None:
+        # Called with arguments, return the decorator
+        return decorator
+    else:
+        # Called without arguments, apply directly
+        return decorator(func)
+
+
+@dataclass
+@safe_auto_docstring(
+    custom_intro="""
+    Base class for Paligemma outputs, with hidden states and attentions.
+    """
+)
+class PaligemmaModelOutputWithPast(BaseModelOutputWithPast):
+    r"""
+    past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+        `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    image_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+    """
+
+    image_hidden_states: torch.FloatTensor | None = None
+
+
+@dataclass
+@safe_auto_docstring(
+    custom_intro="""
+    Base class for PaliGemma causal language model (or autoregressive) outputs.
+    """
+)
+class PaliGemmaCausalLMOutputWithPast(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+        Language modeling loss (for next-token prediction).
+    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.text_config.vocab_size)`):
+        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+    past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+        `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    image_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+        image_hidden_states of the model produced by the vision encoder after projecting last hidden state.
+    """
+
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    past_key_values: list[torch.FloatTensor] | Cache | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    attentions: tuple[torch.FloatTensor] | None = None
+    image_hidden_states: torch.FloatTensor | None = None
+
+
+class PaliGemmaMultiModalProjector(nn.Module):
+    def __init__(self, config: PaliGemmaConfig):
+        super().__init__()
+        self.linear = nn.Linear(
+            config.vision_config.hidden_size, config.vision_config.projection_dim, bias=True
+        )
+
+    def forward(self, image_features):
+        hidden_states = self.linear(image_features)
+
+        return hidden_states
+
+
+@safe_auto_docstring
+class PaliGemmaPreTrainedModel(PreTrainedModel):
+    config_class = PaliGemmaConfig
+    base_model_prefix = ""
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["PaliGemmaMultiModalProjector"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_cache_class = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module):
+        # important: this ported version of PaliGemmaisn't meant for training from scratch - only
+        # inference and fine-tuning
+        std = getattr(self.config, "initializer_range", self.config.get_text_config().initializer_range)
+
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+
+@safe_auto_docstring(
+    custom_intro="""
+    The Base Paligemma model which consists of a vision backbone and a language model without language modeling head.,
+    """
+)
+class PaliGemmaModel(PaliGemmaPreTrainedModel):
+    _checkpoint_conversion_mapping = {"language_model.model": "language_model"}
+    # we are filtering the logits/labels so we shouldn't divide the loss based on num_items_in_batch
+    accepts_loss_kwargs = False
+
+    def __init__(self, config: PaliGemmaConfig):
+        super().__init__(config)
+        self.vision_tower = AutoModel.from_config(config=config.vision_config)
+        self.multi_modal_projector = PaliGemmaMultiModalProjector(config)
+        self.vocab_size = config.text_config.vocab_size
+
+        language_model = AutoModel.from_config(config=config.text_config)
+        self.language_model = language_model
+
+        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
+        self.post_init()
+
+    # Copied from transformers.models.llava.modeling_llava.LlavaModel.get_input_embeddings with Llava->PaliGemma
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    # Copied from transformers.models.llava.modeling_llava.LlavaModel.set_input_embeddings with Llava->PaliGemma
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def set_decoder(self, decoder):
+        self.language_model = decoder
+
+    def get_decoder(self):
+        return self.language_model
+
+    def _update_causal_mask(
+        self,
+        attention_mask,
+        token_type_ids=None,
+        past_key_values=None,
+        cache_position=None,
+        input_tensor=None,
+        is_training: bool | None = None,
+    ):
+        if self.config.text_config._attn_implementation == "flash_attention_2":
+            if attention_mask is not None and 0.0 in attention_mask:
+                return attention_mask
+            return None
+        is_training = is_training if is_training is not None else self.training
+        using_static_cache = isinstance(past_key_values, StaticCache)
+        min_dtype = torch.finfo(self.dtype).min
+        if input_tensor is None:
+            input_tensor = attention_mask
+
+        inputs_lead_dim, sequence_length = input_tensor.shape[:2]
+        if using_static_cache:
+            target_length = past_key_values.get_max_cache_shape()
+        elif isinstance(past_key_values, HybridCache):
+            target_length = past_key_values.get_max_cache_shape()
+        else:
+            target_length = (
+                attention_mask.shape[-1]
+                if isinstance(attention_mask, torch.Tensor)
+                else cache_position[0] + sequence_length + 1
+            )
+
+        if attention_mask is not None and attention_mask.dim() == 4:
+            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+            return attention_mask
+
+        causal_mask = torch.full(
+            (sequence_length, target_length),
+            fill_value=min_dtype,
+            dtype=self.dtype,
+            device=cache_position.device,
+        )
+        # Causal diagonal mask only if training, otherwise attend to the whole prefix. Training-specific attn for prefix is handled below
+        if sequence_length != 1:
+            if is_training:
+                causal_mask = torch.triu(causal_mask, diagonal=1)
+            else:
+                causal_mask[:, :sequence_length] = 0.0
+
+        causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(
+            -1, 1
+        )
+        causal_mask = causal_mask[None, None, :, :].expand(inputs_lead_dim, 1, -1, -1)
+        if attention_mask is not None:
+            causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
+            mask_length = attention_mask.shape[-1]
+
+            # First unmask prefix tokens during training
+            if is_training:
+                if token_type_ids is None:
+                    raise ValueError("Token type ids must be provided during training")
+                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+                    token_type_ids[:, None, None, :].to(causal_mask.device) == 0, 0
+                )
+
+            # Then apply padding mask (will mask pad tokens)
+            padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+                causal_mask.device
+            )
+            padding_mask = padding_mask == 0
+            causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+                padding_mask, min_dtype
+            )
+
+        return causal_mask
+
+    def get_image_features(self, pixel_values: torch.FloatTensor):
+        """
+        Obtains image last hidden states from the vision tower and apply multimodal projection.
+
+        Args:
+            pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`)
+               The tensors corresponding to the input images.
+        Returns:
+            image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
+        """
+        image_outputs = self.vision_tower(pixel_values)
+        selected_image_feature = image_outputs.last_hidden_state
+        image_features = self.multi_modal_projector(selected_image_feature)
+        return image_features
+
+    @can_return_tuple
+    @safe_auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        pixel_values: torch.FloatTensor = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: list[torch.FloatTensor] | Cache | None = None,
+        token_type_ids: torch.LongTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple | PaligemmaModelOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.text_config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
+
+        >>> model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma2-3b-mix-224")
+        >>> processor = AutoProcessor.from_pretrained("google/paligemma2-3b-mix-224")
+
+        >>> prompt = "Where is the cat standing?"
+        >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+
+        >>> inputs = processor(images=image, text=prompt,  return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(**inputs,)
+        >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Where is the cat standing?\nsnow"
+        ```"""
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        is_training = token_type_ids is not None and labels is not None
+
+        # Replace image id with PAD if the image token if OOV, to avoid index-errors
+        if input_ids is not None and self.config.image_token_id >= self.vocab_size:
+            special_image_mask = input_ids == self.config.image_token_id
+            llm_input_ids = input_ids.clone()
+            llm_input_ids[special_image_mask] = 0
+        else:
+            llm_input_ids = input_ids
+
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(llm_input_ids)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0) + 1  # Paligemma positions are 1-indexed
+
+        # Merge text and images
+        if pixel_values is not None:
+            image_features = self.get_image_features(pixel_values)
+
+            if input_ids is None:
+                special_image_mask = inputs_embeds == self.get_input_embeddings()(
+                    torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+                )
+            else:
+                special_image_mask = (input_ids == self.config.image_token_id).unsqueeze(-1)
+                special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
+
+            if (
+                not is_torchdynamo_compiling()
+                and inputs_embeds[special_image_mask].numel() != image_features.numel()
+            ):
+                image_tokens_in_text = (special_image_mask).sum(dim=1).sum(dim=0)[0]
+                raise ValueError(
+                    f"Number of images does not match number of special image tokens in the input text. "
+                    f"Got {image_tokens_in_text} image tokens in the text but {image_features.shape[0] * image_features.shape[1]} "
+                    "tokens from image embeddings."
+                )
+            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+        causal_mask = self._update_causal_mask(
+            attention_mask, token_type_ids, past_key_values, cache_position, inputs_embeds, is_training
+        )
+        outputs = self.language_model(
+            attention_mask=causal_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        return PaligemmaModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=image_features if pixel_values is not None else None,
+        )
+
+
+class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
+
+
+@safe_auto_docstring(
+    custom_intro="""
+    The Base Paligemma model which consists of a vision backbone and a language model without language modeling head.,
+    """
+)
+class PaliGemmaForConditionalGeneration(PaliGemmaPreTrainedModel, GenerationMixin):
+    _checkpoint_conversion_mapping = {
+        "^language_model.model": "model.language_model",
+        "^vision_tower": "model.vision_tower",
+        "^multi_modal_projector": "model.multi_modal_projector",
+        "^language_model.lm_head": "lm_head",
+    }
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: PaliGemmaConfig):
+        super().__init__(config)
+        self.model = PaliGemmaModel(config)
+        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.model.get_decoder()
+
+    def get_image_features(self, pixel_values):
+        return self.model.get_image_features(pixel_values)
+
+    # Make modules available through conditional class for BC
+    @property
+    def language_model(self):
+        return self.model.language_model
+
+    @property
+    def vision_tower(self):
+        return self.model.vision_tower
+
+    @property
+    def multi_modal_projector(self):
+        return self.model.multi_modal_projector
+
+    @can_return_tuple
+    @safe_auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        pixel_values: torch.FloatTensor = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: list[torch.FloatTensor] | Cache | None = None,
+        token_type_ids: torch.LongTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[KwargsForCausalLM],
+    ) -> tuple | PaliGemmaCausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.text_config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
+
+        >>> model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma2-3b-mix-224")
+        >>> processor = AutoProcessor.from_pretrained("google/paligemma2-3b-mix-224")
+
+        >>> prompt = "Where is the cat standing?"
+        >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+
+        >>> inputs = processor(images=image, text=prompt,  return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(**inputs,)
+        >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Where is the cat standing?\nsnow"
+        ```"""
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            token_type_ids=token_type_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            labels=labels,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+            )
+
+        return PaliGemmaCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=outputs.image_hidden_states,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        pixel_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        use_cache=True,
+        logits_to_keep=None,
+        labels=None,
+        **kwargs,
+    ):
+        # Overwritten -- custom `position_ids` and `pixel_values` handling
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            cache_position=cache_position,
+            use_cache=use_cache,
+            logits_to_keep=logits_to_keep,
+            token_type_ids=token_type_ids,
+            **kwargs,
+        )
+
+        # position_ids in Paligemma are 1-indexed
+        if model_inputs.get("position_ids") is not None:
+            model_inputs["position_ids"] += 1
+        # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
+        # Otherwise we need pixel values to be passed to model. NOTE: use_cache=False needs pixel_values always
+        if cache_position[0] == 0:
+            model_inputs["pixel_values"] = pixel_values
+        is_training = token_type_ids is not None and labels is not None
+        if cache_position[0] == 0 and isinstance(past_key_values, HybridCache):
+            input_tensor = inputs_embeds if inputs_embeds is not None else input_ids
+            causal_mask = self.model._update_causal_mask(
+                attention_mask, token_type_ids, past_key_values, cache_position, input_tensor, is_training
+            )
+            model_inputs["attention_mask"] = causal_mask
+
+        return model_inputs
+
+    @staticmethod
+    # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position
+    def _prepare_4d_causal_attention_mask_with_cache_position(
+        attention_mask: torch.Tensor,
+        sequence_length: int,
+        target_length: int,
+        dtype: torch.dtype,
+        cache_position: torch.Tensor,
+        batch_size: int,
+        **kwargs,
+    ):
+        """
+        Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
+        `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
+
+        Args:
+            attention_mask (`torch.Tensor`):
+                A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape
+                `(batch_size, 1, query_length, key_value_length)`.
+            sequence_length (`int`):
+                The sequence length being processed.
+            target_length (`int`):
+                The target length: when generating with static cache, the mask should be as long as the static cache,
+                to account for the 0 padding, the part of the cache that is not filled yet.
+            dtype (`torch.dtype`):
+                The dtype to use for the 4D attention mask.
+            cache_position (`torch.Tensor`):
+                Indices depicting the position of the input sequence tokens in the sequence.
+            batch_size (`torch.Tensor`):
+                Batch size.
+        """
+        if attention_mask is not None and attention_mask.dim() == 4:
+            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+            causal_mask = attention_mask
+        else:
+            min_dtype = torch.finfo(dtype).min
+            causal_mask = torch.full(
+                (sequence_length, target_length),
+                fill_value=min_dtype,
+                dtype=dtype,
+                device=cache_position.device,
+            )
+            if sequence_length != 1:
+                causal_mask = torch.triu(causal_mask, diagonal=1)
+            causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(
+                -1, 1
+            )
+            causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
+            if attention_mask is not None:
+                causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
+                mask_length = attention_mask.shape[-1]
+                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+                    causal_mask.device
+                )
+                padding_mask = padding_mask == 0
+                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+                    padding_mask, min_dtype
+                )
+
+        return causal_mask
+
+
+__all__ = ["PaliGemmaForConditionalGeneration", "PaliGemmaPreTrainedModel", "PaliGemmaModel"]

src/lerobot/policies/pi0_openpi/transformers_replace/models/siglip/check.py

@@ -0,0 +1,5 @@
+import transformers
+
+
+def check_whether_transformers_replace_is_installed_correctly():
+    return transformers.__version__ == "4.53.2"

src/lerobot/record.py

@@ -73,6 +73,7 @@ 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 build_dataset_frame, hw_to_dataset_features
+from lerobot.datasets.video_utils import VideoEncodingManager
 from lerobot.policies.factory import make_policy
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.robots import (  # noqa: F401
@@ -208,7 +209,14 @@ def record_loop(
             (
                 t
                 for t in teleop
-                if isinstance(t, (so100_leader.SO100Leader, so101_leader.SO101Leader, koch_leader.KochLeader))
+                if isinstance(
+                    t,
+                    (
+                        so100_leader.SO100Leader,
+                        so101_leader.SO101Leader,
+                        koch_leader.KochLeader,
+                    ),
+                )
             ),
             None,
         )
@@ -271,8 +279,8 @@ def record_loop(
 
         if dataset is not None:
             action_frame = build_dataset_frame(dataset.features, sent_action, prefix="action")
-            frame = {**observation_frame, **action_frame, "task": single_task}
-            dataset.add_frame(frame)
+            frame = {**observation_frame, **action_frame}
+            dataset.add_frame(frame, task=single_task)
 
         if display_data:
             log_rerun_data(observation, action)
@@ -301,6 +309,7 @@ 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:
@@ -321,6 +330,7 @@ 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
@@ -332,46 +342,47 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
 
     listener, events = init_keyboard_listener()
 
-    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)
+    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)
             record_loop(
                 robot=robot,
                 events=events,
                 fps=cfg.dataset.fps,
                 teleop=teleop,
-                control_time_s=cfg.dataset.reset_time_s,
+                policy=policy,
+                dataset=dataset,
+                control_time_s=cfg.dataset.episode_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
+            # 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,
+                )
 
-        dataset.save_episode()
-        recorded_episodes += 1
+            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
 
     log_say("Stop recording", cfg.play_sounds, blocking=True)
 

src/lerobot/replay.py

@@ -55,6 +55,7 @@ from lerobot.robots import (  # noqa: F401
     hope_jr,
     koch_follower,
     make_robot_from_config,
+    reachy2,
     so100_follower,
     so101_follower,
 )

src/lerobot/robots/bi_so100_follower/config_bi_so100_follower.py

@@ -29,10 +29,10 @@ class BiSO100FollowerConfig(RobotConfig):
 
     # Optional
     left_arm_disable_torque_on_disconnect: bool = True
-    left_arm_max_relative_target: int | None = None
+    left_arm_max_relative_target: float | dict[str, float] | None = None
     left_arm_use_degrees: bool = False
     right_arm_disable_torque_on_disconnect: bool = True
-    right_arm_max_relative_target: int | None = None
+    right_arm_max_relative_target: float | dict[str, float] | None = None
     right_arm_use_degrees: bool = False
 
     # cameras (shared between both arms)

src/lerobot/robots/hope_jr/config_hope_jr.py

@@ -44,8 +44,8 @@ class HopeJrArmConfig(RobotConfig):
     disable_torque_on_disconnect: bool = True
 
     # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
-    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
-    # the number of motors in your follower arms.
-    max_relative_target: int | None = None
+    # Set this to a positive scalar to have the same value for all motors, or a dictionary that maps motor
+    # names to the max_relative_target value for that motor.
+    max_relative_target: float | dict[str, float] | None = None
 
     cameras: dict[str, CameraConfig] = field(default_factory=dict)

src/lerobot/robots/koch_follower/config_koch_follower.py

@@ -28,9 +28,9 @@ class KochFollowerConfig(RobotConfig):
     disable_torque_on_disconnect: bool = True
 
     # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
-    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
-    # the number of motors in your follower arms.
-    max_relative_target: int | None = None
+    # Set this to a positive scalar to have the same value for all motors, or a dictionary that maps motor
+    # names to the max_relative_target value for that motor.
+    max_relative_target: float | dict[str, float] | None = None
 
     # cameras
     cameras: dict[str, CameraConfig] = field(default_factory=dict)

src/lerobot/robots/koch_follower/koch_follower.py

@@ -110,6 +110,7 @@ class KochFollower(Robot):
         return self.bus.is_calibrated
 
     def calibrate(self) -> None:
+        self.bus.disable_torque()
         if self.calibration:
             # Calibration file exists, ask user whether to use it or run new calibration
             user_input = input(
@@ -120,7 +121,6 @@ class KochFollower(Robot):
                 self.bus.write_calibration(self.calibration)
                 return
         logger.info(f"\nRunning calibration of {self}")
-        self.bus.disable_torque()
         for motor in self.bus.motors:
             self.bus.write("Operating_Mode", motor, OperatingMode.EXTENDED_POSITION.value)
 

src/lerobot/robots/lekiwi/config_lekiwi.py

@@ -39,9 +39,9 @@ class LeKiwiConfig(RobotConfig):
     disable_torque_on_disconnect: bool = True
 
     # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
-    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
-    # the number of motors in your follower arms.
-    max_relative_target: int | None = None
+    # Set this to a positive scalar to have the same value for all motors, or a dictionary that maps motor
+    # names to the max_relative_target value for that motor.
+    max_relative_target: float | dict[str, float] | None = None
 
     cameras: dict[str, CameraConfig] = field(default_factory=lekiwi_cameras_config)
 

src/lerobot/robots/reachy2/__init__.py

@@ -0,0 +1,25 @@
+#!/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 .configuration_reachy2 import Reachy2RobotConfig
+from .robot_reachy2 import (
+    REACHY2_ANTENNAS_JOINTS,
+    REACHY2_L_ARM_JOINTS,
+    REACHY2_NECK_JOINTS,
+    REACHY2_R_ARM_JOINTS,
+    REACHY2_VEL,
+    Reachy2Robot,
+)

src/lerobot/robots/reachy2/configuration_reachy2.py

@@ -0,0 +1,107 @@
+# 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.
+
+from dataclasses import dataclass, field
+
+from lerobot.cameras import CameraConfig
+from lerobot.cameras.configs import ColorMode
+from lerobot.cameras.reachy2_camera import Reachy2CameraConfig
+
+from ..config import RobotConfig
+
+
+@RobotConfig.register_subclass("reachy2")
+@dataclass
+class Reachy2RobotConfig(RobotConfig):
+    # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
+    # Set this to a positive scalar to have the same value for all motors.
+    max_relative_target: float | None = None
+
+    # IP address of the Reachy 2 robot
+    ip_address: str | None = "localhost"
+
+    # If True, turn_off_smoothly() will be sent to the robot before disconnecting.
+    disable_torque_on_disconnect: bool = False
+
+    # Tag for external commands control
+    # Set to True if you use an external commands system to control the robot,
+    # such as the official teleoperation application: https://github.com/pollen-robotics/Reachy2Teleoperation
+    # If True, robot.send_action() will not send commands to the robot.
+    use_external_commands: bool = False
+
+    # Robot parts
+    # Set to False to not add the corresponding joints part to the robot list of joints.
+    # By default, all parts are set to True.
+    with_mobile_base: bool = True
+    with_l_arm: bool = True
+    with_r_arm: bool = True
+    with_neck: bool = True
+    with_antennas: bool = True
+
+    # Robot cameras
+    # Set to True if you want to use the corresponding cameras in the observations.
+    # By default, only the teleop cameras are used.
+    with_left_teleop_camera: bool = True
+    with_right_teleop_camera: bool = True
+    with_torso_camera: bool = False
+
+    cameras: dict[str, CameraConfig] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        # Add cameras with same ip_address as the robot
+        if self.with_left_teleop_camera:
+            self.cameras["teleop_left"] = Reachy2CameraConfig(
+                name="teleop",
+                image_type="left",
+                ip_address=self.ip_address,
+                fps=15,
+                width=640,
+                height=480,
+                color_mode=ColorMode.RGB,
+            )
+        if self.with_right_teleop_camera:
+            self.cameras["teleop_right"] = Reachy2CameraConfig(
+                name="teleop",
+                image_type="right",
+                ip_address=self.ip_address,
+                fps=15,
+                width=640,
+                height=480,
+                color_mode=ColorMode.RGB,
+            )
+        if self.with_torso_camera:
+            self.cameras["torso_rgb"] = Reachy2CameraConfig(
+                name="depth",
+                image_type="rgb",
+                ip_address=self.ip_address,
+                fps=15,
+                width=640,
+                height=480,
+                color_mode=ColorMode.RGB,
+            )
+
+        super().__post_init__()
+
+        if not (
+            self.with_mobile_base
+            or self.with_l_arm
+            or self.with_r_arm
+            or self.with_neck
+            or self.with_antennas
+        ):
+            raise ValueError(
+                "No Reachy2Robot part used.\n"
+                "At least one part of the robot must be set to True "
+                "(with_mobile_base, with_l_arm, with_r_arm, with_neck, with_antennas)"
+            )

src/lerobot/robots/reachy2/robot_reachy2.py

@@ -0,0 +1,230 @@
+#!/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 time
+from typing import Any
+
+import numpy as np
+from reachy2_sdk import ReachySDK
+
+from lerobot.cameras.utils import make_cameras_from_configs
+
+from ..robot import Robot
+from ..utils import ensure_safe_goal_position
+from .configuration_reachy2 import Reachy2RobotConfig
+
+# {lerobot_keys: reachy2_sdk_keys}
+REACHY2_NECK_JOINTS = {
+    "neck_yaw.pos": "head.neck.yaw",
+    "neck_pitch.pos": "head.neck.pitch",
+    "neck_roll.pos": "head.neck.roll",
+}
+
+REACHY2_ANTENNAS_JOINTS = {
+    "l_antenna.pos": "head.l_antenna",
+    "r_antenna.pos": "head.r_antenna",
+}
+
+REACHY2_R_ARM_JOINTS = {
+    "r_shoulder_pitch.pos": "r_arm.shoulder.pitch",
+    "r_shoulder_roll.pos": "r_arm.shoulder.roll",
+    "r_elbow_yaw.pos": "r_arm.elbow.yaw",
+    "r_elbow_pitch.pos": "r_arm.elbow.pitch",
+    "r_wrist_roll.pos": "r_arm.wrist.roll",
+    "r_wrist_pitch.pos": "r_arm.wrist.pitch",
+    "r_wrist_yaw.pos": "r_arm.wrist.yaw",
+    "r_gripper.pos": "r_arm.gripper",
+}
+
+REACHY2_L_ARM_JOINTS = {
+    "l_shoulder_pitch.pos": "l_arm.shoulder.pitch",
+    "l_shoulder_roll.pos": "l_arm.shoulder.roll",
+    "l_elbow_yaw.pos": "l_arm.elbow.yaw",
+    "l_elbow_pitch.pos": "l_arm.elbow.pitch",
+    "l_wrist_roll.pos": "l_arm.wrist.roll",
+    "l_wrist_pitch.pos": "l_arm.wrist.pitch",
+    "l_wrist_yaw.pos": "l_arm.wrist.yaw",
+    "l_gripper.pos": "l_arm.gripper",
+}
+
+REACHY2_VEL = {
+    "mobile_base.vx": "vx",
+    "mobile_base.vy": "vy",
+    "mobile_base.vtheta": "vtheta",
+}
+
+
+class Reachy2Robot(Robot):
+    """
+    [Reachy 2](https://www.pollen-robotics.com/reachy/), by Pollen Robotics.
+    """
+
+    config_class = Reachy2RobotConfig
+    name = "reachy2"
+
+    def __init__(self, config: Reachy2RobotConfig):
+        super().__init__(config)
+
+        self.config = config
+        self.robot_type = self.config.type
+        self.use_external_commands = self.config.use_external_commands
+
+        self.reachy: None | ReachySDK = None
+        self.cameras = make_cameras_from_configs(config.cameras)
+
+        self.logs: dict[str, float] = {}
+
+        self.joints_dict: dict[str, str] = self._generate_joints_dict()
+
+    @property
+    def observation_features(self) -> dict[str, Any]:
+        return {**self.motors_features, **self.camera_features}
+
+    @property
+    def action_features(self) -> dict[str, type]:
+        return self.motors_features
+
+    @property
+    def camera_features(self) -> dict[str, tuple[int | None, int | None, int]]:
+        return {cam: (self.cameras[cam].height, self.cameras[cam].width, 3) for cam in self.cameras}
+
+    @property
+    def motors_features(self) -> dict[str, type]:
+        if self.config.with_mobile_base:
+            return {
+                **dict.fromkeys(
+                    self.joints_dict.keys(),
+                    float,
+                ),
+                **dict.fromkeys(
+                    REACHY2_VEL.keys(),
+                    float,
+                ),
+            }
+        else:
+            return dict.fromkeys(self.joints_dict.keys(), float)
+
+    @property
+    def is_connected(self) -> bool:
+        return self.reachy.is_connected() if self.reachy is not None else False
+
+    def connect(self, calibrate: bool = False) -> None:
+        self.reachy = ReachySDK(self.config.ip_address)
+        if not self.is_connected:
+            raise ConnectionError()
+
+        for cam in self.cameras.values():
+            cam.connect()
+
+        self.configure()
+
+    def configure(self) -> None:
+        if self.reachy is not None:
+            self.reachy.turn_on()
+            self.reachy.reset_default_limits()
+
+    @property
+    def is_calibrated(self) -> bool:
+        return True
+
+    def calibrate(self) -> None:
+        pass
+
+    def _generate_joints_dict(self) -> dict[str, str]:
+        joints = {}
+        if self.config.with_neck:
+            joints.update(REACHY2_NECK_JOINTS)
+        if self.config.with_l_arm:
+            joints.update(REACHY2_L_ARM_JOINTS)
+        if self.config.with_r_arm:
+            joints.update(REACHY2_R_ARM_JOINTS)
+        if self.config.with_antennas:
+            joints.update(REACHY2_ANTENNAS_JOINTS)
+        return joints
+
+    def _get_state(self) -> dict[str, float]:
+        if self.reachy is not None:
+            pos_dict = {k: self.reachy.joints[v].present_position for k, v in self.joints_dict.items()}
+            if not self.config.with_mobile_base:
+                return pos_dict
+            vel_dict = {k: self.reachy.mobile_base.odometry[v] for k, v in REACHY2_VEL.items()}
+            return {**pos_dict, **vel_dict}
+        else:
+            return {}
+
+    def get_observation(self) -> dict[str, np.ndarray]:
+        obs_dict: dict[str, Any] = {}
+
+        # Read Reachy 2 state
+        before_read_t = time.perf_counter()
+        obs_dict.update(self._get_state())
+        self.logs["read_pos_dt_s"] = time.perf_counter() - before_read_t
+
+        # Capture images from cameras
+        for cam_key, cam in self.cameras.items():
+            obs_dict[cam_key] = cam.async_read()
+
+        return obs_dict
+
+    def send_action(self, action: dict[str, Any]) -> dict[str, Any]:
+        if self.reachy is not None:
+            if not self.is_connected:
+                raise ConnectionError()
+
+            before_write_t = time.perf_counter()
+
+            vel = {}
+            goal_pos = {}
+            for key, val in action.items():
+                if key not in self.joints_dict:
+                    if key not in REACHY2_VEL:
+                        raise KeyError(f"Key '{key}' is not a valid motor key in Reachy 2.")
+                    else:
+                        vel[REACHY2_VEL[key]] = float(val)
+                else:
+                    if not self.use_external_commands and self.config.max_relative_target is not None:
+                        goal_pos[key] = float(val)
+                        goal_present_pos = {
+                            key: (
+                                goal_pos[key],
+                                self.reachy.joints[self.joints_dict[key]].present_position,
+                            )
+                        }
+                        safe_goal_pos = ensure_safe_goal_position(
+                            goal_present_pos, float(self.config.max_relative_target)
+                        )
+                        val = safe_goal_pos[key]
+                    self.reachy.joints[self.joints_dict[key]].goal_position = float(val)
+
+            if self.config.with_mobile_base:
+                self.reachy.mobile_base.set_goal_speed(vel["vx"], vel["vy"], vel["vtheta"])
+
+            # We don't send the goal positions if we control Reachy 2 externally
+            if not self.use_external_commands:
+                self.reachy.send_goal_positions()
+                if self.config.with_mobile_base:
+                    self.reachy.mobile_base.send_speed_command()
+
+            self.logs["write_pos_dt_s"] = time.perf_counter() - before_write_t
+        return action
+
+    def disconnect(self) -> None:
+        if self.reachy is not None:
+            for cam in self.cameras.values():
+                cam.disconnect()
+            if self.config.disable_torque_on_disconnect:
+                self.reachy.turn_off_smoothly()
+            self.reachy.disconnect()

src/lerobot/robots/so100_follower/config_so100_follower.py

@@ -30,9 +30,9 @@ class SO100FollowerConfig(RobotConfig):
     disable_torque_on_disconnect: bool = True
 
     # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
-    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
-    # the number of motors in your follower arms.
-    max_relative_target: int | None = None
+    # Set this to a positive scalar to have the same value for all motors, or a dictionary that maps motor
+    # names to the max_relative_target value for that motor.
+    max_relative_target: float | dict[str, float] | None = None
 
     # cameras
     cameras: dict[str, CameraConfig] = field(default_factory=dict)

src/lerobot/robots/so100_follower/so100_follower.py

@@ -161,6 +161,11 @@ class SO100Follower(Robot):
                 self.bus.write("I_Coefficient", motor, 0)
                 self.bus.write("D_Coefficient", motor, 32)
 
+                if motor == "gripper":
+                    self.bus.write("Max_Torque_Limit", motor, 500)  # 50% of max torque to avoid burnout
+                    self.bus.write("Protection_Current", motor, 250)  # 50% of max current to avoid burnout
+                    self.bus.write("Overload_Torque", motor, 25)  # 25% torque when overloaded
+
     def setup_motors(self) -> None:
         for motor in reversed(self.bus.motors):
             input(f"Connect the controller board to the '{motor}' motor only and press enter.")

src/lerobot/robots/so101_follower/config_so101_follower.py

@@ -30,9 +30,9 @@ class SO101FollowerConfig(RobotConfig):
     disable_torque_on_disconnect: bool = True
 
     # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
-    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
-    # the number of motors in your follower arms.
-    max_relative_target: int | None = None
+    # Set this to a positive scalar to have the same value for all motors, or a dictionary that maps motor
+    # names to the max_relative_target value for that motor.
+    max_relative_target: float | dict[str, float] | None = None
 
     # cameras
     cameras: dict[str, CameraConfig] = field(default_factory=dict)

src/lerobot/robots/so101_follower/so101_follower.py

@@ -157,6 +157,13 @@ class SO101Follower(Robot):
                 self.bus.write("I_Coefficient", motor, 0)
                 self.bus.write("D_Coefficient", motor, 32)
 
+                if motor == "gripper":
+                    self.bus.write(
+                        "Max_Torque_Limit", motor, 500
+                    )  # 50% of the max torque limit to avoid burnout
+                    self.bus.write("Protection_Current", motor, 250)  # 50% of max current to avoid burnout
+                    self.bus.write("Overload_Torque", motor, 25)  # 25% torque when overloaded
+
     def setup_motors(self) -> None:
         for motor in reversed(self.bus.motors):
             input(f"Connect the controller board to the '{motor}' motor only and press enter.")

src/lerobot/robots/stretch3/configuration_stretch3.py

@@ -24,11 +24,6 @@ from ..config import RobotConfig
 @RobotConfig.register_subclass("stretch3")
 @dataclass
 class Stretch3RobotConfig(RobotConfig):
-    # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
-    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
-    # the number of motors in your follower arms.
-    max_relative_target: int | None = None
-
     # cameras
     cameras: dict[str, CameraConfig] = field(
         default_factory=lambda: {

src/lerobot/robots/utils.py

@@ -61,6 +61,10 @@ def make_robot_from_config(config: RobotConfig) -> Robot:
         from .bi_so100_follower import BiSO100Follower
 
         return BiSO100Follower(config)
+    elif config.type == "reachy2":
+        from .reachy2 import Reachy2Robot
+
+        return Reachy2Robot(config)
     elif config.type == "mock_robot":
         from tests.mocks.mock_robot import MockRobot
 
@@ -70,7 +74,7 @@ def make_robot_from_config(config: RobotConfig) -> Robot:
 
 
 def ensure_safe_goal_position(
-    goal_present_pos: dict[str, tuple[float, float]], max_relative_target: float | dict[float]
+    goal_present_pos: dict[str, tuple[float, float]], max_relative_target: float | dict[str, float]
 ) -> dict[str, float]:
     """Caps relative action target magnitude for safety."""
 

src/lerobot/robots/viperx/config_viperx.py

@@ -28,15 +28,15 @@ class ViperXConfig(RobotConfig):
 
     # /!\ FOR SAFETY, READ THIS /!\
     # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
-    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
-    # the number of motors in your follower arms.
+    # Set this to a positive scalar to have the same value for all motors, or a dictionary that maps motor
+    # names to the max_relative_target value for that motor.
     # For Aloha, for every goal position request, motor rotations are capped at 5 degrees by default.
     # When you feel more confident with teleoperation or running the policy, you can extend
     # this safety limit and even removing it by setting it to `null`.
     # Also, everything is expected to work safely out-of-the-box, but we highly advise to
     # first try to teleoperate the grippers only (by commenting out the rest of the motors in this yaml),
     # then to gradually add more motors (by uncommenting), until you can teleoperate both arms fully
-    max_relative_target: int | None = 5
+    max_relative_target: float | dict[str, float] = 5.0
 
     # cameras
     cameras: dict[str, CameraConfig] = field(default_factory=dict)

src/lerobot/scripts/eval.py

@@ -46,16 +46,19 @@ Note that in both examples, the repo/folder should contain at least `config.json
 You can learn about the CLI options for this script in the `EvalPipelineConfig` in lerobot/configs/eval.py
 """
 
+import concurrent.futures as cf
 import json
 import logging
 import threading
 import time
-from collections.abc import Callable
+from collections import defaultdict
+from collections.abc import Callable, Iterator
 from contextlib import nullcontext
 from copy import deepcopy
 from dataclasses import asdict
 from pathlib import Path
 from pprint import pformat
+from typing import TypedDict
 
 import einops
 import gymnasium as gym
@@ -68,7 +71,11 @@ from tqdm import trange
 from lerobot.configs import parser
 from lerobot.configs.eval import EvalPipelineConfig
 from lerobot.envs.factory import make_env
-from lerobot.envs.utils import add_envs_task, check_env_attributes_and_types, preprocess_observation
+from lerobot.envs.utils import (
+    add_envs_task,
+    check_env_attributes_and_types,
+    preprocess_observation,
+)
 from lerobot.policies.factory import make_policy
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import get_device_from_parameters
@@ -145,7 +152,7 @@ def rollout(
         leave=False,
     )
     check_env_attributes_and_types(env)
-    while not np.all(done):
+    while not np.all(done) and step < max_steps:
         # Numpy array to tensor and changing dictionary keys to LeRobot policy format.
         observation = preprocess_observation(observation)
         if return_observations:
@@ -158,10 +165,8 @@ def rollout(
         # Infer "task" from attributes of environments.
         # TODO: works with SyncVectorEnv but not AsyncVectorEnv
         observation = add_envs_task(env, observation)
-
         with torch.inference_mode():
             action = policy.select_action(observation)
-
         # Convert to CPU / numpy.
         action = action.to("cpu").numpy()
         assert action.ndim == 2, "Action dimensions should be (batch, action_dim)"
@@ -179,7 +184,12 @@ def rollout(
             successes = [False] * env.num_envs
 
         # Keep track of which environments are done so far.
+        # Mark the episode as done if we reach the maximum step limit.
+        # This ensures that the rollout always terminates cleanly at `max_steps`,
+        # and allows logging/saving (e.g., videos) to be triggered consistently.
         done = terminated | truncated | done
+        if step + 1 == max_steps:
+            done = np.ones_like(done, dtype=bool)
 
         all_actions.append(torch.from_numpy(action))
         all_rewards.append(torch.from_numpy(reward))
@@ -402,7 +412,6 @@ def eval_policy(
             "eval_ep_s": (time.time() - start) / n_episodes,
         },
     }
-
     if return_episode_data:
         info["episodes"] = episode_data
 
@@ -463,7 +472,9 @@ def eval_main(cfg: EvalPipelineConfig):
 
     # Check device is available
     device = get_safe_torch_device(cfg.policy.device, log=True)
+    # login to hf
 
+    # login()
     torch.backends.cudnn.benchmark = True
     torch.backends.cuda.matmul.allow_tf32 = True
     set_seed(cfg.seed)
@@ -471,40 +482,212 @@ def eval_main(cfg: EvalPipelineConfig):
     logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
 
     logging.info("Making environment.")
-    env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
+    envs = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
 
     logging.info("Making policy.")
-
     policy = make_policy(
         cfg=cfg.policy,
         env_cfg=cfg.env,
     )
-    policy.eval()
 
+    policy.eval()
     with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
-        info = eval_policy(
-            env,
+        info = eval_policy_all(
+            envs,
             policy,
             cfg.eval.n_episodes,
             max_episodes_rendered=10,
             videos_dir=Path(cfg.output_dir) / "videos",
             start_seed=cfg.seed,
+            max_parallel_tasks=cfg.env.max_parallel_tasks,
+            verbose=False,
         )
-    print(info["aggregated"])
+        print("Overall Aggregated Metrics:")
+        print(info["overall"]["aggregated"])
 
+        # Print per-suite stats
+        for task_group, task_group_info in info.items():
+            if task_group == "overall":
+                continue  # Skip the overall stats since we already printed it
+            print(f"\nAggregated Metrics for {task_group}:")
+            print(task_group_info["aggregated"])
+    # Close all vec envs
+    for _suite, task_map in envs.items():
+        for _vec in task_map.values():
+            _vec.close()
     # Save info
     with open(Path(cfg.output_dir) / "eval_info.json", "w") as f:
         json.dump(info, f, indent=2)
 
-    env.close()
-
     logging.info("End of eval")
 
 
-def main():
-    init_logging()
-    eval_main()
+# ---- typed payload returned by one task eval ----
+class TaskMetrics(TypedDict):
+    sum_rewards: list[float]
+    max_rewards: list[float]
+    successes: list[bool]
+    video_paths: list[str]
+
+
+ACC_KEYS = ("sum_rewards", "max_rewards", "successes", "video_paths")
+
+
+def eval_policy_all(
+    envs: dict[str, dict[int, gym.vector.VectorEnv]],
+    policy: PreTrainedPolicy,
+    n_episodes: int,
+    max_episodes_rendered: int = 0,
+    videos_dir: Path | None = None,
+    return_episode_data: bool = False,
+    start_seed: int | None = None,
+    max_parallel_tasks: int = 5,
+    verbose: bool = True,
+) -> dict:
+    """
+    Evaluate a policy over a dict-of-dicts of vectorized envs:
+    envs[suite_name][task_id] -> gym.vector.VectorEnv
+    Returns a dict with per-suite aggregates and an 'overall' block.
+    """
+    global_start = time.time()
+
+    # inner: evaluate a single (suite, task)
+    def eval_one(
+        task_group: str,
+        task_id: int,
+        env: gym.vector.VectorEnv,
+        *,
+        policy: PreTrainedPolicy,
+        n_episodes: int,
+        max_episodes_rendered: int,
+        videos_dir: Path | None,
+        return_episode_data: bool,
+        start_seed: int | None,
+    ) -> TaskMetrics:
+        """Evaluates one task_id of one suite using the provided vec env."""
+        if verbose:
+            print(f"Evaluating: task_group={task_group}, task_id={task_id} ...")
+
+        task_videos_dir = None
+        if videos_dir is not None:
+            task_videos_dir = videos_dir / f"{task_group}_{task_id}"
+            task_videos_dir.mkdir(parents=True, exist_ok=True)
+
+        task_result = eval_policy(
+            env=env,
+            policy=policy,
+            n_episodes=n_episodes,
+            max_episodes_rendered=max_episodes_rendered,
+            videos_dir=task_videos_dir,
+            return_episode_data=return_episode_data,
+            start_seed=start_seed,
+        )
+
+        per_episode = task_result["per_episode"]
+        return TaskMetrics(
+            sum_rewards=[ep["sum_reward"] for ep in per_episode],
+            max_rewards=[ep["max_reward"] for ep in per_episode],
+            successes=[ep["success"] for ep in per_episode],
+            video_paths=task_result.get("video_paths", []),
+        )
+
+    # result producer: sequential or threaded, same consumer
+    def iter_task_results() -> Iterator[tuple[str, int, TaskMetrics]]:
+        if max_parallel_tasks == 1:
+            for task_group, tasks in envs.items():
+                for task_id, vec in tasks.items():
+                    yield (
+                        task_group,
+                        task_id,
+                        eval_one(
+                            task_group,
+                            task_id,
+                            vec,
+                            policy=policy,
+                            n_episodes=n_episodes,
+                            max_episodes_rendered=max_episodes_rendered,
+                            videos_dir=videos_dir,
+                            return_episode_data=return_episode_data,
+                            start_seed=start_seed,
+                        ),
+                    )
+        else:
+            with cf.ThreadPoolExecutor(max_workers=max_parallel_tasks) as executor:
+                fut2key: dict[cf.Future, tuple[str, int]] = {}
+                for task_group, tasks in envs.items():
+                    for task_id, vec in tasks.items():
+                        fut = executor.submit(
+                            eval_one,
+                            task_group,
+                            task_id,
+                            vec,
+                            policy=policy,
+                            n_episodes=n_episodes,
+                            max_episodes_rendered=max_episodes_rendered,
+                            videos_dir=videos_dir,
+                            return_episode_data=return_episode_data,
+                            start_seed=start_seed,
+                        )
+                        fut2key[fut] = (task_group, task_id)
+                for fut in cf.as_completed(fut2key):
+                    task_group, task_id = fut2key[fut]
+                    yield task_group, task_id, fut.result()
+
+    # single accumulator path on the main thread
+    group_acc: dict[str, dict[str, list]] = defaultdict(lambda: {k: [] for k in ACC_KEYS})
+    overall: dict[str, list] = {k: [] for k in ACC_KEYS}
+
+    for task_group, task_id, metrics in iter_task_results():
+        acc = group_acc[task_group]
+        for k in ACC_KEYS:
+            acc[k].extend(metrics[k])
+            overall[k].extend(metrics[k])
+
+    # build outputs
+    results: dict[str, dict] = {}
+    for task_group, data in group_acc.items():
+        suite_rewards = data["sum_rewards"]
+        suite_max = data["max_rewards"]
+        suite_succ = data["successes"]
+        suite_vids = data["video_paths"]
+
+        suite_eval_s = time.time() - global_start
+        suite_eval_ep_s = suite_eval_s / max(1, len(suite_rewards))
+
+        results[task_group] = {
+            "aggregated": {
+                "avg_sum_reward": float(np.nanmean(suite_rewards)) if suite_rewards else float("nan"),
+                "avg_max_reward": float(np.nanmean(suite_max)) if suite_max else float("nan"),
+                "pc_success": float(np.nanmean(suite_succ) * 100) if suite_succ else float("nan"),
+                "eval_s": suite_eval_s,
+                "eval_ep_s": suite_eval_ep_s,
+            },
+            "video_paths": suite_vids,
+            "episodes": None,
+        }
+
+    global_eval_s = time.time() - global_start
+    global_eval_ep_s = global_eval_s / max(1, len(overall["sum_rewards"]))
+    results["overall"] = {
+        "aggregated": {
+            "avg_sum_reward": float(np.nanmean(overall["sum_rewards"]))
+            if overall["sum_rewards"]
+            else float("nan"),
+            "avg_max_reward": float(np.nanmean(overall["max_rewards"]))
+            if overall["max_rewards"]
+            else float("nan"),
+            "pc_success": float(np.nanmean(overall["successes"]) * 100)
+            if overall["successes"]
+            else float("nan"),
+            "eval_s": global_eval_s,
+            "eval_ep_s": global_eval_ep_s,
+        },
+        "video_paths": overall["video_paths"],
+        "episodes": None,
+    }
+    return results
 
 
 if __name__ == "__main__":
-    main()
+    init_logging()
+    eval_main()

src/lerobot/scripts/rl/crop_dataset_roi.py

@@ -226,8 +226,7 @@ def convert_lerobot_dataset_to_cropper_lerobot_dataset(
                 value = value.unsqueeze(0)
             new_frame[key] = value
 
-        new_frame["task"] = task
-        new_dataset.add_frame(new_frame)
+        new_dataset.add_frame(new_frame, task=task)
 
         if frame["episode_index"].item() != prev_episode_index:
             # Save the episode

src/lerobot/scripts/rl/gym_manipulator.py

@@ -2129,8 +2129,7 @@ def record_dataset(env, policy, cfg):
             frame["complementary_info.discrete_penalty"] = torch.tensor(
                 [info.get("discrete_penalty", 0.0)], dtype=torch.float32
             )
-            frame["task"] = cfg.task
-            dataset.add_frame(frame)
+            dataset.add_frame(frame, task=cfg.task)
 
             # Maintain consistent timing
             if cfg.fps:

src/lerobot/scripts/server/robot_client.py

@@ -302,11 +302,6 @@ class RobotClient:
 
                     self.logger.debug(f"Current latest action: {latest_action}")
 
-                    # Get queue state before changes
-                    old_size, old_timesteps = self._inspect_action_queue()
-                    if not old_timesteps:
-                        old_timesteps = [latest_action]  # queue was empty
-
                     # Get queue state before changes
                     old_size, old_timesteps = self._inspect_action_queue()
                     if not old_timesteps:

src/lerobot/scripts/train.py

@@ -30,11 +30,12 @@ from lerobot.datasets.factory import make_dataset
 from lerobot.datasets.sampler import EpisodeAwareSampler
 from lerobot.datasets.utils import cycle
 from lerobot.envs.factory import make_env
+from lerobot.envs.utils import close_envs
 from lerobot.optim.factory import make_optimizer_and_scheduler
 from lerobot.policies.factory import make_policy
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import get_device_from_parameters
-from lerobot.scripts.eval import eval_policy
+from lerobot.scripts.eval import eval_policy_all
 from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.train_utils import (
@@ -126,7 +127,6 @@ def train(cfg: TrainPipelineConfig):
 
     logging.info("Creating dataset")
     dataset = make_dataset(cfg)
-
     # Create environment used for evaluating checkpoints during training on simulation data.
     # On real-world data, no need to create an environment as evaluations are done outside train.py,
     # using the eval.py instead, with gym_dora environment and dora-rs.
@@ -140,7 +140,6 @@ def train(cfg: TrainPipelineConfig):
         cfg=cfg.policy,
         ds_meta=dataset.meta,
     )
-
     logging.info("Creating optimizer and scheduler")
     optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
     grad_scaler = GradScaler(device.type, enabled=cfg.policy.use_amp)
@@ -166,8 +165,7 @@ def train(cfg: TrainPipelineConfig):
     if hasattr(cfg.policy, "drop_n_last_frames"):
         shuffle = False
         sampler = EpisodeAwareSampler(
-            dataset.meta.episodes["dataset_from_index"],
-            dataset.meta.episodes["dataset_to_index"],
+            dataset.episode_data_index,
             drop_n_last_frames=cfg.policy.drop_n_last_frames,
             shuffle=True,
         )
@@ -187,7 +185,6 @@ def train(cfg: TrainPipelineConfig):
     dl_iter = cycle(dataloader)
 
     policy.train()
-
     train_metrics = {
         "loss": AverageMeter("loss", ":.3f"),
         "grad_norm": AverageMeter("grdn", ":.3f"),
@@ -205,7 +202,6 @@ def train(cfg: TrainPipelineConfig):
         start_time = time.perf_counter()
         batch = next(dl_iter)
         train_tracker.dataloading_s = time.perf_counter() - start_time
-
         for key in batch:
             if isinstance(batch[key], torch.Tensor):
                 batch[key] = batch[key].to(device, non_blocking=device.type == "cuda")
@@ -253,15 +249,27 @@ def train(cfg: TrainPipelineConfig):
                 torch.no_grad(),
                 torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(),
             ):
-                eval_info = eval_policy(
-                    eval_env,
+                eval_info = eval_policy_all(
+                    eval_env,  # dict[suite][task_id] -> vec_env
                     policy,
                     cfg.eval.n_episodes,
-                    videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}",
+                    videos_dir=videos_dir,
                     max_episodes_rendered=4,
                     start_seed=cfg.seed,
+                    max_parallel_tasks=cfg.env.max_parallel_tasks,
+                    verbose=False,
                 )
 
+            # overall metrics (suite-agnostic)
+            aggregated = eval_info["overall"]["aggregated"]
+
+            # optional: per-suite logging
+            for suite, suite_info in eval_info.items():
+                if suite == "overall":
+                    continue
+                logging.info("Suite %s aggregated: %s", suite, suite_info["aggregated"])
+
+            # meters/tracker
             eval_metrics = {
                 "avg_sum_reward": AverageMeter("∑rwrd", ":.3f"),
                 "pc_success": AverageMeter("success", ":.1f"),
@@ -270,17 +278,16 @@ def train(cfg: TrainPipelineConfig):
             eval_tracker = MetricsTracker(
                 cfg.batch_size, dataset.num_frames, dataset.num_episodes, eval_metrics, initial_step=step
             )
-            eval_tracker.eval_s = eval_info["aggregated"].pop("eval_s")
-            eval_tracker.avg_sum_reward = eval_info["aggregated"].pop("avg_sum_reward")
-            eval_tracker.pc_success = eval_info["aggregated"].pop("pc_success")
-            logging.info(eval_tracker)
+            eval_tracker.eval_s = aggregated.get("eval_s", 0.0)
+            eval_tracker.avg_sum_reward = aggregated.get("avg_sum_reward", float("nan"))
+            eval_tracker.pc_success = aggregated.get("pc_success", float("nan"))
             if wandb_logger:
                 wandb_log_dict = {**eval_tracker.to_dict(), **eval_info}
                 wandb_logger.log_dict(wandb_log_dict, step, mode="eval")
                 wandb_logger.log_video(eval_info["video_paths"][0], step, mode="eval")
 
     if eval_env:
-        eval_env.close()
+        close_envs(eval_env)
     logging.info("End of training")
 
     if cfg.policy.push_to_hub:

src/lerobot/scripts/visualize_dataset.py

@@ -79,8 +79,8 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset
 
 class EpisodeSampler(torch.utils.data.Sampler):
     def __init__(self, dataset: LeRobotDataset, episode_index: int):
-        from_idx = dataset.meta.episodes["dataset_from_index"][episode_index]
-        to_idx = dataset.meta.episodes["dataset_to_index"][episode_index]
+        from_idx = dataset.episode_data_index["from"][episode_index].item()
+        to_idx = dataset.episode_data_index["to"][episode_index].item()
         self.frame_ids = range(from_idx, to_idx)
 
     def __iter__(self) -> Iterator:
@@ -283,7 +283,7 @@ def main():
     tolerance_s = kwargs.pop("tolerance_s")
 
     logging.info("Loading dataset")
-    dataset = LeRobotDataset(repo_id, episodes=[args.episode_index], root=root, tolerance_s=tolerance_s)
+    dataset = LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)
 
     visualize_dataset(dataset, **vars(args))
 

src/lerobot/scripts/visualize_dataset_html.py

@@ -152,17 +152,13 @@ def run_server(
         dataset_version = (
             str(dataset.meta._version) if isinstance(dataset, LeRobotDataset) else dataset.codebase_version
         )
-
-        # Check minimum version requirement
         match = re.search(r"v(\d+)\.", dataset_version)
         if match:
             major_version = int(match.group(1))
             if major_version < 2:
                 return "Make sure to convert your LeRobotDataset to v2 & above."
 
-        # Get episode data once
         episode_data_csv_str, columns, ignored_columns = get_episode_data(dataset, episode_id)
-
         dataset_info = {
             "repo_id": f"{dataset_namespace}/{dataset_name}",
             "num_samples": dataset.num_frames
@@ -173,47 +169,19 @@ def run_server(
             else dataset.total_episodes,
             "fps": dataset.fps,
         }
-
         if isinstance(dataset, LeRobotDataset):
-            # Handle local datasets
-            # Determine if this is a chunked video dataset (v3.0+)
-            is_v3_or_later = False
-            match = re.search(r"v(\d+)\.(\d+)", dataset_version)
-            if match:
-                major_version = int(match.group(1))
-                is_v3_or_later = major_version >= 3
-
-            # Create videos_info with unified structure
-            videos_info = []
-
-            for key in dataset.meta.video_keys:
-                video_path = dataset.meta.get_video_file_path(episode_id, key)
-
-                if is_v3_or_later:
-                    # For v3.0+ datasets, get episode timestamps from chunked videos
-                    episode = dataset.meta.episodes[episode_id]
-                    from_timestamp = episode.get(f"videos/{key}/from_timestamp", 0)
-                    to_timestamp = episode.get(f"videos/{key}/to_timestamp", None)
-                    filename = key
-                else:
-                    # For v2.1 and earlier, videos are already per-episode
-                    from_timestamp = None
-                    to_timestamp = None
-                    filename = video_path.parent.name
-
-                videos_info.append(
-                    {
-                        "url": url_for("static", filename=str(video_path).replace("\\", "/")),
-                        "filename": filename,
-                        "start_time": from_timestamp,
-                        "end_time": to_timestamp,
-                        "is_chunked": is_v3_or_later,
-                    }
-                )
-
+            video_paths = [
+                dataset.meta.get_video_file_path(episode_id, key) for key in dataset.meta.video_keys
+            ]
+            videos_info = [
+                {
+                    "url": url_for("static", filename=str(video_path).replace("\\", "/")),
+                    "filename": video_path.parent.name,
+                }
+                for video_path in video_paths
+            ]
             tasks = dataset.meta.episodes[episode_id]["tasks"]
         else:
-            # Handle remote datasets from HF Hub
             video_keys = [key for key, ft in dataset.features.items() if ft["dtype"] == "video"]
             videos_info = [
                 {
@@ -224,9 +192,6 @@ def run_server(
                         episode_index=episode_id,
                     ),
                     "filename": video_key,
-                    "start_time": None,
-                    "end_time": None,
-                    "is_chunked": False,
                 }
                 for video_key in video_keys
             ]
@@ -306,8 +271,8 @@ def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index)
     selected_columns.insert(0, "timestamp")
 
     if isinstance(dataset, LeRobotDataset):
-        from_idx = dataset.meta.episodes["dataset_from_index"][episode_index]
-        to_idx = dataset.meta.episodes["dataset_to_index"][episode_index]
+        from_idx = dataset.episode_data_index["from"][episode_index]
+        to_idx = dataset.episode_data_index["to"][episode_index]
         data = (
             dataset.hf_dataset.select(range(from_idx, to_idx))
             .select_columns(selected_columns)
@@ -343,7 +308,7 @@ def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index)
 
 def get_episode_video_paths(dataset: LeRobotDataset, ep_index: int) -> list[str]:
     # get first frame of episode (hack to get video_path of the episode)
-    first_frame_idx = dataset.meta.episodes["dataset_from_index"][ep_index]
+    first_frame_idx = dataset.episode_data_index["from"][ep_index].item()
     return [
         dataset.hf_dataset.select_columns(key)[first_frame_idx][key]["path"]
         for key in dataset.meta.video_keys
@@ -356,7 +321,7 @@ def get_episode_language_instruction(dataset: LeRobotDataset, ep_index: int) ->
         return None
 
     # get first frame index
-    first_frame_idx = dataset.meta.episodes["dataset_from_index"][ep_index]
+    first_frame_idx = dataset.episode_data_index["from"][ep_index].item()
 
     language_instruction = dataset.hf_dataset[first_frame_idx]["language_instruction"]
     # TODO (michel-aractingi) hack to get the sentence, some strings in openx are badly stored

src/lerobot/teleoperators/koch_leader/koch_leader.py

@@ -88,6 +88,7 @@ class KochLeader(Teleoperator):
         return self.bus.is_calibrated
 
     def calibrate(self) -> None:
+        self.bus.disable_torque()
         if self.calibration:
             # Calibration file exists, ask user whether to use it or run new calibration
             user_input = input(
@@ -98,7 +99,6 @@ class KochLeader(Teleoperator):
                 self.bus.write_calibration(self.calibration)
                 return
         logger.info(f"\nRunning calibration of {self}")
-        self.bus.disable_torque()
         for motor in self.bus.motors:
             self.bus.write("Operating_Mode", motor, OperatingMode.EXTENDED_POSITION.value)
 

src/lerobot/teleoperators/reachy2_teleoperator/__init__.py

@@ -0,0 +1,25 @@
+#!/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 .config_reachy2_teleoperator import Reachy2TeleoperatorConfig
+from .reachy2_teleoperator import (
+    REACHY2_ANTENNAS_JOINTS,
+    REACHY2_L_ARM_JOINTS,
+    REACHY2_NECK_JOINTS,
+    REACHY2_R_ARM_JOINTS,
+    REACHY2_VEL,
+    Reachy2Teleoperator,
+)

src/lerobot/teleoperators/reachy2_teleoperator/config_reachy2_teleoperator.py

@@ -0,0 +1,51 @@
+#!/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.
+
+from dataclasses import dataclass
+
+from ..config import TeleoperatorConfig
+
+
+@TeleoperatorConfig.register_subclass("reachy2_teleoperator")
+@dataclass
+class Reachy2TeleoperatorConfig(TeleoperatorConfig):
+    # IP address of the Reachy 2 robot used as teleoperator
+    ip_address: str | None = "localhost"
+
+    # Whether to use the present position of the joints as actions
+    # if False, the goal position of the joints will be used
+    use_present_position: bool = False
+
+    # Which parts of the robot to use
+    with_mobile_base: bool = True
+    with_l_arm: bool = True
+    with_r_arm: bool = True
+    with_neck: bool = True
+    with_antennas: bool = True
+
+    def __post_init__(self):
+        if not (
+            self.with_mobile_base
+            or self.with_l_arm
+            or self.with_r_arm
+            or self.with_neck
+            or self.with_antennas
+        ):
+            raise ValueError(
+                "No Reachy2Teleoperator part used.\n"
+                "At least one part of the robot must be set to True "
+                "(with_mobile_base, with_l_arm, with_r_arm, with_neck, with_antennas)"
+            )

src/lerobot/teleoperators/reachy2_teleoperator/reachy2_teleoperator.py

@@ -0,0 +1,164 @@
+#!/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 logging
+import time
+
+from reachy2_sdk import ReachySDK
+
+from ..teleoperator import Teleoperator
+from .config_reachy2_teleoperator import Reachy2TeleoperatorConfig
+
+logger = logging.getLogger(__name__)
+
+# {lerobot_keys: reachy2_sdk_keys}
+REACHY2_NECK_JOINTS = {
+    "neck_yaw.pos": "head.neck.yaw",
+    "neck_pitch.pos": "head.neck.pitch",
+    "neck_roll.pos": "head.neck.roll",
+}
+
+REACHY2_ANTENNAS_JOINTS = {
+    "l_antenna.pos": "head.l_antenna",
+    "r_antenna.pos": "head.r_antenna",
+}
+
+REACHY2_R_ARM_JOINTS = {
+    "r_shoulder_pitch.pos": "r_arm.shoulder.pitch",
+    "r_shoulder_roll.pos": "r_arm.shoulder.roll",
+    "r_elbow_yaw.pos": "r_arm.elbow.yaw",
+    "r_elbow_pitch.pos": "r_arm.elbow.pitch",
+    "r_wrist_roll.pos": "r_arm.wrist.roll",
+    "r_wrist_pitch.pos": "r_arm.wrist.pitch",
+    "r_wrist_yaw.pos": "r_arm.wrist.yaw",
+    "r_gripper.pos": "r_arm.gripper",
+}
+
+REACHY2_L_ARM_JOINTS = {
+    "l_shoulder_pitch.pos": "l_arm.shoulder.pitch",
+    "l_shoulder_roll.pos": "l_arm.shoulder.roll",
+    "l_elbow_yaw.pos": "l_arm.elbow.yaw",
+    "l_elbow_pitch.pos": "l_arm.elbow.pitch",
+    "l_wrist_roll.pos": "l_arm.wrist.roll",
+    "l_wrist_pitch.pos": "l_arm.wrist.pitch",
+    "l_wrist_yaw.pos": "l_arm.wrist.yaw",
+    "l_gripper.pos": "l_arm.gripper",
+}
+
+REACHY2_VEL = {
+    "mobile_base.vx": "vx",
+    "mobile_base.vy": "vy",
+    "mobile_base.vtheta": "vtheta",
+}
+
+
+class Reachy2Teleoperator(Teleoperator):
+    """
+    [Reachy 2](https://www.pollen-robotics.com/reachy/), by Pollen Robotics.
+    """
+
+    config_class = Reachy2TeleoperatorConfig
+    name = "reachy2_specific"
+
+    def __init__(self, config: Reachy2TeleoperatorConfig):
+        super().__init__(config)
+        self.config = config
+        self.reachy: None | ReachySDK = None
+
+        self.joints_dict: dict[str, str] = self._generate_joints_dict()
+
+    def _generate_joints_dict(self) -> dict[str, str]:
+        joints = {}
+        if self.config.with_neck:
+            joints.update(REACHY2_NECK_JOINTS)
+        if self.config.with_l_arm:
+            joints.update(REACHY2_L_ARM_JOINTS)
+        if self.config.with_r_arm:
+            joints.update(REACHY2_R_ARM_JOINTS)
+        if self.config.with_antennas:
+            joints.update(REACHY2_ANTENNAS_JOINTS)
+        return joints
+
+    @property
+    def action_features(self) -> dict[str, type]:
+        if self.config.with_mobile_base:
+            return {
+                **dict.fromkeys(
+                    self.joints_dict.keys(),
+                    float,
+                ),
+                **dict.fromkeys(
+                    REACHY2_VEL.keys(),
+                    float,
+                ),
+            }
+        else:
+            return dict.fromkeys(self.joints_dict.keys(), float)
+
+    @property
+    def feedback_features(self) -> dict[str, type]:
+        return {}
+
+    @property
+    def is_connected(self) -> bool:
+        return self.reachy.is_connected() if self.reachy is not None else False
+
+    def connect(self, calibrate: bool = True) -> None:
+        self.reachy = ReachySDK(self.config.ip_address)
+        if not self.is_connected:
+            raise ConnectionError()
+        logger.info(f"{self} connected.")
+
+    @property
+    def is_calibrated(self) -> bool:
+        return True
+
+    def calibrate(self) -> None:
+        pass
+
+    def configure(self) -> None:
+        pass
+
+    def get_action(self) -> dict[str, float]:
+        start = time.perf_counter()
+
+        if self.reachy and self.is_connected:
+            if self.config.use_present_position:
+                joint_action = {
+                    k: self.reachy.joints[v].present_position for k, v in self.joints_dict.items()
+                }
+            else:
+                joint_action = {k: self.reachy.joints[v].goal_position for k, v in self.joints_dict.items()}
+
+            if not self.config.with_mobile_base:
+                dt_ms = (time.perf_counter() - start) * 1e3
+                logger.debug(f"{self} read action: {dt_ms:.1f}ms")
+                return joint_action
+
+            if self.config.use_present_position:
+                vel_action = {k: self.reachy.mobile_base.odometry[v] for k, v in REACHY2_VEL.items()}
+            else:
+                vel_action = {k: self.reachy.mobile_base.last_cmd_vel[v] for k, v in REACHY2_VEL.items()}
+        dt_ms = (time.perf_counter() - start) * 1e3
+        logger.debug(f"{self} read action: {dt_ms:.1f}ms")
+        return {**joint_action, **vel_action}
+
+    def send_feedback(self, feedback: dict[str, float]) -> None:
+        raise NotImplementedError
+
+    def disconnect(self) -> None:
+        if self.reachy and self.is_connected:
+            self.reachy.disconnect()

src/lerobot/teleoperators/utils.py

@@ -65,5 +65,9 @@ def make_teleoperator_from_config(config: TeleoperatorConfig) -> Teleoperator:
         from .bi_so100_leader import BiSO100Leader
 
         return BiSO100Leader(config)
+    elif config.type == "reachy2_teleoperator":
+        from .reachy2_teleoperator import Reachy2Teleoperator
+
+        return Reachy2Teleoperator(config)
     else:
         raise ValueError(config.type)

src/lerobot/utils/buffer.py

@@ -565,7 +565,10 @@ class ReplayBuffer:
         lerobot_dataset.start_image_writer(num_processes=0, num_threads=3)
 
         # Convert transitions into episodes and frames
+        episode_index = 0
+        lerobot_dataset.episode_buffer = lerobot_dataset.create_episode_buffer(episode_index=episode_index)
 
+        frame_idx_in_episode = 0
         for idx in range(self.size):
             actual_idx = (self.position - self.size + idx) % self.capacity
 
@@ -579,7 +582,6 @@ class ReplayBuffer:
             frame_dict["action"] = self.actions[actual_idx].cpu()
             frame_dict["next.reward"] = torch.tensor([self.rewards[actual_idx]], dtype=torch.float32).cpu()
             frame_dict["next.done"] = torch.tensor([self.dones[actual_idx]], dtype=torch.bool).cpu()
-            frame_dict["task"] = task_name
 
             # Add complementary_info if available
             if self.has_complementary_info:
@@ -595,11 +597,19 @@ class ReplayBuffer:
                         frame_dict[f"complementary_info.{key}"] = val
 
             # Add to the dataset's buffer
-            lerobot_dataset.add_frame(frame_dict)
+            lerobot_dataset.add_frame(frame_dict, task=task_name)
+
+            # Move to next frame
+            frame_idx_in_episode += 1
 
             # If we reached an episode boundary, call save_episode, reset counters
             if self.dones[actual_idx] or self.truncateds[actual_idx]:
                 lerobot_dataset.save_episode()
+                episode_index += 1
+                frame_idx_in_episode = 0
+                lerobot_dataset.episode_buffer = lerobot_dataset.create_episode_buffer(
+                    episode_index=episode_index
+                )
 
         # Save any remaining frames in the buffer
         if lerobot_dataset.episode_buffer["size"] > 0:

src/lerobot/utils/utils.py

@@ -274,16 +274,6 @@ def move_cursor_up(lines):
     print(f"\033[{lines}A", end="")
 
 
-def get_elapsed_time_in_days_hours_minutes_seconds(elapsed_time_s: float):
-    days = int(elapsed_time_s // (24 * 3600))
-    elapsed_time_s %= 24 * 3600
-    hours = int(elapsed_time_s // 3600)
-    elapsed_time_s %= 3600
-    minutes = int(elapsed_time_s // 60)
-    seconds = elapsed_time_s % 60
-    return days, hours, minutes, seconds
-
-
 class TimerManager:
     """
     Lightweight utility to measure elapsed time.

test_pi05_openpi.py

@@ -0,0 +1,226 @@
+#!/usr/bin/env python
+
+"""Test script to verify PI0.5 (pi05) support in PI0OpenPI policy."""
+
+import torch
+
+from lerobot.policies.pi0_openpi.configuration_pi0openpi import PI0OpenPIConfig
+from lerobot.policies.pi0_openpi.modeling_pi0openpi import PI0OpenPIPolicy
+from lerobot.policies.pi05_openpi import PI05OpenPIConfig, PI05OpenPIPolicy
+
+
+def test_pi05_model_architecture():
+    """Test that pi05=True creates the correct model architecture."""
+    print("Testing PI0.5 model architecture...")
+
+    # Create config
+    config = PI05OpenPIConfig(
+        action_dim=7,
+        state_dim=14,
+        dtype="float32",
+    )
+
+    # Verify tokenizer max length is set correctly
+    assert config.tokenizer_max_length == 200, (
+        f"Expected tokenizer_max_length=200 for pi05, got {config.tokenizer_max_length}"
+    )
+    print(f"✓ Tokenizer max length correctly set to {config.tokenizer_max_length}")
+
+    # Verify discrete_state_input defaults to pi05
+    assert config.discrete_state_input == True, (  # noqa: E712
+        f"Expected discrete_state_input=True for pi05, got {config.discrete_state_input}"
+    )
+    print(f"✓ discrete_state_input correctly defaults to pi05 value: {config.discrete_state_input}")
+
+    # Create dummy dataset stats
+    dataset_stats = {
+        "observation.state": {
+            "mean": torch.zeros(14),
+            "std": torch.ones(14),
+        },
+        "action": {
+            "mean": torch.zeros(7),
+            "std": torch.ones(7),
+        },
+    }
+
+    # Instantiate policy
+    policy = PI05OpenPIPolicy(config, dataset_stats)
+
+    # Verify pi05 model components exist
+
+    # Check that time_mlp layers exist (for AdaRMS conditioning)
+    assert hasattr(policy.model, "time_mlp_in"), "Missing time_mlp_in layer for pi05"
+    assert hasattr(policy.model, "time_mlp_out"), "Missing time_mlp_out layer for pi05"
+    print("✓ Time MLP layers present for AdaRMS conditioning")
+
+    # Check that action_time_mlp layers don't exist (pi0 only)
+    assert not hasattr(policy.model, "action_time_mlp_in"), "action_time_mlp_in should not exist in pi05 mode"
+    assert not hasattr(policy.model, "action_time_mlp_out"), (
+        "action_time_mlp_out should not exist in pi05 mode"
+    )
+    print("✓ Action-time MLP layers correctly absent")
+
+    # Check that state_proj doesn't exist in pi05 mode
+    assert not hasattr(policy.model, "state_proj"), "state_proj should not exist in pi05 mode"
+    print("✓ State projection layer correctly absent")
+
+    # Check AdaRMS configuration in the underlying model
+    adarms_config = policy.model.paligemma_with_expert.paligemma.config.text_config.use_adarms
+    assert adarms_config == False, f"PaliGemma should not use AdaRMS, got {adarms_config}"  # noqa: E712
+
+    adarms_expert_config = policy.model.paligemma_with_expert.gemma_expert.config.use_adarms
+    assert adarms_expert_config == True, (  # noqa: E712
+        f"Action expert should use AdaRMS in pi05, got {adarms_expert_config}"
+    )
+    print("✓ AdaRMS correctly configured: PaliGemma=False, Expert=True")
+
+    return True
+
+
+def test_pi05_forward_pass():
+    """Test forward pass with"""
+    print("\nTesting PI0.5 forward pass...")
+
+    # Create config
+    config = PI05OpenPIConfig(
+        action_dim=7,
+        state_dim=14,
+        dtype="float32",
+        chunk_size=16,  # Shorter chunk_size for testing
+        n_action_steps=16,  # Shorter action steps for testing
+    )
+
+    # Create dummy dataset stats
+    dataset_stats = {
+        "observation.state": {
+            "mean": torch.zeros(14),
+            "std": torch.ones(14),
+        },
+        "action": {
+            "mean": torch.zeros(7),
+            "std": torch.ones(7),
+        },
+    }
+
+    # Instantiate policy
+    policy = PI05OpenPIPolicy(config, dataset_stats)
+
+    # Create test batch
+    batch_size = 2
+    device = next(policy.parameters()).device
+    batch = {
+        "observation.state": torch.randn(batch_size, 14, dtype=torch.float32, device=device),
+        "action": torch.randn(batch_size, config.chunk_size, 7, dtype=torch.float32, device=device),
+        "observation.images.base_0_rgb": torch.rand(
+            batch_size, 3, 224, 224, dtype=torch.float32, device=device
+        ),
+        "task": ["Pick up the object"] * batch_size,
+    }
+
+    # Test forward pass
+    try:
+        loss, loss_dict = policy.forward(batch)
+        print(f"✓ Forward pass successful. Loss: {loss_dict['loss']:.4f}")
+        assert not torch.isnan(loss), "Loss is NaN"
+        assert loss.item() >= 0, "Loss should be non-negative"
+    except Exception as e:
+        print(f"✗ Forward pass failed: {e}")
+        return False
+
+    # Test action prediction
+    try:
+        with torch.no_grad():
+            action = policy.select_action(batch)
+        print(f"✓ Action prediction successful. Action shape: {action.shape}")
+        # When batch_size > 1, select_action returns (batch_size, action_dim)
+        assert action.shape == (batch_size, 7), f"Expected action shape ({batch_size}, 7), got {action.shape}"
+        assert not torch.isnan(action).any(), "Action contains NaN values"
+    except Exception as e:
+        print(f"✗ Action prediction failed: {e}")
+        return False
+
+    return True
+
+
+def test_pi0_vs_pi05_differences():
+    """Test key differences between pi0 and pi05 modes."""
+    print("\nComparing PI0 vs PI0.5 architectures...")
+
+    # Create both configurations
+    config_pi0 = PI0OpenPIConfig(action_dim=7, state_dim=14, dtype="float32")
+    config_pi05 = PI05OpenPIConfig(action_dim=7, state_dim=14, dtype="float32")
+
+    dataset_stats = {
+        "observation.state": {"mean": torch.zeros(14), "std": torch.ones(14)},
+        "action": {"mean": torch.zeros(7), "std": torch.ones(7)},
+    }
+
+    # Create both models
+    policy_pi0 = PI0OpenPIPolicy(config_pi0, dataset_stats)
+    policy_pi05 = PI05OpenPIPolicy(config_pi05, dataset_stats)
+
+    print("\nPI0 Model:")
+    print(f"  - Tokenizer max length: {config_pi0.tokenizer_max_length}")
+    print(f"  - Has state_proj: {hasattr(policy_pi0.model, 'state_proj')}")
+    print(f"  - Has action_time_mlp: {hasattr(policy_pi0.model, 'action_time_mlp_in')}")
+    print(f"  - Has time_mlp: {hasattr(policy_pi0.model, 'time_mlp_in')}")
+    print(f"  - Uses AdaRMS: {policy_pi0.model.paligemma_with_expert.gemma_expert.config.use_adarms}")
+
+    print("\nPI0.5 Model:")
+    print(f"  - Tokenizer max length: {config_pi05.tokenizer_max_length}")
+    print(f"  - discrete_state_input: {config_pi05.discrete_state_input}")
+    print(f"  - Has state_proj: {hasattr(policy_pi05.model, 'state_proj')}")
+    print(f"  - Has action_time_mlp: {hasattr(policy_pi05.model, 'action_time_mlp_in')}")
+    print(f"  - Has time_mlp: {hasattr(policy_pi05.model, 'time_mlp_in')}")
+    print(f"  - Uses AdaRMS: {policy_pi05.model.paligemma_with_expert.gemma_expert.config.use_adarms}")
+
+    # Count parameters
+    pi0_params = sum(p.numel() for p in policy_pi0.parameters())
+    pi05_params = sum(p.numel() for p in policy_pi05.parameters())
+
+    print("\nParameter counts:")
+    print(f"  - PI0: {pi0_params:,}")
+    print(f"  - PI0.5: {pi05_params:,}")
+    print(f"  - Difference: {pi0_params - pi05_params:,} (PI0.5 has fewer params due to no state embedding)")
+
+    return True
+
+
+def main():
+    """Run all PI0.5 tests."""
+    print("=" * 60)
+    print("PI0.5 Support Test Suite")
+    print("=" * 60)
+
+    tests = [
+        ("Model Architecture", test_pi05_model_architecture),
+        ("Forward Pass", test_pi05_forward_pass),
+        ("PI0 vs PI0.5 Comparison", test_pi0_vs_pi05_differences),
+    ]
+
+    all_passed = True
+    for test_name, test_func in tests:
+        print(f"\n[{test_name}]")
+        print("-" * 40)
+        try:
+            if not test_func():
+                all_passed = False
+                print(f"✗ {test_name} failed")
+        except Exception as e:
+            all_passed = False
+            print(f"✗ {test_name} failed with exception: {e}")
+            import traceback
+
+            traceback.print_exc()
+
+    print("\n" + "=" * 60)
+    if all_passed:
+        print("✅ All PI0.5 tests passed!")
+    else:
+        print("❌ Some tests failed.")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()

test_pi0_openpi.py

@@ -0,0 +1,109 @@
+#!/usr/bin/env python
+
+"""Test script to verify PI0OpenPI policy integration with LeRobot."""
+
+import torch
+
+from lerobot.policies.factory import make_policy_config
+from lerobot.policies.pi0_openpi import PI0OpenPIConfig, PI0OpenPIPolicy
+
+
+def test_policy_instantiation():
+    """Test basic policy instantiation."""
+    print("Testing PI0OpenPI policy instantiation...")
+
+    # Create config
+    config = PI0OpenPIConfig(action_dim=7, state_dim=14, dtype="float32")
+
+    # Create dummy dataset stats
+    dataset_stats = {
+        "observation.state": {
+            "mean": torch.zeros(14),
+            "std": torch.ones(14),
+        },
+        "action": {
+            "mean": torch.zeros(7),
+            "std": torch.ones(7),
+        },
+    }
+
+    # Instantiate policy
+    policy = PI0OpenPIPolicy(config, dataset_stats)
+    print(f"Policy created successfully: {policy.name}")
+
+    # Test forward pass with dummy data
+    batch_size = 1
+    device = policy.device if hasattr(policy, "device") else "cpu"
+    batch = {
+        "observation.state": torch.randn(batch_size, 14, dtype=torch.float32, device=device),
+        "action": torch.randn(batch_size, config.chunk_size, 7, dtype=torch.float32, device=device),
+        "observation.images.base_0_rgb": torch.rand(
+            batch_size, 3, 224, 224, dtype=torch.float32, device=device
+        ),  # Use rand for [0,1] range
+        "task": ["Pick up the object"] * batch_size,
+    }
+
+    print("\nTesting forward pass...")
+    try:
+        loss, loss_dict = policy.forward(batch)
+        print(f"✓ Forward pass successful. Loss: {loss_dict['loss']:.4f}")
+    except Exception as e:
+        print(f"✗ Forward pass failed: {e}")
+        return False
+
+    print("\nTesting action prediction...")
+    try:
+        with torch.no_grad():
+            action = policy.select_action(batch)
+        print(f"✓ Action prediction successful. Action shape: {action.shape}")
+    except Exception as e:
+        print(f"✗ Action prediction failed: {e}")
+        return False
+
+    return True
+
+
+def test_config_creation():
+    """Test policy config creation through factory."""
+    print("\nTesting config creation through factory...")
+
+    try:
+        config = make_policy_config(
+            policy_type="pi0_openpi",
+            action_dim=7,
+            state_dim=14,
+        )
+        print("✓ Config created successfully through factory")
+        print(f"  Config type: {type(config).__name__}")
+        print(f"  PaliGemma variant: {config.paligemma_variant}")
+        print(f"  Action expert variant: {config.action_expert_variant}")
+        return True
+    except Exception as e:
+        print(f"✗ Config creation failed: {e}")
+        return False
+
+
+def main():
+    """Run all tests."""
+    print("=" * 60)
+    print("PI0OpenPI Policy Integration Test")
+    print("=" * 60)
+
+    # Test config creation
+    config_test = test_config_creation()
+
+    print("\n" + "-" * 60)
+
+    # Test policy instantiation
+    policy_test = test_policy_instantiation()
+
+    print("\n" + "=" * 60)
+    if config_test and policy_test:
+        print("✓ All tests passed!")
+    else:
+        print("✗ Some tests failed.")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()

test_pi0_original_vs_lerobot.py

@@ -0,0 +1,396 @@
+"""Test script to verify PI0OpenPI policy integration with LeRobot vs the original implementation."""
+
+import os
+
+import torch
+from openpi.models_pytorch import preprocessing_pytorch as openpi_preprocessing
+
+# NOTE: Assumes PYTHONPATH is set to include OpenPI src as per instructions.
+from openpi.models_pytorch.pi0_pytorch import PI0Pytorch
+from transformers import AutoTokenizer
+
+from lerobot.policies.pi0_openpi import PI0OpenPIConfig, PI0OpenPIPolicy
+
+DUMMY_ACTION_DIM = 32
+DUMMY_STATE_DIM = 32
+DUMMY_ACTION_HORIZON = 50
+DUMMY_MAX_TOKEN_LEN = 48  # Default for PI0 (non-pi05)
+DEVICE = "cpu"  # Use CPU to avoid memory issues for testing
+
+DUMMY_DATASET_STATS = {
+    "observation.state": {
+        "mean": torch.zeros(DUMMY_STATE_DIM),
+        "std": torch.ones(DUMMY_STATE_DIM),
+    },
+    "action": {
+        "mean": torch.zeros(DUMMY_ACTION_DIM),
+        "std": torch.ones(DUMMY_ACTION_DIM),
+    },
+    "images": {
+        "base_0_rgb": {
+            "mean": torch.zeros(3, 224, 224),
+            "std": torch.ones(3, 224, 224),
+        },
+        "left_wrist_0_rgb": {
+            "mean": torch.zeros(3, 224, 224),
+            "std": torch.ones(3, 224, 224),
+        },
+        "right_wrist_0_rgb": {
+            "mean": torch.zeros(3, 224, 224),
+            "std": torch.ones(3, 224, 224),
+        },
+    },
+}
+
+
+class PI0BaseOriginalConfig:
+    action_dim: int = DUMMY_ACTION_DIM
+    action_horizon: int = DUMMY_ACTION_HORIZON
+    paligemma_variant: str = "gemma_2b"
+    action_expert_variant: str = "gemma_300m"
+    precision: str = "float32"
+    pi05: bool = False
+    dtype: str = "float32"
+
+
+def instantiate_lerobot_pi0(from_pretrained: bool = False):
+    if from_pretrained:
+        # Load the policy first
+        policy = PI0OpenPIPolicy.from_pretrained(
+            pretrained_name_or_path="pepijn223/pi0_base_fp32", strict=True
+        )
+        # Then reinitialize the normalization with proper stats
+        from lerobot.policies.normalize import Normalize, Unnormalize
+
+        policy.normalize_inputs = Normalize(
+            policy.config.input_features, policy.config.normalization_mapping, DUMMY_DATASET_STATS
+        )
+        policy.normalize_targets = Normalize(
+            policy.config.output_features, policy.config.normalization_mapping, DUMMY_DATASET_STATS
+        )
+        policy.unnormalize_outputs = Unnormalize(
+            policy.config.output_features, policy.config.normalization_mapping, DUMMY_DATASET_STATS
+        )
+    else:
+        config = PI0OpenPIConfig(action_dim=DUMMY_ACTION_DIM, state_dim=DUMMY_STATE_DIM, dtype="float32")
+        policy = PI0OpenPIPolicy(config, DUMMY_DATASET_STATS)
+    policy.to(DEVICE)
+    return policy
+
+
+def instantiate_original_pi0(from_pretrained: bool = False, model_path: str = None):
+    config = PI0BaseOriginalConfig()
+    policy = PI0Pytorch(config)
+
+    if from_pretrained:
+        try:
+            print("Loading converted PyTorch weights from HuggingFace Hub (pepijn223/pi0_base_fp32)...")
+
+            # Download the model from HuggingFace Hub
+            import safetensors.torch
+            from huggingface_hub import snapshot_download
+
+            # Download the entire repository
+            if model_path and os.path.exists(model_path):
+                cache_dir = model_path
+                print(f"Using cached model from: {cache_dir}")
+            else:
+                cache_dir = snapshot_download(repo_id="pepijn223/pi0_base_fp32", repo_type="model")
+                print(f"Downloaded model to: {cache_dir}")
+
+            # Try to load safetensors format first
+            model_file = os.path.join(cache_dir, "model.safetensors")
+            if os.path.exists(model_file):
+                state_dict = safetensors.torch.load_file(model_file)
+                print(f"Loaded {len(state_dict)} parameters from safetensors")
+            else:
+                raise FileNotFoundError(f"No safetensors file found in {cache_dir}")
+
+            # Load the state dict into the model
+            missing_keys, unexpected_keys = policy.load_state_dict(state_dict, strict=False)
+
+            if missing_keys:
+                print(f"Missing keys: {len(missing_keys)}")
+                if len(missing_keys) <= 5:
+                    for key in missing_keys:
+                        print(f"    - {key}")
+                else:
+                    for key in missing_keys[:5]:
+                        print(f"    - {key}")
+                    print(f"    ... and {len(missing_keys) - 5} more")
+
+            if unexpected_keys:
+                print(f"Unexpected keys: {len(unexpected_keys)}")
+                if len(unexpected_keys) <= 5:
+                    for key in unexpected_keys:
+                        print(f"    - {key}")
+                else:
+                    for key in unexpected_keys[:5]:
+                        print(f"    - {key}")
+                    print(f"    ... and {len(unexpected_keys) - 5} more")
+
+            if not missing_keys and not unexpected_keys:
+                print("All pretrained weights loaded successfully!")
+            else:
+                print("Pretrained weights loaded with some missing/unexpected keys (this may be normal)")
+
+        except Exception as e:
+            print(f"Failed to load pretrained weights: {e}")
+            print("   Using randomly initialized weights...")
+            import traceback
+
+            traceback.print_exc()
+
+    policy.to(DEVICE)
+    return policy
+
+
+def create_dummy_data():
+    batch_size = 2  # Reduce batch size for testing
+    device = DEVICE
+
+    # Use the exact same prompt for both implementations
+    prompt = "Pick up the red block and place it in the bin"
+
+    batch = {
+        "observation.state": torch.randn(batch_size, DUMMY_STATE_DIM, dtype=torch.float32, device=device),
+        "action": torch.randn(
+            batch_size, DUMMY_ACTION_HORIZON, DUMMY_ACTION_DIM, dtype=torch.float32, device=device
+        ),
+        # Create images in [0, 1] range as expected by LeRobot (will be converted to [-1, 1] internally)
+        "observation.images.base_0_rgb": torch.rand(
+            batch_size, 3, 224, 224, dtype=torch.float32, device=device
+        ),
+        "observation.images.left_wrist_0_rgb": torch.rand(
+            batch_size, 3, 224, 224, dtype=torch.float32, device=device
+        ),
+        "observation.images.right_wrist_0_rgb": torch.rand(
+            batch_size, 3, 224, 224, dtype=torch.float32, device=device
+        ),
+        # Add the task prompt for LeRobot - provide as list with single element to trigger expansion
+        "task": [prompt],
+    }
+    return batch
+
+
+def extract_lerobot_processed_inputs(lerobot_pi0, batch):
+    """Extract the exact same processed inputs that LeRobot uses internally."""
+    # Get the tokenized language from LeRobot's internal method
+    lang_tokens, lang_masks = lerobot_pi0._tokenize_language(batch)
+
+    # Get the preprocessed images from LeRobot's internal method
+    images, img_masks = lerobot_pi0._preprocess_images(batch, train=False)
+
+    # Create dummy token_ar_mask and token_loss_mask for original implementation
+    token_ar_mask = torch.zeros_like(lang_tokens, dtype=torch.int32)
+    token_loss_mask = torch.ones_like(lang_masks, dtype=torch.bool)
+
+    return images, img_masks, lang_tokens, lang_masks, token_ar_mask, token_loss_mask
+
+
+class PI0Observation:
+    """Observation class that matches the original OpenPI format."""
+
+    def __init__(
+        self,
+        state,
+        images,
+        image_masks,
+        tokenized_prompt,
+        tokenized_prompt_mask,
+        token_ar_mask,
+        token_loss_mask,
+    ):
+        self.state = state
+        self.images = images
+        self.image_masks = image_masks
+        self.tokenized_prompt = tokenized_prompt
+        self.tokenized_prompt_mask = tokenized_prompt_mask
+        self.token_ar_mask = token_ar_mask
+        self.token_loss_mask = token_loss_mask
+
+
+def create_original_observation_with_openpi_preprocessing(batch):
+    """Create observation object for OpenPI using OpenPI's own preprocessing."""
+    batch_size = batch["observation.state"].shape[0]
+    device = batch["observation.state"].device
+
+    # Create tokenizer for OpenPI (same as LeRobot uses)
+    tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224")
+
+    # Get task description
+    if "task" in batch:
+        tasks = batch["task"]
+        if isinstance(tasks, str):
+            tasks = [tasks]
+        elif isinstance(tasks, list) and len(tasks) == 1:
+            # Expand to batch size
+            tasks = tasks * batch_size
+    else:
+        # Default task if not provided
+        tasks = ["Pick up the object"] * batch_size
+
+    # Tokenize with max_length padding to match OpenPI's expected format
+    tokenized = tokenizer(
+        tasks,
+        padding="max_length",
+        padding_side="right",
+        truncation=True,
+        max_length=DUMMY_MAX_TOKEN_LEN,
+        return_tensors="pt",
+    )
+
+    lang_tokens = tokenized["input_ids"].to(device)
+    lang_masks = tokenized["attention_mask"].to(device, dtype=torch.bool)
+
+    # Create dummy token_ar_mask and token_loss_mask for OpenPI
+    token_ar_mask = torch.zeros_like(lang_tokens, dtype=torch.int32)
+    token_loss_mask = torch.ones_like(lang_masks, dtype=torch.bool)
+
+    # Convert LeRobot images format to OpenPI format (convert [0,1] to [-1,1] range)
+    image_dict = {
+        "base_0_rgb": batch["observation.images.base_0_rgb"] * 2.0 - 1.0,
+        "left_wrist_0_rgb": batch["observation.images.left_wrist_0_rgb"] * 2.0 - 1.0,
+        "right_wrist_0_rgb": batch["observation.images.right_wrist_0_rgb"] * 2.0 - 1.0,
+    }
+
+    # Create image masks (all ones for real images)
+    image_masks_dict = {}
+    for key in image_dict:
+        image_masks_dict[key] = torch.ones(batch_size, dtype=torch.bool, device=device)
+
+    # Create raw observation object (before preprocessing)
+    raw_observation = PI0Observation(
+        state=batch["observation.state"],
+        images=image_dict,
+        image_masks=image_masks_dict,
+        tokenized_prompt=lang_tokens,
+        tokenized_prompt_mask=lang_masks,
+        token_ar_mask=token_ar_mask,
+        token_loss_mask=token_loss_mask,
+    )
+
+    # Now use OpenPI's preprocessing
+    processed_obs = openpi_preprocessing.preprocess_observation_pytorch(raw_observation, train=False)
+
+    return processed_obs
+
+
+def create_original_observation_from_lerobot(lerobot_pi0, batch):
+    """Create observation object compatible with original OpenPI using the exact same inputs as LeRobot."""
+    _batch_size = batch["observation.state"].shape[0]
+    _device = batch["observation.state"].device
+
+    # Extract the exact same processed inputs that LeRobot uses
+    images, img_masks, lang_tokens, lang_masks, token_ar_mask, token_loss_mask = (
+        extract_lerobot_processed_inputs(lerobot_pi0, batch)
+    )
+
+    # Convert images list to dict with original OpenPI keys
+    image_dict = {
+        "base_0_rgb": images[0],
+        "left_wrist_0_rgb": images[1],
+        "right_wrist_0_rgb": images[2],
+    }
+
+    # Convert image masks list to dict with original OpenPI keys
+    image_masks_dict = {
+        "base_0_rgb": img_masks[0],
+        "left_wrist_0_rgb": img_masks[1],
+        "right_wrist_0_rgb": img_masks[2],
+    }
+
+    return PI0Observation(
+        state=batch["observation.state"],
+        images=image_dict,
+        image_masks=image_masks_dict,
+        tokenized_prompt=lang_tokens,
+        tokenized_prompt_mask=lang_masks,
+        token_ar_mask=token_ar_mask,
+        token_loss_mask=token_loss_mask,
+    )
+
+
+def main():
+    print("Initializing models...")
+    lerobot_pi0 = instantiate_lerobot_pi0(from_pretrained=True)  # Load pretrained LeRobot model
+    original_pi0 = instantiate_original_pi0(
+        from_pretrained=True
+    )  # Load pretrained OpenPI model from HuggingFace Hub
+
+    print("Creating dummy data...")
+    batch = create_dummy_data()
+
+    # Test 1: Each model with its own preprocessing (more realistic end-to-end test)
+    print("\n=== TEST 1: Each model with its own preprocessing ===")
+    print("Creating observation for OpenPI using OpenPI's own preprocessing...")
+    pi0_obs_openpi = create_original_observation_with_openpi_preprocessing(batch)
+
+    print(f"Task prompt: '{batch['task'][0]}'")
+    print(f"OpenPI tokenized prompt shape: {pi0_obs_openpi.tokenized_prompt.shape}")
+    print(f"OpenPI image shapes: {[img.shape for img in pi0_obs_openpi.images.values()]}")
+    print(f"OpenPI state shape: {pi0_obs_openpi.state.shape}")
+
+    print("Testing OpenPI with own preprocessing...")
+    original_pi0.eval()
+    torch.manual_seed(42)  # Set seed for reproducibility
+    batch_size = batch["observation.state"].shape[0]
+    noise_shape = (batch_size, DUMMY_ACTION_HORIZON, DUMMY_ACTION_DIM)
+    fixed_noise = torch.randn(noise_shape, dtype=torch.float32, device=DEVICE)
+
+    with torch.no_grad():
+        openpi_actions = original_pi0.sample_actions(
+            device=DEVICE, observation=pi0_obs_openpi, noise=fixed_noise, num_steps=10
+        )
+    print(f"OpenPI (own preprocessing) Actions shape: {openpi_actions.shape}")
+    print(f"OpenPI (own preprocessing) Actions mean: {openpi_actions.mean().item():.6f}")
+    print(f"OpenPI (own preprocessing) Actions std: {openpi_actions.std().item():.6f}")
+
+    print("Testing LeRobot with own preprocessing...")
+    lerobot_pi0.eval()
+    torch.manual_seed(42)  # Set the same seed
+    with torch.no_grad():
+        lerobot_actions_own = lerobot_pi0.predict_action_chunk(batch)
+    print(f"LeRobot (own preprocessing) Actions shape: {lerobot_actions_own.shape}")
+    print(f"LeRobot (own preprocessing) Actions mean: {lerobot_actions_own.mean().item():.6f}")
+    print(f"LeRobot (own preprocessing) Actions std: {lerobot_actions_own.std().item():.6f}")
+
+    print("\nComparing end-to-end implementations:")
+    print(f"Actions close (atol=1e-4): {torch.allclose(lerobot_actions_own, openpi_actions, atol=1e-4)}")
+    print(f"Actions close (atol=1e-2): {torch.allclose(lerobot_actions_own, openpi_actions, atol=1e-2)}")
+    print(f"Max absolute difference: {torch.abs(lerobot_actions_own - openpi_actions).max().item():.6f}")
+
+    # Test 2: Both models with LeRobot preprocessing (isolates model differences)
+    print("\n=== TEST 2: Both models with LeRobot preprocessing (model comparison) ===")
+    print("Creating observation for OpenPI using LeRobot's preprocessing...")
+    pi0_obs_lerobot = create_original_observation_from_lerobot(lerobot_pi0, batch)
+
+    print("Testing OpenPI with LeRobot preprocessing...")
+    torch.manual_seed(42)  # Set seed for reproducibility
+    with torch.no_grad():
+        openpi_actions_lerobot_preproc = original_pi0.sample_actions(
+            device=DEVICE, observation=pi0_obs_lerobot, noise=fixed_noise, num_steps=10
+        )
+    print(f"OpenPI (LeRobot preprocessing) Actions shape: {openpi_actions_lerobot_preproc.shape}")
+    print(f"OpenPI (LeRobot preprocessing) Actions mean: {openpi_actions_lerobot_preproc.mean().item():.6f}")
+    print(f"OpenPI (LeRobot preprocessing) Actions std: {openpi_actions_lerobot_preproc.std().item():.6f}")
+
+    print("\nComparing models with same preprocessing:")
+    print(
+        f"Actions close (atol=1e-4): {torch.allclose(lerobot_actions_own, openpi_actions_lerobot_preproc, atol=1e-4)}"
+    )
+    print(
+        f"Actions close (atol=1e-2): {torch.allclose(lerobot_actions_own, openpi_actions_lerobot_preproc, atol=1e-2)}"
+    )
+    print(
+        f"Max absolute difference: {torch.abs(lerobot_actions_own - openpi_actions_lerobot_preproc).max().item():.6f}"
+    )
+
+    print("\n=== SUMMARY ===")
+    print("Test 1 compares end-to-end pipelines (each model with its own preprocessing)")
+    print("Test 2 isolates model differences (both models with LeRobot preprocessing)")
+    print("Both tests completed successfully!")
+
+
+if __name__ == "__main__":
+    main()

test_pi0_pi05_hub.py

@@ -0,0 +1,273 @@
+#!/usr/bin/env python
+
+"""Test script to load PI0OpenPI model from HuggingFace hub and run inference."""
+
+import torch
+
+from lerobot.policies.pi0_openpi import PI0OpenPIPolicy
+from lerobot.policies.pi05_openpi.modeling_pi05openpi import PI05OpenPIPolicy
+
+
+def create_dummy_stats(config):
+    """Create dummy dataset statistics for testing."""
+    dummy_stats = {
+        "observation.state": {
+            "mean": torch.zeros(config.state_dim),
+            "std": torch.ones(config.state_dim),
+        },
+        "action": {
+            "mean": torch.zeros(config.action_dim),
+            "std": torch.ones(config.action_dim),
+        },
+    }
+
+    # Add stats for image keys if they exist
+    for key in config.image_keys:
+        dummy_stats[key] = {
+            "mean": torch.zeros(3, config.image_resolution[0], config.image_resolution[1]),
+            "std": torch.ones(3, config.image_resolution[0], config.image_resolution[1]),
+        }
+
+    return dummy_stats
+
+
+def test_hub_loading(model_id="pepijn223/pi0_base_fp32", model_name="PI0"):
+    """Test loading model from HuggingFace hub.
+
+    Args:
+        model_id: HuggingFace model ID to load
+        model_name: Display name for the model (e.g., "PI0", "PI0.5")
+    """
+    print("=" * 60)
+    print(f"{model_name} OpenPI HuggingFace Hub Loading Test")
+    print("=" * 60)
+
+    print(f"\nLoading model from: {model_id}")
+    print("-" * 60)
+
+    try:
+        # Load the model from HuggingFace hub with strict mode
+        if model_name == "PI0.5":
+            policy = PI05OpenPIPolicy.from_pretrained(
+                model_id,
+                strict=True,  # Ensure all weights are loaded correctly,
+            )
+        else:
+            policy = PI0OpenPIPolicy.from_pretrained(
+                model_id,
+                strict=True,  # Ensure all weights are loaded correctly,
+            )
+
+        print("✓ Model loaded successfully from HuggingFace hub")
+
+        # Inject dummy stats since they aren't loaded from the hub
+        print("Creating dummy dataset stats for testing...")
+        device = next(policy.parameters()).device
+        dummy_stats = create_dummy_stats(policy.config)
+
+        # Move dummy stats to device
+        for key, stats in dummy_stats.items():
+            dummy_stats[key] = {
+                "mean": stats["mean"].to(device),
+                "std": stats["std"].to(device),
+            }
+
+        # Initialize normalization layers with dummy stats if they have NaN/inf values
+        print("✓ Dummy stats created and moved to device")
+
+        # Get model info
+        print("\nModel configuration:")
+        print(f"  - Model type: {model_name}")
+        print(f"  - PaliGemma variant: {policy.config.paligemma_variant}")
+        print(f"  - Action expert variant: {policy.config.action_expert_variant}")
+        print(f"  - Action dimension: {policy.config.action_dim}")
+        print(f"  - State dimension: {policy.config.state_dim}")
+        print(f"  - Chunk_size: {policy.config.chunk_size}")
+        print(f"  - Tokenizer max length: {policy.config.tokenizer_max_length}")
+        if model_name == "PI0.5":
+            print(f"  - discrete_state_input: {policy.config.discrete_state_input}")
+        print(f"  - Device: {device}")
+        print(f"  - Dtype: {next(policy.parameters()).dtype}")
+
+        # Check model-specific features
+        if model_name == "PI0.5":
+            print("\nPI0.5 specific features:")
+            print(f"  - Has time_mlp layers: {hasattr(policy.model, 'time_mlp_in')}")
+            print(f"  - Has state_proj: {hasattr(policy.model, 'state_proj')} (should be False)")
+            print(f"  - Uses AdaRMS: {policy.model.paligemma_with_expert.gemma_expert.config.use_adarms}")
+
+            # Verify PI0.5 architecture
+            assert hasattr(policy.model, "time_mlp_in"), "PI0.5 should have time_mlp_in"
+            assert hasattr(policy.model, "time_mlp_out"), "PI0.5 should have time_mlp_out"
+            assert not hasattr(policy.model, "state_proj"), "PI0.5 should not have state_proj"
+            assert not hasattr(policy.model, "action_time_mlp_in"), "PI0.5 should not have action_time_mlp_in"
+            print("  ✓ PI0.5 architecture verified")
+        else:
+            print("\nPI0 specific features:")
+            print(f"  - Has action_time_mlp layers: {hasattr(policy.model, 'action_time_mlp_in')}")
+            print(f"  - Has state_proj: {hasattr(policy.model, 'state_proj')} (should be True)")
+            print(
+                f"  - Uses AdaRMS: {policy.model.paligemma_with_expert.gemma_expert.config.use_adarms} (should be False)"
+            )
+
+            # Verify PI0 architecture
+            assert hasattr(policy.model, "action_time_mlp_in"), "PI0 should have action_time_mlp_in"
+            assert hasattr(policy.model, "action_time_mlp_out"), "PI0 should have action_time_mlp_out"
+            assert hasattr(policy.model, "state_proj"), "PI0 should have state_proj"
+            assert not hasattr(policy.model, "time_mlp_in"), "PI0 should not have time_mlp_in"
+            print("  ✓ PI0 architecture verified")
+
+    except Exception as e:
+        print(f"✗ Failed to load model: {e}")
+        return False
+
+    print("\n" + "-" * 60)
+    print("Testing forward pass with loaded model...")
+
+    # Create dummy batch for testing
+    batch_size = 1
+
+    # Check if normalization layers have invalid stats and replace with dummy stats if needed
+    try:
+        # Check if the normalize_inputs has valid stats
+        if hasattr(policy.normalize_inputs, "stats"):
+            obs_state_mean = policy.normalize_inputs.stats.get("observation.state", {}).get("mean")
+            if obs_state_mean is not None and (
+                torch.isinf(obs_state_mean).any() or torch.isnan(obs_state_mean).any()
+            ):
+                print("⚠️  Found invalid normalization stats, replacing with dummy stats...")
+
+                # Replace with dummy stats
+                from lerobot.policies.normalize import Normalize, Unnormalize
+
+                policy.normalize_inputs = Normalize(
+                    policy.config.input_features, policy.config.normalization_mapping, dummy_stats
+                )
+                policy.normalize_targets = Normalize(
+                    policy.config.output_features, policy.config.normalization_mapping, dummy_stats
+                )
+                policy.unnormalize_outputs = Unnormalize(
+                    policy.config.output_features, policy.config.normalization_mapping, dummy_stats
+                )
+                print("✓ Normalization layers updated with dummy stats")
+    except Exception as e:
+        print(f"⚠️  Error checking normalization stats, creating new ones: {e}")
+        # Fallback: create new normalization layers
+        from lerobot.policies.normalize import Normalize, Unnormalize
+
+        policy.normalize_inputs = Normalize(
+            policy.config.input_features, policy.config.normalization_mapping, dummy_stats
+        )
+        policy.normalize_targets = Normalize(
+            policy.config.output_features, policy.config.normalization_mapping, dummy_stats
+        )
+        policy.unnormalize_outputs = Unnormalize(
+            policy.config.output_features, policy.config.normalization_mapping, dummy_stats
+        )
+
+    # Create test batch
+    batch = {
+        "observation.state": torch.randn(
+            batch_size, policy.config.state_dim, dtype=torch.float32, device=device
+        ),
+        "action": torch.randn(
+            batch_size,
+            policy.config.chunk_size,
+            policy.config.action_dim,
+            dtype=torch.float32,
+            device=device,
+        ),
+        "task": ["Pick up the object"] * batch_size,
+    }
+
+    # Add images if they're in the config
+    for key in policy.config.image_keys:
+        batch[key] = torch.rand(batch_size, 3, 224, 224, dtype=torch.float32, device=device)
+
+    try:
+        # Test forward pass
+        policy.train()  # Set to training mode for forward pass with loss
+        loss, loss_dict = policy.forward(batch)
+        print("✓ Forward pass successful")
+        print(f"  - Loss: {loss_dict['loss']:.4f}")
+        print(f"  - Loss shape: {loss.shape if hasattr(loss, 'shape') else 'scalar'}")
+
+    except Exception as e:
+        print(f"✗ Forward pass failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+        return False
+
+    print("\n" + "-" * 60)
+    print("Testing inference with loaded model...")
+
+    try:
+        # Test action prediction
+        policy.eval()  # Set to evaluation mode for inference
+        with torch.no_grad():
+            action = policy.select_action(batch)
+        print("✓ Action prediction successful")
+        print(f"  - Action shape: {action.shape}")
+        print(f"  - Action range: [{action.min().item():.3f}, {action.max().item():.3f}]")
+
+    except Exception as e:
+        print(f"✗ Action prediction failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+        return False
+
+    print("\n" + "=" * 60)
+    print(f"✓ All tests passed for {model_name}!")
+    print("=" * 60)
+    return True
+
+
+def main():
+    """Run tests for both PI0 and PI0.5 models."""
+    print("\n")
+    print("╔" + "═" * 58 + "╗")
+    print("║" + "  PI0 & PI0.5 HuggingFace Hub Loading Test Suite  ".center(58) + "║")
+    print("╚" + "═" * 58 + "╝")
+    print()
+
+    results = []
+
+    # Test PI0 model
+    print("\n[Test 1/2] Testing PI0 model...")
+    print("─" * 60)
+    pi0_success = test_hub_loading(model_id="pepijn223/pi0_base_fp32", model_name="PI0")
+    results.append(("PI0", pi0_success))
+
+    # Test PI0.5 model
+    print("\n\n[Test 2/2] Testing PI0.5 model...")
+    print("─" * 60)
+    pi05_success = test_hub_loading(model_id="pepijn223/pi05_base_fp32", model_name="PI0.5")
+    results.append(("PI0.5", pi05_success))
+
+    # Summary
+    print("\n\n")
+    print("╔" + "═" * 58 + "╗")
+    print("║" + "  TEST SUMMARY  ".center(58) + "║")
+    print("╚" + "═" * 58 + "╝")
+
+    all_passed = True
+    for model_name, success in results:
+        status = "✅ PASSED" if success else "❌ FAILED"
+        print(f"  {model_name:10} : {status}")
+        if not success:
+            all_passed = False
+
+    print()
+    if all_passed:
+        print("🎉 All models loaded and tested successfully!")
+    else:
+        print("⚠️  Some tests failed. Check the output above for details.")
+
+    return all_passed
+
+
+if __name__ == "__main__":
+    success = main()
+    exit(0 if success else 1)

tests/artifacts/datasets/save_dataset_to_safetensors.py

@@ -47,26 +47,38 @@ def save_dataset_to_safetensors(output_dir, repo_id="lerobot/pusht"):
     )
 
     # save 2 first frames of first episode
-    i = dataset.meta.episodes["dataset_from_index"][0].item()
+    i = dataset.episode_data_index["from"][0].item()
     save_file(dataset[i], repo_dir / f"frame_{i}.safetensors")
     save_file(dataset[i + 1], repo_dir / f"frame_{i + 1}.safetensors")
 
     # save 2 frames at the middle of first episode
-    i = int(
-        (
-            dataset.meta.episodes["dataset_to_index"][0].item()
-            - dataset.meta.episodes["dataset_from_index"][0].item()
-        )
-        / 2
-    )
+    i = int((dataset.episode_data_index["to"][0].item() - dataset.episode_data_index["from"][0].item()) / 2)
     save_file(dataset[i], repo_dir / f"frame_{i}.safetensors")
     save_file(dataset[i + 1], repo_dir / f"frame_{i + 1}.safetensors")
 
     # save 2 last frames of first episode
-    i = dataset.meta.episodes["dataset_to_index"][0].item()
+    i = dataset.episode_data_index["to"][0].item()
     save_file(dataset[i - 2], repo_dir / f"frame_{i - 2}.safetensors")
     save_file(dataset[i - 1], repo_dir / f"frame_{i - 1}.safetensors")
 
+    # TODO(rcadene): Enable testing on second and last episode
+    # We currently cant because our test dataset only contains the first episode
+
+    # # save 2 first frames of second episode
+    # i = dataset.episode_data_index["from"][1].item()
+    # save_file(dataset[i], repo_dir / f"frame_{i}.safetensors")
+    # save_file(dataset[i + 1], repo_dir / f"frame_{i+1}.safetensors")
+
+    # # save 2 last frames of second episode
+    # i = dataset.episode_data_index["to"][1].item()
+    # save_file(dataset[i - 2], repo_dir / f"frame_{i-2}.safetensors")
+    # save_file(dataset[i - 1], repo_dir / f"frame_{i-1}.safetensors")
+
+    # # save 2 last frames of last episode
+    # i = dataset.episode_data_index["to"][-1].item()
+    # save_file(dataset[i - 2], repo_dir / f"frame_{i-2}.safetensors")
+    # save_file(dataset[i - 1], repo_dir / f"frame_{i-1}.safetensors")
+
 
 if __name__ == "__main__":
     for dataset in [

tests/cameras/test_reachy2_camera.py

@@ -0,0 +1,177 @@
+#!/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 time
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+
+from lerobot.cameras.reachy2_camera import Reachy2Camera, Reachy2CameraConfig
+from lerobot.errors import DeviceNotConnectedError
+
+PARAMS = [
+    ("teleop", "left"),
+    ("teleop", "right"),
+    ("depth", "rgb"),
+    # ("depth", "depth"),  # Depth camera is not available yet
+]
+
+
+def _make_cam_manager_mock():
+    c = MagicMock(name="CameraManagerMock")
+
+    teleop = MagicMock(name="TeleopCam")
+    teleop.width = 640
+    teleop.height = 480
+    teleop.get_frame = MagicMock(
+        side_effect=lambda *_, **__: (
+            np.zeros((480, 640, 3), dtype=np.uint8),
+            time.time(),
+        )
+    )
+
+    depth = MagicMock(name="DepthCam")
+    depth.width = 640
+    depth.height = 480
+    depth.get_frame = MagicMock(
+        side_effect=lambda *_, **__: (
+            np.zeros((480, 640, 3), dtype=np.uint8),
+            time.time(),
+        )
+    )
+
+    c.is_connected.return_value = True
+    c.teleop = teleop
+    c.depth = depth
+
+    def _connect():
+        c.teleop = teleop
+        c.depth = depth
+        c.is_connected.return_value = True
+
+    def _disconnect():
+        c.teleop = None
+        c.depth = None
+        c.is_connected.return_value = False
+
+    c.connect = MagicMock(side_effect=_connect)
+    c.disconnect = MagicMock(side_effect=_disconnect)
+
+    # Mock methods
+    c.initialize_cameras = MagicMock()
+
+    return c
+
+
+@pytest.fixture(
+    params=PARAMS,
+    # ids=["teleop-left", "teleop-right", "torso-rgb", "torso-depth"],
+    ids=["teleop-left", "teleop-right", "torso-rgb"],
+)
+def camera(request):
+    name, image_type = request.param
+    with (
+        patch(
+            "lerobot.cameras.reachy2_camera.reachy2_camera.CameraManager",
+            side_effect=lambda *a, **k: _make_cam_manager_mock(),
+        ),
+    ):
+        config = Reachy2CameraConfig(name=name, image_type=image_type)
+        cam = Reachy2Camera(config)
+        yield cam
+        if cam.is_connected:
+            cam.disconnect()
+
+
+def test_connect(camera):
+    camera.connect()
+    assert camera.is_connected
+    camera.cam_manager.initialize_cameras.assert_called_once()
+
+
+def test_read(camera):
+    camera.connect()
+
+    img = camera.read()
+    if camera.config.name == "teleop":
+        camera.cam_manager.teleop.get_frame.assert_called_once()
+    elif camera.config.name == "depth":
+        camera.cam_manager.depth.get_frame.assert_called_once()
+    assert isinstance(img, np.ndarray)
+    assert img.shape == (480, 640, 3)
+
+
+def test_disconnect(camera):
+    camera.connect()
+
+    camera.disconnect()
+    assert not camera.is_connected
+
+
+def test_async_read(camera):
+    camera.connect()
+    try:
+        img = camera.async_read()
+
+        assert camera.thread is not None
+        assert camera.thread.is_alive()
+        assert isinstance(img, np.ndarray)
+    finally:
+        if camera.is_connected:
+            camera.disconnect()
+
+
+def test_async_read_timeout(camera):
+    camera.connect()
+    try:
+        with pytest.raises(TimeoutError):
+            camera.async_read(timeout_ms=0)
+    finally:
+        if camera.is_connected:
+            camera.disconnect()
+
+
+def test_read_before_connect(camera):
+    with pytest.raises(DeviceNotConnectedError):
+        _ = camera.read()
+
+
+def test_disconnect_before_connect(camera):
+    with pytest.raises(DeviceNotConnectedError):
+        camera.disconnect()
+
+
+def test_async_read_before_connect(camera):
+    with pytest.raises(DeviceNotConnectedError):
+        _ = camera.async_read()
+
+
+def test_wrong_camera_name():
+    with pytest.raises(ValueError):
+        _ = Reachy2CameraConfig(name="wrong-name", image_type="left")
+
+
+def test_wrong_image_type():
+    with pytest.raises(ValueError):
+        _ = Reachy2CameraConfig(name="teleop", image_type="rgb")
+    with pytest.raises(ValueError):
+        _ = Reachy2CameraConfig(name="depth", image_type="left")
+
+
+def test_wrong_color_mode():
+    with pytest.raises(ValueError):
+        _ = Reachy2CameraConfig(name="teleop", image_type="left", color_mode="wrong-color")

tests/conftest.py

@@ -28,6 +28,7 @@ pytest_plugins = [
     "tests.fixtures.files",
     "tests.fixtures.hub",
     "tests.fixtures.optimizers",
+    "tests.plugins.reachy2_sdk",
 ]
 
 

tests/datasets/test_aggregate.py

@@ -1,292 +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 unittest.mock import patch
-
-import torch
-
-from lerobot.datasets.aggregate import aggregate_datasets
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from tests.fixtures.constants import DUMMY_REPO_ID
-
-
-def assert_episode_and_frame_counts(aggr_ds, expected_episodes, expected_frames):
-    """Test that total number of episodes and frames are correctly aggregated."""
-    assert aggr_ds.num_episodes == expected_episodes, (
-        f"Expected {expected_episodes} episodes, got {aggr_ds.num_episodes}"
-    )
-    assert aggr_ds.num_frames == expected_frames, (
-        f"Expected {expected_frames} frames, got {aggr_ds.num_frames}"
-    )
-
-
-def assert_dataset_content_integrity(aggr_ds, ds_0, ds_1):
-    """Test that the content of both datasets is preserved correctly in the aggregated dataset."""
-    keys_to_ignore = ["episode_index", "index", "timestamp"]
-
-    # Test first part of dataset corresponds to ds_0, check first item (index 0) matches ds_0[0]
-    aggr_first_item = aggr_ds[0]
-    ds_0_first_item = ds_0[0]
-
-    # Compare all keys except episode_index and index which should be updated
-    for key in ds_0_first_item:
-        if key not in keys_to_ignore:
-            # Handle both tensor and non-tensor data
-            if torch.is_tensor(aggr_first_item[key]) and torch.is_tensor(ds_0_first_item[key]):
-                assert torch.allclose(aggr_first_item[key], ds_0_first_item[key], atol=1e-6), (
-                    f"First item key '{key}' doesn't match between aggregated and ds_0"
-                )
-            else:
-                assert aggr_first_item[key] == ds_0_first_item[key], (
-                    f"First item key '{key}' doesn't match between aggregated and ds_0"
-                )
-
-    # Check last item of ds_0 part (index len(ds_0)-1) matches ds_0[-1]
-    aggr_ds_0_last_item = aggr_ds[len(ds_0) - 1]
-    ds_0_last_item = ds_0[-1]
-
-    for key in ds_0_last_item:
-        if key not in keys_to_ignore:
-            # Handle both tensor and non-tensor data
-            if torch.is_tensor(aggr_ds_0_last_item[key]) and torch.is_tensor(ds_0_last_item[key]):
-                assert torch.allclose(aggr_ds_0_last_item[key], ds_0_last_item[key], atol=1e-6), (
-                    f"Last ds_0 item key '{key}' doesn't match between aggregated and ds_0"
-                )
-            else:
-                assert aggr_ds_0_last_item[key] == ds_0_last_item[key], (
-                    f"Last ds_0 item key '{key}' doesn't match between aggregated and ds_0"
-                )
-
-    # Test second part of dataset corresponds to ds_1
-    # Check first item of ds_1 part (index len(ds_0)) matches ds_1[0]
-    aggr_ds_1_first_item = aggr_ds[len(ds_0)]
-    ds_1_first_item = ds_1[0]
-
-    for key in ds_1_first_item:
-        if key not in keys_to_ignore:
-            # Handle both tensor and non-tensor data
-            if torch.is_tensor(aggr_ds_1_first_item[key]) and torch.is_tensor(ds_1_first_item[key]):
-                assert torch.allclose(aggr_ds_1_first_item[key], ds_1_first_item[key], atol=1e-6), (
-                    f"First ds_1 item key '{key}' doesn't match between aggregated and ds_1"
-                )
-            else:
-                assert aggr_ds_1_first_item[key] == ds_1_first_item[key], (
-                    f"First ds_1 item key '{key}' doesn't match between aggregated and ds_1"
-                )
-
-    # Check last item matches ds_1[-1]
-    aggr_last_item = aggr_ds[-1]
-    ds_1_last_item = ds_1[-1]
-
-    for key in ds_1_last_item:
-        if key not in keys_to_ignore:
-            # Handle both tensor and non-tensor data
-            if torch.is_tensor(aggr_last_item[key]) and torch.is_tensor(ds_1_last_item[key]):
-                assert torch.allclose(aggr_last_item[key], ds_1_last_item[key], atol=1e-6), (
-                    f"Last item key '{key}' doesn't match between aggregated and ds_1"
-                )
-            else:
-                assert aggr_last_item[key] == ds_1_last_item[key], (
-                    f"Last item key '{key}' doesn't match between aggregated and ds_1"
-                )
-
-
-def assert_metadata_consistency(aggr_ds, ds_0, ds_1):
-    """Test that metadata is correctly aggregated."""
-    # Test basic info
-    assert aggr_ds.fps == ds_0.fps == ds_1.fps, "FPS should be the same across all datasets"
-    assert aggr_ds.meta.info["robot_type"] == ds_0.meta.info["robot_type"] == ds_1.meta.info["robot_type"], (
-        "Robot type should be the same"
-    )
-
-    # Test features are the same
-    assert aggr_ds.features == ds_0.features == ds_1.features, "Features should be the same"
-
-    # Test tasks aggregation
-    expected_tasks = set(ds_0.meta.tasks.index) | set(ds_1.meta.tasks.index)
-    actual_tasks = set(aggr_ds.meta.tasks.index)
-    assert actual_tasks == expected_tasks, f"Expected tasks {expected_tasks}, got {actual_tasks}"
-
-
-def assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1):
-    """Test that episode indices are correctly updated after aggregation."""
-    # ds_0 episodes should have episode_index 0 to ds_0.num_episodes-1
-    for i in range(len(ds_0)):
-        assert aggr_ds[i]["episode_index"] < ds_0.num_episodes, (
-            f"Episode index {aggr_ds[i]['episode_index']} at position {i} should be < {ds_0.num_episodes}"
-        )
-
-    def ds1_episodes_condition(ep_idx):
-        return (ep_idx >= ds_0.num_episodes) and (ep_idx < ds_0.num_episodes + ds_1.num_episodes)
-
-    # ds_1 episodes should have episode_index ds_0.num_episodes to total_episodes-1
-    for i in range(len(ds_0), len(ds_0) + len(ds_1)):
-        expected_min_episode_idx = ds_0.num_episodes
-        assert ds1_episodes_condition(aggr_ds[i]["episode_index"]), (
-            f"Episode index {aggr_ds[i]['episode_index']} at position {i} should be >= {expected_min_episode_idx}"
-        )
-
-
-def assert_video_frames_integrity(aggr_ds, ds_0, ds_1):
-    """Test that video frames are correctly preserved and frame indices are updated."""
-
-    def visual_frames_equal(frame1, frame2):
-        return torch.allclose(frame1, frame2)
-
-    video_keys = list(
-        filter(
-            lambda key: aggr_ds.meta.info["features"][key]["dtype"] == "video",
-            aggr_ds.meta.info["features"].keys(),
-        )
-    )
-
-    # Test the section corresponding to the first dataset (ds_0)
-    for i in range(len(ds_0)):
-        assert aggr_ds[i]["index"] == i, (
-            f"Frame index at position {i} should be {i}, but got {aggr_ds[i]['index']}"
-        )
-        for key in video_keys:
-            assert visual_frames_equal(aggr_ds[i][key], ds_0[i][key]), (
-                f"Visual frames at position {i} should be equal between aggregated and ds_0"
-            )
-
-    # Test the section corresponding to the second dataset (ds_1)
-    for i in range(len(ds_0), len(ds_0) + len(ds_1)):
-        # The frame index in the aggregated dataset should also match its position.
-        assert aggr_ds[i]["index"] == i, (
-            f"Frame index at position {i} should be {i}, but got {aggr_ds[i]['index']}"
-        )
-        for key in video_keys:
-            assert visual_frames_equal(aggr_ds[i][key], ds_1[i - len(ds_0)][key]), (
-                f"Visual frames at position {i} should be equal between aggregated and ds_1"
-            )
-
-
-def assert_dataset_iteration_works(aggr_ds):
-    """Test that we can iterate through the entire dataset without errors."""
-    for _ in aggr_ds:
-        pass
-
-
-def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
-    """Test basic aggregation functionality with standard parameters."""
-    ds_0_num_frames = 400
-    ds_1_num_frames = 800
-    ds_0_num_episodes = 10
-    ds_1_num_episodes = 25
-
-    # Create two datasets with different number of frames and episodes
-    ds_0 = lerobot_dataset_factory(
-        root=tmp_path / "test_0",
-        repo_id=f"{DUMMY_REPO_ID}_0",
-        total_episodes=ds_0_num_episodes,
-        total_frames=ds_0_num_frames,
-    )
-    ds_1 = lerobot_dataset_factory(
-        root=tmp_path / "test_1",
-        repo_id=f"{DUMMY_REPO_ID}_1",
-        total_episodes=ds_1_num_episodes,
-        total_frames=ds_1_num_frames,
-    )
-
-    aggregate_datasets(
-        repo_ids=[ds_0.repo_id, ds_1.repo_id],
-        roots=[ds_0.root, ds_1.root],
-        aggr_repo_id=f"{DUMMY_REPO_ID}_aggr",
-        aggr_root=tmp_path / "test_aggr",
-    )
-
-    # Mock the revision to prevent Hub calls during dataset loading
-    with (
-        patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
-        patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
-    ):
-        mock_get_safe_version.return_value = "v3.0"
-        mock_snapshot_download.return_value = str(tmp_path / "test_aggr")
-        aggr_ds = LeRobotDataset(f"{DUMMY_REPO_ID}_aggr", root=tmp_path / "test_aggr")
-
-    # Run all assertion functions
-    expected_total_episodes = ds_0.num_episodes + ds_1.num_episodes
-    expected_total_frames = ds_0.num_frames + ds_1.num_frames
-
-    assert_episode_and_frame_counts(aggr_ds, expected_total_episodes, expected_total_frames)
-    assert_dataset_content_integrity(aggr_ds, ds_0, ds_1)
-    assert_metadata_consistency(aggr_ds, ds_0, ds_1)
-    assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1)
-    assert_video_frames_integrity(aggr_ds, ds_0, ds_1)
-    assert_dataset_iteration_works(aggr_ds)
-
-
-def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory):
-    """Test aggregation with small file size limits to force file rotation/sharding."""
-    ds_0_num_episodes = ds_1_num_episodes = 10
-    ds_0_num_frames = ds_1_num_frames = 400
-
-    ds_0 = lerobot_dataset_factory(
-        root=tmp_path / "small_0",
-        repo_id=f"{DUMMY_REPO_ID}_small_0",
-        total_episodes=ds_0_num_episodes,
-        total_frames=ds_0_num_frames,
-    )
-    ds_1 = lerobot_dataset_factory(
-        root=tmp_path / "small_1",
-        repo_id=f"{DUMMY_REPO_ID}_small_1",
-        total_episodes=ds_1_num_episodes,
-        total_frames=ds_1_num_frames,
-    )
-
-    # Use the new configurable parameters to force file rotation
-    aggregate_datasets(
-        repo_ids=[ds_0.repo_id, ds_1.repo_id],
-        roots=[ds_0.root, ds_1.root],
-        aggr_repo_id=f"{DUMMY_REPO_ID}_small_aggr",
-        aggr_root=tmp_path / "small_aggr",
-        # Tiny file size to trigger new file instantiation
-        data_files_size_in_mb=0.01,
-        video_files_size_in_mb=0.1,
-    )
-
-    # Mock the revision to prevent Hub calls during dataset loading
-    with (
-        patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
-        patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
-    ):
-        mock_get_safe_version.return_value = "v3.0"
-        mock_snapshot_download.return_value = str(tmp_path / "small_aggr")
-        aggr_ds = LeRobotDataset(f"{DUMMY_REPO_ID}_small_aggr", root=tmp_path / "small_aggr")
-
-    # Verify aggregation worked correctly despite file size constraints
-    expected_total_episodes = ds_0_num_episodes + ds_1_num_episodes
-    expected_total_frames = ds_0_num_frames + ds_1_num_frames
-
-    assert_episode_and_frame_counts(aggr_ds, expected_total_episodes, expected_total_frames)
-    assert_dataset_content_integrity(aggr_ds, ds_0, ds_1)
-    assert_metadata_consistency(aggr_ds, ds_0, ds_1)
-    assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1)
-    assert_video_frames_integrity(aggr_ds, ds_0, ds_1)
-    assert_dataset_iteration_works(aggr_ds)
-
-    # Check that multiple files were actually created due to small size limits
-    data_dir = tmp_path / "small_aggr" / "data"
-    video_dir = tmp_path / "small_aggr" / "videos"
-
-    if data_dir.exists():
-        parquet_files = list(data_dir.rglob("*.parquet"))
-        assert len(parquet_files) > 1, "Small file size limits should create multiple parquet files"
-
-    if video_dir.exists():
-        video_files = list(video_dir.rglob("*.mp4"))
-        assert len(video_files) > 1, "Small file size limits should create multiple video files"