Merge branch 'main' into feat/add_macos_ci

chore(deps): libero dep pointing to main (#2201 )
feat(sim): add metaworld env (#2088 )
2026-06-17 08:17:02 +00:00 · 2025-10-14 18:52:04 +02:00 · 2025-10-14 18:19:49 +02:00 · 2025-10-14 17:21:18 +02:00 · 2025-10-14 16:19:50 +02:00 · 2025-10-14 16:00:42 +02:00
129 changed files with 11099 additions and 4019 deletions
@@ -57,7 +57,11 @@ jobs:
  # It runs everytime we commit to a PR or push to main
  fast-pytest-tests:
    name: Fast Pytest Tests
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-latest]
    env:
      MUJOCO_GL: egl
    steps:
@@ -67,12 +71,21 @@ jobs:
          lfs: true

      # TODO(Steven): Evaluate the need of these dependencies
-      - name: Install apt dependencies
+      - name: Install dependencies
        run: |
-          sudo apt-get update && sudo apt-get install -y build-essential git \
-          curl libglib2.0-0 libegl1-mesa-dev ffmpeg \
-          libusb-1.0-0-dev speech-dispatcher libgeos-dev portaudio19-dev
-
+          if [[ "${{ matrix.os }}" == 'ubuntu-latest' ]]; then
+            sudo apt-get update && sudo apt-get install -y build-essential \
+            git curl libglib2.0-0 libegl1-mesa-dev ffmpeg libusb-1.0-0-dev \
+            speech-dispatcher libgeos-dev portaudio19-dev
+          elif [[ "${{ matrix.os }}" == 'macos-latest' ]]; then
+            brew update && brew install git geos portaudio ffmpeg@7
+            # Add ffmpeg@7 paths for subsequent steps
+            echo "PATH=/opt/homebrew/opt/ffmpeg@7/bin:$PATH" >> $GITHUB_ENV
+            echo "LDFLAGS=-L/opt/homebrew/opt/ffmpeg@7/lib" >> $GITHUB_ENV
+            echo "CPPFLAGS=-I/opt/homebrew/opt/ffmpeg@7/include" >> $GITHUB_ENV
+            echo "PKG_CONFIG_PATH=/opt/homebrew/opt/ffmpeg@7/lib/pkgconfig" >> $GITHUB_ENV
+            echo "DYLD_LIBRARY_PATH=/opt/homebrew/opt/ffmpeg@7/lib:/opt/homebrew/lib:/usr/local/lib:$DYLD_LIBRARY_PATH" >> $GITHUB_ENV
+          fi
      - name: Setup uv and Python
        uses: astral-sh/setup-uv@v6 # zizmor: ignore[unpinned-uses]
        with:
@@ -51,7 +51,11 @@ jobs:
  # It runs everytime a PR is approved or a push to main
  full-tests:
    name: Full Tests
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-latest]
    if: |
      (github.event_name == 'pull_request_review' && github.event.review.state == 'approved') ||
      github.event_name == 'push' ||
@@ -64,11 +68,16 @@ jobs:
          lfs: true
          persist-credentials: false

-      - name: Install apt dependencies
+      - name: Install dependencies
        run: |
-          sudo apt-get update && sudo apt-get install -y build-essential \
-          git curl libglib2.0-0 libegl1-mesa-dev ffmpeg libusb-1.0-0-dev \
-          speech-dispatcher libgeos-dev portaudio19-dev
+          if [[ "${{ matrix.os }}" == 'ubuntu-latest' ]]; then
+            sudo apt-get update && sudo apt-get install -y build-essential \
+            git curl libglib2.0-0 libegl1-mesa-dev ffmpeg libusb-1.0-0-dev \
+            speech-dispatcher libgeos-dev portaudio19-dev
+          elif [[ "${{ matrix.os }}" == 'macos-latest' ]]; then
+            brew update && brew install git geos portaudio ffmpeg@7
+            echo "DYLD_LIBRARY_PATH=/opt/homebrew/opt/ffmpeg@7/lib:/opt/homebrew/lib:/usr/local/lib:$DYLD_LIBRARY_PATH" >> $GITHUB_ENV
+          fi

      - name: Setup uv and Python
        uses: astral-sh/setup-uv@v6 # zizmor: ignore[unpinned-uses]
@@ -120,7 +120,11 @@ jobs:
  test-release:
    name: Test Release
    needs: [build-and-publish]
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-latest]
    permissions:
      contents: read
    env:
@@ -130,11 +134,16 @@ jobs:
        with:
          lfs: true
          persist-credentials: false
-      - name: Install apt dependencies
+      - name: Install dependencies
        run: |
-          sudo apt-get update && sudo apt-get install -y build-essential \
-          git curl libglib2.0-0 libegl1-mesa-dev ffmpeg libusb-1.0-0-dev \
-          speech-dispatcher libgeos-dev portaudio19-dev
+          if [[ "${{ matrix.os }}" == 'ubuntu-latest' ]]; then
+            sudo apt-get update && sudo apt-get install -y build-essential \
+            git curl libglib2.0-0 libegl1-mesa-dev ffmpeg libusb-1.0-0-dev \
+            speech-dispatcher libgeos-dev portaudio19-dev
+          elif [[ "${{ matrix.os }}" == 'macos-latest' ]]; then
+            brew update && brew install git geos portaudio ffmpeg@7
+            echo "DYLD_LIBRARY_PATH=/opt/homebrew/opt/ffmpeg@7/lib:/opt/homebrew/lib:/usr/local/lib:$DYLD_LIBRARY_PATH" >> $GITHUB_ENV
+          fi
      - name: Setup uv and Python
        uses: astral-sh/setup-uv@v6 # zizmor: ignore[unpinned-uses]
        with:
@@ -0,0 +1,192 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This workflow handles full testing with unboud dependencies versions.
+name: Unbound Dependency Tests
+
+on:
+  # Allows running this workflow manually from the Actions tab
+  workflow_dispatch:
+
+  # Run on the 1st and 15th of every month at 09:00 UTC
+  schedule:
+    - cron: '0 2 1,15 * *'
+
+permissions:
+  contents: read
+
+# Sets up the environment variables
+env:
+  UV_VERSION: "0.8.0"
+  PYTHON_VERSION: "3.10"
+  DOCKER_IMAGE_NAME: huggingface/lerobot-gpu:unbound
+
+# Ensures that only the latest action is built, canceling older runs.
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+
+  # This job runs the E2E tests + pytest with all unbound extras
+  full-tests:
+    name: Full Unbound Tests
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+    env:
+      MUJOCO_GL: egl
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          lfs: true
+          persist-credentials: false
+
+      - name: Install dependencies
+        run: |
+          if [[ "${{ matrix.os }}" == 'ubuntu-latest' ]]; then
+            sudo apt-get update && sudo apt-get install -y build-essential \
+            git curl libglib2.0-0 libegl1-mesa-dev ffmpeg libusb-1.0-0-dev \
+            speech-dispatcher libgeos-dev portaudio19-dev
+          elif [[ "${{ matrix.os }}" == 'macos-latest' ]]; then
+            brew update && brew install git geos portaudio ffmpeg@7
+            echo "DYLD_LIBRARY_PATH=/opt/homebrew/opt/ffmpeg@7/lib:/opt/homebrew/lib:/usr/local/lib:$DYLD_LIBRARY_PATH" >> $GITHUB_ENV
+          fi
+
+      - name: Setup uv and Python
+        uses: astral-sh/setup-uv@v6 # zizmor: ignore[unpinned-uses]
+        with:
+          enable-cache: true
+          version: ${{ env.UV_VERSION }}
+          python-version: ${{ env.PYTHON_VERSION }}
+
+      - name: Unbound dependencies
+        run: |
+          sed -i 's/,[[:space:]]*<[0-9\.]*//g' pyproject.toml
+          echo "Dependencies unbound:" && cat pyproject.toml
+
+      - name: Install lerobot with all extras
+        run: uv sync --all-extras
+
+      - name: Run pytest (all extras)
+        run: uv run pytest tests -vv
+
+      - name: Run end-to-end tests
+        run: uv run make test-end-to-end
+
+  # This job builds a GPU enabled image for testing
+  build-and-push-docker:
+    name: Build and Push Docker
+    runs-on:
+      group: aws-general-8-plus
+    outputs:
+      image_tag: ${{ env.DOCKER_IMAGE_NAME }}
+    env:
+      GITHUB_REF: ${{ github.ref }}
+    steps:
+      - name: Install Git LFS
+        run: |
+          sudo apt-get update
+          sudo apt-get install git-lfs
+          git lfs install
+      - uses: actions/checkout@v4
+        with:
+          lfs: true
+          persist-credentials: false
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3 # zizmor: ignore[unpinned-uses]
+        with:
+          cache-binary: false
+      - name: Login to Docker Hub
+        uses: docker/login-action@v3 # zizmor: ignore[unpinned-uses]
+        with:
+          username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
+      - name: Build and push Docker image
+        uses: docker/build-push-action@v6 # zizmor: ignore[unpinned-uses]
+        with:
+          context: .
+          file: ./docker/Dockerfile.internal
+          push: true
+          tags: ${{ env.DOCKER_IMAGE_NAME }}
+          build-args: |
+            UNBOUND_DEPS=true
+
+  # This job runs pytest with all unbound extras in a GPU enabled host
+  # It runs everytime a test image is created
+  gpu-tests:
+    name: GPU Unbound Tests
+    needs: [build-and-push-docker]
+    runs-on:
+      group: aws-g6-4xlarge-plus
+    env:
+      HF_HOME: /home/user_lerobot/.cache/huggingface
+      HF_LEROBOT_HOME: /home/user_lerobot/.cache/huggingface/lerobot
+      TORCH_HOME: /home/user_lerobot/.cache/torch
+      TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton
+    container:
+      image: ${{ needs.build-and-push-docker.outputs.image_tag }} # zizmor: ignore[unpinned-images]
+      options: --gpus all --shm-size "16gb"
+      credentials:
+        username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
+        password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
+    defaults:
+      run:
+        shell: bash
+        working-directory: /lerobot
+    steps:
+      - name: Run pytest on GPU
+        run: pytest tests -vv
+      - name: Run end-to-end tests
+        run: make test-end-to-end
+
+  # This job deletes the test image recently created
+  # It runs everytime after the gpu-tests have finished
+  delete-unbound-image:
+    name: Delete Unbound Image
+    needs: [gpu-tests, build-and-push-docker]
+    if: always() && needs.build-and-push-docker.result == 'success'
+    runs-on: ubuntu-latest
+    steps:
+      - name: Get Docker Hub Token and Delete Image
+        # zizmor: ignore[template-injection]
+        run: |
+          IMAGE_NAME=$(echo "${{ needs.build-and-push-docker.outputs.image_tag }}" | cut -d':' -f1)
+          IMAGE_TAG=$(echo "${{ needs.build-and-push-docker.outputs.image_tag }}" | cut -d':' -f2)
+
+          echo "Attempting to delete image: $IMAGE_NAME:$IMAGE_TAG"
+
+          TOKEN=$(curl -s -H "Content-Type: application/json" \
+                       -X POST \
+                       -d '{"username": "${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}", "password": "${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}"}' \
+                       https://hub.docker.com/v2/users/login/ | jq -r .token)
+
+          if [ "$TOKEN" == "null" ] || [ -z "$TOKEN" ]; then
+            echo "::error::Failed to get Docker Hub token."
+            exit 1
+          fi
+
+          HTTP_RESPONSE=$(curl -s -o /dev/null -w "%{http_code}" \
+                               -H "Authorization: JWT ${TOKEN}" \
+                               -X DELETE \
+                               https://hub.docker.com/v2/repositories/${IMAGE_NAME}/tags/${IMAGE_TAG}/)
+
+          if [ "$HTTP_RESPONSE" -eq 204 ]; then
+            echo "Successfully deleted Docker image tag: $IMAGE_NAME:$IMAGE_TAG"
+          else
+            echo "::error::Failed to delete Docker image. HTTP status: $HTTP_RESPONSE"
+            exit 1
+          fi
@@ -86,11 +86,12 @@ repos:

  # TODO(Steven): Uncomment when ready to use
  ##### Static Analysis & Typing #####
-  # - repo: https://github.com/pre-commit/mirrors-mypy
-  #   rev: v1.16.0
-  #   hooks:
-  #     - id: mypy
-  #       args: [--python-version=3.10]
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.16.0
+    hooks:
+      - id: mypy
+        args: [--config-file=pyproject.toml]
+        exclude: ^(examples|benchmarks|tests)/

  ##### Docstring Checks #####
  # - repo: https://github.com/akaihola/darglint2
@@ -72,7 +72,6 @@ post it.

 Look at our implementations for [datasets](./src/lerobot/datasets/), [policies](./src/lerobot/policies/),
 environments ([aloha](https://github.com/huggingface/gym-aloha),
-[xarm](https://github.com/huggingface/gym-xarm),
 [pusht](https://github.com/huggingface/gym-pusht))
 and follow the same api design.

@@ -119,10 +119,9 @@ test-tdmpc-ete-train:
 		--policy.type=tdmpc \
 		--policy.device=$(DEVICE) \
 		--policy.push_to_hub=false \
-		--env.type=xarm \
-		--env.task=XarmLift-v0 \
+		--env.type=pusht \
 		--env.episode_length=5 \
-		--dataset.repo_id=lerobot/xarm_lift_medium \
+		--dataset.repo_id=lerobot/pusht_image \
 		--dataset.image_transforms.enable=true \
 		--dataset.episodes="[0]" \
 		--batch_size=2 \
@@ -140,9 +139,10 @@ test-tdmpc-ete-eval:
 	lerobot-eval \
 		--policy.path=tests/outputs/tdmpc/checkpoints/000002/pretrained_model \
 		--policy.device=$(DEVICE) \
-		--env.type=xarm \
+		--env.type=pusht \
 		--env.episode_length=5 \
-		--env.task=XarmLift-v0 \
+		--env.observation_height=96 \
+        --env.observation_width=96 \
 		--eval.n_episodes=1 \
 		--eval.batch_size=1

@@ -197,7 +197,7 @@ wandb login

 ### Visualize datasets

-Check out [example 1](https://github.com/huggingface/lerobot/blob/main/examples/1_load_lerobot_dataset.py) that illustrates how to use our dataset class which automatically downloads data from the Hugging Face hub.
+Check out [example 1](https://github.com/huggingface/lerobot/blob/main/examples/dataset/load_lerobot_dataset.py) that illustrates how to use our dataset class which automatically downloads data from the Hugging Face hub.

 You can also locally visualize episodes from a dataset on the hub by executing our script from the command line:

@@ -75,6 +75,14 @@ RUN uv venv --python python${PYTHON_VERSION}
 # Install Python dependencies for caching
 COPY --chown=user_lerobot:user_lerobot pyproject.toml README.md MANIFEST.in ./
 COPY --chown=user_lerobot:user_lerobot src/ src/
+
+ARG UNBOUND_DEPS=false
+
+RUN if [ "$UNBOUND_DEPS" = "true" ]; then \
+    sed -i 's/,[[:space:]]*<[0-9\.]*//g' pyproject.toml; \
+    echo "Dependencies unbound:" && cat pyproject.toml; \
+    fi
+
 RUN uv pip install --no-cache ".[all]"

 # Copy the rest of the application source code
@@ -61,6 +61,14 @@ RUN uv venv
 # Install Python dependencies for caching
 COPY --chown=user_lerobot:user_lerobot pyproject.toml README.md MANIFEST.in ./
 COPY --chown=user_lerobot:user_lerobot src/ src/
+
+ARG UNBOUND_DEPS=false
+
+RUN if [ "$UNBOUND_DEPS" = "true" ]; then \
+    sed -i 's/,[[:space:]]*<[0-9\.]*//g' pyproject.toml; \
+    echo "Dependencies unbound:" && cat pyproject.toml; \
+    fi
+
 RUN uv pip install --no-cache ".[all]"

 # Copy the rest of the application code
@@ -7,8 +7,6 @@
 - sections:
  - local: il_robots
    title: Imitation Learning for Robots
-  - local: il_sim
-    title: Imitation Learning in Sim
  - local: cameras
    title: Cameras
  - local: integrate_hardware
@@ -25,14 +23,27 @@
    title: Using LeRobotDataset
  - local: porting_datasets_v3
    title: Porting Large Datasets
+  - local: using_dataset_tools
+    title: Using the Dataset Tools
  title: "Datasets"
 - sections:
+  - local: act
+    title: ACT
  - local: smolvla
-    title: Finetune SmolVLA
+    title: SmolVLA
+  - local: pi0
+    title: π₀ (Pi0)
+  - local: pi05
+    title: π₀.₅ (Pi05)
+  title: "Policies"
+- sections:
+  - local: il_sim
+    title: Imitation Learning in Sim
  - local: libero
    title: Using Libero
-  title: "Policies"
-
+  - local: metaworld
+    title: Using MetaWorld
+  title: "Simulation"
 - sections:
  - local: introduction_processors
    title: Introduction to Robot Processors
@@ -0,0 +1,92 @@
+# ACT (Action Chunking with Transformers)
+
+ACT is a **lightweight and efficient policy for imitation learning**, especially well-suited for fine-grained manipulation tasks. It's the **first model we recommend when you're starting out** with LeRobot due to its fast training time, low computational requirements, and strong performance.
+
+<div class="video-container">
+  <iframe
+    width="100%"
+    height="415"
+    src="https://www.youtube.com/embed/ft73x0LfGpM"
+    title="LeRobot ACT Tutorial"
+    frameborder="0"
+    allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture"
+    allowfullscreen
+  ></iframe>
+</div>
+
+_Watch this tutorial from the LeRobot team to learn how ACT works: [LeRobot ACT Tutorial](https://www.youtube.com/watch?v=ft73x0LfGpM)_
+
+## Model Overview
+
+Action Chunking with Transformers (ACT) was introduced in the paper [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://arxiv.org/abs/2304.13705) by Zhao et al. The policy was designed to enable precise, contact-rich manipulation tasks using affordable hardware and minimal demonstration data.
+
+### Why ACT is Great for Beginners
+
+ACT stands out as an excellent starting point for several reasons:
+
+- **Fast Training**: Trains in a few hours on a single GPU
+- **Lightweight**: Only ~80M parameters, making it efficient and easy to work with
+- **Data Efficient**: Often achieves high success rates with just 50 demonstrations
+
+### Architecture
+
+ACT uses a transformer-based architecture with three main components:
+
+1. **Vision Backbone**: ResNet-18 processes images from multiple camera viewpoints
+2. **Transformer Encoder**: Synthesizes information from camera features, joint positions, and a learned latent variable
+3. **Transformer Decoder**: Generates coherent action sequences using cross-attention
+
+The policy takes as input:
+
+- Multiple RGB images (e.g., from wrist cameras, front/top cameras)
+- Current robot joint positions
+- A latent style variable `z` (learned during training, set to zero during inference)
+
+And outputs a chunk of `k` future action sequences.
+
+## Installation Requirements
+
+1. Install LeRobot by following our [Installation Guide](./installation).
+2. ACT is included in the base LeRobot installation, so no additional dependencies are needed!
+
+## Training ACT
+
+ACT works seamlessly with the standard LeRobot training pipeline. Here's a complete example for training ACT on your dataset:
+
+```bash
+lerobot-train \
+  --dataset.repo_id=${HF_USER}/your_dataset \
+  --policy.type=act \
+  --output_dir=outputs/train/act_your_dataset \
+  --job_name=act_your_dataset \
+  --policy.device=cuda \
+  --wandb.enable=true \
+  --policy.repo_id=${HF_USER}/act_policy
+```
+
+### Training Tips
+
+1. **Start with defaults**: ACT's default hyperparameters work well for most tasks
+2. **Training duration**: Expect a few hours for 100k training steps on a single GPU
+3. **Batch size**: Start with batch size 8 and adjust based on your GPU memory
+
+### Train using Google Colab
+
+If your local computer doesn't have a powerful GPU, you can utilize Google Colab to train your model by following the [ACT training notebook](./notebooks#training-act).
+
+## Evaluating ACT
+
+Once training is complete, you can evaluate your ACT policy using the `lerobot-record` command with your trained policy. This will run inference and record evaluation episodes:
+
+```bash
+lerobot-record \
+  --robot.type=so100_follower \
+  --robot.port=/dev/ttyACM0 \
+  --robot.id=my_robot \
+  --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
+  --display_data=true \
+  --dataset.repo_id=${HF_USER}/eval_act_your_dataset \
+  --dataset.num_episodes=10 \
+  --dataset.single_task="Your task description" \
+  --policy.path=${HF_USER}/act_policy
+```
@@ -31,15 +31,15 @@ Then, spin up a policy server (in one terminal, or in a separate machine) specif
 You can spin up a policy server running:

 ```shell
-python src/lerobot/async_inference/policy_server.py \
-    --host=127.0.0.1 \
-    --port=8080 \
+python -m lerobot.async_inference.policy_server \
+     --host=127.0.0.1 \
+     --port=8080
 ```

 This will start a policy server listening on `127.0.0.1:8080` (`localhost`, port 8080). At this stage, the policy server is empty, as all information related to which policy to run and with which parameters are specified during the first handshake with the client. Spin up a client with:

 ```shell
-python src/lerobot/async_inference/robot_client.py \
+python -m lerobot.async_inference.robot_client \
    --server_address=127.0.0.1:8080 \ # SERVER: the host address and port of the policy server
    --robot.type=so100_follower \ # ROBOT: your robot type
    --robot.port=/dev/tty.usbmodem585A0076841 \ # ROBOT: your robot port
@@ -113,9 +113,9 @@ As such, spinning up a policy server is as easy as specifying the host address a
 <hfoptions id="start_policy_server">
 <hfoption id="Command">
 ```bash
-python -m lerobot.scripts.server.policy_server \
-    --host="localhost" \
-    --port=8080
+python -m lerobot.async_inference.policy_server \
+     --host=127.0.0.1 \
+     --port=8080
 ```
 </hfoption>
 <hfoption id="API example">
@@ -148,7 +148,7 @@ The `RobotClient` streams observations to the `PolicyServer`, and receives actio
 <hfoptions id="start_robot_client">
 <hfoption id="Command">
 ```bash
-python src/lerobot/async_inference/robot_client.py \
+python -m lerobot.async_inference.robot_client \
    --server_address=127.0.0.1:8080 \ # SERVER: the host address and port of the policy server
    --robot.type=so100_follower \ # ROBOT: your robot type
    --robot.port=/dev/tty.usbmodem585A0076841 \ # ROBOT: your robot port
@@ -513,13 +513,14 @@ from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.policies.act.modeling_act import ACTPolicy
+from lerobot.policies.factory import make_pre_post_processors
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
+from lerobot.scripts.lerobot_record import record_loop
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
 from lerobot.utils.visualization_utils import init_rerun
-from lerobot.record import record_loop
-from lerobot.policies.factory import make_processor
+

 NUM_EPISODES = 5
 FPS = 30
@@ -562,7 +563,7 @@ init_rerun(session_name="recording")
 # Connect the robot
 robot.connect()

-preprocessor, postprocessor = make_processor(
+preprocessor, postprocessor = make_pre_post_processors(
    policy_cfg=policy,
    pretrained_path=HF_MODEL_ID,
    dataset_stats=dataset.meta.stats,
@@ -91,7 +91,7 @@ LeRobot provides optional extras for specific functionalities. Multiple extras c

 ### Simulations

-Install environment packages: `aloha` ([gym-aloha](https://github.com/huggingface/gym-aloha)), `xarm` ([gym-xarm](https://github.com/huggingface/gym-xarm)), or `pusht` ([gym-pusht](https://github.com/huggingface/gym-pusht))
+Install environment packages: `aloha` ([gym-aloha](https://github.com/huggingface/gym-aloha)), or `pusht` ([gym-pusht](https://github.com/huggingface/gym-pusht))
 Example:

 ```bash
@@ -8,7 +8,7 @@ To that end, we provide the [`Robot`](https://github.com/huggingface/lerobot/blo

 - Your own robot which exposes a communication interface (e.g. serial, CAN, TCP)
 - A way to read sensor data and send motor commands programmatically, e.g. manufacturer's SDK or API, or your own protocol implementation.
- LeRobot installed in your environment. Follow our [Installation Guide](./installation.mdx).
+- LeRobot installed in your environment. Follow our [Installation Guide](./installation).

 ## Choose your motors

@@ -65,7 +65,7 @@ class MyCoolRobotConfig(RobotConfig):
 ```
 <!-- prettier-ignore-end -->

-[Cameras tutorial](./cameras.mdx) to understand how to detect and add your camera.
+[Cameras tutorial](./cameras) to understand how to detect and add your camera.

 Next, we'll create our actual robot class which inherits from `Robot`. This abstract class defines a contract you must follow for your robot to be usable with the rest of the LeRobot tools.

@@ -335,6 +335,134 @@ For implementing teleoperation devices, we also provide a [`Teleoperator`](https

 The main differences are in the I/O functions: a teleoperator allows you to produce action via `get_action` and can receive feedback actions via `send_feedback`. Feedback could be anything controllable on the teleoperation device that could help the person controlling it understand the consequences of the actions sent. Think motion/force feedback on a leader arm, vibrations on a gamepad controller for example. To implement a teleoperator, you can follow this same tutorial and adapt it for these two methods.

+## Using Your Own `LeRobot` Devices 🔌
+
+You can easily extend `lerobot` with your own custom hardware—be it a camera, robot, or teleoperation device—by creating a separate, installable Python package. If you follow a few simple conventions, the `lerobot` command-line tools (like `lerobot-teleop` and `lerobot-record`) will **automatically discover and integrate your creations** without requiring any changes to the `lerobot` source code.
+
+This guide outlines the conventions your plugin must follow.
+
+### The 4 Core Conventions
+
+To ensure your custom device is discoverable, you must adhere to the following four rules.
+
+#### 1\. Create an Installable Package with a Specific Prefix
+
+Your project must be a standard, installable Python package. Crucially, the name of your package (as defined in `pyproject.toml` or `setup.py`) must begin with one of these prefixes:
+
+- `lerobot_robot_` for a robot.
+- `lerobot_camera_` for a camera.
+- `lerobot_teleoperator_` for a teleoperation device.
+
+This prefix system is how `lerobot` automatically finds your plugin in the Python environment.
+
+#### 2\. Follow the `SomethingConfig`/`Something` Naming Pattern
+
+Your device's implementation class must be named after its configuration class, simply by removing the `Config` suffix.
+
+- **Config Class:** `MyAwesomeTeleopConfig`
+- **Device Class:** `MyAwesomeTeleop`
+
+#### 3\. Place Your Files in a Predictable Structure
+
+The device class (`MyAwesomeTeleop`) must be located in a predictable module relative to its configuration class (`MyAwesomeTeleopConfig`). `lerobot` will automatically search in these locations:
+
+- In the **same module** as the config class.
+- In a **submodule named after the device** (e.g., `my_awesome_teleop.py`).
+
+The recommended and simplest structure is to place them in separate, clearly named files within the same directory.
+
+#### 4\. Expose Classes in `__init__.py`
+
+Your package's `__init__.py` file should import and expose both the configuration and the device classes, making them easily accessible.
+
+### Putting It All Together: A Complete Example
+
+Let's create a new teleoperator called `my_awesome_teleop`.
+
+#### Directory Structure
+
+Here is what the project folder should look like. The package name, `lerobot_teleoperator_my_awesome_teleop`, follows **Convention \#1**.
+
+```
+lerobot_teleoperator_my_awesome_teleop/
+├── pyproject.toml # (or setup.py) lists lerobot as a dependency
+└── lerobot_teleoperator_my_awesome_teleop/
+    ├── __init__.py
+    ├── config_my_awesome_teleop.py
+    └── my_awesome_teleop.py
+```
+
+#### File Contents
+
+- **`config_my_awesome_teleop.py`**: Defines the configuration class. Note the `Config` suffix (**Convention \#2**).
+
+  ```python
+  from dataclasses import dataclass
+
+  from lerobot.teleoperators.config import TeleoperatorConfig
+
+  @TeleoperatorConfig.register_subclass("my_awesome_teleop")
+  @dataclass
+  class MyAwesomeTeleopConfig(TeleoperatorConfig):
+      # Your configuration fields go here
+      port: str = "192.168.1.1"
+  ```
+
+- **`my_awesome_teleop.py`**: Implements the device. The class name `MyAwesomeTeleop` matches its config class name (**Convention \#2**). This file structure adheres to **Convention \#3**.
+
+  ```python
+  from lerobot.teleoperators.teleoperator import Teleoperator
+
+  from .config_my_awesome_teleop import MyAwesomeTeleopConfig
+
+  class MyAwesomeTeleop(Teleoperator):
+      config_class = MyAwesomeTeleopConfig
+      name = "my_awesome_teleop"
+
+      def __init__(self, config: MyAwesomeTeleopConfig):
+          super().__init__(config)
+          self.config = config
+
+      # Your device logic (e.g., connect) goes here
+  ```
+
+- **`__init__.py`**: Exposes the key classes (**Convention \#4**).
+
+  ```python
+  from .config_my_awesome_teleop import MyAwesomeTeleopConfig
+  from .my_awesome_teleop import MyAwesomeTeleop
+  ```
+
+### Installation and Usage
+
+1.  **Install your new plugin in your Python environment.** You can install your local plugin package using `pip`'s editable mode or from PyPi.
+
+    ```bash
+    # Locally
+    # Navigate to your plugin's root directory and install it
+    cd lerobot_teleoperator_my_awesome_teleop
+    pip install -e .
+
+    # From PyPi
+    pip install lerobot_teleoperator_my_awesome_teleop
+    ```
+
+2.  **Use it directly from the command line.** Now, you can use your custom device by referencing its type.
+
+    ```bash
+    lerobot-teleoperate --teleop.type=my_awesome_teleop \
+    # other arguments
+    ```
+
+And that's it\! Your custom device is now fully integrated.
+
+### Looking for an example ?
+
+Check out these two packages from the community:
+
+- https://github.com/SpesRobotics/lerobot-robot-xarm
+- https://github.com/SpesRobotics/lerobot-teleoperator-teleop
+
 ## Wrapping Up

 Once your robot class is complete, you can leverage the LeRobot ecosystem:
@@ -297,9 +297,9 @@ LeRobot provides many registered processor steps. Here are the most commonly use

 ### Next Steps

- **[Implement Your Own Processor](implement_your_own_processor.mdx)** - Create custom processor steps
- **[Debug Your Pipeline](debug_processor_pipeline.mdx)** - Troubleshoot and optimize pipelines
- **[Processors for Robots and Teleoperators](processors_robots_teleop.mdx)** - Real-world integration patterns
+- **[Implement Your Own Processor](./implement_your_own_processor)** - Create custom processor steps
+- **[Debug Your Pipeline](./debug_processor_pipeline)** - Troubleshoot and optimize pipelines
+- **[Processors for Robots and Teleoperators](./processors_robots_teleop)** - Real-world integration patterns

 ## Summary

@@ -279,3 +279,36 @@ python -m lerobot.datasets.v30.convert_dataset_v21_to_v30 --repo-id=<HF_USER/DAT
 - Aggregates parquet files: `episode-0000.parquet`, `episode-0001.parquet`, … → **`file-0000.parquet`**, …
 - Aggregates mp4 files: `episode-0000.mp4`, `episode-0001.mp4`, … → **`file-0000.mp4`**, …
 - Updates `meta/episodes/*` (chunked Parquet) with per‑episode lengths, tasks, and byte/frame offsets.
+
+## Common Issues
+
+### Always call `finalize()` before pushing
+
+When creating or recording datasets, you **must** call `dataset.finalize()` to properly close parquet writers. See the [PR #1903](https://github.com/huggingface/lerobot/pull/1903) for more details.
+
+```python
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+# Create dataset and record episodes
+dataset = LeRobotDataset.create(...)
+
+for episode in range(num_episodes):
+    # Record frames
+    for frame in episode_data:
+        dataset.add_frame(frame)
+    dataset.save_episode()
+
+# Call finalize() when done recording and before push_to_hub()
+dataset.finalize()  # Closes parquet writers, writes metadata footers
+dataset.push_to_hub()
+```
+
+**Why is this necessary?**
+
+Dataset v3.0 uses incremental parquet writing with buffered metadata for efficiency. The `finalize()` method:
+
+- Flushes any buffered episode metadata to disk
+- Closes parquet writers to write footer metadata, otherwise the parquet files will be corrupt
+- Ensures the dataset is valid for loading
+
+Without calling `finalize()`, your parquet files will be incomplete and the dataset won't load properly.
@@ -125,3 +125,42 @@ lerobot-train \
 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)
+
+## Reproducing π₀.₅ results
+
+We reproduce the results of π₀.₅ on the LIBERO benchmark using the LeRobot implementation. We take the Physical Intelligence LIBERO base model (`pi05_libero`) and finetune for an additional 6k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
+
+The finetuned model can be found here:
+
+- **π₀.₅ LIBERO**: [lerobot/pi05_libero_finetuned](https://huggingface.co/lerobot/pi05_libero_finetuned)
+
+We then evaluate the finetuned model using the LeRobot LIBERO implementation, by running the following command:
+
+```bash
+lerobot-eval \
+  --output_dir=/logs/ \
+  --env.type=libero \
+  --env.task=libero_spatial,libero_object,libero_goal,libero_10 \
+  --eval.batch_size=1 \
+  --eval.n_episodes=10 \
+  --policy.path=pi05_libero_finetuned \
+  --policy.n_action_steps=10 \
+  --output_dir=./eval_logs/ \
+  --env.max_parallel_tasks=1
+```
+
+**Note:** We set `n_action_steps=10`, similar to the original OpenPI implementation.
+
+### Results
+
+We obtain the following results on the LIBERO benchmark:
+
+| Model    | LIBERO Spatial | LIBERO Object | LIBERO Goal | LIBERO 10 | Average  |
+| -------- | -------------- | ------------- | ----------- | --------- | -------- |
+| **π₀.₅** | 97.0           | 99.0          | 98.0        | 96.0      | **97.5** |
+
+These results are consistent with the original [results](https://github.com/Physical-Intelligence/openpi/tree/main/examples/libero#results) reported by Physical Intelligence:
+
+| Model    | LIBERO Spatial | LIBERO Object | LIBERO Goal | LIBERO 10 | Average   |
+| -------- | -------------- | ------------- | ----------- | --------- | --------- |
+| **π₀.₅** | 98.8           | 98.2          | 98.0        | 92.4      | **96.85** |
@@ -0,0 +1,80 @@
+# Meta-World
+
+Meta-World is a well-designed, open-source simulation benchmark for multi-task and meta reinforcement learning in continuous-control robotic manipulation. It gives researchers a shared, realistic playground to test whether algorithms can _learn many different tasks_ and _generalize quickly to new ones_ — two central challenges for real-world robotics.
+
+- 📄 [MetaWorld paper](https://arxiv.org/pdf/1910.10897)
+- 💻 [Original MetaWorld repo](https://github.com/Farama-Foundation/Metaworld)
+
+![MetaWorld MT10 demo](https://meta-world.github.io/figures/ml45.gif)
+
+## Why Meta-World matters
+
+- **Diverse, realistic tasks.** Meta-World bundles a large suite of simulated manipulation tasks (50 in the MT50 suite) using everyday objects and a common tabletop Sawyer arm. This diversity exposes algorithms to a wide variety of dynamics, contacts and goal specifications while keeping a consistent control and observation structure.
+- **Focus on generalization and multi-task learning.** By evaluating across task distributions that share structure but differ in goals and objects, Meta-World reveals whether an agent truly learns transferable skills rather than overfitting to a narrow task.
+- **Standardized evaluation protocol.** It provides clear evaluation modes and difficulty splits, so different methods can be compared fairly across easy, medium, hard and very-hard regimes.
+- **Empirical insight.** Past evaluations on Meta-World show impressive progress on some fronts, but also highlight that current multi-task and meta-RL methods still struggle with large, diverse task sets. That gap points to important research directions.
+
+## What it enables in LeRobot
+
+In LeRobot, you can evaluate any policy or vision-language-action (VLA) model on Meta-World tasks and get a clear success-rate measure. The integration is designed to be straightforward:
+
+- We provide a LeRobot-ready dataset for Meta-World (MT50) on the HF Hub: `https://huggingface.co/datasets/lerobot/metaworld_mt50`.
+  - This dataset is formatted for the MT50 evaluation that uses all 50 tasks (the most challenging multi-task setting).
+  - MT50 gives the policy a one-hot task vector and uses fixed object/goal positions for consistency.
+
+- Task descriptions and the exact keys required for evaluation are available in the repo/dataset — use these to ensure your policy outputs the right success signals.
+
+## Quick start, train a SmolVLA policy on Meta-World
+
+Example command to train a SmolVLA policy on a subset of tasks:
+
+```bash
+lerobot-train \
+  --policy.type=smolvla \
+  --policy.repo_id=${HF_USER}/metaworld-test \
+  --policy.load_vlm_weights=true \
+  --dataset.repo_id=lerobot/metaworld_mt50 \
+  --env.type=metaworld \
+  --env.task=assembly-v3,dial-turn-v3,handle-press-side-v3 \
+  --output_dir=./outputs/ \
+  --steps=100000 \
+  --batch_size=4 \
+  --eval.batch_size=1 \
+  --eval.n_episodes=1 \
+  --eval_freq=1000
+```
+
+Notes:
+
+- `--env.task` accepts explicit task lists (comma separated) or difficulty groups (e.g., `env.task="hard"`).
+- Adjust `batch_size`, `steps`, and `eval_freq` to match your compute budget.
+- **Gymnasium Assertion Error**: if you encounter an error like
+  `AssertionError: ['human', 'rgb_array', 'depth_array']` when running MetaWorld environments, this comes from a mismatch between MetaWorld and your Gymnasium version.
+  We recommend using:
+
+```bash
+  pip install "gymnasium==1.1.0"
+```
+
+to ensure proper compatibility.
+
+## Quick start — evaluate a trained policy
+
+To evaluate a trained policy on the Meta-World medium difficulty split:
+
+```bash
+lerobot-eval \
+  --policy.path="your-policy-id" \
+  --env.type=metaworld \
+  --env.task=medium \
+  --eval.batch_size=1 \
+  --eval.n_episodes=2
+```
+
+This will run episodes and return per-task success rates using the standard Meta-World evaluation keys.
+
+## Practical tips
+
+- If you care about generalization, run on the full MT50 suite — it’s intentionally challenging and reveals strengths/weaknesses better than a few narrow tasks.
+- Use the one-hot task conditioning for multi-task training (MT10 / MT50 conventions) so policies have explicit task context.
+- Inspect the dataset task descriptions and the `info["is_success"]` keys when writing post-processing or logging so your success metrics line up with the benchmark.
@@ -79,7 +79,7 @@ After running the example:
 - Android: after starting the script, open the printed local URL on your phone, tap Start, then press and hold Move.
 - iOS: open HEBI Mobile I/O first; B1 enables motion. A3 controls the gripper.

-Additionally you can customize mapping or safety limits by editing the processor steps shown in the examples. You can also remap inputs (e.g., use a different analog input) or adapt the pipeline to other robots (e.g., LeKiwi) by modifying the input and kinematics steps. More about this in the [Processors for Robots and Teleoperators](./processors_robots_teleop.mdx) guide.
+Additionally you can customize mapping or safety limits by editing the processor steps shown in the examples. You can also remap inputs (e.g., use a different analog input) or adapt the pipeline to other robots (e.g., LeKiwi) by modifying the input and kinematics steps. More about this in the [Processors for Robots and Teleoperators](./processors_robots_teleop) guide.

 - Run this example to record a dataset, which saves absolute end effector observations and actions:

@@ -0,0 +1,79 @@
+# π₀ (Pi0)
+
+π₀ is a **Vision-Language-Action model for general robot control**, from Physical Intelligence. The LeRobot implementation is adapted from their open source [OpenPI](https://github.com/Physical-Intelligence/openpi) repository.
+
+## Model Overview
+
+π₀ represents a breakthrough in robotics as the first general-purpose robot foundation model developed by [Physical Intelligence](https://www.physicalintelligence.company/blog/pi0). Unlike traditional robot programs that are narrow specialists programmed for repetitive motions, π₀ is designed to be a generalist policy that can understand visual inputs, interpret natural language instructions, and control a variety of different robots across diverse tasks.
+
+### The Vision for Physical Intelligence
+
+As described by Physical Intelligence, while AI has achieved remarkable success in digital domains, from chess-playing to drug discovery, human intelligence still dramatically outpaces AI in the physical world. To paraphrase Moravec's paradox, winning a game of chess represents an "easy" problem for AI, but folding a shirt or cleaning up a table requires solving some of the most difficult engineering problems ever conceived. π₀ represents a first step toward developing artificial physical intelligence that enables users to simply ask robots to perform any task they want, just like they can with large language models.
+
+### Architecture and Approach
+
+π₀ combines several key innovations:
+
+- **Flow Matching**: Uses a novel method to augment pre-trained VLMs with continuous action outputs via flow matching (a variant of diffusion models)
+- **Cross-Embodiment Training**: Trained on data from 8 distinct robot platforms including UR5e, Bimanual UR5e, Franka, Bimanual Trossen, Bimanual ARX, Mobile Trossen, and Mobile Fibocom
+- **Internet-Scale Pre-training**: Inherits semantic knowledge from a pre-trained 3B parameter Vision-Language Model
+- **High-Frequency Control**: Outputs motor commands at up to 50 Hz for real-time dexterous manipulation
+
+## Installation Requirements
+
+1. Install LeRobot by following our [Installation Guide](./installation).
+2. Install Pi0 dependencies by running:
+
+   ```bash
+   pip install -e ".[pi]"
+   ```
+
+## Training Data and Capabilities
+
+π₀ is trained on the largest robot interaction dataset to date, combining three key data sources:
+
+1. **Internet-Scale Pre-training**: Vision-language data from the web for semantic understanding
+2. **Open X-Embodiment Dataset**: Open-source robot manipulation datasets
+3. **Physical Intelligence Dataset**: Large and diverse dataset of dexterous tasks across 8 distinct robots
+
+## Usage
+
+To use π₀ in LeRobot, specify the policy type as:
+
+```python
+policy.type=pi0
+```
+
+## Training
+
+For training π₀, you can use the standard LeRobot training script with the appropriate configuration:
+
+```bash
+python src/lerobot/scripts/lerobot_train.py \
+    --dataset.repo_id=your_dataset \
+    --policy.type=pi0 \
+    --output_dir=./outputs/pi0_training \
+    --job_name=pi0_training \
+    --policy.pretrained_path=lerobot/pi0_base \
+    --policy.repo_id=your_repo_id \
+    --policy.compile_model=true \
+    --policy.gradient_checkpointing=true \
+    --policy.dtype=bfloat16 \
+    --steps=3000 \
+    --policy.device=cuda \
+    --batch_size=32
+```
+
+### Key Training Parameters
+
+- **`--policy.compile_model=true`**: Enables model compilation for faster training
+- **`--policy.gradient_checkpointing=true`**: Reduces memory usage significantly during training
+- **`--policy.dtype=bfloat16`**: Use mixed precision training for efficiency
+- **`--batch_size=32`**: Batch size for training, adapt this based on your GPU memory
+- **`--policy.pretrained_path=lerobot/pi0_base`**: The base π₀ model you want to finetune, options are:
+  - [lerobot/pi0_base](https://huggingface.co/lerobot/pi0_base)
+  - [lerobot/pi0_libero](https://huggingface.co/lerobot/pi0_libero) (specifically trained on the Libero dataset)
+
+## License
+
+This model follows the **Apache 2.0 License**, consistent with the original [OpenPI repository](https://github.com/Physical-Intelligence/openpi).
@@ -0,0 +1,107 @@
+# π₀.₅ (Pi05) Policy
+
+π₀.₅ is a **Vision-Language-Action model with open-world generalization**, from Physical Intelligence. The LeRobot implementation is adapted from their open source [OpenPI](https://github.com/Physical-Intelligence/openpi) repository.
+
+## Model Overview
+
+π₀.₅ represents a significant evolution from π₀, developed by [Physical Intelligence](https://www.physicalintelligence.company/blog/pi05) to address a big challenge in robotics: **open-world generalization**. While robots can perform impressive tasks in controlled environments, π₀.₅ is designed to generalize to entirely new environments and situations that were never seen during training.
+
+### The Generalization Challenge
+
+As Physical Intelligence explains, the fundamental challenge isn't performing tasks of agility or dexterity, but generalization, the ability to correctly perform tasks in new settings with new objects. Consider a robot cleaning different homes: each home has different objects in different places. Generalization must occur at multiple levels:
+
+- **Physical Level**: Understanding how to pick up a spoon (by the handle) or plate (by the edge), even with unseen objects in cluttered environments
+- **Semantic Level**: Understanding task semantics, where to put clothes and shoes (laundry hamper, not on the bed), and what tools are appropriate for cleaning spills
+- **Environmental Level**: Adapting to "messy" real-world environments like homes, grocery stores, offices, and hospitals
+
+### Co-Training on Heterogeneous Data
+
+The breakthrough innovation in π₀.₅ is **co-training on heterogeneous data sources**. The model learns from:
+
+1. **Multimodal Web Data**: Image captioning, visual question answering, object detection
+2. **Verbal Instructions**: Humans coaching robots through complex tasks step-by-step
+3. **Subtask Commands**: High-level semantic behavior labels (e.g., "pick up the pillow" for an unmade bed)
+4. **Cross-Embodiment Robot Data**: Data from various robot platforms with different capabilities
+5. **Multi-Environment Data**: Static robots deployed across many different homes
+6. **Mobile Manipulation Data**: ~400 hours of mobile robot demonstrations
+
+This diverse training mixture creates a "curriculum" that enables generalization across physical, visual, and semantic levels simultaneously.
+
+## Installation Requirements
+
+1. Install LeRobot by following our [Installation Guide](./installation).
+2. Install Pi0.5 dependencies by running:
+
+   ```bash
+   pip install -e ".[pi]"
+   ```
+
+## Usage
+
+To use π₀.₅ in your LeRobot configuration, specify the policy type as:
+
+```python
+policy.type=pi05
+```
+
+## Training
+
+### Training Command Example
+
+Here's a complete training command for finetuning the base π₀.₅ model on your own dataset:
+
+```bash
+python src/lerobot/scripts/lerobot_train.py\
+    --dataset.repo_id=your_dataset \
+    --policy.type=pi05 \
+    --output_dir=./outputs/pi05_training \
+    --job_name=pi05_training \
+    --policy.repo_id=your_repo_id \
+    --policy.pretrained_path=lerobot/pi05_base \
+    --policy.compile_model=true \
+    --policy.gradient_checkpointing=true \
+    --wandb.enable=true \
+    --policy.dtype=bfloat16 \
+    --steps=3000 \
+    --policy.device=cuda \
+    --batch_size=32
+```
+
+### Key Training Parameters
+
+- **`--policy.compile_model=true`**: Enables model compilation for faster training
+- **`--policy.gradient_checkpointing=true`**: Reduces memory usage significantly during training
+- **`--policy.dtype=bfloat16`**: Use mixed precision training for efficiency
+- **`--batch_size=32`**: Batch size for training, adapt this based on your GPU memory
+- **`--policy.pretrained_path=lerobot/pi05_base`**: The base π₀.₅ model you want to finetune, options are:
+  - [lerobot/pi05_base](https://huggingface.co/lerobot/pi05_base)
+  - [lerobot/pi05_libero](https://huggingface.co/lerobot/pi05_libero) (specifically trained on the Libero dataset)
+
+If your dataset is not converted with `quantiles`, you can convert it with the following command:
+
+```bash
+python src/lerobot/datasets/v30/augment_dataset_quantile_stats.py \
+    --repo-id=your_dataset \
+```
+
+Or train pi05 with this normalization mapping: `--policy.normalization_mapping='{"ACTION": "MEAN_STD", "STATE": "MEAN_STD", "VISUAL": "IDENTITY"}'`
+
+## Performance Results
+
+### Libero Benchmark Results
+
+π₀.₅ has demonstrated strong performance on the Libero benchmark suite. To compare and test its LeRobot implementation, we finetuned the libero base model for an additional 6k steps on the Libero dataset and compared the results to the OpenPI reference results.
+
+| Benchmark          | LeRobot Implementation | OpenPI Reference |
+| ------------------ | ---------------------- | ---------------- |
+| **Libero Spatial** | 97.0%                  | 98.8%            |
+| **Libero Object**  | 99.0%                  | 98.2%            |
+| **Libero Goal**    | 98.0%                  | 98.0%            |
+| **Libero 10**      | 96.0%                  | 92.4%            |
+| **Average**        | 97.5%                  | 96.85%           |
+
+These results demonstrate π₀.₅'s strong generalization capabilities across diverse robotic manipulation tasks. To reproduce these results, you can follow the instructions in the [Libero](https://huggingface.co/docs/lerobot/libero) section.
+
+## License
+
+This model follows the **Apache 2.0 License**, consistent with the original [OpenPI repository](https://github.com/Physical-Intelligence/openpi).
@@ -1,4 +1,4 @@
-# Finetune SmolVLA
+# SmolVLA

 SmolVLA is Hugging Face’s lightweight foundation model for robotics. Designed for easy fine-tuning on LeRobot datasets, it helps accelerate your development!

@@ -0,0 +1,102 @@
+# Using Dataset Tools
+
+This guide covers the dataset tools utilities available in LeRobot for modifying and editing existing datasets.
+
+## Overview
+
+LeRobot provides several utilities for manipulating datasets:
+
+1. **Delete Episodes** - Remove specific episodes from a dataset
+2. **Split Dataset** - Divide a dataset into multiple smaller datasets
+3. **Merge Datasets** - Combine multiple datasets into one. The datasets must have identical features, and episodes are concatenated in the order specified in `repo_ids`
+4. **Add Features** - Add new features to a dataset
+5. **Remove Features** - Remove features from a dataset
+
+The core implementation is in `lerobot.datasets.dataset_tools`.
+An example script detailing how to use the tools API is available in `examples/dataset/use_dataset_tools.py`.
+
+## Command-Line Tool: lerobot-edit-dataset
+
+`lerobot-edit-dataset` is a command-line script for editing datasets. It can be used to delete episodes, split datasets, merge datasets, add features, and remove features.
+
+Run `lerobot-edit-dataset --help` for more information on the configuration of each operation.
+
+### Usage Examples
+
+#### Delete Episodes
+
+Remove specific episodes from a dataset. This is useful for filtering out undesired data.
+
+```bash
+# Delete episodes 0, 2, and 5 (modifies original dataset)
+lerobot-edit-dataset \
+    --repo_id lerobot/pusht \
+    --operation.type delete_episodes \
+    --operation.episode_indices "[0, 2, 5]"
+
+# Delete episodes and save to a new dataset (preserves original dataset)
+lerobot-edit-dataset \
+    --repo_id lerobot/pusht \
+    --new_repo_id lerobot/pusht_after_deletion \
+    --operation.type delete_episodes \
+    --operation.episode_indices "[0, 2, 5]"
+```
+
+#### Split Dataset
+
+Divide a dataset into multiple subsets.
+
+```bash
+# Split by fractions (e.g. 80% train, 20% test, 20% val)
+lerobot-edit-dataset \
+    --repo_id lerobot/pusht \
+    --operation.type split \
+    --operation.splits '{"train": 0.8, "test": 0.2, "val": 0.2}'
+
+# Split by specific episode indices
+lerobot-edit-dataset \
+    --repo_id lerobot/pusht \
+    --operation.type split \
+    --operation.splits '{"task1": [0, 1, 2, 3], "task2": [4, 5]}'
+```
+
+There are no constraints on the split names, they can be determined by the user. Resulting datasets are saved under the repo id with the split name appended, e.g. `lerobot/pusht_train`, `lerobot/pusht_task1`, `lerobot/pusht_task2`.
+
+#### Merge Datasets
+
+Combine multiple datasets into a single dataset.
+
+```bash
+# Merge train and validation splits back into one dataset
+lerobot-edit-dataset \
+    --repo_id lerobot/pusht_merged \
+    --operation.type merge \
+    --operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']"
+```
+
+#### Remove Features
+
+Remove features from a dataset.
+
+```bash
+# Remove a camera feature
+lerobot-edit-dataset \
+    --repo_id lerobot/pusht \
+    --operation.type remove_feature \
+    --operation.feature_names "['observation.images.top']"
+```
+
+### Push to Hub
+
+Add the `--push_to_hub` flag to any command to automatically upload the resulting dataset to the Hugging Face Hub:
+
+```bash
+lerobot-edit-dataset \
+    --repo_id lerobot/pusht \
+    --new_repo_id lerobot/pusht_after_deletion \
+    --operation.type delete_episodes \
+    --operation.episode_indices "[0, 2, 5]" \
+    --push_to_hub
+```
+
+There is also a tool for adding features to a dataset that is not yet covered in `lerobot-edit-dataset`.
@@ -0,0 +1,124 @@
+#!/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.
+
+"""
+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/dataset/use_dataset_tools.py
+"""
+
+import numpy as np
+
+from lerobot.datasets.dataset_tools import (
+    add_features,
+    delete_episodes,
+    merge_datasets,
+    modify_features,
+    remove_feature,
+    split_dataset,
+)
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+
+def main():
+    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())}")
+
+    print("\n1. Deleting episodes 0 and 2...")
+    filtered_dataset = delete_episodes(dataset, episode_indices=[0, 2], repo_id="lerobot/pusht_filtered")
+    print(f"Filtered dataset: {filtered_dataset.meta.total_episodes} episodes")
+
+    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")
+
+    print("\n3. Adding features...")
+
+    reward_values = np.random.randn(dataset.meta.total_frames).astype(np.float32)
+
+    def compute_success(row_dict, episode_index, frame_index):
+        episode_length = 10
+        return float(frame_index >= episode_length - 10)
+
+    dataset_with_features = add_features(
+        dataset,
+        features={
+            "reward": (
+                reward_values,
+                {"dtype": "float32", "shape": (1,), "names": None},
+            ),
+            "success": (
+                compute_success,
+                {"dtype": "float32", "shape": (1,), "names": None},
+            ),
+        },
+        repo_id="lerobot/pusht_with_features",
+    )
+
+    print(f"New features: {list(dataset_with_features.meta.features.keys())}")
+
+    print("\n4. Removing the success feature...")
+    dataset_cleaned = remove_feature(
+        dataset_with_features, feature_names="success", repo_id="lerobot/pusht_cleaned"
+    )
+    print(f"Features after removal: {list(dataset_cleaned.meta.features.keys())}")
+
+    print("\n5. Using modify_features to add and remove features simultaneously...")
+    dataset_modified = modify_features(
+        dataset_with_features,
+        add_features={
+            "discount": (
+                np.ones(dataset.meta.total_frames, dtype=np.float32) * 0.99,
+                {"dtype": "float32", "shape": (1,), "names": None},
+            ),
+        },
+        remove_features="reward",
+        repo_id="lerobot/pusht_modified",
+    )
+    print(f"Modified features: {list(dataset_modified.meta.features.keys())}")
+
+    print("\n6. Merging train and val splits back together...")
+    merged = merge_datasets([splits["train"], splits["val"]], output_repo_id="lerobot/pusht_merged")
+    print(f"Merged dataset: {merged.meta.total_episodes} episodes")
+
+    print("\n7. Complex workflow example...")
+
+    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()
@@ -133,4 +133,6 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
 log_say("Stop recording")
 robot.disconnect()
 listener.stop()
+
+dataset.finalize()
 dataset.push_to_hub()
@@ -130,4 +130,6 @@ robot.disconnect()
 leader_arm.disconnect()
 keyboard.disconnect()
 listener.stop()
+
+dataset.finalize()
 dataset.push_to_hub()
@@ -194,4 +194,6 @@ for episode_idx in range(NUM_EPISODES):
 log_say("Stop recording")
 robot.disconnect()
 listener.stop()
+
+dataset.finalize()
 dataset.push_to_hub()
@@ -200,4 +200,6 @@ log_say("Stop recording")
 robot.disconnect()
 phone.disconnect()
 listener.stop()
+
+dataset.finalize()
 dataset.push_to_hub()
@@ -362,6 +362,8 @@ def port_droid(
        lerobot_dataset.save_episode()
        logging.info("Save_episode")

+    lerobot_dataset.finalize()
+
    if push_to_hub:
        lerobot_dataset.push_to_hub(
            # Add openx tag, since it belongs to the openx collection of datasets
@@ -195,4 +195,6 @@ for episode_idx in range(NUM_EPISODES):
 log_say("Stop recording")
 robot.disconnect()
 listener.stop()
+
+dataset.finalize()
 dataset.push_to_hub()
@@ -199,4 +199,6 @@ log_say("Stop recording")
 leader.disconnect()
 follower.disconnect()
 listener.stop()
+
+dataset.finalize()
 dataset.push_to_hub()
@@ -59,27 +59,28 @@ keywords = ["lerobot", "huggingface", "robotics",  "machine learning", "artifici
 dependencies = [

    # Hugging Face dependencies
-    "datasets>=4.0.0",
-    "diffusers>=0.27.2",
-    "huggingface-hub[hf-transfer,cli]>=0.34.2",
+    "datasets>=4.0.0,<4.2.0",
+    "diffusers>=0.27.2,<0.36.0",
+    "huggingface-hub[hf-transfer,cli]>=0.34.2,<0.36.0",

    # Core dependencies
-    "cmake>=3.29.0.1",
-    "einops>=0.8.0",
-    "opencv-python-headless>=4.9.0",
-    "av>=14.2.0",
-    "jsonlines>=4.0.0",
-    "packaging>=24.2",
-    "pynput>=1.7.7",
-    "pyserial>=3.5",
-    "wandb>=0.20.0",
+    "setuptools>=71.0.0,<81.0.0",
+    "cmake>=3.29.0.1,<4.2.0",
+    "einops>=0.8.0,<0.9.0",
+    "opencv-python-headless>=4.9.0,<4.13.0",
+    "av>=15.0.0,<16.0.0",
+    "jsonlines>=4.0.0,<5.0.0",
+    "packaging>=24.2,<26.0",
+    "pynput>=1.7.7,<1.9.0",
+    "pyserial>=3.5,<4.0",
+    "wandb>=0.20.0,<0.23.0",

    "torch>=2.2.1,<2.8.0", # TODO: Bumb dependency
    "torchcodec>=0.2.1,<0.6.0; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')", # TODO: Bumb dependency
    "torchvision>=0.21.0,<0.23.0", # TODO: Bumb dependency

    "draccus==0.10.0", # TODO: Remove ==
-    "gymnasium>=0.29.1,<1.0.0", # TODO: Bumb dependency
+    "gymnasium>=1.0.0",
    "rerun-sdk>=0.21.0,<0.23.0", # TODO: Bumb dependency

    # Support dependencies
@@ -92,26 +93,26 @@ dependencies = [
 [project.optional-dependencies]

 # Common
-pygame-dep = ["pygame>=2.5.1"]
-placo-dep = ["placo>=0.9.6"]
-transformers-dep = ["transformers>=4.52.0"]
+pygame-dep = ["pygame>=2.5.1,<2.7.0"]
+placo-dep = ["placo>=0.9.6,<0.10.0"]
+transformers-dep = ["transformers>=4.53.0,<5.0.0"]
 grpcio-dep = ["grpcio==1.73.1", "protobuf==6.31.0"]

 # Motors
-feetech = ["feetech-servo-sdk>=1.0.0"]
-dynamixel = ["dynamixel-sdk>=3.7.31"]
+feetech = ["feetech-servo-sdk>=1.0.0,<2.0.0"]
+dynamixel = ["dynamixel-sdk>=3.7.31,<3.9.0"]

 # Robots
-gamepad = ["lerobot[pygame-dep]", "hidapi>=0.14.0"]
+gamepad = ["lerobot[pygame-dep]", "hidapi>=0.14.0,<0.15.0"]
 hopejr = ["lerobot[feetech]", "lerobot[pygame-dep]"]
-lekiwi = ["lerobot[feetech]", "pyzmq>=26.2.1"]
-reachy2 = ["reachy2_sdk>=1.0.14"]
+lekiwi = ["lerobot[feetech]", "pyzmq>=26.2.1,<28.0.0"]
+reachy2 = ["reachy2_sdk>=1.0.14,<1.1.0"]
 kinematics = ["lerobot[placo-dep]"]
 intelrealsense = [
-    "pyrealsense2>=2.55.1.6486 ; sys_platform != 'darwin'",
-    "pyrealsense2-macosx>=2.54 ; sys_platform == 'darwin'",
+    "pyrealsense2>=2.55.1.6486,<2.57.0 ; sys_platform != 'darwin'",
+    "pyrealsense2-macosx>=2.54,<2.55.0 ; sys_platform == 'darwin'",
 ]
-phone = ["hebi-py>=2.8.0", "teleop>=0.1.0"]
+phone = ["hebi-py>=2.8.0,<2.12.0", "teleop>=0.1.0,<0.2.0"]
 # stretch = [
 #     "hello-robot-stretch-body>=0.7.27 ; sys_platform == 'linux'",
 #     "pyrender @ git+https://github.com/mmatl/pyrender.git ; sys_platform == 'linux'",
@@ -119,24 +120,23 @@ phone = ["hebi-py>=2.8.0", "teleop>=0.1.0"]
 # ] # TODO: Currently not supported

 # Policies
-pi0 = ["lerobot[transformers-dep]"]
-smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14", "accelerate>=1.7.0", "safetensors>=0.4.3"]
-hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.11", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
+pi = ["transformers @ git+https://github.com/huggingface/transformers.git@fix/lerobot_openpi"]
+smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14,<0.6.0", "accelerate>=1.7.0,<2.0.0", "safetensors>=0.4.3,<1.0.0"]
+hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.11,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]

 # Features
-async = ["lerobot[grpcio-dep]", "matplotlib>=3.10.3"]
+async = ["lerobot[grpcio-dep]", "matplotlib>=3.10.3,<4.0.0"]

 # Development
-dev = ["pre-commit>=3.7.0", "debugpy>=1.8.1", "lerobot[grpcio-dep]", "grpcio-tools==1.73.1"]
-test = ["pytest>=8.1.0", "pytest-timeout>=2.4.0", "pytest-cov>=5.0.0", "mock-serial>=0.0.1 ; sys_platform != 'win32'"]
-video_benchmark = ["scikit-image>=0.23.2", "pandas>=2.2.2"]
+dev = ["pre-commit>=3.7.0,<5.0.0", "debugpy>=1.8.1,<1.9.0", "lerobot[grpcio-dep]", "grpcio-tools==1.73.1"]
+test = ["pytest>=8.1.0,<9.0.0", "pytest-timeout>=2.4.0,<3.0.0", "pytest-cov>=5.0.0,<8.0.0", "mock-serial>=0.0.1,<0.1.0 ; sys_platform != 'win32'"]
+video_benchmark = ["scikit-image>=0.23.2,<0.26.0", "pandas>=2.2.2,<2.4.0"]

 # Simulation
-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"]
+aloha = ["gym-aloha>=0.1.2,<0.2.0"]
+pusht = ["gym-pusht>=0.1.5,<0.2.0", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead
 libero = ["lerobot[transformers-dep]", "libero @ git+https://github.com/huggingface/lerobot-libero.git@main#egg=libero"]
-
+metaworld = ["metaworld>=3.0.0"]

 # All
 all = [
@@ -147,7 +147,7 @@ all = [
    "lerobot[reachy2]",
    "lerobot[kinematics]",
    "lerobot[intelrealsense]",
-    "lerobot[pi0]",
+    "lerobot[pi]",
    "lerobot[smolvla]",
    "lerobot[hilserl]",
    "lerobot[async]",
@@ -156,9 +156,9 @@ all = [
    "lerobot[video_benchmark]",
    "lerobot[aloha]",
    "lerobot[pusht]",
-    "lerobot[xarm]",
    "lerobot[phone]",
    "lerobot[libero]",
+    "lerobot[metaworld]",
 ]

 [project.scripts]
@@ -175,6 +175,7 @@ lerobot-dataset-viz="lerobot.scripts.lerobot_dataset_viz:main"
 lerobot-info="lerobot.scripts.lerobot_info:main"
 lerobot-find-joint-limits="lerobot.scripts.lerobot_find_joint_limits:main"
 lerobot-imgtransform-viz="lerobot.scripts.lerobot_imgtransform_viz:main"
+lerobot-edit-dataset="lerobot.scripts.lerobot_edit_dataset:main"

 # ---------------- Tool Configurations ----------------
 [tool.setuptools.packages.find]
@@ -270,80 +271,84 @@ default.extend-ignore-identifiers-re = [
 # TODO: Enable mypy gradually module by module across multiple PRs
 # Uncomment [tool.mypy] first, then uncomment individual module overrides as they get proper type annotations

-# [tool.mypy]
-# python_version = "3.10"
+[tool.mypy]
+python_version = "3.10"
+ignore_missing_imports = true
+follow_imports = "skip"
 # warn_return_any = true
 # warn_unused_configs = true
-# ignore_missing_imports = false
 # strict = true
 # disallow_untyped_defs = true
 # disallow_incomplete_defs = true
 # check_untyped_defs = true

+[[tool.mypy.overrides]]
+module = "lerobot.*"
+ignore_errors = true
+
+[[tool.mypy.overrides]]
+module = "lerobot.envs.*"
+# Enable type checking only for the envs module
+ignore_errors = false
+
+
 # [[tool.mypy.overrides]]
 # module = "lerobot.utils.*"
-# # include = "src/lerobot/utils/**/*.py"
+# ignore_errors = false

 # [[tool.mypy.overrides]]
 # module = "lerobot.configs.*"
-# # include = "src/lerobot/configs/**/*.py"
+# ignore_errors = false

-# # Data processing modules
-# [[tool.mypy.overrides]]
-# module = "lerobot.processor.*"
-# # include = "src/lerobot/processor/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.datasets.*"
-# # include = "src/lerobot/datasets/**/*.py"
-
-# # Core machine learning modules
 # [[tool.mypy.overrides]]
 # module = "lerobot.optim.*"
-# # include = "src/lerobot/optim/**/*.py"
+# ignore_errors = false

 # [[tool.mypy.overrides]]
 # module = "lerobot.model.*"
-# # include = "src/lerobot/model/**/*.py"
+# ignore_errors = false
+
+# [[tool.mypy.overrides]]
+# module = "lerobot.processor.*"
+# ignore_errors = false
+
+# [[tool.mypy.overrides]]
+# module = "lerobot.datasets.*"
+# ignore_errors = false

-# # Hardware interfaces
 # [[tool.mypy.overrides]]
 # module = "lerobot.cameras.*"
-# # include = "src/lerobot/cameras/**/*.py"
+# ignore_errors = false

 # [[tool.mypy.overrides]]
 # module = "lerobot.motors.*"
-# # include = "src/lerobot/motors/**/*.py"
+# ignore_errors = false

 # [[tool.mypy.overrides]]
 # module = "lerobot.robots.*"
-# # include = "src/lerobot/robots/**/*.py"
+# ignore_errors = false

 # [[tool.mypy.overrides]]
 # module = "lerobot.teleoperators.*"
-# # include = "src/lerobot/teleoperators/**/*.py"
+# ignore_errors = false

-# # Complex modules (enable these last)
 # [[tool.mypy.overrides]]
 # module = "lerobot.policies.*"
-# # include = "src/lerobot/policies/**/*.py"
+# ignore_errors = false

 # [[tool.mypy.overrides]]
 # module = "lerobot.rl.*"
-# # include = "src/lerobot/rl/**/*.py"
+# ignore_errors = false

-# [[tool.mypy.overrides]]
-# module = "lerobot.envs.*"
-# # include = "src/lerobot/envs/**/*.py"

 # [[tool.mypy.overrides]]
 # module = "lerobot.async_inference.*"
-# # include = "src/lerobot/async_inference/**/*.py"
+# ignore_errors = false

 # [[tool.mypy.overrides]]
 # module = "lerobot.transport.*"
-# # include = "src/lerobot/transport/**/*.py"
+# ignore_errors = false

 # [[tool.mypy.overrides]]
 # module = "lerobot.scripts.*"
-# # include = "src/lerobot/scripts/**/*.py"
+# ignore_errors = false
@@ -57,7 +57,6 @@ available_tasks_per_env = {
        "AlohaTransferCube-v0",
    ],
    "pusht": ["PushT-v0"],
-    "xarm": ["XarmLift-v0"],
 }
 available_envs = list(available_tasks_per_env.keys())

@@ -75,16 +74,6 @@ available_datasets_per_env = {
    # TODO(alexander-soare): Add "lerobot/pusht_keypoints". Right now we can't because this is too tightly
    # coupled with tests.
    "pusht": ["lerobot/pusht", "lerobot/pusht_image"],
-    "xarm": [
-        "lerobot/xarm_lift_medium",
-        "lerobot/xarm_lift_medium_replay",
-        "lerobot/xarm_push_medium",
-        "lerobot/xarm_push_medium_replay",
-        "lerobot/xarm_lift_medium_image",
-        "lerobot/xarm_lift_medium_replay_image",
-        "lerobot/xarm_push_medium_image",
-        "lerobot/xarm_push_medium_replay_image",
-    ],
 }

 available_real_world_datasets = [
@@ -195,7 +184,6 @@ available_motors = [
 available_policies_per_env = {
    "aloha": ["act"],
    "pusht": ["diffusion", "vqbet"],
-    "xarm": ["tdmpc"],
    "koch_real": ["act_koch_real"],
    "aloha_real": ["act_aloha_real"],
 }
@@ -142,11 +142,6 @@ class RobotClientConfig:
        default=False, metadata={"help": "Visualize the action queue size"}
    )

-    # Verification configuration
-    verify_robot_cameras: bool = field(
-        default=True, metadata={"help": "Verify that the robot cameras match the policy cameras"}
-    )
-
    @property
    def environment_dt(self) -> float:
        """Environment time step, in seconds"""
@@ -23,7 +23,7 @@ DEFAULT_INFERENCE_LATENCY = 1 / DEFAULT_FPS
 DEFAULT_OBS_QUEUE_TIMEOUT = 2

 # All action chunking policies
-SUPPORTED_POLICIES = ["act", "smolvla", "diffusion", "pi0", "tdmpc", "vqbet"]
+SUPPORTED_POLICIES = ["act", "smolvla", "diffusion", "tdmpc", "vqbet", "pi0", "pi05"]

 # TODO: Add all other robots
-SUPPORTED_ROBOTS = ["so100_follower", "so101_follower"]
+SUPPORTED_ROBOTS = ["so100_follower", "so101_follower", "bi_so100_follower"]
@@ -25,7 +25,14 @@ from lerobot.configs.types import PolicyFeature
 from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features

 # NOTE: Configs need to be loaded for the client to be able to instantiate the policy config
-from lerobot.policies import ACTConfig, DiffusionConfig, PI0Config, SmolVLAConfig, VQBeTConfig  # noqa: F401
+from lerobot.policies import (  # noqa: F401
+    ACTConfig,
+    DiffusionConfig,
+    PI0Config,
+    PI05Config,
+    SmolVLAConfig,
+    VQBeTConfig,
+)
 from lerobot.robots.robot import Robot
 from lerobot.utils.constants import OBS_IMAGES, OBS_STATE, OBS_STR
 from lerobot.utils.utils import init_logging
@@ -55,15 +62,6 @@ def visualize_action_queue_size(action_queue_size: list[int]) -> None:
    plt.show()


-def validate_robot_cameras_for_policy(
-    lerobot_observation_features: dict[str, dict], policy_image_features: dict[str, PolicyFeature]
-) -> None:
-    image_keys = list(filter(is_image_key, lerobot_observation_features))
-    assert set(image_keys) == set(policy_image_features.keys()), (
-        f"Policy image features must match robot cameras! Received {list(policy_image_features.keys())} != {image_keys}"
-    )
-
-
 def map_robot_keys_to_lerobot_features(robot: Robot) -> dict[str, dict]:
    return hw_to_dataset_features(robot.observation_features, OBS_STR, use_video=False)

@@ -85,11 +83,11 @@ def resize_robot_observation_image(image: torch.tensor, resize_dims: tuple[int,
    return resized.squeeze(0)


+# TODO(Steven): Consider implementing a pipeline step for this
 def raw_observation_to_observation(
    raw_observation: RawObservation,
    lerobot_features: dict[str, dict],
    policy_image_features: dict[str, PolicyFeature],
-    device: str,
 ) -> Observation:
    observation = {}

@@ -98,9 +96,7 @@ def raw_observation_to_observation(
        if isinstance(v, torch.Tensor):  # VLAs present natural-language instructions in observations
            if "image" in k:
                # Policy expects images in shape (B, C, H, W)
-                observation[k] = prepare_image(v).unsqueeze(0).to(device)
-            else:
-                observation[k] = v.to(device)
+                observation[k] = prepare_image(v).unsqueeze(0)
        else:
            observation[k] = v

@@ -15,7 +15,7 @@
 """
 Example:
 ```shell
-python src/lerobot/async_inference/policy_server.py \
+python -m lerobot.async_inference.policy_server \
     --host=127.0.0.1 \
     --port=8080 \
     --fps=30 \
@@ -32,12 +32,17 @@ from concurrent import futures
 from dataclasses import asdict
 from pprint import pformat
 from queue import Empty, Queue
+from typing import Any

 import draccus
 import grpc
 import torch

-from lerobot.policies.factory import get_policy_class
+from lerobot.policies.factory import get_policy_class, make_pre_post_processors
+from lerobot.processor import (
+    PolicyAction,
+    PolicyProcessorPipeline,
+)
 from lerobot.transport import (
    services_pb2,  # type: ignore
    services_pb2_grpc,  # type: ignore
@@ -82,6 +87,8 @@ class PolicyServer(services_pb2_grpc.AsyncInferenceServicer):
        self.lerobot_features = None
        self.actions_per_chunk = None
        self.policy = None
+        self.preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]] | None = None
+        self.postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction] | None = None

    @property
    def running(self):
@@ -146,6 +153,16 @@ class PolicyServer(services_pb2_grpc.AsyncInferenceServicer):
        start = time.perf_counter()
        self.policy = policy_class.from_pretrained(policy_specs.pretrained_name_or_path)
        self.policy.to(self.device)
+
+        # Load preprocessor and postprocessor, overriding device to match requested device
+        device_override = {"device": self.device}
+        self.preprocessor, self.postprocessor = make_pre_post_processors(
+            self.policy.config,
+            pretrained_path=policy_specs.pretrained_name_or_path,
+            preprocessor_overrides={"device_processor": device_override},
+            postprocessor_overrides={"device_processor": device_override},
+        )
+
        end = time.perf_counter()

        self.logger.info(f"Time taken to put policy on {self.device}: {end - start:.4f} seconds")
@@ -173,7 +190,7 @@ class PolicyServer(services_pb2_grpc.AsyncInferenceServicer):
        # Calculate FPS metrics
        fps_metrics = self.fps_tracker.calculate_fps_metrics(obs_timestamp)

-        self.logger.info(
+        self.logger.debug(
            f"Received observation #{obs_timestep} | "
            f"Avg FPS: {fps_metrics['avg_fps']:.2f} | "  # fps at which observations are received from client
            f"Target: {fps_metrics['target_fps']:.2f} | "
@@ -189,7 +206,7 @@ class PolicyServer(services_pb2_grpc.AsyncInferenceServicer):
        if not self._enqueue_observation(
            timed_observation  # wrapping a RawObservation
        ):
-            self.logger.info(f"Observation #{obs_timestep} has been filtered out")
+            self.logger.debug(f"Observation #{obs_timestep} has been filtered out")

        return services_pb2.Empty()

@@ -301,23 +318,6 @@ class PolicyServer(services_pb2_grpc.AsyncInferenceServicer):
            for i, action in enumerate(action_chunk)
        ]

-    def _prepare_observation(self, observation_t: TimedObservation) -> Observation:
-        """
-        Prepare observation, ready for policy inference.
-        E.g.: To keep observation sampling rate high (and network packet tiny) we send int8 [0,255] images from the
-        client and then convert them to float32 [0,1] images here, before running inference.
-        """
-        # RawObservation from robot.get_observation() - wrong keys, wrong dtype, wrong image shape
-        observation: Observation = raw_observation_to_observation(
-            observation_t.get_observation(),
-            self.lerobot_features,
-            self.policy_image_features,
-            self.device,
-        )
-        # processed Observation - right keys, right dtype, right image shape
-
-        return observation
-
    def _get_action_chunk(self, observation: dict[str, torch.Tensor]) -> torch.Tensor:
        """Get an action chunk from the policy. The chunk contains only"""
        chunk = self.policy.predict_action_chunk(observation)
@@ -327,44 +327,76 @@ class PolicyServer(services_pb2_grpc.AsyncInferenceServicer):
        return chunk[:, : self.actions_per_chunk, :]

    def _predict_action_chunk(self, observation_t: TimedObservation) -> list[TimedAction]:
-        """Predict an action chunk based on an observation"""
-        inference_starts = time.perf_counter()
+        """Predict an action chunk based on an observation.

+        Pipeline:
+        1. Convert raw observation to LeRobot format
+        2. Apply preprocessor (tokenization, normalization, batching, device placement)
+        3. Run policy inference to get action chunk
+        4. Apply postprocessor (unnormalization, device movement)
+        5. Convert to TimedAction list
+        """
        """1. Prepare observation"""
-        start_time = time.perf_counter()
-        observation = self._prepare_observation(observation_t)
-        preprocessing_time = time.perf_counter() - start_time
+        start_prepare = time.perf_counter()
+        observation: Observation = raw_observation_to_observation(
+            observation_t.get_observation(),
+            self.lerobot_features,
+            self.policy_image_features,
+        )
+        prepare_time = time.perf_counter() - start_prepare

+        """2. Apply preprocessor"""
+        start_preprocess = time.perf_counter()
+        observation = self.preprocessor(observation)
        self.last_processed_obs: TimedObservation = observation_t
+        preprocessing_time = time.perf_counter() - start_preprocess

-        """2. Get action chunk"""
-        start_time = time.perf_counter()
+        """3. Get action chunk"""
+        start_inference = time.perf_counter()
        action_tensor = self._get_action_chunk(observation)
-        inference_time = time.perf_counter() - start_time
+        inference_time = time.perf_counter() - start_inference
+        self.logger.info(
+            f"Preprocessing and inference took {inference_time:.4f}s, action shape: {action_tensor.shape}"
+        )

-        """3. Post-inference processing"""
-        start_time = time.perf_counter()
-        # Move to CPU before serializing
-        action_tensor = action_tensor.cpu().squeeze(0)
+        """4. Apply postprocessor"""
+        # Apply postprocessor (handles unnormalization and device movement)
+        # Postprocessor expects (B, action_dim) per action, but we have (B, chunk_size, action_dim)
+        # So we process each action in the chunk individually
+        start_postprocess = time.perf_counter()
+        _, chunk_size, _ = action_tensor.shape

+        # Process each action in the chunk
+        processed_actions = []
+        for i in range(chunk_size):
+            # Extract action at timestep i: (B, action_dim)
+            single_action = action_tensor[:, i, :]
+            processed_action = self.postprocessor(single_action)
+            processed_actions.append(processed_action)
+
+        # Stack back to (B, chunk_size, action_dim), then remove batch dim
+        action_tensor = torch.stack(processed_actions, dim=1).squeeze(0)
+        self.logger.debug(f"Postprocessed action shape: {action_tensor.shape}")
+
+        """5. Convert to TimedAction list"""
        action_chunk = self._time_action_chunk(
            observation_t.get_timestamp(), list(action_tensor), observation_t.get_timestep()
        )
-        postprocessing_time = time.perf_counter() - start_time
-        inference_stops = time.perf_counter()
+        postprocess_stops = time.perf_counter()
+        postprocessing_time = postprocess_stops - start_postprocess

        self.logger.info(
-            f"Observation {observation_t.get_timestep()} |"
-            f"Inference time: {1000 * (inference_stops - inference_starts):.2f}ms"
+            f"Observation {observation_t.get_timestep()} | "
+            f"Total time: {1000 * (postprocess_stops - start_prepare):.2f}ms"
        )

-        # full-process latency breakdown for debugging purposes
        self.logger.debug(
            f"Observation {observation_t.get_timestep()} | "
-            f"Preprocessing time: {1000 * (preprocessing_time - inference_starts):.2f}ms | "
-            f"Inference time: {1000 * (inference_time - preprocessing_time):.2f}ms | "
-            f"Postprocessing time: {1000 * (postprocessing_time - inference_time):.2f}ms | "
-            f"Total time: {1000 * (postprocessing_time - inference_starts):.2f}ms"
+            f"Prepare time: {1000 * prepare_time:.2f}ms | "
+            f"Preprocessing time: {1000 * preprocessing_time:.2f}ms | "
+            f"Inference time: {1000 * inference_time:.2f}ms | "
+            f"Postprocessing time: {1000 * postprocessing_time:.2f}ms | "
+            f"Total time: {1000 * (postprocess_stops - start_prepare):.2f}ms"
        )

        return action_chunk
@@ -48,10 +48,10 @@ import torch

 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
-from lerobot.configs.policies import PreTrainedConfig
 from lerobot.robots import (  # noqa: F401
    Robot,
    RobotConfig,
+    bi_so100_follower,
    koch_follower,
    make_robot_from_config,
    so100_follower,
@@ -75,7 +75,6 @@ from .helpers import (
    TimedObservation,
    get_logger,
    map_robot_keys_to_lerobot_features,
-    validate_robot_cameras_for_policy,
    visualize_action_queue_size,
 )

@@ -97,14 +96,6 @@ class RobotClient:

        lerobot_features = map_robot_keys_to_lerobot_features(self.robot)

-        if config.verify_robot_cameras:
-            # Load policy config for validation
-            policy_config = PreTrainedConfig.from_pretrained(config.pretrained_name_or_path)
-            policy_image_features = policy_config.image_features
-
-            # The cameras specified for inference must match the one supported by the policy chosen
-            validate_robot_cameras_for_policy(lerobot_features, policy_image_features)
-
        # Use environment variable if server_address is not provided in config
        self.server_address = config.server_address

@@ -214,7 +205,7 @@ class RobotClient:
            )
            _ = self.stub.SendObservations(observation_iterator)
            obs_timestep = obs.get_timestep()
-            self.logger.info(f"Sent observation #{obs_timestep} | ")
+            self.logger.debug(f"Sent observation #{obs_timestep} | ")

            return True

@@ -467,7 +458,7 @@ class RobotClient:
            if self._ready_to_send_observation():
                _captured_observation = self.control_loop_observation(task, verbose)

-            self.logger.info(f"Control loop (ms): {(time.perf_counter() - control_loop_start) * 1000:.2f}")
+            self.logger.debug(f"Control loop (ms): {(time.perf_counter() - control_loop_start) * 1000:.2f}")
            # Dynamically adjust sleep time to maintain the desired control frequency
            time.sleep(max(0, self.config.environment_dt - (time.perf_counter() - control_loop_start)))

@@ -15,15 +15,19 @@
 # limitations under the License.

 import platform
+from typing import cast
+
+from lerobot.utils.import_utils import make_device_from_device_class

 from .camera import Camera
 from .configs import CameraConfig, Cv2Rotation


 def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> dict[str, Camera]:
-    cameras = {}
+    cameras: dict[str, Camera] = {}

    for key, cfg in camera_configs.items():
+        # TODO(Steven): Consider just using the make_device_from_device_class for all types
        if cfg.type == "opencv":
            from .opencv import OpenCVCamera

@@ -40,7 +44,10 @@ def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> dict[s
            cameras[key] = Reachy2Camera(cfg)

        else:
-            raise ValueError(f"The camera type '{cfg.type}' is not valid.")
+            try:
+                cameras[key] = cast(Camera, make_device_from_device_class(cfg))
+            except Exception as e:
+                raise ValueError(f"Error creating camera {key} with config {cfg}: {e}") from e

    return cameras

@@ -71,9 +71,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}'.")
@@ -35,6 +35,8 @@ class NormalizationMode(str, Enum):
    MIN_MAX = "MIN_MAX"
    MEAN_STD = "MEAN_STD"
    IDENTITY = "IDENTITY"
+    QUANTILES = "QUANTILES"
+    QUANTILE10 = "QUANTILE10"


@dataclass
@@ -31,15 +31,15 @@ from lerobot.datasets.utils import (
    DEFAULT_EPISODES_PATH,
    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
    DEFAULT_VIDEO_PATH,
+    get_file_size_in_mb,
    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 concatenate_video_files
+from lerobot.datasets.video_utils import concatenate_video_files, get_video_duration_in_s


 def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
@@ -130,10 +130,34 @@ def update_meta_data(
    df["data/chunk_index"] = df["data/chunk_index"] + data_idx["chunk"]
    df["data/file_index"] = df["data/file_index"] + data_idx["file"]
    for key, video_idx in videos_idx.items():
-        df[f"videos/{key}/chunk_index"] = df[f"videos/{key}/chunk_index"] + video_idx["chunk"]
-        df[f"videos/{key}/file_index"] = df[f"videos/{key}/file_index"] + video_idx["file"]
-        df[f"videos/{key}/from_timestamp"] = df[f"videos/{key}/from_timestamp"] + video_idx["latest_duration"]
-        df[f"videos/{key}/to_timestamp"] = df[f"videos/{key}/to_timestamp"] + video_idx["latest_duration"]
+        # Store original video file indices before updating
+        orig_chunk_col = f"videos/{key}/chunk_index"
+        orig_file_col = f"videos/{key}/file_index"
+        df["_orig_chunk"] = df[orig_chunk_col].copy()
+        df["_orig_file"] = df[orig_file_col].copy()
+
+        # Update chunk and file indices to point to destination
+        df[orig_chunk_col] = video_idx["chunk"]
+        df[orig_file_col] = video_idx["file"]
+
+        # Apply per-source-file timestamp offsets
+        src_to_offset = video_idx.get("src_to_offset", {})
+        if src_to_offset:
+            # Apply offset based on original source file
+            for idx in df.index:
+                src_key = (df.at[idx, "_orig_chunk"], df.at[idx, "_orig_file"])
+                offset = src_to_offset.get(src_key, 0)
+                df.at[idx, f"videos/{key}/from_timestamp"] += offset
+                df.at[idx, f"videos/{key}/to_timestamp"] += offset
+        else:
+            # Fallback to simple offset (for backward compatibility)
+            df[f"videos/{key}/from_timestamp"] = (
+                df[f"videos/{key}/from_timestamp"] + video_idx["latest_duration"]
+            )
+            df[f"videos/{key}/to_timestamp"] = df[f"videos/{key}/to_timestamp"] + video_idx["latest_duration"]
+
+        # Clean up temporary columns
+        df = df.drop(columns=["_orig_chunk", "_orig_file"])

    df["dataset_from_index"] = df["dataset_from_index"] + dst_meta.info["total_frames"]
    df["dataset_to_index"] = df["dataset_to_index"] + dst_meta.info["total_frames"]
@@ -193,6 +217,10 @@ def aggregate_datasets(
        robot_type=robot_type,
        features=features,
        root=aggr_root,
+        use_videos=len(video_keys) > 0,
+        chunks_size=chunk_size,
+        data_files_size_in_mb=data_files_size_in_mb,
+        video_files_size_in_mb=video_files_size_in_mb,
    )

    logging.info("Find all tasks")
@@ -236,6 +264,11 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu
    Returns:
        dict: Updated videos_idx with current chunk and file indices.
    """
+    for key in videos_idx:
+        videos_idx[key]["episode_duration"] = 0
+        # Track offset for each source (chunk, file) pair
+        videos_idx[key]["src_to_offset"] = {}
+
    for key, video_idx in videos_idx.items():
        unique_chunk_file_pairs = {
            (chunk, file)
@@ -249,6 +282,7 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu

        chunk_idx = video_idx["chunk"]
        file_idx = video_idx["file"]
+        current_offset = video_idx["latest_duration"]

        for src_chunk_idx, src_file_idx in unique_chunk_file_pairs:
            src_path = src_meta.root / DEFAULT_VIDEO_PATH.format(
@@ -263,21 +297,25 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu
                file_index=file_idx,
            )

-            # If a new file is created, we don't want to increment the latest_duration
-            update_latest_duration = False
+            src_duration = get_video_duration_in_s(src_path)

            if not dst_path.exists():
-                # First write to this destination file
+                # Store offset before incrementing
+                videos_idx[key]["src_to_offset"][(src_chunk_idx, src_file_idx)] = current_offset
                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
+                videos_idx[key]["episode_duration"] += src_duration
+                current_offset += src_duration
+                continue

            # Check file sizes before appending
-            src_size = get_video_size_in_mb(src_path)
-            dst_size = get_video_size_in_mb(dst_path)
+            src_size = get_file_size_in_mb(src_path)
+            dst_size = get_file_size_in_mb(dst_path)

            if dst_size + src_size >= video_files_size_in_mb:
-                # Rotate to a new chunk/file
+                # Rotate to a new file, this source becomes start of new destination
+                # So its offset should be 0
+                videos_idx[key]["src_to_offset"][(src_chunk_idx, src_file_idx)] = 0
                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,
@@ -286,25 +324,22 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu
                )
                dst_path.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy(str(src_path), str(dst_path))
+                # Reset offset for next file
+                current_offset = src_duration
            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
+                # Append to existing video file - use current accumulated offset
+                videos_idx[key]["src_to_offset"][(src_chunk_idx, src_file_idx)] = current_offset
                concatenate_video_files(
                    [dst_path, src_path],
                    dst_path,
                )
-                # Update the latest_duration when appending (shifts timestamps!)
-                update_latest_duration = not update_latest_duration
+                current_offset += src_duration
+
+            videos_idx[key]["episode_duration"] += src_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


@@ -389,9 +424,6 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_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,
@@ -403,6 +435,10 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
            aggr_root=dst_meta.root,
        )

+    # Increment latest_duration by the total duration added from this source dataset
+    for k in videos_idx:
+        videos_idx[k]["latest_duration"] += videos_idx[k]["episode_duration"]
+
    return meta_idx


@@ -23,6 +23,9 @@ Please, update your dataset to the new format using this command:
 python -m lerobot.datasets.v30.convert_dataset_v21_to_v30 --repo-id={repo_id}
 ```

+If you already have a converted version uploaded to the hub, then this error might be because of
+an older version in your local cache. Consider deleting the cached version and retrying.
+
 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).
 """
@@ -17,6 +17,179 @@ import numpy as np

 from lerobot.datasets.utils import load_image_as_numpy

+DEFAULT_QUANTILES = [0.01, 0.10, 0.50, 0.90, 0.99]
+
+
+class RunningQuantileStats:
+    """
+    Maintains running statistics for batches of vectors, including mean,
+    standard deviation, min, max, and approximate quantiles.
+
+    Statistics are computed per feature dimension and updated incrementally
+    as new batches are observed. Quantiles are estimated using histograms,
+    which adapt dynamically if the observed data range expands.
+    """
+
+    def __init__(self, quantile_list: list[float] | None = None, num_quantile_bins: int = 5000):
+        self._count = 0
+        self._mean = None
+        self._mean_of_squares = None
+        self._min = None
+        self._max = None
+        self._histograms = None
+        self._bin_edges = None
+        self._num_quantile_bins = num_quantile_bins
+
+        self._quantile_list = quantile_list
+        if self._quantile_list is None:
+            self._quantile_list = DEFAULT_QUANTILES
+        self._quantile_keys = [f"q{int(q * 100):02d}" for q in self._quantile_list]
+
+    def update(self, batch: np.ndarray) -> None:
+        """Update the running statistics with a batch of vectors.
+
+        Args:
+            batch: An array where all dimensions except the last are batch dimensions.
+        """
+        batch = batch.reshape(-1, batch.shape[-1])
+        num_elements, vector_length = batch.shape
+
+        if self._count == 0:
+            self._mean = np.mean(batch, axis=0)
+            self._mean_of_squares = np.mean(batch**2, axis=0)
+            self._min = np.min(batch, axis=0)
+            self._max = np.max(batch, axis=0)
+            self._histograms = [np.zeros(self._num_quantile_bins) for _ in range(vector_length)]
+            self._bin_edges = [
+                np.linspace(self._min[i] - 1e-10, self._max[i] + 1e-10, self._num_quantile_bins + 1)
+                for i in range(vector_length)
+            ]
+        else:
+            if vector_length != self._mean.size:
+                raise ValueError("The length of new vectors does not match the initialized vector length.")
+
+            new_max = np.max(batch, axis=0)
+            new_min = np.min(batch, axis=0)
+            max_changed = np.any(new_max > self._max)
+            min_changed = np.any(new_min < self._min)
+            self._max = np.maximum(self._max, new_max)
+            self._min = np.minimum(self._min, new_min)
+
+            if max_changed or min_changed:
+                self._adjust_histograms()
+
+        self._count += num_elements
+
+        batch_mean = np.mean(batch, axis=0)
+        batch_mean_of_squares = np.mean(batch**2, axis=0)
+
+        # Update running mean and mean of squares
+        self._mean += (batch_mean - self._mean) * (num_elements / self._count)
+        self._mean_of_squares += (batch_mean_of_squares - self._mean_of_squares) * (
+            num_elements / self._count
+        )
+
+        self._update_histograms(batch)
+
+    def get_statistics(self) -> dict[str, np.ndarray]:
+        """Compute and return the statistics of the vectors processed so far.
+
+        Args:
+            quantiles: List of quantiles to compute (e.g., [0.01, 0.10, 0.50, 0.90, 0.99]). If None, no quantiles computed.
+
+        Returns:
+            Dictionary containing the computed statistics.
+        """
+        if self._count < 2:
+            raise ValueError("Cannot compute statistics for less than 2 vectors.")
+
+        variance = self._mean_of_squares - self._mean**2
+
+        stddev = np.sqrt(np.maximum(0, variance))
+
+        stats = {
+            "min": self._min.copy(),
+            "max": self._max.copy(),
+            "mean": self._mean.copy(),
+            "std": stddev,
+            "count": np.array([self._count]),
+        }
+
+        quantile_results = self._compute_quantiles()
+        for i, q in enumerate(self._quantile_keys):
+            stats[q] = quantile_results[i]
+
+        return stats
+
+    def _adjust_histograms(self):
+        """Adjust histograms when min or max changes."""
+        for i in range(len(self._histograms)):
+            old_edges = self._bin_edges[i]
+            old_hist = self._histograms[i]
+
+            # Create new edges with small padding to ensure range coverage
+            padding = (self._max[i] - self._min[i]) * 1e-10
+            new_edges = np.linspace(
+                self._min[i] - padding, self._max[i] + padding, self._num_quantile_bins + 1
+            )
+
+            # Redistribute existing histogram counts to new bins
+            # We need to map each old bin center to the new bins
+            old_centers = (old_edges[:-1] + old_edges[1:]) / 2
+            new_hist = np.zeros(self._num_quantile_bins)
+
+            for old_center, count in zip(old_centers, old_hist, strict=False):
+                if count > 0:
+                    # Find which new bin this old center belongs to
+                    bin_idx = np.searchsorted(new_edges, old_center) - 1
+                    bin_idx = max(0, min(bin_idx, self._num_quantile_bins - 1))
+                    new_hist[bin_idx] += count
+
+            self._histograms[i] = new_hist
+            self._bin_edges[i] = new_edges
+
+    def _update_histograms(self, batch: np.ndarray) -> None:
+        """Update histograms with new vectors."""
+        for i in range(batch.shape[1]):
+            hist, _ = np.histogram(batch[:, i], bins=self._bin_edges[i])
+            self._histograms[i] += hist
+
+    def _compute_quantiles(self) -> list[np.ndarray]:
+        """Compute quantiles based on histograms."""
+        results = []
+        for q in self._quantile_list:
+            target_count = q * self._count
+            q_values = []
+
+            for hist, edges in zip(self._histograms, self._bin_edges, strict=True):
+                q_value = self._compute_single_quantile(hist, edges, target_count)
+                q_values.append(q_value)
+
+            results.append(np.array(q_values))
+        return results
+
+    def _compute_single_quantile(self, hist: np.ndarray, edges: np.ndarray, target_count: float) -> float:
+        """Compute a single quantile value from histogram and bin edges."""
+        cumsum = np.cumsum(hist)
+        idx = np.searchsorted(cumsum, target_count)
+
+        if idx == 0:
+            return edges[0]
+        if idx >= len(cumsum):
+            return edges[-1]
+
+        # If not edge case, interpolate within the bin
+        count_before = cumsum[idx - 1]
+        count_in_bin = cumsum[idx] - count_before
+
+        # If no samples in this bin, use the bin edge
+        if count_in_bin == 0:
+            return edges[idx]
+
+        # Linear interpolation within the bin
+        fraction = (target_count - count_before) / count_in_bin
+        return edges[idx] + fraction * (edges[idx + 1] - edges[idx])
+

 def estimate_num_samples(
    dataset_len: int, min_num_samples: int = 100, max_num_samples: int = 10_000, power: float = 0.75
@@ -72,33 +245,282 @@ def sample_images(image_paths: list[str]) -> np.ndarray:
    return images


-def get_feature_stats(array: np.ndarray, axis: tuple, keepdims: bool) -> dict[str, np.ndarray]:
-    return {
-        "min": np.min(array, axis=axis, keepdims=keepdims),
-        "max": np.max(array, axis=axis, keepdims=keepdims),
-        "mean": np.mean(array, axis=axis, keepdims=keepdims),
-        "std": np.std(array, axis=axis, keepdims=keepdims),
-        "count": np.array([len(array)]),
+def _reshape_stats_by_axis(
+    stats: dict[str, np.ndarray],
+    axis: int | tuple[int, ...] | None,
+    keepdims: bool,
+    original_shape: tuple[int, ...],
+) -> dict[str, np.ndarray]:
+    """Reshape all statistics to match NumPy's output conventions.
+
+    Applies consistent reshaping to all statistics (except 'count') based on the
+    axis and keepdims parameters. This ensures statistics have the correct shape
+    for broadcasting with the original data.
+
+    Args:
+        stats: Dictionary of computed statistics
+        axis: Axis or axes along which statistics were computed
+        keepdims: Whether to keep reduced dimensions as size-1 dimensions
+        original_shape: Shape of the original array
+
+    Returns:
+        Dictionary with reshaped statistics
+
+    Note:
+        The 'count' statistic is never reshaped as it represents metadata
+        rather than per-feature statistics.
+    """
+    if axis == (1,) and not keepdims:
+        return stats
+
+    result = {}
+    for key, value in stats.items():
+        if key == "count":
+            result[key] = value
+        else:
+            result[key] = _reshape_single_stat(value, axis, keepdims, original_shape)
+
+    return result
+
+
+def _reshape_for_image_stats(value: np.ndarray, keepdims: bool) -> np.ndarray:
+    """Reshape statistics for image data (axis=(0,2,3))."""
+    if keepdims and value.ndim == 1:
+        return value.reshape(1, -1, 1, 1)
+    return value
+
+
+def _reshape_for_vector_stats(
+    value: np.ndarray, keepdims: bool, original_shape: tuple[int, ...]
+) -> np.ndarray:
+    """Reshape statistics for vector data (axis=0 or axis=(0,))."""
+    if not keepdims:
+        return value
+
+    if len(original_shape) == 1 and value.ndim > 0:
+        return value.reshape(1)
+    elif len(original_shape) >= 2 and value.ndim == 1:
+        return value.reshape(1, -1)
+    return value
+
+
+def _reshape_for_feature_stats(value: np.ndarray, keepdims: bool) -> np.ndarray:
+    """Reshape statistics for feature-wise computation (axis=(1,))."""
+    if not keepdims:
+        return value
+
+    if value.ndim == 0:
+        return value.reshape(1, 1)
+    elif value.ndim == 1:
+        return value.reshape(-1, 1)
+    return value
+
+
+def _reshape_for_global_stats(
+    value: np.ndarray, keepdims: bool, original_shape: tuple[int, ...]
+) -> np.ndarray | float:
+    """Reshape statistics for global reduction (axis=None)."""
+    if keepdims:
+        target_shape = tuple(1 for _ in original_shape)
+        return value.reshape(target_shape)
+    # Keep at least 1-D arrays to satisfy validator
+    return np.atleast_1d(value)
+
+
+def _reshape_single_stat(
+    value: np.ndarray, axis: int | tuple[int, ...] | None, keepdims: bool, original_shape: tuple[int, ...]
+) -> np.ndarray | float:
+    """Apply appropriate reshaping to a single statistic array.
+
+    This function transforms statistic arrays to match expected output shapes
+    based on the axis configuration and keepdims parameter.
+
+    Args:
+        value: The statistic array to reshape
+        axis: Axis or axes that were reduced during computation
+        keepdims: Whether to maintain reduced dimensions as size-1 dimensions
+        original_shape: Shape of the original data before reduction
+
+    Returns:
+        Reshaped array following NumPy broadcasting conventions
+
+    """
+    if axis == (0, 2, 3):
+        return _reshape_for_image_stats(value, keepdims)
+
+    if axis in [0, (0,)]:
+        return _reshape_for_vector_stats(value, keepdims, original_shape)
+
+    if axis == (1,):
+        return _reshape_for_feature_stats(value, keepdims)
+
+    if axis is None:
+        return _reshape_for_global_stats(value, keepdims, original_shape)
+
+    return value
+
+
+def _prepare_array_for_stats(array: np.ndarray, axis: int | tuple[int, ...] | None) -> tuple[np.ndarray, int]:
+    """Prepare array for statistics computation by reshaping according to axis.
+
+    Args:
+        array: Input data array
+        axis: Axis or axes along which to compute statistics
+
+    Returns:
+        Tuple of (reshaped_array, sample_count)
+    """
+    if axis == (0, 2, 3):  # Image data
+        batch_size, channels, height, width = array.shape
+        reshaped = array.transpose(0, 2, 3, 1).reshape(-1, channels)
+        return reshaped, batch_size
+
+    if axis == 0 or axis == (0,):  # Vector data
+        reshaped = array
+        if array.ndim == 1:
+            reshaped = array.reshape(-1, 1)
+        return reshaped, array.shape[0]
+
+    if axis == (1,):  # Feature-wise statistics
+        return array.T, array.shape[1]
+
+    if axis is None:  # Global statistics
+        reshaped = array.reshape(-1, 1)
+        # For backward compatibility, count represents the first dimension size
+        return reshaped, array.shape[0] if array.ndim > 0 else 1
+
+    raise ValueError(f"Unsupported axis configuration: {axis}")
+
+
+def _compute_basic_stats(
+    array: np.ndarray, sample_count: int, quantile_list: list[float] | None = None
+) -> dict[str, np.ndarray]:
+    """Compute basic statistics for arrays with insufficient samples for quantiles.
+
+    Args:
+        array: Reshaped array ready for statistics computation
+        sample_count: Number of samples represented in the data
+
+    Returns:
+        Dictionary with basic statistics and quantiles set to mean values
+    """
+    if quantile_list is None:
+        quantile_list = DEFAULT_QUANTILES
+    quantile_list_keys = [f"q{int(q * 100):02d}" for q in quantile_list]
+
+    stats = {
+        "min": np.min(array, axis=0),
+        "max": np.max(array, axis=0),
+        "mean": np.mean(array, axis=0),
+        "std": np.std(array, axis=0),
+        "count": np.array([sample_count]),
    }

+    for q in quantile_list_keys:
+        stats[q] = stats["mean"].copy()
+
+    return stats
+
+
+def get_feature_stats(
+    array: np.ndarray,
+    axis: int | tuple[int, ...] | None,
+    keepdims: bool,
+    quantile_list: list[float] | None = None,
+) -> dict[str, np.ndarray]:
+    """Compute comprehensive statistics for array features along specified axes.
+
+    This function calculates min, max, mean, std, and quantiles (1%, 10%, 50%, 90%, 99%)
+    for the input array along the specified axes. It handles different data layouts:
+    - Image data: axis=(0,2,3) computes per-channel statistics
+    - Vector data: axis=0 computes per-feature statistics
+    - Feature-wise: axis=1 computes statistics across features
+    - Global: axis=None computes statistics over entire array
+
+    Args:
+        array: Input data array with shape appropriate for the specified axis
+        axis: Axis or axes along which to compute statistics
+            - (0, 2, 3): For image data (batch, channels, height, width)
+            - 0 or (0,): For vector/tabular data (samples, features)
+            - (1,): For computing across features
+            - None: For global statistics over entire array
+        keepdims: If True, reduced axes are kept as dimensions with size 1
+
+    Returns:
+        Dictionary containing:
+            - 'min': Minimum values
+            - 'max': Maximum values
+            - 'mean': Mean values
+            - 'std': Standard deviation
+            - 'count': Number of samples (always shape (1,))
+            - 'q01', 'q10', 'q50', 'q90', 'q99': Quantile values
+
+    """
+    if quantile_list is None:
+        quantile_list = DEFAULT_QUANTILES
+
+    original_shape = array.shape
+    reshaped, sample_count = _prepare_array_for_stats(array, axis)
+
+    if reshaped.shape[0] < 2:
+        stats = _compute_basic_stats(reshaped, sample_count, quantile_list)
+    else:
+        running_stats = RunningQuantileStats()
+        running_stats.update(reshaped)
+        stats = running_stats.get_statistics()
+        stats["count"] = np.array([sample_count])
+
+    stats = _reshape_stats_by_axis(stats, axis, keepdims, original_shape)
+    return stats
+
+
+def compute_episode_stats(
+    episode_data: dict[str, list[str] | np.ndarray],
+    features: dict,
+    quantile_list: list[float] | None = None,
+) -> dict:
+    """Compute comprehensive statistics for all features in an episode.
+
+    Processes different data types appropriately:
+    - Images/videos: Samples from paths, computes per-channel stats, normalizes to [0,1]
+    - Numerical arrays: Computes per-feature statistics
+    - Strings: Skipped (no statistics computed)
+
+    Args:
+        episode_data: Dictionary mapping feature names to data
+            - For images/videos: list of file paths
+            - For numerical data: numpy arrays
+        features: Dictionary describing each feature's dtype and shape
+
+    Returns:
+        Dictionary mapping feature names to their statistics dictionaries.
+        Each statistics dictionary contains min, max, mean, std, count, and quantiles.
+
+    Note:
+        Image statistics are normalized to [0,1] range and have shape (3,1,1) for
+        per-channel values when dtype is 'image' or 'video'.
+    """
+    if quantile_list is None:
+        quantile_list = DEFAULT_QUANTILES

-def compute_episode_stats(episode_data: dict[str, list[str] | np.ndarray], features: dict) -> dict:
    ep_stats = {}
    for key, data in episode_data.items():
        if features[key]["dtype"] == "string":
-            continue  # HACK: we should receive np.arrays of strings
-        elif features[key]["dtype"] in ["image", "video"]:
-            ep_ft_array = sample_images(data)  # data is a list of image paths
-            axes_to_reduce = (0, 2, 3)  # keep channel dim
+            continue
+
+        if features[key]["dtype"] in ["image", "video"]:
+            ep_ft_array = sample_images(data)
+            axes_to_reduce = (0, 2, 3)
            keepdims = True
        else:
-            ep_ft_array = data  # data is already a np.ndarray
-            axes_to_reduce = 0  # compute stats over the first axis
-            keepdims = data.ndim == 1  # keep as np.array
+            ep_ft_array = data
+            axes_to_reduce = 0
+            keepdims = data.ndim == 1

-        ep_stats[key] = get_feature_stats(ep_ft_array, axis=axes_to_reduce, keepdims=keepdims)
+        ep_stats[key] = get_feature_stats(
+            ep_ft_array, axis=axes_to_reduce, keepdims=keepdims, quantile_list=quantile_list
+        )

-        # finally, we normalize and remove batch dim for images
        if features[key]["dtype"] in ["image", "video"]:
            ep_stats[key] = {
                k: v if k == "count" else np.squeeze(v / 255.0, axis=0) for k, v in ep_stats[key].items()
@@ -107,20 +529,37 @@ def compute_episode_stats(episode_data: dict[str, list[str] | np.ndarray], featu
    return ep_stats


+def _validate_stat_value(value: np.ndarray, key: str, feature_key: str) -> None:
+    """Validate a single statistic value."""
+    if not isinstance(value, np.ndarray):
+        raise ValueError(
+            f"Stats must be composed of numpy array, but key '{key}' of feature '{feature_key}' "
+            f"is of type '{type(value)}' instead."
+        )
+
+    if value.ndim == 0:
+        raise ValueError("Number of dimensions must be at least 1, and is 0 instead.")
+
+    if key == "count" and value.shape != (1,):
+        raise ValueError(f"Shape of 'count' must be (1), but is {value.shape} instead.")
+
+    if "image" in feature_key and key != "count" and value.shape != (3, 1, 1):
+        raise ValueError(f"Shape of quantile '{key}' must be (3,1,1), but is {value.shape} instead.")
+
+
 def _assert_type_and_shape(stats_list: list[dict[str, dict]]):
-    for i in range(len(stats_list)):
-        for fkey in stats_list[i]:
-            for k, v in stats_list[i][fkey].items():
-                if not isinstance(v, np.ndarray):
-                    raise ValueError(
-                        f"Stats must be composed of numpy array, but key '{k}' of feature '{fkey}' is of type '{type(v)}' instead."
-                    )
-                if v.ndim == 0:
-                    raise ValueError("Number of dimensions must be at least 1, and is 0 instead.")
-                if k == "count" and v.shape != (1,):
-                    raise ValueError(f"Shape of 'count' must be (1), but is {v.shape} instead.")
-                if "image" in fkey and k != "count" and v.shape != (3, 1, 1):
-                    raise ValueError(f"Shape of '{k}' must be (3,1,1), but is {v.shape} instead.")
+    """Validate that all statistics have correct types and shapes.
+
+    Args:
+        stats_list: List of statistics dictionaries to validate
+
+    Raises:
+        ValueError: If any statistic has incorrect type or shape
+    """
+    for stats in stats_list:
+        for feature_key, feature_stats in stats.items():
+            for stat_key, stat_value in feature_stats.items():
+                _validate_stat_value(stat_value, stat_key, feature_key)


 def aggregate_feature_stats(stats_ft_list: list[dict[str, dict]]) -> dict[str, dict[str, np.ndarray]]:
@@ -143,7 +582,7 @@ def aggregate_feature_stats(stats_ft_list: list[dict[str, dict]]) -> dict[str, d
    weighted_variances = (variances + delta_means**2) * counts
    total_variance = weighted_variances.sum(axis=0) / total_count

-    return {
+    aggregated = {
        "min": np.min(np.stack([s["min"] for s in stats_ft_list]), axis=0),
        "max": np.max(np.stack([s["max"] for s in stats_ft_list]), axis=0),
        "mean": total_mean,
@@ -151,6 +590,17 @@ def aggregate_feature_stats(stats_ft_list: list[dict[str, dict]]) -> dict[str, d
        "count": total_count,
    }

+    if stats_ft_list:
+        quantile_keys = [k for k in stats_ft_list[0] if k.startswith("q") and k[1:].isdigit()]
+
+        for q_key in quantile_keys:
+            if all(q_key in s for s in stats_ft_list):
+                quantile_values = np.stack([s[q_key] for s in stats_ft_list])
+                weighted_quantiles = quantile_values * counts
+                aggregated[q_key] = weighted_quantiles.sum(axis=0) / total_count
+
+    return aggregated
+

 def aggregate_stats(stats_list: list[dict[str, dict]]) -> dict[str, dict[str, np.ndarray]]:
    """Aggregate stats from multiple compute_stats outputs into a single set of stats.
@@ -68,7 +68,30 @@ def image_array_to_pil_image(image_array: np.ndarray, range_check: bool = True)
    return PIL.Image.fromarray(image_array)


-def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path):
+def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1):
+    """
+    Saves a NumPy array or PIL Image to a file.
+
+    This function handles both NumPy arrays and PIL Image objects, converting
+    the former to a PIL Image before saving. It includes error handling for
+    the save operation.
+
+    Args:
+        image (np.ndarray | PIL.Image.Image): The image data to save.
+        fpath (Path): The destination file path for the image.
+        compress_level (int, optional): The compression level for the saved
+            image, as used by PIL.Image.save(). Defaults to 1.
+            Refer to: https://github.com/huggingface/lerobot/pull/2135
+            for more details on the default value rationale.
+
+    Raises:
+        TypeError: If the input 'image' is not a NumPy array or a
+            PIL.Image.Image object.
+
+    Side Effects:
+        Prints an error message to the console if the image writing process
+        fails for any reason.
+    """
    try:
        if isinstance(image, np.ndarray):
            img = image_array_to_pil_image(image)
@@ -76,7 +99,7 @@ def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path):
            img = image
        else:
            raise TypeError(f"Unsupported image type: {type(image)}")
-        img.save(fpath)
+        img.save(fpath, compress_level=compress_level)
    except Exception as e:
        print(f"Error writing image {fpath}: {e}")

@@ -14,7 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import contextlib
-import gc
 import logging
 import shutil
 import tempfile
@@ -26,6 +25,8 @@ import numpy as np
 import packaging.version
 import pandas as pd
 import PIL.Image
+import pyarrow as pa
+import pyarrow.parquet as pq
 import torch
 import torch.utils
 from huggingface_hub import HfApi, snapshot_download
@@ -46,13 +47,9 @@ from lerobot.datasets.utils import (
    embed_images,
    flatten_dict,
    get_delta_indices,
-    get_hf_dataset_cache_dir,
-    get_hf_dataset_size_in_mb,
+    get_file_size_in_mb,
    get_hf_features_from_features,
-    get_parquet_file_size_in_mb,
-    get_parquet_num_frames,
    get_safe_version,
-    get_video_size_in_mb,
    hf_transform_to_torch,
    is_valid_version,
    load_episodes,
@@ -60,7 +57,6 @@ from lerobot.datasets.utils import (
    load_nested_dataset,
    load_stats,
    load_tasks,
-    to_parquet_with_hf_images,
    update_chunk_file_indices,
    validate_episode_buffer,
    validate_frame,
@@ -90,10 +86,15 @@ class LeRobotDatasetMetadata:
        root: str | Path | None = None,
        revision: str | None = None,
        force_cache_sync: bool = False,
+        metadata_buffer_size: int = 10,
    ):
        self.repo_id = repo_id
        self.revision = revision if revision else CODEBASE_VERSION
        self.root = Path(root) if root is not None else HF_LEROBOT_HOME / repo_id
+        self.writer = None
+        self.latest_episode = None
+        self.metadata_buffer: list[dict] = []
+        self.metadata_buffer_size = metadata_buffer_size

        try:
            if force_cache_sync:
@@ -107,6 +108,54 @@ class LeRobotDatasetMetadata:
            self.pull_from_repo(allow_patterns="meta/")
            self.load_metadata()

+    def _flush_metadata_buffer(self) -> None:
+        """Write all buffered episode metadata to parquet file."""
+        if not hasattr(self, "metadata_buffer") or len(self.metadata_buffer) == 0:
+            return
+
+        combined_dict = {}
+        for episode_dict in self.metadata_buffer:
+            for key, value in episode_dict.items():
+                if key not in combined_dict:
+                    combined_dict[key] = []
+                # Extract value and serialize numpy arrays
+                # because PyArrow's from_pydict function doesn't support numpy arrays
+                val = value[0] if isinstance(value, list) else value
+                combined_dict[key].append(val.tolist() if isinstance(val, np.ndarray) else val)
+
+        first_ep = self.metadata_buffer[0]
+        chunk_idx = first_ep["meta/episodes/chunk_index"][0]
+        file_idx = first_ep["meta/episodes/file_index"][0]
+
+        table = pa.Table.from_pydict(combined_dict)
+
+        if not self.writer:
+            path = Path(self.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx))
+            path.parent.mkdir(parents=True, exist_ok=True)
+            self.writer = pq.ParquetWriter(
+                path, schema=table.schema, compression="snappy", use_dictionary=True
+            )
+
+        self.writer.write_table(table)
+
+        self.latest_episode = self.metadata_buffer[-1]
+        self.metadata_buffer.clear()
+
+    def _close_writer(self) -> None:
+        """Close and cleanup the parquet writer if it exists."""
+        self._flush_metadata_buffer()
+
+        writer = getattr(self, "writer", None)
+        if writer is not None:
+            writer.close()
+            self.writer = None
+
+    def __del__(self):
+        """
+        Trust the user to call .finalize() but as an added safety check call the parquet writer to stop when calling the destructor
+        """
+        self._close_writer()
+
    def load_metadata(self):
        self.info = load_info(self.root)
        check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
@@ -138,6 +187,12 @@ class LeRobotDatasetMetadata:
        return packaging.version.parse(self.info["codebase_version"])

    def get_data_file_path(self, ep_index: int) -> Path:
+        if self.episodes is None:
+            self.episodes = load_episodes(self.root)
+        if ep_index >= len(self.episodes):
+            raise IndexError(
+                f"Episode index {ep_index} out of range. Episodes: {len(self.episodes) if self.episodes else 0}"
+            )
        ep = self.episodes[ep_index]
        chunk_idx = ep["data/chunk_index"]
        file_idx = ep["data/file_index"]
@@ -145,6 +200,12 @@ class LeRobotDatasetMetadata:
        return Path(fpath)

    def get_video_file_path(self, ep_index: int, vid_key: str) -> Path:
+        if self.episodes is None:
+            self.episodes = load_episodes(self.root)
+        if ep_index >= len(self.episodes):
+            raise IndexError(
+                f"Episode index {ep_index} out of range. Episodes: {len(self.episodes) if self.episodes else 0}"
+            )
        ep = self.episodes[ep_index]
        chunk_idx = ep[f"videos/{vid_key}/chunk_index"]
        file_idx = ep[f"videos/{vid_key}/file_index"]
@@ -260,72 +321,75 @@ class LeRobotDatasetMetadata:
            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.
+        """Buffer episode metadata and write to parquet in batches for efficiency.

-        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.
+        This function accumulates episode metadata in a buffer and flushes it when the buffer
+        reaches the configured size. This reduces I/O overhead by writing multiple episodes
+        at once instead of one row at a time.

        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
+        # Convert to list format for each value
        episode_dict = {key: [value] for key, value in episode_dict.items()}
-        ep_dataset = datasets.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]

-        if self.episodes is None:
+        if self.latest_episode 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"]
+            if self.episodes is not None and len(self.episodes) > 0:
+                # It means we are resuming recording, so we need to load the latest episode
+                # Update the indices to avoid overwriting the latest episode
+                chunk_idx = self.episodes[-1]["meta/episodes/chunk_index"]
+                file_idx = self.episodes[-1]["meta/episodes/file_index"]
+                latest_num_frames = self.episodes[-1]["dataset_to_index"]
+                episode_dict["dataset_from_index"] = [latest_num_frames]
+                episode_dict["dataset_to_index"] = [latest_num_frames + num_frames]

-            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
+                # When resuming, move to the next file
                chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, self.chunks_size)
+            else:
+                episode_dict["dataset_from_index"] = [0]
+                episode_dict["dataset_to_index"] = [num_frames]
+
+            episode_dict["meta/episodes/chunk_index"] = [chunk_idx]
+            episode_dict["meta/episodes/file_index"] = [file_idx]
+        else:
+            chunk_idx = self.latest_episode["meta/episodes/chunk_index"][0]
+            file_idx = self.latest_episode["meta/episodes/file_index"][0]
+
+            latest_path = (
+                self.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
+                if self.writer is None
+                else self.writer.where
+            )
+
+            if Path(latest_path).exists():
+                latest_size_in_mb = get_file_size_in_mb(Path(latest_path))
+                latest_num_frames = self.latest_episode["episode_index"][0]
+
+                av_size_per_frame = latest_size_in_mb / latest_num_frames if latest_num_frames > 0 else 0.0
+
+                if latest_size_in_mb + av_size_per_frame * num_frames >= self.data_files_size_in_mb:
+                    # Size limit is reached, flush buffer and prepare new parquet file
+                    self._flush_metadata_buffer()
+                    chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, self.chunks_size)
+                    self._close_writer()

            # 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]
+            episode_dict["meta/episodes/chunk_index"] = [chunk_idx]
+            episode_dict["meta/episodes/file_index"] = [file_idx]
+            episode_dict["dataset_from_index"] = [self.latest_episode["dataset_to_index"][0]]
+            episode_dict["dataset_to_index"] = [self.latest_episode["dataset_to_index"][0] + 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)
+        # Add to buffer
+        self.metadata_buffer.append(episode_dict)
+        self.latest_episode = episode_dict

-                # Memort optimization
-                del latest_df
-                gc.collect()
-
-        # 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)
-
-        if self.episodes is not None:
-            # Remove the episodes cache directory, necessary to avoid cache bloat
-            cached_dir = get_hf_dataset_cache_dir(self.episodes)
-            if cached_dir is not None:
-                shutil.rmtree(cached_dir)
-
-        self.episodes = load_episodes(self.root)
+        if len(self.metadata_buffer) >= self.metadata_buffer_size:
+            self._flush_metadata_buffer()

    def save_episode(
        self,
@@ -438,6 +502,10 @@ class LeRobotDatasetMetadata:
        robot_type: str | None = None,
        root: str | Path | None = None,
        use_videos: bool = True,
+        metadata_buffer_size: int = 10,
+        chunks_size: int | None = None,
+        data_files_size_in_mb: int | None = None,
+        video_files_size_in_mb: int | None = None,
    ) -> "LeRobotDatasetMetadata":
        """Creates metadata for a LeRobotDataset."""
        obj = cls.__new__(cls)
@@ -452,11 +520,24 @@ class LeRobotDatasetMetadata:
        obj.tasks = None
        obj.episodes = None
        obj.stats = None
-        obj.info = create_empty_dataset_info(CODEBASE_VERSION, fps, features, use_videos, robot_type)
+        obj.info = create_empty_dataset_info(
+            CODEBASE_VERSION,
+            fps,
+            features,
+            use_videos,
+            robot_type,
+            chunks_size,
+            data_files_size_in_mb,
+            video_files_size_in_mb,
+        )
        if len(obj.video_keys) > 0 and not use_videos:
            raise ValueError()
        write_json(obj.info, obj.root / INFO_PATH)
        obj.revision = None
+        obj.writer = None
+        obj.latest_episode = None
+        obj.metadata_buffer = []
+        obj.metadata_buffer_size = metadata_buffer_size
        return obj


@@ -603,6 +684,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
        # Unused attributes
        self.image_writer = None
        self.episode_buffer = None
+        self.writer = None
+        self.latest_episode = None

        self.root.mkdir(exist_ok=True, parents=True)

@@ -611,6 +694,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
            self.repo_id, self.root, self.revision, force_cache_sync=force_cache_sync
        )

+        # Track dataset state for efficient incremental writing
+        self._lazy_loading = False
+        self._recorded_frames = self.meta.total_frames
+        self._writer_closed_for_reading = False
+
        # Load actual data
        try:
            if force_cache_sync:
@@ -629,6 +717,19 @@ class LeRobotDataset(torch.utils.data.Dataset):
            check_delta_timestamps(self.delta_timestamps, self.fps, self.tolerance_s)
            self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps)

+    def _close_writer(self) -> None:
+        """Close and cleanup the parquet writer if it exists."""
+        writer = getattr(self, "writer", None)
+        if writer is not None:
+            writer.close()
+            self.writer = None
+
+    def __del__(self):
+        """
+        Trust the user to call .finalize() but as an added safety check call the parquet writer to stop when calling the destructor
+        """
+        self._close_writer()
+
    def push_to_hub(
        self,
        branch: str | None = None,
@@ -769,8 +870,15 @@ class LeRobotDataset(torch.utils.data.Dataset):

    @property
    def num_frames(self) -> int:
-        """Number of frames in selected episodes."""
-        return len(self.hf_dataset) if self.hf_dataset is not None else self.meta.total_frames
+        """Number of frames in selected episodes.
+
+        Note: When episodes a subset of the full dataset is requested, we must return the
+        actual loaded data length (len(self.hf_dataset)) rather than metadata total_frames.
+        self.meta.total_frames is the total number of frames in the full dataset.
+        """
+        if self.episodes is not None and self.hf_dataset is not None:
+            return len(self.hf_dataset)
+        return self.meta.total_frames

    @property
    def num_episodes(self) -> int:
@@ -848,10 +956,22 @@ class LeRobotDataset(torch.utils.data.Dataset):

        return item

+    def _ensure_hf_dataset_loaded(self):
+        """Lazy load the HF dataset only when needed for reading."""
+        if self._lazy_loading or self.hf_dataset is None:
+            # Close the writer before loading to ensure parquet file is properly finalized
+            if self.writer is not None:
+                self._close_writer()
+                self._writer_closed_for_reading = True
+            self.hf_dataset = self.load_hf_dataset()
+            self._lazy_loading = False
+
    def __len__(self):
        return self.num_frames

    def __getitem__(self, idx) -> dict:
+        # Ensure dataset is loaded when we actually need to read from it
+        self._ensure_hf_dataset_loaded()
        item = self.hf_dataset[idx]
        ep_idx = item["episode_index"].item()

@@ -890,6 +1010,14 @@ class LeRobotDataset(torch.utils.data.Dataset):
            "})',\n"
        )

+    def finalize(self):
+        """
+        Close the parquet writers. This function needs to be called after data collection/conversion, else footer metadata won't be written to the parquet files.
+        The dataset won't be valid and can't be loaded as ds = LeRobotDataset(repo_id=repo, root=HF_LEROBOT_HOME.joinpath(repo))
+        """
+        self._close_writer()
+        self.meta._close_writer()
+
    def create_episode_buffer(self, episode_index: int | None = None) -> dict:
        current_ep_idx = self.meta.total_episodes if episode_index is None else episode_index
        ep_buffer = {}
@@ -1097,74 +1225,101 @@ class LeRobotDataset(torch.utils.data.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:
+        if self.latest_episode 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
+            global_frame_index = 0
+            # However, if the episodes already exists
+            # It means we are resuming recording, so we need to load the latest episode
+            # Update the indices to avoid overwriting the latest episode
+            if self.meta.episodes is not None and len(self.meta.episodes) > 0:
+                latest_ep = self.meta.episodes[-1]
+                global_frame_index = latest_ep["dataset_to_index"]
+                chunk_idx = latest_ep["data/chunk_index"]
+                file_idx = latest_ep["data/file_index"]
+
+                # When resuming, move to the next file
+                chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, self.meta.chunks_size)
        else:
            # Retrieve information from the latest parquet file
-            latest_ep = self.meta.episodes[-1]
+            latest_ep = self.latest_episode
            chunk_idx = latest_ep["data/chunk_index"]
            file_idx = latest_ep["data/file_index"]
+            global_frame_index = latest_ep["index"][-1] + 1

            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)
+            latest_size_in_mb = get_file_size_in_mb(latest_path)
+
+            frames_in_current_file = global_frame_index - latest_ep["dataset_from_index"]
+            av_size_per_frame = (
+                latest_size_in_mb / frames_in_current_file if frames_in_current_file > 0 else 0
+            )

            # 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
+            if (
+                latest_size_in_mb + av_size_per_frame * ep_num_frames >= self.meta.data_files_size_in_mb
+                or self._writer_closed_for_reading
+            ):
+                # Size limit is reached or writer was closed for reading, 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)
+                self._close_writer()
+                self._writer_closed_for_reading = False

-                # Memort optimization
-                del latest_df
-                gc.collect()
+        ep_dict["data/chunk_index"] = chunk_idx
+        ep_dict["data/file_index"] = file_idx

        # 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)

-        if self.hf_dataset is not None:
-            # Remove hf dataset cache directory, necessary to avoid cache bloat
-            cached_dir = get_hf_dataset_cache_dir(self.hf_dataset)
-            if cached_dir is not None:
-                shutil.rmtree(cached_dir)
-
-        self.hf_dataset = self.load_hf_dataset()
+        table = ep_dataset.with_format("arrow")[:]
+        if not self.writer:
+            self.writer = pq.ParquetWriter(
+                path, schema=table.schema, compression="snappy", use_dictionary=True
+            )
+        self.writer.write_table(table)

        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,
+            "dataset_from_index": global_frame_index,
+            "dataset_to_index": global_frame_index + ep_num_frames,
        }
+
+        # Store metadata with episode data for next episode
+        self.latest_episode = {**ep_dict, **metadata}
+
+        # Mark that the HF dataset needs reloading (lazy loading approach)
+        # This avoids expensive reloading during sequential recording
+        self._lazy_loading = True
+        # Update recorded frames count for efficient length tracking
+        self._recorded_frames += ep_num_frames
+
        return metadata

    def _save_episode_video(self, video_key: str, episode_index: int) -> dict:
        # 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_size_in_mb = get_file_size_in_mb(ep_path)
        ep_duration_in_s = get_video_duration_in_s(ep_path)

-        if self.meta.episodes is None or (
-            f"videos/{video_key}/chunk_index" not in self.meta.episodes.column_names
-            or f"videos/{video_key}/file_index" not in self.meta.episodes.column_names
+        if (
+            episode_index == 0
+            or self.meta.latest_episode is None
+            or f"videos/{video_key}/chunk_index" not in self.meta.latest_episode
        ):
            # Initialize indices for a new dataset made of the first episode data
            chunk_idx, file_idx = 0, 0
+            if self.meta.episodes is not None and len(self.meta.episodes) > 0:
+                # It means we are resuming recording, so we need to load the latest episode
+                # Update the indices to avoid overwriting the latest episode
+                old_chunk_idx = self.meta.episodes[-1][f"videos/{video_key}/chunk_index"]
+                old_file_idx = self.meta.episodes[-1][f"videos/{video_key}/file_index"]
+                chunk_idx, file_idx = update_chunk_file_indices(
+                    old_chunk_idx, old_file_idx, self.meta.chunks_size
+                )
            latest_duration_in_s = 0.0
            new_path = self.root / self.meta.video_path.format(
                video_key=video_key, chunk_index=chunk_idx, file_index=file_idx
@@ -1172,16 +1327,16 @@ class LeRobotDataset(torch.utils.data.Dataset):
            new_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.move(str(ep_path), str(new_path))
        else:
-            # Retrieve information from the latest updated video file (possibly several episodes ago)
-            latest_ep = self.meta.episodes[episode_index - 1]
-            chunk_idx = latest_ep[f"videos/{video_key}/chunk_index"]
-            file_idx = latest_ep[f"videos/{video_key}/file_index"]
+            # Retrieve information from the latest updated video file using latest_episode
+            latest_ep = self.meta.latest_episode
+            chunk_idx = latest_ep[f"videos/{video_key}/chunk_index"][0]
+            file_idx = latest_ep[f"videos/{video_key}/file_index"][0]

            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)
+            latest_size_in_mb = get_file_size_in_mb(latest_path)
+            latest_duration_in_s = latest_ep[f"videos/{video_key}/to_timestamp"][0]

            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
@@ -1315,6 +1470,12 @@ class LeRobotDataset(torch.utils.data.Dataset):
        obj.delta_timestamps = None
        obj.delta_indices = None
        obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec()
+        obj.writer = None
+        obj.latest_episode = None
+        # Initialize tracking for incremental recording
+        obj._lazy_loading = False
+        obj._recorded_frames = 0
+        obj._writer_closed_for_reading = False
        return obj


@@ -30,7 +30,7 @@ import pandas
 import pandas as pd
 import pyarrow.parquet as pq
 import torch
-from datasets import Dataset, concatenate_datasets
+from datasets import Dataset
 from datasets.table import embed_table_storage
 from huggingface_hub import DatasetCard, DatasetCardData, HfApi
 from huggingface_hub.errors import RevisionNotFoundError
@@ -44,7 +44,7 @@ from lerobot.datasets.backward_compatibility import (
    ForwardCompatibilityError,
 )
 from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_STR
-from lerobot.utils.utils import is_valid_numpy_dtype_string
+from lerobot.utils.utils import SuppressProgressBars, 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
@@ -94,12 +94,6 @@ def get_hf_dataset_size_in_mb(hf_ds: Dataset) -> int:
    return hf_ds.data.nbytes // (1024**2)


-def get_hf_dataset_cache_dir(hf_ds: Dataset) -> Path | None:
-    if hf_ds.cache_files is None or len(hf_ds.cache_files) == 0:
-        return None
-    return Path(hf_ds.cache_files[0]["filename"]).parents[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
@@ -123,8 +117,9 @@ def load_nested_dataset(pq_dir: Path, features: datasets.Features | None = None)
        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)
+    with SuppressProgressBars():
+        datasets = Dataset.from_parquet([str(path) for path in paths], features=features)
+    return datasets


 def get_parquet_num_frames(parquet_path: str | Path) -> int:
@@ -132,10 +127,14 @@ def get_parquet_num_frames(parquet_path: str | Path) -> int:
    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_file_size_in_mb(file_path: Path) -> float:
+    """Get file size on disk in megabytes.
+
+    Args:
+        file_path (Path): Path to the file.
+    """
+    file_size_bytes = file_path.stat().st_size
+    return file_size_bytes / (1024**2)


 def flatten_dict(d: dict, parent_key: str = "", sep: str = "/") -> dict:
@@ -0,0 +1,260 @@
+#!/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 augments existing LeRobot datasets with quantile statistics.
+
+Most datasets created before the quantile feature was added do not contain
+quantile statistics (q01, q10, q50, q90, q99) in their metadata. This script:
+
+1. Loads an existing LeRobot dataset in v3.0 format
+2. Checks if it already contains quantile statistics
+3. If missing, computes quantile statistics for all features
+4. Updates the dataset metadata with the new quantile statistics
+
+Usage:
+
+```bash
+python src/lerobot/datasets/v30/augment_dataset_quantile_stats.py \
+    --repo-id=lerobot/pusht \
+```
+"""
+
+import argparse
+import concurrent.futures
+import logging
+from pathlib import Path
+
+import numpy as np
+import torch
+from huggingface_hub import HfApi
+from requests import HTTPError
+from tqdm import tqdm
+
+from lerobot.datasets.compute_stats import DEFAULT_QUANTILES, aggregate_stats, get_feature_stats
+from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
+from lerobot.datasets.utils import write_stats
+from lerobot.utils.utils import init_logging
+
+
+def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[str] | None = None) -> bool:
+    """Check if dataset statistics already contain quantile information.
+
+    Args:
+        stats: Dataset statistics dictionary
+
+    Returns:
+        True if quantile statistics are present, False otherwise
+    """
+    if quantile_list_keys is None:
+        quantile_list_keys = [f"q{int(q * 100):02d}" for q in DEFAULT_QUANTILES]
+
+    if stats is None:
+        return False
+
+    for feature_stats in stats.values():
+        if any(q_key in feature_stats for q_key in quantile_list_keys):
+            return True
+
+    return False
+
+
+def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
+    """Process a single episode and return its statistics.
+
+    Args:
+        dataset: The LeRobot dataset
+        episode_idx: Index of the episode to process
+
+    Returns:
+        Dictionary containing episode statistics
+    """
+    logging.info(f"Computing stats for episode {episode_idx}")
+
+    start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
+    end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]
+
+    collected_data: dict[str, list] = {}
+    for idx in range(start_idx, end_idx):
+        item = dataset[idx]
+        for key, value in item.items():
+            if key not in dataset.features:
+                continue
+
+            if key not in collected_data:
+                collected_data[key] = []
+            collected_data[key].append(value)
+
+    ep_stats = {}
+    for key, data_list in collected_data.items():
+        if dataset.features[key]["dtype"] == "string":
+            continue
+
+        data = torch.stack(data_list).cpu().numpy()
+        if dataset.features[key]["dtype"] in ["image", "video"]:
+            if data.dtype == np.uint8:
+                data = data.astype(np.float32) / 255.0
+
+            axes_to_reduce = (0, 2, 3)
+            keepdims = True
+        else:
+            axes_to_reduce = 0
+            keepdims = data.ndim == 1
+
+        ep_stats[key] = get_feature_stats(
+            data, axis=axes_to_reduce, keepdims=keepdims, quantile_list=DEFAULT_QUANTILES
+        )
+
+        if dataset.features[key]["dtype"] in ["image", "video"]:
+            ep_stats[key] = {
+                k: v if k == "count" else np.squeeze(v, axis=0) for k, v in ep_stats[key].items()
+            }
+
+    return ep_stats
+
+
+def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dict]:
+    """Compute quantile statistics for all episodes in the dataset.
+
+    Args:
+        dataset: The LeRobot dataset to compute statistics for
+
+    Returns:
+        Dictionary containing aggregated statistics with quantiles
+
+    Note:
+        Video decoding operations are not thread-safe, so we process episodes sequentially
+        when video keys are present. For datasets without videos, we use parallel processing
+        with ThreadPoolExecutor for better performance.
+    """
+    logging.info(f"Computing quantile statistics for dataset with {dataset.num_episodes} episodes")
+
+    episode_stats_list = []
+    has_videos = len(dataset.meta.video_keys) > 0
+
+    if has_videos:
+        logging.info("Dataset contains video keys - using sequential processing for thread safety")
+        for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
+            ep_stats = process_single_episode(dataset, episode_idx)
+            episode_stats_list.append(ep_stats)
+    else:
+        logging.info("Dataset has no video keys - using parallel processing for better performance")
+        max_workers = min(dataset.num_episodes, 16)
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+            future_to_episode = {
+                executor.submit(process_single_episode, dataset, episode_idx): episode_idx
+                for episode_idx in range(dataset.num_episodes)
+            }
+
+            episode_results = {}
+            with tqdm(total=dataset.num_episodes, desc="Processing episodes") as pbar:
+                for future in concurrent.futures.as_completed(future_to_episode):
+                    episode_idx = future_to_episode[future]
+                    ep_stats = future.result()
+                    episode_results[episode_idx] = ep_stats
+                    pbar.update(1)
+
+        for episode_idx in range(dataset.num_episodes):
+            if episode_idx in episode_results:
+                episode_stats_list.append(episode_results[episode_idx])
+
+    if not episode_stats_list:
+        raise ValueError("No episode data found for computing statistics")
+
+    logging.info(f"Aggregating statistics from {len(episode_stats_list)} episodes")
+    return aggregate_stats(episode_stats_list)
+
+
+def augment_dataset_with_quantile_stats(
+    repo_id: str,
+    root: str | Path | None = None,
+    overwrite: bool = False,
+) -> None:
+    """Augment a dataset with quantile statistics if they are missing.
+
+    Args:
+        repo_id: Repository ID of the dataset
+        root: Local root directory for the dataset
+        overwrite: Overwrite existing quantile statistics if they already exist
+    """
+    logging.info(f"Loading dataset: {repo_id}")
+    dataset = LeRobotDataset(
+        repo_id=repo_id,
+        root=root,
+    )
+
+    if not overwrite and has_quantile_stats(dataset.meta.stats):
+        logging.info("Dataset already contains quantile statistics. No action needed.")
+        return
+
+    logging.info("Dataset does not contain quantile statistics. Computing them now...")
+
+    new_stats = compute_quantile_stats_for_dataset(dataset)
+
+    logging.info("Updating dataset metadata with new quantile statistics")
+    dataset.meta.stats = new_stats
+
+    write_stats(new_stats, dataset.meta.root)
+
+    logging.info("Successfully updated dataset with quantile statistics")
+    dataset.push_to_hub()
+
+    hub_api = HfApi()
+    try:
+        hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
+    except HTTPError as e:
+        logging.info(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
+        pass
+    hub_api.create_tag(repo_id, tag=CODEBASE_VERSION, revision=None, repo_type="dataset")
+
+
+def main():
+    """Main function to run the augmentation script."""
+    parser = argparse.ArgumentParser(description="Augment LeRobot dataset with quantile statistics")
+
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        required=True,
+        help="Repository ID of the dataset (e.g., 'lerobot/pusht')",
+    )
+
+    parser.add_argument(
+        "--root",
+        type=str,
+        help="Local root directory for the dataset",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Overwrite existing quantile statistics if they already exist",
+    )
+
+    args = parser.parse_args()
+    root = Path(args.root) if args.root else None
+
+    init_logging()
+
+    augment_dataset_with_quantile_stats(
+        repo_id=args.repo_id,
+        root=root,
+        overwrite=args.overwrite,
+    )
+
+
+if __name__ == "__main__":
+    main()
@@ -26,11 +26,20 @@ This script will help you convert any LeRobot dataset already pushed to the hub

 Usage:

+Convert a dataset from the hub:
 ```bash
 python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
    --repo-id=lerobot/pusht
 ```

+Convert a local dataset (works in place):
+```bash
+python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
+    --repo-id=lerobot/pusht \
+    --root=/path/to/local/dataset/directory
+    --push-to-hub=false
+```
+
 """

 import argparse
@@ -60,9 +69,9 @@ from lerobot.datasets.utils import (
    LEGACY_TASKS_PATH,
    cast_stats_to_numpy,
    flatten_dict,
+    get_file_size_in_mb,
    get_parquet_file_size_in_mb,
    get_parquet_num_frames,
-    get_video_size_in_mb,
    load_info,
    update_chunk_file_indices,
    write_episodes,
@@ -75,7 +84,7 @@ from lerobot.utils.constants import HF_LEROBOT_HOME
 from lerobot.utils.utils import init_logging

 V21 = "v2.1"
-
+V30 = "v3.0"

 """
 -------------------------
@@ -145,6 +154,17 @@ def legacy_load_tasks(local_dir: Path) -> tuple[dict, dict]:
    return tasks, task_to_task_index


+def validate_local_dataset_version(local_path: Path) -> None:
+    """Validate that the local dataset has the expected v2.1 version."""
+    info = load_info(local_path)
+    dataset_version = info.get("codebase_version", "unknown")
+    if dataset_version != V21:
+        raise ValueError(
+            f"Local dataset has codebase version '{dataset_version}', expected '{V21}'. "
+            f"This script is specifically for converting v2.1 datasets to v3.0."
+        )
+
+
 def convert_tasks(root, new_root):
    logging.info(f"Converting tasks from {root} to {new_root}")
    tasks, _ = legacy_load_tasks(root)
@@ -290,7 +310,7 @@ def convert_videos_of_camera(root: Path, new_root: Path, video_key: str, video_f
    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_size_in_mb = get_file_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
@@ -407,13 +427,13 @@ def convert_episodes_metadata(root, new_root, episodes_metadata, episodes_video_

 def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb):
    info = load_info(root)
-    info["codebase_version"] = "v3.0"
+    info["codebase_version"] = V30
    del info["total_chunks"]
    del info["total_videos"]
    info["data_files_size_in_mb"] = data_file_size_in_mb
    info["video_files_size_in_mb"] = video_file_size_in_mb
    info["data_path"] = DEFAULT_DATA_PATH
-    info["video_path"] = DEFAULT_VIDEO_PATH
+    info["video_path"] = DEFAULT_VIDEO_PATH if info["video_path"] is not None else None
    info["fps"] = int(info["fps"])
    logging.info(f"Converting info from {root} to {new_root}")
    for key in info["features"]:
@@ -429,16 +449,36 @@ def convert_dataset(
    branch: str | None = None,
    data_file_size_in_mb: int | None = None,
    video_file_size_in_mb: int | None = None,
+    root: str | Path | None = None,
+    push_to_hub: bool = True,
+    force_conversion: bool = False,
 ):
-    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 data_file_size_in_mb is None:
        data_file_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
    if video_file_size_in_mb is None:
        video_file_size_in_mb = DEFAULT_VIDEO_FILE_SIZE_IN_MB

+    # First check if the dataset already has a v3.0 version
+    if root is None and not force_conversion:
+        try:
+            print("Trying to download v3.0 version of the dataset from the hub...")
+            snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id)
+            return
+        except Exception:
+            print("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
+
+    # Set root based on whether local dataset path is provided
+    use_local_dataset = False
+    root = HF_LEROBOT_HOME / repo_id if root is None else Path(root) / repo_id
+    if root.exists():
+        validate_local_dataset_version(root)
+        use_local_dataset = True
+        print(f"Using local dataset at {root}")
+
+    old_root = root.parent / f"{root.name}_old"
+    new_root = root.parent / f"{root.name}_v30"
+
+    # Handle old_root cleanup if both old_root and root exist
    if old_root.is_dir() and root.is_dir():
        shutil.rmtree(str(root))
        shutil.move(str(old_root), str(root))
@@ -446,12 +486,13 @@ def convert_dataset(
    if new_root.is_dir():
        shutil.rmtree(new_root)

-    snapshot_download(
-        repo_id,
-        repo_type="dataset",
-        revision=V21,
-        local_dir=root,
-    )
+    if not use_local_dataset:
+        snapshot_download(
+            repo_id,
+            repo_type="dataset",
+            revision=V21,
+            local_dir=root,
+        )

    convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb)
    convert_tasks(root, new_root)
@@ -462,21 +503,22 @@ def convert_dataset(
    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")
+    if push_to_hub:
+        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()
+        LeRobotDataset(repo_id).push_to_hub()


 if __name__ == "__main__":
@@ -507,6 +549,23 @@ if __name__ == "__main__":
        default=None,
        help="File size in MB. Defaults to 100 for data and 500 for videos.",
    )
+    parser.add_argument(
+        "--root",
+        type=str,
+        default=None,
+        help="Local directory to use for downloading/writing the dataset.",
+    )
+    parser.add_argument(
+        "--push-to-hub",
+        type=lambda input: input.lower() == "true",
+        default=True,
+        help="Push the converted dataset to the hub.",
+    )
+    parser.add_argument(
+        "--force-conversion",
+        action="store_true",
+        help="Force conversion even if the dataset already has a v3.0 version.",
+    )

    args = parser.parse_args()
    convert_dataset(**vars(args))
@@ -451,11 +451,9 @@ def concatenate_video_files(
            stream_map[input_stream.index] = output_container.add_stream_from_template(
                template=input_stream, opaque=True
            )
-            stream_map[
-                input_stream.index
-            ].time_base = (
-                input_stream.time_base
-            )  # set the time base to the input stream time base (missing in the codec context)
+
+            # set the time base to the input stream time base (missing in the codec context)
+            stream_map[input_stream.index].time_base = input_stream.time_base

    # Demux + remux packets (no re-encode)
    for packet in input_container.demux():
@@ -644,6 +642,9 @@ class VideoEncodingManager:
            )
            self.dataset._batch_save_episode_video(start_ep, end_ep)

+        # Finalize the dataset to properly close all writers
+        self.dataset.finalize()
+
        # Clean up episode images if recording was interrupted
        if exc_type is not None:
            interrupted_episode_index = self.dataset.num_episodes
@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

-from .configs import AlohaEnv, EnvConfig, PushtEnv, XarmEnv  # noqa: F401
+from .configs import AlohaEnv, EnvConfig, PushtEnv  # noqa: F401
@@ -50,6 +50,8 @@ class AlohaEnv(EnvConfig):
    fps: int = 50
    episode_length: int = 400
    obs_type: str = "pixels_agent_pos"
+    observation_height: int = 480
+    observation_width: int = 640
    render_mode: str = "rgb_array"
    features: dict[str, PolicyFeature] = field(
        default_factory=lambda: {
@@ -67,10 +69,14 @@ class AlohaEnv(EnvConfig):

    def __post_init__(self):
        if self.obs_type == "pixels":
-            self.features["top"] = PolicyFeature(type=FeatureType.VISUAL, shape=(480, 640, 3))
+            self.features["top"] = PolicyFeature(
+                type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
+            )
        elif self.obs_type == "pixels_agent_pos":
            self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(14,))
-            self.features["pixels/top"] = PolicyFeature(type=FeatureType.VISUAL, shape=(480, 640, 3))
+            self.features["pixels/top"] = PolicyFeature(
+                type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
+            )

    @property
    def gym_kwargs(self) -> dict:
@@ -91,6 +97,8 @@ class PushtEnv(EnvConfig):
    render_mode: str = "rgb_array"
    visualization_width: int = 384
    visualization_height: int = 384
+    observation_height: int = 384
+    observation_width: int = 384
    features: dict[str, PolicyFeature] = field(
        default_factory=lambda: {
            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
@@ -108,7 +116,9 @@ class PushtEnv(EnvConfig):

    def __post_init__(self):
        if self.obs_type == "pixels_agent_pos":
-            self.features["pixels"] = PolicyFeature(type=FeatureType.VISUAL, shape=(384, 384, 3))
+            self.features["pixels"] = PolicyFeature(
+                type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
+            )
        elif self.obs_type == "environment_state_agent_pos":
            self.features["environment_state"] = PolicyFeature(type=FeatureType.ENV, shape=(16,))

@@ -123,45 +133,6 @@ class PushtEnv(EnvConfig):
        }


-@EnvConfig.register_subclass("xarm")
-@dataclass
-class XarmEnv(EnvConfig):
-    task: str | None = "XarmLift-v0"
-    fps: int = 15
-    episode_length: int = 200
-    obs_type: str = "pixels_agent_pos"
-    render_mode: str = "rgb_array"
-    visualization_width: int = 384
-    visualization_height: int = 384
-    features: dict[str, PolicyFeature] = field(
-        default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(4,)),
-            "pixels": PolicyFeature(type=FeatureType.VISUAL, shape=(84, 84, 3)),
-        }
-    )
-    features_map: dict[str, str] = field(
-        default_factory=lambda: {
-            ACTION: ACTION,
-            "agent_pos": OBS_STATE,
-            "pixels": OBS_IMAGE,
-        }
-    )
-
-    def __post_init__(self):
-        if self.obs_type == "pixels_agent_pos":
-            self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(4,))
-
-    @property
-    def gym_kwargs(self) -> dict:
-        return {
-            "obs_type": self.obs_type,
-            "render_mode": self.render_mode,
-            "visualization_width": self.visualization_width,
-            "visualization_height": self.visualization_height,
-            "max_episode_steps": self.episode_length,
-        }
-
-
@dataclass
 class ImagePreprocessingConfig:
    crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
@@ -254,7 +225,9 @@ class LiberoEnv(EnvConfig):
    render_mode: str = "rgb_array"
    camera_name: str = "agentview_image,robot0_eye_in_hand_image"
    init_states: bool = True
-    camera_name_mapping: dict[str, str] | None = (None,)
+    camera_name_mapping: dict[str, str] | None = None
+    observation_height: int = 360
+    observation_width: int = 360
    features: dict[str, PolicyFeature] = field(
        default_factory=lambda: {
            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
@@ -272,18 +245,18 @@ class LiberoEnv(EnvConfig):
    def __post_init__(self):
        if self.obs_type == "pixels":
            self.features["pixels/agentview_image"] = PolicyFeature(
-                type=FeatureType.VISUAL, shape=(360, 360, 3)
+                type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
            )
            self.features["pixels/robot0_eye_in_hand_image"] = PolicyFeature(
-                type=FeatureType.VISUAL, shape=(360, 360, 3)
+                type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 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)
+                type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
            )
            self.features["pixels/robot0_eye_in_hand_image"] = PolicyFeature(
-                type=FeatureType.VISUAL, shape=(360, 360, 3)
+                type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
            )
        else:
            raise ValueError(f"Unsupported obs_type: {self.obs_type}")
@@ -294,3 +267,45 @@ class LiberoEnv(EnvConfig):
            "obs_type": self.obs_type,
            "render_mode": self.render_mode,
        }
+
+
+@EnvConfig.register_subclass("metaworld")
+@dataclass
+class MetaworldEnv(EnvConfig):
+    task: str = "metaworld-push-v2"  # add all tasks
+    fps: int = 80
+    episode_length: int = 400
+    obs_type: str = "pixels_agent_pos"
+    render_mode: str = "rgb_array"
+    multitask_eval: bool = True
+    features: dict[str, PolicyFeature] = field(
+        default_factory=lambda: {
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(4,)),
+        }
+    )
+    features_map: dict[str, str] = field(
+        default_factory=lambda: {
+            "action": ACTION,
+            "agent_pos": OBS_STATE,
+            "top": f"{OBS_IMAGE}",
+            "pixels/top": f"{OBS_IMAGE}",
+        }
+    )
+
+    def __post_init__(self):
+        if self.obs_type == "pixels":
+            self.features["top"] = PolicyFeature(type=FeatureType.VISUAL, shape=(480, 480, 3))
+
+        elif self.obs_type == "pixels_agent_pos":
+            self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(4,))
+            self.features["pixels/top"] = PolicyFeature(type=FeatureType.VISUAL, shape=(480, 480, 3))
+
+        else:
+            raise ValueError(f"Unsupported obs_type: {self.obs_type}")
+
+    @property
+    def gym_kwargs(self) -> dict:
+        return {
+            "obs_type": self.obs_type,
+            "render_mode": self.render_mode,
+        }
@@ -17,7 +17,7 @@ import importlib

 import gymnasium as gym

-from lerobot.envs.configs import AlohaEnv, EnvConfig, LiberoEnv, PushtEnv, XarmEnv
+from lerobot.envs.configs import AlohaEnv, EnvConfig, LiberoEnv, PushtEnv


 def make_env_config(env_type: str, **kwargs) -> EnvConfig:
@@ -25,8 +25,6 @@ def make_env_config(env_type: str, **kwargs) -> EnvConfig:
        return AlohaEnv(**kwargs)
    elif env_type == "pusht":
        return PushtEnv(**kwargs)
-    elif env_type == "xarm":
-        return XarmEnv(**kwargs)
    elif env_type == "libero":
        return LiberoEnv(**kwargs)
    else:
@@ -63,6 +61,9 @@ def make_env(
    if "libero" in cfg.type:
        from lerobot.envs.libero import create_libero_envs

+        if cfg.task is None:
+            raise ValueError("LiberoEnv requires a task to be specified")
+
        return create_libero_envs(
            task=cfg.task,
            n_envs=n_envs,
@@ -71,7 +72,18 @@ def make_env(
            gym_kwargs=cfg.gym_kwargs,
            env_cls=env_cls,
        )
+    elif "metaworld" in cfg.type:
+        from lerobot.envs.metaworld import create_metaworld_envs

+        if cfg.task is None:
+            raise ValueError("MetaWorld requires a task to be specified")
+
+        return create_metaworld_envs(
+            task=cfg.task,
+            n_envs=n_envs,
+            gym_kwargs=cfg.gym_kwargs,
+            env_cls=env_cls,
+        )
    package_name = f"gym_{cfg.type}"
    try:
        importlib.import_module(package_name)
@@ -84,7 +96,7 @@ def make_env(
    def _make_one():
        return gym.make(gym_handle, disable_env_checker=cfg.disable_env_checker, **(cfg.gym_kwargs or {}))

-    vec = env_cls([_make_one for _ in range(n_envs)])
+    vec = env_cls([_make_one for _ in range(n_envs)], autoreset_mode=gym.vector.AutoresetMode.SAME_STEP)

    # normalize to {suite: {task_id: vec_env}} for consistency
    suite_name = cfg.type  # e.g., "pusht", "aloha"
@@ -260,19 +260,23 @@ class LiberoEnv(gym.Env):

        is_success = self._env.check_success()
        terminated = done or is_success
-        info["is_success"] = is_success
-
+        info.update(
+            {
+                "task": self.task,
+                "task_id": self.task_id,
+                "done": done,
+                "is_success": is_success,
+            }
+        )
        observation = self._format_raw_obs(raw_obs)
-        if done:
+        if terminated:
+            info["final_info"] = {
+                "task": self.task,
+                "task_id": self.task_id,
+                "done": bool(done),
+                "is_success": bool(is_success),
+            }
            self.reset()
-            info.update(
-                {
-                    "task": self.task,
-                    "task_id": self.task_id,
-                    "done": done,
-                    "is_success": is_success,
-                }
-            )
        truncated = False
        return observation, reward, terminated, truncated, info

@@ -0,0 +1,313 @@
+#!/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 json
+from collections import defaultdict
+from collections.abc import Callable, Sequence
+from pathlib import Path
+from typing import Any
+
+import gymnasium as gym
+import metaworld
+import metaworld.policies as policies
+import numpy as np
+from gymnasium import spaces
+
+# ---- Load configuration data from the external JSON file ----
+CONFIG_PATH = Path(__file__).parent / "metaworld_config.json"
+try:
+    with open(CONFIG_PATH) as f:
+        data = json.load(f)
+except FileNotFoundError as err:
+    raise FileNotFoundError(
+        "Could not find 'metaworld_config.json'. "
+        "Please ensure the configuration file is in the same directory as the script."
+    ) from err
+except json.JSONDecodeError as err:
+    raise ValueError(
+        "Failed to decode 'metaworld_config.json'. Please ensure it is a valid JSON file."
+    ) from err
+
+# ---- Process the loaded data ----
+
+# extract and type-check top-level dicts
+task_descriptions_obj = data.get("TASK_DESCRIPTIONS")
+if not isinstance(task_descriptions_obj, dict):
+    raise TypeError("Expected TASK_DESCRIPTIONS to be a dict[str, str]")
+TASK_DESCRIPTIONS: dict[str, str] = task_descriptions_obj
+
+task_name_to_id_obj = data.get("TASK_NAME_TO_ID")
+if not isinstance(task_name_to_id_obj, dict):
+    raise TypeError("Expected TASK_NAME_TO_ID to be a dict[str, int]")
+TASK_NAME_TO_ID: dict[str, int] = task_name_to_id_obj
+
+# difficulty -> tasks mapping
+difficulty_to_tasks = data.get("DIFFICULTY_TO_TASKS")
+if not isinstance(difficulty_to_tasks, dict):
+    raise TypeError("Expected 'DIFFICULTY_TO_TASKS' to be a dict[str, list[str]]")
+DIFFICULTY_TO_TASKS: dict[str, list[str]] = difficulty_to_tasks
+
+# convert policy strings -> actual policy classes
+task_policy_mapping = data.get("TASK_POLICY_MAPPING")
+if not isinstance(task_policy_mapping, dict):
+    raise TypeError("Expected 'TASK_POLICY_MAPPING' to be a dict[str, str]")
+TASK_POLICY_MAPPING: dict[str, Any] = {
+    task_name: getattr(policies, policy_class_name)
+    for task_name, policy_class_name in task_policy_mapping.items()
+}
+ACTION_DIM = 4
+OBS_DIM = 4
+
+
+class MetaworldEnv(gym.Env):
+    metadata = {"render_modes": ["rgb_array"], "render_fps": 80}
+
+    def __init__(
+        self,
+        task,
+        camera_name="corner2",
+        obs_type="pixels",
+        render_mode="rgb_array",
+        observation_width=480,
+        observation_height=480,
+        visualization_width=640,
+        visualization_height=480,
+    ):
+        super().__init__()
+        self.task = task.replace("metaworld-", "")
+        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.camera_name = camera_name
+
+        self._env = self._make_envs_task(self.task)
+        self._max_episode_steps = self._env.max_path_length
+        self.task_description = TASK_DESCRIPTIONS[self.task]
+
+        self.expert_policy = TASK_POLICY_MAPPING[self.task]()
+
+        if self.obs_type == "state":
+            raise NotImplementedError()
+        elif self.obs_type == "pixels":
+            self.observation_space = spaces.Dict(
+                {
+                    "pixels": spaces.Box(
+                        low=0,
+                        high=255,
+                        shape=(self.observation_height, self.observation_width, 3),
+                        dtype=np.uint8,
+                    )
+                }
+            )
+        elif self.obs_type == "pixels_agent_pos":
+            self.observation_space = spaces.Dict(
+                {
+                    "pixels": spaces.Box(
+                        low=0,
+                        high=255,
+                        shape=(self.observation_height, self.observation_width, 3),
+                        dtype=np.uint8,
+                    ),
+                    "agent_pos": spaces.Box(
+                        low=-1000.0,
+                        high=1000.0,
+                        shape=(OBS_DIM,),
+                        dtype=np.float64,
+                    ),
+                }
+            )
+
+        self.action_space = spaces.Box(low=-1, high=1, shape=(ACTION_DIM,), dtype=np.float32)
+
+    def render(self) -> np.ndarray:
+        """
+        Render the current environment frame.
+
+        Returns:
+            np.ndarray: The rendered RGB image from the environment.
+        """
+        image = self._env.render()
+        if self.camera_name == "corner2":
+            # Images from this camera are flipped — correct them
+            image = np.flip(image, (0, 1))
+        return image
+
+    def _make_envs_task(self, env_name: str):
+        mt1 = metaworld.MT1(env_name, seed=42)
+        env = mt1.train_classes[env_name](render_mode="rgb_array", camera_name=self.camera_name)
+        env.set_task(mt1.train_tasks[0])
+        if self.camera_name == "corner2":
+            env.model.cam_pos[2] = [
+                0.75,
+                0.075,
+                0.7,
+            ]  # corner2 position, similar to https://arxiv.org/pdf/2206.14244
+        env.reset()
+        env._freeze_rand_vec = False  # otherwise no randomization
+        return env
+
+    def _format_raw_obs(self, raw_obs: np.ndarray) -> dict[str, Any]:
+        image = None
+        if self._env is not None:
+            image = self._env.render()
+            if self.camera_name == "corner2":
+                # NOTE: The "corner2" camera in MetaWorld environments outputs images with both axes inverted.
+                image = np.flip(image, (0, 1))
+        agent_pos = raw_obs[:4]
+        if self.obs_type == "state":
+            raise NotImplementedError(
+                "'state' obs_type not implemented for MetaWorld. Use pixel modes instead."
+            )
+
+        elif self.obs_type in ("pixels", "pixels_agent_pos"):
+            assert image is not None, (
+                "Expected `image` to be rendered before constructing pixel-based observations. "
+                "This likely means `env.render()` returned None or the environment was not provided."
+            )
+
+            if self.obs_type == "pixels":
+                obs = {"pixels": image.copy()}
+
+            else:  # pixels_agent_pos
+                obs = {
+                    "pixels": image.copy(),
+                    "agent_pos": agent_pos,
+                }
+        else:
+            raise ValueError(f"Unknown obs_type: {self.obs_type}")
+        return obs
+
+    def reset(
+        self,
+        seed: int | None = None,
+        **kwargs,
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """
+        Reset the environment to its initial state.
+
+        Args:
+            seed (Optional[int]): Random seed for environment initialization.
+
+        Returns:
+            observation (Dict[str, Any]): The initial formatted observation.
+            info (Dict[str, Any]): Additional info about the reset state.
+        """
+        super().reset(seed=seed)
+
+        raw_obs, info = self._env.reset(seed=seed)
+
+        observation = self._format_raw_obs(raw_obs)
+
+        info = {"is_success": False}
+        return observation, info
+
+    def step(self, action: np.ndarray) -> tuple[dict[str, Any], float, bool, bool, dict[str, Any]]:
+        """
+        Perform one environment step.
+
+        Args:
+            action (np.ndarray): The action to execute, must be 1-D with shape (action_dim,).
+
+        Returns:
+            observation (Dict[str, Any]): The formatted observation after the step.
+            reward (float): The scalar reward for this step.
+            terminated (bool): Whether the episode terminated successfully.
+            truncated (bool): Whether the episode was truncated due to a time limit.
+            info (Dict[str, Any]): Additional environment info.
+        """
+        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, truncated, info = self._env.step(action)
+
+        # Determine whether the task was successful
+        is_success = bool(info.get("success", 0))
+        terminated = done or is_success
+        info.update(
+            {
+                "task": self.task,
+                "done": done,
+                "is_success": is_success,
+            }
+        )
+
+        # Format the raw observation into the expected structure
+        observation = self._format_raw_obs(raw_obs)
+        if terminated:
+            info["final_info"] = {
+                "task": self.task,
+                "done": bool(done),
+                "is_success": bool(is_success),
+            }
+            self.reset()
+
+        return observation, reward, terminated, truncated, info
+
+    def close(self):
+        self._env.close()
+
+
+# ---- Main API ----------------------------------------------------------------
+
+
+def create_metaworld_envs(
+    task: str,
+    n_envs: int,
+    gym_kwargs: dict[str, Any] | None = None,
+    env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
+) -> dict[str, dict[int, Any]]:
+    """
+    Create vectorized Meta-World environments with a consistent return shape.
+
+    Returns:
+        dict[task_group][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 difficulty group (e.g., "easy", "medium", "hard") or a comma-separated list.
+        - If a task name is not in DIFFICULTY_TO_TASKS, we treat it as a single custom task.
+    """
+    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_groups = [t.strip() for t in task.split(",") if t.strip()]
+    if not task_groups:
+        raise ValueError("`task` must contain at least one Meta-World task or difficulty group.")
+
+    print(f"Creating Meta-World envs | task_groups={task_groups} | n_envs(per task)={n_envs}")
+
+    out: dict[str, dict[int, Any]] = defaultdict(dict)
+
+    for group in task_groups:
+        # if not in difficulty presets, treat it as a single custom task
+        tasks = DIFFICULTY_TO_TASKS.get(group, [group])
+
+        for tid, task_name in enumerate(tasks):
+            print(f"Building vec env | group={group} | task_id={tid} | task={task_name}")
+
+            # build n_envs factories
+            fns = [(lambda tn=task_name: MetaworldEnv(task=tn, **gym_kwargs)) for _ in range(n_envs)]
+
+            out[group][tid] = env_cls(fns)
+
+    # return a plain dict for consistency
+    return {group: dict(task_map) for group, task_map in out.items()}
@@ -0,0 +1,121 @@
+{
+  "TASK_DESCRIPTIONS": {
+    "assembly-v3": "Pick up a nut and place it onto a peg",
+    "basketball-v3": "Dunk the basketball into the basket",
+    "bin-picking-v3": "Grasp the puck from one bin and place it into another bin",
+    "box-close-v3": "Grasp the cover and close the box with it",
+    "button-press-topdown-v3": "Press a button from the top",
+    "button-press-topdown-wall-v3": "Bypass a wall and press a button from the top",
+    "button-press-v3": "Press a button",
+    "button-press-wall-v3": "Bypass a wall and press a button",
+    "coffee-button-v3": "Push a button on the coffee machine",
+    "coffee-pull-v3": "Pull a mug from a coffee machine",
+    "coffee-push-v3": "Push a mug under a coffee machine",
+    "dial-turn-v3": "Rotate a dial 180 degrees",
+    "disassemble-v3": "Pick a nut out of a peg",
+    "door-close-v3": "Close a door with a revolving joint",
+    "door-lock-v3": "Lock the door by rotating the lock clockwise",
+    "door-open-v3": "Open a door with a revolving joint",
+    "door-unlock-v3": "Unlock the door by rotating the lock counter-clockwise",
+    "hand-insert-v3": "Insert the gripper into a hole",
+    "drawer-close-v3": "Push and close a drawer",
+    "drawer-open-v3": "Open a drawer",
+    "faucet-open-v3": "Rotate the faucet counter-clockwise",
+    "faucet-close-v3": "Rotate the faucet clockwise",
+    "hammer-v3": "Hammer a screw on the wall",
+    "handle-press-side-v3": "Press a handle down sideways",
+    "handle-press-v3": "Press a handle down",
+    "handle-pull-side-v3": "Pull a handle up sideways",
+    "handle-pull-v3": "Pull a handle up",
+    "lever-pull-v3": "Pull a lever down 90 degrees",
+    "peg-insert-side-v3": "Insert a peg sideways",
+    "pick-place-wall-v3": "Pick a puck, bypass a wall and place the puck",
+    "pick-out-of-hole-v3": "Pick up a puck from a hole",
+    "reach-v3": "Reach a goal position",
+    "push-back-v3": "Push the puck to a goal",
+    "push-v3": "Push the puck to a goal",
+    "pick-place-v3": "Pick and place a puck to a goal",
+    "plate-slide-v3": "Slide a plate into a cabinet",
+    "plate-slide-side-v3": "Slide a plate into a cabinet sideways",
+    "plate-slide-back-v3": "Get a plate from the cabinet",
+    "plate-slide-back-side-v3": "Get a plate from the cabinet sideways",
+    "peg-unplug-side-v3": "Unplug a peg sideways",
+    "soccer-v3": "Kick a soccer into the goal",
+    "stick-push-v3": "Grasp a stick and push a box using the stick",
+    "stick-pull-v3": "Grasp a stick and pull a box with the stick",
+    "push-wall-v3": "Bypass a wall and push a puck to a goal",
+    "reach-wall-v3": "Bypass a wall and reach a goal",
+    "shelf-place-v3": "Pick and place a puck onto a shelf",
+    "sweep-into-v3": "Sweep a puck into a hole",
+    "sweep-v3": "Sweep a puck off the table",
+    "window-open-v3": "Push and open a window",
+    "window-close-v3": "Push and close a window"
+  },
+  "TASK_NAME_TO_ID": {
+    "assembly-v3": 0, "basketball-v3": 1, "bin-picking-v3": 2, "box-close-v3": 3,
+    "button-press-topdown-v3": 4, "button-press-topdown-wall-v3": 5, "button-press-v3": 6,
+    "button-press-wall-v3": 7, "coffee-button-v3": 8, "coffee-pull-v3": 9, "coffee-push-v3": 10,
+    "dial-turn-v3": 11, "disassemble-v3": 12, "door-close-v3": 13, "door-lock-v3": 14,
+    "door-open-v3": 15, "door-unlock-v3": 16, "drawer-close-v3": 17, "drawer-open-v3": 18,
+    "faucet-close-v3": 19, "faucet-open-v3": 20, "hammer-v3": 21, "hand-insert-v3": 22,
+    "handle-press-side-v3": 23, "handle-press-v3": 24, "handle-pull-side-v3": 25,
+    "handle-pull-v3": 26, "lever-pull-v3": 27, "peg-insert-side-v3": 28, "peg-unplug-side-v3": 29,
+    "pick-out-of-hole-v3": 30, "pick-place-v3": 31, "pick-place-wall-v3": 32,
+    "plate-slide-back-side-v3": 33, "plate-slide-back-v3": 34, "plate-slide-side-v3": 35,
+    "plate-slide-v3": 36, "push-back-v3": 37, "push-v3": 38, "push-wall-v3": 39, "reach-v3": 40,
+    "reach-wall-v3": 41, "shelf-place-v3": 42, "soccer-v3": 43, "stick-pull-v3": 44,
+    "stick-push-v3": 45, "sweep-into-v3": 46, "sweep-v3": 47, "window-open-v3": 48,
+    "window-close-v3": 49
+  },
+  "DIFFICULTY_TO_TASKS": {
+    "easy": [
+      "button-press-v3", "button-press-topdown-v3", "button-press-topdown-wall-v3",
+      "button-press-wall-v3", "coffee-button-v3", "dial-turn-v3", "door-close-v3",
+      "door-lock-v3", "door-open-v3", "door-unlock-v3", "drawer-close-v3", "drawer-open-v3",
+      "faucet-close-v3", "faucet-open-v3", "handle-press-v3", "handle-press-side-v3",
+      "handle-pull-v3", "handle-pull-side-v3", "lever-pull-v3", "plate-slide-v3",
+      "plate-slide-back-v3", "plate-slide-back-side-v3", "plate-slide-side-v3", "reach-v3",
+      "reach-wall-v3", "window-close-v3", "window-open-v3", "peg-unplug-side-v3"
+    ],
+    "medium": [
+      "basketball-v3", "bin-picking-v3", "box-close-v3", "coffee-pull-v3", "coffee-push-v3",
+      "hammer-v3", "peg-insert-side-v3", "push-wall-v3", "soccer-v3", "sweep-v3", "sweep-into-v3"
+    ],
+    "hard": [
+      "assembly-v3", "hand-insert-v3", "pick-out-of-hole-v3", "pick-place-v3", "push-v3", "push-back-v3"
+    ],
+    "very_hard": [
+      "shelf-place-v3", "disassemble-v3", "stick-pull-v3", "stick-push-v3", "pick-place-wall-v3"
+    ]
+  },
+  "TASK_POLICY_MAPPING": {
+    "assembly-v3": "SawyerAssemblyV3Policy", "basketball-v3": "SawyerBasketballV3Policy",
+    "bin-picking-v3": "SawyerBinPickingV3Policy", "box-close-v3": "SawyerBoxCloseV3Policy",
+    "button-press-topdown-v3": "SawyerButtonPressTopdownV3Policy",
+    "button-press-topdown-wall-v3": "SawyerButtonPressTopdownWallV3Policy",
+    "button-press-v3": "SawyerButtonPressV3Policy", "button-press-wall-v3": "SawyerButtonPressWallV3Policy",
+    "coffee-button-v3": "SawyerCoffeeButtonV3Policy", "coffee-pull-v3": "SawyerCoffeePullV3Policy",
+    "coffee-push-v3": "SawyerCoffeePushV3Policy", "dial-turn-v3": "SawyerDialTurnV3Policy",
+    "disassemble-v3": "SawyerDisassembleV3Policy", "door-close-v3": "SawyerDoorCloseV3Policy",
+    "door-lock-v3": "SawyerDoorLockV3Policy", "door-open-v3": "SawyerDoorOpenV3Policy",
+    "door-unlock-v3": "SawyerDoorUnlockV3Policy", "drawer-close-v3": "SawyerDrawerCloseV3Policy",
+    "drawer-open-v3": "SawyerDrawerOpenV3Policy", "faucet-close-v3": "SawyerFaucetCloseV3Policy",
+    "faucet-open-v3": "SawyerFaucetOpenV3Policy", "hammer-v3": "SawyerHammerV3Policy",
+    "hand-insert-v3": "SawyerHandInsertV3Policy", "handle-press-side-v3": "SawyerHandlePressSideV3Policy",
+    "handle-press-v3": "SawyerHandlePressV3Policy", "handle-pull-side-v3": "SawyerHandlePullSideV3Policy",
+    "handle-pull-v3": "SawyerHandlePullV3Policy", "lever-pull-v3": "SawyerLeverPullV3Policy",
+    "peg-insert-side-v3": "SawyerPegInsertionSideV3Policy", "peg-unplug-side-v3": "SawyerPegUnplugSideV3Policy",
+    "pick-out-of-hole-v3": "SawyerPickOutOfHoleV3Policy", "pick-place-v3": "SawyerPickPlaceV3Policy",
+    "pick-place-wall-v3": "SawyerPickPlaceWallV3Policy",
+    "plate-slide-back-side-v3": "SawyerPlateSlideBackSideV3Policy",
+    "plate-slide-back-v3": "SawyerPlateSlideBackV3Policy",
+    "plate-slide-side-v3": "SawyerPlateSlideSideV3Policy", "plate-slide-v3": "SawyerPlateSlideV3Policy",
+    "push-back-v3": "SawyerPushBackV3Policy", "push-v3": "SawyerPushV3Policy",
+    "push-wall-v3": "SawyerPushWallV3Policy", "reach-v3": "SawyerReachV3Policy",
+    "reach-wall-v3": "SawyerReachWallV3Policy", "shelf-place-v3": "SawyerShelfPlaceV3Policy",
+    "soccer-v3": "SawyerSoccerV3Policy", "stick-pull-v3": "SawyerStickPullV3Policy",
+    "stick-push-v3": "SawyerStickPushV3Policy", "sweep-into-v3": "SawyerSweepIntoV3Policy",
+    "sweep-v3": "SawyerSweepV3Policy", "window-open-v3": "SawyerWindowOpenV3Policy",
+    "window-close-v3": "SawyerWindowCloseV3Policy"
+  }
+}
@@ -48,25 +48,25 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten

        for imgkey, img in imgs.items():
            # TODO(aliberts, rcadene): use transforms.ToTensor()?
-            img = torch.from_numpy(img)
+            img_tensor = torch.from_numpy(img)

            # When preprocessing observations in a non-vectorized environment, we need to add a batch dimension.
            # This is the case for human-in-the-loop RL where there is only one environment.
-            if img.ndim == 3:
-                img = img.unsqueeze(0)
+            if img_tensor.ndim == 3:
+                img_tensor = img_tensor.unsqueeze(0)
            # sanity check that images are channel last
-            _, h, w, c = img.shape
-            assert c < h and c < w, f"expect channel last images, but instead got {img.shape=}"
+            _, h, w, c = img_tensor.shape
+            assert c < h and c < w, f"expect channel last images, but instead got {img_tensor.shape=}"

            # sanity check that images are uint8
-            assert img.dtype == torch.uint8, f"expect torch.uint8, but instead {img.dtype=}"
+            assert img_tensor.dtype == torch.uint8, f"expect torch.uint8, but instead {img_tensor.dtype=}"

            # convert to channel first of type float32 in range [0,1]
-            img = einops.rearrange(img, "b h w c -> b c h w").contiguous()
-            img = img.type(torch.float32)
-            img /= 255
+            img_tensor = einops.rearrange(img_tensor, "b h w c -> b c h w").contiguous()
+            img_tensor = img_tensor.type(torch.float32)
+            img_tensor /= 255

-            return_observations[imgkey] = img
+            return_observations[imgkey] = img_tensor

    if "environment_state" in observations:
        env_state = torch.from_numpy(observations["environment_state"]).float()
@@ -1 +1,17 @@
+#!/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 .motors_bus import Motor, MotorCalibration, MotorNormMode, MotorsBus
@@ -15,7 +15,7 @@
 from .act.configuration_act import ACTConfig as ACTConfig
 from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
 from .pi0.configuration_pi0 import PI0Config as PI0Config
-from .pi0.processor_pi0 import Pi0NewLineProcessor
+from .pi05.configuration_pi05 import PI05Config as PI05Config
 from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
 from .smolvla.processor_smolvla import SmolVLANewLineProcessor
 from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
@@ -25,6 +25,7 @@ __all__ = [
    "ACTConfig",
    "DiffusionConfig",
    "PI0Config",
+    "PI05Config",
    "SmolVLAConfig",
    "TDMPCConfig",
    "VQBeTConfig",
@@ -90,16 +90,16 @@ class DiffusionPolicy(PreTrainedPolicy):
            self._queues[OBS_ENV_STATE] = deque(maxlen=self.config.n_obs_steps)

    @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
        """Predict a chunk of actions given environment observations."""
        # stack n latest observations from the queue
        batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
-        actions = self.diffusion.generate_actions(batch)
+        actions = self.diffusion.generate_actions(batch, noise=noise)

        return actions

    @torch.no_grad()
-    def select_action(self, batch: dict[str, Tensor]) -> Tensor:
+    def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
        """Select a single action given environment observations.

        This method handles caching a history of observations and an action trajectory generated by the
@@ -131,7 +131,7 @@ class DiffusionPolicy(PreTrainedPolicy):
        self._queues = populate_queues(self._queues, batch)

        if len(self._queues[ACTION]) == 0:
-            actions = self.predict_action_chunk(batch)
+            actions = self.predict_action_chunk(batch, noise=noise)
            self._queues[ACTION].extend(actions.transpose(0, 1))

        action = self._queues[ACTION].popleft()
@@ -199,17 +199,25 @@ class DiffusionModel(nn.Module):

    # ========= inference  ============
    def conditional_sample(
-        self, batch_size: int, global_cond: Tensor | None = None, generator: torch.Generator | None = None
+        self,
+        batch_size: int,
+        global_cond: Tensor | None = None,
+        generator: torch.Generator | None = None,
+        noise: Tensor | None = None,
    ) -> Tensor:
        device = get_device_from_parameters(self)
        dtype = get_dtype_from_parameters(self)

        # Sample prior.
-        sample = torch.randn(
-            size=(batch_size, self.config.horizon, self.config.action_feature.shape[0]),
-            dtype=dtype,
-            device=device,
-            generator=generator,
+        sample = (
+            noise
+            if noise is not None
+            else torch.randn(
+                size=(batch_size, self.config.horizon, self.config.action_feature.shape[0]),
+                dtype=dtype,
+                device=device,
+                generator=generator,
+            )
        )

        self.noise_scheduler.set_timesteps(self.num_inference_steps)
@@ -264,7 +272,7 @@ class DiffusionModel(nn.Module):
        # Concatenate features then flatten to (B, global_cond_dim).
        return torch.cat(global_cond_feats, dim=-1).flatten(start_dim=1)

-    def generate_actions(self, batch: dict[str, Tensor]) -> Tensor:
+    def generate_actions(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
        """
        This function expects `batch` to have:
        {
@@ -282,7 +290,7 @@ class DiffusionModel(nn.Module):
        global_cond = self._prepare_global_conditioning(batch)  # (B, global_cond_dim)

        # run sampling
-        actions = self.conditional_sample(batch_size, global_cond=global_cond)
+        actions = self.conditional_sample(batch_size, global_cond=global_cond, noise=noise)

        # Extract `n_action_steps` steps worth of actions (from the current observation).
        start = n_obs_steps - 1
@@ -31,7 +31,7 @@ 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.pi0fast.configuration_pi0fast import PI0FASTConfig
+from lerobot.policies.pi05.configuration_pi05 import PI05Config
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.configuration_sac import SACConfig
 from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
@@ -57,7 +57,7 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:

    Args:
        name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
-              "vqbet", "pi0", "pi0fast", "sac", "reward_classifier", "smolvla".
+              "vqbet", "pi0", "pi05", "sac", "reward_classifier", "smolvla".

    Returns:
        The policy class corresponding to the given name.
@@ -85,10 +85,10 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
        from lerobot.policies.pi0.modeling_pi0 import PI0Policy

        return PI0Policy
-    elif name == "pi0fast":
-        from lerobot.policies.pi0fast.modeling_pi0fast import PI0FASTPolicy
+    elif name == "pi05":
+        from lerobot.policies.pi05.modeling_pi05 import PI05Policy

-        return PI0FASTPolicy
+        return PI05Policy
    elif name == "sac":
        from lerobot.policies.sac.modeling_sac import SACPolicy

@@ -114,7 +114,7 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:

    Args:
        policy_type: The type of the policy. Supported types include "tdmpc",
-                     "diffusion", "act", "vqbet", "pi0", "pi0fast", "sac", "smolvla",
+                     "diffusion", "act", "vqbet", "pi0", "pi05", "sac", "smolvla",
                     "reward_classifier".
        **kwargs: Keyword arguments to be passed to the configuration class constructor.

@@ -134,8 +134,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
        return VQBeTConfig(**kwargs)
    elif policy_type == "pi0":
        return PI0Config(**kwargs)
-    elif policy_type == "pi0fast":
-        return PI0FASTConfig(**kwargs)
+    elif policy_type == "pi05":
+        return PI05Config(**kwargs)
    elif policy_type == "sac":
        return SACConfig(**kwargs)
    elif policy_type == "smolvla":
@@ -261,10 +261,10 @@ def make_pre_post_processors(
            dataset_stats=kwargs.get("dataset_stats"),
        )

-    elif isinstance(policy_cfg, PI0FASTConfig):
-        from lerobot.policies.pi0fast.processor_pi0fast import make_pi0fast_pre_post_processors
+    elif isinstance(policy_cfg, PI05Config):
+        from lerobot.policies.pi05.processor_pi05 import make_pi05_pre_post_processors

-        processors = make_pi0fast_pre_post_processors(
+        processors = make_pi05_pre_post_processors(
            config=policy_cfg,
            dataset_stats=kwargs.get("dataset_stats"),
        )
@@ -0,0 +1,49 @@
+# π₀ (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 model for general robot control**.
+
+---
+
+## Model Overview
+
+| Feature              | π₀                                                     | π₀.₅                                      |
+| -------------------- | ------------------------------------------------------ | ----------------------------------------- |
+| 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 (Uses `state_proj` layer)                        | 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).
@@ -0,0 +1,21 @@
+#!/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_pi0 import PI0Config
+from .modeling_pi0 import PI0Policy
+from .processor_pi0 import make_pi0_pre_post_processors
+
+__all__ = ["PI0Config", "PI0Policy", "make_pi0_pre_post_processors"]
@@ -1,4 +1,6 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#!/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.
@@ -17,20 +19,40 @@ from dataclasses import dataclass, field
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.optim.optimizers import AdamWConfig
-from lerobot.optim.schedulers import (
-    CosineDecayWithWarmupSchedulerConfig,
-)
+from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
 from lerobot.utils.constants import OBS_IMAGES


@PreTrainedConfig.register_subclass("pi0")
@dataclass
 class PI0Config(PreTrainedConfig):
-    # Input / output structure.
-    n_obs_steps: int = 1
-    chunk_size: int = 50
-    n_action_steps: int = 50
+    paligemma_variant: str = "gemma_2b"
+    action_expert_variant: str = "gemma_300m"
+    dtype: str = "float32"  # Options: "bfloat16", "float32"

+    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
+    max_action_dim: int = 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
+    time_sampling_beta_beta: float = 1.0
+    time_sampling_scale: float = 0.999
+    time_sampling_offset: float = 0.001
+    min_period: float = 4e-3
+    max_period: float = 4.0
+
+    image_resolution: tuple[int, int] = (224, 224)  # see openpi `preprocessing_pytorch.py`
+
+    # Add empty images. Used to add empty cameras when no image features are present.
+    empty_cameras: int = 0
+
+    # Normalization
    normalization_mapping: dict[str, NormalizationMode] = field(
        default_factory=lambda: {
            "VISUAL": NormalizationMode.IDENTITY,
@@ -39,94 +61,75 @@ class PI0Config(PreTrainedConfig):
        }
    )

-    # Shorter state and action vectors will be padded
-    max_state_dim: int = 32
-    max_action_dim: int = 32
+    # 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)

-    # Image preprocessing
-    resize_imgs_with_padding: tuple[int, int] = (224, 224)
-
-    # Add empty images. Used by pi0_aloha_sim which adds the empty
-    # left and right wrist cameras in addition to the top camera.
-    empty_cameras: int = 0
-
-    # Converts the joint and gripper values from the standard Aloha space to
-    # the space used by the pi internal runtime which was used to train the base model.
-    adapt_to_pi_aloha: bool = False
-
-    # Converts joint dimensions to deltas with respect to the current state before passing to the model.
-    # Gripper dimensions will remain in absolute values.
-    use_delta_joint_actions_aloha: bool = False
-
-    # Tokenizer
-    tokenizer_max_length: int = 48
-
-    # Projector
-    proj_width: int = 1024
-
-    # Decoding
-    num_steps: int = 10
-
-    # Attention utils
-    use_cache: bool = True
-    attention_implementation: str = "eager"  # or fa2, flex
-
-    # Finetuning settings
-    freeze_vision_encoder: bool = True
-    train_expert_only: bool = False
-    train_state_proj: bool = True
-
-    # Training presets
-    optimizer_lr: float = 2.5e-5
+    # Optimizer settings: see openpi `AdamW``
+    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 = 1e-10
+    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

-    # TODO: Add EMA
+    tokenizer_max_length: int = 48  # see openpi `__post_init__`

    def __post_init__(self):
        super().__post_init__()

-        # TODO(Steven): Validate device and amp? in all policy configs?
-        """Input validation (not exhaustive)."""
+        # Validate configuration
        if self.n_action_steps > self.chunk_size:
            raise ValueError(
-                f"The chunk size is the upper bound for the number of action steps per model invocation. Got "
-                f"{self.n_action_steps} for `n_action_steps` and {self.chunk_size} for `chunk_size`."
-            )
-        if self.n_obs_steps != 1:
-            raise ValueError(
-                f"Multiple observation steps not handled yet. Got `nobs_steps={self.n_obs_steps}`"
+                f"n_action_steps ({self.n_action_steps}) cannot be greater than chunk_size ({self.chunk_size})"
            )

-        if self.use_delta_joint_actions_aloha:
-            raise NotImplementedError(
-                "`use_delta_joint_actions_aloha` is used by pi0 for aloha real models. It is not ported yet in LeRobot."
-            )
+        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:
-        # TODO: implement value error
-        # if not self.image_features and not self.env_state_feature:
-        #     raise ValueError("You must provide at least one image or the environment state among the inputs.")
-
+        """Validate and set up input/output features."""
        for i in range(self.empty_cameras):
            key = f"{OBS_IMAGES}.empty_camera_{i}"
            empty_camera = PolicyFeature(
                type=FeatureType.VISUAL,
-                shape=(3, 480, 640),
+                shape=(3, *self.image_resolution),  # Use configured image resolution
            )
            self.input_features[key] = empty_camera

+        if "observation.state" not in self.input_features:
+            state_feature = PolicyFeature(
+                type=FeatureType.STATE,
+                shape=(self.max_state_dim,),  # Padded to max_state_dim
+            )
+            self.input_features["observation.state"] = state_feature
+
+        if "action" not in self.output_features:
+            action_feature = PolicyFeature(
+                type=FeatureType.ACTION,
+                shape=(self.max_action_dim,),  # Padded to max_action_dim
+            )
+            self.output_features["action"] = action_feature
+
    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):
@@ -1,82 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import torch
-
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.policies.factory import make_policy
-
-torch.backends.cudnn.benchmark = True
-
-
-def main():
-    device = "cuda"
-    dataset_repo_id = "danaaubakirova/koch_test"
-    # model_name = "pi0_base"
-    # ckpt_torch_dir = Path.home() / f".cache/openpi/openpi-assets/checkpoints/{model_name}_pytorch"
-    ckpt_torch_dir = "lerobot/pi0"
-
-    dataset = LeRobotDataset(dataset_repo_id, episodes=[0])
-
-    dataloader = torch.utils.data.DataLoader(
-        dataset,
-        num_workers=0,
-        batch_size=1,
-    )
-
-    batch = next(iter(dataloader))
-
-    # To device
-    for k in batch:
-        if isinstance(batch[k], torch.Tensor):
-            batch[k] = batch[k].to(device=device, dtype=torch.float32)
-
-    cfg = PreTrainedConfig.from_pretrained(ckpt_torch_dir)
-    cfg.pretrained_path = ckpt_torch_dir
-    policy = make_policy(cfg, ds_meta=dataset.meta)
-
-    # policy = torch.compile(policy, mode="reduce-overhead")
-
-    warmup_iters = 10
-    benchmark_iters = 30
-
-    # Warmup
-    for _ in range(warmup_iters):
-        torch.cuda.synchronize()
-        policy.select_action(batch)
-        policy.reset()
-        torch.cuda.synchronize()
-
-    # Benchmark
-    start_event = torch.cuda.Event(enable_timing=True)
-    end_event = torch.cuda.Event(enable_timing=True)
-
-    start_event.record()
-    for _ in range(benchmark_iters):
-        policy.select_action(batch)
-        policy.reset()
-    end_event.record()
-
-    # Synchronize and measure time
-    torch.cuda.synchronize()
-    elapsed_time_ms = start_event.elapsed_time(end_event)
-
-    avg_time_per_iter = elapsed_time_ms / benchmark_iters
-    print(f"Average execution time per iteration: {avg_time_per_iter:.3f} ms")
-
-
-if __name__ == "__main__":
-    with torch.inference_mode():
-        main()
@@ -1,132 +0,0 @@
-# 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 json
-import pickle
-from pathlib import Path
-
-import torch
-
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
-from lerobot.policies.factory import make_policy
-from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
-
-
-def display(tensor: torch.Tensor):
-    if tensor.dtype == torch.bool:
-        tensor = tensor.float()
-    print(f"Shape: {tensor.shape}")
-    print(f"Mean: {tensor.mean().item()}")
-    print(f"Std: {tensor.std().item()}")
-    print(f"Min: {tensor.min().item()}")
-    print(f"Max: {tensor.max().item()}")
-
-
-def main():
-    num_motors = 14
-    device = "cuda"
-    # model_name = "pi0_aloha_towel"
-    model_name = "pi0_aloha_sim"
-
-    if model_name == "pi0_aloha_towel":
-        dataset_repo_id = "lerobot/aloha_static_towel"
-    else:
-        dataset_repo_id = "lerobot/aloha_sim_transfer_cube_human"
-
-    ckpt_torch_dir = Path.home() / f".cache/openpi/openpi-assets/checkpoints/{model_name}_pytorch"
-    ckpt_jax_dir = Path.home() / f".cache/openpi/openpi-assets/checkpoints/{model_name}"
-    save_dir = Path(f"../openpi/data/{model_name}/save")
-
-    with open(save_dir / "example.pkl", "rb") as f:
-        example = pickle.load(f)
-    with open(save_dir / "outputs.pkl", "rb") as f:
-        outputs = pickle.load(f)
-    with open(save_dir / "noise.pkl", "rb") as f:
-        noise = pickle.load(f)
-
-    with open(ckpt_jax_dir / "assets/norm_stats.json") as f:
-        norm_stats = json.load(f)
-
-    # Override stats
-    dataset_meta = LeRobotDatasetMetadata(dataset_repo_id)
-    dataset_meta.stats[OBS_STATE]["mean"] = torch.tensor(
-        norm_stats["norm_stats"]["state"]["mean"][:num_motors], dtype=torch.float32
-    )
-    dataset_meta.stats[OBS_STATE]["std"] = torch.tensor(
-        norm_stats["norm_stats"]["state"]["std"][:num_motors], dtype=torch.float32
-    )
-
-    # Create LeRobot batch from Jax
-    batch = {}
-    for cam_key, uint_chw_array in example["images"].items():
-        batch[f"{OBS_IMAGES}.{cam_key}"] = torch.from_numpy(uint_chw_array) / 255.0
-    batch[OBS_STATE] = torch.from_numpy(example["state"])
-    batch[ACTION] = torch.from_numpy(outputs["actions"])
-    batch["task"] = example["prompt"]
-
-    if model_name == "pi0_aloha_towel":
-        del batch[f"{OBS_IMAGES}.cam_low"]
-    elif model_name == "pi0_aloha_sim":
-        batch[f"{OBS_IMAGES}.top"] = batch[f"{OBS_IMAGES}.cam_high"]
-        del batch[f"{OBS_IMAGES}.cam_high"]
-
-    # Batchify
-    for key in batch:
-        if isinstance(batch[key], torch.Tensor):
-            batch[key] = batch[key].unsqueeze(0)
-        elif isinstance(batch[key], str):
-            batch[key] = [batch[key]]
-        else:
-            raise ValueError(f"{key}, {batch[key]}")
-
-    # To device
-    for k in batch:
-        if isinstance(batch[k], torch.Tensor):
-            batch[k] = batch[k].to(device=device, dtype=torch.float32)
-
-    noise = torch.from_numpy(noise).to(device=device, dtype=torch.float32)
-
-    from lerobot import policies  # noqa
-
-    cfg = PreTrainedConfig.from_pretrained(ckpt_torch_dir)
-    cfg.pretrained_path = ckpt_torch_dir
-    policy = make_policy(cfg, dataset_meta)
-
-    # loss_dict = policy.forward(batch, noise=noise, time=time_beta)
-    # loss_dict["loss"].backward()
-    # print("losses")
-    # display(loss_dict["losses_after_forward"])
-    # print("pi_losses")
-    # display(pi_losses)
-
-    actions = []
-    for _ in range(50):
-        action = policy.select_action(batch, noise=noise)
-        actions.append(action)
-
-    actions = torch.stack(actions, dim=1)
-    pi_actions = batch[ACTION]
-    print("actions")
-    display(actions)
-    print()
-    print("pi_actions")
-    display(pi_actions)
-    print("atol=3e-2", torch.allclose(actions, pi_actions, atol=3e-2))
-    print("atol=2e-2", torch.allclose(actions, pi_actions, atol=2e-2))
-    print("atol=1e-2", torch.allclose(actions, pi_actions, atol=1e-2))
-
-
-if __name__ == "__main__":
-    main()
@@ -1,84 +0,0 @@
-# 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 transformers import GemmaConfig, PaliGemmaConfig
-
-
-def get_paligemma_config(precision: str):
-    config = {
-        "image_token_index": None,
-        "pad_token_id": 0,
-        "bos_token_id": 2,
-        "eos_token_id": 1,
-    }
-
-    # image_sizes = {"2b-test": 224, "3b-224px": 224, "3b-448px": 448, "3b-896px": 896}
-
-    image_size = 224  # image_sizes[variant]
-    patch_size = 14
-    num_image_tokens = (image_size**2) // (patch_size**2)
-
-    config["image_token_index"] = 257152
-    text_config = {
-        "vocab_size": 257152,
-        "num_hidden_layers": 18,
-        "num_key_value_heads": 1,
-        "head_dim": 256,
-        "torch_dtype": precision,
-        "hidden_size": 2048,
-        "hidden_activation": "gelu_pytorch_tanh",
-        "num_attention_heads": 8,
-        "intermediate_size": 16384,
-        "is_encoder_decoder": False,
-    }
-    vision_config = {
-        "torch_dtype": precision,
-        "image_size": image_size,
-        "patch_size": patch_size,
-        "num_image_tokens": num_image_tokens,
-        "hidden_size": 1152,
-        "intermediate_size": 4304,
-        "num_hidden_layers": 27,
-        "num_attention_heads": 16,
-        "projector_hidden_act": "gelu_fast",
-        "vision_use_head": False,
-    }
-    final_config = PaliGemmaConfig(text_config=text_config, vision_config=vision_config, **config)
-    return final_config
-
-
-def get_gemma_config(precision: str):
-    config = {
-        "image_token_index": None,
-        "pad_token_id": 0,
-        "bos_token_id": 2,
-        "eos_token_id": 1,
-    }
-
-    config["image_token_index"] = 257152
-    text_config = {
-        "vocab_size": 257152,
-        "num_hidden_layers": 18,
-        "num_key_value_heads": 1,
-        "head_dim": 256,
-        "torch_dtype": precision,
-        "hidden_size": 1024,
-        "hidden_activation": "gelu_pytorch_tanh",
-        "num_attention_heads": 8,
-        "intermediate_size": 4096,
-        "is_encoder_decoder": False,
-    }
-    final_config = GemmaConfig()
-    final_config.update(text_config)
-    return final_config
@@ -1,437 +0,0 @@
-# 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.
-
-"""
-Convert pi0 parameters from Jax to Pytorch
-
-Follow [README of openpi](https://github.com/Physical-Intelligence/openpi) to create a new environment
-and install the required libraries.
-
-```bash
-cd ~/code/openpi
-source .venv/bin/activate
-```
-
-Example downloading parameters:
-```bash
-python
->>> import openpi.shared.download as download
->>> path='s3://openpi-assets/checkpoints/pi0_base/params'
->>> download.maybe_download(path)
-```
-
-Converting pi0_base:
-```python
-python -m lerobot.policies.pi0.conversion_scripts.convert_pi0_to_hf_lerobot \
-    --checkpoint_dir /home/remi_cadene/.cache/openpi/openpi-assets/checkpoints/pi0_base/params \
-    --output_path /home/remi_cadene/.cache/openpi/openpi-assets/checkpoints/pi0_base_pytorch
-```
-
-```python
-python -m lerobot.policies.pi0.conversion_scripts.convert_pi0_to_hf_lerobot \
-    --checkpoint_dir /home/remi_cadene/.cache/openpi/openpi-assets/checkpoints/pi0_aloha_sim/params \
-    --output_path /home/remi_cadene/.cache/openpi/openpi-assets/checkpoints/pi0_aloha_sim_pytorch
-```
-"""
-
-import argparse
-import pathlib
-
-import jax
-import numpy as np
-import orbax.checkpoint as ocp
-import torch
-from jax.sharding import SingleDeviceSharding
-
-from lerobot.policies.pi0.configuration_pi0 import PI0Config
-from lerobot.policies.pi0.conversion_scripts.conversion_utils import (
-    get_gemma_config,
-    get_paligemma_config,
-)
-from lerobot.policies.pi0.modeling_pi0 import PI0Policy
-
-PRECISIONS = {"bfloat16": torch.bfloat16, "float32": torch.float32, "float16": torch.float16}
-
-
-def slice_paligemma_state_dict(state_dict, config):
-    suffix = "/value" if "img/embedding/kernel/value" in state_dict else ""
-
-    # fmt: off
-    # patch embeddings
-    state_dict["paligemma.vision_tower.vision_model.embeddings.patch_embedding.weight"] = state_dict.pop(f"img/embedding/kernel{suffix}").transpose(
-        3, 2, 0, 1
-    )
-    state_dict["paligemma.vision_tower.vision_model.embeddings.patch_embedding.bias"] = state_dict.pop(f"img/embedding/bias{suffix}")
-    # positional embeddings
-    state_dict["paligemma.vision_tower.vision_model.embeddings.position_embedding.weight"] = state_dict.pop(f"img/pos_embedding{suffix}").reshape(
-        -1, config.vision_config.hidden_size
-    )
-
-    # extract vision layers to be sliced at index 0. There are 27 layers in the base model.
-    encoderblock_layernorm0_scale = state_dict.pop(f"img/Transformer/encoderblock/LayerNorm_0/scale{suffix}")
-    encoderblock_layernorm0_bias = state_dict.pop(f"img/Transformer/encoderblock/LayerNorm_0/bias{suffix}")
-    encoderblock_layernorm1_scale = state_dict.pop(f"img/Transformer/encoderblock/LayerNorm_1/scale{suffix}")
-    encoderblock_layernorm1_bias = state_dict.pop(f"img/Transformer/encoderblock/LayerNorm_1/bias{suffix}")
-
-    encoderblock_mlp_dense0_kernel= state_dict.pop(f"img/Transformer/encoderblock/MlpBlock_0/Dense_0/kernel{suffix}")
-    encoderblock_mlp_dense0_bias= state_dict.pop(f"img/Transformer/encoderblock/MlpBlock_0/Dense_0/bias{suffix}")
-    encoderblock_mlp_dense1_kernel= state_dict.pop(f"img/Transformer/encoderblock/MlpBlock_0/Dense_1/kernel{suffix}")
-    encoderblock_mlp_dense1_bias= state_dict.pop(f"img/Transformer/encoderblock/MlpBlock_0/Dense_1/bias{suffix}")
-
-    encoderblock_attention_0_key_kernel = state_dict.pop(f"img/Transformer/encoderblock/MultiHeadDotProductAttention_0/key/kernel{suffix}")
-    encoderblock_attention_0_key_bias = state_dict.pop(f"img/Transformer/encoderblock/MultiHeadDotProductAttention_0/key/bias{suffix}")
-    encoderblock_attention_0_value_kernel = state_dict.pop(f"img/Transformer/encoderblock/MultiHeadDotProductAttention_0/value/kernel{suffix}")
-    encoderblock_attention_0_value_bias = state_dict.pop(f"img/Transformer/encoderblock/MultiHeadDotProductAttention_0/value/bias{suffix}")
-    encoderblock_attention_0_query_kernel = state_dict.pop(f"img/Transformer/encoderblock/MultiHeadDotProductAttention_0/query/kernel{suffix}")
-    encoderblock_attention_0_query_bias = state_dict.pop(f"img/Transformer/encoderblock/MultiHeadDotProductAttention_0/query/bias{suffix}")
-    encoderblock_attention_0_out_kernel = state_dict.pop(f"img/Transformer/encoderblock/MultiHeadDotProductAttention_0/out/kernel{suffix}")
-    encoderblock_attention_0_out_bias = state_dict.pop(f"img/Transformer/encoderblock/MultiHeadDotProductAttention_0/out/bias{suffix}")
-
-    for i in range(config.vision_config.num_hidden_layers):
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.layer_norm1.weight"] = encoderblock_layernorm0_scale[i].transpose()
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.layer_norm1.bias"] = encoderblock_layernorm0_bias[i]
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.layer_norm2.weight"] = encoderblock_layernorm1_scale[i].transpose()
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.layer_norm2.bias"] = encoderblock_layernorm1_bias[i]
-
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.mlp.fc1.weight"] = encoderblock_mlp_dense0_kernel[i].transpose()
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.mlp.fc1.bias"] = encoderblock_mlp_dense0_bias[i]
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.mlp.fc2.weight"] = encoderblock_mlp_dense1_kernel[i].transpose()
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.mlp.fc2.bias"] = encoderblock_mlp_dense1_bias[i]
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.self_attn.k_proj.weight"] = encoderblock_attention_0_key_kernel[i].reshape(-1, config.vision_config.hidden_size).transpose()
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.self_attn.k_proj.bias"] = encoderblock_attention_0_key_bias[i].reshape(-1, config.vision_config.hidden_size).reshape(-1)
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.self_attn.v_proj.weight"] = encoderblock_attention_0_value_kernel[i].reshape(-1, config.vision_config.hidden_size).transpose()
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.self_attn.v_proj.bias"] = encoderblock_attention_0_value_bias[i].reshape(-1, config.vision_config.hidden_size).reshape(-1)
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.self_attn.q_proj.weight"] = encoderblock_attention_0_query_kernel[i].reshape(-1, config.vision_config.hidden_size).transpose()
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.self_attn.q_proj.bias"] = encoderblock_attention_0_query_bias[i].reshape(-1, config.vision_config.hidden_size).reshape(-1)
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.self_attn.out_proj.weight"] = encoderblock_attention_0_out_kernel[i].reshape(-1, config.vision_config.hidden_size).transpose()
-        state_dict[f"paligemma.vision_tower.vision_model.encoder.layers.{i}.self_attn.out_proj.bias"] = encoderblock_attention_0_out_bias[i].reshape(-1, config.vision_config.hidden_size).reshape(-1)
-
-    state_dict["paligemma.vision_tower.vision_model.post_layernorm.weight"] = state_dict.pop(f"img/Transformer/encoder_norm/scale{suffix}").transpose()
-    state_dict["paligemma.vision_tower.vision_model.post_layernorm.bias"] = state_dict.pop(f"img/Transformer/encoder_norm/bias{suffix}")
-
-    # multimodal projector
-
-    state_dict['paligemma.multi_modal_projector.linear.weight'] = state_dict.pop(f"img/head/kernel{suffix}").transpose()
-    state_dict['paligemma.multi_modal_projector.linear.bias'] = state_dict.pop(f"img/head/bias{suffix}")
-
-    # text decoder (gemma)
-    embedding_vector = state_dict.pop(f"llm/embedder/input_embedding{suffix}")
-    state_dict["paligemma.language_model.model.embed_tokens.weight"] = embedding_vector
-
-    # pop the einsum attention + mlp representations. There are 18 layers in gemma-2b.
-
-    llm_attention_attn_vec_einsum = state_dict.pop(f"llm/layers/attn/attn_vec_einsum/w{suffix}")
-    llm_attention_kv_einsum = state_dict.pop(f"llm/layers/attn/kv_einsum/w{suffix}")
-    llm_attention_q_einsum = state_dict.pop(f"llm/layers/attn/q_einsum/w{suffix}")
-
-    llm_mlp_gating_einsum = state_dict.pop(f"llm/layers/mlp/gating_einsum{suffix}")
-    llm_mlp_linear = state_dict.pop(f"llm/layers/mlp/linear{suffix}")
-    # TODO verify correctness of layer norm loading
-
-    llm_input_layernorm = state_dict.pop(f"llm/layers/pre_attention_norm/scale{suffix}")
-    llm_post_attention_layernorm = state_dict.pop(f"llm/layers/pre_ffw_norm/scale{suffix}")
-
-    for i in range(config.text_config.num_hidden_layers):
-        # llm_attention_q_einsum[i].shape = (8, 2048, 256)
-        q_proj_weight_reshaped = llm_attention_q_einsum[i].transpose(0, 2, 1).reshape(config.text_config.num_attention_heads * config.text_config.head_dim, config.text_config.hidden_size)
-
-        state_dict[f"paligemma.language_model.model.layers.{i}.self_attn.q_proj.weight"] = q_proj_weight_reshaped
-
-        # llm_attention_kv_einsum[i, 0, 0].shape = (2048, 256)
-        k_proj_weight_reshaped = llm_attention_kv_einsum[i, 0, 0].transpose()
-        state_dict[f"paligemma.language_model.model.layers.{i}.self_attn.k_proj.weight"] = k_proj_weight_reshaped
-        # llm_attention_kv_einsum[i, 1, 0].shape = (2048, 256)
-        v_proj_weight_reshaped = llm_attention_kv_einsum[i, 1, 0].transpose()
-        state_dict[f"paligemma.language_model.model.layers.{i}.self_attn.v_proj.weight"] = v_proj_weight_reshaped
-
-        # output projection.
-
-        # llm_attention_attn_vec_einsum[i].shape = (8, 256, 2048)
-        o_proj_weight_reshaped = llm_attention_attn_vec_einsum[i].transpose(2, 0, 1).reshape(config.text_config.num_attention_heads * config.text_config.head_dim, config.text_config.hidden_size)
-
-        state_dict[f"paligemma.language_model.model.layers.{i}.self_attn.o_proj.weight"] = o_proj_weight_reshaped
-        # mlp layers
-        gate_proj_weight = llm_mlp_gating_einsum[i, 0]
-        state_dict[f"paligemma.language_model.model.layers.{i}.mlp.gate_proj.weight"] = gate_proj_weight.transpose()
-        up_proj_weight = llm_mlp_gating_einsum[i, 1]
-        state_dict[f"paligemma.language_model.model.layers.{i}.mlp.up_proj.weight"] = up_proj_weight.transpose()
-        state_dict[f"paligemma.language_model.model.layers.{i}.mlp.down_proj.weight"] = llm_mlp_linear[i].transpose()
-        state_dict[f"paligemma.language_model.model.layers.{i}.input_layernorm.weight"] = llm_input_layernorm[i]
-        state_dict[f"paligemma.language_model.model.layers.{i}.post_attention_layernorm.weight"] = llm_post_attention_layernorm[i]
-
-    state_dict["paligemma.language_model.model.norm.weight"] = state_dict.pop(f"llm/final_norm/scale{suffix}")
-    state_dict["paligemma.language_model.lm_head.weight"] = embedding_vector # weights are tied.
-
-    # fmt: on
-    expert_dict = {}
-    final_state_dict = {}
-    for key, value in state_dict.items():
-        if key not in [
-            f"llm/final_norm_1/scale{suffix}",
-            f"llm/layers/attn/attn_vec_einsum_1/w{suffix}",
-            f"llm/layers/attn/kv_einsum_1/w{suffix}",
-            f"llm/layers/attn/q_einsum_1/w{suffix}",
-            f"llm/layers/mlp_1/gating_einsum{suffix}",
-            f"llm/layers/mlp_1/linear{suffix}",
-            f"llm/layers/pre_attention_norm_1/scale{suffix}",
-            f"llm/layers/pre_ffw_norm_1/scale{suffix}",
-        ]:
-            final_state_dict[key] = torch.from_numpy(value)
-        else:
-            expert_dict[key] = value
-
-    return final_state_dict, expert_dict
-
-
-def slice_gemma_state_dict(state_dict, config, num_expert=1):
-    # fmt: off
-    # text decoder (gemma)
-    # no embedding vector, the expert just has the decoder layers
-
-    embedding_vector = torch.zeros([config.vocab_size, config.hidden_size])
-    state_dict["gemma_expert.model.embed_tokens.weight"] = embedding_vector
-
-    # pop the einsum attention + mlp representations. There are 18 layers in gemma-2b.
-
-    suffix = "/value" if f"llm/layers/attn/attn_vec_einsum_{num_expert}/w/value" in state_dict else ""
-
-    llm_attention_attn_vec_einsum = state_dict.pop(f"llm/layers/attn/attn_vec_einsum_{num_expert}/w{suffix}")
-    llm_attention_kv_einsum = state_dict.pop(f"llm/layers/attn/kv_einsum_{num_expert}/w{suffix}")
-    llm_attention_q_einsum = state_dict.pop(f"llm/layers/attn/q_einsum_{num_expert}/w{suffix}")
-
-    llm_mlp_gating_einsum = state_dict.pop(f"llm/layers/mlp_{num_expert}/gating_einsum{suffix}")
-    llm_mlp_linear = state_dict.pop(f"llm/layers/mlp_{num_expert}/linear{suffix}")
-    # TODO verify correctness of layer norm loading
-
-    llm_input_layernorm = state_dict.pop(f"llm/layers/pre_attention_norm_{num_expert}/scale{suffix}")
-    llm_post_attention_layernorm = state_dict.pop(f"llm/layers/pre_ffw_norm_{num_expert}/scale{suffix}")
-
-    for i in range(config.num_hidden_layers):
-        q_proj_weight_reshaped = llm_attention_q_einsum[i].transpose(0, 2, 1).reshape(config.num_attention_heads * config.head_dim, config.hidden_size)
-
-        state_dict[f"gemma_expert.model.layers.{i}.self_attn.q_proj.weight"] = q_proj_weight_reshaped
-
-        k_proj_weight_reshaped = llm_attention_kv_einsum[i, 0, 0].transpose()
-        state_dict[f"gemma_expert.model.layers.{i}.self_attn.k_proj.weight"] = k_proj_weight_reshaped
-        v_proj_weight_reshaped = llm_attention_kv_einsum[i, 1, 0].transpose()
-        state_dict[f"gemma_expert.model.layers.{i}.self_attn.v_proj.weight"] = v_proj_weight_reshaped
-
-        # output projection.
-
-        # llm_attention_attn_vec_einsum[i].shape = (8, 256, 1024)
-        o_proj_weight_reshaped = llm_attention_attn_vec_einsum[i].reshape(config.num_attention_heads * config.head_dim, config.hidden_size).transpose(1,0)# .transpose(2, 0, 1).reshape(config.num_attention_heads * config.head_dim, config.hidden_size).transpose(1, 0)
-
-        state_dict[f"gemma_expert.model.layers.{i}.self_attn.o_proj.weight"] = o_proj_weight_reshaped
-        # mlp layers
-        gate_proj_weight = llm_mlp_gating_einsum[i, 0]
-        state_dict[f"gemma_expert.model.layers.{i}.mlp.gate_proj.weight"] = gate_proj_weight.transpose()
-        up_proj_weight = llm_mlp_gating_einsum[i, 1]
-        state_dict[f"gemma_expert.model.layers.{i}.mlp.up_proj.weight"] = up_proj_weight.transpose()
-        state_dict[f"gemma_expert.model.layers.{i}.mlp.down_proj.weight"] = llm_mlp_linear[i].transpose()
-        state_dict[f"gemma_expert.model.layers.{i}.input_layernorm.weight"] = llm_input_layernorm[i]
-        state_dict[f"gemma_expert.model.layers.{i}.post_attention_layernorm.weight"] = llm_post_attention_layernorm[i]
-
-    state_dict["gemma_expert.model.norm.weight"] = state_dict.pop(f"llm/final_norm_{num_expert}/scale{suffix}")
-    state_dict["gemma_expert.lm_head.weight"] = embedding_vector # weights are tied. (and zeros here)
-
-    # fmt: on
-    final_state_dict = {}
-    for key, value in state_dict.items():
-        if not isinstance(value, torch.Tensor):
-            final_state_dict[key] = torch.from_numpy(value)
-        else:
-            final_state_dict[key] = value
-    return final_state_dict
-
-
-def flatten_for_memory(tree, parent_key=""):
-    out = {}
-    for k, v in tree.items():
-        new_key = f"{parent_key}/{k}" if parent_key else k
-        if isinstance(v, dict):
-            out.update(flatten_for_memory(v, new_key))
-        else:
-            out[new_key] = np.array(v)  # Ensure conversion to np.array for consistency
-    return out
-
-
-def flatten_for_npz(tree, parent_key=""):
-    out = {}
-    for k, v in tree.items():
-        new_key = f"{parent_key}/{k}" if parent_key else k
-        if isinstance(v, dict):
-            out.update(flatten_for_npz(v, new_key))
-        else:
-            # bf16/f32 here?
-            out[new_key] = np.array(v)
-    return out
-
-
-def slice_initial_orbax_checkpoint(checkpoint_dir: str):
-    params_path = pathlib.Path(checkpoint_dir).resolve()
-    checkpointer = ocp.PyTreeCheckpointer()
-
-    metadata = checkpointer.metadata(params_path)
-    print("Metadata keys:", list(metadata.keys()))
-
-    params_name = "params"
-
-    item = {params_name: metadata[params_name]}
-    device = jax.local_devices()[0]  # Use the first local device
-    sharding = SingleDeviceSharding(device)
-    restored = checkpointer.restore(
-        params_path,
-        ocp.args.PyTreeRestore(
-            item=item,
-            restore_args=jax.tree_util.tree_map(
-                lambda _: ocp.ArrayRestoreArgs(
-                    restore_type=jax.Array,  # or np.ndarray, but bf16 is annoying about it
-                    sharding=sharding,
-                ),
-                item,
-            ),
-            transforms={},
-        ),
-    )
-    params = restored[params_name]
-
-    # get params for PaliGemma
-    pali_params = params["PaliGemma"]
-    del params["PaliGemma"]
-    pali_params_flat = flatten_for_npz(pali_params)
-    return {"paligemma_params": pali_params_flat, "projection_params": params}
-
-
-def update_keys_with_prefix(d: dict, prefix: str) -> dict:
-    """Update dictionary keys by adding a prefix."""
-    return {f"{prefix}{key}": value for key, value in d.items()}
-
-
-def convert_pi0_checkpoint(checkpoint_dir: str, precision: str, tokenizer_id: str, output_path: str):
-    # Break down orbax ckpts - they are in OCDBT
-    initial_params = slice_initial_orbax_checkpoint(checkpoint_dir=checkpoint_dir)
-    # process projection params
-    keys = [
-        "state_proj",
-        "action_in_proj",
-        "action_out_proj",
-        "action_time_mlp_in",
-        "action_time_mlp_out",
-    ]
-
-    projection_params = {}
-    for key in keys:
-        kernel_params = initial_params["projection_params"][key]["kernel"]
-        bias_params = initial_params["projection_params"][key]["bias"]
-        if isinstance(kernel_params, dict):
-            weight = kernel_params["value"]
-            bias = bias_params["value"]
-        else:
-            weight = kernel_params
-            bias = bias_params
-        projection_params[f"{key}.weight"] = torch.from_numpy(np.array(weight)).T
-        projection_params[f"{key}.bias"] = torch.from_numpy(np.array(bias))
-
-    # Process PaliGemma weights
-    paligemma_config = get_paligemma_config(precision)
-    paligemma_params, gemma_raw_dictionary = slice_paligemma_state_dict(
-        initial_params["paligemma_params"], paligemma_config
-    )
-
-    # Process Gemma  weights (at this stage they are unused)
-    gemma_config = get_gemma_config(precision)
-    gemma_params = slice_gemma_state_dict(gemma_raw_dictionary, config=gemma_config)
-
-    # Instantiate model from configs
-
-    if "pi0_aloha_sim" in checkpoint_dir:
-        pi0_config = PI0Config(
-            empty_cameras=2,
-            adapt_to_pi_aloha=True,
-            use_delta_joint_actions_aloha=False,
-        )
-    elif "pi0_aloha_towel" in checkpoint_dir:
-        pi0_config = PI0Config(
-            adapt_to_pi_aloha=True,
-            use_delta_joint_actions_aloha=True,
-        )
-    elif "pi0_base" in checkpoint_dir:
-        pi0_config = PI0Config(
-            empty_cameras=0,
-            adapt_to_pi_aloha=False,
-            use_delta_joint_actions_aloha=False,
-        )
-    else:
-        raise ValueError()
-
-    # gemma_config=gemma_config, paligemma_config=paligemma_config)
-    pi0_model = PI0Policy(pi0_config)
-
-    paligemma_params = update_keys_with_prefix(paligemma_params, "model.paligemma_with_expert.")
-    gemma_params = update_keys_with_prefix(gemma_params, "model.paligemma_with_expert.")
-    projection_params = update_keys_with_prefix(projection_params, "model.")
-
-    # load state dict
-    torch_dtype = PRECISIONS[precision]
-    pi0_model.load_state_dict({**paligemma_params, **gemma_params, **projection_params})
-    pi0_model = pi0_model.to(torch_dtype)
-    # pi0_tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
-
-    pi0_model.save_pretrained(output_path, safe_serialization=True)
-    # pi0_tokenizer.save_pretrained(output_path, dtype=torch_dtype)
-
-    # assert that model loads properly
-    del pi0_model
-    PI0Policy.from_pretrained(output_path)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--checkpoint_dir",
-        default="/raid/pablo/.cache/openpi/openpi-assets/checkpoints/pi0_aloha_sim/params",
-        type=str,
-        help="Path to the ocdbt checkpoint",
-    )
-
-    parser.add_argument(
-        "--precision",
-        choices=["float32", "bfloat16", "float16"],
-        default="float32",
-        type=str,
-        help="Precision identifier for model conversion - should match the base checkpoint precision.",
-    )
-    # tokenizer is identical to paligemma, it appears
-
-    parser.add_argument(
-        "--tokenizer_hub_id",
-        default="google/paligemma-3b-pt-224",
-        type=str,
-        help="Hub path to the tokenizer to save",
-    )
-
-    parser.add_argument(
-        "--output_path",
-        required=True,
-        type=str,
-        help="Path to save converted weights to",
-    )
-
-    args = parser.parse_args()
-    convert_pi0_checkpoint(
-        checkpoint_dir=args.checkpoint_dir,
-        precision=args.precision,
-        tokenizer_id=args.tokenizer_hub_id,
-        output_path=args.output_path,
-    )
@@ -1,141 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import torch
-import torch.nn.functional as F  # noqa: N812
-from packaging.version import Version
-
-if Version(torch.__version__) > Version("2.5.0"):
-    # Ffex attention is only available from torch 2.5 onwards
-    from torch.nn.attention.flex_attention import (
-        _mask_mod_signature,
-        _round_up_to_multiple,
-        create_block_mask,
-        create_mask,
-        flex_attention,
-    )
-
-
-# @torch.compile(dynamic=False)
-def flex_attention_forward(
-    attention_mask: torch.Tensor,
-    batch_size: int,
-    head_dim: int,
-    query_states: torch.Tensor,
-    key_states: torch.Tensor,
-    value_states: torch.Tensor,
-    scaling=None,
-):
-    """
-    This is defined out of classes to make compile happy.
-    """
-
-    original_dtype = query_states.dtype
-    num_att_heads = 8
-    num_key_value_heads = 1
-    num_key_value_groups = num_att_heads // num_key_value_heads
-
-    key_states = key_states[:, :, :, None, :]
-    key_states = key_states.expand(
-        batch_size, key_states.shape[1], num_key_value_heads, num_key_value_groups, head_dim
-    )
-    key_states = key_states.reshape(
-        batch_size, key_states.shape[1], num_key_value_heads * num_key_value_groups, head_dim
-    )
-
-    value_states = value_states[:, :, :, None, :]
-    value_states = value_states.expand(
-        batch_size, value_states.shape[1], num_key_value_heads, num_key_value_groups, head_dim
-    )
-    value_states = value_states.reshape(
-        batch_size, value_states.shape[1], num_key_value_heads * num_key_value_groups, head_dim
-    )
-
-    query_states = query_states.transpose(1, 2)
-    key_states = key_states.transpose(1, 2)
-    value_states = value_states.transpose(1, 2)
-
-    query_states = query_states.to(torch.float32)
-    key_states = key_states.to(torch.float32)
-    value_states = value_states.to(torch.float32)
-
-    causal_mask = attention_mask
-    if causal_mask is not None:
-        causal_mask = causal_mask[:, None, :, : key_states.shape[2]]
-
-        if causal_mask.shape[1] == 1 and query_states.shape[1] > 1:
-            causal_mask = causal_mask.expand(-1, query_states.shape[1], -1, -1)
-
-    def precomputed_mask_factory(precomputed_mask: torch.Tensor) -> _mask_mod_signature:
-        def mask_mod(b, h, q_idx, kv_idx):
-            # Danger zone: if b,h,q_idx,kv_idx exceed the shape, device-side assert occurs.
-            return precomputed_mask[b][h][q_idx][kv_idx]
-
-        return mask_mod
-
-    b_mask, h_mask, q_len, kv_len = causal_mask.shape  # The shape of your mask
-
-    block_size = 128
-    q_len_rounded = _round_up_to_multiple(q_len, block_size)
-    kv_len_rounded = _round_up_to_multiple(kv_len, block_size)
-
-    # *CRITICAL* we do need to expand here, else we get a CUDA index error
-
-    pad_q = q_len_rounded - q_len
-    pad_k = kv_len_rounded - kv_len
-
-    padded_causal_mask = F.pad(causal_mask, (0, pad_k, 0, pad_q), value=0.0)
-    mask_mod_fn_orig = precomputed_mask_factory(padded_causal_mask)
-
-    mask_4d = create_mask(
-        mod_fn=mask_mod_fn_orig,
-        B=b_mask,
-        H=h_mask,
-        Q_LEN=q_len_rounded,
-        KV_LEN=kv_len_rounded,
-        device=causal_mask.device,
-        _compile=False,
-    )
-
-    mask_mod_fn_padded = precomputed_mask_factory(mask_4d)
-    block_mask = create_block_mask(
-        mask_mod=mask_mod_fn_padded,
-        B=b_mask,
-        H=h_mask,
-        Q_LEN=q_len_rounded,
-        KV_LEN=kv_len_rounded,
-        BLOCK_SIZE=block_size,
-        device=causal_mask.device,
-        _compile=False,
-    )
-
-    #  mask is applied inside the kernel, ideally more efficiently than score_mod.
-    attn_output, attention_weights = flex_attention(
-        query_states,
-        key_states,
-        value_states,
-        block_mask=block_mask,
-        enable_gqa=True,  # because we shaped query/key states for GQA
-        scale=head_dim**-0.5 if scaling is None else scaling,
-        return_lse=True,
-    )
-
-    attn_output = attn_output.to(dtype=original_dtype)
-    attn_output = attn_output.transpose(1, 2).contiguous()  # [B, Q_LEN, H, head_dim]
-    attn_output = attn_output.reshape(
-        batch_size,
-        -1,
-        attn_output.shape[2] * attn_output.shape[3],  # merges [H, head_dim]
-    )
-    return attn_output
@@ -1,420 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-import torch
-import torch.version
-from pytest import Cache
-from torch import nn
-from transformers import (
-    AutoConfig,
-    GemmaForCausalLM,
-    PaliGemmaForConditionalGeneration,
-    PretrainedConfig,
-    PreTrainedModel,
-)
-from transformers.models.auto import CONFIG_MAPPING
-
-from lerobot.policies.pi0.flex_attention import flex_attention_forward
-
-
-def apply_rope(x, positions, max_wavelength=10_000):
-    """
-    Applies RoPE positions [B, L] to x [B, L, H, D].
-    """
-    d_half = x.shape[-1] // 2
-    device = x.device
-    dtype = x.dtype
-    x = x.to(torch.float32)
-
-    freq_exponents = (2.0 / x.shape[-1]) * torch.arange(d_half, dtype=torch.float32, device=device)
-    timescale = max_wavelength**freq_exponents
-    radians = positions[..., None].to(torch.float32) / timescale[None, None, :].to(torch.float32)
-
-    radians = radians[..., None, :]
-
-    sin = torch.sin(radians)  # .to(dtype=dtype)
-    cos = torch.cos(radians)  # .to(dtype=dtype)
-
-    x1, x2 = x.split(d_half, dim=-1)
-    res = torch.empty_like(x)
-    res[..., :d_half] = x1 * cos - x2 * sin
-    res[..., d_half:] = x2 * cos + x1 * sin
-
-    return res.to(dtype)
-
-
-class PaliGemmaWithExpertConfig(PretrainedConfig):
-    model_type = "PaliGemmaWithExpertModel"
-    sub_configs = {"paligemma_config": AutoConfig, "gemma_expert_config": AutoConfig}
-
-    def __init__(
-        self,
-        paligemma_config: dict | None = None,
-        gemma_expert_config: dict | None = None,
-        freeze_vision_encoder: bool = True,
-        train_expert_only: bool = True,
-        attention_implementation: str = "eager",
-        **kwargs,
-    ):
-        self.freeze_vision_encoder = freeze_vision_encoder
-        self.train_expert_only = train_expert_only
-        self.attention_implementation = attention_implementation
-
-        if paligemma_config is None:
-            # Default config from Pi0
-            self.paligemma_config = CONFIG_MAPPING["paligemma"](
-                transformers_version="4.48.1",
-                _vocab_size=257152,
-                bos_token_id=2,
-                eos_token_id=1,
-                hidden_size=2048,
-                image_token_index=257152,
-                model_type="paligemma",
-                pad_token_id=0,
-                projection_dim=2048,
-                text_config={
-                    "hidden_activation": "gelu_pytorch_tanh",
-                    "hidden_size": 2048,
-                    "intermediate_size": 16384,
-                    "model_type": "gemma",
-                    "num_attention_heads": 8,
-                    "num_hidden_layers": 18,
-                    "num_image_tokens": 256,
-                    "num_key_value_heads": 1,
-                    "torch_dtype": "float32",
-                    "vocab_size": 257152,
-                },
-                vision_config={
-                    "hidden_size": 1152,
-                    "intermediate_size": 4304,
-                    "model_type": "siglip_vision_model",
-                    "num_attention_heads": 16,
-                    "num_hidden_layers": 27,
-                    "num_image_tokens": 256,
-                    "patch_size": 14,
-                    "projection_dim": 2048,
-                    "projector_hidden_act": "gelu_fast",
-                    "torch_dtype": "float32",
-                    "vision_use_head": False,
-                },
-            )
-        elif isinstance(self.paligemma_config, dict):
-            # Override Pi0 default config for PaliGemma
-            if "model_type" not in gemma_expert_config:
-                paligemma_config["model_type"] = "paligemma"
-
-            cfg_cls = CONFIG_MAPPING[paligemma_config["model_type"]]
-            self.paligemma_config = cfg_cls(**paligemma_config)
-
-        if gemma_expert_config is None:
-            # Default config from Pi0
-            self.gemma_expert_config = CONFIG_MAPPING["gemma"](
-                attention_bias=False,
-                attention_dropout=0.0,
-                bos_token_id=2,
-                eos_token_id=1,
-                head_dim=256,
-                hidden_act="gelu_pytorch_tanh",
-                hidden_activation="gelu_pytorch_tanh",
-                hidden_size=1024,
-                initializer_range=0.02,
-                intermediate_size=4096,
-                max_position_embeddings=8192,
-                model_type="gemma",
-                num_attention_heads=8,
-                num_hidden_layers=18,
-                num_key_value_heads=1,
-                pad_token_id=0,
-                rms_norm_eps=1e-06,
-                rope_theta=10000.0,
-                torch_dtype="float32",
-                transformers_version="4.48.1",
-                use_cache=True,
-                vocab_size=257152,
-            )
-        elif isinstance(self.gemma_expert_config, dict):
-            # Override Pi0 default config for Gemma Expert
-            if "model_type" not in gemma_expert_config:
-                gemma_expert_config["model_type"] = "gemma"
-
-            cfg_cls = CONFIG_MAPPING[paligemma_config["model_type"]]
-            self.gemma_expert_config = cfg_cls(**gemma_expert_config)
-
-        super().__init__(**kwargs)
-
-    def __post_init__(self):
-        super().__post_init__()
-        if self.train_expert_only and not self.freeze_vision_encoder:
-            raise ValueError(
-                "You set `freeze_vision_encoder=False` and `train_expert_only=True` which are not compatible."
-            )
-
-        if self.attention_implementation not in ["eager", "fa2", "flex"]:
-            raise ValueError(
-                f"Wrong value provided for `attention_implementation` ({self.attention_implementation}). Expected 'eager', 'fa2' or 'flex'."
-            )
-
-
-class PaliGemmaWithExpertModel(PreTrainedModel):
-    config_class = PaliGemmaWithExpertConfig
-
-    def __init__(self, config: PaliGemmaWithExpertConfig):
-        super().__init__(config=config)
-        self.config = config
-        self.paligemma = PaliGemmaForConditionalGeneration(config=config.paligemma_config)
-        self.gemma_expert = GemmaForCausalLM(config=config.gemma_expert_config)
-        # Remove unused embed_tokens
-        self.gemma_expert.model.embed_tokens = None
-
-        self.to_bfloat16_like_physical_intelligence()
-        self.set_requires_grad()
-
-    def set_requires_grad(self):
-        if self.config.freeze_vision_encoder:
-            self.paligemma.vision_tower.eval()
-            for params in self.paligemma.vision_tower.parameters():
-                params.requires_grad = False
-
-        if self.config.train_expert_only:
-            self.paligemma.eval()
-            for params in self.paligemma.parameters():
-                params.requires_grad = False
-
-    def train(self, mode: bool = True):
-        super().train(mode)
-
-        if self.config.freeze_vision_encoder:
-            self.paligemma.vision_tower.eval()
-
-        if self.config.train_expert_only:
-            self.paligemma.eval()
-
-    def to_bfloat16_like_physical_intelligence(self):
-        self.paligemma = self.paligemma.to(dtype=torch.bfloat16)
-
-        params_to_change_dtype = [
-            "language_model.model.layers",
-            "gemma_expert.model.layers",
-            "vision_tower",
-            "multi_modal",
-        ]
-        for name, param in self.named_parameters():
-            if any(selector in name for selector in params_to_change_dtype):
-                param.data = param.data.to(dtype=torch.bfloat16)
-
-    def embed_image(self, image: torch.Tensor):
-        # Handle different transformers versions
-        if hasattr(self.paligemma, "get_image_features"):
-            return self.paligemma.get_image_features(image)
-        else:
-            return self.paligemma.model.get_image_features(image)
-
-    def embed_language_tokens(self, tokens: torch.Tensor):
-        return self.paligemma.language_model.embed_tokens(tokens)
-
-    # TODO: break down this huge forward into modules or functions
-    def forward(
-        self,
-        attention_mask: torch.Tensor | None = None,
-        position_ids: torch.LongTensor | None = None,
-        past_key_values: list[torch.FloatTensor] | Cache | None = None,
-        inputs_embeds: list[torch.FloatTensor] = None,
-        use_cache: bool | None = None,
-        fill_kv_cache: bool | None = None,
-    ):
-        models = [self.paligemma.language_model, self.gemma_expert.model]
-
-        for hidden_states in inputs_embeds:
-            # TODO this is very inefficient
-            # dtype is always the same, batch size too (if > 1 len)
-            # device could be trickier in multi gpu edge cases but that's it
-            if hidden_states is None:
-                continue
-            batch_size = hidden_states.shape[0]
-
-        # RMSNorm
-        num_layers = self.paligemma.config.text_config.num_hidden_layers
-        head_dim = self.paligemma.config.text_config.head_dim
-        for layer_idx in range(num_layers):
-            query_states = []
-            key_states = []
-            value_states = []
-            for i, hidden_states in enumerate(inputs_embeds):
-                if hidden_states is None:
-                    continue
-                layer = models[i].layers[layer_idx]
-                # normalizer = torch.tensor(models[i].config.hidden_size**0.5, dtype=hidden_states.dtype)
-                # hidden_states = hidden_states * normalizer
-                hidden_states = layer.input_layernorm(hidden_states)
-
-                input_shape = hidden_states.shape[:-1]
-                hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
-
-                hidden_states = hidden_states.to(dtype=torch.bfloat16)
-                query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape)
-                key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape)
-                value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape)
-
-                query_states.append(query_state)
-                key_states.append(key_state)
-                value_states.append(value_state)
-
-            # B,L,H,D with L sequence length, H number of heads, D head dim
-            # concatenate on the number of embeddings/tokens
-            query_states = torch.cat(query_states, dim=1)
-            key_states = torch.cat(key_states, dim=1)
-            value_states = torch.cat(value_states, dim=1)
-
-            query_states = apply_rope(query_states, position_ids)
-            key_states = apply_rope(key_states, position_ids)
-
-            if use_cache and past_key_values is None:
-                past_key_values = {}
-
-            if use_cache:
-                if fill_kv_cache:
-                    past_key_values[layer_idx] = {
-                        "key_states": key_states,
-                        "value_states": value_states,
-                    }
-                else:
-                    # TODO here, some optimization can be done - similar to a `StaticCache` we can declare the `max_len` before.
-                    # so we create an empty cache, with just one cuda malloc, and if (in autoregressive case) we reach
-                    # the max len, then we (for instance) double the cache size. This implementation already exists
-                    # in `transformers`. (molbap)
-                    key_states = torch.cat([past_key_values[layer_idx]["key_states"], key_states], dim=1)
-                    value_states = torch.cat(
-                        [past_key_values[layer_idx]["value_states"], value_states], dim=1
-                    )
-
-            attention_interface = self.get_attention_interface()
-            att_output = attention_interface(
-                attention_mask, batch_size, head_dim, query_states, key_states, value_states
-            )
-            att_output = att_output.to(dtype=torch.bfloat16)
-
-            # first part of att_output is prefix (up to sequence length, [:, 0:prefix_seq_len])
-            outputs_embeds = []
-            start = 0
-            for i, hidden_states in enumerate(inputs_embeds):
-                layer = models[i].layers[layer_idx]
-
-                if hidden_states is not None:
-                    end = start + hidden_states.shape[1]
-
-                    if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
-                        att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
-                    out_emb = layer.self_attn.o_proj(att_output[:, start:end])
-
-                    # TODO: first dropout (by default 0.0)
-
-                    # first residual
-                    out_emb += hidden_states
-                    after_first_residual = out_emb.clone()
-
-                    out_emb = layer.post_attention_layernorm(out_emb)
-                    out_emb = layer.mlp(out_emb)
-
-                    # TODO: second dropout (by default 0.0)
-
-                    # second residual
-                    out_emb += after_first_residual
-
-                    outputs_embeds.append(out_emb)
-
-                    start = end
-                else:
-                    outputs_embeds.append(None)
-
-            inputs_embeds = outputs_embeds
-
-        # final norm
-        outputs_embeds = []
-        for i, hidden_states in enumerate(inputs_embeds):
-            if hidden_states is not None:
-                out_emb = models[i].norm(hidden_states)
-                outputs_embeds.append(out_emb)
-            else:
-                outputs_embeds.append(None)
-
-        return outputs_embeds, past_key_values
-
-    def get_attention_interface(self):
-        if self.config.attention_implementation == "fa2":
-            attention_interface = self.flash_attention_forward
-        elif self.config.attention_implementation == "flex":
-            attention_interface = flex_attention_forward
-        else:
-            attention_interface = self.eager_attention_forward
-        return attention_interface
-
-    def flash_attention_forward(
-        self, attention_mask, batch_size, head_dim, query_states, key_states, value_states
-    ):
-        raise NotImplementedError("FA2 is not implemented (yet)")
-
-    def eager_attention_forward(
-        self, attention_mask, batch_size, head_dim, query_states, key_states, value_states
-    ):
-        num_att_heads = self.config.paligemma_config.text_config.num_attention_heads
-        num_key_value_heads = self.config.paligemma_config.text_config.num_key_value_heads
-        num_key_value_groups = num_att_heads // num_key_value_heads
-
-        # query_states: batch_size, sequence_length, num_att_head, head_dim
-        # key_states: batch_size, sequence_length, num_key_value_head, head_dim
-        # value_states: batch_size, sequence_length, num_key_value_head, head_dim
-        sequence_length = key_states.shape[1]
-
-        key_states = key_states[:, :, :, None, :].expand(
-            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
-        )
-        key_states = key_states.reshape(
-            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
-        )
-
-        value_states = value_states[:, :, :, None, :].expand(
-            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
-        )
-        value_states = value_states.reshape(
-            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
-        )
-
-        # Attention here is upcasted to float32 to match the original eager implementation.
-
-        query_states = query_states.to(dtype=torch.float32)
-        key_states = key_states.to(dtype=torch.float32)
-
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-
-        att_weights = torch.matmul(query_states, key_states.transpose(2, 3))
-        att_weights *= head_dim**-0.5
-        big_neg = -2.3819763e38  # See gemma/modules.py
-
-        masked_att_weights = torch.where(attention_mask[:, None, :, :], att_weights, big_neg)
-
-        probs = nn.functional.softmax(masked_att_weights, dim=-1)
-        probs = probs.to(dtype=value_states.dtype)
-
-        # probs: batch_size, num_key_value_head, num_att_head, sequence_length, sequence_length
-        # value_states: batch_size, sequence_length, num_att_heads, head_dim
-
-        att_output = torch.matmul(probs, value_states.permute(0, 2, 1, 3))
-
-        att_output = att_output.permute(0, 2, 1, 3)
-        # we use -1 because sequence length can change
-        att_output = att_output.reshape(batch_size, -1, num_key_value_heads * num_key_value_groups * head_dim)
-
-        return att_output
@@ -0,0 +1,49 @@
+# π₀.₅ (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**.
+
+---
+
+## Model Overview
+
+| Feature              | π₀                                                     | π₀.₅                                      |
+| -------------------- | ------------------------------------------------------ | ----------------------------------------- |
+| 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 (Uses `state_proj` layer)                        | 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).
@@ -0,0 +1,21 @@
+#!/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_pi05 import PI05Config
+from .modeling_pi05 import PI05Policy
+from .processor_pi05 import make_pi05_pre_post_processors
+
+__all__ = ["PI05Config", "PI05Policy", "make_pi05_pre_post_processors"]
@@ -0,0 +1,153 @@
+#!/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 FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim.optimizers import AdamWConfig
+from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
+
+
+@PreTrainedConfig.register_subclass("pi05")
+@dataclass
+class PI05Config(PreTrainedConfig):
+    paligemma_variant: str = "gemma_2b"
+    action_expert_variant: str = "gemma_300m"
+    dtype: str = "float32"  # Options: "bfloat16", "float32"
+
+    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
+    max_action_dim: int = 32
+
+    # Flow matching parameters: see openpi `PI0Pytorch`
+    num_inference_steps: int = 10
+    time_sampling_beta_alpha: float = 1.5
+    time_sampling_beta_beta: float = 1.0
+    time_sampling_scale: float = 0.999
+    time_sampling_offset: float = 0.001
+    min_period: float = 4e-3
+    max_period: float = 4.0
+
+    image_resolution: tuple[int, int] = (224, 224)  # see openpi `preprocessing_pytorch.py`
+
+    # Add empty images. Used to add empty cameras when no image features are present.
+    empty_cameras: int = 0
+
+    tokenizer_max_length: int = 200  # see openpi `__post_init__`
+
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.IDENTITY,
+            "STATE": NormalizationMode.QUANTILES,  # Pi0.5 uses quantiles for state
+            "ACTION": NormalizationMode.QUANTILES,  # Pi0.5 uses quantiles for action
+        }
+    )
+
+    # 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`
+    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."""
+        for i in range(self.empty_cameras):
+            key = f"observation.images.empty_camera_{i}"
+            empty_camera = PolicyFeature(
+                type=FeatureType.VISUAL,
+                shape=(3, *self.image_resolution),  # Use configured image resolution
+            )
+            self.input_features[key] = empty_camera
+
+        if "observation.state" not in self.input_features:
+            state_feature = PolicyFeature(
+                type=FeatureType.STATE,
+                shape=(self.max_state_dim,),  # Padded to max_state_dim
+            )
+            self.input_features["observation.state"] = state_feature
+
+        if "action" not in self.output_features:
+            action_feature = PolicyFeature(
+                type=FeatureType.ACTION,
+                shape=(self.max_action_dim,),  # Padded to max_action_dim
+            )
+            self.output_features["action"] = action_feature
+
+    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
@@ -0,0 +1,171 @@
+#!/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 copy import deepcopy
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+
+from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.policies.pi05.configuration_pi05 import PI05Config
+from lerobot.policies.pi05.modeling_pi05 import pad_vector
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    ProcessorStep,
+    ProcessorStepRegistry,
+    RenameObservationsProcessorStep,
+    TokenizerProcessorStep,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+from lerobot.processor.core import EnvTransition, TransitionKey
+from lerobot.utils.constants import (
+    OBS_STATE,
+    POLICY_POSTPROCESSOR_DEFAULT_NAME,
+    POLICY_PREPROCESSOR_DEFAULT_NAME,
+)
+
+
+@ProcessorStepRegistry.register(name="pi05_prepare_state_tokenizer_processor_step")
+@dataclass
+class Pi05PrepareStateTokenizerProcessorStep(ProcessorStep):
+    """
+    Processor step to prepare the state and tokenize the language input.
+    """
+
+    max_state_dim: int = 32
+    task_key: str = "task"
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        transition = transition.copy()
+
+        state = transition.get(TransitionKey.OBSERVATION, {}).get(OBS_STATE)
+        if state is None:
+            raise ValueError("State is required for PI05")
+        tasks = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}).get(self.task_key)
+        if tasks is None:
+            raise ValueError("No task found in complementary data")
+
+        # TODO: check if this necessary
+        state = deepcopy(state)
+
+        # Prepare state (pad to max_state_dim)
+        state = pad_vector(state, self.max_state_dim)
+
+        # State should already be normalized to [-1, 1] by the NormalizerProcessorStep that runs before this step
+        # Discretize into 256 bins (see openpi `PaligemmaTokenizer.tokenize()`)
+        state_np = state.cpu().numpy()
+        discretized_states = np.digitize(state_np, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1
+
+        full_prompts = []
+        for i, task in enumerate(tasks):
+            cleaned_text = task.strip().replace("_", " ").replace("\n", " ")
+            state_str = " ".join(map(str, discretized_states[i]))
+            full_prompt = f"Task: {cleaned_text}, State: {state_str};\nAction: "
+            full_prompts.append(full_prompt)
+
+        transition[TransitionKey.COMPLEMENTARY_DATA][self.task_key] = full_prompts
+        # Normalize state to [-1, 1] range if needed (assuming it's already normalized by normalizer processor step!!)
+        # Discretize into 256 bins (see openpi `PaligemmaTokenizer.tokenize()`)
+        return transition
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        """
+        This step does not alter the feature definitions.
+        """
+        return features
+
+
+def make_pi05_pre_post_processors(
+    config: PI05Config,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """
+    Constructs pre-processor and post-processor pipelines for the PI0 policy.
+
+    The pre-processing pipeline prepares input data for the model by:
+    1. Renaming features to match pretrained configurations.
+    2. Normalizing input and output features based on dataset statistics.
+    3. Adding a batch dimension.
+    4. Appending a newline character to the task description for tokenizer compatibility.
+    5. Tokenizing the text prompt using the PaliGemma tokenizer.
+    6. Moving all data to the specified device.
+
+    The post-processing pipeline handles the model's output by:
+    1. Moving data to the CPU.
+    2. Unnormalizing the output features to their original scale.
+
+    Args:
+        config: The configuration object for the PI0 policy.
+        dataset_stats: A dictionary of statistics for normalization.
+        preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
+        postprocessor_kwargs: Additional arguments for the post-processor pipeline.
+
+    Returns:
+        A tuple containing the configured pre-processor and post-processor pipelines.
+    """
+
+    # Add remaining processors
+    input_steps: list[ProcessorStep] = [
+        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
+        AddBatchDimensionProcessorStep(),
+        # NOTE: NormalizerProcessorStep MUST come before Pi05PrepareStateTokenizerProcessorStep
+        # because the tokenizer step expects normalized state in [-1, 1] range for discretization
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+        Pi05PrepareStateTokenizerProcessorStep(max_state_dim=config.max_state_dim),
+        TokenizerProcessorStep(
+            tokenizer_name="google/paligemma-3b-pt-224",
+            max_length=config.tokenizer_max_length,
+            padding_side="right",
+            padding="max_length",
+        ),
+        DeviceProcessorStep(device=config.device),
+    ]
+
+    output_steps: list[ProcessorStep] = [
+        UnnormalizerProcessorStep(
+            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
+        DeviceProcessorStep(device="cpu"),
+    ]
+
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )
@@ -1,137 +0,0 @@
-from dataclasses import dataclass, field
-
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
-from lerobot.optim.optimizers import AdamWConfig
-from lerobot.optim.schedulers import (
-    CosineDecayWithWarmupSchedulerConfig,
-)
-from lerobot.utils.constants import OBS_IMAGES
-
-
-@PreTrainedConfig.register_subclass("pi0fast")
-@dataclass
-class PI0FASTConfig(PreTrainedConfig):
-    # Input / output structure.
-    n_obs_steps: int = 1
-    chunk_size: int = 10
-    n_action_steps: int = 5
-
-    normalization_mapping: dict[str, NormalizationMode] = field(
-        default_factory=lambda: {
-            "VISUAL": NormalizationMode.IDENTITY,
-            "STATE": NormalizationMode.MEAN_STD,
-            "ACTION": NormalizationMode.MEAN_STD,
-        }
-    )
-
-    # Shorter state and action vectors will be padded
-    max_state_dim: int = 32  # 32
-    max_action_dim: int = 32  # 32
-
-    # Image preprocessing
-    resize_imgs_with_padding: tuple[int, int] = (224, 224)
-    interpolate_like_pi: bool = False
-
-    # Add empty images. Used by pi0_aloha_sim which adds the empty
-    # left and right wrist cameras in addition to the top camera.
-    empty_cameras: int = 0
-
-    # Converts the joint and gripper values from the standard Aloha space to
-    # the space used by the pi internal runtime which was used to train the base model.
-    adapt_to_pi_aloha: bool = False
-
-    # Converts joint dimensions to deltas with respect to the current state before passing to the model.
-    # Gripper dimensions will remain in absolute values.
-    use_delta_joint_actions_aloha: bool = False
-
-    # Tokenizer
-    tokenizer_max_length: int = 48
-
-    # Projector
-    proj_width: int = 1024
-
-    # Decoding
-    max_decoding_steps: int = 256
-    fast_skip_tokens: int = 128  # Skip last 128 tokens in PaliGemma vocab since they are special tokens
-    max_input_seq_len: int = 256  # 512
-
-    # Utils
-    use_cache: bool = True
-
-    # Frozen parameters
-    freeze_vision_encoder: bool = True
-    freeze_lm_head: bool = True
-
-    # Training presets
-    optimizer_lr: float = 1e-4
-    optimizer_betas: tuple[float, float] = (0.9, 0.95)
-    optimizer_eps: float = 1e-8
-    optimizer_weight_decay: float = 1e-5
-
-    scheduler_warmup_steps: int = 1_000
-    scheduler_decay_steps: int = 30_000
-    scheduler_decay_lr: float = 2.5e-6
-
-    checkpoint_path: str = None
-
-    padding_side: str = "right"
-
-    precision: str = "bfloat16"
-    grad_clip_norm: float = 1
-
-    # Allows padding/truncation of generated action tokens during detokenization to ensure decoding.
-    # In the original version, tensors of 0s were generated if shapes didn't match for stable decoding.
-    relaxed_action_decoding: bool = True
-
-    def __post_init__(self):
-        super().__post_init__()
-
-        """Input validation (not exhaustive)."""
-        if self.n_action_steps > self.chunk_size:
-            raise ValueError(
-                f"The chunk size is the upper bound for the number of action steps per model invocation. Got "
-                f"{self.n_action_steps} for `n_action_steps` and {self.chunk_size} for `chunk_size`."
-            )
-        if self.n_obs_steps != 1:
-            raise ValueError(
-                f"Multiple observation steps not handled yet. Got `nobs_steps={self.n_obs_steps}`"
-            )
-
-    def validate_features(self) -> None:
-        for i in range(self.empty_cameras):
-            key = f"{OBS_IMAGES}.empty_camera_{i}"
-            empty_camera = PolicyFeature(
-                type=FeatureType.VISUAL,
-                shape=(3, 480, 640),
-            )
-            self.input_features[key] = empty_camera
-
-    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.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
@@ -1,980 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 Physical Intelligence and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-π0+FAST: Efficient Action Tokenization for Vision-Language-Action Models
-
-[Paper](https://huggingface.co/papers/2501.09747)
-[Jax code](https://github.com/Physical-Intelligence/openpi)
-
-Designed by Physical Intelligence. Ported from Jax by Hugging Face.
-Disclaimer: It is not expected to perform as well as the original implementation.
-
-Example of finetuning the pi0+FAST pretrained model (`pi0_fast_base` in `openpi`):
-```bash
-lerobot-train \
--policy.path=lerobot/pi0fast_base \
--dataset.repo_id=danaaubakirova/koch_test
-```
-
-Example of training the pi0+FAST neural network with from scratch:
-```bash
-lerobot-train \
--policy.type=pi0fast \
--dataset.repo_id=danaaubakirova/koch_test
-```
-
-Example of using the pi0 pretrained model outside LeRobot training framework:
-```python
-policy = PI0FASTPolicy.from_pretrained("lerobot/pi0fast_base")
-```
-
-"""
-
-from collections import deque
-from functools import partial
-
-import numpy as np
-import torch
-import torch.nn.functional as F  # noqa: N812
-from PIL import Image
-from scipy.fft import idct
-from torch import Tensor, nn
-from transformers import AutoProcessor, AutoTokenizer, PaliGemmaForConditionalGeneration
-from transformers.cache_utils import HybridCache, StaticCache
-from transformers.models.auto import CONFIG_MAPPING
-
-from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
-from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.utils.constants import ACTION, OBS_STATE
-
-PRECISION = {
-    "float16": torch.float16,
-    "float32": torch.float32,
-    "bfloat16": torch.bfloat16,
-}
-
-
-def normalize(x, min_val, max_val):
-    return (x - min_val) / (max_val - min_val)
-
-
-def unnormalize(x, min_val, max_val):
-    return x * (max_val - min_val) + min_val
-
-
-def safe_arcsin(value):
-    # This ensures that the input stays within
-    # [−1,1] to avoid invalid values for arcsin
-    return torch.arcsin(torch.clamp(value, -1.0, 1.0))
-
-
-def aloha_gripper_to_angular(value):
-    # Aloha transforms the gripper positions into a linear space. The following code
-    # reverses this transformation to be consistent with pi0 which is pretrained in
-    # angular space.
-    #
-    # These values are coming from the Aloha code:
-    # PUPPET_GRIPPER_POSITION_OPEN, PUPPET_GRIPPER_POSITION_CLOSED
-    value = unnormalize(value, min_val=0.01844, max_val=0.05800)
-
-    # This is the inverse of the angular to linear transformation inside the Interbotix code.
-    def linear_to_radian(linear_position, arm_length, horn_radius):
-        value = (horn_radius**2 + linear_position**2 - arm_length**2) / (2 * horn_radius * linear_position)
-        return safe_arcsin(value)
-
-    # The constants are taken from the Interbotix code.
-    value = linear_to_radian(value, arm_length=0.036, horn_radius=0.022)
-
-    # Normalize to [0, 1].
-    # The values 0.4 and 1.5 were measured on an actual Trossen robot.
-    return normalize(value, min_val=0.4, max_val=1.5)
-
-
-def aloha_gripper_from_angular(value):
-    # Convert from the gripper position used by pi0 to the gripper position that is used by Aloha.
-    # Note that the units are still angular but the range is different.
-
-    # The values 0.4 and 1.5 were measured on an actual Trossen robot.
-    value = unnormalize(value, min_val=0.4, max_val=1.5)
-
-    # These values are coming from the Aloha code:
-    # PUPPET_GRIPPER_JOINT_OPEN, PUPPET_GRIPPER_JOINT_CLOSE
-    return normalize(value, min_val=-0.6213, max_val=1.4910)
-
-
-def aloha_gripper_from_angular_inv(value):
-    # Directly inverts the gripper_from_angular function.
-    value = unnormalize(value, min_val=-0.6213, max_val=1.4910)
-    return normalize(value, min_val=0.4, max_val=1.5)
-
-
-class PI0FASTPolicy(PreTrainedPolicy):
-    """Wrapper class around PI0FAST tokenizer and model to train and run inference within LeRobot."""
-
-    config_class = PI0FASTConfig
-    name = "pi0fast"
-
-    def __init__(
-        self,
-        config: PI0FASTConfig,
-        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
-    ):
-        """
-        Args:
-            config: Policy configuration class instance or None, in which case the default instantiation of
-                    the configuration class is used.
-            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
-                that they will be passed with a call to `load_state_dict` before the policy is used.
-        """
-
-        super().__init__(config)
-        config.validate_features()
-        self.config = config
-
-        self.language_tokenizer = AutoProcessor.from_pretrained("google/paligemma-3b-pt-224")
-        self.model = PI0FAST(config)
-
-        self.reset()
-
-    def reset(self):
-        """This should be called whenever the environment is reset."""
-        self._action_queue = deque([], maxlen=self.config.n_action_steps)
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        """Override the from_pretrained method to display important disclaimer."""
-        print(
-            "⚠️  DISCLAIMER: The PI0FAST model is ported from JAX by the Hugging Face team. \n"
-            "   It is not expected to perform as well as the original implementation. \n"
-            "   Original implementation: https://github.com/Physical-Intelligence/openpi"
-        )
-        return super().from_pretrained(*args, **kwargs)
-
-    def get_optim_params(self) -> dict:
-        return self.parameters()
-
-    def _pi_aloha_decode_state(self, state):
-        # Flip the joints.
-        for motor_idx in [1, 2, 8, 9]:
-            state[:, motor_idx] *= -1
-        # Reverse the gripper transformation that is being applied by the Aloha runtime.
-        for motor_idx in [6, 13]:
-            state[:, motor_idx] = aloha_gripper_to_angular(state[:, motor_idx])
-        return state
-
-    def _pi_aloha_encode_actions(self, actions):
-        # Flip the joints.
-        for motor_idx in [1, 2, 8, 9]:
-            actions[:, :, motor_idx] *= -1
-        # Reverse the gripper transformation that is being applied by the Aloha runtime.
-        for motor_idx in [6, 13]:
-            actions[:, :, motor_idx] = aloha_gripper_from_angular(actions[:, :, motor_idx])
-        return actions
-
-    def _pi_aloha_encode_actions_inv(self, actions):
-        # Flip the joints again.
-        for motor_idx in [1, 2, 8, 9]:
-            actions[:, :, motor_idx] *= -1
-        # Reverse the gripper transformation that is being applied by the Aloha runtime.
-        for motor_idx in [6, 13]:
-            actions[:, :, motor_idx] = aloha_gripper_from_angular_inv(actions[:, :, motor_idx])
-        return actions
-
-    @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
-        """Predict a chunk of actions given environment observations."""
-        raise NotImplementedError("Currently not implemented for PI0FAST")
-
-    @torch.no_grad()
-    def select_action(self, batch: dict[str, Tensor]) -> Tensor:
-        """Select a single action given environment observations.
-
-        This method wraps `select_actions` in order to return one action at a time for execution in the
-        environment. It works by managing the actions in a queue and only calling `select_actions` when the
-        queue is empty.
-        """
-        self.eval()
-
-        if self.config.adapt_to_pi_aloha:
-            batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
-
-        # Action queue logic for n_action_steps > 1. When the action_queue is depleted, populate it by
-        # querying the policy.
-        if len(self._action_queue) == 0:
-            actions = self.model.generate_actions(batch)
-
-            actions = actions[:, : self.config.n_action_steps]
-
-            original_action_dim = self.config.action_feature.shape[
-                0
-            ]  # self.config.max_action_dim  # self.config.action_feature.shape[0]
-            actions = actions[:, :, :original_action_dim]
-
-            if self.config.adapt_to_pi_aloha:
-                actions = self._pi_aloha_encode_actions(actions)
-
-            # `self.model.forward` returns a (batch_size, n_action_steps, action_dim) tensor, but the queue
-            # effectively has shape (n_action_steps, batch_size, *), hence the transpose.
-            self._action_queue.extend(actions.transpose(0, 1))
-        return self._action_queue.popleft()
-
-    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
-        if self.config.adapt_to_pi_aloha:
-            batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
-            batch[ACTION] = self._pi_aloha_encode_actions_inv(batch[ACTION])
-        loss_dict = self.model.forward(batch)
-        return loss_dict["loss"], loss_dict
-
-
-def block_causal_update_causal_mask(
-    attention_mask,
-    token_type_ids=None,
-    past_key_values=None,
-    cache_position=None,
-    input_tensor=None,
-    attn_implementation: str = "eager",
-    dtype: torch.dtype = "float32",
-):
-    """
-    Update the causal mask during training and generation. It can be customized to different attention masks.
-    """
-    if attn_implementation == "flash_attention_2":
-        if attention_mask is not None and 0.0 in attention_mask:
-            return attention_mask
-        return None
-    using_static_cache = isinstance(past_key_values, StaticCache)
-    min_dtype = torch.finfo(dtype).min
-
-    if input_tensor is None:
-        input_tensor = attention_mask
-
-    inputs_lead_dim, sequence_length = input_tensor.shape[:2]
-
-    if using_static_cache or 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
-        )
-
-    # Handle precomputed attention masks
-    if attention_mask is not None and attention_mask.dim() == 4:
-        return attention_mask
-
-    # Causal mask initialization
-    causal_mask = torch.full(
-        (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device
-    )
-
-    # Standard causal masking (triu ensures tokens can only attend to past)
-    if sequence_length != 1:
-        causal_mask = torch.triu(causal_mask, diagonal=1)
-
-        # Apply block causal mask
-        if token_type_ids is not None:
-            token_type_ids = token_type_ids.to(causal_mask.device).bool()
-            cumsum = torch.cumsum(token_type_ids, dim=1)
-            block_causal_mask = cumsum[:, None, :] <= cumsum[:, :, None]
-
-            # Combine causal_mask with block-wise attention mask
-            causal_mask = torch.where(block_causal_mask, 0.0, causal_mask)
-            causal_mask = causal_mask[:, None, :, :]
-        else:
-            # Apply past cache position constraint
-            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)
-    else:
-        # Apply past cache position constraint
-        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 edits
-        mask_length = attention_mask.shape[-1]
-
-        # Apply padding mask
-        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 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,
-    num_logits_to_keep=None,
-    labels=None,
-    self=None,
-    **kwargs,
-):
-    # create block causal attention
-    if cache_position[0] > 0 and input_ids.shape[1] > 0:
-        input_tensor = input_ids[:, -1:]
-        new_positions = (
-            torch.ones(
-                (position_ids.shape[0], input_ids.shape[1]),
-                dtype=position_ids.dtype,
-                device=position_ids.device,
-            ).cumsum(-1)
-            + position_ids[:, -1:]
-        )
-        position_ids = torch.cat([position_ids, new_positions], dim=-1)
-    else:
-        input_tensor = inputs_embeds
-    attention_mask = block_causal_update_causal_mask(
-        attention_mask=attention_mask,
-        past_key_values=past_key_values,
-        cache_position=cache_position,
-        input_tensor=input_tensor,
-        token_type_ids=token_type_ids,
-        dtype=self.dtype,
-        attn_implementation=self.config.text_config._attn_implementation,
-    )
-    # Overwritten -- custom `position_ids` and `pixel_values` handling
-    model_inputs = self.language_model.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,
-        num_logits_to_keep=num_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._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
-
-
-class PI0FAST(nn.Module):
-    def __init__(self, config: PI0FASTConfig):
-        super().__init__()
-        self.config = config
-
-        # TODO: move tokenizers in Policy
-        fast_tokenizer_path = "physical-intelligence/fast"
-        pi0_paligemma_path = "google/paligemma-3b-pt-224"
-        self.paligemma_tokenizer = AutoTokenizer.from_pretrained(pi0_paligemma_path)
-        self.processor = AutoProcessor.from_pretrained(pi0_paligemma_path)
-        self.fast_tokenizer = AutoProcessor.from_pretrained(fast_tokenizer_path, trust_remote_code=True)
-        self.fast_skip_tokens = self.config.fast_skip_tokens
-        self.max_input_seq_len = self.config.max_input_seq_len
-        self.action_horizon = self.config.chunk_size
-        self.action_dim = self.config.action_feature.shape[
-            0
-        ]  # self.config.max_action_dim  # self.config.action_feature.shape[0]
-        precision = config.precision
-        torch_precision = PRECISION.get(precision, torch.float32)
-        self.pad_token_id = (
-            self.paligemma_tokenizer.pad_token_id
-            if hasattr(self.paligemma_tokenizer, "pad_token_id")
-            else self.paligemma_tokenizer.eos_token_id
-        )
-
-        paligemma_config = CONFIG_MAPPING["paligemma"](
-            transformers_version="4.48.1",
-            _vocab_size=257152,
-            bos_token_id=2,
-            eos_token_id=1,
-            hidden_size=2048,
-            image_token_index=257152,
-            model_type="paligemma",
-            pad_token_id=0,
-            projection_dim=2048,
-            text_config={
-                "hidden_activation": "gelu_pytorch_tanh",
-                "hidden_size": 2048,
-                "intermediate_size": 16384,
-                "model_type": "gemma",
-                "num_attention_heads": 8,
-                "num_hidden_layers": 18,
-                "num_image_tokens": 256,
-                "num_key_value_heads": 1,
-                "torch_dtype": precision,
-                "vocab_size": 257152,
-                "_attn_implementation": "eager",
-            },
-            vision_config={
-                "hidden_size": 1152,
-                "intermediate_size": 4304,
-                "model_type": "siglip_vision_model",
-                "num_attention_heads": 16,
-                "num_hidden_layers": 27,
-                "num_image_tokens": 256,
-                "patch_size": 14,
-                "projection_dim": 2048,
-                "projector_hidden_act": "gelu_pytorch_tanh",
-                "torch_dtype": precision,
-                "vision_use_head": False,
-            },
-        )
-        self.pi0_paligemma = PaliGemmaForConditionalGeneration(config=paligemma_config)
-
-        self.pi0_paligemma.prepare_inputs_for_generation = partial(
-            prepare_inputs_for_generation, self=self.pi0_paligemma
-        )
-        # change important stuff in bf16
-        params_to_change_dtype = [
-            "language_model",
-            "vision_tower",
-            "multi_modal",
-        ]
-        for name, param in self.pi0_paligemma.named_parameters():
-            if any(selector in name for selector in params_to_change_dtype):
-                param.data = param.data.to(dtype=torch_precision)
-        self.set_requires_grad()
-        self.image_keys = self.config.image_features.keys()
-        # TODO: Remove this once we bump transformers to >4.52.0 because the attribute will be removed
-        # AttributeError: 'PaliGemmaConfig' object has no attribute 'ignore_index'
-        self.ignore_index = self.pi0_paligemma.config.ignore_index
-        self.padding_side = self.config.padding_side
-
-    def set_requires_grad(self):
-        if self.config.freeze_vision_encoder:
-            self.pi0_paligemma.vision_tower.eval()
-            for params in self.pi0_paligemma.vision_tower.parameters():
-                params.requires_grad = False
-        # To avoid unused params issue with distributed training
-        if self.config.freeze_lm_head:
-            for name, params in self.pi0_paligemma.named_parameters():
-                if "embed_tokens" in name:  # lm heads and embedding layer are tied
-                    params.requires_grad = False
-
-    def embed_tokens(self, tokens: torch.Tensor):
-        return self.pi0_paligemma.language_model.model.embed_tokens(tokens)
-
-    def prepare_inputs_for_generation(self, *args, **kwargs):
-        return self.pi0_paligemma.prepare_inputs_for_generation(*args, **kwargs)
-
-    def prepare_images(self, batch):
-        """Preprocess LeRobot batch into Pi0 inputs"""
-        images = []
-        img_masks = []
-        present_img_keys = [key for key in self.image_keys if key in batch]
-        if len(present_img_keys) == 0:
-            raise ValueError(
-                f"All image features are missing from the batch. At least one expected. (batch: {batch.keys()}) (image_features:{self.config.image_features})"
-            )
-
-        # Preprocess image features present in the batch
-        num_empty_cameras = 0
-        for key in self.image_keys:
-            if key in present_img_keys:
-                img = batch[key]
-
-                if self.config.resize_imgs_with_padding is not None:
-                    img = resize_with_pad(
-                        img,
-                        *self.config.resize_imgs_with_padding,
-                        pad_value=0,
-                        interpolate_like_pi=self.config.interpolate_like_pi,
-                    )
-
-                # Normalize from range [0,1] to [-1,1] as expected by siglip
-                img = img * 2.0 - 1.0
-
-                bsize = img.shape[0]
-                device = img.device
-                mask = torch.ones(bsize, dtype=torch.bool, device=device)
-            else:
-                if num_empty_cameras >= self.config.empty_cameras:
-                    continue
-                img = torch.ones_like(img) * -1
-                bsize = img.shape[0]
-                device = img.device
-                mask = torch.ones(bsize, dtype=torch.bool, device=device)
-                num_empty_cameras += 1
-
-            images.append(img)
-            img_masks.append(mask)
-        return images, img_masks
-
-    def normalize_actions(self, actions: torch.Tensor) -> torch.Tensor:
-        mins = actions.amin(dim=(1, 2), keepdim=True)  # [0]
-        maxs = actions.amax(dim=(1, 2), keepdim=True)  # [0]
-        return 2 * (actions - mins) / (maxs - mins + 1e-8) - 1
-
-    def _act_tokens_to_paligemma_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
-        out = self.paligemma_tokenizer.vocab_size - 1 - self.fast_skip_tokens - tokens
-        return out
-
-    def fast_tokenizer_wrapper(self, actions_norm):
-        """
-        A wrapper for self.fast_tokenizer that ensures batch processing,
-        conversion to PyTorch tensors, and returns a dictionary without padding.
-        """
-        batch_tokens = self.fast_tokenizer(actions_norm)
-        fast_out = self.processor.tokenizer.pad({"input_ids": batch_tokens}, return_tensors="pt")
-
-        return fast_out
-
-    def create_token_type_ids(self, padded_mask: torch.Tensor, prefix_len: int) -> torch.Tensor:
-        token_type_ids = torch.zeros_like(padded_mask, dtype=torch.bool)
-        # Compute cumulative sum mask
-        cumsum_mask = (padded_mask != 0).cumsum(dim=1)
-        # Suffix block (everything after prefix_len)
-        suffix_mask = cumsum_mask > prefix_len
-        token_type_ids = suffix_mask
-        return token_type_ids
-
-    def create_input_tokens(self, state, lang_text, actions=None):
-        bsize = state.shape[0]
-        device = state.device
-        bins = torch.linspace(-1, 1, 256 + 1, device=device)[:-1]
-        discretized = torch.bucketize(state, bins) - 1
-        discretized = discretized[:, :32]
-
-        prefix_texts = []
-        state_text = []
-        for txt, disc in zip(lang_text, discretized, strict=False):
-            cleaned = txt.lower().strip().replace("_", " ")
-            state_str = " ".join(str(val.item()) for val in disc)
-            prefix_texts.append(f"Task: {cleaned}, State: {state_str};\n")
-            state_text.append(f"State: {state_str};\n")
-
-        prefix_out = self.paligemma_tokenizer(
-            prefix_texts, add_special_tokens=True, return_tensors="pt", padding="longest", truncation=False
-        )
-        prefix_ids = prefix_out["input_ids"].to(device)
-        prefix_mask = prefix_out["attention_mask"].to(device)
-        prefix_lens = prefix_mask.sum(dim=1)[:, None].cpu()
-
-        if actions is not None:
-            actions_norm = self.normalize_actions(actions)
-            actions_pad = F.pad(
-                actions_norm, (0, max(0, self.config.max_action_dim - actions_norm.shape[2])), value=0
-            )[:, :, : self.config.max_action_dim]
-            fast_out = self.fast_tokenizer_wrapper(
-                actions_pad.cpu(),
-            )
-            act_ids = fast_out["input_ids"]
-            act_mask = fast_out["attention_mask"].to(device)
-
-            act_ids = self._act_tokens_to_paligemma_tokens(act_ids).to(device)
-            # Replace action with 0 to pad tokens
-            act_ids = torch.where(
-                act_ids == self.paligemma_tokenizer.vocab_size - 1 - self.fast_skip_tokens,
-                self.pad_token_id,
-                act_ids,
-            )
-
-            eos_token = torch.tensor(
-                [self.paligemma_tokenizer.eos_token_id], dtype=torch.long, device=device
-            ).expand(bsize, -1)
-            eos_mask = torch.tensor([1], dtype=torch.long, device=device).expand(bsize, -1)
-            bos = self.paligemma_tokenizer("Action: ", add_special_tokens=False, return_tensors="pt")
-            bos_token = bos["input_ids"].expand(act_ids.shape[0], -1).to(device)
-            bos_mask = bos["attention_mask"].expand(act_ids.shape[0], -1).to(device)
-            act_ids = torch.cat([bos_token, act_ids, eos_token], dim=1)
-            act_mask = torch.cat([bos_mask, act_mask, eos_mask], dim=1)
-            act_mask = act_mask.to(device)
-        else:
-            act_ids = torch.empty(bsize, self.pad_token_id, dtype=torch.long, device=device)
-            act_mask = torch.empty(bsize, 0, dtype=torch.long, device=device)
-        final_ids = torch.cat([prefix_ids, act_ids], dim=1)
-
-        final_mask = torch.cat([prefix_mask, act_mask], dim=1)
-        batch_inputs = {"input_ids": final_ids.tolist(), "attention_mask": final_mask.tolist()}
-
-        # Use tokenizer pad function
-        padded_output = self.paligemma_tokenizer.pad(
-            batch_inputs, padding="longest", max_length=180, return_tensors="pt"
-        )
-        padded_mask = padded_output["attention_mask"]
-
-        # define tensor of padding lengths
-        att_mask = (padded_mask != 0).cumsum(dim=1) > prefix_lens
-
-        token_type_ids = self.create_token_type_ids(padded_mask=padded_mask, prefix_len=prefix_lens)
-
-        padded_output["padded_mask"] = padded_output.pop("attention_mask")
-        padded_output["attention_mask"] = att_mask
-        # loss is computed not on prefix, and not on padding
-        padded_output["loss_mask"] = att_mask & padded_output["padded_mask"]
-        padded_output["token_type_ids"] = token_type_ids
-        return padded_output
-
-    def shift_padding_side(
-        self,
-        tokens: torch.Tensor,
-        ar_mask: torch.Tensor,
-        padding_mask: torch.Tensor,
-        loss_mask: torch.Tensor,
-        targets: torch.Tensor,
-        token_type_ids: torch.Tensor,
-        padding_side: str = "right",
-    ) -> tuple[torch.Tensor]:
-        if padding_side not in ["right", "left"]:
-            return tokens, ar_mask, padding_mask, loss_mask, targets, token_type_ids
-
-        new_tokens = torch.empty_like(tokens)
-        new_ar_masks = torch.empty_like(ar_mask)
-        new_padding_mask = torch.empty_like(padding_mask)
-        new_loss_mask = torch.empty_like(loss_mask)
-        new_targets = torch.empty_like(targets)
-        new_token_type_ids = torch.empty_like(token_type_ids)
-        batch_size = tokens.shape[0]
-        for i in range(batch_size):
-            padding_indices = torch.where(padding_mask[i] == 0)[0]
-            non_padding_indices = torch.where(padding_mask[i] == 1)[0]
-            if padding_side == "left":
-                new_indices = torch.cat((padding_indices, non_padding_indices), dim=0)
-            else:
-                new_indices = torch.cat((non_padding_indices, padding_indices), dim=0)
-            new_tokens[i] = tokens[i].index_select(0, new_indices)
-            new_ar_masks[i] = ar_mask[i].index_select(0, new_indices)
-            new_padding_mask[i] = padding_mask[i].index_select(0, new_indices)
-            new_loss_mask[i] = loss_mask[i].index_select(0, new_indices)
-            new_targets[i] = targets[i].index_select(0, new_indices)
-            new_token_type_ids[i] = token_type_ids[i].index_select(0, new_indices)
-
-        return new_tokens, new_ar_masks, new_padding_mask, new_loss_mask, new_targets, new_token_type_ids
-
-    def forward(self, batch: dict[str, Tensor]):
-        device = batch[OBS_STATE].device
-        # TODO: keep like this or move to the policy .forward
-        images, img_masks = self.prepare_images(batch)
-
-        padded_outs = self.create_input_tokens(
-            state=batch[OBS_STATE],
-            lang_text=batch["task"],
-            actions=batch[ACTION],
-        )
-
-        embs, pad_masks, _, targets, loss_mask, token_type_ids = self.embed_inputs(
-            images,
-            img_masks,
-            padded_outs["input_ids"],
-            padded_outs["padded_mask"],
-            padded_outs["attention_mask"],
-            padded_outs["loss_mask"],
-            padded_outs["token_type_ids"],
-            padding_side=self.padding_side,
-        )
-        position_ids = torch.cumsum(pad_masks, dim=1) - 1
-        token_type_ids = token_type_ids.to(dtype=torch.int64)
-        past_seen_tokens = 0
-        cache_position = torch.arange(past_seen_tokens, past_seen_tokens + embs.shape[1], device=embs.device)
-        pad_masks = block_causal_update_causal_mask(
-            attention_mask=pad_masks,
-            past_key_values=None,
-            cache_position=cache_position,
-            input_tensor=embs,
-            token_type_ids=token_type_ids,
-            dtype=self.pi0_paligemma.dtype,
-            attn_implementation=self.pi0_paligemma.config.text_config._attn_implementation,
-        )
-        outputs = self.pi0_paligemma.forward(
-            input_ids=None,
-            token_type_ids=None,
-            attention_mask=pad_masks,
-            position_ids=position_ids,
-            past_key_values=None,
-            inputs_embeds=embs,
-            use_cache=False,
-            labels=None,
-        )
-
-        logits = outputs.logits
-
-        loss_fct = nn.CrossEntropyLoss(reduction="none")
-
-        # Shift left for next-step prediction
-        logits = logits[:, :-1, :]
-        targets = targets[:, 1:].to(device)  # Shift targets
-        loss_mask = loss_mask[:, 1:].to(device)  # Ensure correct shape
-
-        # Compute per-token loss
-        token_loss = loss_fct(logits.reshape(-1, logits.shape[-1]), targets.reshape(-1))
-
-        # Apply loss mask
-        token_loss = token_loss * loss_mask.reshape(-1)
-
-        # Compute final loss
-        loss = token_loss.sum() / torch.clamp(loss_mask.sum(), min=1)
-
-        # Return loss dictionary
-        loss_dict = {"ce_loss": loss.item(), "loss": loss}
-        return loss_dict
-
-    def decode_actions_with_fast(
-        self,
-        tokens: list[list[int]],
-        *,
-        time_horizon: int | None = None,
-        action_dim: int | None = None,
-        relaxed_decoding: bool = True,
-    ) -> np.array:
-        """
-        Adapt original decoding in FAST to always return actions instead of zeros.
-        """
-        self.time_horizon = (
-            time_horizon or self.fast_tokenizer.time_horizon or self.fast_tokenizer.called_time_horizon
-        )
-        self.action_dim = (
-            action_dim or self.fast_tokenizer.action_dim or self.fast_tokenizer.called_action_dim
-        )
-
-        # Cache the time horizon and action dimension for the next call
-        self.called_time_horizon = self.time_horizon
-        self.called_action_dim = self.action_dim
-
-        assert self.time_horizon is not None and self.action_dim is not None, (
-            "Tokenizer not initialized, call encode() once or pass in time_horizon and action_dim."
-        )
-
-        decoded_actions = []
-        for token in tokens:
-            try:
-                decoded_tokens = self.fast_tokenizer.bpe_tokenizer.decode(token)
-                decoded_dct_coeff = np.array(list(map(ord, decoded_tokens))) + self.fast_tokenizer.min_token
-                if relaxed_decoding:
-                    # Expected sequence length
-                    expected_seq_len = self.time_horizon * self.action_dim
-                    diff = expected_seq_len - decoded_dct_coeff.shape[0]
-                    # Apply truncation if too long
-                    if diff < 0:
-                        decoded_dct_coeff = decoded_dct_coeff[:expected_seq_len]  # Truncate on the right
-                    # Apply padding if too short
-                    elif diff > 0:
-                        decoded_dct_coeff = np.pad(
-                            decoded_dct_coeff, (0, diff), mode="constant", constant_values=0
-                        )
-
-                decoded_dct_coeff = decoded_dct_coeff.reshape(-1, self.action_dim)
-                assert decoded_dct_coeff.shape == (
-                    self.time_horizon,
-                    self.action_dim,
-                ), (
-                    f"Decoded DCT coefficients have shape {decoded_dct_coeff.shape}, expected ({self.time_horizon}, {self.action_dim})"
-                )
-            except Exception as e:
-                print(f"Error decoding tokens: {e}")
-                print(f"Tokens: {token}")
-                decoded_dct_coeff = np.zeros((self.time_horizon, self.action_dim))
-            decoded_actions.append(idct(decoded_dct_coeff / self.fast_tokenizer.scale, axis=0, norm="ortho"))
-        return np.stack(decoded_actions)
-
-    def extract_actions(self, tokens: torch.Tensor, action_horizon: int, action_dim: int) -> torch.Tensor:
-        """
-        Extracts actions from predicted output tokens using the FAST model.
-
-        Args:
-            tokens (torch.Tensor): The input tensor of tokenized outputs.
-            action_horizon (int): The number of timesteps for actions.
-            action_dim (int): The dimensionality of each action.
-
-        Returns:
-            torch.Tensor: The extracted actions as a tensor of shape (action_horizon, action_dim).
-        """
-        # Decode predicted output tokens
-        decoded_tokens = self.paligemma_tokenizer.batch_decode(tokens, skip_special_tokens=True)
-        cleaned_tokens = [
-            tokens_sequence.replace("Action:", "").replace(":", "").strip().split("|")[0].strip()
-            for tokens_sequence in decoded_tokens
-        ]
-        raw_action_tokens = [
-            self.processor.tokenizer.encode(sample_tokens, return_tensors="pt", padding=False)
-            for sample_tokens in cleaned_tokens
-        ]  # something like this should be robust #looks good
-        action_tokens = [
-            self._act_tokens_to_paligemma_tokens(raw_action_token) for raw_action_token in raw_action_tokens
-        ]
-        # returns the tensor of decoded actions per sample in a list
-        decoded_actions = [
-            torch.tensor(
-                self.decode_actions_with_fast(
-                    tok.tolist(),
-                    time_horizon=action_horizon,
-                    action_dim=action_dim,
-                    relaxed_decoding=self.config.relaxed_action_decoding,
-                ),
-                device=tokens.device,
-            ).squeeze(0)
-            for tok in action_tokens
-        ]
-
-        return torch.stack(
-            decoded_actions,
-            dim=0,
-        )
-
-    def generate_actions(self, batch: dict[str, Tensor]):
-        # TODO: keep like this or move to the policy .forward
-        images, img_masks = self.prepare_images(batch)
-
-        padded_outs = self.create_input_tokens(state=batch[OBS_STATE], lang_text=batch["task"], actions=None)
-        embs, pad_masks, att_masks2, targets, loss_mask, token_type_ids = self.embed_inputs(
-            images,
-            img_masks,
-            padded_outs["input_ids"],
-            padded_outs["padded_mask"],
-            padded_outs["attention_mask"],
-            padded_outs["loss_mask"],
-            padded_outs["token_type_ids"],
-            padding_side="left",
-        )
-        token_type_ids = token_type_ids.to(dtype=torch.int64)
-        prefix_position_ids = torch.cumsum(pad_masks, dim=1) - 1
-        output_tokens = self.pi0_paligemma.generate(
-            input_ids=None,
-            attention_mask=pad_masks,
-            position_ids=prefix_position_ids,
-            past_key_values=None,
-            inputs_embeds=embs,
-            use_cache=self.config.use_cache,
-            max_new_tokens=self.config.max_decoding_steps,
-            do_sample=False,
-            num_beams=1,
-            token_type_ids=token_type_ids,
-        )
-        actions = self.extract_actions(output_tokens, self.action_horizon, self.action_dim)
-        return actions
-
-    def embed_image(self, image: torch.Tensor):
-        # Handle different transformers versions
-        if hasattr(self.pi0_paligemma, "get_image_features"):
-            return self.pi0_paligemma.get_image_features(image)
-        else:
-            return self.pi0_paligemma.model.get_image_features(image)
-
-    def embed_inputs(
-        self,
-        images,
-        img_masks,
-        tokens,
-        pad_mask,
-        ar_mask,
-        loss_mask,
-        token_type_ids,
-        padding_side: str = "right",
-    ):
-        # TODO: avoid list in python and torch.cat ; prefer pre-allocation with torch.empty
-        # images are a list of same size
-        # vectorizing everything!
-        device = images[0].device
-        image_embedding_dim = images[0].shape[-1]  # TODO should be from self.config
-        all_images = torch.stack(images, dim=1).to(device)
-        b, n, c, h, w = all_images.shape
-        all_images = all_images.view(b * n, c, h, w)
-        embedded = self.embed_image(all_images).to(device)
-        b_n, p, image_embedding_dim = embedded.shape  # Extract current dimensions
-        m = b_n // b  # Compute the number of images per sample dynamically
-
-        # Reshape dynamically
-        embedded = embedded.view(b, m, p, image_embedding_dim)
-        tokens_embs = self.embed_tokens(tokens.to(device))
-
-        img_masks = torch.stack(img_masks, dim=1).unsqueeze(-1).to(device)
-        num_img_emb = embedded.shape[2]
-        img_pad_masks = img_masks.repeat(1, 1, num_img_emb).view(b, -1)
-        img_att_masks = torch.zeros((b, n, num_img_emb), dtype=torch.long, device=device).reshape(b, -1)
-
-        image_target_tokens = (
-            torch.ones((b, n, num_img_emb), dtype=torch.long, device=device) * self.pad_token_id
-        ).reshape(b, -1)
-        image_loss_mask = torch.zeros((b, n, num_img_emb), dtype=torch.long, device=device).reshape(b, -1)
-
-        embedded = embedded.reshape(b, n * num_img_emb, image_embedding_dim)  # Shape: (B, N*P, D)
-
-        embs = torch.cat([embedded, tokens_embs], dim=1).to(device)
-        pad_masks = torch.cat([img_pad_masks, pad_mask.to(device)], dim=1)
-        att_masks = torch.cat([img_att_masks, ar_mask.to(device)], dim=1)
-        loss_masks = torch.cat([image_loss_mask, loss_mask.to(device)], dim=1)
-        targets = torch.cat([image_target_tokens, tokens.to(device)], dim=1)
-        token_type_ids = torch.cat([img_att_masks, token_type_ids.to(device)], dim=1)
-
-        # Shift pad tokens to the left (.generate()) or right (.train())
-        embs, att_masks, pad_masks, loss_masks, targets, token_type_ids = self.shift_padding_side(
-            embs, att_masks, pad_masks, loss_masks, targets, token_type_ids, padding_side=padding_side
-        )
-
-        targets = torch.where(targets == self.pad_token_id, self.ignore_index, targets)
-        return embs, pad_masks, att_masks, targets, loss_masks, token_type_ids
-
-
-def resize_with_pad(img, width, height, pad_value=0, interpolate_like_pi=True):
-    # assume no-op when width height fits already
-    if img.ndim != 4:
-        raise ValueError(f"(b,c,h,w) expected, but {img.shape}")
-
-    cur_height, cur_width = img.shape[2:]
-
-    ratio = max(cur_width / width, cur_height / height)
-    resized_height = int(cur_height / ratio)
-    resized_width = int(cur_width / ratio)
-
-    if interpolate_like_pi:
-        img = (img * 255.0).to(dtype=torch.uint8)
-        img = img.permute(0, 2, 3, 1)
-        original_device = img.device
-        img = img.to(device="cpu").numpy()
-        imgs = []
-        for sub_img in img:
-            sub_img = Image.fromarray(sub_img)
-            resized_img = sub_img.resize((resized_width, resized_height), resample=2)
-            resized_img = torch.from_numpy(np.array(resized_img))
-            imgs.append(resized_img)
-        img = torch.stack(imgs, dim=0)
-        img = img.permute(0, 3, 1, 2)
-        resized_img = img.to(device=original_device, dtype=torch.float32) / 255.0
-    else:
-        resized_img = F.interpolate(
-            img, size=(resized_height, resized_width), mode="bilinear", align_corners=False
-        )
-
-    pad_height = max(0, int(height - resized_height))
-    pad_width = max(0, int(width - resized_width))
-
-    # pad on left and top of image
-    padded_img = F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value)
-    return padded_img
@@ -1,92 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 Physical Intelligence and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Any
-
-import torch
-
-from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
-from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-)
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-
-
-def make_pi0fast_pre_post_processors(
-    config: PI0FASTConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """
-    Constructs pre-processor and post-processor pipelines for the PI0Fast policy.
-
-    The pre-processing pipeline prepares input data for the model by:
-    1. Renaming features to match pretrained configurations.
-    2. Normalizing input and output features based on dataset statistics.
-    3. Adding a batch dimension.
-    4. Moving all data to the specified device.
-
-    The post-processing pipeline handles the model's output by:
-    1. Moving data to the CPU.
-    2. Unnormalizing the output features to their original scale.
-
-    Args:
-        config: The configuration object for the PI0Fast policy.
-        dataset_stats: A dictionary of statistics for normalization.
-        preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
-        postprocessor_kwargs: Additional arguments for the post-processor pipeline.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-    ]
-    output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-        DeviceProcessorStep(device="cpu"),
-    ]
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )
@@ -18,7 +18,7 @@ import os
 from importlib.resources import files
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TypeVar
+from typing import TypedDict, TypeVar

 import packaging
 import safetensors
@@ -27,6 +27,7 @@ from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from huggingface_hub.errors import HfHubHTTPError
 from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor
 from torch import Tensor, nn
+from typing_extensions import Unpack

 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.train import TrainPipelineConfig
@@ -36,6 +37,10 @@ from lerobot.utils.hub import HubMixin
 T = TypeVar("T", bound="PreTrainedPolicy")


+class ActionSelectKwargs(TypedDict, total=False):
+    noise: Tensor | None
+
+
 class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
    """
    Base class for policy models.
@@ -181,7 +186,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
        raise NotImplementedError

    @abc.abstractmethod
-    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
        """Returns the action chunk (for action chunking policies) for a given observation, potentially in batch mode.

        Child classes using action chunking should use this method within `select_action` to form the action chunk
@@ -190,7 +195,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
        raise NotImplementedError

    @abc.abstractmethod
-    def select_action(self, batch: dict[str, Tensor]) -> Tensor:
+    def select_action(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
        """Return one action to run in the environment (potentially in batch mode).

        When the model uses a history of observations, or outputs a sequence of actions, this method deals
@@ -16,10 +16,16 @@

 import logging
 from collections import deque
+from typing import Any

+import numpy as np
 import torch
 from torch import nn

+from lerobot.datasets.utils import build_dataset_frame
+from lerobot.processor import PolicyAction, RobotAction, RobotObservation
+from lerobot.utils.constants import ACTION, OBS_STR
+

 def populate_queues(
    queues: dict[str, deque], batch: dict[str, torch.Tensor], exclude_keys: list[str] | None = None
@@ -85,3 +91,110 @@ def log_model_loading_keys(missing_keys: list[str], unexpected_keys: list[str])
        logging.warning(f"Missing key(s) when loading model: {missing_keys}")
    if unexpected_keys:
        logging.warning(f"Unexpected key(s) when loading model: {unexpected_keys}")
+
+
+# TODO(Steven): Move this function to a proper preprocessor step
+def prepare_observation_for_inference(
+    observation: dict[str, np.ndarray],
+    device: torch.device,
+    task: str | None = None,
+    robot_type: str | None = None,
+) -> RobotObservation:
+    """Converts observation data to model-ready PyTorch tensors.
+
+    This function takes a dictionary of NumPy arrays, performs necessary
+    preprocessing, and prepares it for model inference. The steps include:
+    1. Converting NumPy arrays to PyTorch tensors.
+    2. Normalizing and permuting image data (if any).
+    3. Adding a batch dimension to each tensor.
+    4. Moving all tensors to the specified compute device.
+    5. Adding task and robot type information to the dictionary.
+
+    Args:
+        observation: A dictionary mapping observation names (str) to NumPy
+            array data. For images, the format is expected to be (H, W, C).
+        device: The PyTorch device (e.g., 'cpu' or 'cuda') to which the
+            tensors will be moved.
+        task: An optional string identifier for the current task.
+        robot_type: An optional string identifier for the robot being used.
+
+    Returns:
+        A dictionary where values are PyTorch tensors preprocessed for
+        inference, residing on the target device. Image tensors are reshaped
+        to (C, H, W) and normalized to a [0, 1] range.
+    """
+    for name in observation:
+        observation[name] = torch.from_numpy(observation[name])
+        if "image" in name:
+            observation[name] = observation[name].type(torch.float32) / 255
+            observation[name] = observation[name].permute(2, 0, 1).contiguous()
+        observation[name] = observation[name].unsqueeze(0)
+        observation[name] = observation[name].to(device)
+
+    observation["task"] = task if task else ""
+    observation["robot_type"] = robot_type if robot_type else ""
+
+    return observation
+
+
+def build_inference_frame(
+    observation: dict[str, Any],
+    device: torch.device,
+    ds_features: dict[str, dict],
+    task: str | None = None,
+    robot_type: str | None = None,
+) -> RobotObservation:
+    """Constructs a model-ready observation tensor dict from a raw observation.
+
+    This utility function orchestrates the process of converting a raw,
+    unstructured observation from an environment into a structured,
+    tensor-based format suitable for passing to a policy model.
+
+    Args:
+        observation: The raw observation dictionary, which may contain
+            superfluous keys.
+        device: The target PyTorch device for the final tensors.
+        ds_features: A configuration dictionary that specifies which features
+            to extract from the raw observation.
+        task: An optional string identifier for the current task.
+        robot_type: An optional string identifier for the robot being used.
+
+    Returns:
+        A dictionary of preprocessed tensors ready for model inference.
+    """
+    # Extracts the correct keys from the incoming raw observation
+    observation = build_dataset_frame(ds_features, observation, prefix=OBS_STR)
+
+    # Performs the necessary conversions to the observation
+    observation = prepare_observation_for_inference(observation, device, task, robot_type)
+
+    return observation
+
+
+def make_robot_action(action_tensor: PolicyAction, ds_features: dict[str, dict]) -> RobotAction:
+    """Converts a policy's output tensor into a dictionary of named actions.
+
+    This function translates the numerical output from a policy model into a
+    human-readable and robot-consumable format, where each dimension of the
+    action tensor is mapped to a named motor or actuator command.
+
+    Args:
+        action_tensor: A PyTorch tensor representing the policy's action,
+            typically with a batch dimension (e.g., shape [1, action_dim]).
+        ds_features: A configuration dictionary containing metadata, including
+            the names corresponding to each index of the action tensor.
+
+    Returns:
+        A dictionary mapping action names (e.g., "joint_1_motor") to their
+        corresponding floating-point values, ready to be sent to a robot
+        controller.
+    """
+    # TODO(Steven): Check if these steps are already in all postprocessor policies
+    action_tensor = action_tensor.squeeze(0)
+    action_tensor = action_tensor.to("cpu")
+
+    action_names = ds_features[ACTION]["names"]
+    act_processed_policy: RobotAction = {
+        f"{name}": float(action_tensor[i]) for i, name in enumerate(action_names)
+    }
+    return act_processed_policy
@@ -303,6 +303,65 @@ def clean_state_dict(
    return new_state_dict


+def load_state_dict_with_missing_key_handling(
+    policy: torch.nn.Module,
+    state_dict: dict[str, torch.Tensor],
+    policy_type: str,
+    known_missing_keys_whitelist: dict[str, list[str]],
+) -> list[str]:
+    """
+    Load state dict into policy with graceful handling of missing keys.
+
+    This function loads the state dict with strict=False, filters out whitelisted
+    missing keys, and provides detailed reporting about any issues found.
+
+    Args:
+        policy: The policy model to load the state dict into.
+        state_dict: The cleaned state dictionary to load.
+        policy_type: The type of policy (used for whitelist lookup).
+        known_missing_keys_whitelist: Dictionary mapping policy types to lists of
+                                     known acceptable missing keys.
+
+    Returns:
+        List of problematic missing keys that weren't in the whitelist.
+    """
+    # Load the cleaned state dict with strict=False to capture missing/unexpected keys
+    load_result = policy.load_state_dict(state_dict, strict=False)
+
+    # Check for missing keys
+    missing_keys = load_result.missing_keys
+    unexpected_keys = load_result.unexpected_keys
+
+    # Filter out whitelisted missing keys
+    policy_type_lower = policy_type.lower()
+    whitelisted_keys = known_missing_keys_whitelist.get(policy_type_lower, [])
+    problematic_missing_keys = [key for key in missing_keys if key not in whitelisted_keys]
+
+    if missing_keys:
+        if problematic_missing_keys:
+            print(f"WARNING: Found {len(problematic_missing_keys)} unexpected missing keys:")
+            for key in problematic_missing_keys:
+                print(f"   - {key}")
+
+        if len(missing_keys) > len(problematic_missing_keys):
+            whitelisted_missing = [key for key in missing_keys if key in whitelisted_keys]
+            print(f"INFO: Found {len(whitelisted_missing)} expected missing keys (whitelisted):")
+            for key in whitelisted_missing:
+                print(f"   - {key}")
+
+    if unexpected_keys:
+        print(f"WARNING: Found {len(unexpected_keys)} unexpected keys:")
+        for key in unexpected_keys:
+            print(f"   - {key}")
+
+    if not missing_keys and not unexpected_keys:
+        print("Successfully loaded cleaned state dict into policy model (all keys matched)")
+    else:
+        print("State dict loaded with some missing/unexpected keys (see details above)")
+
+    return problematic_missing_keys
+
+
 def convert_features_to_policy_features(features_dict: dict[str, dict]) -> dict[str, PolicyFeature]:
    """
    Converts a feature dictionary from the old config format to the new `PolicyFeature` format.
@@ -336,9 +395,45 @@ def convert_features_to_policy_features(features_dict: dict[str, dict]) -> dict[
    return converted_features


+def display_migration_summary_with_warnings(problematic_missing_keys: list[str]) -> None:
+    """
+    Display final migration summary with warnings about problematic missing keys.
+
+    Args:
+        problematic_missing_keys: List of missing keys that weren't in the whitelist.
+    """
+    if not problematic_missing_keys:
+        return
+
+    print("\n" + "=" * 60)
+    print("IMPORTANT: MIGRATION COMPLETED WITH WARNINGS")
+    print("=" * 60)
+    print(
+        f"The migration was successful, but {len(problematic_missing_keys)} unexpected missing keys were found:"
+    )
+    print()
+    for key in problematic_missing_keys:
+        print(f"   - {key}")
+    print()
+    print("These missing keys may indicate:")
+    print("  • The model architecture has changed")
+    print("  • Some components were not properly saved in the original model")
+    print("  • The migration script needs to be updated for this policy type")
+    print()
+    print("What to do next:")
+    print("  1. Test your migrated model carefully to ensure it works as expected")
+    print("  2. If you encounter issues, please open an issue at:")
+    print("     https://github.com/huggingface/lerobot/issues")
+    print("  3. Include this migration log and the missing keys listed above")
+    print()
+    print("If the model works correctly despite these warnings, the missing keys")
+    print("might be expected for your policy type and can be added to the whitelist.")
+    print("=" * 60)
+
+
 def load_model_from_hub(
    repo_id: str, revision: str | None = None
-) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, Any]]:
+) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, Any] | None]:
    """
    Downloads and loads a model's state_dict and configs from the Hugging Face Hub.

@@ -348,13 +443,12 @@ def load_model_from_hub(

    Returns:
        A tuple containing the model's state dictionary, the policy configuration,
-        and the training configuration.
+        and the training configuration (None if train_config.json is not found).
    """
    # Download files.
    safetensors_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", revision=revision)

    config_path = hf_hub_download(repo_id=repo_id, filename="config.json", revision=revision)
-    train_config_path = hf_hub_download(repo_id=repo_id, filename="train_config.json", revision=revision)

    # Load state_dict
    state_dict = load_safetensors(safetensors_path)
@@ -363,8 +457,14 @@ def load_model_from_hub(
    with open(config_path) as f:
        config = json.load(f)

-    with open(train_config_path) as f:
-        train_config = json.load(f)
+    # Try to load train_config (optional)
+    train_config = None
+    try:
+        train_config_path = hf_hub_download(repo_id=repo_id, filename="train_config.json", revision=revision)
+        with open(train_config_path) as f:
+            train_config = json.load(f)
+    except FileNotFoundError:
+        print("train_config.json not found - continuing without training configuration")

    return state_dict, config, train_config

@@ -410,8 +510,15 @@ def main():
        state_dict = load_safetensors(os.path.join(args.pretrained_path, "model.safetensors"))
        with open(os.path.join(args.pretrained_path, "config.json")) as f:
            config = json.load(f)
-        with open(os.path.join(args.pretrained_path, "train_config.json")) as f:
-            train_config = json.load(f)
+
+        # Try to load train_config (optional)
+        train_config = None
+        train_config_path = os.path.join(args.pretrained_path, "train_config.json")
+        if os.path.exists(train_config_path):
+            with open(train_config_path) as f:
+                train_config = json.load(f)
+        else:
+            print("train_config.json not found - continuing without training configuration")
    else:
        # Hub repository
        state_dict, config, train_config = load_model_from_hub(args.pretrained_path, args.revision)
@@ -488,10 +595,20 @@ def main():
    policy_class = get_policy_class(policy_type)
    policy = policy_class(policy_config)

-    # Load the cleaned state dict
-    policy.load_state_dict(new_state_dict, strict=True)
-    print("Successfully loaded cleaned state dict into policy model")
+    # Define whitelist of known missing keys that are acceptable (for example weight tie) for certain policy types
+    known_missing_keys_whitelist = {
+        "pi0": ["model.paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"],
+        # Add other policy types and their known missing keys here as needed
+    }

+    # Load state dict with graceful missing key handling
+    problematic_missing_keys = load_state_dict_with_missing_key_handling(
+        policy=policy,
+        state_dict=new_state_dict,
+        policy_type=policy_type,
+        known_missing_keys_whitelist=known_missing_keys_whitelist,
+    )
+    policy.to(torch.float32)
    # Create preprocessor and postprocessor using the factory
    print("Creating preprocessor and postprocessor using make_pre_post_processors...")
    preprocessor, postprocessor = make_pre_post_processors(policy_cfg=policy_config, dataset_stats=stats)
@@ -521,7 +638,9 @@ def main():
    # Generate and save model card
    print("Generating model card...")
    # Get metadata from original config
-    dataset_repo_id = train_config.get("repo_id", "unknown")
+    dataset_repo_id = "unknown"
+    if train_config is not None:
+        dataset_repo_id = train_config.get("repo_id", "unknown")
    license = config.get("license", "apache-2.0")

    tags = config.get("tags", ["robotics", "lerobot", policy_type]) or ["robotics", "lerobot", policy_type]
@@ -552,25 +671,25 @@ def main():

        if create_pr:
            # Separate commit description for PR body
-            commit_description = """🤖 **Automated Policy Migration to PolicyProcessorPipeline**
+            commit_description = """**Automated Policy Migration to PolicyProcessorPipeline**

 This PR migrates your model to the new LeRobot policy format using the modern PolicyProcessorPipeline architecture.

 ## What Changed

-### ✨ **New Architecture - PolicyProcessorPipeline**
+### **New Architecture - PolicyProcessorPipeline**
 Your model now uses external PolicyProcessorPipeline components for data processing instead of built-in normalization layers. This provides:
 - **Modularity**: Separate preprocessing and postprocessing pipelines
 - **Flexibility**: Easy to swap, configure, and debug processing steps
 - **Compatibility**: Works with the latest LeRobot ecosystem

-### 🔧 **Normalization Extraction**
+### **Normalization Extraction**
 We've extracted normalization statistics from your model's state_dict and removed the built-in normalization layers:
 - **Extracted patterns**: `normalize_inputs.*`, `unnormalize_outputs.*`, `normalize.*`, `unnormalize.*`, `input_normalizer.*`, `output_normalizer.*`
 - **Statistics preserved**: Mean, std, min, max values for all features
 - **Clean model**: State dict now contains only core model weights

-### 📦 **Files Added**
+### **Files Added**
 - **preprocessor_config.json**: Configuration for input preprocessing pipeline
 - **postprocessor_config.json**: Configuration for output postprocessing pipeline
 - **model.safetensors**: Clean model weights without normalization layers
@@ -578,13 +697,13 @@ We've extracted normalization statistics from your model's state_dict and remove
 - **train_config.json**: Training configuration
 - **README.md**: Updated model card with migration information

-### 🚀 **Benefits**
+### **Benefits**
 - **Backward Compatible**: Your model behavior remains identical
 - **Future Ready**: Compatible with latest LeRobot features and updates
 - **Debuggable**: Easy to inspect and modify processing steps
 - **Portable**: Processors can be shared and reused across models

-### 💻 **Usage**
+### **Usage**
 ```python
 # Load your migrated model
 from lerobot.policies import get_policy_class
@@ -642,6 +761,9 @@ final_action = postprocessor(action)
        else:
            print(f"\nView the changes at: https://huggingface.co/{hub_repo_id}")

+    # Display final summary about any problematic missing keys
+    display_migration_summary_with_warnings(problematic_missing_keys)
+

 if __name__ == "__main__":
    main()
@@ -281,8 +281,14 @@ class _NormalizationMixin:
        """
        Core logic to apply a normalization or unnormalization transformation to a tensor.

-        This method selects the appropriate normalization mode (e.g., mean/std, min/max)
-        based on the feature type and applies the corresponding mathematical operation.
+        This method selects the appropriate normalization mode based on the feature type
+        and applies the corresponding mathematical operation.
+
+        Normalization Modes:
+          - MEAN_STD: Centers data around zero with unit variance.
+          - MIN_MAX: Scales data to [-1, 1] range using actual min/max values.
+          - QUANTILES: Scales data to [-1, 1] range using 1st and 99th percentiles (q01/q99).
+          - QUANTILE10: Scales data to [-1, 1] range using 10th and 90th percentiles (q10/q90).

        Args:
            tensor: The input tensor to transform.
@@ -300,7 +306,12 @@ class _NormalizationMixin:
        if norm_mode == NormalizationMode.IDENTITY or key not in self._tensor_stats:
            return tensor

-        if norm_mode not in (NormalizationMode.MEAN_STD, NormalizationMode.MIN_MAX):
+        if norm_mode not in (
+            NormalizationMode.MEAN_STD,
+            NormalizationMode.MIN_MAX,
+            NormalizationMode.QUANTILES,
+            NormalizationMode.QUANTILE10,
+        ):
            raise ValueError(f"Unsupported normalization mode: {norm_mode}")

        # For Accelerate compatibility: Ensure stats are on the same device and dtype as the input tensor
@@ -311,7 +322,14 @@ class _NormalizationMixin:

        stats = self._tensor_stats[key]

-        if norm_mode == NormalizationMode.MEAN_STD and "mean" in stats and "std" in stats:
+        if norm_mode == NormalizationMode.MEAN_STD:
+            mean = stats.get("mean", None)
+            std = stats.get("std", None)
+            if mean is None or std is None:
+                raise ValueError(
+                    "MEAN_STD normalization mode requires mean and std stats, please update the dataset with the correct stats"
+                )
+
            mean, std = stats["mean"], stats["std"]
            # Avoid division by zero by adding a small epsilon.
            denom = std + self.eps
@@ -319,7 +337,14 @@ class _NormalizationMixin:
                return tensor * std + mean
            return (tensor - mean) / denom

-        if norm_mode == NormalizationMode.MIN_MAX and "min" in stats and "max" in stats:
+        if norm_mode == NormalizationMode.MIN_MAX:
+            min_val = stats.get("min", None)
+            max_val = stats.get("max", None)
+            if min_val is None or max_val is None:
+                raise ValueError(
+                    "MIN_MAX normalization mode requires min and max stats, please update the dataset with the correct stats"
+                )
+
            min_val, max_val = stats["min"], stats["max"]
            denom = max_val - min_val
            # When min_val == max_val, substitute the denominator with a small epsilon
@@ -334,6 +359,40 @@ class _NormalizationMixin:
            # Map from [min, max] to [-1, 1]
            return 2 * (tensor - min_val) / denom - 1

+        if norm_mode == NormalizationMode.QUANTILES:
+            q01 = stats.get("q01", None)
+            q99 = stats.get("q99", None)
+            if q01 is None or q99 is None:
+                raise ValueError(
+                    "QUANTILES normalization mode requires q01 and q99 stats, please update the dataset with the correct stats using the `augment_dataset_quantile_stats.py` script"
+                )
+
+            denom = q99 - q01
+            # Avoid division by zero by adding epsilon when quantiles are identical
+            denom = torch.where(
+                denom == 0, torch.tensor(self.eps, device=tensor.device, dtype=tensor.dtype), denom
+            )
+            if inverse:
+                return (tensor + 1.0) * denom / 2.0 + q01
+            return 2.0 * (tensor - q01) / denom - 1.0
+
+        if norm_mode == NormalizationMode.QUANTILE10:
+            q10 = stats.get("q10", None)
+            q90 = stats.get("q90", None)
+            if q10 is None or q90 is None:
+                raise ValueError(
+                    "QUANTILE10 normalization mode requires q10 and q90 stats, please update the dataset with the correct stats using the `augment_dataset_quantile_stats.py` script"
+                )
+
+            denom = q90 - q10
+            # Avoid division by zero by adding epsilon when quantiles are identical
+            denom = torch.where(
+                denom == 0, torch.tensor(self.eps, device=tensor.device, dtype=tensor.dtype), denom
+            )
+            if inverse:
+                return (tensor + 1.0) * denom / 2.0 + q10
+            return 2.0 * (tensor - q10) / denom - 1.0
+
        # If necessary stats are missing, return input unchanged.
        return tensor

@@ -1,3 +1,19 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from dataclasses import asdict, dataclass
 from typing import Any

@@ -607,6 +607,7 @@ class ReplayBuffer:
            lerobot_dataset.save_episode()

        lerobot_dataset.stop_image_writer()
+        lerobot_dataset.finalize()

        return lerobot_dataset

@@ -696,7 +696,7 @@ def control_loop(
            episode_idx += 1

            if dataset is not None:
-                if transition[TransitionKey.INFO].get("rerecord_episode", False):
+                if transition[TransitionKey.INFO].get(TeleopEvents.RERECORD_EPISODE, False):
                    logging.info(f"Re-recording episode {episode_idx}")
                    dataset.clear_episode_buffer()
                    episode_idx -= 1
@@ -1,3 +1,19 @@
+#!/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 import RobotConfig
 from .robot import Robot
 from .utils import make_robot_from_config
@@ -14,13 +14,16 @@

 import logging
 from pprint import pformat
+from typing import cast

-from lerobot.robots import RobotConfig
+from lerobot.utils.import_utils import make_device_from_device_class

+from .config import RobotConfig
 from .robot import Robot


 def make_robot_from_config(config: RobotConfig) -> Robot:
+    # TODO(Steven): Consider just using the make_device_from_device_class for all types
    if config.type == "koch_follower":
        from .koch_follower import KochFollower

@@ -66,7 +69,10 @@ def make_robot_from_config(config: RobotConfig) -> Robot:

        return MockRobot(config)
    else:
-        raise ValueError(config.type)
+        try:
+            return cast(Robot, make_device_from_device_class(config))
+        except Exception as e:
+            raise ValueError(f"Error creating robot with config {config}: {e}") from e


 # TODO(pepijn): Move to pipeline step to make sure we don't have to do this in the robot code and send action to robot is clean for use in dataset
@@ -52,6 +52,7 @@ from lerobot.teleoperators import (  # noqa: F401
    so100_leader,
    so101_leader,
 )
+from lerobot.utils.import_utils import register_third_party_devices
 from lerobot.utils.utils import init_logging


@@ -83,6 +84,7 @@ def calibrate(cfg: CalibrateConfig):


 def main():
+    register_third_party_devices()
    calibrate()


@@ -0,0 +1,286 @@
+#!/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.
+
+"""
+Edit LeRobot datasets using various transformation tools.
+
+This script allows you to delete episodes, split datasets, merge datasets,
+and remove features. When new_repo_id is specified, creates a new dataset.
+
+Usage Examples:
+
+Delete episodes 0, 2, and 5 from a dataset:
+    python -m lerobot.scripts.lerobot_edit_dataset \
+        --repo_id lerobot/pusht \
+        --operation.type delete_episodes \
+        --operation.episode_indices "[0, 2, 5]"
+
+Delete episodes and save to a new dataset:
+    python -m lerobot.scripts.lerobot_edit_dataset \
+        --repo_id lerobot/pusht \
+        --new_repo_id lerobot/pusht_filtered \
+        --operation.type delete_episodes \
+        --operation.episode_indices "[0, 2, 5]"
+
+Split dataset by fractions:
+    python -m lerobot.scripts.lerobot_edit_dataset \
+        --repo_id lerobot/pusht \
+        --operation.type split \
+        --operation.splits '{"train": 0.8, "val": 0.2}'
+
+Split dataset by episode indices:
+    python -m lerobot.scripts.lerobot_edit_dataset \
+        --repo_id lerobot/pusht \
+        --operation.type split \
+        --operation.splits '{"train": [0, 1, 2, 3], "val": [4, 5]}'
+
+Split into more than two splits:
+    python -m lerobot.scripts.lerobot_edit_dataset \
+        --repo_id lerobot/pusht \
+        --operation.type split \
+        --operation.splits '{"train": 0.6, "val": 0.2, "test": 0.2}'
+
+Merge multiple datasets:
+    python -m lerobot.scripts.lerobot_edit_dataset \
+        --repo_id lerobot/pusht_merged \
+        --operation.type merge \
+        --operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']"
+
+Remove camera feature:
+    python -m lerobot.scripts.lerobot_edit_dataset \
+        --repo_id lerobot/pusht \
+        --operation.type remove_feature \
+        --operation.feature_names "['observation.images.top']"
+
+Using JSON config file:
+    python -m lerobot.scripts.lerobot_edit_dataset \
+        --config_path path/to/edit_config.json
+"""
+
+import logging
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+
+from lerobot.configs import parser
+from lerobot.datasets.dataset_tools import (
+    delete_episodes,
+    merge_datasets,
+    remove_feature,
+    split_dataset,
+)
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.utils.constants import HF_LEROBOT_HOME
+from lerobot.utils.utils import init_logging
+
+
+@dataclass
+class DeleteEpisodesConfig:
+    type: str = "delete_episodes"
+    episode_indices: list[int] | None = None
+
+
+@dataclass
+class SplitConfig:
+    type: str = "split"
+    splits: dict[str, float | list[int]] | None = None
+
+
+@dataclass
+class MergeConfig:
+    type: str = "merge"
+    repo_ids: list[str] | None = None
+
+
+@dataclass
+class RemoveFeatureConfig:
+    type: str = "remove_feature"
+    feature_names: list[str] | None = None
+
+
+@dataclass
+class EditDatasetConfig:
+    repo_id: str
+    operation: DeleteEpisodesConfig | SplitConfig | MergeConfig | RemoveFeatureConfig
+    root: str | None = None
+    new_repo_id: str | None = None
+    push_to_hub: bool = False
+
+
+def get_output_path(repo_id: str, new_repo_id: str | None, root: Path | None) -> tuple[str, Path]:
+    if new_repo_id:
+        output_repo_id = new_repo_id
+        output_dir = root / new_repo_id if root else HF_LEROBOT_HOME / new_repo_id
+    else:
+        output_repo_id = repo_id
+        dataset_path = root / repo_id if root else HF_LEROBOT_HOME / repo_id
+        old_path = Path(str(dataset_path) + "_old")
+
+        if dataset_path.exists():
+            if old_path.exists():
+                shutil.rmtree(old_path)
+            shutil.move(str(dataset_path), str(old_path))
+
+        output_dir = dataset_path
+
+    return output_repo_id, output_dir
+
+
+def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
+    if not isinstance(cfg.operation, DeleteEpisodesConfig):
+        raise ValueError("Operation config must be DeleteEpisodesConfig")
+
+    if not cfg.operation.episode_indices:
+        raise ValueError("episode_indices must be specified for delete_episodes operation")
+
+    dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
+    output_repo_id, output_dir = get_output_path(
+        cfg.repo_id, cfg.new_repo_id, Path(cfg.root) if cfg.root else None
+    )
+
+    if cfg.new_repo_id is None:
+        dataset.root = Path(str(dataset.root) + "_old")
+
+    logging.info(f"Deleting episodes {cfg.operation.episode_indices} from {cfg.repo_id}")
+    new_dataset = delete_episodes(
+        dataset,
+        episode_indices=cfg.operation.episode_indices,
+        output_dir=output_dir,
+        repo_id=output_repo_id,
+    )
+
+    logging.info(f"Dataset saved to {output_dir}")
+    logging.info(f"Episodes: {new_dataset.meta.total_episodes}, Frames: {new_dataset.meta.total_frames}")
+
+    if cfg.push_to_hub:
+        logging.info(f"Pushing to hub as {output_repo_id}")
+        LeRobotDataset(output_repo_id, root=output_dir).push_to_hub()
+
+
+def handle_split(cfg: EditDatasetConfig) -> None:
+    if not isinstance(cfg.operation, SplitConfig):
+        raise ValueError("Operation config must be SplitConfig")
+
+    if not cfg.operation.splits:
+        raise ValueError(
+            "splits dict must be specified with split names as keys and fractions/episode lists as values"
+        )
+
+    dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
+
+    logging.info(f"Splitting dataset {cfg.repo_id} with splits: {cfg.operation.splits}")
+    split_datasets = split_dataset(dataset, splits=cfg.operation.splits)
+
+    for split_name, split_ds in split_datasets.items():
+        split_repo_id = f"{cfg.repo_id}_{split_name}"
+        logging.info(
+            f"{split_name}: {split_ds.meta.total_episodes} episodes, {split_ds.meta.total_frames} frames"
+        )
+
+        if cfg.push_to_hub:
+            logging.info(f"Pushing {split_name} split to hub as {split_repo_id}")
+            LeRobotDataset(split_ds.repo_id, root=split_ds.root).push_to_hub()
+
+
+def handle_merge(cfg: EditDatasetConfig) -> None:
+    if not isinstance(cfg.operation, MergeConfig):
+        raise ValueError("Operation config must be MergeConfig")
+
+    if not cfg.operation.repo_ids:
+        raise ValueError("repo_ids must be specified for merge operation")
+
+    if not cfg.repo_id:
+        raise ValueError("repo_id must be specified as the output repository for merged dataset")
+
+    logging.info(f"Loading {len(cfg.operation.repo_ids)} datasets to merge")
+    datasets = [LeRobotDataset(repo_id, root=cfg.root) for repo_id in cfg.operation.repo_ids]
+
+    output_dir = Path(cfg.root) / cfg.repo_id if cfg.root else HF_LEROBOT_HOME / cfg.repo_id
+
+    logging.info(f"Merging datasets into {cfg.repo_id}")
+    merged_dataset = merge_datasets(
+        datasets,
+        output_repo_id=cfg.repo_id,
+        output_dir=output_dir,
+    )
+
+    logging.info(f"Merged dataset saved to {output_dir}")
+    logging.info(
+        f"Episodes: {merged_dataset.meta.total_episodes}, Frames: {merged_dataset.meta.total_frames}"
+    )
+
+    if cfg.push_to_hub:
+        logging.info(f"Pushing to hub as {cfg.repo_id}")
+        LeRobotDataset(merged_dataset.repo_id, root=output_dir).push_to_hub()
+
+
+def handle_remove_feature(cfg: EditDatasetConfig) -> None:
+    if not isinstance(cfg.operation, RemoveFeatureConfig):
+        raise ValueError("Operation config must be RemoveFeatureConfig")
+
+    if not cfg.operation.feature_names:
+        raise ValueError("feature_names must be specified for remove_feature operation")
+
+    dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
+    output_repo_id, output_dir = get_output_path(
+        cfg.repo_id, cfg.new_repo_id, Path(cfg.root) if cfg.root else None
+    )
+
+    if cfg.new_repo_id is None:
+        dataset.root = Path(str(dataset.root) + "_old")
+
+    logging.info(f"Removing features {cfg.operation.feature_names} from {cfg.repo_id}")
+    new_dataset = remove_feature(
+        dataset,
+        feature_names=cfg.operation.feature_names,
+        output_dir=output_dir,
+        repo_id=output_repo_id,
+    )
+
+    logging.info(f"Dataset saved to {output_dir}")
+    logging.info(f"Remaining features: {list(new_dataset.meta.features.keys())}")
+
+    if cfg.push_to_hub:
+        logging.info(f"Pushing to hub as {output_repo_id}")
+        LeRobotDataset(output_repo_id, root=output_dir).push_to_hub()
+
+
+@parser.wrap()
+def edit_dataset(cfg: EditDatasetConfig) -> None:
+    operation_type = cfg.operation.type
+
+    if operation_type == "delete_episodes":
+        handle_delete_episodes(cfg)
+    elif operation_type == "split":
+        handle_split(cfg)
+    elif operation_type == "merge":
+        handle_merge(cfg)
+    elif operation_type == "remove_feature":
+        handle_remove_feature(cfg)
+    else:
+        raise ValueError(
+            f"Unknown operation type: {operation_type}\n"
+            f"Available operations: delete_episodes, split, merge, remove_feature"
+        )
+
+
+def main() -> None:
+    init_logging()
+    edit_dataset()
+
+
+if __name__ == "__main__":
+    main()
@@ -180,9 +180,15 @@ def rollout(
            render_callback(env)

        # VectorEnv stores is_success in `info["final_info"][env_index]["is_success"]`. "final_info" isn't
-        # available of none of the envs finished.
+        # available if none of the envs finished.
        if "final_info" in info:
-            successes = [info["is_success"] if info is not None else False for info in info["final_info"]]
+            final_info = info["final_info"]
+            if not isinstance(final_info, dict):
+                raise RuntimeError(
+                    "Unsupported `final_info` format: expected dict (Gymnasium >= 1.0). "
+                    "You're likely using an older version of gymnasium (< 1.0). Please upgrade."
+                )
+            successes = final_info["is_success"].tolist()
        else:
            successes = [False] * env.num_envs

@@ -79,6 +79,7 @@ from lerobot.datasets.utils import build_dataset_frame, combine_feature_dicts
 from lerobot.datasets.video_utils import VideoEncodingManager
 from lerobot.policies.factory import make_policy, make_pre_post_processors
 from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.utils import make_robot_action
 from lerobot.processor import (
    PolicyAction,
    PolicyProcessorPipeline,
@@ -117,6 +118,7 @@ from lerobot.utils.control_utils import (
    sanity_check_dataset_name,
    sanity_check_dataset_robot_compatibility,
 )
+from lerobot.utils.import_utils import register_third_party_devices
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import (
    get_safe_torch_device,
@@ -315,10 +317,7 @@ def record_loop(
                robot_type=robot.robot_type,
            )

-            action_names = dataset.features[ACTION]["names"]
-            act_processed_policy: RobotAction = {
-                f"{name}": float(action_values[i]) for i, name in enumerate(action_names)
-            }
+            act_processed_policy: RobotAction = make_robot_action(action_values, dataset.features)

        elif policy is None and isinstance(teleop, Teleoperator):
            act = teleop.get_action()
@@ -513,6 +512,7 @@ def record(cfg: RecordConfig) -> LeRobotDataset:


 def main():
+    register_third_party_devices()
    record()


@@ -61,6 +61,7 @@ from lerobot.robots import (  # noqa: F401
    so101_follower,
 )
 from lerobot.utils.constants import ACTION
+from lerobot.utils.import_utils import register_third_party_devices
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import (
    init_logging,
@@ -126,6 +127,7 @@ def replay(cfg: ReplayConfig):


 def main():
+    register_third_party_devices()
    replay()


@@ -88,6 +88,7 @@ from lerobot.teleoperators import (  # noqa: F401
    so100_leader,
    so101_leader,
 )
+from lerobot.utils.import_utils import register_third_party_devices
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import init_logging, move_cursor_up
 from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
@@ -215,6 +216,7 @@ def teleoperate(cfg: TeleoperateConfig):


 def main():
+    register_third_party_devices()
    teleoperate()


@@ -180,21 +180,33 @@ def train(cfg: TrainPipelineConfig):

    # Create processors - only provide dataset_stats if not resuming from saved processors
    processor_kwargs = {}
-    if not (cfg.resume and cfg.policy.pretrained_path):
+    postprocessor_kwargs = {}
+    if (cfg.policy.pretrained_path and not cfg.resume) or not cfg.policy.pretrained_path:
        # Only provide dataset_stats when not resuming from saved processor state
        processor_kwargs["dataset_stats"] = dataset.meta.stats

    if cfg.policy.pretrained_path is not None:
        processor_kwargs["preprocessor_overrides"] = {
            "device_processor": {"device": device.type},
-            "normalizer_processor": {"stats": dataset.meta.stats},
+            "normalizer_processor": {
+                "stats": dataset.meta.stats,
+                "features": {**policy.config.input_features, **policy.config.output_features},
+                "norm_map": policy.config.normalization_mapping,
+            },
        }
-        processor_kwargs["postprocessor_overrides"] = {
-            "unnormalizer_processor": {"stats": dataset.meta.stats},
+        postprocessor_kwargs["postprocessor_overrides"] = {
+            "unnormalizer_processor": {
+                "stats": dataset.meta.stats,
+                "features": policy.config.output_features,
+                "norm_map": policy.config.normalization_mapping,
+            },
        }

    preprocessor, postprocessor = make_pre_post_processors(
-        policy_cfg=cfg.policy, pretrained_path=cfg.policy.pretrained_path, **processor_kwargs
+        policy_cfg=cfg.policy,
+        pretrained_path=cfg.policy.pretrained_path,
+        **processor_kwargs,
+        **postprocessor_kwargs,
    )

    logging.info("Creating optimizer and scheduler")
@@ -13,6 +13,9 @@
 # limitations under the License.

 from enum import Enum
+from typing import cast
+
+from lerobot.utils.import_utils import make_device_from_device_class

 from .config import TeleoperatorConfig
 from .teleoperator import Teleoperator
@@ -29,6 +32,7 @@ class TeleopEvents(Enum):


 def make_teleoperator_from_config(config: TeleoperatorConfig) -> Teleoperator:
+    # TODO(Steven): Consider just using the make_device_from_device_class for all types
    if config.type == "keyboard":
        from .keyboard import KeyboardTeleop

@@ -82,4 +86,7 @@ def make_teleoperator_from_config(config: TeleoperatorConfig) -> Teleoperator:

        return Reachy2Teleoperator(config)
    else:
-        raise ValueError(config.type)
+        try:
+            return cast(Teleoperator, make_device_from_device_class(config))
+        except Exception as e:
+            raise ValueError(f"Error creating robot with config {config}: {e}") from e
@@ -20,9 +20,25 @@
 {% elif model_name == "vqbet" %}
 [VQ-BET](https://huggingface.co/papers/2403.03181) combines vector-quantised action tokens with Behaviour Transformers to discretise control and achieve data-efficient imitation across diverse skills.
 {% elif model_name == "pi0" %}
-[Pi0](https://huggingface.co/papers/2410.24164) is a generalist vision-language-action transformer that converts multimodal observations and text instructions into robot actions for zero-shot task transfer.
-{% elif model_name == "pi0fast" %}
-[Pi0-Fast](https://huggingface.co/papers/2501.09747) is a variant of Pi0 that uses a new tokenization method called FAST, which enables training of an autoregressive vision-language-action policy for high-frequency robotic tasks with improved performance and reduced training time.
+**π₀ (Pi0)**
+
+π₀ is a Vision-Language-Action model for general robot control, from Physical Intelligence. The LeRobot implementation is adapted from their open source OpenPI repository.
+
+**Model Overview**
+
+π₀ represents a breakthrough in robotics as the first general-purpose robot foundation model developed by Physical Intelligence. Unlike traditional robots that are narrow specialists programmed for repetitive motions, π₀ is designed to be a generalist policy that can understand visual inputs, interpret natural language instructions, and control a variety of different robots across diverse tasks.
+
+For more details, see the [Physical Intelligence π₀ blog post](https://www.physicalintelligence.company/blog/pi0).
+{% elif model_name == "pi05" %}
+**π₀.₅ (Pi05) Policy**
+
+π₀.₅ is a Vision-Language-Action model with open-world generalization, from Physical Intelligence. The LeRobot implementation is adapted from their open source OpenPI repository.
+
+**Model Overview**
+
+π₀.₅ represents a significant evolution from π₀, developed by Physical Intelligence to address a big challenge in robotics: open-world generalization. While robots can perform impressive tasks in controlled environments, π₀.₅ is designed to generalize to entirely new environments and situations that were never seen during training.
+
+For more details, see the [Physical Intelligence π₀.₅ blog post](https://www.physicalintelligence.company/blog/pi05).
 {% elif model_name == "sac" %}
 [Soft Actor-Critic (SAC)](https://huggingface.co/papers/1801.01290) is an entropy-regularised actor-critic algorithm offering stable, sample-efficient learning in continuous-control environments.
 {% elif model_name == "reward_classifier" %}
--- a/Show More
+++ b/Show More