pi0 hack

* reshape noise after taking it as output of the network

.github/workflows/full_tests.yml

@@ -78,7 +78,7 @@ jobs:
           python-version: ${{ env.PYTHON_VERSION }}
 
       - name: Install lerobot with all extras
-        run: uv sync --all-extras --no-extra groot # TODO(Steven): Make flash-attn optional
+        run: uv sync --all-extras
 
       - name: Run pytest (all extras)
         run: uv run pytest tests -vv --maxfail=10

.github/workflows/nightly.yml

@@ -119,7 +119,6 @@ jobs:
       TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton
     container:
       image: ${{ needs.build-docker-cpu-nightly.outputs.image_tag }} # zizmor: ignore[unpinned-images]
-      options: --shm-size "16gb"
       credentials:
         username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
         password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
@@ -159,36 +158,3 @@ jobs:
         run: pytest tests -vv --maxfail=10
       - name: Run end-to-end tests
         run: make test-end-to-end
-
-  # This job runs multi-GPU training tests with 4 GPUs
-  nightly-multi-gpu-tests:
-    name: Nightly Multi-GPU Tests
-    needs: [build-docker-gpu-nightly]
-    runs-on:
-      group: aws-g4dn-12xlarge  # Instance with 4 GPUs
-    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
-      CUDA_VISIBLE_DEVICES: "0,1,2,3"
-    container:
-      image: ${{ needs.build-docker-gpu-nightly.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: Verify GPU availability
-        run: |
-          nvidia-smi
-          python -c "import torch; print(f'PyTorch CUDA available: {torch.cuda.is_available()}'); print(f'Number of GPUs: {torch.cuda.device_count()}')"
-
-      - name: Run multi-GPU training tests
-      # TODO(Steven): Investigate why motors tests are failing in multi-GPU setup
-        run: pytest tests -vv --maxfail=10 --ignore=tests/motors/
-        timeout-minutes: 10

.github/workflows/release.yml

@@ -82,14 +82,6 @@ jobs:
             exit 1
           fi
 
-      - name: Remove Tags with Git dependencies
-        # TODO(Steven): Temporary patch to remove pi from PyPi 0.4.0 release due to its reliance on git dependencies.
-        run: |
-          echo "::info:: Checking for Git dependencies to remove from pyproject.toml..."
-          grep -E '@ git\+https|lerobot\[pi\]' pyproject.toml | sed 's/^/::warning:: Removing line: /' || true
-          sed -E -i '/@ git\+https|lerobot\[pi\]/d' pyproject.toml
-          echo "::info:: Git dependencies removed. Proceeding with build."
-
       - name: Install build dependencies
         run: python -m pip install build
 
@@ -111,7 +103,7 @@ jobs:
       - name: Publish to TestPyPI for pre-releases
         # True for tags like 'v0.2.0-rc1'
         if: startsWith(github.ref, 'refs/tags/v') && contains(github.ref, '-')
-        uses: pypa/gh-action-pypi-publish@v1.13.0 # zizmor: ignore[unpinned-uses, use-trusted-publishing]
+        uses: pypa/gh-action-pypi-publish@v1.12.4 # zizmor: ignore[unpinned-uses, use-trusted-publishing]
         with:
           repository-url: https://test.pypi.org/legacy/
           verbose: true
@@ -119,7 +111,7 @@ jobs:
 
       - name: Publish to PyPI
         if: startsWith(github.ref, 'refs/tags/v') && !contains(github.ref, '-')
-        uses: pypa/gh-action-pypi-publish@v1.13.0 # zizmor: ignore[unpinned-uses, use-trusted-publishing]
+        uses: pypa/gh-action-pypi-publish@v1.12.4 # zizmor: ignore[unpinned-uses, use-trusted-publishing]
         with:
           verbose: true
           print-hash: true
@@ -146,7 +138,7 @@ jobs:
       - name: Setup uv and Python
         uses: astral-sh/setup-uv@v6 # zizmor: ignore[unpinned-uses]
         with:
-          enable-cache: true # zizmor: ignore[cache-poisoning]
+          enable-cache: true
           version: ${{ env.UV_VERSION }}
           python-version: ${{ env.PYTHON_VERSION }}
       - name: Create uv virtual environment

.github/workflows/stale.yml

@@ -27,17 +27,15 @@ env:
     This issue was closed because it has been stalled for 14 days with no activity.
     Feel free to reopen if is still relevant, or to ping a collaborator if you have any questions.
   CLOSE_PR_MESSAGE: >
-    This PR was closed because it has been stalled for 21 days with no activity.
+    This PR was closed because it has been stalled for 14 days with no activity.
     Feel free to reopen if is still relevant, or to ping a collaborator if you have any questions.
   WARN_ISSUE_MESSAGE: >
     This issue has been automatically marked as stale because it has not had
     recent activity (6 months). It will be closed if no further activity occurs.
-    Any change, comment or update to this issue will reset this count.
     Thank you for your contributions.
   WARN_PR_MESSAGE: >
     This PR has been automatically marked as stale because it has not had
-    recent activity (1 year). It will be closed if no further activity occurs.
-    Any change, comment or update to this PR will reset this count.
+    recent activity (6 months). It will be closed if no further activity occurs.
     Thank you for your contributions.
 
 jobs:
@@ -58,10 +56,10 @@ jobs:
           stale-pr-label: stale
           exempt-issue-labels: never-stale
           exempt-pr-labels: never-stale
-          days-before-issue-stale: 180
+          days-before-issue-stale: 180 # TODO(Steven): Will modify this to 90 after initial cleanup
           days-before-issue-close: 14
-          days-before-pr-stale: 365
-          days-before-pr-close: 21
+          days-before-pr-stale: 180
+          days-before-pr-close: 14
           delete-branch: true
           close-issue-message: ${{ env.CLOSE_ISSUE_MESSAGE }}
           close-pr-message: ${{ env.CLOSE_PR_MESSAGE }}

.github/workflows/unbound_deps_tests.yml

@@ -70,7 +70,7 @@ jobs:
           echo "Dependencies unbound:" && cat pyproject.toml
 
       - name: Install lerobot with all extras
-        run: uv sync --all-extras --no-extra groot # TODO(Steven): Make flash-attn optional
+        run: uv sync --all-extras
 
       - name: Run pytest (all extras)
         run: uv run pytest tests -vv

.pre-commit-config.yaml

@@ -26,7 +26,7 @@ repos:
 
    ##### General Code Quality & Formatting #####
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v6.0.0
+    rev: v5.0.0
     hooks:
       - id: check-added-large-files
         args: ['--maxkb=1024']
@@ -39,20 +39,20 @@ repos:
       - id: trailing-whitespace
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.14.1
+    rev: v0.12.4
     hooks:
       - id: ruff-format
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]
 
   - repo: https://github.com/adhtruong/mirrors-typos
-    rev: v1.38.1
+    rev: v1.34.0
     hooks:
       - id: typos
         args: [--force-exclude]
 
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.21.0
+    rev: v3.20.0
     hooks:
     -   id: pyupgrade
         args: [--py310-plus]
@@ -68,12 +68,12 @@ repos:
 
   ##### Security #####
   - repo: https://github.com/gitleaks/gitleaks
-    rev: v8.28.0
+    rev: v8.27.2
     hooks:
       - id: gitleaks
 
   - repo: https://github.com/woodruffw/zizmor-pre-commit
-    rev: v1.15.2
+    rev: v1.11.0
     hooks:
       - id: zizmor
 
@@ -87,7 +87,7 @@ repos:
   # TODO(Steven): Uncomment when ready to use
   ##### Static Analysis & Typing #####
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.18.2
+    rev: v1.16.0
     hooks:
       - id: mypy
         args: [--config-file=pyproject.toml]

CONTRIBUTING.md

@@ -72,6 +72,7 @@ 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.
 
@@ -137,7 +138,7 @@ Follow these steps to start contributing:
 4. for development, we advise to use a tool like `poetry` or `uv` instead of just `pip` to easily track our dependencies.
    Follow the instructions to [install poetry](https://python-poetry.org/docs/#installation) (use a version >=2.1.0) or to [install uv](https://docs.astral.sh/uv/getting-started/installation/#installation-methods) if you don't have one of them already.
 
-   Set up a development environment with conda:
+   Set up a development environment with conda or miniconda:
 
    ```bash
    conda create -y -n lerobot-dev python=3.10 && conda activate lerobot-dev

Makefile

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

README.md

@@ -104,14 +104,14 @@ LeRobot works with Python 3.10+ and PyTorch 2.2+.
 
 ### Environment Setup
 
-Create a virtual environment with Python 3.10 and activate it, e.g. with [`miniforge`](https://conda-forge.org/download/):
+Create a virtual environment with Python 3.10 and activate it, e.g. with [`miniconda`](https://docs.anaconda.com/free/miniconda/index.html):
 
 ```bash
 conda create -y -n lerobot python=3.10
 conda activate lerobot
 ```
 
-When using `conda`, install `ffmpeg` in your environment:
+When using `miniconda`, install `ffmpeg` in your environment:
 
 ```bash
 conda install ffmpeg -c conda-forge
@@ -185,11 +185,6 @@ _Replace `[...]` with your desired features._
 For a full list of optional dependencies, see:
 https://pypi.org/project/lerobot/
 
-> [!NOTE]
-> For lerobot 0.4.0, if you want to install pi tags, you will have to do: `pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git"`.
->
-> This will be solved in the next patch release
-
 ### Weights & Biases
 
 To use [Weights and Biases](https://docs.wandb.ai/quickstart) for experiment tracking, log in with
@@ -212,13 +207,13 @@ lerobot-dataset-viz \
     --episode-index 0
 ```
 
-or from a dataset in a local folder with the `root` option and the `--mode local` (in the following case the dataset will be searched for in `./my_local_data_dir/lerobot/pusht`)
+or from a dataset in a local folder with the `root` option and the `--local-files-only` (in the following case the dataset will be searched for in `./my_local_data_dir/lerobot/pusht`)
 
 ```bash
 lerobot-dataset-viz \
     --repo-id lerobot/pusht \
     --root ./my_local_data_dir \
-    --mode local \
+    --local-files-only 1 \
     --episode-index 0
 ```
 
@@ -315,7 +310,7 @@ To upload these to the hub, run the following:
 huggingface-cli upload ${hf_user}/${repo_name} path/to/pretrained_model
 ```
 
-See [lerobot_eval.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/lerobot_eval.py) for an example of how other people may use your policy.
+See [eval.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/eval.py) for an example of how other people may use your policy.
 
 ### Acknowledgment
 
@@ -342,3 +337,7 @@ If you want, you can cite this work with:
 ## Star History
 
 [![Star History Chart](https://api.star-history.com/svg?repos=huggingface/lerobot&type=Timeline)](https://star-history.com/#huggingface/lerobot&Timeline)
+
+```
+
+```

docs/source/_toctree.yml

@@ -7,6 +7,8 @@
 - sections:
   - local: il_robots
     title: Imitation Learning for Robots
+  - local: il_sim
+    title: Imitation Learning in Sim
   - local: cameras
     title: Cameras
   - local: integrate_hardware
@@ -17,8 +19,6 @@
     title: Train RL in Simulation
   - local: async
     title: Use Async Inference
-  - local: multi_gpu_training
-    title: Multi GPU training
   title: "Tutorials"
 - sections:
   - local: lerobot-dataset-v3
@@ -37,19 +37,9 @@
     title: π₀ (Pi0)
   - local: pi05
     title: π₀.₅ (Pi05)
-  - local: groot
-    title: NVIDIA GR00T N1.5
-  title: "Policies"
-- sections:
-  - local: envhub
-    title: Environments from the Hub
-  - local: il_sim
-    title: Imitation Learning in Sim
   - local: libero
     title: Using Libero
-  - local: metaworld
-    title: Using MetaWorld
-  title: "Simulation"
+  title: "Policies"
 - sections:
   - local: introduction_processors
     title: Introduction to Robot Processors

docs/source/envhub.mdx

@@ -1,424 +0,0 @@
-# Loading Environments from the Hub
-
-The **EnvHub** feature allows you to load simulation environments directly from the Hugging Face Hub with a single line of code. This unlocks a powerful new model for collaboration: instead of environments being locked away inside monolithic libraries, anyone can publish custom environments and share them with the community.
-
-## Overview
-
-With EnvHub, you can:
-
-- Load environments from the Hub instantly
-- Share your custom simulation tasks with the community
-- Version control your environments using Git
-- Distribute complex physics simulations without packaging hassles
-
-## Quick Start
-
-Loading an environment from the Hub is as simple as:
-
-```python
-from lerobot.envs.factory import make_env
-
-# Load a hub environment (requires explicit consent to run remote code)
-env = make_env("lerobot/cartpole-env", trust_remote_code=True)
-```
-
-<Tip warning={true}>
-  **Security Notice**: Loading environments from the Hub executes Python code
-  from third-party repositories. Only use `trust_remote_code=True` with
-  repositories you trust. We strongly recommend pinning to a specific commit
-  hash for reproducibility and security.
-</Tip>
-
-## What is EnvHub?
-
-EnvHub is a framework that allows researchers and developers to:
-
-1. **Publish environments** to the Hugging Face Hub as Git repositories
-2. **Load environments** dynamically without installing them as packages
-3. **Version and track** environment changes using Git semantics
-4. **Discover** new simulation tasks shared by the community
-
-This design means you can go from discovering an interesting environment on the Hub to running experiments in seconds, without worrying about dependency conflicts or complex installation procedures.
-
-## Repository Structure
-
-To make your environment loadable from the Hub, your repository must contain at minimum:
-
-### Required Files
-
-**`env.py`** (or custom Python file)
-
-- Must expose a `make_env(n_envs: int, use_async_envs: bool)` function
-- This function should return one of:
-  - A `gym.vector.VectorEnv` (most common)
-  - A single `gym.Env` (will be automatically wrapped)
-  - A dict mapping `{suite_name: {task_id: VectorEnv}}` (for multi-task benchmarks)
-
-### Optional Files
-
-**`requirements.txt`**
-
-- List any additional dependencies your environment needs
-- Users will need to install these manually before loading your environment
-
-**`README.md`**
-
-- Document your environment: what task it implements, observation/action spaces, rewards, etc.
-- Include usage examples and any special setup instructions
-
-**`.gitignore`**
-
-- Exclude unnecessary files from your repository
-
-### Example Repository Structure
-
-```
-my-environment-repo/
-├── env.py                 # Main environment definition (required)
-├── requirements.txt       # Dependencies (optional)
-├── README.md             # Documentation (recommended)
-├── assets/               # Images, videos, etc. (optional)
-│   └── demo.gif
-└── configs/              # Config files if needed (optional)
-    └── task_config.yaml
-```
-
-## Creating Your Environment Repository
-
-### Step 1: Define Your Environment
-
-Create an `env.py` file with a `make_env` function:
-
-```python
-# env.py
-import gymnasium as gym
-
-def make_env(n_envs: int = 1, use_async_envs: bool = False):
-    """
-    Create vectorized environments for your custom task.
-
-    Args:
-        n_envs: Number of parallel environments
-        use_async_envs: Whether to use AsyncVectorEnv or SyncVectorEnv
-
-    Returns:
-        gym.vector.VectorEnv or dict mapping suite names to vectorized envs
-    """
-    def _make_single_env():
-        # Create your custom environment
-        return gym.make("CartPole-v1")
-
-    # Choose vector environment type
-    env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
-
-    # Create vectorized environment
-    vec_env = env_cls([_make_single_env for _ in range(n_envs)])
-
-    return vec_env
-```
-
-### Step 2: Test Locally
-
-Before uploading, test your environment locally:
-
-```python
-from lerobot.envs.utils import _load_module_from_path, _call_make_env, _normalize_hub_result
-
-# Load your module
-module = _load_module_from_path("./env.py")
-
-# Test the make_env function
-result = _call_make_env(module, n_envs=2, use_async_envs=False)
-normalized = _normalize_hub_result(result)
-
-# Verify it works
-suite_name = next(iter(normalized))
-env = normalized[suite_name][0]
-obs, info = env.reset()
-print(f"Observation shape: {obs.shape if hasattr(obs, 'shape') else type(obs)}")
-env.close()
-```
-
-### Step 3: Upload to the Hub
-
-Upload your repository to Hugging Face:
-
-```bash
-# Install huggingface_hub if needed
-pip install huggingface_hub
-
-# Login to Hugging Face
-huggingface-cli login
-
-# Create a new repository
-huggingface-cli repo create my-custom-env --type space --org my-org
-
-# Initialize git and push
-git init
-git add .
-git commit -m "Initial environment implementation"
-git remote add origin https://huggingface.co/my-org/my-custom-env
-git push -u origin main
-```
-
-Alternatively, use the `huggingface_hub` Python API:
-
-```python
-from huggingface_hub import HfApi
-
-api = HfApi()
-
-# Create repository
-api.create_repo("my-custom-env", repo_type="space")
-
-# Upload files
-api.upload_folder(
-    folder_path="./my-env-folder",
-    repo_id="username/my-custom-env",
-    repo_type="space",
-)
-```
-
-## Loading Environments from the Hub
-
-### Basic Usage
-
-```python
-from lerobot.envs.factory import make_env
-
-# Load from the hub
-envs_dict = make_env(
-    "username/my-custom-env",
-    n_envs=4,
-    trust_remote_code=True
-)
-
-# Access the environment
-suite_name = next(iter(envs_dict))
-env = envs_dict[suite_name][0]
-
-# Use it like any gym environment
-obs, info = env.reset()
-action = env.action_space.sample()
-obs, reward, terminated, truncated, info = env.step(action)
-```
-
-### Advanced: Pinning to Specific Versions
-
-For reproducibility and security, pin to a specific Git revision:
-
-```python
-# Pin to a specific branch
-env = make_env("username/my-env@main", trust_remote_code=True)
-
-# Pin to a specific commit (recommended for papers/experiments)
-env = make_env("username/my-env@abc123def456", trust_remote_code=True)
-
-# Pin to a tag
-env = make_env("username/my-env@v1.0.0", trust_remote_code=True)
-```
-
-### Custom File Paths
-
-If your environment definition is not in `env.py`:
-
-```python
-# Load from a custom file
-env = make_env("username/my-env:custom_env.py", trust_remote_code=True)
-
-# Combine with version pinning
-env = make_env("username/my-env@v1.0:envs/task_a.py", trust_remote_code=True)
-```
-
-### Async Environments
-
-For better performance with multiple environments:
-
-```python
-envs_dict = make_env(
-    "username/my-env",
-    n_envs=8,
-    use_async_envs=True,  # Use AsyncVectorEnv for parallel execution
-    trust_remote_code=True
-)
-```
-
-## URL Format Reference
-
-The hub URL format supports several patterns:
-
-| Pattern              | Description                    | Example                                |
-| -------------------- | ------------------------------ | -------------------------------------- |
-| `user/repo`          | Load `env.py` from main branch | `make_env("lerobot/pusht-env")`        |
-| `user/repo@revision` | Load from specific revision    | `make_env("lerobot/pusht-env@main")`   |
-| `user/repo:path`     | Load custom file               | `make_env("lerobot/envs:pusht.py")`    |
-| `user/repo@rev:path` | Revision + custom file         | `make_env("lerobot/envs@v1:pusht.py")` |
-
-## Multi-Task Environments
-
-For benchmarks with multiple tasks (like LIBERO), return a nested dictionary:
-
-```python
-def make_env(n_envs: int = 1, use_async_envs: bool = False):
-    env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
-
-    # Return dict: {suite_name: {task_id: VectorEnv}}
-    return {
-        "suite_1": {
-            0: env_cls([lambda: gym.make("Task1-v0") for _ in range(n_envs)]),
-            1: env_cls([lambda: gym.make("Task2-v0") for _ in range(n_envs)]),
-        },
-        "suite_2": {
-            0: env_cls([lambda: gym.make("Task3-v0") for _ in range(n_envs)]),
-        }
-    }
-```
-
-## Security Considerations
-
-<Tip warning={true}>
-  **Important**: The `trust_remote_code=True` flag is required to execute
-  environment code from the Hub. This is by design for security.
-</Tip>
-
-When loading environments from the Hub:
-
-1. **Review the code first**: Visit the repository and inspect `env.py` before loading
-2. **Pin to commits**: Use specific commit hashes for reproducibility
-3. **Check dependencies**: Review `requirements.txt` for suspicious packages
-4. **Use trusted sources**: Prefer official organizations or well-known researchers
-5. **Sandbox if needed**: Run untrusted code in isolated environments (containers, VMs)
-
-Example of safe usage:
-
-```python
-# ❌ BAD: Loading without inspection
-env = make_env("random-user/untrusted-env", trust_remote_code=True)
-
-# ✅ GOOD: Review code, then pin to specific commit
-# 1. Visit https://huggingface.co/trusted-org/verified-env
-# 2. Review the env.py file
-# 3. Copy the commit hash
-env = make_env("trusted-org/verified-env@a1b2c3d4", trust_remote_code=True)
-```
-
-## Example: CartPole from the Hub
-
-Here's a complete example using the reference CartPole environment:
-
-```python
-from lerobot.envs.factory import make_env
-import numpy as np
-
-# Load the environment
-envs_dict = make_env("lerobot/cartpole-env", n_envs=4, trust_remote_code=True)
-
-# Get the vectorized environment
-suite_name = next(iter(envs_dict))
-env = envs_dict[suite_name][0]
-
-# Run a simple episode
-obs, info = env.reset()
-done = np.zeros(env.num_envs, dtype=bool)
-total_reward = np.zeros(env.num_envs)
-
-while not done.all():
-    # Random policy
-    action = env.action_space.sample()
-    obs, reward, terminated, truncated, info = env.step(action)
-    total_reward += reward
-    done = terminated | truncated
-
-print(f"Average reward: {total_reward.mean():.2f}")
-env.close()
-```
-
-## Benefits of EnvHub
-
-### For Environment Authors
-
-- **Easy distribution**: No PyPI packaging required
-- **Version control**: Use Git for environment versioning
-- **Rapid iteration**: Push updates instantly
-- **Documentation**: Hub README renders beautifully
-- **Community**: Reach LeRobot users directly
-
-### For Researchers
-
-- **Quick experiments**: Load any environment in one line
-- **Reproducibility**: Pin to specific commits
-- **Discovery**: Browse environments on the Hub
-- **No conflicts**: No need to install conflicting packages
-
-### For the Community
-
-- **Growing ecosystem**: More diverse simulation tasks
-- **Standardization**: Common `make_env` API
-- **Collaboration**: Fork and improve existing environments
-- **Accessibility**: Lower barrier to sharing research
-
-## Troubleshooting
-
-### "Refusing to execute remote code"
-
-You must explicitly pass `trust_remote_code=True`:
-
-```python
-env = make_env("user/repo", trust_remote_code=True)
-```
-
-### "Module X not found"
-
-The hub environment has dependencies you need to install:
-
-```bash
-# Check the repo's requirements.txt and install dependencies
-pip install gymnasium numpy
-```
-
-### "make_env not found in module"
-
-Your `env.py` must expose a `make_env` function:
-
-```python
-def make_env(n_envs: int, use_async_envs: bool):
-    # Your implementation
-    pass
-```
-
-### Environment returns wrong type
-
-The `make_env` function must return:
-
-- A `gym.vector.VectorEnv`, or
-- A single `gym.Env`, or
-- A dict `{suite_name: {task_id: VectorEnv}}`
-
-## Best Practices
-
-1. **Document your environment**: Include observation/action space descriptions, reward structure, and termination conditions in your README
-2. **Add requirements.txt**: List all dependencies with versions
-3. **Test thoroughly**: Verify your environment works locally before pushing
-4. **Use semantic versioning**: Tag releases with version numbers
-5. **Add examples**: Include usage examples in your README
-6. **Keep it simple**: Minimize dependencies when possible
-7. **License your work**: Add a LICENSE file to clarify usage terms
-
-## Future Directions
-
-The EnvHub ecosystem enables exciting possibilities:
-
-- **GPU-accelerated physics**: Share Isaac Gym or Brax environments
-- **Photorealistic rendering**: Distribute environments with advanced graphics
-- **Multi-agent scenarios**: Complex interaction tasks
-- **Real-world simulators**: Digital twins of physical setups
-- **Procedural generation**: Infinite task variations
-- **Domain randomization**: Pre-configured DR pipelines
-
-As more researchers and developers contribute, the diversity and quality of available environments will grow, benefiting the entire robotics learning community.
-
-## See Also
-
-- [Hugging Face Hub Documentation](https://huggingface.co/docs/hub/en/index)
-- [Gymnasium Documentation](https://gymnasium.farama.org/index.html)
-- [Example Hub Environment](https://huggingface.co/lerobot/cartpole-env)

docs/source/groot.mdx

@@ -1,125 +0,0 @@
-# GR00T N1.5 Policy
-
-GR00T N1.5 is an open foundation model from NVIDIA designed for generalized humanoid robot reasoning and skills. It is a cross-embodiment model that accepts multimodal input, including language and images, to perform manipulation tasks in diverse environments.
-
-This document outlines the specifics of its integration and usage within the LeRobot framework.
-
-## Model Overview
-
-NVIDIA Isaac GR00T N1.5 is an upgraded version of the GR00T N1 foundation model. It is built to improve generalization and language-following abilities for humanoid robots.
-
-Developers and researchers can post-train GR00T N1.5 with their own real or synthetic data to adapt it for specific humanoid robots or tasks.
-
-GR00T N1.5 (specifically the GR00T-N1.5-3B model) is built using pre-trained vision and language encoders. It utilizes a flow matching action transformer to model a chunk of actions, conditioned on vision, language, and proprioception.
-
-Its strong performance comes from being trained on an expansive and diverse humanoid dataset, which includes:
-
-- Real captured data from robots.
-- Synthetic data generated using NVIDIA Isaac GR00T Blueprint.
-- Internet-scale video data.
-
-This approach allows the model to be highly adaptable through post-training for specific embodiments, tasks, and environments.
-
-## Installation Requirements
-
-As of today, GR00T N1.5 requires flash attention for it's internal working.
-
-We are working on making this optional, but in the meantime that means that we require an extra installation step and it can only be used in CUDA enabled devices.
-
-1. Following the Environment Setup of our [Installation Guide](./installation). **Attention** don't install `lerobot` in this step.
-2. Install [Flash Attention](https://github.com/Dao-AILab/flash-attention) by running:
-
-```bash
-# Check https://pytorch.org/get-started/locally/ for your system
-pip install "torch>=2.2.1,<2.8.0" "torchvision>=0.21.0,<0.23.0" # --index-url https://download.pytorch.org/whl/cu1XX
-pip install ninja "packaging>=24.2,<26.0" # flash attention dependencies
-pip install "flash-attn>=2.5.9,<3.0.0" --no-build-isolation
-python -c "import flash_attn; print(f'Flash Attention {flash_attn.__version__} imported successfully')"
-```
-
-3. Install LeRobot by running:
-
-```bash
-pip install lerobot[groot]
-```
-
-## Usage
-
-To use GR00T in your LeRobot configuration, specify the policy type as:
-
-```python
-policy.type=groot
-```
-
-## Training
-
-### Training Command Example
-
-Here's a complete training command for finetuning the base GR00T model on your own dataset:
-
-```bash
-# Using a multi-GPU setup
-accelerate launch \
-  --multi_gpu \
-  --num_processes=$NUM_GPUS \
-  $(which lerobot-train) \
-  --output_dir=$OUTPUT_DIR \
-  --save_checkpoint=true \
-  --batch_size=$BATCH_SIZE \
-  --steps=$NUM_STEPS \
-  --save_freq=$SAVE_FREQ \
-  --log_freq=$LOG_FREQ \
-  --policy.push_to_hub=true \
-  --policy.type=groot \
-  --policy.repo_id=$REPO_ID \
-  --policy.tune_diffusion_model=false \
-  --dataset.repo_id=$DATASET_ID \
-  --wandb.enable=true \
-  --wandb.disable_artifact=true \
-  --job_name=$JOB_NAME
-```
-
-## Performance Results
-
-### Libero Benchmark Results
-
-> [!NOTE]
-> Follow our instructions for Libero usage: [Libero](./libero)
-
-GR00T has demonstrated strong performance on the Libero benchmark suite. To compare and test its LeRobot implementation, we finetuned the GR00T N1.5 model for 30k steps on the Libero dataset and compared the results to the GR00T reference results.
-
-| Benchmark          | LeRobot Implementation | GR00T Reference |
-| ------------------ | ---------------------- | --------------- |
-| **Libero Spatial** | 82.0%                  | 92.0%           |
-| **Libero Object**  | 99.0%                  | 92.0%           |
-| **Libero Long**    | 82.0%                  | 76.0%           |
-| **Average**        | 87.0%                  | 87.0%           |
-
-These results demonstrate GR00T'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.
-
-### Evaluate in your hardware setup
-
-Once you have trained your model using your parameters you can run inference in your downstream task. Follow the instructions in [Imitation Learning for Robots](./il_robots). For example:
-
-```bash
-lerobot-record \
-  --robot.type=bi_so100_follower \
-  --robot.left_arm_port=/dev/ttyACM1 \
-  --robot.right_arm_port=/dev/ttyACM0 \
-  --robot.id=bimanual_follower \
-  --robot.cameras='{ right: {"type": "opencv", "index_or_path": 0, "width": 640, "height": 480, "fps": 30},
-    left: {"type": "opencv", "index_or_path": 2, "width": 640, "height": 480, "fps": 30},
-    top: {"type": "opencv", "index_or_path": 4, "width": 640, "height": 480, "fps": 30},
-  }' \
-  --display_data=true \
-  --dataset.repo_id=<user>/eval_groot-bimanual  \
-  --dataset.num_episodes=10 \
-  --dataset.single_task="Grab and handover the red cube to the other arm"
-  --policy.path=<user>/groot-bimanual # your trained model
-  --dataset.episode_time_s=30
-  --dataset.reset_time_s=10
-```
-
-## License
-
-This model follows the **Apache 2.0 License**, consistent with the original [GR00T repository](https://github.com/NVIDIA/Isaac-GR00T).

docs/source/il_robots.mdx

@@ -165,7 +165,7 @@ huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
 Then store your Hugging Face repository name in a variable:
 
 ```bash
-HF_USER=$(hf auth whoami | head -n 1)
+HF_USER=$(huggingface-cli whoami | head -n 1)
 echo $HF_USER
 ```
 
@@ -513,14 +513,13 @@ 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
@@ -563,7 +562,7 @@ init_rerun(session_name="recording")
 # Connect the robot
 robot.connect()
 
-preprocessor, postprocessor = make_pre_post_processors(
+preprocessor, postprocessor = make_processor(
     policy_cfg=policy,
     pretrained_path=HF_MODEL_ID,
     dataset_stats=dataset.meta.stats,

docs/source/installation.mdx

@@ -1,15 +1,8 @@
 # Installation
 
-## Install [`miniforge`](https://conda-forge.org/download/)
-
-```bash
-wget "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh"
-bash Miniforge3-$(uname)-$(uname -m).sh
-```
-
 ## Environment Setup
 
-Create a virtual environment with Python 3.10, using conda:
+Create a virtual environment with Python 3.10, using [`Miniconda`](https://docs.anaconda.com/miniconda/install/#quick-command-line-install)
 
 ```bash
 conda create -y -n lerobot python=3.10
@@ -21,7 +14,7 @@ Then activate your conda environment, you have to do this each time you open a s
 conda activate lerobot
 ```
 
-When using `conda`, install `ffmpeg` in your environment:
+When using `miniconda`, install `ffmpeg` in your environment:
 
 ```bash
 conda install ffmpeg -c conda-forge
@@ -81,9 +74,6 @@ _Replace `[...]` with your desired features._
 For a full list of optional dependencies, see:
 https://pypi.org/project/lerobot/
 
-> [!NOTE]
-> For lerobot 0.4.0, if you want to install pi, you will have to do: `pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git"`
-
 ### Troubleshooting
 
 If you encounter build errors, you may need to install additional dependencies: `cmake`, `build-essential`, and `ffmpeg libs`.
@@ -101,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)), or `pusht` ([gym-pusht](https://github.com/huggingface/gym-pusht))
+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))
 Example:
 
 ```bash

docs/source/integrate_hardware.mdx

@@ -208,36 +208,34 @@ LeRobot supports saving and loading calibration data automatically. This is usef
 
 <!-- prettier-ignore-start -->
 ```python
-@property
-def is_calibrated(self) -> bool:
-    return True
-
-def calibrate(self) -> None:
-    pass
-```
-<!-- prettier-ignore-end -->
+> @property
+> def is_calibrated(self) -> bool:
+>    return True
+>
+> def calibrate(self) -> None:
+>    pass
+> ```
 
 ### `is_calibrated`
 
 This should reflect whether your robot has the required calibration loaded.
 
-<!-- prettier-ignore-start -->
-```python
+```
+<!-- prettier-ignore-end -->python
 @property
 def is_calibrated(self) -> bool:
     return self.bus.is_calibrated
 ```
-<!-- prettier-ignore-end -->
 
 ### `calibrate()`
 
 The goal of the calibration is twofold:
-
-- Know the physical range of motion of each motors in order to only send commands within this range.
-- Normalize raw motors positions to sensible continuous values (e.g. percentages, degrees) instead of arbitrary discrete value dependant on the specific motor used that will not replicate elsewhere.
+    - Know the physical range of motion of each motors in order to only send commands within this range.
+    - Normalize raw motors positions to sensible continuous values (e.g. percentages, degrees) instead of arbitrary discrete value dependant on the specific motor used that will not replicate elsewhere.
 
 It should implement the logic for calibration (if relevant) and update the `self.calibration` dictionary. If you are using Feetech or Dynamixel motors, our bus interfaces already include methods to help with this.
 
+
 <!-- prettier-ignore-start -->
 ```python
 def calibrate(self) -> None:

docs/source/lerobot-dataset-v3.mdx

@@ -279,36 +279,3 @@ 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.

docs/source/libero.mdx

@@ -137,7 +137,7 @@ The finetuned model can be found here:
 We then evaluate the finetuned model using the LeRobot LIBERO implementation, by running the following command:
 
 ```bash
-lerobot-eval \
+python src/lerobot/scripts/eval.py \
   --output_dir=/logs/ \
   --env.type=libero \
   --env.task=libero_spatial,libero_object,libero_goal,libero_10 \

docs/source/metaworld.mdx

@@ -1,80 +0,0 @@
-# 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.

docs/source/multi_gpu_training.mdx

@@ -1,125 +0,0 @@
-# Multi-GPU Training
-
-This guide shows you how to train policies on multiple GPUs using [Hugging Face Accelerate](https://huggingface.co/docs/accelerate).
-
-## Installation
-
-First, ensure you have accelerate installed:
-
-```bash
-pip install accelerate
-```
-
-## Training with Multiple GPUs
-
-You can launch training in two ways:
-
-### Option 1: Without config (specify parameters directly)
-
-You can specify all parameters directly in the command without running `accelerate config`:
-
-```bash
-accelerate launch \
-  --multi_gpu \
-  --num_processes=2 \
-  $(which lerobot-train) \
-  --dataset.repo_id=${HF_USER}/my_dataset \
-  --policy.type=act \
-  --policy.repo_id=${HF_USER}/my_trained_policy \
-  --output_dir=outputs/train/act_multi_gpu \
-  --job_name=act_multi_gpu \
-  --wandb.enable=true
-```
-
-**Key accelerate parameters:**
-
-- `--multi_gpu`: Enable multi-GPU training
-- `--num_processes=2`: Number of GPUs to use
-- `--mixed_precision=fp16`: Use fp16 mixed precision (or `bf16` if supported)
-
-### Option 2: Using accelerate config
-
-If you prefer to save your configuration, you can optionally configure accelerate for your hardware setup by running:
-
-```bash
-accelerate config
-```
-
-This interactive setup will ask you questions about your training environment (number of GPUs, mixed precision settings, etc.) and saves the configuration for future use. For a simple multi-GPU setup on a single machine, you can use these recommended settings:
-
-- Compute environment: This machine
-- Number of machines: 1
-- Number of processes: (number of GPUs you want to use)
-- GPU ids to use: (leave empty to use all)
-- Mixed precision: fp16 or bf16 (recommended for faster training)
-
-Then launch training with:
-
-```bash
-accelerate launch $(which lerobot-train) \
-  --dataset.repo_id=${HF_USER}/my_dataset \
-  --policy.type=act \
-  --policy.repo_id=${HF_USER}/my_trained_policy \
-  --output_dir=outputs/train/act_multi_gpu \
-  --job_name=act_multi_gpu \
-  --wandb.enable=true
-```
-
-## How It Works
-
-When you launch training with accelerate:
-
-1. **Automatic detection**: LeRobot automatically detects if it's running under accelerate
-2. **Data distribution**: Your batch is automatically split across GPUs
-3. **Gradient synchronization**: Gradients are synchronized across GPUs during backpropagation
-4. **Single process logging**: Only the main process logs to wandb and saves checkpoints
-
-## Learning Rate and Training Steps Scaling
-
-**Important:** LeRobot does **NOT** automatically scale learning rates or training steps based on the number of GPUs. This gives you full control over your training hyperparameters.
-
-### Why No Automatic Scaling?
-
-Many distributed training frameworks automatically scale the learning rate by the number of GPUs (e.g., `lr = base_lr × num_gpus`).
-However, LeRobot keeps the learning rate exactly as you specify it.
-
-### When and How to Scale
-
-If you want to scale your hyperparameters when using multiple GPUs, you should do it manually:
-
-**Learning Rate Scaling:**
-
-```bash
-# Example: 2 GPUs with linear LR scaling
-# Base LR: 1e-4, with 2 GPUs -> 2e-4
-accelerate launch --num_processes=2 $(which lerobot-train) \
-  --optimizer.lr=2e-4 \
-  --dataset.repo_id=lerobot/pusht \
-  --policy=act
-```
-
-**Training Steps Scaling:**
-
-Since the effective batch size `bs` increases with multiple GPUs (batch_size × num_gpus), you may want to reduce the number of training steps proportionally:
-
-```bash
-# Example: 2 GPUs with effective batch size 2x larger
-# Original: batch_size=8, steps=100000
-# With 2 GPUs: batch_size=8 (16 in total), steps=50000
-accelerate launch --num_processes=2 $(which lerobot-train) \
-  --batch_size=8 \
-  --steps=50000 \
-  --dataset.repo_id=lerobot/pusht \
-  --policy=act
-```
-
-## Notes
-
-- The `--policy.use_amp` flag in `lerobot-train` is only used when **not** running with accelerate. When using accelerate, mixed precision is controlled by accelerate's configuration.
-- Training logs, checkpoints, and hub uploads are only done by the main process to avoid conflicts. Non-main processes have console logging disabled to prevent duplicate output.
-- The effective batch size is `batch_size × num_gpus`. If you use 4 GPUs with `--batch_size=8`, your effective batch size is 32.
-- Learning rate scheduling is handled correctly across multiple processes—LeRobot sets `step_scheduler_with_optimizer=False` to prevent accelerate from adjusting scheduler steps based on the number of processes.
-- When saving or pushing models, LeRobot automatically unwraps the model from accelerate's distributed wrapper to ensure compatibility.
-- WandB integration automatically initializes only on the main process, preventing multiple runs from being created.
-
-For more advanced configurations and troubleshooting, see the [Accelerate documentation](https://huggingface.co/docs/accelerate). If you want to learn more about how to train on a large number of GPUs, checkout this awesome guide: [Ultrascale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook).

docs/source/pi0.mdx

@@ -28,11 +28,6 @@ As described by Physical Intelligence, while AI has achieved remarkable success
    pip install -e ".[pi]"
    ```
 
-   > [!NOTE]
-   > For lerobot 0.4.0, if you want to install pi tag, you will have to do: `pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git"`.
-   >
-   > This will be solved in the next patch release
-
 ## Training Data and Capabilities
 
 π₀ is trained on the largest robot interaction dataset to date, combining three key data sources:

docs/source/pi05.mdx

@@ -36,11 +36,6 @@ This diverse training mixture creates a "curriculum" that enables generalization
    pip install -e ".[pi]"
    ```
 
-   > [!NOTE]
-   > For lerobot 0.4.0, if you want to install pi tag, you will have to do: `pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git"`.
-   >
-   > This will be solved in the next patch release
-
 ## Usage
 
 To use π₀.₅ in your LeRobot configuration, specify the policy type as:

docs/source/policy_groot_README.md

@@ -1,27 +0,0 @@
-## Research Paper
-
-Paper: https://research.nvidia.com/labs/gear/gr00t-n1_5/
-
-## Repository
-
-Code: https://github.com/NVIDIA/Isaac-GR00T
-
-## Citation
-
-```bibtex
-@inproceedings{gr00tn1_2025,
-  archivePrefix = {arxiv},
-  eprint     = {2503.14734},
-  title      = {{GR00T} {N1}: An Open Foundation Model for Generalist Humanoid Robots},
-  author     = {NVIDIA and Johan Bjorck andFernando Castañeda, Nikita Cherniadev and Xingye Da and Runyu Ding and Linxi "Jim" Fan and Yu Fang and Dieter Fox and Fengyuan Hu and Spencer Huang and Joel Jang and Zhenyu Jiang and Jan Kautz and Kaushil Kundalia and Lawrence Lao and Zhiqi Li and Zongyu Lin and Kevin Lin and Guilin Liu and Edith Llontop and Loic Magne and Ajay Mandlekar and Avnish Narayan and Soroush Nasiriany and Scott Reed and You Liang Tan and Guanzhi Wang and Zu Wang and Jing Wang and Qi Wang and Jiannan Xiang and Yuqi Xie and Yinzhen Xu and Zhenjia Xu and Seonghyeon Ye and Zhiding Yu and Ao Zhang and Hao Zhang and Yizhou Zhao and Ruijie Zheng and Yuke Zhu},
-  month      = {March},
-  year       = {2025},
-  booktitle  = {ArXiv Preprint},
-}
-```
-
-## Additional Resources
-
-Blog: https://developer.nvidia.com/isaac/gr00t
-
-Hugging Face Model: https://huggingface.co/nvidia/GR00T-N1.5-3B

examples/dataset/load_lerobot_dataset.py

@@ -132,15 +132,17 @@ print(f"\n{dataset[0][camera_key].shape=}")  # (4, c, h, w)
 print(f"{dataset[0]['observation.state'].shape=}")  # (6, c)
 print(f"{dataset[0]['action'].shape=}\n")  # (64, c)
 
-if __name__ == "__main__":
-    dataloader = torch.utils.data.DataLoader(
-        dataset,
-        num_workers=4,
-        batch_size=32,
-        shuffle=True,
-    )
-    for batch in dataloader:
-        print(f"{batch[camera_key].shape=}")  # (32, 4, c, h, w)
-        print(f"{batch['observation.state'].shape=}")  # (32, 6, c)
-        print(f"{batch['action'].shape=}")  # (32, 64, c)
-        break
+# Finally, our datasets are fully compatible with PyTorch dataloaders and samplers because they are just
+# PyTorch datasets.
+dataloader = torch.utils.data.DataLoader(
+    dataset,
+    num_workers=4,
+    batch_size=32,
+    shuffle=True,
+)
+
+for batch in dataloader:
+    print(f"{batch[camera_key].shape=}")  # (32, 4, c, h, w)
+    print(f"{batch['observation.state'].shape=}")  # (32, 6, c)
+    print(f"{batch['action'].shape=}")  # (32, 64, c)
+    break

examples/dataset/use_dataset_tools.py

@@ -30,10 +30,9 @@ Usage:
 import numpy as np
 
 from lerobot.datasets.dataset_tools import (
-    add_features,
+    add_feature,
     delete_episodes,
     merge_datasets,
-    modify_features,
     remove_feature,
     split_dataset,
 )
@@ -58,56 +57,50 @@ def main():
     print(f"Train split: {splits['train'].meta.total_episodes} episodes")
     print(f"Val split: {splits['val'].meta.total_episodes} episodes")
 
-    print("\n3. Adding features...")
+    print("\n3. Adding a reward feature...")
 
     reward_values = np.random.randn(dataset.meta.total_frames).astype(np.float32)
+    dataset_with_reward = add_feature(
+        dataset,
+        feature_name="reward",
+        feature_values=reward_values,
+        feature_info={
+            "dtype": "float32",
+            "shape": (1,),
+            "names": None,
+        },
+        repo_id="lerobot/pusht_with_reward",
+    )
 
     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},
-            ),
+    dataset_with_success = add_feature(
+        dataset_with_reward,
+        feature_name="success",
+        feature_values=compute_success,
+        feature_info={
+            "dtype": "float32",
+            "shape": (1,),
+            "names": None,
         },
-        repo_id="lerobot/pusht_with_features",
+        repo_id="lerobot/pusht_with_reward_and_success",
     )
 
-    print(f"New features: {list(dataset_with_features.meta.features.keys())}")
+    print(f"New features: {list(dataset_with_success.meta.features.keys())}")
 
     print("\n4. Removing the success feature...")
     dataset_cleaned = remove_feature(
-        dataset_with_features, feature_names="success", repo_id="lerobot/pusht_cleaned"
+        dataset_with_success, 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...")
+    print("\n5. 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...")
+    print("\n6. Complex workflow example...")
 
     if len(dataset.meta.camera_keys) > 1:
         camera_to_remove = dataset.meta.camera_keys[0]

examples/lekiwi/evaluate.py

@@ -133,6 +133,4 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
 log_say("Stop recording")
 robot.disconnect()
 listener.stop()
-
-dataset.finalize()
 dataset.push_to_hub()

examples/lekiwi/record.py

@@ -130,6 +130,4 @@ robot.disconnect()
 leader_arm.disconnect()
 keyboard.disconnect()
 listener.stop()
-
-dataset.finalize()
 dataset.push_to_hub()

examples/phone_to_so100/evaluate.py

@@ -194,6 +194,4 @@ for episode_idx in range(NUM_EPISODES):
 log_say("Stop recording")
 robot.disconnect()
 listener.stop()
-
-dataset.finalize()
 dataset.push_to_hub()

examples/phone_to_so100/record.py

@@ -200,6 +200,4 @@ log_say("Stop recording")
 robot.disconnect()
 phone.disconnect()
 listener.stop()
-
-dataset.finalize()
 dataset.push_to_hub()

examples/port_datasets/port_droid.py

@@ -362,8 +362,6 @@ 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

examples/rtc/PROFILING_GUIDE.md

@@ -1,263 +0,0 @@
-# RTC Profiling Guide
-
-This guide explains how to profile RTC (Real-Time Chunking) performance to identify bottlenecks and understand why RTC might be slower than expected.
-
-## Quick Start
-
-### 1. Profile with Real Robot (Profiled Version)
-
-Use `eval_with_real_robot_profiled.py` to profile actual robot execution:
-
-```bash
-# With RTC enabled
-uv run examples/rtc/eval_with_real_robot_profiled.py \
-    --policy.path=helper2424/pi05_check_rtc \
-    --policy.device=mps \
-    --rtc.enabled=true \
-    --rtc.execution_horizon=20 \
-    --robot.type=so100_follower \
-    --robot.port=/dev/tty.usbmodem58FA0834591 \
-    --robot.id=so100_follower \
-    --robot.cameras="{ gripper: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}, front: {type: opencv, index_or_path: 1, width: 640, height: 480, fps: 30}}" \
-    --task="Move green small object into the purple platform" \
-    --duration=30
-
-# Without RTC for comparison
-uv run examples/rtc/eval_with_real_robot_profiled.py \
-    --policy.path=helper2424/pi05_check_rtc \
-    --policy.device=mps \
-    --rtc.enabled=false \
-    --robot.type=so100_follower \
-    --robot.port=/dev/tty.usbmodem58FA0834591 \
-    --robot.id=so100_follower \
-    --robot.cameras="{ gripper: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}, front: {type: opencv, index_or_path: 1, width: 640, height: 480, fps: 30}}" \
-    --task="Move green small object into the purple platform" \
-    --duration=30
-```
-
-**Output**: At the end of execution, you'll see a detailed breakdown of timing for each component:
-- `get_actions.policy_inference` - Time spent in policy inference
-- `get_actions.preprocessing` - Time spent preprocessing observations
-- `get_actions.postprocessing` - Time spent postprocessing actions
-- `get_actions.action_queue_merge` - Time spent merging actions with RTC
-- `robot.get_observation` - Time to get observations from robot
-- `robot.send_action` - Time to send actions to robot
-- And more...
-
-### 2. Profile Without Robot (Comparison Script)
-
-Use `profile_rtc_comparison.py` to profile just the policy inference without needing a robot:
-
-```bash
-uv run examples/rtc/profile_rtc_comparison.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=50 \
-    --execution_horizon=20
-```
-
-**Output**: Side-by-side comparison of performance with and without RTC, including:
-- Mean/min/max inference times
-- Throughput (iterations per second)
-- Verdict on whether RTC is faster or slower
-
-### 3. Enable Detailed Method-Level Profiling
-
-For even more granular profiling, add the `--enable_detailed_profiling` flag:
-
-```bash
-uv run examples/rtc/profile_rtc_comparison.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=50 \
-    --execution_horizon=20 \
-    --enable_detailed_profiling
-```
-
-This will show timing for individual methods within the policy.
-
-## Understanding the Output
-
-### Key Metrics to Look At
-
-1. **get_actions.policy_inference** - This should be the largest component
-   - If RTC is enabled, this includes the RTC guidance overhead
-   - Compare this with/without RTC to see the overhead
-
-2. **get_actions.preprocessing** - Image preprocessing and normalization
-   - Should be relatively fast
-   - If slow, consider optimizing image processing
-
-3. **get_actions.postprocessing** - Action denormalization
-   - Should be minimal
-   - If slow, check postprocessor implementation
-
-4. **get_actions.action_queue_merge** - RTC-specific merging logic
-   - Only present when RTC is enabled
-   - If this is taking significant time, the RTC algorithm may need optimization
-
-5. **robot.get_observation** - Robot communication overhead
-   - If slow, check camera/sensor latency
-   - Consider reducing image resolution
-
-6. **robot.send_action** - Action execution overhead
-   - Should be very fast
-   - If slow, check robot communication
-
-### Expected Performance
-
-For a typical Pi0 policy on Apple Silicon (MPS):
-- **Without RTC**: ~100-200ms per inference
-- **With RTC**: Should be similar or slightly faster due to action reuse
-- **Preprocessing**: ~5-20ms depending on number of cameras
-- **Postprocessing**: ~1-5ms
-
-If RTC is significantly slower, likely causes:
-1. **RTC overhead exceeds benefits** - The guidance computation is expensive
-2. **Execution horizon too small** - Not reusing enough actions to amortize overhead
-3. **No compilation** - Try with `--use_torch_compile`
-4. **Large prev_actions buffer** - Copying/processing previous actions is slow
-
-## Profiling Your Own Code
-
-### Using the Profiling Decorator
-
-Add profiling to your own methods:
-
-```python
-from lerobot.utils.profiling import profile_method, enable_profiling, print_profiling_summary
-
-# Enable profiling
-enable_profiling()
-
-# Decorate methods you want to profile
-@profile_method
-def my_slow_function(x):
-    # ... your code ...
-    return result
-
-# At end of execution
-print_profiling_summary()
-```
-
-### Using Profile Context Manager
-
-For profiling specific code blocks:
-
-```python
-from lerobot.utils.profiling import profile_section, enable_profiling
-
-enable_profiling()
-
-with profile_section("data_loading"):
-    data = load_data()
-
-with profile_section("model_inference"):
-    output = model(data)
-```
-
-### Adding Profiling to Policy Methods
-
-To profile specific parts of the Pi0 policy, you can add decorators:
-
-```python
-# In src/lerobot/policies/pi0/modeling_pi0.py
-from lerobot.utils.profiling import profile_method, profile_section
-
-class Pi0Policy:
-    @profile_method
-    def predict_action_chunk(self, obs, inference_delay=0, prev_chunk_left_over=None):
-        # ... existing code ...
-        pass
-
-    def _generate_actions_with_rtc(self, ...):
-        with profile_section("rtc.guidance_computation"):
-            # ... guidance code ...
-            pass
-        
-        with profile_section("rtc.action_merging"):
-            # ... merging code ...
-            pass
-```
-
-## Analyzing Results
-
-### Comparison Checklist
-
-When comparing RTC vs non-RTC performance, check:
-
-- [ ] Is `policy_inference` time higher with RTC?
-- [ ] Is `action_queue_merge` taking significant time?
-- [ ] Are you running enough iterations to amortize warmup?
-- [ ] Is torch.compile enabled for fair comparison?
-- [ ] Is the execution horizon large enough? (should be >= 10-20)
-- [ ] Are you testing on the same hardware/device?
-
-### Common Bottlenecks
-
-1. **Image preprocessing dominates** 
-   - Solution: Reduce image resolution, use fewer cameras, or optimize preprocessing
-
-2. **Action queue operations are slow**
-   - Solution: Review queue implementation, consider using ring buffer
-
-3. **RTC guidance is expensive**
-   - Solution: Reduce guidance weight, simplify guidance computation, use torch.compile
-
-4. **Robot communication is slow**
-   - Solution: Increase baud rate, reduce action frequency, optimize protocol
-
-5. **Memory allocation overhead**
-   - Solution: Pre-allocate buffers, reuse tensors, avoid unnecessary copies
-
-## Advanced: Adding Custom Metrics
-
-You can add custom timing metrics to the profiled script:
-
-```python
-from lerobot.utils.profiling import record_timing
-
-start = time.perf_counter()
-# ... your code ...
-duration = time.perf_counter() - start
-record_timing("my_custom_metric", duration)
-```
-
-## Troubleshooting
-
-### Profiling shows RTC is slower by >50%
-
-1. Check if torch.compile is enabled: `--use_torch_compile`
-2. Increase execution horizon: `--rtc.execution_horizon=30`
-3. Verify inference_delay is calculated correctly
-4. Profile with `--enable_detailed_profiling` to find exact bottleneck
-
-### Profiling output is empty
-
-1. Make sure profiling is enabled with `enable_profiling()`
-2. Verify you're running enough iterations (at least 10)
-3. Check that code is actually executing (not short-circuited)
-
-### Inconsistent results between runs
-
-1. Run more iterations: `--num_iterations=100`
-2. Increase warmup iterations
-3. Check for thermal throttling on device
-4. Ensure no other processes competing for resources
-
-## Next Steps
-
-1. Run both profiling scripts (with/without robot)
-2. Compare timing breakdowns
-3. Identify the largest bottleneck
-4. Focus optimization efforts on that component
-5. Re-run profiling to verify improvements
-
-## Questions?
-
-If profiling reveals unexpected bottlenecks or you need help interpreting results, please share:
-- The full profiling output
-- Your configuration (RTC enabled/disabled, execution horizon, etc.)
-- Hardware specs (device type, memory, etc.)
-- Policy type and size
-

examples/rtc/PROFILING_QUICK_START.md

@@ -1,208 +0,0 @@
-# RTC Profiling - Quick Start
-
-Quick reference for profiling Pi0 with RTC to identify performance bottlenecks.
-
-## 🚀 Quick Commands
-
-### 1. Profile with Real Robot
-
-```bash
-# With RTC enabled (profiled version)
-uv run examples/rtc/eval_with_real_robot_profiled.py \
-    --policy.path=helper2424/pi05_check_rtc \
-    --policy.device=mps \
-    --rtc.enabled=true \
-    --rtc.execution_horizon=20 \
-    --robot.type=so100_follower \
-    --robot.port=/dev/tty.usbmodem58FA0834591 \
-    --robot.cameras="{ gripper: {type: opencv, index_or_path: 0}, front: {type: opencv, index_or_path: 1}}" \
-    --task="Pick up object" \
-    --duration=30
-```
-
-### 2. Compare RTC vs No-RTC (No Robot Needed)
-
-```bash
-uv run examples/rtc/profile_rtc_comparison.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=50 \
-    --execution_horizon=20
-```
-
-### 3. Detailed RTC Method Profiling
-
-```bash
-uv run examples/rtc/profile_pi0_rtc_detailed.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=20 \
-    --execution_horizon=20 \
-    --enable_rtc_profiling
-```
-
-## 📊 What Each Tool Does
-
-| Tool | Purpose | Needs Robot? |
-|------|---------|--------------|
-| `eval_with_real_robot_profiled.py` | Profile actual robot execution with RTC | ✅ Yes |
-| `profile_rtc_comparison.py` | Compare RTC vs no-RTC side-by-side | ❌ No |
-| `profile_pi0_rtc_detailed.py` | Deep dive into RTC internals | ❌ No |
-
-## 🔍 Key Metrics to Watch
-
-### Overall Performance
-- **iteration.policy_inference** - Total policy inference time
-- **iteration.preprocessing** - Image preprocessing time
-- **iteration.postprocessing** - Action denormalization time
-
-### RTC-Specific (with `--enable_rtc_profiling`)
-- **rtc.denoise_step.base_denoising** - Time without RTC overhead
-- **rtc.denoise_step.autograd_correction** - Gradient computation time
-- **rtc.denoise_step.guidance_computation** - Total RTC guidance overhead
-
-### Robot Communication
-- **robot.get_observation** - Time to get robot state
-- **robot.send_action** - Time to send action command
-
-## 🎯 Quick Diagnosis
-
-### RTC is slower than expected?
-
-1. **Check if torch.compile is enabled**
-   ```bash
-   # Add this flag
-   --use_torch_compile
-   ```
-
-2. **Try larger execution horizon**
-   ```bash
-   # Increase to amortize RTC overhead
-   --rtc.execution_horizon=30
-   ```
-
-3. **Profile to find bottleneck**
-   ```bash
-   uv run examples/rtc/profile_pi0_rtc_detailed.py \
-       --policy_path=helper2424/pi05_check_rtc \
-       --device=mps \
-       --enable_rtc_profiling
-   ```
-
-### Preprocessing is slow?
-
-- Reduce image resolution in robot config
-- Use fewer cameras
-- Check camera FPS settings
-
-### Policy inference is slow?
-
-- Enable torch.compile
-- Check device (MPS vs CUDA vs CPU)
-- Try smaller model if available
-
-## 📈 Expected Performance
-
-### Typical timings on Apple Silicon (MPS):
-
-| Component | Time (ms) | Notes |
-|-----------|-----------|-------|
-| Policy inference | 100-200 | Depends on model size |
-| Preprocessing | 5-20 | Depends on #cameras |
-| Postprocessing | 1-5 | Usually fast |
-| RTC overhead | 10-50 | Should be < 50% of base |
-
-### When RTC helps:
-- ✅ Execution horizon ≥ 10
-- ✅ Inference time > action execution rate
-- ✅ Using torch.compile
-- ✅ Proper inference_delay calculation
-
-### When RTC might not help:
-- ❌ Very fast inference already
-- ❌ Small execution horizon (< 5)
-- ❌ No compilation (interpreted mode)
-- ❌ Inference delay not accounted for
-
-## 🛠️ Adding Profiling to Your Code
-
-### Quick snippet:
-
-```python
-from lerobot.utils.profiling import enable_profiling, print_profiling_summary, profile_section
-
-# Enable at start
-enable_profiling()
-
-# Profile sections
-with profile_section("my_operation"):
-    # ... your code ...
-    pass
-
-# Print at end
-print_profiling_summary()
-```
-
-### Profile specific methods:
-
-```python
-from lerobot.utils.profiling import profile_method
-
-@profile_method
-def my_slow_function():
-    # ... your code ...
-    pass
-```
-
-## 📝 Example Output
-
-```
-PROFILING SUMMARY
-================================================================================
-Function                                                    Count    Mean (ms)
---------------------------------------------------------------------------------
-iteration.policy_inference                                    20       150.23
-iteration.preprocessing                                       20        12.45
-rtc.denoise_step.guidance_computation                        200        15.67
-rtc.denoise_step.autograd_correction                         200         8.23
-rtc.denoise_step.base_denoising                             200       120.45
-================================================================================
-```
-
-## 🚨 Common Issues
-
-### "No profiling data available"
-- Did you call `enable_profiling()`?
-- Running enough iterations?
-
-### Inconsistent results
-- Increase `--num_iterations`
-- Check for thermal throttling
-- Close other applications
-
-### Can't find bottleneck
-- Enable `--enable_rtc_profiling` for detailed breakdown
-- Check both preprocessing and inference
-- Compare with and without RTC
-
-## 📖 More Details
-
-See `PROFILING_GUIDE.md` for comprehensive documentation.
-
-## 🤔 Still Slow?
-
-1. Run comparison: `profile_rtc_comparison.py`
-2. Run detailed profiling: `profile_pi0_rtc_detailed.py --enable_rtc_profiling`
-3. Share output for help (include device, model, settings)
-
-## ✅ Quick Checklist
-
-Before asking for help, verify:
-
-- [ ] Ran comparison script (with/without RTC)
-- [ ] Tried torch.compile
-- [ ] Tested different execution horizons (10, 20, 30)
-- [ ] Profiled with detailed RTC profiling
-- [ ] Checked preprocessing vs inference split
-- [ ] Verified hardware (device type, thermal state)
-

examples/rtc/PROFILING_README.md

@@ -1,352 +0,0 @@
-# RTC Profiling Toolkit
-
-Complete toolkit for profiling Pi0 with RTC to identify performance bottlenecks.
-
-## 📦 What's Included
-
-### Scripts
-
-1. **`eval_with_real_robot_profiled.py`**
-   - Profiled version of the real robot eval script
-   - Adds timing measurements throughout execution
-   - Works with actual robot hardware
-   - Same usage as original but with profiling output
-
-2. **`profile_rtc_comparison.py`**
-   - Side-by-side comparison of RTC vs no-RTC
-   - No robot needed (uses mock observations)
-   - Shows clear verdict on whether RTC is helping
-   - Great for quick performance checks
-
-3. **`profile_pi0_rtc_detailed.py`**
-   - Most detailed profiling available
-   - Can enable RTC method-level profiling
-   - Provides insights and recommendations
-   - Perfect for deep-dive investigations
-
-4. **`add_rtc_profiling.py`**
-   - Monkey-patching utility for RTC internals
-   - Profiles individual RTC operations
-   - Can be applied without modifying source
-   - Shows exactly where RTC spends time
-
-### Utilities
-
-5. **`src/lerobot/utils/profiling.py`**
-   - Core profiling utilities
-   - Decorators for method profiling
-   - Context managers for code blocks
-   - Statistics collection and reporting
-
-### Documentation
-
-6. **`PROFILING_GUIDE.md`** - Comprehensive guide
-7. **`PROFILING_QUICK_START.md`** - Quick reference
-
-## 🚀 Quick Start
-
-### Step 1: Compare Performance
-
-Run this first to see if RTC is actually slower:
-
-```bash
-uv run examples/rtc/profile_rtc_comparison.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=50 \
-    --execution_horizon=20
-```
-
-**Expected output:**
-```
-COMPARISON SUMMARY
-================================================================================
-Metric                         Without RTC        With RTC      Difference
---------------------------------------------------------------------------------
-Mean time (ms)                       150.23         165.45          +15.22
-Throughput (iter/s)                    6.66           6.05           -0.61
-================================================================================
-VERDICT
-✗ RTC is SLOWER by 10.1%
-  Mean time increased by 15.22 ms
-  
-  Possible reasons:
-  - RTC overhead exceeds benefits at current execution horizon
-  - No torch.compile enabled
-```
-
-### Step 2: Identify Bottleneck
-
-If RTC is slower, find out why:
-
-```bash
-uv run examples/rtc/profile_pi0_rtc_detailed.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=20 \
-    --execution_horizon=20 \
-    --enable_rtc_profiling
-```
-
-**Expected output:**
-```
-PROFILING SUMMARY
-================================================================================
-Function                                             Count    Mean (ms)    Total (s)
-------------------------------------------------------------------------------------
-iteration.policy_inference                              20      150.23         3.00
-rtc.denoise_step.guidance_computation                  200       15.67         3.13
-rtc.denoise_step.autograd_correction                   200        8.23         1.65
-iteration.preprocessing                                 20       12.45         0.25
-================================================================================
-
-KEY INSIGHTS
-================================================================================
-Time breakdown:
-  Policy inference:  150.23 ms (87.2%)
-  Preprocessing:     12.45 ms (7.2%)
-  Postprocessing:    2.10 ms (1.2%)
-
-RTC breakdown:
-  Base denoising:    120.45 ms
-  Guidance compute:  15.67 ms
-  Autograd correct:  8.23 ms
-  RTC overhead:      23.90 ms (19.8% of base)
-
-Recommendations:
-  ⚠ RTC autograd overhead is significant
-    → This is expected, but consider increasing execution_horizon
-    → Try torch.compile if not already enabled
-  💡 torch.compile not enabled
-    → Try --use_torch_compile for potential speedup
-================================================================================
-```
-
-### Step 3: Try Optimizations
-
-Based on recommendations:
-
-```bash
-# Try with torch.compile
-uv run examples/rtc/profile_rtc_comparison.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=50 \
-    --execution_horizon=20 \
-    --use_torch_compile
-
-# Try larger execution horizon
-uv run examples/rtc/profile_rtc_comparison.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=50 \
-    --execution_horizon=30
-```
-
-### Step 4: Profile Real Robot (Optional)
-
-Test with actual hardware:
-
-```bash
-uv run examples/rtc/eval_with_real_robot_profiled.py \
-    --policy.path=helper2424/pi05_check_rtc \
-    --policy.device=mps \
-    --rtc.enabled=true \
-    --rtc.execution_horizon=20 \
-    --robot.type=so100_follower \
-    --robot.port=/dev/tty.usbmodem58FA0834591 \
-    --robot.cameras="{...}" \
-    --task="Pick up object" \
-    --duration=30
-```
-
-## 🎯 Common Scenarios
-
-### "RTC is 2x slower!"
-
-This usually means:
-- RTC overhead is high but not getting benefits
-- Need to enable torch.compile
-- Execution horizon too small
-- Inference delay not calculated correctly
-
-**Try:**
-1. `--use_torch_compile`
-2. Increase `--execution_horizon` to 30+
-3. Check inference_delay calculation
-
-### "RTC is only slightly slower"
-
-This is expected! RTC overhead is about 10-30% typically.
-The benefit comes during **execution**, not single inference:
-- Actions are reused across chunks
-- Overall system latency is reduced
-- Robot gets smoother actions
-
-### "Want to optimize specific part"
-
-Use the profiling utilities:
-
-```python
-from lerobot.utils.profiling import enable_profiling, profile_section, print_profiling_summary
-
-enable_profiling()
-
-with profile_section("my_custom_operation"):
-    # Your code here
-    pass
-
-print_profiling_summary()
-```
-
-## 📊 Understanding Results
-
-### Key Metrics
-
-**Policy Inference Time**
-- Time for forward pass through model
-- Should be largest component (70-90%)
-- Includes RTC guidance if enabled
-
-**Preprocessing Time**
-- Image normalization, resizing
-- Should be < 20% of total
-- If high: reduce image resolution
-
-**RTC Guidance Overhead**
-- Extra time for RTC guidance computation
-- Typically 10-30% of base inference
-- If > 50%: RTC may not be beneficial at current settings
-
-**Autograd Correction**
-- Time computing gradients for RTC
-- Usually 5-15% of base inference
-- Can be reduced with torch.compile
-
-### Expected Ranges (Apple Silicon MPS)
-
-| Metric | Good | Acceptable | Poor |
-|--------|------|------------|------|
-| Policy inference | 100-150ms | 150-250ms | >250ms |
-| Preprocessing | <20ms | 20-50ms | >50ms |
-| RTC overhead | 10-30% | 30-50% | >50% |
-
-## 🔧 Optimization Guide
-
-### If RTC overhead is too high:
-
-1. **Enable compilation:**
-   ```bash
-   --use_torch_compile
-   ```
-   Expected improvement: 20-40% faster
-
-2. **Increase execution horizon:**
-   ```bash
-   --execution_horizon=30  # or higher
-   ```
-   Amortizes RTC cost over more actions
-
-3. **Check guidance weight:**
-   ```python
-   # In config
-   rtc.max_guidance_weight=1.0  # try 0.5 for less overhead
-   ```
-
-### If preprocessing is slow:
-
-1. **Reduce image resolution:**
-   ```python
-   # In robot config
-   cameras={
-       "gripper": {"width": 320, "height": 240}  # instead of 640x480
-   }
-   ```
-
-2. **Use fewer cameras:**
-   - Profile which cameras are essential
-   - Remove unnecessary views
-
-### If inference is generally slow:
-
-1. Use torch.compile (if not already)
-2. Check device is correct (MPS vs CUDA)
-3. Verify model is in eval mode
-4. Check for unnecessary gradient tracking
-
-## 🐛 Troubleshooting
-
-### Empty profiling output
-```python
-# Make sure to enable profiling!
-from lerobot.utils.profiling import enable_profiling
-enable_profiling()
-```
-
-### Inconsistent timings
-- Run more iterations (50-100)
-- Check thermal throttling
-- Close background apps
-- Use `--warmup_iterations=10`
-
-### Can't find bottleneck
-1. Start with `profile_rtc_comparison.py`
-2. Then run `profile_pi0_rtc_detailed.py --enable_rtc_profiling`
-3. Compare with/without RTC
-4. Check each component separately
-
-## 📖 Full Documentation
-
-- **`PROFILING_GUIDE.md`** - Complete reference with examples
-- **`PROFILING_QUICK_START.md`** - Quick commands and tips
-
-## 🤝 Getting Help
-
-If you're still experiencing issues:
-
-1. Run comparison script and save output
-2. Run detailed profiling and save output
-3. Include:
-   - Policy path
-   - Device type
-   - RTC settings (execution_horizon, etc.)
-   - Hardware specs
-   - Full profiling output
-
-## 🎓 Learning More
-
-### Profiling your own code:
-
-```python
-from lerobot.utils.profiling import profile_method, enable_profiling
-
-enable_profiling()
-
-@profile_method
-def my_function():
-    # Automatically profiled
-    pass
-```
-
-### RTC internals:
-
-```python
-from examples.rtc.add_rtc_profiling import monkey_patch_rtc_profiling
-
-enable_profiling()
-monkey_patch_rtc_profiling()
-
-# Now RTC methods are profiled
-policy.predict_action_chunk(...)
-```
-
-## ✨ Next Steps
-
-1. Run `profile_rtc_comparison.py` to establish baseline
-2. Use `profile_pi0_rtc_detailed.py` to find bottlenecks
-3. Apply optimizations (torch.compile, larger horizon)
-4. Re-run comparison to verify improvements
-5. Test with real robot using profiled version
-
-Happy profiling! 🚀
-

examples/rtc/README.md

@@ -1,251 +0,0 @@
-# Real-Time Chunking (RTC) Examples
-
-This directory contains examples and evaluation scripts for Real-Time Chunking (RTC), a technique for improving action chunking policies in real-time robot control.
-
-## Overview
-
-Real-Time Chunking addresses the challenge of maintaining consistency and reactivity when using action chunking policies with non-negligible inference latency. It uses a guidance technique during diffusion sampling to blend new action predictions with previously planned actions.
-
-**Key Benefits:**
-
-- Maintains consistency between consecutive action chunks
-- Reduces jitter and improves smoothness
-- Adapts to inference delays dynamically
-
-**Reference:** [Physical Intelligence - Real-Time Chunking](https://www.physicalintelligence.company/download/real_time_chunking.pdf)
-
-## Scripts
-
-### 1. `eval_dataset.py`
-
-Offline evaluation on dataset samples with detailed visualization and validation.
-
-**Features:**
-
-- Compare RTC vs non-RTC predictions on two random dataset samples
-- Validate RTC behavior (delay region, blend region, post-horizon region)
-- Generate debug visualizations:
-  - Denoising step comparisons (x_t, v_t, x1_t, corrections)
-  - Final action predictions comparison
-- Support for torch.compile() optimization
-- Memory-efficient sequential policy loading for large models
-
-**Usage:**
-
-```bash
-# Basic usage with SmolVLA policy
-uv run python examples/rtc/eval_dataset.py \
-    --policy.path=helper2424/smolvla_check_rtc_last3 \
-    --dataset.repo_id=helper2424/check_rtc \
-    --rtc.execution_horizon=8 \
-    --device=mps \
-    --rtc.max_guidance_weight=10.0 \
-    --seed=10
-
-# With Pi0.5 policy on CUDA
-uv run python examples/rtc/eval_dataset.py \
-    --policy.path=lerobot/pi05_libero_finetuned \
-    --dataset.repo_id=HuggingFaceVLA/libero \
-    --rtc.execution_horizon=8 \
-    --device=cuda
-
-# With Pi0 policy
-uv run python examples/rtc/eval_dataset.py \
-    --policy.path=lerobot/pi0_libero_finetuned \
-    --dataset.repo_id=HuggingFaceVLA/libero \
-    --rtc.execution_horizon=8 \
-    --device=cuda
-
-# With torch.compile for faster inference
-uv run python examples/rtc/eval_dataset.py \
-    --policy.path=helper2424/smolvla_check_rtc_last3 \
-    --dataset.repo_id=helper2424/check_rtc \
-    --rtc.execution_horizon=8 \
-    --device=cuda \
-    --use_torch_compile=true \
-    --torch_compile_mode=max-autotune
-
-# Enable CUDA graphs (advanced - may cause tensor aliasing errors)
-uv run python examples/rtc/eval_dataset.py \
-    --policy.path=helper2424/smolvla_check_rtc_last3 \
-    --dataset.repo_id=helper2424/check_rtc \
-    --use_torch_compile=true \
-    --torch_compile_backend=inductor \
-    --torch_compile_mode=max-autotune \
-    --torch_compile_disable_cudagraphs=false
-```
-
-**Key Parameters:**
-
-- `--policy.path`: Path to pretrained policy
-- `--dataset.repo_id`: Dataset to evaluate on
-- `--rtc.execution_horizon`: Number of steps to maintain consistency (default: 20)
-- `--rtc.max_guidance_weight`: Maximum guidance weight (default: 10.0)
-- `--rtc.prefix_attention_schedule`: Schedule type (ZEROS, ONES, LINEAR, EXP)
-- `--inference_delay`: Inference delay for RTC (default: 4)
-- `--seed`: Random seed for reproducibility (default: 42)
-- `--output_dir`: Directory to save visualizations (default: rtc_debug_output)
-- `--device`: Device to use (cuda, cpu, mps, auto)
-- `--use_torch_compile`: Enable torch.compile() for faster inference
-
-**Output:**
-
-The script generates several visualization files in `rtc_debug_output/`:
-
-- `denoising_xt_comparison.png` - Noisy state evolution during denoising
-- `denoising_vt_comparison.png` - Velocity predictions during denoising
-- `denoising_x1t_comparison.png` - Predicted final states during denoising
-- `denoising_correction_comparison.png` - RTC guidance corrections applied
-- `final_actions_comparison.png` - Final action predictions (prev_chunk, no_rtc, rtc)
-
-The script also validates RTC behavior and reports:
-
-- ✅ Delay region [0:inference_delay]: RTC = prev_chunk
-- ✅ Blend region [inference_delay:execution_horizon]: prev_chunk ≤ RTC ≤ no_rtc
-- ✅ Post-horizon [execution_horizon:]: RTC = no_rtc
-
-### 2. `eval_with_real_robot.py`
-
-Real-time evaluation on physical robots or simulation environments.
-
-**Features:**
-
-- Run policy with RTC on real robot or simulation
-- Multi-threaded action execution and inference
-- Action queue management with proper timing
-- Latency tracking and adaptive inference delay
-- Support for both robots and gym environments
-- Support for torch.compile() optimization
-
-**Usage:**
-
-```bash
-# With real robot
-uv run python examples/rtc/eval_with_real_robot.py \
-    --policy.path=lerobot/smolvla_base \
-    --robot.type=so100 \
-    --task="pick up the cup" \
-    --duration=30.0
-
-# With simulation environment
-uv run python examples/rtc/eval_with_real_robot.py \
-    --policy.path=lerobot/smolvla_base \
-    --env.type=pusht \
-    --duration=60.0
-
-# With policy compilation (CUDA only, not MPS)
-uv run python examples/rtc/eval_with_real_robot.py \
-    --policy.path=lerobot/smolvla_base \
-    --robot.type=so100 \
-    --use_torch_compile=true \
-    --torch_compile_mode=max-autotune
-```
-
-**Key Parameters:**
-
-- `--policy.path`: Path to pretrained policy
-- `--robot.type` or `--env.type`: Robot or environment to use
-- `--task`: Task description (for VLA models)
-- `--rtc.execution_horizon`: Number of steps to maintain consistency (default: 10)
-- `--rtc.max_guidance_weight`: Maximum guidance weight (default: 1.0)
-- `--rtc.prefix_attention_schedule`: Schedule type (ZEROS, ONES, LINEAR, EXP)
-- `--duration`: How long to run (seconds, default: 30.0)
-- `--fps`: Action execution frequency (Hz, default: 10.0)
-- `--action_queue_size_to_get_new_actions`: Queue size threshold to request new actions (default: 30)
-- `--device`: Device to use (cuda, cpu, mps, auto)
-- `--use_torch_compile`: Enable torch.compile() for faster inference
-
-## Understanding RTC Parameters
-
-### `execution_horizon`
-
-Number of timesteps from previous chunk to maintain consistency with. Higher values mean more consistency but potentially less reactivity.
-
-**Typical values:** 8-12 steps for dataset evaluation, 10 steps for real-time execution
-
-### `max_guidance_weight`
-
-Upper bound on guidance strength. Higher values give stronger consistency but may over-constrain new predictions.
-
-**Typical values:**
-
-- Dataset evaluation: 10.0-100.0 (can be higher for analysis)
-- Real-time execution: 1.0-10.0 (more conservative)
-
-### `prefix_attention_schedule`
-
-How to weight consistency across the overlap region:
-
-- `ZEROS`: Binary (full weight up to inference_delay, then zero)
-- `ONES`: Full weight across entire execution_horizon
-- `LINEAR`: Linear decay from inference_delay to execution_horizon
-- `EXP`: Exponential decay (recommended)
-
-**Recommended:** `EXP`
-
-### `inference_delay`
-
-Number of timesteps from the prefix to use for guidance. Typically calculated dynamically based on inference latency in real-time execution, but fixed for dataset evaluation.
-
-**Typical values:** 3-5 steps for dataset evaluation
-
-### `action_queue_size_to_get_new_actions` (real-time only)
-
-Threshold for requesting new action chunks. Should be higher than `inference_delay + execution_horizon` to ensure smooth operation.
-
-**Typical values:** 20-30 steps
-
-## Validation Rules (Dataset Evaluation)
-
-The dataset evaluation script validates that RTC behavior matches expectations:
-
-1. **Delay Region [0:inference_delay]**: RTC actions should equal previous chunk
-   - Ensures consistency during the inference delay period
-
-2. **Blend Region [inference_delay:execution_horizon]**: RTC should be between prev_chunk and no_rtc
-   - Smooth transition from previous plan to new predictions
-
-3. **Post-Horizon [execution_horizon:]**: RTC should equal no_rtc
-   - Full adoption of new predictions after execution horizon
-
-## Tips
-
-1. **Start with dataset evaluation** (`eval_dataset.py`) to understand RTC behavior and tune parameters before running on robot
-2. **Use visualizations** to debug unexpected behavior - check denoising steps and final actions
-3. **Tune execution_horizon** based on your inference latency and action frequency
-4. **Monitor validation output** - failures indicate potential implementation issues or misconfigured parameters
-5. **Compare different schedules** - EXP usually works best but LINEAR can be more interpretable
-
-## Troubleshooting
-
-### Validation fails in delay region
-
-- Check that `prev_chunk_left_over` is properly passed to the policy
-- Verify RTC guidance is being applied during denoising
-- Look at denoising visualizations to see where guidance diverges
-
-### Validation fails in post-horizon region
-
-- RTC and no_rtc use different noise - verify same noise is being used for comparison
-- Check that weights are correctly zeroed out after execution horizon
-- Review prefix_attention_schedule visualization
-
-### Poor performance on real robot
-
-- Increase `action_queue_size_to_get_new_actions` if you see warnings
-- Reduce `max_guidance_weight` if robot is too conservative
-- Try different `prefix_attention_schedule` values
-- Enable torch.compile() for faster inference (CUDA only)
-
-### Memory issues with large models
-
-- The dataset evaluation script loads policies sequentially to minimize memory
-- For real-time execution, only one policy is loaded
-- Use smaller batch sizes if needed
-
-## Related Documentation
-
-- [RTC Implementation](../../src/lerobot/policies/rtc/modeling_rtc.py)
-- [RTC Configuration](../../src/lerobot/policies/rtc/configuration_rtc.py)
-- [Action Queue](../../src/lerobot/policies/rtc/action_queue.py)
-- [Physical Intelligence Paper](https://www.physicalintelligence.company/download/real_time_chunking.pdf)

examples/rtc/add_rtc_profiling.py

@@ -1,202 +0,0 @@
-#!/usr/bin/env python
-
-"""
-Script to add profiling instrumentation to RTCProcessor.
-
-This script shows which methods to profile in the RTC code to identify bottlenecks.
-You can either:
-1. Apply these changes directly to modeling_rtc.py
-2. Use monkey patching to add profiling without modifying source
-3. Use as reference for manual instrumentation
-
-Usage:
-    # Option 1: Monkey patch (no source changes)
-    python examples/rtc/add_rtc_profiling.py
-
-    # Option 2: Apply changes to source
-    # Copy the profiled methods below into src/lerobot/policies/rtc/modeling_rtc.py
-"""
-
-import logging
-
-import torch
-from torch import Tensor
-
-from lerobot.policies.rtc.modeling_rtc import RTCProcessor
-from lerobot.utils.profiling import ProfileContext, enable_profiling, is_profiling_enabled
-
-logger = logging.getLogger(__name__)
-
-
-def profile_denoise_step(self, x_t, prev_chunk_left_over, inference_delay, time, original_denoise_step_partial, execution_horizon=None) -> Tensor:
-    """Profiled version of denoise_step."""
-    
-    if not is_profiling_enabled():
-        # Call original implementation if profiling disabled
-        return self._original_denoise_step(x_t, prev_chunk_left_over, inference_delay, time, original_denoise_step_partial, execution_horizon)
-    
-    with ProfileContext("rtc.denoise_step.total"):
-        # In the original implementation, the time goes from 0 to 1 and
-        # In our implementation, the time goes from 1 to 0
-        # So we need to invert the time
-        tau = 1 - time
-
-        if prev_chunk_left_over is None:
-            # First step, no guidance - return v_t
-            with ProfileContext("rtc.denoise_step.base_denoising"):
-                v_t = original_denoise_step_partial(x_t)
-            return v_t
-
-        with ProfileContext("rtc.denoise_step.setup"):
-            x_t = x_t.clone().detach()
-
-            squeezed = False
-            if len(x_t.shape) < 3:
-                x_t = x_t.unsqueeze(0)
-                squeezed = True
-
-            if len(prev_chunk_left_over.shape) < 3:
-                prev_chunk_left_over = prev_chunk_left_over.unsqueeze(0)
-
-            if execution_horizon is None:
-                execution_horizon = self.rtc_config.execution_horizon
-
-            if execution_horizon > prev_chunk_left_over.shape[1]:
-                execution_horizon = prev_chunk_left_over.shape[1]
-
-            batch_size = x_t.shape[0]
-            action_chunk_size = x_t.shape[1]
-            action_dim = x_t.shape[2]
-
-        # Padding
-        with ProfileContext("rtc.denoise_step.padding"):
-            if prev_chunk_left_over.shape[1] < action_chunk_size or prev_chunk_left_over.shape[2] < action_dim:
-                padded = torch.zeros(batch_size, action_chunk_size, action_dim).to(x_t.device)
-                padded[:, : prev_chunk_left_over.shape[1], : prev_chunk_left_over.shape[2]] = prev_chunk_left_over
-                prev_chunk_left_over = padded
-
-        # Get prefix weights
-        with ProfileContext("rtc.denoise_step.get_prefix_weights"):
-            weights = (
-                self.get_prefix_weights(inference_delay, execution_horizon, action_chunk_size)
-                .to(x_t.device)
-                .unsqueeze(0)
-                .unsqueeze(-1)
-            )
-
-        # Main RTC guidance computation
-        with ProfileContext("rtc.denoise_step.guidance_computation"):
-            with torch.enable_grad():
-                # Base denoising
-                with ProfileContext("rtc.denoise_step.base_denoising"):
-                    v_t = original_denoise_step_partial(x_t)
-                
-                x_t.requires_grad_(True)
-
-                # Compute x1_t
-                with ProfileContext("rtc.denoise_step.compute_x1_t"):
-                    x1_t = x_t - time * v_t
-
-                # Compute error
-                with ProfileContext("rtc.denoise_step.compute_error"):
-                    err = (prev_chunk_left_over - x1_t) * weights
-                    grad_outputs = err.clone().detach()
-
-                # Compute correction via autograd
-                with ProfileContext("rtc.denoise_step.autograd_correction"):
-                    correction = torch.autograd.grad(x1_t, x_t, grad_outputs, retain_graph=False)[0]
-
-        # Compute guidance weight
-        with ProfileContext("rtc.denoise_step.compute_guidance_weight"):
-            max_guidance_weight = torch.as_tensor(self.rtc_config.max_guidance_weight)
-            tau_tensor = torch.as_tensor(tau)
-            squared_one_minus_tau = (1 - tau_tensor) ** 2
-            inv_r2 = (squared_one_minus_tau + tau_tensor**2) / (squared_one_minus_tau)
-            c = torch.nan_to_num((1 - tau_tensor) / tau_tensor, posinf=max_guidance_weight)
-            guidance_weight = torch.nan_to_num(c * inv_r2, posinf=max_guidance_weight)
-            guidance_weight = torch.minimum(guidance_weight, max_guidance_weight)
-
-        # Apply guidance
-        with ProfileContext("rtc.denoise_step.apply_guidance"):
-            result = v_t - guidance_weight * correction
-
-        # Cleanup
-        with ProfileContext("rtc.denoise_step.cleanup"):
-            if squeezed:
-                result = result.squeeze(0)
-                correction = correction.squeeze(0)
-                x1_t = x1_t.squeeze(0)
-                err = err.squeeze(0)
-
-            self.track(
-                time=time,
-                x1_t=x1_t,
-                correction=correction,
-                err=err,
-                weights=weights,
-                guidance_weight=guidance_weight,
-                inference_delay=inference_delay,
-                execution_horizon=execution_horizon,
-            )
-
-        return result
-
-
-def monkey_patch_rtc_profiling():
-    """Apply profiling to RTCProcessor via monkey patching.
-    
-    This modifies the RTCProcessor class at runtime to add profiling
-    without changing source files.
-    """
-    logger.info("Applying RTC profiling monkey patch...")
-    
-    # Save original method
-    RTCProcessor._original_denoise_step = RTCProcessor.denoise_step
-    
-    # Replace with profiled version
-    RTCProcessor.denoise_step = profile_denoise_step
-    
-    logger.info("✓ RTC profiling enabled")
-
-
-def print_usage():
-    """Print usage instructions."""
-    print("\n" + "="*80)
-    print("RTC PROFILING INSTRUMENTATION")
-    print("="*80)
-    print("\nThis script provides profiling for RTCProcessor methods.")
-    print("\nOption 1: Monkey Patch (Recommended)")
-    print("-" * 40)
-    print("Add to your script:")
-    print("""
-    from lerobot.utils.profiling import enable_profiling, print_profiling_summary
-    from examples.rtc.add_rtc_profiling import monkey_patch_rtc_profiling
-
-    # Enable profiling
-    enable_profiling()
-    monkey_patch_rtc_profiling()
-
-    # ... run your code ...
-
-    # Print results
-    print_profiling_summary()
-    """)
-    
-    print("\nOption 2: Manual Source Modification")
-    print("-" * 40)
-    print("1. Copy profile_denoise_step() from this file")
-    print("2. Replace denoise_step() in src/lerobot/policies/rtc/modeling_rtc.py")
-    print("3. Add profiling imports at top of file")
-    
-    print("\nKey Metrics to Watch:")
-    print("-" * 40)
-    print("- rtc.denoise_step.base_denoising     - Time for base policy inference")
-    print("- rtc.denoise_step.autograd_correction - Time computing gradients")
-    print("- rtc.denoise_step.guidance_computation - Total guidance overhead")
-    print("- rtc.denoise_step.get_prefix_weights  - Time computing weights")
-    print("="*80 + "\n")
-
-
-if __name__ == "__main__":
-    print_usage()
-

examples/rtc/eval_with_real_robot.py

@@ -1,549 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Demo script showing how to use Real-Time Chunking (RTC) with action chunking policies on real robots.
-
-This script demonstrates:
-1. Creating a robot and policy (SmolVLA, Pi0, etc.) with RTC
-2. Consuming actions from the policy while the robot executes
-3. Periodically requesting new action chunks in the background using threads
-4. Managing action buffers and timing for real-time operation
-
-For simulation environments, see eval_with_simulation.py
-
-Usage:
-    # Run RTC with Real robot with RTC
-    uv run examples/rtc/eval_with_real_robot.py \
-        --policy.path=helper2424/smolvla_check_rtc_last3 \
-        --policy.device=mps \
-        --rtc.enabled=true \
-        --rtc.execution_horizon=20 \
-        --robot.type=so100_follower \
-        --robot.port=/dev/tty.usbmodem58FA0834591 \
-        --robot.id=so100_follower \
-        --robot.cameras="{ gripper: {type: opencv, index_or_path: 1, width: 640, height: 480, fps: 30}, front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
-        --task="Move green small object into the purple platform" \
-        --duration=120
-
-    # Run RTC with Real robot without RTC
-    uv run examples/rtc/eval_with_real_robot.py \
-        --policy.path=helper2424/smolvla_check_rtc_last3 \
-        --policy.device=mps \
-        --rtc.enabled=false \
-        --robot.type=so100_follower \
-        --robot.port=/dev/tty.usbmodem58FA0834591 \
-        --robot.id=so100_follower \
-        --robot.cameras="{ gripper: {type: opencv, index_or_path: 1, width: 640, height: 480, fps: 30}, front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
-        --task="Move green small object into the purple platform" \
-        --duration=120
-
-    # Run RTC with Real robot with pi0.5 policy
-    uv run examples/rtc/eval_with_real_robot.py \
-        --policy.path=helper2424/pi05_check_rtc \
-        --policy.device=mps \
-        --rtc.enabled=true \
-        --rtc.execution_horizon=20 \
-        --robot.type=so100_follower \
-        --robot.port=/dev/tty.usbmodem58FA0834591 \
-        --robot.id=so100_follower \
-        --robot.cameras="{ gripper: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}, front: {type: opencv, index_or_path: 1, width: 640, height: 480, fps: 30}}" \
-        --task="Move green small object into the purple platform" \
-        --duration=120
-"""
-
-import logging
-import math
-import sys
-import time
-import traceback
-from dataclasses import dataclass, field
-from threading import Event, Lock, Thread
-
-import torch
-from torch import Tensor
-
-from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
-from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
-from lerobot.configs import parser
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import RTCAttentionSchedule
-from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
-from lerobot.policies.factory import get_policy_class, make_pre_post_processors
-from lerobot.policies.rtc.action_queue import ActionQueue
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.policies.rtc.latency_tracker import LatencyTracker
-from lerobot.processor.factory import (
-    make_default_robot_action_processor,
-    make_default_robot_observation_processor,
-)
-from lerobot.rl.process import ProcessSignalHandler
-from lerobot.robots import (  # noqa: F401
-    Robot,
-    RobotConfig,
-    koch_follower,
-    so100_follower,
-    so101_follower,
-)
-from lerobot.robots.utils import make_robot_from_config
-from lerobot.utils.constants import OBS_IMAGES
-from lerobot.utils.hub import HubMixin
-from lerobot.utils.utils import init_logging
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-class RobotWrapper:
-    def __init__(self, robot: Robot):
-        self.robot = robot
-        self.lock = Lock()
-
-    def get_observation(self) -> dict[str, Tensor]:
-        with self.lock:
-            return self.robot.get_observation()
-
-    def send_action(self, action: Tensor):
-        with self.lock:
-            self.robot.send_action(action)
-
-    def observation_features(self) -> list[str]:
-        with self.lock:
-            return self.robot.observation_features
-
-    def action_features(self) -> list[str]:
-        with self.lock:
-            return self.robot.action_features
-
-
-@dataclass
-class RTCDemoConfig(HubMixin):
-    """Configuration for RTC demo with action chunking policies and real robots."""
-
-    # Policy configuration
-    policy: PreTrainedConfig | None = None
-
-    # Robot configuration
-    robot: RobotConfig | None = None
-
-    # RTC configuration
-    rtc: RTCConfig = field(
-        default_factory=lambda: RTCConfig(
-            execution_horizon=10,
-            max_guidance_weight=1.0,
-            prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        )
-    )
-
-    # Demo parameters
-    duration: float = 30.0  # Duration to run the demo (seconds)
-    fps: float = 10.0  # Action execution frequency (Hz)
-
-    # Compute device
-    device: str | None = None  # Device to run on (cuda, cpu, auto)
-
-    # Get new actions horizon. The amount of executed steps after which will be requested new actions.
-    # It should be higher than inference delay + execution horizon.
-    action_queue_size_to_get_new_actions: int = 30
-
-    # Task to execute
-    task: str = field(default="", metadata={"help": "Task to execute"})
-
-    # Torch compile configuration
-    use_torch_compile: bool = field(
-        default=False,
-        metadata={"help": "Use torch.compile for faster inference (PyTorch 2.0+)"},
-    )
-
-    torch_compile_backend: str = field(
-        default="inductor",
-        metadata={"help": "Backend for torch.compile (inductor, aot_eager, cudagraphs)"},
-    )
-
-    torch_compile_mode: str = field(
-        default="default",
-        metadata={"help": "Compilation mode (default, reduce-overhead, max-autotune)"},
-    )
-
-    torch_compile_disable_cudagraphs: bool = field(
-        default=True,
-        metadata={
-            "help": "Disable CUDA graphs in torch.compile. Required due to in-place tensor "
-            "operations in denoising loop (x_t += dt * v_t) which cause tensor aliasing issues."
-        },
-    )
-
-    def __post_init__(self):
-        # HACK: We parse again the cli args here to get the pretrained path if there was one.
-        policy_path = parser.get_path_arg("policy")
-        if policy_path:
-            cli_overrides = parser.get_cli_overrides("policy")
-            self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
-            self.policy.pretrained_path = policy_path
-        else:
-            raise ValueError("Policy path is required")
-
-        # Validate that robot configuration is provided
-        if self.robot is None:
-            raise ValueError("Robot configuration must be provided")
-
-    @classmethod
-    def __get_path_fields__(cls) -> list[str]:
-        """This enables the parser to load config from the policy using `--policy.path=local/dir`"""
-        return ["policy"]
-
-
-def is_image_key(k: str) -> bool:
-    return k.startswith(OBS_IMAGES)
-
-
-def get_actions(
-    policy,
-    robot: RobotWrapper,
-    robot_observation_processor,
-    action_queue: ActionQueue,
-    shutdown_event: Event,
-    cfg: RTCDemoConfig,
-):
-    """Thread function to request action chunks from the policy.
-
-    Args:
-        policy: The policy instance (SmolVLA, Pi0, etc.)
-        robot: The robot instance for getting observations
-        robot_observation_processor: Processor for raw robot observations
-        action_queue: Queue to put new action chunks
-        shutdown_event: Event to signal shutdown
-        cfg: Demo configuration
-    """
-    try:
-        logger.info("[GET_ACTIONS] Starting get actions thread")
-
-        latency_tracker = LatencyTracker()  # Track latency of action chunks
-        fps = cfg.fps
-        time_per_chunk = 1.0 / fps
-
-        dataset_features = hw_to_dataset_features(robot.observation_features(), "observation")
-        policy_device = policy.config.device
-
-        # Load preprocessor and postprocessor from pretrained files
-        # The stats are embedded in the processor .safetensors files
-        logger.info(f"[GET_ACTIONS] Loading preprocessor/postprocessor from {cfg.policy.pretrained_path}")
-
-        preprocessor, postprocessor = make_pre_post_processors(
-            policy_cfg=cfg.policy,
-            pretrained_path=cfg.policy.pretrained_path,
-            dataset_stats=None,  # Will load from pretrained processor files
-            preprocessor_overrides={
-                "device_processor": {"device": cfg.policy.device},
-            },
-        )
-
-        logger.info("[GET_ACTIONS] Preprocessor/postprocessor loaded successfully with embedded stats")
-
-        get_actions_threshold = cfg.action_queue_size_to_get_new_actions
-
-        if not cfg.rtc.enabled:
-            get_actions_threshold = 0
-
-        while not shutdown_event.is_set():
-            if action_queue.qsize() <= get_actions_threshold:
-                current_time = time.perf_counter()
-                action_index_before_inference = action_queue.get_action_index()
-                prev_actions = action_queue.get_left_over()
-
-                inference_latency = latency_tracker.max()
-                inference_delay = math.ceil(inference_latency / time_per_chunk)
-
-                obs = robot.get_observation()
-
-                # Apply robot observation processor
-                obs_processed = robot_observation_processor(obs)
-
-                obs_with_policy_features = build_dataset_frame(
-                    dataset_features, obs_processed, prefix="observation"
-                )
-
-                for name in obs_with_policy_features:
-                    obs_with_policy_features[name] = torch.from_numpy(obs_with_policy_features[name])
-                    if "image" in name:
-                        obs_with_policy_features[name] = (
-                            obs_with_policy_features[name].type(torch.float32) / 255
-                        )
-                        obs_with_policy_features[name] = (
-                            obs_with_policy_features[name].permute(2, 0, 1).contiguous()
-                        )
-                    obs_with_policy_features[name] = obs_with_policy_features[name].unsqueeze(0)
-                    obs_with_policy_features[name] = obs_with_policy_features[name].to(policy_device)
-
-                obs_with_policy_features["task"] = [cfg.task]  # Task should be a list, not a string!
-                obs_with_policy_features["robot_type"] = (
-                    robot.robot.name if hasattr(robot.robot, "name") else ""
-                )
-
-                preproceseded_obs = preprocessor(obs_with_policy_features)
-
-                # Generate actions WITH RTC
-                actions = policy.predict_action_chunk(
-                    preproceseded_obs,
-                    inference_delay=inference_delay,
-                    prev_chunk_left_over=prev_actions,
-                )
-
-                # Store original actions (before postprocessing) for RTC
-                original_actions = actions.squeeze(0).clone()
-
-                postprocessed_actions = postprocessor(actions)
-
-                postprocessed_actions = postprocessed_actions.squeeze(0)
-
-                new_latency = time.perf_counter() - current_time
-                new_delay = math.ceil(new_latency / time_per_chunk)
-                latency_tracker.add(new_latency)
-
-                if cfg.action_queue_size_to_get_new_actions < cfg.rtc.execution_horizon + new_delay:
-                    logger.warning(
-                        "[GET_ACTIONS] cfg.action_queue_size_to_get_new_actions Too small, It should be higher than inference delay + execution horizon."
-                    )
-
-                action_queue.merge(
-                    original_actions, postprocessed_actions, new_delay, action_index_before_inference
-                )
-            else:
-                # Small sleep to prevent busy waiting
-                time.sleep(0.1)
-
-        logger.info("[GET_ACTIONS] get actions thread shutting down")
-    except Exception as e:
-        logger.error(f"[GET_ACTIONS] Fatal exception in get_actions thread: {e}")
-        logger.error(traceback.format_exc())
-        sys.exit(1)
-
-
-def actor_control(
-    robot: RobotWrapper,
-    robot_action_processor,
-    action_queue: ActionQueue,
-    shutdown_event: Event,
-    cfg: RTCDemoConfig,
-):
-    """Thread function to execute actions on the robot.
-
-    Args:
-        robot: The robot instance
-        action_queue: Queue to get actions from
-        shutdown_event: Event to signal shutdown
-        cfg: Demo configuration
-    """
-    try:
-        logger.info("[ACTOR] Starting actor thread")
-
-        action_count = 0
-        action_interval = 1.0 / cfg.fps
-
-        while not shutdown_event.is_set():
-            start_time = time.perf_counter()
-
-            # Try to get an action from the queue with timeout
-            action = action_queue.get()
-
-            if action is not None:
-                action = action.cpu()
-                action_dict = {key: action[i].item() for i, key in enumerate(robot.action_features())}
-                action_processed = robot_action_processor((action_dict, None))
-                robot.send_action(action_processed)
-
-                action_count += 1
-
-            dt_s = time.perf_counter() - start_time
-            time.sleep(max(0, (action_interval - dt_s) - 0.001))
-
-        logger.info(f"[ACTOR] Actor thread shutting down. Total actions executed: {action_count}")
-    except Exception as e:
-        logger.error(f"[ACTOR] Fatal exception in actor_control thread: {e}")
-        logger.error(traceback.format_exc())
-        sys.exit(1)
-
-
-def _apply_torch_compile(policy, cfg: RTCDemoConfig):
-    """Apply torch.compile to the policy's predict_action_chunk method.
-
-    Args:
-        policy: Policy instance to compile
-        cfg: Configuration containing torch compile settings
-
-    Returns:
-        Policy with compiled predict_action_chunk method
-    """
-
-    # PI models handle their own compilation
-    if policy.type == "pi05" or policy.type == "pi0":
-        return policy
-
-    try:
-        # Check if torch.compile is available (PyTorch 2.0+)
-        if not hasattr(torch, "compile"):
-            logger.warning(
-                f"torch.compile is not available. Requires PyTorch 2.0+. "
-                f"Current version: {torch.__version__}. Skipping compilation."
-            )
-            return policy
-
-        logger.info("Applying torch.compile to predict_action_chunk...")
-        logger.info(f"  Backend: {cfg.torch_compile_backend}")
-        logger.info(f"  Mode: {cfg.torch_compile_mode}")
-        logger.info(f"  Disable CUDA graphs: {cfg.torch_compile_disable_cudagraphs}")
-
-        # Compile the predict_action_chunk method
-        # - CUDA graphs disabled to prevent tensor aliasing from in-place ops (x_t += dt * v_t)
-        compile_kwargs = {
-            "backend": cfg.torch_compile_backend,
-            "mode": cfg.torch_compile_mode,
-        }
-
-        # Disable CUDA graphs if requested (prevents tensor aliasing issues)
-        if cfg.torch_compile_disable_cudagraphs:
-            compile_kwargs["options"] = {"triton.cudagraphs": False}
-
-        original_method = policy.predict_action_chunk
-        compiled_method = torch.compile(original_method, **compile_kwargs)
-        policy.predict_action_chunk = compiled_method
-        logger.info("✓ Successfully compiled predict_action_chunk")
-
-    except Exception as e:
-        logger.error(f"Failed to apply torch.compile: {e}")
-        logger.warning("Continuing without torch.compile")
-
-    return policy
-
-
-@parser.wrap()
-def demo_cli(cfg: RTCDemoConfig):
-    """Main entry point for RTC demo with draccus configuration."""
-
-    # Initialize logging
-    init_logging()
-
-    logger.info(f"Using device: {cfg.device}")
-
-    # Setup signal handler for graceful shutdown
-    signal_handler = ProcessSignalHandler(use_threads=True, display_pid=False)
-    shutdown_event = signal_handler.shutdown_event
-
-    policy = None
-    robot = None
-    get_actions_thread = None
-    actor_thread = None
-
-    policy_class = get_policy_class(cfg.policy.type)
-
-    # Load config and set compile_model for pi0/pi05 models
-    config = PreTrainedConfig.from_pretrained(cfg.policy.pretrained_path)
-
-    if cfg.policy.type == "pi05" or cfg.policy.type == "pi0":
-        config.compile_model = cfg.use_torch_compile
-
-    policy = policy_class.from_pretrained(cfg.policy.pretrained_path, config=config)
-
-    # Turn on RTC
-    policy.config.rtc_config = cfg.rtc
-
-    # Init RTC processort, as by default if RTC disabled in the config
-    # The processor won't be created
-    policy.init_rtc_processor()
-
-    assert policy.name in ["smolvla", "pi05", "pi0"], "Only smolvla, pi05, and pi0 are supported for RTC"
-
-    policy = policy.to(cfg.device)
-    policy.eval()
-
-    # Apply torch.compile to predict_action_chunk method if enabled
-    if cfg.use_torch_compile:
-        policy = _apply_torch_compile(policy, cfg)
-
-    # Create robot
-    logger.info(f"Initializing robot: {cfg.robot.type}")
-    robot = make_robot_from_config(cfg.robot)
-    robot.connect()
-    robot_wrapper = RobotWrapper(robot)
-
-    # Create robot observation processor
-    robot_observation_processor = make_default_robot_observation_processor()
-    robot_action_processor = make_default_robot_action_processor()
-
-    # Create action queue for communication between threads
-    action_queue = ActionQueue(cfg.rtc)
-
-    # Start chunk requester thread
-    get_actions_thread = Thread(
-        target=get_actions,
-        args=(policy, robot_wrapper, robot_observation_processor, action_queue, shutdown_event, cfg),
-        daemon=True,
-        name="GetActions",
-    )
-    get_actions_thread.start()
-    logger.info("Started get actions thread")
-
-    # Start action executor thread
-    actor_thread = Thread(
-        target=actor_control,
-        args=(robot_wrapper, robot_action_processor, action_queue, shutdown_event, cfg),
-        daemon=True,
-        name="Actor",
-    )
-    actor_thread.start()
-    logger.info("Started actor thread")
-
-    logger.info("Started stop by duration thread")
-
-    # Main thread monitors for duration or shutdown
-    logger.info(f"Running demo for {cfg.duration} seconds...")
-    start_time = time.time()
-
-    while not shutdown_event.is_set() and (time.time() - start_time) < cfg.duration:
-        time.sleep(10)
-
-        # Log queue status periodically
-        if int(time.time() - start_time) % 5 == 0:
-            logger.info(f"[MAIN] Action queue size: {action_queue.qsize()}")
-
-        if time.time() - start_time > cfg.duration:
-            break
-
-    logger.info("Demo duration reached or shutdown requested")
-
-    # Signal shutdown
-    shutdown_event.set()
-
-    # Wait for threads to finish
-    if get_actions_thread and get_actions_thread.is_alive():
-        logger.info("Waiting for chunk requester thread to finish...")
-        get_actions_thread.join()
-
-    if actor_thread and actor_thread.is_alive():
-        logger.info("Waiting for action executor thread to finish...")
-        actor_thread.join()
-
-    # Cleanup robot
-    if robot:
-        robot.disconnect()
-        logger.info("Robot disconnected")
-
-    logger.info("Cleanup completed")
-
-
-if __name__ == "__main__":
-    demo_cli()
-    logging.info("RTC demo finished")

examples/rtc/eval_with_real_robot_profiled.py

@@ -1,631 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Profiled version of eval_with_real_robot.py for performance analysis.
-
-This version adds detailed timing measurements for:
-- Policy inference
-- Preprocessing
-- Postprocessing
-- Action queue operations
-- Robot communication
-- Thread execution times
-
-Usage: Same as eval_with_real_robot.py but with profiling output.
-"""
-
-import logging
-import math
-import sys
-import time
-import traceback
-from collections import defaultdict
-from dataclasses import dataclass, field
-from threading import Event, Lock, Thread
-
-import torch
-from torch import Tensor
-
-from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
-from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
-from lerobot.configs import parser
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import RTCAttentionSchedule
-from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
-from lerobot.policies.factory import get_policy_class, make_pre_post_processors
-from lerobot.policies.rtc.action_queue import ActionQueue
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.policies.rtc.latency_tracker import LatencyTracker
-from lerobot.processor.factory import (
-    make_default_robot_action_processor,
-    make_default_robot_observation_processor,
-)
-from lerobot.rl.process import ProcessSignalHandler
-from lerobot.robots import (  # noqa: F401
-    Robot,
-    RobotConfig,
-    koch_follower,
-    so100_follower,
-    so101_follower,
-)
-from lerobot.robots.utils import make_robot_from_config
-from lerobot.utils.constants import OBS_IMAGES
-from lerobot.utils.hub import HubMixin
-from lerobot.utils.utils import init_logging
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-class ProfileTimer:
-    """Context manager and utility class for timing code sections."""
-
-    def __init__(self, name: str, stats_dict: dict):
-        self.name = name
-        self.stats_dict = stats_dict
-        self.start_time = None
-
-    def __enter__(self):
-        self.start_time = time.perf_counter()
-        return self
-
-    def __exit__(self, *args):
-        elapsed = time.perf_counter() - self.start_time
-        if self.name not in self.stats_dict:
-            self.stats_dict[self.name] = []
-        self.stats_dict[self.name].append(elapsed)
-
-
-class ProfilingStats:
-    """Global profiling statistics collector."""
-
-    def __init__(self):
-        self.stats = defaultdict(list)
-        self.lock = Lock()
-
-    def record(self, name: str, duration: float):
-        with self.lock:
-            self.stats[name].append(duration)
-
-    def timer(self, name: str):
-        """Return a context manager for timing."""
-        return ProfileTimer(name, self.stats)
-
-    def get_summary(self) -> dict[str, dict[str, float]]:
-        """Get summary statistics for all timings."""
-        with self.lock:
-            summary = {}
-            for name, times in self.stats.items():
-                if times:
-                    summary[name] = {
-                        "count": len(times),
-                        "mean": sum(times) / len(times),
-                        "min": min(times),
-                        "max": max(times),
-                        "total": sum(times),
-                    }
-            return summary
-
-    def print_summary(self):
-        """Print formatted summary of all timings."""
-        summary = self.get_summary()
-        
-        logger.info("\n" + "=" * 80)
-        logger.info("PROFILING SUMMARY")
-        logger.info("=" * 80)
-        
-        # Sort by total time (descending)
-        sorted_items = sorted(summary.items(), key=lambda x: x[1]["total"], reverse=True)
-        
-        for name, stats in sorted_items:
-            logger.info(f"\n{name}:")
-            logger.info(f"  Count:     {stats['count']}")
-            logger.info(f"  Mean:      {stats['mean']*1000:.2f} ms")
-            logger.info(f"  Min:       {stats['min']*1000:.2f} ms")
-            logger.info(f"  Max:       {stats['max']*1000:.2f} ms")
-            logger.info(f"  Total:     {stats['total']:.2f} s")
-            logger.info(f"  Hz:        {stats['count']/stats['total']:.2f}")
-        
-        logger.info("\n" + "=" * 80)
-
-
-# Global profiling stats
-profiling_stats = ProfilingStats()
-
-
-class RobotWrapper:
-    def __init__(self, robot: Robot):
-        self.robot = robot
-        self.lock = Lock()
-
-    def get_observation(self) -> dict[str, Tensor]:
-        with profiling_stats.timer("robot.get_observation"):
-            with self.lock:
-                return self.robot.get_observation()
-
-    def send_action(self, action: Tensor):
-        with profiling_stats.timer("robot.send_action"):
-            with self.lock:
-                self.robot.send_action(action)
-
-    def observation_features(self) -> list[str]:
-        with self.lock:
-            return self.robot.observation_features
-
-    def action_features(self) -> list[str]:
-        with self.lock:
-            return self.robot.action_features
-
-
-@dataclass
-class RTCDemoConfig(HubMixin):
-    """Configuration for RTC demo with action chunking policies and real robots."""
-
-    # Policy configuration
-    policy: PreTrainedConfig | None = None
-
-    # Robot configuration
-    robot: RobotConfig | None = None
-
-    # RTC configuration
-    rtc: RTCConfig = field(
-        default_factory=lambda: RTCConfig(
-            execution_horizon=10,
-            max_guidance_weight=1.0,
-            prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        )
-    )
-
-    # Demo parameters
-    duration: float = 30.0  # Duration to run the demo (seconds)
-    fps: float = 10.0  # Action execution frequency (Hz)
-
-    # Compute device
-    device: str | None = None  # Device to run on (cuda, cpu, auto)
-
-    # Get new actions horizon. The amount of executed steps after which will be requested new actions.
-    # It should be higher than inference delay + execution horizon.
-    action_queue_size_to_get_new_actions: int = 30
-
-    # Task to execute
-    task: str = field(default="", metadata={"help": "Task to execute"})
-
-    # Torch compile configuration
-    use_torch_compile: bool = field(
-        default=False,
-        metadata={"help": "Use torch.compile for faster inference (PyTorch 2.0+)"},
-    )
-
-    torch_compile_backend: str = field(
-        default="inductor",
-        metadata={"help": "Backend for torch.compile (inductor, aot_eager, cudagraphs)"},
-    )
-
-    torch_compile_mode: str = field(
-        default="default",
-        metadata={"help": "Compilation mode (default, reduce-overhead, max-autotune)"},
-    )
-
-    torch_compile_disable_cudagraphs: bool = field(
-        default=True,
-        metadata={
-            "help": "Disable CUDA graphs in torch.compile. Required due to in-place tensor "
-            "operations in denoising loop (x_t += dt * v_t) which cause tensor aliasing issues."
-        },
-    )
-
-    def __post_init__(self):
-        # HACK: We parse again the cli args here to get the pretrained path if there was one.
-        policy_path = parser.get_path_arg("policy")
-        if policy_path:
-            cli_overrides = parser.get_cli_overrides("policy")
-            self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
-            self.policy.pretrained_path = policy_path
-        else:
-            raise ValueError("Policy path is required")
-
-        # Validate that robot configuration is provided
-        if self.robot is None:
-            raise ValueError("Robot configuration must be provided")
-
-    @classmethod
-    def __get_path_fields__(cls) -> list[str]:
-        """This enables the parser to load config from the policy using `--policy.path=local/dir`"""
-        return ["policy"]
-
-
-def is_image_key(k: str) -> bool:
-    return k.startswith(OBS_IMAGES)
-
-
-def get_actions(
-    policy,
-    robot: RobotWrapper,
-    robot_observation_processor,
-    action_queue: ActionQueue,
-    shutdown_event: Event,
-    cfg: RTCDemoConfig,
-):
-    """Thread function to request action chunks from the policy with profiling.
-
-    Args:
-        policy: The policy instance (SmolVLA, Pi0, etc.)
-        robot: The robot instance for getting observations
-        robot_observation_processor: Processor for raw robot observations
-        action_queue: Queue to put new action chunks
-        shutdown_event: Event to signal shutdown
-        cfg: Demo configuration
-    """
-    try:
-        logger.info("[GET_ACTIONS] Starting get actions thread")
-
-        latency_tracker = LatencyTracker()  # Track latency of action chunks
-        fps = cfg.fps
-        time_per_chunk = 1.0 / fps
-
-        dataset_features = hw_to_dataset_features(robot.observation_features(), "observation")
-        policy_device = policy.config.device
-
-        # Load preprocessor and postprocessor from pretrained files
-        logger.info(f"[GET_ACTIONS] Loading preprocessor/postprocessor from {cfg.policy.pretrained_path}")
-
-        preprocessor, postprocessor = make_pre_post_processors(
-            policy_cfg=cfg.policy,
-            pretrained_path=cfg.policy.pretrained_path,
-            dataset_stats=None,  # Will load from pretrained processor files
-            preprocessor_overrides={
-                "device_processor": {"device": cfg.policy.device},
-            },
-        )
-
-        logger.info("[GET_ACTIONS] Preprocessor/postprocessor loaded successfully with embedded stats")
-
-        get_actions_threshold = cfg.action_queue_size_to_get_new_actions
-
-        if not cfg.rtc.enabled:
-            get_actions_threshold = 0
-
-        inference_count = 0
-
-        while not shutdown_event.is_set():
-            if action_queue.qsize() <= get_actions_threshold:
-                with profiling_stats.timer("get_actions.total_iteration"):
-                    inference_count += 1
-                    logger.info(f"[GET_ACTIONS] Starting inference #{inference_count}")
-
-                    current_time = time.perf_counter()
-                    action_index_before_inference = action_queue.get_action_index()
-                    
-                    with profiling_stats.timer("get_actions.get_prev_actions"):
-                        prev_actions = action_queue.get_left_over()
-
-                    inference_latency = latency_tracker.max()
-                    inference_delay = math.ceil(inference_latency / time_per_chunk)
-
-                    # Get observation
-                    obs = robot.get_observation()
-
-                    # Apply robot observation processor
-                    with profiling_stats.timer("get_actions.robot_obs_processing"):
-                        obs_processed = robot_observation_processor(obs)
-
-                    # Build dataset frame
-                    with profiling_stats.timer("get_actions.build_dataset_frame"):
-                        obs_with_policy_features = build_dataset_frame(
-                            dataset_features, obs_processed, prefix="observation"
-                        )
-
-                    # Convert to tensors and normalize
-                    with profiling_stats.timer("get_actions.tensor_conversion"):
-                        for name in obs_with_policy_features:
-                            obs_with_policy_features[name] = torch.from_numpy(obs_with_policy_features[name])
-                            if "image" in name:
-                                obs_with_policy_features[name] = (
-                                    obs_with_policy_features[name].type(torch.float32) / 255
-                                )
-                                obs_with_policy_features[name] = (
-                                    obs_with_policy_features[name].permute(2, 0, 1).contiguous()
-                                )
-                            obs_with_policy_features[name] = obs_with_policy_features[name].unsqueeze(0)
-                            obs_with_policy_features[name] = obs_with_policy_features[name].to(policy_device)
-
-                        obs_with_policy_features["task"] = [cfg.task]
-                        obs_with_policy_features["robot_type"] = (
-                            robot.robot.name if hasattr(robot.robot, "name") else ""
-                        )
-
-                    # Preprocessing
-                    with profiling_stats.timer("get_actions.preprocessing"):
-                        preproceseded_obs = preprocessor(obs_with_policy_features)
-
-                    # Policy inference
-                    with profiling_stats.timer("get_actions.policy_inference"):
-                        actions = policy.predict_action_chunk(
-                            preproceseded_obs,
-                            inference_delay=inference_delay,
-                            prev_chunk_left_over=prev_actions,
-                        )
-
-                    # Clone for RTC
-                    with profiling_stats.timer("get_actions.clone_actions"):
-                        original_actions = actions.squeeze(0).clone()
-
-                    # Postprocessing
-                    with profiling_stats.timer("get_actions.postprocessing"):
-                        postprocessed_actions = postprocessor(actions)
-                        postprocessed_actions = postprocessed_actions.squeeze(0)
-
-                    # Update latency tracker
-                    new_latency = time.perf_counter() - current_time
-                    new_delay = math.ceil(new_latency / time_per_chunk)
-                    latency_tracker.add(new_latency)
-
-                    logger.info(
-                        f"[GET_ACTIONS] Inference #{inference_count} completed in {new_latency*1000:.2f}ms "
-                        f"(delay={new_delay} chunks)"
-                    )
-
-                    if cfg.action_queue_size_to_get_new_actions < cfg.rtc.execution_horizon + new_delay:
-                        logger.warning(
-                            "[GET_ACTIONS] cfg.action_queue_size_to_get_new_actions Too small, "
-                            "It should be higher than inference delay + execution horizon."
-                        )
-
-                    # Merge into action queue
-                    with profiling_stats.timer("get_actions.action_queue_merge"):
-                        action_queue.merge(
-                            original_actions, postprocessed_actions, new_delay, action_index_before_inference
-                        )
-            else:
-                # Small sleep to prevent busy waiting
-                time.sleep(0.1)
-
-        logger.info("[GET_ACTIONS] get actions thread shutting down")
-    except Exception as e:
-        logger.error(f"[GET_ACTIONS] Fatal exception in get_actions thread: {e}")
-        logger.error(traceback.format_exc())
-        sys.exit(1)
-
-
-def actor_control(
-    robot: RobotWrapper,
-    robot_action_processor,
-    action_queue: ActionQueue,
-    shutdown_event: Event,
-    cfg: RTCDemoConfig,
-):
-    """Thread function to execute actions on the robot with profiling.
-
-    Args:
-        robot: The robot instance
-        action_queue: Queue to get actions from
-        shutdown_event: Event to signal shutdown
-        cfg: Demo configuration
-    """
-    try:
-        logger.info("[ACTOR] Starting actor thread")
-
-        action_count = 0
-        action_interval = 1.0 / cfg.fps
-
-        while not shutdown_event.is_set():
-            start_time = time.perf_counter()
-
-            with profiling_stats.timer("actor.total_iteration"):
-                # Get action from queue
-                with profiling_stats.timer("actor.queue_get"):
-                    action = action_queue.get()
-
-                if action is not None:
-                    # Process action
-                    with profiling_stats.timer("actor.action_processing"):
-                        action = action.cpu()
-                        action_dict = {key: action[i].item() for i, key in enumerate(robot.action_features())}
-                        action_processed = robot_action_processor((action_dict, None))
-                    
-                    # Send to robot (includes robot.send_action timing)
-                    robot.send_action(action_processed)
-                    action_count += 1
-
-            # Sleep to maintain target FPS
-            dt_s = time.perf_counter() - start_time
-            sleep_time = max(0, (action_interval - dt_s) - 0.001)
-            if sleep_time > 0:
-                time.sleep(sleep_time)
-
-        logger.info(f"[ACTOR] Actor thread shutting down. Total actions executed: {action_count}")
-    except Exception as e:
-        logger.error(f"[ACTOR] Fatal exception in actor_control thread: {e}")
-        logger.error(traceback.format_exc())
-        sys.exit(1)
-
-
-def _apply_torch_compile(policy, cfg: RTCDemoConfig):
-    """Apply torch.compile to the policy's predict_action_chunk method.
-
-    Args:
-        policy: Policy instance to compile
-        cfg: Configuration containing torch compile settings
-
-    Returns:
-        Policy with compiled predict_action_chunk method
-    """
-
-    # PI models handle their own compilation
-    if policy.type == "pi05" or policy.type == "pi0":
-        return policy
-
-    try:
-        # Check if torch.compile is available (PyTorch 2.0+)
-        if not hasattr(torch, "compile"):
-            logger.warning(
-                f"torch.compile is not available. Requires PyTorch 2.0+. "
-                f"Current version: {torch.__version__}. Skipping compilation."
-            )
-            return policy
-
-        logger.info("Applying torch.compile to predict_action_chunk...")
-        logger.info(f"  Backend: {cfg.torch_compile_backend}")
-        logger.info(f"  Mode: {cfg.torch_compile_mode}")
-        logger.info(f"  Disable CUDA graphs: {cfg.torch_compile_disable_cudagraphs}")
-
-        # Compile the predict_action_chunk method
-        compile_kwargs = {
-            "backend": cfg.torch_compile_backend,
-            "mode": cfg.torch_compile_mode,
-        }
-
-        # Disable CUDA graphs if requested (prevents tensor aliasing issues)
-        if cfg.torch_compile_disable_cudagraphs:
-            compile_kwargs["options"] = {"triton.cudagraphs": False}
-
-        original_method = policy.predict_action_chunk
-        compiled_method = torch.compile(original_method, **compile_kwargs)
-        policy.predict_action_chunk = compiled_method
-        logger.info("✓ Successfully compiled predict_action_chunk")
-
-    except Exception as e:
-        logger.error(f"Failed to apply torch.compile: {e}")
-        logger.warning("Continuing without torch.compile")
-
-    return policy
-
-
-@parser.wrap()
-def demo_cli(cfg: RTCDemoConfig):
-    """Main entry point for RTC demo with profiling."""
-
-    # Initialize logging
-    init_logging()
-
-    logger.info(f"Using device: {cfg.device}")
-    logger.info("=" * 80)
-    logger.info("PROFILING MODE ENABLED")
-    logger.info("=" * 80)
-
-    # Setup signal handler for graceful shutdown
-    signal_handler = ProcessSignalHandler(use_threads=True, display_pid=False)
-    shutdown_event = signal_handler.shutdown_event
-
-    policy = None
-    robot = None
-    get_actions_thread = None
-    actor_thread = None
-
-    policy_class = get_policy_class(cfg.policy.type)
-
-    # Load config and set compile_model for pi0/pi05 models
-    config = PreTrainedConfig.from_pretrained(cfg.policy.pretrained_path)
-
-    if cfg.policy.type == "pi05" or cfg.policy.type == "pi0":
-        config.compile_model = cfg.use_torch_compile
-
-    policy = policy_class.from_pretrained(cfg.policy.pretrained_path, config=config)
-
-    # Turn on RTC
-    policy.config.rtc_config = cfg.rtc
-
-    # Init RTC processor
-    policy.init_rtc_processor()
-
-    assert policy.name in ["smolvla", "pi05", "pi0"], "Only smolvla, pi05, and pi0 are supported for RTC"
-
-    policy = policy.to(cfg.device)
-    policy.eval()
-
-    # Apply torch.compile to predict_action_chunk method if enabled
-    if cfg.use_torch_compile:
-        policy = _apply_torch_compile(policy, cfg)
-
-    # Create robot
-    logger.info(f"Initializing robot: {cfg.robot.type}")
-    robot = make_robot_from_config(cfg.robot)
-    robot.connect()
-    robot_wrapper = RobotWrapper(robot)
-
-    # Create robot observation processor
-    robot_observation_processor = make_default_robot_observation_processor()
-    robot_action_processor = make_default_robot_action_processor()
-
-    # Create action queue for communication between threads
-    action_queue = ActionQueue(cfg.rtc)
-
-    # Start chunk requester thread
-    get_actions_thread = Thread(
-        target=get_actions,
-        args=(policy, robot_wrapper, robot_observation_processor, action_queue, shutdown_event, cfg),
-        daemon=True,
-        name="GetActions",
-    )
-    get_actions_thread.start()
-    logger.info("Started get actions thread")
-
-    # Start action executor thread
-    actor_thread = Thread(
-        target=actor_control,
-        args=(robot_wrapper, robot_action_processor, action_queue, shutdown_event, cfg),
-        daemon=True,
-        name="Actor",
-    )
-    actor_thread.start()
-    logger.info("Started actor thread")
-
-    logger.info("Started stop by duration thread")
-
-    # Main thread monitors for duration or shutdown
-    logger.info(f"Running demo for {cfg.duration} seconds...")
-    start_time = time.time()
-
-    while not shutdown_event.is_set() and (time.time() - start_time) < cfg.duration:
-        time.sleep(10)
-
-        # Log queue status periodically
-        if int(time.time() - start_time) % 5 == 0:
-            logger.info(f"[MAIN] Action queue size: {action_queue.qsize()}")
-
-        if time.time() - start_time > cfg.duration:
-            break
-
-    logger.info("Demo duration reached or shutdown requested")
-
-    # Signal shutdown
-    shutdown_event.set()
-
-    # Wait for threads to finish
-    if get_actions_thread and get_actions_thread.is_alive():
-        logger.info("Waiting for chunk requester thread to finish...")
-        get_actions_thread.join()
-
-    if actor_thread and actor_thread.is_alive():
-        logger.info("Waiting for action executor thread to finish...")
-        actor_thread.join()
-
-    # Cleanup robot
-    if robot:
-        robot.disconnect()
-        logger.info("Robot disconnected")
-
-    # Print profiling summary
-    profiling_stats.print_summary()
-
-    logger.info("Cleanup completed")
-
-
-if __name__ == "__main__":
-    demo_cli()
-    logging.info("RTC demo finished")
-

examples/rtc/profile_pi0_rtc_detailed.py

@@ -1,358 +0,0 @@
-#!/usr/bin/env python
-
-"""
-Comprehensive profiling script for Pi0 with RTC.
-
-This script demonstrates how to use all the profiling tools to identify
-bottlenecks in Pi0 policy inference with RTC enabled.
-
-It profiles:
-1. Overall inference time
-2. RTC-specific operations (guidance, weights, etc.)
-3. Preprocessing/postprocessing
-4. Individual method timings
-
-Usage:
-    uv run examples/rtc/profile_pi0_rtc_detailed.py \
-        --policy_path=helper2424/pi05_check_rtc \
-        --device=mps \
-        --num_iterations=20 \
-        --execution_horizon=20 \
-        --enable_rtc_profiling
-"""
-
-import argparse
-import logging
-import sys
-import time
-
-import numpy as np
-import torch
-
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import RTCAttentionSchedule
-from lerobot.policies.factory import get_policy_class, make_pre_post_processors
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.utils.profiling import (
-    ProfileContext,
-    clear_profiling_stats,
-    enable_profiling,
-    get_profiling_stats,
-    print_profiling_summary,
-)
-
-# Import monkey patching for RTC profiling
-try:
-    from examples.rtc.add_rtc_profiling import monkey_patch_rtc_profiling
-except ImportError:
-    logging.warning("Could not import add_rtc_profiling, detailed RTC profiling disabled")
-    monkey_patch_rtc_profiling = None
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-def create_mock_observation(policy_config, device: str) -> dict:
-    """Create a mock observation matching policy requirements.
-    
-    Args:
-        policy_config: Policy configuration
-        device: Device to create tensors on
-        
-    Returns:
-        Mock observation dictionary
-    """
-    obs = {}
-    
-    # Create mock state observation
-    state_dim = 10  # Typical robot state dimension
-    obs["observation.state"] = torch.randn(1, state_dim, device=device)
-    
-    # Create mock images if needed
-    # For Pi0, we typically need at least one image
-    image_height = 224
-    image_width = 224
-    
-    # Common image keys for Pi0
-    image_keys = ["observation.images.gripper", "observation.images.front"]
-    
-    for key in image_keys:
-        # Images should be [B, C, H, W] and normalized to [0, 1]
-        obs[key] = torch.rand(1, 3, image_height, image_width, device=device)
-    
-    # Add task
-    obs["task"] = ["Pick up the object"]
-    
-    # Add language tokens and attention mask (required for Pi0)
-    # These are mock values - in real usage they come from tokenizer
-    max_seq_len = 32
-    obs["observation.language_tokens"] = torch.randint(0, 1000, (1, max_seq_len), device=device)
-    obs["observation.language_attention_mask"] = torch.ones(1, max_seq_len, device=device)
-    
-    return obs
-
-
-def profile_single_iteration(
-    policy,
-    preprocessor,
-    postprocessor,
-    observation: dict,
-    prev_actions: torch.Tensor | None,
-    use_rtc: bool,
-    inference_delay: int = 0,
-) -> tuple[torch.Tensor, torch.Tensor | None, dict]:
-    """Profile a single inference iteration.
-    
-    Args:
-        policy: Policy instance
-        preprocessor: Observation preprocessor
-        postprocessor: Action postprocessor
-        observation: Input observation
-        prev_actions: Previous action chunk (for RTC)
-        use_rtc: Whether RTC is enabled
-        inference_delay: Inference delay in timesteps
-        
-    Returns:
-        Tuple of (actions, new_prev_actions, timings)
-    """
-    timings = {}
-    
-    with ProfileContext("iteration.total"):
-        # Preprocessing
-        with ProfileContext("iteration.preprocessing"):
-            preprocessed_obs = preprocessor(observation)
-        
-        # Policy inference
-        with ProfileContext("iteration.policy_inference"):
-            if use_rtc:
-                actions = policy.predict_action_chunk(
-                    preprocessed_obs,
-                    inference_delay=inference_delay,
-                    prev_chunk_left_over=prev_actions,
-                )
-            else:
-                actions = policy.predict_action_chunk(preprocessed_obs)
-        
-        # Clone for next iteration (if RTC)
-        new_prev_actions = None
-        if use_rtc:
-            with ProfileContext("iteration.prepare_prev_actions"):
-                execution_horizon = policy.config.rtc_config.execution_horizon
-                if actions.shape[1] > execution_horizon:
-                    new_prev_actions = actions[:, execution_horizon:].clone()
-        
-        # Postprocessing
-        with ProfileContext("iteration.postprocessing"):
-            processed_actions = postprocessor(actions)
-    
-    return processed_actions, new_prev_actions, timings
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Detailed profiling for Pi0 with RTC")
-    parser.add_argument("--policy_path", type=str, required=True, help="Path to pretrained policy")
-    parser.add_argument("--device", type=str, default="cuda", help="Device (cuda/cpu/mps)")
-    parser.add_argument("--num_iterations", type=int, default=20, help="Number of iterations")
-    parser.add_argument("--execution_horizon", type=int, default=10, help="RTC execution horizon")
-    parser.add_argument("--warmup_iterations", type=int, default=5, help="Warmup iterations")
-    parser.add_argument("--enable_rtc_profiling", action="store_true", help="Enable detailed RTC profiling")
-    parser.add_argument("--use_torch_compile", action="store_true", help="Use torch.compile")
-    
-    args = parser.parse_args()
-    
-    logger.info("="*80)
-    logger.info("DETAILED PI0 RTC PROFILING")
-    logger.info("="*80)
-    logger.info(f"Policy: {args.policy_path}")
-    logger.info(f"Device: {args.device}")
-    logger.info(f"Iterations: {args.num_iterations}")
-    logger.info(f"Execution Horizon: {args.execution_horizon}")
-    logger.info(f"RTC Profiling: {args.enable_rtc_profiling}")
-    logger.info("="*80 + "\n")
-    
-    # Enable profiling
-    enable_profiling()
-    
-    # Apply RTC profiling if requested
-    if args.enable_rtc_profiling:
-        if monkey_patch_rtc_profiling is not None:
-            monkey_patch_rtc_profiling()
-            logger.info("✓ Detailed RTC profiling enabled\n")
-        else:
-            logger.warning("⚠ Could not enable detailed RTC profiling\n")
-    
-    # Load policy
-    logger.info("Loading policy...")
-    config = PreTrainedConfig.from_pretrained(args.policy_path)
-    
-    if hasattr(config, "compile_model"):
-        config.compile_model = args.use_torch_compile
-    
-    policy_class = get_policy_class(config.type)
-    policy = policy_class.from_pretrained(args.policy_path, config=config)
-    
-    # Configure RTC
-    policy.config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=args.execution_horizon,
-        max_guidance_weight=1.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-    )
-    policy.init_rtc_processor()
-    
-    policy = policy.to(args.device)
-    policy.eval()
-    
-    logger.info(f"✓ Policy loaded: {config.type}\n")
-    
-    # Create preprocessor and postprocessor
-    logger.info("Loading preprocessor/postprocessor...")
-    preprocessor, postprocessor = make_pre_post_processors(
-        policy_cfg=config,
-        pretrained_path=args.policy_path,
-        dataset_stats=None,
-        preprocessor_overrides={
-            "device_processor": {"device": args.device},
-        },
-    )
-    logger.info("✓ Preprocessor/postprocessor loaded\n")
-    
-    # Create mock observation
-    logger.info("Creating mock observation...")
-    observation = create_mock_observation(config, args.device)
-    logger.info("✓ Mock observation created\n")
-    
-    # Warmup
-    logger.info(f"Warming up ({args.warmup_iterations} iterations)...")
-    prev_actions = None
-    for i in range(args.warmup_iterations):
-        with torch.no_grad():
-            _, prev_actions, _ = profile_single_iteration(
-                policy=policy,
-                preprocessor=preprocessor,
-                postprocessor=postprocessor,
-                observation=observation,
-                prev_actions=prev_actions,
-                use_rtc=True,
-                inference_delay=0,
-            )
-    
-    # Clear warmup stats
-    clear_profiling_stats()
-    logger.info("✓ Warmup complete\n")
-    
-    # Profiled run WITH RTC
-    logger.info(f"Running profiled iterations WITH RTC ({args.num_iterations} iterations)...")
-    prev_actions = None
-    iteration_times = []
-    
-    for i in range(args.num_iterations):
-        start = time.perf_counter()
-        
-        with torch.no_grad():
-            _, prev_actions, _ = profile_single_iteration(
-                policy=policy,
-                preprocessor=preprocessor,
-                postprocessor=postprocessor,
-                observation=observation,
-                prev_actions=prev_actions,
-                use_rtc=True,
-                inference_delay=0,
-            )
-        
-        # Sync CUDA if needed
-        if args.device.startswith("cuda"):
-            torch.cuda.synchronize()
-        
-        elapsed = time.perf_counter() - start
-        iteration_times.append(elapsed)
-        
-        if (i + 1) % 5 == 0:
-            logger.info(f"  Completed {i+1}/{args.num_iterations}")
-    
-    logger.info("✓ Profiling complete\n")
-    
-    # Print summary statistics
-    logger.info("\n" + "="*80)
-    logger.info("ITERATION TIMING SUMMARY")
-    logger.info("="*80)
-    
-    times_arr = np.array(iteration_times)
-    logger.info(f"Mean time:       {np.mean(times_arr)*1000:.2f} ms")
-    logger.info(f"Median time:     {np.median(times_arr)*1000:.2f} ms")
-    logger.info(f"Std dev:         {np.std(times_arr)*1000:.2f} ms")
-    logger.info(f"Min time:        {np.min(times_arr)*1000:.2f} ms")
-    logger.info(f"Max time:        {np.max(times_arr)*1000:.2f} ms")
-    logger.info(f"Total time:      {np.sum(times_arr):.2f} s")
-    logger.info(f"Throughput:      {len(times_arr)/np.sum(times_arr):.2f} iter/s")
-    logger.info("="*80 + "\n")
-    
-    # Print detailed profiling breakdown
-    print_profiling_summary(sort_by="total")
-    
-    # Print key insights
-    stats = get_profiling_stats()
-    
-    logger.info("\n" + "="*80)
-    logger.info("KEY INSIGHTS")
-    logger.info("="*80)
-    
-    # Find bottlenecks
-    if stats:
-        policy_inference_time = stats.get("iteration.policy_inference", {}).get("mean", 0)
-        preprocessing_time = stats.get("iteration.preprocessing", {}).get("mean", 0)
-        postprocessing_time = stats.get("iteration.postprocessing", {}).get("mean", 0)
-        
-        total_time = policy_inference_time + preprocessing_time + postprocessing_time
-        
-        if total_time > 0:
-            logger.info(f"\nTime breakdown:")
-            logger.info(f"  Policy inference:  {policy_inference_time*1000:.2f} ms ({policy_inference_time/total_time*100:.1f}%)")
-            logger.info(f"  Preprocessing:     {preprocessing_time*1000:.2f} ms ({preprocessing_time/total_time*100:.1f}%)")
-            logger.info(f"  Postprocessing:    {postprocessing_time*1000:.2f} ms ({postprocessing_time/total_time*100:.1f}%)")
-        
-        # RTC-specific insights
-        if args.enable_rtc_profiling:
-            rtc_guidance = stats.get("rtc.denoise_step.guidance_computation", {}).get("mean", 0)
-            rtc_autograd = stats.get("rtc.denoise_step.autograd_correction", {}).get("mean", 0)
-            rtc_base = stats.get("rtc.denoise_step.base_denoising", {}).get("mean", 0)
-            
-            if rtc_guidance > 0:
-                logger.info(f"\nRTC breakdown:")
-                logger.info(f"  Base denoising:    {rtc_base*1000:.2f} ms")
-                logger.info(f"  Guidance compute:  {rtc_guidance*1000:.2f} ms")
-                logger.info(f"  Autograd correct:  {rtc_autograd*1000:.2f} ms")
-                logger.info(f"  RTC overhead:      {(rtc_guidance - rtc_base)*1000:.2f} ms")
-        
-        # Recommendations
-        logger.info("\nRecommendations:")
-        
-        if preprocessing_time > policy_inference_time * 0.3:
-            logger.info("  ⚠ Preprocessing is taking >30% of time")
-            logger.info("    → Consider reducing image resolution")
-            logger.info("    → Consider using fewer cameras")
-        
-        if args.enable_rtc_profiling and rtc_autograd > rtc_base * 0.5:
-            logger.info("  ⚠ RTC autograd overhead is significant")
-            logger.info("    → This is expected, but consider increasing execution_horizon")
-            logger.info("    → Try torch.compile if not already enabled")
-        
-        if not args.use_torch_compile:
-            logger.info("  💡 torch.compile not enabled")
-            logger.info("    → Try --use_torch_compile for potential speedup")
-    
-    logger.info("="*80 + "\n")
-
-
-if __name__ == "__main__":
-    try:
-        main()
-    except KeyboardInterrupt:
-        logger.info("\n\nProfiling interrupted by user")
-        sys.exit(0)
-    except Exception as e:
-        logger.error(f"\n\nError during profiling: {e}")
-        import traceback
-        traceback.print_exc()
-        sys.exit(1)
-

examples/rtc/profile_rtc_comparison.py

@@ -1,347 +0,0 @@
-#!/usr/bin/env python
-
-"""
-Script to compare performance with and without RTC enabled.
-
-This script helps identify whether RTC is actually improving or degrading performance
-by running multiple inference passes and collecting detailed timing statistics.
-
-Usage:
-    # Profile with mock data (no robot needed)
-    uv run examples/rtc/profile_rtc_comparison.py \
-        --policy_path=helper2424/pi05_check_rtc \
-        --device=mps \
-        --num_iterations=50
-
-    # Profile with specific RTC config
-    uv run examples/rtc/profile_rtc_comparison.py \
-        --policy_path=helper2424/pi05_check_rtc \
-        --device=mps \
-        --num_iterations=50 \
-        --execution_horizon=20
-"""
-
-import argparse
-import logging
-import time
-from dataclasses import dataclass
-
-import numpy as np
-import torch
-
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import RTCAttentionSchedule
-from lerobot.policies.factory import get_policy_class, make_pre_post_processors
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.utils.profiling import (
-    clear_profiling_stats,
-    enable_profiling,
-    get_profiling_stats,
-    print_profiling_summary,
-)
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class ProfileResults:
-    """Results from profiling run."""
-
-    mode: str  # "with_rtc" or "without_rtc"
-    mean_time: float
-    std_time: float
-    min_time: float
-    max_time: float
-    times: list[float]
-    throughput: float  # iterations per second
-
-
-def create_mock_observation(policy, device: str) -> dict:
-    """Create a mock observation for testing.
-
-    Args:
-        policy: Policy instance
-        device: Device to create tensors on
-
-    Returns:
-        Mock observation dictionary
-    """
-    # Get expected input shapes from policy config
-    # This is a simplified version - adjust based on actual policy requirements
-    obs = {}
-
-    # Mock image observations (if needed)
-    if hasattr(policy.config, "input_shapes"):
-        for key, shape in policy.config.input_shapes.items():
-            if "image" in key:
-                # Typical image shape: (batch, channels, height, width)
-                obs[key] = torch.randn(1, *shape, device=device)
-            else:
-                obs[key] = torch.randn(1, *shape, device=device)
-
-    # Add task if needed
-    if "task" in policy.config.__dict__ or hasattr(policy, "accepts_task"):
-        obs["task"] = ["Pick up the object"]
-
-    # Mock state observation
-    obs["observation.state"] = torch.randn(1, 10, device=device)  # Adjust size as needed
-
-    return obs
-
-
-def profile_inference(
-    policy, observation: dict, num_iterations: int, use_rtc: bool, execution_horizon: int = 10
-) -> ProfileResults:
-    """Profile policy inference with or without RTC.
-
-    Args:
-        policy: Policy instance
-        observation: Observation dictionary
-        num_iterations: Number of inference iterations to run
-        use_rtc: Whether to enable RTC
-        execution_horizon: Execution horizon for RTC
-
-    Returns:
-        ProfileResults with timing statistics
-    """
-    mode = "with_rtc" if use_rtc else "without_rtc"
-    logger.info(f"\n{'='*80}")
-    logger.info(f"Profiling: {mode.upper()}")
-    logger.info(f"{'='*80}")
-
-    # Configure RTC
-    if use_rtc:
-        policy.config.rtc_config.enabled = True
-        policy.config.rtc_config.execution_horizon = execution_horizon
-        policy.init_rtc_processor()
-    else:
-        policy.config.rtc_config.enabled = False
-
-    times = []
-    prev_actions = None
-
-    # Warmup
-    logger.info("Warming up (5 iterations)...")
-    for _ in range(5):
-        with torch.no_grad():
-            if use_rtc:
-                _ = policy.predict_action_chunk(
-                    observation, inference_delay=0, prev_chunk_left_over=prev_actions
-                )
-            else:
-                _ = policy.predict_action_chunk(observation)
-
-    # Actual profiling
-    logger.info(f"Running {num_iterations} profiled iterations...")
-    for i in range(num_iterations):
-        start = time.perf_counter()
-
-        with torch.no_grad():
-            if use_rtc:
-                actions = policy.predict_action_chunk(
-                    observation, inference_delay=0, prev_chunk_left_over=prev_actions
-                )
-                # Simulate consuming some actions for next iteration
-                if actions.shape[1] > execution_horizon:
-                    prev_actions = actions[:, execution_horizon:].clone()
-                else:
-                    prev_actions = None
-            else:
-                actions = policy.predict_action_chunk(observation)
-
-        # Synchronize if using CUDA
-        if observation["observation.state"].device.type == "cuda":
-            torch.cuda.synchronize()
-
-        elapsed = time.perf_counter() - start
-        times.append(elapsed)
-
-        if (i + 1) % 10 == 0:
-            logger.info(f"  Completed {i+1}/{num_iterations} iterations")
-
-    # Calculate statistics
-    times_arr = np.array(times)
-    results = ProfileResults(
-        mode=mode,
-        mean_time=float(np.mean(times_arr)),
-        std_time=float(np.std(times_arr)),
-        min_time=float(np.min(times_arr)),
-        max_time=float(np.max(times_arr)),
-        times=times,
-        throughput=num_iterations / sum(times),
-    )
-
-    logger.info(f"\nResults for {mode}:")
-    logger.info(f"  Mean time: {results.mean_time*1000:.2f} ms")
-    logger.info(f"  Std dev:   {results.std_time*1000:.2f} ms")
-    logger.info(f"  Min time:  {results.min_time*1000:.2f} ms")
-    logger.info(f"  Max time:  {results.max_time*1000:.2f} ms")
-    logger.info(f"  Throughput: {results.throughput:.2f} iter/s")
-
-    return results
-
-
-def compare_results(results_without_rtc: ProfileResults, results_with_rtc: ProfileResults):
-    """Compare and print results from both runs.
-
-    Args:
-        results_without_rtc: Results from run without RTC
-        results_with_rtc: Results from run with RTC
-    """
-    logger.info(f"\n{'='*80}")
-    logger.info("COMPARISON SUMMARY")
-    logger.info(f"{'='*80}")
-
-    mean_diff = results_with_rtc.mean_time - results_without_rtc.mean_time
-    mean_diff_pct = (mean_diff / results_without_rtc.mean_time) * 100
-
-    throughput_diff = results_with_rtc.throughput - results_without_rtc.throughput
-    throughput_diff_pct = (throughput_diff / results_without_rtc.throughput) * 100
-
-    logger.info(f"\n{'Metric':<30} {'Without RTC':>15} {'With RTC':>15} {'Difference':>15}")
-    logger.info("-" * 80)
-    logger.info(
-        f"{'Mean time (ms)':<30} "
-        f"{results_without_rtc.mean_time*1000:>15.2f} "
-        f"{results_with_rtc.mean_time*1000:>15.2f} "
-        f"{mean_diff*1000:>+15.2f}"
-    )
-    logger.info(
-        f"{'Std dev (ms)':<30} "
-        f"{results_without_rtc.std_time*1000:>15.2f} "
-        f"{results_with_rtc.std_time*1000:>15.2f} "
-        f"{(results_with_rtc.std_time - results_without_rtc.std_time)*1000:>+15.2f}"
-    )
-    logger.info(
-        f"{'Min time (ms)':<30} "
-        f"{results_without_rtc.min_time*1000:>15.2f} "
-        f"{results_with_rtc.min_time*1000:>15.2f} "
-        f"{(results_with_rtc.min_time - results_without_rtc.min_time)*1000:>+15.2f}"
-    )
-    logger.info(
-        f"{'Max time (ms)':<30} "
-        f"{results_without_rtc.max_time*1000:>15.2f} "
-        f"{results_with_rtc.max_time*1000:>15.2f} "
-        f"{(results_with_rtc.max_time - results_without_rtc.max_time)*1000:>+15.2f}"
-    )
-    logger.info(
-        f"{'Throughput (iter/s)':<30} "
-        f"{results_without_rtc.throughput:>15.2f} "
-        f"{results_with_rtc.throughput:>15.2f} "
-        f"{throughput_diff:>+15.2f}"
-    )
-
-    logger.info(f"\n{'='*80}")
-    logger.info("VERDICT")
-    logger.info(f"{'='*80}")
-
-    if mean_diff_pct < -5:
-        logger.info(f"✓ RTC is FASTER by {abs(mean_diff_pct):.1f}%")
-        logger.info(f"  Mean time reduced by {abs(mean_diff)*1000:.2f} ms")
-    elif mean_diff_pct > 5:
-        logger.info(f"✗ RTC is SLOWER by {mean_diff_pct:.1f}%")
-        logger.info(f"  Mean time increased by {mean_diff*1000:.2f} ms")
-        logger.info("\n  Possible reasons:")
-        logger.info("  - RTC overhead exceeds benefits at current execution horizon")
-        logger.info("  - Inference delay calculation not accounting for RTC processing")
-        logger.info("  - Additional tensor operations in RTC guidance")
-    else:
-        logger.info(f"≈ Performance is SIMILAR (difference: {mean_diff_pct:+.1f}%)")
-
-    logger.info(f"{'='*80}\n")
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Profile RTC performance")
-    parser.add_argument(
-        "--policy_path", type=str, required=True, help="Path to pretrained policy"
-    )
-    parser.add_argument(
-        "--device", type=str, default="cuda", help="Device to run on (cuda/cpu/mps)"
-    )
-    parser.add_argument(
-        "--num_iterations", type=int, default=50, help="Number of inference iterations"
-    )
-    parser.add_argument(
-        "--execution_horizon", type=int, default=10, help="RTC execution horizon"
-    )
-    parser.add_argument(
-        "--enable_detailed_profiling",
-        action="store_true",
-        help="Enable detailed method-level profiling",
-    )
-    parser.add_argument(
-        "--use_torch_compile", action="store_true", help="Use torch.compile for faster inference"
-    )
-
-    args = parser.parse_args()
-
-    # Load policy
-    logger.info(f"Loading policy from {args.policy_path}")
-    config = PreTrainedConfig.from_pretrained(args.policy_path)
-    policy_class = get_policy_class(config.type)
-
-    # Set compile flag if needed
-    if hasattr(config, "compile_model"):
-        config.compile_model = args.use_torch_compile
-
-    policy = policy_class.from_pretrained(args.policy_path, config=config)
-
-    # Initialize RTC config
-    policy.config.rtc_config = RTCConfig(
-        execution_horizon=args.execution_horizon,
-        max_guidance_weight=1.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-    )
-
-    policy = policy.to(args.device)
-    policy.eval()
-
-    logger.info(f"Policy loaded: {config.type}")
-    logger.info(f"Device: {args.device}")
-    logger.info(f"Execution horizon: {args.execution_horizon}")
-
-    # Create mock observation
-    logger.info("Creating mock observation...")
-    observation = create_mock_observation(policy, args.device)
-
-    # Enable detailed profiling if requested
-    if args.enable_detailed_profiling:
-        enable_profiling()
-        logger.info("Detailed profiling enabled")
-
-    # Profile without RTC
-    results_without_rtc = profile_inference(
-        policy=policy,
-        observation=observation,
-        num_iterations=args.num_iterations,
-        use_rtc=False,
-        execution_horizon=args.execution_horizon,
-    )
-
-    if args.enable_detailed_profiling:
-        logger.info("\nDetailed profiling stats (WITHOUT RTC):")
-        print_profiling_summary()
-        clear_profiling_stats()
-
-    # Profile with RTC
-    results_with_rtc = profile_inference(
-        policy=policy,
-        observation=observation,
-        num_iterations=args.num_iterations,
-        use_rtc=True,
-        execution_horizon=args.execution_horizon,
-    )
-
-    if args.enable_detailed_profiling:
-        logger.info("\nDetailed profiling stats (WITH RTC):")
-        print_profiling_summary()
-
-    # Compare results
-    compare_results(results_without_rtc, results_with_rtc)
-
-
-if __name__ == "__main__":
-    main()
-

examples/so100_to_so100_EE/evaluate.py

@@ -195,6 +195,4 @@ for episode_idx in range(NUM_EPISODES):
 log_say("Stop recording")
 robot.disconnect()
 listener.stop()
-
-dataset.finalize()
 dataset.push_to_hub()

examples/so100_to_so100_EE/record.py

@@ -199,6 +199,4 @@ log_say("Stop recording")
 leader.disconnect()
 follower.disconnect()
 listener.stop()
-
-dataset.finalize()
 dataset.push_to_hub()

examples/tutorial/act/act_training_example.py

@@ -1,98 +0,0 @@
-"""This script demonstrates how to train ACT Policy on a real-world dataset."""
-
-from pathlib import Path
-
-import torch
-
-from lerobot.configs.types import FeatureType
-from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
-from lerobot.datasets.utils import dataset_to_policy_features
-from lerobot.policies.act.configuration_act import ACTConfig
-from lerobot.policies.act.modeling_act import ACTPolicy
-from lerobot.policies.factory import make_pre_post_processors
-
-
-def make_delta_timestamps(delta_indices: list[int] | None, fps: int) -> list[float]:
-    if delta_indices is None:
-        return [0]
-
-    return [i / fps for i in delta_indices]
-
-
-output_directory = Path("outputs/robot_learning_tutorial/act")
-output_directory.mkdir(parents=True, exist_ok=True)
-
-# Select your device
-device = torch.device("mps")  # or "cuda" or "cpu"
-
-dataset_id = "lerobot/svla_so101_pickplace"
-
-# This specifies the inputs the model will be expecting and the outputs it will produce
-dataset_metadata = LeRobotDatasetMetadata(dataset_id)
-features = dataset_to_policy_features(dataset_metadata.features)
-
-output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
-input_features = {key: ft for key, ft in features.items() if key not in output_features}
-
-cfg = ACTConfig(input_features=input_features, output_features=output_features)
-policy = ACTPolicy(cfg)
-preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
-
-policy.train()
-policy.to(device)
-
-# To perform action chunking, ACT expects a given number of actions as targets
-delta_timestamps = {
-    "action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps),
-}
-
-# add image features if they are present
-delta_timestamps |= {
-    k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps) for k in cfg.image_features
-}
-
-# Instantiate the dataset
-dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps)
-
-# Create the optimizer and dataloader for offline training
-optimizer = cfg.get_optimizer_preset().build(policy.parameters())
-batch_size = 32
-dataloader = torch.utils.data.DataLoader(
-    dataset,
-    batch_size=batch_size,
-    shuffle=True,
-    pin_memory=device.type != "cpu",
-    drop_last=True,
-)
-
-# Number of training steps and logging frequency
-training_steps = 1
-log_freq = 1
-
-# Run training loop
-step = 0
-done = False
-while not done:
-    for batch in dataloader:
-        batch = preprocessor(batch)
-        loss, _ = policy.forward(batch)
-        loss.backward()
-        optimizer.step()
-        optimizer.zero_grad()
-
-        if step % log_freq == 0:
-            print(f"step: {step} loss: {loss.item():.3f}")
-        step += 1
-        if step >= training_steps:
-            done = True
-            break
-
-# Save the policy checkpoint, alongside the pre/post processors
-policy.save_pretrained(output_directory)
-preprocessor.save_pretrained(output_directory)
-postprocessor.save_pretrained(output_directory)
-
-# Save all assets to the Hub
-policy.push_to_hub("fracapuano/robot_learning_tutorial_act")
-preprocessor.push_to_hub("fracapuano/robot_learning_tutorial_act")
-postprocessor.push_to_hub("fracapuano/robot_learning_tutorial_act")

examples/tutorial/act/act_using_example.py

@@ -1,57 +0,0 @@
-import torch
-
-from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
-from lerobot.policies.act.modeling_act import ACTPolicy
-from lerobot.policies.factory import make_pre_post_processors
-from lerobot.policies.utils import build_inference_frame, make_robot_action
-from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-
-device = torch.device("mps")  # or "cuda" or "cpu"
-model_id = "fracapuano/robot_learning_tutorial_act"
-model = ACTPolicy.from_pretrained(model_id)
-
-dataset_id = "lerobot/svla_so101_pickplace"
-# This only downloads the metadata for the dataset, ~10s of MB even for large-scale datasets
-dataset_metadata = LeRobotDatasetMetadata(dataset_id)
-preprocess, postprocess = make_pre_post_processors(model.config, dataset_stats=dataset_metadata.stats)
-
-# # find ports using lerobot-find-port
-follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
-
-# # the robot ids are used the load the right calibration files
-follower_id = ...  # something like "follower_so100"
-
-MAX_EPISODES = 5
-MAX_STEPS_PER_EPISODE = 20
-
-# Robot and environment configuration
-# Camera keys must match the name and resolutions of the ones used for training!
-# You can check the camera keys expected by a model in the info.json card on the model card on the Hub
-camera_config = {
-    "side": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
-    "up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
-}
-
-robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
-robot = SO100Follower(robot_cfg)
-robot.connect()
-
-for _ in range(MAX_EPISODES):
-    for _ in range(MAX_STEPS_PER_EPISODE):
-        obs = robot.get_observation()
-        obs_frame = build_inference_frame(
-            observation=obs, ds_features=dataset_metadata.features, device=device
-        )
-
-        obs = preprocess(obs_frame)
-
-        action = model.select_action(obs)
-        action = postprocess(action)
-
-        action = make_robot_action(action, dataset_metadata.features)
-
-        robot.send_action(action)
-
-    print("Episode finished! Starting new episode...")

examples/tutorial/async-inf/policy_server.py

@@ -1,11 +0,0 @@
-from lerobot.async_inference.configs import PolicyServerConfig
-from lerobot.async_inference.policy_server import serve
-
-host = ...  # something like "127.0.0.1" if you're exposing to localhost
-port = ...  # something like 8080
-
-config = PolicyServerConfig(
-    host=host,
-    port=port,
-)
-serve(config)

examples/tutorial/async-inf/robot_client.py

@@ -1,55 +0,0 @@
-import threading
-
-from lerobot.async_inference.configs import RobotClientConfig
-from lerobot.async_inference.helpers import visualize_action_queue_size
-from lerobot.async_inference.robot_client import RobotClient
-from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.robots.so100_follower import SO100FollowerConfig
-
-# these cameras must match the ones expected by the policy - find your cameras with lerobot-find-cameras
-# check the config.json on the Hub for the policy you are using to see the expected camera specs
-camera_cfg = {
-    "up": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
-    "side": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
-}
-
-# # find ports using lerobot-find-port
-follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
-
-# # the robot ids are used the load the right calibration files
-follower_id = ...  # something like "follower_so100"
-
-robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_cfg)
-
-server_address = ...  # something like "127.0.0.1:8080" if using localhost
-
-# 3. Create client configuration
-client_cfg = RobotClientConfig(
-    robot=robot_cfg,
-    server_address=server_address,
-    policy_device="mps",
-    policy_type="act",
-    pretrained_name_or_path="fracapuano/robot_learning_tutorial_act",
-    chunk_size_threshold=0.5,  # g
-    actions_per_chunk=50,  # make sure this is less than the max actions of the policy
-)
-
-# 4. Create and start client
-client = RobotClient(client_cfg)
-
-# 5. Provide a textual description of the task
-task = ...
-
-if client.start():
-    # Start action receiver thread
-    action_receiver_thread = threading.Thread(target=client.receive_actions, daemon=True)
-    action_receiver_thread.start()
-
-    try:
-        # Run the control loop
-        client.control_loop(task)
-    except KeyboardInterrupt:
-        client.stop()
-        action_receiver_thread.join()
-        # (Optionally) plot the action queue size
-        visualize_action_queue_size(client.action_queue_size)

examples/tutorial/diffusion/diffusion_training_example.py

@@ -1,99 +0,0 @@
-"""This script demonstrates how to train Diffusion Policy on a real-world dataset."""
-
-from pathlib import Path
-
-import torch
-
-from lerobot.configs.types import FeatureType
-from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
-from lerobot.datasets.utils import dataset_to_policy_features
-from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
-from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
-from lerobot.policies.factory import make_pre_post_processors
-
-
-def make_delta_timestamps(delta_indices: list[int] | None, fps: int) -> list[float]:
-    if delta_indices is None:
-        return [0]
-
-    return [i / fps for i in delta_indices]
-
-
-output_directory = Path("outputs/robot_learning_tutorial/diffusion")
-output_directory.mkdir(parents=True, exist_ok=True)
-
-# Select your device
-device = torch.device("mps")  # or "cuda" or "cpu"
-
-dataset_id = "lerobot/svla_so101_pickplace"
-
-# This specifies the inputs the model will be expecting and the outputs it will produce
-dataset_metadata = LeRobotDatasetMetadata(dataset_id)
-features = dataset_to_policy_features(dataset_metadata.features)
-
-output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
-input_features = {key: ft for key, ft in features.items() if key not in output_features}
-
-cfg = DiffusionConfig(input_features=input_features, output_features=output_features)
-policy = DiffusionPolicy(cfg)
-preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
-
-policy.train()
-policy.to(device)
-
-# To perform action chunking, ACT expects a given number of actions as targets
-delta_timestamps = {
-    "observation.state": make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps),
-    "action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps),
-}
-
-# add image features if they are present
-delta_timestamps |= {
-    k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps) for k in cfg.image_features
-}
-
-# Instantiate the dataset
-dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps)
-
-# Create the optimizer and dataloader for offline training
-optimizer = cfg.get_optimizer_preset().build(policy.parameters())
-batch_size = 32
-dataloader = torch.utils.data.DataLoader(
-    dataset,
-    batch_size=batch_size,
-    shuffle=True,
-    pin_memory=device.type != "cpu",
-    drop_last=True,
-)
-
-# Number of training steps and logging frequency
-training_steps = 1
-log_freq = 1
-
-# Run training loop
-step = 0
-done = False
-while not done:
-    for batch in dataloader:
-        batch = preprocessor(batch)
-        loss, _ = policy.forward(batch)
-        loss.backward()
-        optimizer.step()
-        optimizer.zero_grad()
-
-        if step % log_freq == 0:
-            print(f"step: {step} loss: {loss.item():.3f}")
-        step += 1
-        if step >= training_steps:
-            done = True
-            break
-
-# Save the policy checkpoint, alongside the pre/post processors
-policy.save_pretrained(output_directory)
-preprocessor.save_pretrained(output_directory)
-postprocessor.save_pretrained(output_directory)
-
-# Save all assets to the Hub
-policy.push_to_hub("fracapuano/robot_learning_tutorial_diffusion")
-preprocessor.push_to_hub("fracapuano/robot_learning_tutorial_diffusion")
-postprocessor.push_to_hub("fracapuano/robot_learning_tutorial_diffusion")

examples/tutorial/diffusion/diffusion_using_example.py

@@ -1,60 +0,0 @@
-import torch
-
-from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
-from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
-from lerobot.policies.factory import make_pre_post_processors
-from lerobot.policies.utils import build_inference_frame, make_robot_action
-from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-
-device = torch.device("mps")  # or "cuda" or "cpu"
-model_id = "fracapuano/robot_learning_tutorial_diffusion"
-
-model = DiffusionPolicy.from_pretrained(model_id)
-
-dataset_id = "lerobot/svla_so101_pickplace"
-# This only downloads the metadata for the dataset, ~10s of MB even for large-scale datasets
-dataset_metadata = LeRobotDatasetMetadata(dataset_id)
-preprocess, postprocess = make_pre_post_processors(
-    model.config, model_id, dataset_stats=dataset_metadata.stats
-)
-
-MAX_EPISODES = 5
-MAX_STEPS_PER_EPISODE = 20
-
-
-# # find ports using lerobot-find-port
-follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
-
-# # the robot ids are used the load the right calibration files
-follower_id = ...  # something like "follower_so100"
-
-# Robot and environment configuration
-# Camera keys must match the name and resolutions of the ones used for training!
-# You can check the camera keys expected by a model in the info.json card on the model card on the Hub
-camera_config = {
-    "side": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
-    "up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
-}
-
-robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
-robot = SO100Follower(robot_cfg)
-robot.connect()
-
-
-for _ in range(MAX_EPISODES):
-    for _ in range(MAX_STEPS_PER_EPISODE):
-        obs = robot.get_observation()
-        obs_frame = build_inference_frame(
-            observation=obs, ds_features=dataset_metadata.features, device=device
-        )
-
-        obs = preprocess(obs_frame)
-
-        action = model.select_action(obs)
-        action = postprocess(action)
-        action = make_robot_action(action, dataset_metadata.features)
-        robot.send_action(action)
-
-    print("Episode finished! Starting new episode...")

examples/tutorial/pi0/using_pi0_example.py

@@ -1,67 +0,0 @@
-import torch
-
-from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.datasets.utils import hw_to_dataset_features
-from lerobot.policies.factory import make_pre_post_processors
-from lerobot.policies.pi0.modeling_pi0 import PI0Policy
-from lerobot.policies.utils import build_inference_frame, make_robot_action
-from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-
-MAX_EPISODES = 5
-MAX_STEPS_PER_EPISODE = 20
-
-device = torch.device("mps")  # or "cuda" or "cpu"
-model_id = "lerobot/pi0_base"
-
-model = PI0Policy.from_pretrained(model_id)
-
-preprocess, postprocess = make_pre_post_processors(
-    model.config,
-    model_id,
-    # This overrides allows to run on MPS, otherwise defaults to CUDA (if available)
-    preprocessor_overrides={"device_processor": {"device": str(device)}},
-)
-
-# find ports using lerobot-find-port
-follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
-
-# the robot ids are used the load the right calibration files
-follower_id = ...  # something like "follower_so100"
-
-# Robot and environment configuration
-# Camera keys must match the name and resolutions of the ones used for training!
-# You can check the camera keys expected by a model in the info.json card on the model card on the Hub
-camera_config = {
-    "base_0_rgb": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
-    "left_wrist_0_rgb": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
-    "right_wrist_0_rgb": OpenCVCameraConfig(index_or_path=2, width=640, height=480, fps=30),
-}
-
-robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
-robot = SO100Follower(robot_cfg)
-robot.connect()
-
-task = ""  # something like "pick the red block"
-robot_type = ""  # something like "so100_follower" for multi-embodiment datasets
-
-# This is used to match the raw observation keys to the keys expected by the policy
-action_features = hw_to_dataset_features(robot.action_features, "action")
-obs_features = hw_to_dataset_features(robot.observation_features, "observation")
-dataset_features = {**action_features, **obs_features}
-
-for _ in range(MAX_EPISODES):
-    for _ in range(MAX_STEPS_PER_EPISODE):
-        obs = robot.get_observation()
-        obs_frame = build_inference_frame(
-            observation=obs, ds_features=dataset_features, device=device, task=task, robot_type=robot_type
-        )
-
-        obs = preprocess(obs_frame)
-
-        action = model.select_action(obs)
-        action = postprocess(action)
-        action = make_robot_action(action, dataset_features)
-        robot.send_action(action)
-
-    print("Episode finished! Starting new episode...")

examples/tutorial/rl/hilserl_example.py

@@ -1,345 +0,0 @@
-import multiprocessing as mp
-import signal
-from pathlib import Path
-from queue import Empty, Full
-
-import torch
-import torch.optim as optim
-
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.utils import hw_to_dataset_features
-from lerobot.envs.configs import HILSerlProcessorConfig, HILSerlRobotEnvConfig
-from lerobot.policies.sac.configuration_sac import SACConfig
-from lerobot.policies.sac.modeling_sac import SACPolicy
-from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
-from lerobot.rl.buffer import ReplayBuffer
-from lerobot.rl.gym_manipulator import make_robot_env
-from lerobot.robots.so100_follower import SO100FollowerConfig
-from lerobot.teleoperators.so100_leader import SO100LeaderConfig
-from lerobot.teleoperators.utils import TeleopEvents
-
-LOG_EVERY = 10
-SEND_EVERY = 10
-
-
-def run_learner(
-    transitions_queue: mp.Queue,
-    parameters_queue: mp.Queue,
-    shutdown_event: mp.Event,
-    policy_learner: SACPolicy,
-    online_buffer: ReplayBuffer,
-    offline_buffer: ReplayBuffer,
-    lr: float = 3e-4,
-    batch_size: int = 32,
-    device: torch.device = "mps",
-):
-    """The learner process - trains SAC policy on transitions streamed from the actor, updating parameters
-    for the actor to adopt."""
-    policy_learner.train()
-    policy_learner.to(device)
-
-    # Create Adam optimizer from scratch - simple and clean
-    optimizer = optim.Adam(policy_learner.parameters(), lr=lr)
-
-    print(f"[LEARNER] Online buffer capacity: {online_buffer.capacity}")
-    print(f"[LEARNER] Offline buffer capacity: {offline_buffer.capacity}")
-
-    training_step = 0
-
-    while not shutdown_event.is_set():
-        # retrieve incoming transitions from the actor process
-        try:
-            transitions = transitions_queue.get(timeout=0.1)
-            for transition in transitions:
-                # HIL-SERL: Add ALL transitions to online buffer
-                online_buffer.add(**transition)
-
-                # HIL-SERL: Add ONLY human intervention transitions to offline buffer
-                is_intervention = transition.get("complementary_info", {}).get("is_intervention", False)
-                if is_intervention:
-                    offline_buffer.add(**transition)
-                    print(
-                        f"[LEARNER] Human intervention detected! Added to offline buffer (now {len(offline_buffer)} transitions)"
-                    )
-
-        except Empty:
-            pass  # No transitions available, continue
-
-        # Train if we have enough data
-        if len(online_buffer) >= policy_learner.config.online_step_before_learning:
-            # Sample from online buffer (autonomous + human data)
-            online_batch = online_buffer.sample(batch_size // 2)
-
-            # Sample from offline buffer (human demonstrations only, either precollected or at runtime)
-            offline_batch = offline_buffer.sample(batch_size // 2)
-
-            # Combine batches - this is the key HIL-SERL mechanism!
-            batch = {}
-            for key in online_batch:
-                if key in offline_batch:
-                    batch[key] = torch.cat([online_batch[key], offline_batch[key]], dim=0)
-                else:
-                    batch[key] = online_batch[key]
-
-            loss, _ = policy_learner.forward(batch)
-
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-            training_step += 1
-
-            if training_step % LOG_EVERY == 0:
-                print(
-                    f"[LEARNER] Training step {training_step}, Loss: {loss.item():.4f}, "
-                    f"Buffers: Online={len(online_buffer)}, Offline={len(offline_buffer)}"
-                )
-
-            # Send updated parameters to actor every 10 training steps
-            if training_step % SEND_EVERY == 0:
-                try:
-                    state_dict = {k: v.cpu() for k, v in policy_learner.state_dict().items()}
-                    parameters_queue.put_nowait(state_dict)
-                    print("[LEARNER] Sent updated parameters to actor")
-                except Full:
-                    # Missing write due to queue not being consumed (should happen rarely)
-                    pass
-
-    print("[LEARNER] Learner process finished")
-
-
-def run_actor(
-    transitions_queue: mp.Queue,
-    parameters_queue: mp.Queue,
-    shutdown_event: mp.Event,
-    policy_actor: SACPolicy,
-    reward_classifier: Classifier,
-    env_cfg: HILSerlRobotEnvConfig,
-    device: torch.device = "mps",
-    output_directory: Path | None = None,
-):
-    """The actor process - interacts with environment and collects data.
-    The policy is frozen and only the parameters are updated, popping the most recent ones from a queue."""
-    policy_actor.eval()
-    policy_actor.to(device)
-
-    reward_classifier.eval()
-    reward_classifier.to(device)
-
-    # Create robot environment inside the actor process
-    env, teleop_device = make_robot_env(env_cfg)
-
-    try:
-        for episode in range(MAX_EPISODES):
-            if shutdown_event.is_set():
-                break
-
-            obs, _info = env.reset()
-            episode_reward = 0.0
-            step = 0
-            episode_transitions = []
-
-            print(f"[ACTOR] Starting episode {episode + 1}")
-
-            while step < MAX_STEPS_PER_EPISODE and not shutdown_event.is_set():
-                try:
-                    new_params = parameters_queue.get_nowait()
-                    policy_actor.load_state_dict(new_params)
-                    print("[ACTOR] Updated policy parameters from learner")
-                except Empty:  # No new updated parameters available from learner, waiting
-                    pass
-
-                # Get action from policy
-                policy_obs = make_policy_obs(obs, device=device)
-                action_tensor = policy_actor.select_action(policy_obs)  # predicts a single action
-                action = action_tensor.squeeze(0).cpu().numpy()
-
-                # Step environment
-                next_obs, _env_reward, terminated, truncated, _info = env.step(action)
-                done = terminated or truncated
-
-                # Predict reward
-                policy_next_obs = make_policy_obs(next_obs, device=device)
-                reward = reward_classifier.predict_reward(policy_next_obs)
-
-                if reward >= 1.0 and not done:  # success detected! halt episode
-                    terminated = True
-                    done = True
-
-                # In HIL-SERL, human interventions come from the teleop device
-                is_intervention = False
-                if hasattr(teleop_device, "get_teleop_events"):
-                    # Real intervention detection from teleop device
-                    teleop_events = teleop_device.get_teleop_events()
-                    is_intervention = teleop_events.get(TeleopEvents.IS_INTERVENTION, False)
-
-                # Store transition with intervention metadata
-                transition = {
-                    "state": policy_obs,
-                    "action": action,
-                    "reward": float(reward) if hasattr(reward, "item") else reward,
-                    "next_state": policy_next_obs,
-                    "done": done,
-                    "truncated": truncated,
-                    "complementary_info": {
-                        "is_intervention": is_intervention,
-                    },
-                }
-
-                episode_transitions.append(transition)
-
-                episode_reward += reward
-                step += 1
-
-                obs = next_obs
-
-                if done:
-                    break
-
-            # Send episode transitions to learner
-            transitions_queue.put_nowait(episode_transitions)
-
-    except KeyboardInterrupt:
-        print("[ACTOR] Interrupted by user")
-    finally:
-        # Clean up
-        if hasattr(env, "robot") and env.robot.is_connected:
-            env.robot.disconnect()
-        if teleop_device and hasattr(teleop_device, "disconnect"):
-            teleop_device.disconnect()
-        if output_directory is not None:
-            policy_actor.save_pretrained(output_directory)
-            print(f"[ACTOR] Latest actor policy saved at: {output_directory}")
-
-        print("[ACTOR] Actor process finished")
-
-
-def make_policy_obs(obs, device: torch.device = "cpu"):
-    return {
-        "observation.state": torch.from_numpy(obs["agent_pos"]).float().unsqueeze(0).to(device),
-        **{
-            f"observation.image.{k}": torch.from_numpy(obs["pixels"][k]).float().unsqueeze(0).to(device)
-            for k in obs["pixels"]
-        },
-    }
-
-
-"""Main function - coordinates actor and learner processes."""
-
-device = "mps"  # or "cuda" or "cpu"
-output_directory = Path("outputs/robot_learning_tutorial/hil_serl")
-output_directory.mkdir(parents=True, exist_ok=True)
-
-# find ports using lerobot-find-port
-follower_port = ...
-leader_port = ...
-
-# the robot ids are used the load the right calibration files
-follower_id = ...
-leader_id = ...
-
-# A pretrained model (to be used in-distribution!)
-reward_classifier_id = "fracapuano/reward_classifier_hil_serl_example"
-reward_classifier = Classifier.from_pretrained(reward_classifier_id)
-
-reward_classifier.to(device)
-reward_classifier.eval()
-
-MAX_EPISODES = 5
-MAX_STEPS_PER_EPISODE = 20
-
-# Robot and environment configuration
-robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id)
-teleop_cfg = SO100LeaderConfig(port=leader_port, id=leader_id)
-processor_cfg = HILSerlProcessorConfig(control_mode="leader")
-
-env_cfg = HILSerlRobotEnvConfig(robot=robot_cfg, teleop=teleop_cfg, processor=processor_cfg)
-
-# Create robot environment
-env, teleop_device = make_robot_env(env_cfg)
-
-obs_features = hw_to_dataset_features(env.robot.observation_features, "observation")
-action_features = hw_to_dataset_features(env.robot.action_features, "action")
-
-# Create SAC policy for action selection
-policy_cfg = SACConfig(
-    device=device,
-    input_features=obs_features,
-    output_features=action_features,
-)
-
-policy_actor = SACPolicy(policy_cfg)
-policy_learner = SACPolicy(policy_cfg)
-
-demonstrations_repo_id = "lerobot/example_hil_serl_dataset"
-offline_dataset = LeRobotDataset(repo_id=demonstrations_repo_id)
-
-# Online buffer: initialized from scratch
-online_replay_buffer = ReplayBuffer(device=device, state_keys=list(obs_features.keys()))
-# Offline buffer: Created from dataset (pre-populated it with demonstrations)
-offline_replay_buffer = ReplayBuffer.from_lerobot_dataset(
-    lerobot_dataset=offline_dataset, device=device, state_keys=list(obs_features.keys())
-)
-
-# Create communication channels between learner and actor processes
-transitions_queue = mp.Queue(maxsize=10)
-parameters_queue = mp.Queue(maxsize=2)
-shutdown_event = mp.Event()
-
-
-# Signal handler for graceful shutdown
-def signal_handler(sig):
-    print(f"\nSignal {sig} received, shutting down...")
-    shutdown_event.set()
-
-
-signal.signal(signal.SIGINT, signal_handler)
-signal.signal(signal.SIGTERM, signal_handler)
-
-# Create processes
-learner_process = mp.Process(
-    target=run_learner,
-    args=(
-        transitions_queue,
-        parameters_queue,
-        shutdown_event,
-        policy_learner,
-        online_replay_buffer,
-        offline_replay_buffer,
-    ),
-    kwargs={"device": device},  # can run on accelerated hardware for training
-)
-
-actor_process = mp.Process(
-    target=run_actor,
-    args=(
-        transitions_queue,
-        parameters_queue,
-        shutdown_event,
-        policy_actor,
-        reward_classifier,
-        env_cfg,
-        output_directory,
-    ),
-    kwargs={"device": "cpu"},  # actor is frozen, can run on CPU or accelerate for inference
-)
-
-learner_process.start()
-actor_process.start()
-
-try:
-    # Wait for actor to finish (it controls the episode loop)
-    actor_process.join()
-    shutdown_event.set()
-    learner_process.join(timeout=10)
-
-except KeyboardInterrupt:
-    print("Main process interrupted")
-    shutdown_event.set()
-    actor_process.join(timeout=5)
-    learner_process.join(timeout=10)
-
-finally:
-    if learner_process.is_alive():
-        learner_process.terminate()
-    if actor_process.is_alive():
-        actor_process.terminate()

examples/tutorial/rl/reward_classifier_example.py

@@ -1,62 +0,0 @@
-import torch
-
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.policies.factory import make_policy, make_pre_post_processors
-from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
-
-# Device to use for training
-device = "mps"  # or "cuda", or "cpu"
-
-# Load the dataset used for training
-repo_id = "lerobot/example_hil_serl_dataset"
-dataset = LeRobotDataset(repo_id)
-
-# Configure the policy to extract features from the image frames
-camera_keys = dataset.meta.camera_keys
-
-config = RewardClassifierConfig(
-    num_cameras=len(camera_keys),
-    device=device,
-    # backbone model to extract features from the image frames
-    model_name="microsoft/resnet-18",
-)
-
-# Make policy, preprocessor, and optimizer
-policy = make_policy(config, ds_meta=dataset.meta)
-optimizer = config.get_optimizer_preset().build(policy.parameters())
-preprocessor, _ = make_pre_post_processors(policy_cfg=config, dataset_stats=dataset.meta.stats)
-
-
-classifier_id = "fracapuano/reward_classifier_hil_serl_example"
-
-# Instantiate a dataloader
-dataloader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True)
-
-# Training loop
-num_epochs = 5
-for epoch in range(num_epochs):
-    total_loss = 0
-    total_accuracy = 0
-    for batch in dataloader:
-        # Preprocess the batch and move it to the correct device.
-        batch = preprocessor(batch)
-
-        # Forward pass
-        loss, output_dict = policy.forward(batch)
-
-        # Backward pass and optimization
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-
-        total_loss += loss.item()
-        total_accuracy += output_dict["accuracy"]
-
-    avg_loss = total_loss / len(dataloader)
-    avg_accuracy = total_accuracy / len(dataloader)
-    print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}, Accuracy: {avg_accuracy:.2f}%")
-
-print("Training finished!")
-
-# You can now save the trained policy.
-policy.push_to_hub(classifier_id)

examples/tutorial/smolvla/using_smolvla_example.py

@@ -1,66 +0,0 @@
-import torch
-
-from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.datasets.utils import hw_to_dataset_features
-from lerobot.policies.factory import make_pre_post_processors
-from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
-from lerobot.policies.utils import build_inference_frame, make_robot_action
-from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-
-MAX_EPISODES = 5
-MAX_STEPS_PER_EPISODE = 20
-
-device = torch.device("mps")  # or "cuda" or "cpu"
-model_id = "lerobot/smolvla_base"
-
-model = SmolVLAPolicy.from_pretrained(model_id)
-
-preprocess, postprocess = make_pre_post_processors(
-    model.config,
-    model_id,
-    # This overrides allows to run on MPS, otherwise defaults to CUDA (if available)
-    preprocessor_overrides={"device_processor": {"device": str(device)}},
-)
-
-# find ports using lerobot-find-port
-follower_port = ...  # something like "/dev/tty.usbmodem58760431631"
-
-# the robot ids are used the load the right calibration files
-follower_id = ...  # something like "follower_so100"
-
-# Robot and environment configuration
-# Camera keys must match the name and resolutions of the ones used for training!
-# You can check the camera keys expected by a model in the info.json card on the model card on the Hub
-camera_config = {
-    "camera1": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
-    "camera2": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
-}
-
-robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
-robot = SO100Follower(robot_cfg)
-robot.connect()
-
-task = ""  # something like "pick the red block"
-robot_type = ""  # something like "so100_follower" for multi-embodiment datasets
-
-# This is used to match the raw observation keys to the keys expected by the policy
-action_features = hw_to_dataset_features(robot.action_features, "action")
-obs_features = hw_to_dataset_features(robot.observation_features, "observation")
-dataset_features = {**action_features, **obs_features}
-
-for _ in range(MAX_EPISODES):
-    for _ in range(MAX_STEPS_PER_EPISODE):
-        obs = robot.get_observation()
-        obs_frame = build_inference_frame(
-            observation=obs, ds_features=dataset_features, device=device, task=task, robot_type=robot_type
-        )
-
-        obs = preprocess(obs_frame)
-
-        action = model.select_action(obs)
-        action = postprocess(action)
-        action = make_robot_action(action, dataset_features)
-        robot.send_action(action)
-
-    print("Episode finished! Starting new episode...")

pyproject.toml

@@ -25,7 +25,7 @@ discord = "https://discord.gg/s3KuuzsPFb"
 
 [project]
 name = "lerobot"
-version = "0.4.2"
+version = "0.3.4"
 description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch"
 readme = "README.md"
 license = { text = "Apache-2.0" }
@@ -62,10 +62,8 @@ dependencies = [
     "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",
-    "accelerate>=1.10.0,<2.0.0",
 
     # Core dependencies
-    "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",
@@ -74,15 +72,15 @@ dependencies = [
     "packaging>=24.2,<26.0",
     "pynput>=1.7.7,<1.9.0",
     "pyserial>=3.5,<4.0",
-    "wandb>=0.20.0,<0.22.0", # TODO: Bumb dependency (compatible with protobuf)
+    "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>=1.1.1,<2.0.0",
-    "rerun-sdk>=0.24.0,<0.27.0",
+    "gymnasium>=0.29.1,<1.0.0", # TODO: Bumb dependency
+    "rerun-sdk>=0.21.0,<0.23.0", # TODO: Bumb dependency
 
     # Support dependencies
     "deepdiff>=7.0.1,<9.0.0",
@@ -97,8 +95,7 @@ dependencies = [
 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"] # TODO: Bumb dependency (compatible with wandb)
-matplotlib-dep = ["matplotlib>=3.10.3,<4.0.0"]
+grpcio-dep = ["grpcio==1.73.1", "protobuf==6.31.0"]
 
 # Motors
 feetech = ["feetech-servo-sdk>=1.0.0,<2.0.0"]
@@ -114,26 +111,20 @@ intelrealsense = [
     "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,<2.12.0", "teleop>=0.1.0,<0.2.0", "fastapi<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'",
+#     "pyrealsense2>=2.55.1.6486 ; sys_platform != 'darwin'"
+# ] # TODO: Currently not supported
 
 # Policies
 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"]
-groot = [
-    "lerobot[transformers-dep]",
-    "peft>=0.13.0,<1.0.0",
-    "dm-tree>=0.1.8,<1.0.0",
-    "timm>=1.0.0,<1.1.0",
-    "safetensors>=0.4.3,<1.0.0",
-    "Pillow>=10.0.0,<13.0.0",
-    "decord>=0.6.0,<1.0.0; (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
-    "ninja>=1.11.1,<2.0.0",
-    "flash-attn>=2.5.9,<3.0.0 ; sys_platform != 'darwin'"
-]
-hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.13,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
+hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.11,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
 
 # Features
-async = ["lerobot[grpcio-dep]", "lerobot[matplotlib-dep]"]
+async = ["lerobot[grpcio-dep]", "matplotlib>=3.10.3,<4.0.0"]
 
 # Development
 dev = ["pre-commit>=3.7.0,<5.0.0", "debugpy>=1.8.1,<1.9.0", "lerobot[grpcio-dep]", "grpcio-tools==1.73.1"]
@@ -141,10 +132,11 @@ test = ["pytest>=8.1.0,<9.0.0", "pytest-timeout>=2.4.0,<3.0.0", "pytest-cov>=5.0
 video_benchmark = ["scikit-image>=0.23.2,<0.26.0", "pandas>=2.2.2,<2.4.0"]
 
 # Simulation
-aloha = ["gym-aloha>=0.1.2,<0.2.0"]
+aloha = ["gym-aloha>=0.1.1,<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]", "hf-libero>=0.1.3,<0.2.0"]
-metaworld = ["metaworld==3.0.0"]
+xarm = ["gym-xarm>=0.1.1,<0.2.0"]
+libero = ["lerobot[transformers-dep]", "libero @ git+https://github.com/huggingface/lerobot-libero.git@main#egg=libero"]
+
 
 # All
 all = [
@@ -157,7 +149,6 @@ all = [
     "lerobot[intelrealsense]",
     "lerobot[pi]",
     "lerobot[smolvla]",
-    # "lerobot[groot]", TODO(Steven): Gr00t requires specific installation instructions for flash-attn
     "lerobot[hilserl]",
     "lerobot[async]",
     "lerobot[dev]",
@@ -165,9 +156,9 @@ all = [
     "lerobot[video_benchmark]",
     "lerobot[aloha]",
     "lerobot[pusht]",
+    "lerobot[xarm]",
     "lerobot[phone]",
     "lerobot[libero]",
-    "lerobot[metaworld]",
 ]
 
 [project.scripts]
@@ -242,6 +233,9 @@ exclude_dirs = [
     "tests",
     "benchmarks",
     "src/lerobot/datasets/push_dataset_to_hub",
+    "src/lerobot/datasets/v2/convert_dataset_v1_to_v2",
+    "src/lerobot/policies/pi0/conversion_scripts",
+    "src/lerobot/scripts/push_dataset_to_hub.py",
 ]
 skips = ["B101", "B311", "B404", "B603", "B615"]
 
@@ -256,8 +250,6 @@ default.extend-ignore-identifiers-re = [
     "pn",
     "ser",
     "ein",
-    "thw",
-    "inpt",
 ]
 
 # TODO: Uncomment when ready to use
@@ -296,6 +288,7 @@ ignore_errors = true
 
 [[tool.mypy.overrides]]
 module = "lerobot.envs.*"
+# Enable type checking only for the envs module
 ignore_errors = false
 
 
@@ -303,22 +296,17 @@ ignore_errors = false
 # module = "lerobot.utils.*"
 # ignore_errors = false
 
-[[tool.mypy.overrides]]
-module = "lerobot.configs.*"
-ignore_errors = false
-
-# extra strictness for configs
-disallow_untyped_defs = true
-disallow_incomplete_defs = true
-check_untyped_defs = true
+# [[tool.mypy.overrides]]
+# module = "lerobot.configs.*"
+# ignore_errors = false
 
 # [[tool.mypy.overrides]]
 # module = "lerobot.optim.*"
 # ignore_errors = false
 
-[[tool.mypy.overrides]]
-module = "lerobot.model.*"
-ignore_errors = false
+# [[tool.mypy.overrides]]
+# module = "lerobot.model.*"
+# ignore_errors = false
 
 # [[tool.mypy.overrides]]
 # module = "lerobot.processor.*"
@@ -328,9 +316,9 @@ ignore_errors = false
 # module = "lerobot.datasets.*"
 # ignore_errors = false
 
-[[tool.mypy.overrides]]
-module = "lerobot.cameras.*"
-ignore_errors = false
+# [[tool.mypy.overrides]]
+# module = "lerobot.cameras.*"
+# ignore_errors = false
 
 # [[tool.mypy.overrides]]
 # module = "lerobot.motors.*"

requirements-macos.txt

@@ -1,4 +1,3 @@
-#
 # This file is autogenerated by pip-compile with Python 3.10
 # by the following command:
 #
@@ -13,62 +12,47 @@ absl-py==2.3.1
     #   dm-tree
     #   labmaze
     #   mujoco
-    #   tensorboard
-accelerate==1.11.0
-    # via
-    #   lerobot
-    #   peft
+accelerate==1.9.0
+    # via lerobot
 aiohappyeyeballs==2.6.1
     # via aiohttp
-aiohttp==3.13.1
+aiohttp==3.12.15
     # via fsspec
 aiosignal==1.4.0
     # via aiohttp
 annotated-types==0.7.0
     # via pydantic
-antlr4-python3-runtime==4.9.3
-    # via
-    #   hydra-core
-    #   omegaconf
-anyio==4.11.0
-    # via
-    #   starlette
-    #   watchfiles
 asttokens==3.0.0
     # via stack-data
 async-timeout==5.0.1
     # via aiohttp
-attrs==25.4.0
+attrs==25.3.0
     # via
     #   aiohttp
     #   dm-tree
     #   jsonlines
-    #   jsonschema
-    #   referencing
     #   rerun-sdk
-av==15.1.0
+av==15.0.0
     # via lerobot
-bddl==1.0.1
-    # via libero
-certifi==2025.10.5
+blinker==1.9.0
+    # via flask
+certifi==2025.7.14
     # via
     #   requests
     #   sentry-sdk
-cffi==2.0.0
+cffi==1.17.1
     # via pymunk
 cfgv==3.4.0
     # via pre-commit
-charset-normalizer==3.4.4
+charset-normalizer==3.4.2
     # via requests
-click==8.3.0
+click==8.2.1
     # via
-    #   uvicorn
+    #   flask
     #   wandb
 cloudpickle==3.1.1
-    # via
-    #   gymnasium
-    #   libero
-cmake==4.1.0
+    # via gymnasium
+cmake==4.0.3
     # via lerobot
 cmeel==0.57.3
     # via
@@ -110,27 +94,27 @@ coal-library==3.0.1
     # via pin
 contourpy==1.3.2
     # via matplotlib
-coverage[toml]==7.11.0
+coverage[toml]==7.10.1
     # via pytest-cov
 cycler==0.12.1
     # via matplotlib
-datasets==4.1.1
+datasets==3.6.0
     # via lerobot
-debugpy==1.8.17
+debugpy==1.8.15
     # via lerobot
 decorator==5.2.1
     # via ipython
-deepdiff==8.6.1
+deepdiff==8.5.0
     # via lerobot
-diffusers==0.35.2
+diffusers==0.34.0
     # via lerobot
-dill==0.4.0
+dill==0.3.8
     # via
     #   datasets
     #   multiprocess
 distlib==0.4.0
     # via virtualenv
-dm-control==1.0.34
+dm-control==1.0.14
     # via gym-aloha
 dm-env==1.6
     # via dm-control
@@ -138,45 +122,29 @@ dm-tree==0.1.9
     # via
     #   dm-control
     #   dm-env
-    #   lerobot
 docopt==0.6.2
     # via num2words
 draccus==0.10.0
     # via lerobot
-dynamixel-sdk==3.8.4
+dynamixel-sdk==3.7.31
     # via lerobot
-easydict==1.13
-    # via libero
-egl-probe @ git+https://github.com/huggingface/egl_probe.git
-    # via
-    #   libero
-    #   robomimic
 eigenpy==3.10.3
     # via coal-library
 einops==0.8.1
-    # via
-    #   lerobot
-    #   libero
+    # via lerobot
 eiquadprog==1.2.9
     # via placo
-etils[epath,epy]==1.13.0
-    # via mujoco
 exceptiongroup==1.3.0
     # via
-    #   anyio
     #   ipython
     #   pytest
-executing==2.2.1
+executing==2.2.0
     # via stack-data
 farama-notifications==0.0.4
     # via gymnasium
-fastapi==0.119.1
-    # via teleop
-fastjsonschema==2.21.2
-    # via nbformat
 feetech-servo-sdk==1.0.0
     # via lerobot
-filelock==3.20.0
+filelock==3.18.0
     # via
     #   datasets
     #   diffusers
@@ -184,25 +152,24 @@ filelock==3.20.0
     #   torch
     #   transformers
     #   virtualenv
-fonttools==4.60.1
+flask==3.1.1
+    # via lerobot
+fonttools==4.59.0
     # via matplotlib
-frozenlist==1.8.0
+frozenlist==1.7.0
     # via
     #   aiohttp
     #   aiosignal
-fsspec[http]==2025.9.0
+fsspec[http]==2025.3.0
     # via
     #   datasets
-    #   etils
     #   huggingface-hub
     #   torch
-future==1.0.0
-    # via libero
 gitdb==4.0.12
     # via gitpython
 gitpython==3.1.45
     # via wandb
-glfw==2.10.0
+glfw==2.9.0
     # via
     #   dm-control
     #   mujoco
@@ -210,79 +177,61 @@ grpcio==1.73.1
     # via
     #   grpcio-tools
     #   lerobot
-    #   reachy2-sdk
-    #   reachy2-sdk-api
-    #   tensorboard
 grpcio-tools==1.73.1
-    # via
-    #   lerobot
-    #   reachy2-sdk-api
-gym-aloha==0.1.3
     # via lerobot
-gym-hil==0.1.13
+gym-aloha==0.1.1
     # via lerobot
-gym-pusht==0.1.6
+gym-hil==0.1.10
     # via lerobot
-gymnasium==1.2.1
+gym-pusht==0.1.5
+    # via lerobot
+gym-xarm==0.1.1
+    # via lerobot
+gymnasium==0.29.1
     # via
     #   gym-aloha
     #   gym-hil
     #   gym-pusht
+    #   gym-xarm
+    #   gymnasium-robotics
     #   lerobot
-    #   libero
-    #   metaworld
-h11==0.16.0
-    # via uvicorn
-h5py==3.15.1
-    # via robomimic
-hebi-py==2.11.0
-    # via lerobot
+    #   pettingzoo
+gymnasium-robotics==1.2.4
+    # via gym-xarm
 hf-transfer==0.1.9
     # via huggingface-hub
-hf-xet==1.1.10
+hf-xet==1.1.5
     # via huggingface-hub
 hidapi==0.14.0.post4
     # via
     #   gym-hil
     #   lerobot
-httptools==0.7.1
-    # via uvicorn
-huggingface-hub[cli,hf-transfer]==0.35.3
+huggingface-hub[cli,hf-transfer]==0.34.3
     # via
     #   accelerate
     #   datasets
     #   diffusers
     #   lerobot
-    #   peft
-    #   timm
     #   tokenizers
     #   transformers
-hydra-core==1.3.2
-    # via libero
-identify==2.6.15
+identify==2.6.12
     # via pre-commit
-idna==3.11
+idna==3.10
     # via
-    #   anyio
     #   requests
     #   yarl
 imageio[ffmpeg]==2.37.0
     # via
     #   gym-aloha
     #   gym-hil
+    #   gymnasium-robotics
     #   lerobot
-    #   metaworld
-    #   robomimic
     #   scikit-image
 imageio-ffmpeg==0.6.0
-    # via
-    #   imageio
-    #   robomimic
+    # via imageio
 importlib-metadata==8.7.0
     # via diffusers
-importlib-resources==6.5.2
-    # via etils
-iniconfig==2.3.0
+iniconfig==2.1.0
     # via pytest
 inquirerpy==0.3.4
     # via huggingface-hub
@@ -290,71 +239,50 @@ ipython==8.37.0
     # via meshcat
 ischedule==1.2.7
     # via placo
+itsdangerous==2.2.0
+    # via flask
 jedi==0.19.2
     # via ipython
 jinja2==3.1.6
-    # via torch
+    # via
+    #   flask
+    #   gymnasium-robotics
+    #   torch
 jsonlines==4.0.0
     # via lerobot
-jsonschema==4.25.1
-    # via nbformat
-jsonschema-specifications==2025.9.1
-    # via jsonschema
-jupyter-core==5.9.1
-    # via nbformat
-jupytext==1.18.1
-    # via bddl
-kiwisolver==1.4.9
+kiwisolver==1.4.8
     # via matplotlib
 labmaze==1.0.6
     # via dm-control
 lazy-loader==0.4
     # via scikit-image
-libero @ git+https://github.com/huggingface/lerobot-libero.git@main
-    # via lerobot
-llvmlite==0.45.1
-    # via numba
-lxml==6.0.2
+lxml==6.0.0
     # via dm-control
-markdown==3.9
-    # via tensorboard
-markdown-it-py==4.0.0
-    # via
-    #   jupytext
-    #   mdit-py-plugins
-markupsafe==3.0.3
+markupsafe==3.0.2
     # via
+    #   flask
     #   jinja2
     #   werkzeug
-matplotlib==3.10.7
-    # via
-    #   lerobot
-    #   libero
-matplotlib-inline==0.2.1
+matplotlib==3.10.5
+    # via lerobot
+matplotlib-inline==0.1.7
     # via ipython
-mdit-py-plugins==0.5.0
-    # via jupytext
-mdurl==0.1.2
-    # via markdown-it-py
 mergedeep==1.3.4
     # via draccus
 meshcat==0.3.2
     # via placo
-metaworld==3.0.0
-    # via lerobot
 mock-serial==0.0.1
     # via lerobot
 mpmath==1.3.0
     # via sympy
-mujoco==3.3.7
+mujoco==2.3.7
     # via
     #   dm-control
     #   gym-aloha
     #   gym-hil
-    #   libero
-    #   metaworld
-    #   robosuite
-multidict==6.7.0
+    #   gym-xarm
+    #   gymnasium-robotics
+multidict==6.6.3
     # via
     #   aiohttp
     #   yarl
@@ -362,25 +290,17 @@ multiprocess==0.70.16
     # via datasets
 mypy-extensions==1.1.0
     # via typing-inspect
-nbformat==5.10.4
-    # via jupytext
 networkx==3.4.2
     # via
-    #   bddl
     #   scikit-image
     #   torch
-ninja==1.13.0
-    # via lerobot
 nodeenv==1.9.1
     # via pre-commit
 num2words==0.5.14
     # via lerobot
-numba==0.62.1
-    # via robosuite
 numpy==2.2.6
     # via
     #   accelerate
-    #   bddl
     #   cmeel-boost
     #   contourpy
     #   datasets
@@ -389,43 +309,25 @@ numpy==2.2.6
     #   dm-env
     #   dm-tree
     #   gymnasium
-    #   h5py
-    #   hebi-py
+    #   gymnasium-robotics
     #   imageio
     #   labmaze
-    #   libero
     #   matplotlib
     #   meshcat
-    #   metaworld
     #   mujoco
-    #   numba
     #   opencv-python
     #   opencv-python-headless
     #   pandas
-    #   peft
-    #   pyquaternion
-    #   reachy2-sdk
+    #   pettingzoo
     #   rerun-sdk
-    #   robomimic
-    #   robosuite
     #   scikit-image
     #   scipy
     #   shapely
-    #   teleop
-    #   tensorboard
-    #   tensorboardx
     #   tifffile
     #   torchvision
     #   transformers
-    #   transforms3d
-omegaconf==2.3.0
-    # via hydra-core
 opencv-python==4.12.0.88
-    # via
-    #   gym-pusht
-    #   libero
-    #   reachy2-sdk
-    #   robosuite
+    # via gym-pusht
 opencv-python-headless==4.12.0.88
     # via lerobot
 orderly-set==5.5.0
@@ -435,63 +337,53 @@ packaging==25.0
     #   accelerate
     #   datasets
     #   huggingface-hub
-    #   hydra-core
-    #   jupytext
     #   lazy-loader
     #   lerobot
     #   matplotlib
-    #   peft
     #   pytest
-    #   reachy2-sdk
     #   scikit-image
-    #   tensorboard
-    #   tensorboardx
     #   transformers
     #   wandb
-pandas==2.3.3
+pandas==2.3.1
     # via
     #   datasets
     #   lerobot
-parso==0.8.5
+parso==0.8.4
     # via jedi
-peft==0.17.1
-    # via lerobot
+pettingzoo==1.24.3
+    # via gymnasium-robotics
 pexpect==4.9.0
     # via ipython
 pfzy==0.3.4
     # via inquirerpy
-pillow==12.0.0
+pillow==11.3.0
     # via
     #   diffusers
     #   imageio
-    #   lerobot
     #   matplotlib
     #   meshcat
     #   rerun-sdk
-    #   robosuite
     #   scikit-image
-    #   tensorboard
     #   torchvision
 pin==3.4.0
     # via placo
 placo==0.9.14
     # via lerobot
-platformdirs==4.5.0
+platformdirs==4.3.8
     # via
-    #   jupyter-core
     #   virtualenv
     #   wandb
 pluggy==1.6.0
     # via
     #   pytest
     #   pytest-cov
-pre-commit==4.3.0
+pre-commit==4.2.0
     # via lerobot
-prompt-toolkit==3.0.52
+prompt-toolkit==3.0.51
     # via
     #   inquirerpy
     #   ipython
-propcache==0.4.1
+propcache==0.3.2
     # via
     #   aiohttp
     #   yarl
@@ -500,17 +392,11 @@ protobuf==6.31.0
     #   dm-control
     #   grpcio-tools
     #   lerobot
-    #   reachy2-sdk
-    #   reachy2-sdk-api
-    #   tensorboard
-    #   tensorboardx
     #   wandb
-psutil==7.1.1
+psutil==7.0.0
     # via
     #   accelerate
     #   imageio
-    #   peft
-    #   robomimic
 ptyprocess==0.7.0
     # via pexpect
 pure-eval==0.2.3
@@ -519,13 +405,11 @@ pyarrow==21.0.0
     # via
     #   datasets
     #   rerun-sdk
-pycparser==2.23
+pycparser==2.22
     # via cffi
-pydantic==2.12.3
-    # via
-    #   fastapi
-    #   wandb
-pydantic-core==2.41.4
+pydantic==2.11.7
+    # via wandb
+pydantic-core==2.33.2
     # via pydantic
 pygame==2.6.1
     # via
@@ -540,42 +424,40 @@ pymunk==6.11.1
     # via
     #   gym-pusht
     #   lerobot
-pyngrok==7.4.1
+pyngrok==7.2.12
     # via meshcat
 pynput==1.8.1
     # via
     #   gym-hil
     #   lerobot
-pyobjc-core==12.0
+pyobjc-core==11.1
     # via
     #   pyobjc-framework-applicationservices
     #   pyobjc-framework-cocoa
     #   pyobjc-framework-coretext
     #   pyobjc-framework-quartz
-pyobjc-framework-applicationservices==12.0
+pyobjc-framework-applicationservices==11.1
     # via pynput
-pyobjc-framework-cocoa==12.0
+pyobjc-framework-cocoa==11.1
     # via
     #   pyobjc-framework-applicationservices
     #   pyobjc-framework-coretext
     #   pyobjc-framework-quartz
-pyobjc-framework-coretext==12.0
+pyobjc-framework-coretext==11.1
     # via pyobjc-framework-applicationservices
-pyobjc-framework-quartz==12.0
+pyobjc-framework-quartz==11.1
     # via
     #   pynput
     #   pyobjc-framework-applicationservices
     #   pyobjc-framework-coretext
-pyopengl==3.1.10
+pyopengl==3.1.9
     # via
     #   dm-control
     #   mujoco
-pyparsing==3.2.5
+pyparsing==3.2.3
     # via
     #   dm-control
     #   matplotlib
-pyquaternion==0.9.9
-    # via reachy2-sdk
 pyrealsense2-macosx==2.54.2
     # via lerobot
 pyserial==3.5
@@ -583,14 +465,12 @@ pyserial==3.5
     #   dynamixel-sdk
     #   feetech-servo-sdk
     #   lerobot
-pytest==8.4.2
+pytest==8.4.1
     # via
-    #   bddl
     #   lerobot
     #   pytest-cov
     #   pytest-timeout
-    #   teleop
-pytest-cov==7.0.0
+pytest-cov==6.2.1
     # via lerobot
 pytest-timeout==2.4.0
     # via lerobot
@@ -598,73 +478,46 @@ python-dateutil==2.9.0.post0
     # via
     #   matplotlib
     #   pandas
-python-dotenv==1.1.1
-    # via uvicorn
 pytz==2025.2
     # via pandas
-pyyaml==6.0.3
+pyyaml==6.0.2
     # via
     #   accelerate
     #   datasets
     #   draccus
-    #   hebi-py
     #   huggingface-hub
-    #   jupytext
-    #   omegaconf
-    #   peft
     #   pre-commit
     #   pyngrok
     #   pyyaml-include
-    #   timm
     #   transformers
-    #   uvicorn
     #   wandb
 pyyaml-include==1.4.1
     # via draccus
-pyzmq==27.1.0
+pyzmq==27.0.0
     # via
     #   lerobot
     #   meshcat
-reachy2-sdk==1.0.14
-    # via lerobot
-reachy2-sdk-api==1.0.21
-    # via reachy2-sdk
-referencing==0.37.0
-    # via
-    #   jsonschema
-    #   jsonschema-specifications
-regex==2025.10.23
+regex==2025.7.34
     # via
     #   diffusers
     #   transformers
-requests==2.32.5
+requests==2.32.4
     # via
     #   datasets
     #   diffusers
     #   dm-control
     #   huggingface-hub
-    #   teleop
     #   transformers
     #   wandb
-rerun-sdk==0.26.1
+rerun-sdk==0.22.1
     # via lerobot
 rhoban-cmeel-jsoncpp==1.9.4.9
     # via placo
-robomimic==0.2.0
-    # via libero
-robosuite==1.4.0
-    # via libero
-rpds-py==0.28.0
-    # via
-    #   jsonschema
-    #   referencing
-safetensors==0.6.2
+safetensors==0.5.3
     # via
     #   accelerate
     #   diffusers
     #   lerobot
-    #   peft
-    #   timm
     #   transformers
 scikit-image==0.25.2
     # via
@@ -673,12 +526,10 @@ scikit-image==0.25.2
 scipy==1.15.3
     # via
     #   dm-control
-    #   metaworld
-    #   robosuite
     #   scikit-image
-sentry-sdk==2.42.1
+sentry-sdk==2.34.1
     # via wandb
-shapely==2.1.2
+shapely==2.1.1
     # via gym-pusht
 six==1.17.0
     # via
@@ -686,106 +537,64 @@ six==1.17.0
     #   python-dateutil
 smmap==5.0.2
     # via gitdb
-sniffio==1.3.1
-    # via anyio
 stack-data==0.6.3
     # via ipython
-starlette==0.48.0
-    # via fastapi
 sympy==1.14.0
     # via torch
-teleop==0.1.2
-    # via lerobot
-tensorboard==2.20.0
-    # via robomimic
-tensorboard-data-server==0.7.2
-    # via tensorboard
-tensorboardx==2.6.4
-    # via robomimic
 termcolor==3.1.0
-    # via
-    #   lerobot
-    #   robomimic
-thop==0.1.1.post2209072238
-    # via libero
+    # via lerobot
 tifffile==2025.5.10
     # via scikit-image
-timm==1.0.20
-    # via lerobot
-tokenizers==0.22.1
+tokenizers==0.21.4
     # via transformers
 toml==0.10.2
     # via draccus
-tomli==2.3.0
+tomli==2.2.1
     # via
     #   cmeel
     #   coverage
-    #   jupytext
     #   pytest
 torch==2.7.1
     # via
     #   accelerate
     #   lerobot
-    #   peft
-    #   robomimic
-    #   thop
-    #   timm
     #   torchvision
 torchcodec==0.5
     # via lerobot
 torchvision==0.22.1
-    # via
-    #   lerobot
-    #   robomimic
-    #   timm
-tornado==6.5.2
+    # via lerobot
+tornado==6.5.1
     # via meshcat
 tqdm==4.67.1
     # via
     #   datasets
     #   dm-control
     #   huggingface-hub
-    #   peft
-    #   robomimic
     #   transformers
 traitlets==5.14.3
     # via
     #   ipython
-    #   jupyter-core
     #   matplotlib-inline
-    #   nbformat
-transformers==4.57.1
-    # via
-    #   lerobot
-    #   libero
-    #   peft
-transforms3d==0.4.2
-    # via teleop
-typing-extensions==4.15.0
+transformers==4.51.3
+    # via lerobot
+typing-extensions==4.14.1
     # via
     #   aiosignal
-    #   anyio
-    #   etils
     #   exceptiongroup
-    #   fastapi
     #   gymnasium
     #   huggingface-hub
     #   ipython
     #   multidict
     #   pydantic
     #   pydantic-core
-    #   referencing
     #   rerun-sdk
-    #   starlette
     #   torch
     #   typing-inspect
     #   typing-inspection
-    #   uvicorn
-    #   virtualenv
     #   wandb
 typing-inspect==0.9.0
     # via draccus
-typing-inspection==0.4.2
+typing-inspection==0.4.1
     # via pydantic
 tzdata==2025.2
     # via pandas
@@ -795,36 +604,22 @@ urllib3==2.5.0
     # via
     #   requests
     #   sentry-sdk
-uvicorn[standard]==0.38.0
-    # via teleop
-uvloop==0.22.1
-    # via uvicorn
-virtualenv==20.35.3
+virtualenv==20.32.0
     # via pre-commit
-wandb==0.21.4
-    # via
-    #   lerobot
-    #   libero
-watchfiles==1.1.1
-    # via uvicorn
-wcwidth==0.2.14
+wandb==0.21.0
+    # via lerobot
+wcwidth==0.2.13
     # via prompt-toolkit
-websocket-client==1.9.0
-    # via teleop
-websockets==15.0.1
-    # via uvicorn
 werkzeug==3.1.3
-    # via tensorboard
-wrapt==2.0.0
+    # via flask
+wrapt==1.17.2
     # via dm-tree
-xxhash==3.6.0
+xxhash==3.5.0
     # via datasets
-yarl==1.22.0
+yarl==1.20.1
     # via aiohttp
 zipp==3.23.0
-    # via
-    #   etils
-    #   importlib-metadata
+    # via importlib-metadata
 
 # The following packages are considered to be unsafe in a requirements file:
 # setuptools

requirements-ubuntu.txt

@@ -13,62 +13,47 @@ absl-py==2.3.1
     #   dm-tree
     #   labmaze
     #   mujoco
-    #   tensorboard
-accelerate==1.11.0
-    # via
-    #   lerobot
-    #   peft
+accelerate==1.9.0
+    # via lerobot
 aiohappyeyeballs==2.6.1
     # via aiohttp
-aiohttp==3.13.1
+aiohttp==3.12.15
     # via fsspec
 aiosignal==1.4.0
     # via aiohttp
 annotated-types==0.7.0
     # via pydantic
-antlr4-python3-runtime==4.9.3
-    # via
-    #   hydra-core
-    #   omegaconf
-anyio==4.11.0
-    # via
-    #   starlette
-    #   watchfiles
 asttokens==3.0.0
     # via stack-data
 async-timeout==5.0.1
     # via aiohttp
-attrs==25.4.0
+attrs==25.3.0
     # via
     #   aiohttp
     #   dm-tree
     #   jsonlines
-    #   jsonschema
-    #   referencing
     #   rerun-sdk
-av==15.1.0
+av==15.0.0
     # via lerobot
-bddl==1.0.1
-    # via libero
-certifi==2025.10.5
+blinker==1.9.0
+    # via flask
+certifi==2025.7.14
     # via
     #   requests
     #   sentry-sdk
-cffi==2.0.0
+cffi==1.17.1
     # via pymunk
 cfgv==3.4.0
     # via pre-commit
-charset-normalizer==3.4.4
+charset-normalizer==3.4.2
     # via requests
-click==8.3.0
+click==8.2.1
     # via
-    #   uvicorn
+    #   flask
     #   wandb
 cloudpickle==3.1.1
-    # via
-    #   gymnasium
-    #   libero
-cmake==4.1.0
+    # via gymnasium
+cmake==4.0.3
     # via lerobot
 cmeel==0.57.3
     # via
@@ -110,29 +95,27 @@ coal-library==3.0.1
     # via pin
 contourpy==1.3.2
     # via matplotlib
-coverage[toml]==7.11.0
+coverage[toml]==7.10.1
     # via pytest-cov
 cycler==0.12.1
     # via matplotlib
-datasets==4.1.1
+datasets==3.6.0
     # via lerobot
-debugpy==1.8.17
+debugpy==1.8.15
     # via lerobot
 decorator==5.2.1
     # via ipython
-decord==0.6.0
+deepdiff==8.5.0
     # via lerobot
-deepdiff==8.6.1
+diffusers==0.34.0
     # via lerobot
-diffusers==0.35.2
-    # via lerobot
-dill==0.4.0
+dill==0.3.8
     # via
     #   datasets
     #   multiprocess
 distlib==0.4.0
     # via virtualenv
-dm-control==1.0.34
+dm-control==1.0.14
     # via gym-aloha
 dm-env==1.6
     # via dm-control
@@ -140,48 +123,31 @@ dm-tree==0.1.9
     # via
     #   dm-control
     #   dm-env
-    #   lerobot
 docopt==0.6.2
     # via num2words
 draccus==0.10.0
     # via lerobot
-dynamixel-sdk==3.8.4
+dynamixel-sdk==3.7.31
     # via lerobot
-easydict==1.13
-    # via libero
-egl-probe @ git+https://github.com/huggingface/egl_probe.git
-    # via
-    #   libero
-    #   robomimic
 eigenpy==3.10.3
     # via coal-library
 einops==0.8.1
-    # via
-    #   flash-attn
-    #   lerobot
-    #   libero
+    # via lerobot
 eiquadprog==1.2.9
     # via placo
-etils[epath,epy]==1.13.0
-    # via mujoco
 evdev==1.9.2
     # via pynput
 exceptiongroup==1.3.0
     # via
-    #   anyio
     #   ipython
     #   pytest
-executing==2.2.1
+executing==2.2.0
     # via stack-data
 farama-notifications==0.0.4
     # via gymnasium
-fastapi==0.119.1
-    # via teleop
-fastjsonschema==2.21.2
-    # via nbformat
 feetech-servo-sdk==1.0.0
     # via lerobot
-filelock==3.20.0
+filelock==3.18.0
     # via
     #   datasets
     #   diffusers
@@ -189,27 +155,24 @@ filelock==3.20.0
     #   torch
     #   transformers
     #   virtualenv
-flash-attn==2.8.3
+flask==3.1.1
     # via lerobot
-fonttools==4.60.1
+fonttools==4.59.0
     # via matplotlib
-frozenlist==1.8.0
+frozenlist==1.7.0
     # via
     #   aiohttp
     #   aiosignal
-fsspec[http]==2025.9.0
+fsspec[http]==2025.3.0
     # via
     #   datasets
-    #   etils
     #   huggingface-hub
     #   torch
-future==1.0.0
-    # via libero
 gitdb==4.0.12
     # via gitpython
 gitpython==3.1.45
     # via wandb
-glfw==2.10.0
+glfw==2.9.0
     # via
     #   dm-control
     #   mujoco
@@ -217,79 +180,61 @@ grpcio==1.73.1
     # via
     #   grpcio-tools
     #   lerobot
-    #   reachy2-sdk
-    #   reachy2-sdk-api
-    #   tensorboard
 grpcio-tools==1.73.1
-    # via
-    #   lerobot
-    #   reachy2-sdk-api
-gym-aloha==0.1.3
     # via lerobot
-gym-hil==0.1.13
+gym-aloha==0.1.1
     # via lerobot
-gym-pusht==0.1.6
+gym-hil==0.1.10
     # via lerobot
-gymnasium==1.2.1
+gym-pusht==0.1.5
+    # via lerobot
+gym-xarm==0.1.1
+    # via lerobot
+gymnasium==0.29.1
     # via
     #   gym-aloha
     #   gym-hil
     #   gym-pusht
+    #   gym-xarm
+    #   gymnasium-robotics
     #   lerobot
-    #   libero
-    #   metaworld
-h11==0.16.0
-    # via uvicorn
-h5py==3.15.1
-    # via robomimic
-hebi-py==2.11.0
-    # via lerobot
+    #   pettingzoo
+gymnasium-robotics==1.2.4
+    # via gym-xarm
 hf-transfer==0.1.9
     # via huggingface-hub
-hf-xet==1.1.10
+hf-xet==1.1.5
     # via huggingface-hub
 hidapi==0.14.0.post4
     # via
     #   gym-hil
     #   lerobot
-httptools==0.7.1
-    # via uvicorn
-huggingface-hub[cli,hf-transfer]==0.35.3
+huggingface-hub[cli,hf-transfer]==0.34.3
     # via
     #   accelerate
     #   datasets
     #   diffusers
     #   lerobot
-    #   peft
-    #   timm
     #   tokenizers
     #   transformers
-hydra-core==1.3.2
-    # via libero
-identify==2.6.15
+identify==2.6.12
     # via pre-commit
-idna==3.11
+idna==3.10
     # via
-    #   anyio
     #   requests
     #   yarl
 imageio[ffmpeg]==2.37.0
     # via
     #   gym-aloha
     #   gym-hil
+    #   gymnasium-robotics
     #   lerobot
-    #   metaworld
-    #   robomimic
     #   scikit-image
 imageio-ffmpeg==0.6.0
-    # via
-    #   imageio
-    #   robomimic
+    # via imageio
 importlib-metadata==8.7.0
     # via diffusers
-importlib-resources==6.5.2
-    # via etils
-iniconfig==2.3.0
+iniconfig==2.1.0
     # via pytest
 inquirerpy==0.3.4
     # via huggingface-hub
@@ -297,71 +242,50 @@ ipython==8.37.0
     # via meshcat
 ischedule==1.2.7
     # via placo
+itsdangerous==2.2.0
+    # via flask
 jedi==0.19.2
     # via ipython
 jinja2==3.1.6
-    # via torch
+    # via
+    #   flask
+    #   gymnasium-robotics
+    #   torch
 jsonlines==4.0.0
     # via lerobot
-jsonschema==4.25.1
-    # via nbformat
-jsonschema-specifications==2025.9.1
-    # via jsonschema
-jupyter-core==5.9.1
-    # via nbformat
-jupytext==1.18.1
-    # via bddl
-kiwisolver==1.4.9
+kiwisolver==1.4.8
     # via matplotlib
 labmaze==1.0.6
     # via dm-control
 lazy-loader==0.4
     # via scikit-image
-libero @ git+https://github.com/huggingface/lerobot-libero.git@main
-    # via lerobot
-llvmlite==0.45.1
-    # via numba
-lxml==6.0.2
+lxml==6.0.0
     # via dm-control
-markdown==3.9
-    # via tensorboard
-markdown-it-py==4.0.0
-    # via
-    #   jupytext
-    #   mdit-py-plugins
-markupsafe==3.0.3
+markupsafe==3.0.2
     # via
+    #   flask
     #   jinja2
     #   werkzeug
-matplotlib==3.10.7
-    # via
-    #   lerobot
-    #   libero
-matplotlib-inline==0.2.1
+matplotlib==3.10.5
+    # via lerobot
+matplotlib-inline==0.1.7
     # via ipython
-mdit-py-plugins==0.5.0
-    # via jupytext
-mdurl==0.1.2
-    # via markdown-it-py
 mergedeep==1.3.4
     # via draccus
 meshcat==0.3.2
     # via placo
-metaworld==3.0.0
-    # via lerobot
 mock-serial==0.0.1
     # via lerobot
 mpmath==1.3.0
     # via sympy
-mujoco==3.3.7
+mujoco==2.3.7
     # via
     #   dm-control
     #   gym-aloha
     #   gym-hil
-    #   libero
-    #   metaworld
-    #   robosuite
-multidict==6.7.0
+    #   gym-xarm
+    #   gymnasium-robotics
+multidict==6.6.3
     # via
     #   aiohttp
     #   yarl
@@ -369,63 +293,42 @@ multiprocess==0.70.16
     # via datasets
 mypy-extensions==1.1.0
     # via typing-inspect
-nbformat==5.10.4
-    # via jupytext
 networkx==3.4.2
     # via
-    #   bddl
     #   scikit-image
     #   torch
-ninja==1.13.0
-    # via lerobot
 nodeenv==1.9.1
     # via pre-commit
 num2words==0.5.14
     # via lerobot
-numba==0.62.1
-    # via robosuite
 numpy==2.2.6
     # via
     #   accelerate
-    #   bddl
     #   cmeel-boost
     #   contourpy
     #   datasets
-    #   decord
     #   diffusers
     #   dm-control
     #   dm-env
     #   dm-tree
     #   gymnasium
-    #   h5py
-    #   hebi-py
+    #   gymnasium-robotics
     #   imageio
     #   labmaze
-    #   libero
     #   matplotlib
     #   meshcat
-    #   metaworld
     #   mujoco
-    #   numba
     #   opencv-python
     #   opencv-python-headless
     #   pandas
-    #   peft
-    #   pyquaternion
-    #   reachy2-sdk
+    #   pettingzoo
     #   rerun-sdk
-    #   robomimic
-    #   robosuite
     #   scikit-image
     #   scipy
     #   shapely
-    #   teleop
-    #   tensorboard
-    #   tensorboardx
     #   tifffile
     #   torchvision
     #   transformers
-    #   transforms3d
 nvidia-cublas-cu12==12.6.4.1
     # via
     #   nvidia-cudnn-cu12
@@ -463,14 +366,8 @@ nvidia-nvjitlink-cu12==12.6.85
     #   torch
 nvidia-nvtx-cu12==12.6.77
     # via torch
-omegaconf==2.3.0
-    # via hydra-core
 opencv-python==4.12.0.88
-    # via
-    #   gym-pusht
-    #   libero
-    #   reachy2-sdk
-    #   robosuite
+    # via gym-pusht
 opencv-python-headless==4.12.0.88
     # via lerobot
 orderly-set==5.5.0
@@ -480,63 +377,53 @@ packaging==25.0
     #   accelerate
     #   datasets
     #   huggingface-hub
-    #   hydra-core
-    #   jupytext
     #   lazy-loader
     #   lerobot
     #   matplotlib
-    #   peft
     #   pytest
-    #   reachy2-sdk
     #   scikit-image
-    #   tensorboard
-    #   tensorboardx
     #   transformers
     #   wandb
-pandas==2.3.3
+pandas==2.3.1
     # via
     #   datasets
     #   lerobot
-parso==0.8.5
+parso==0.8.4
     # via jedi
-peft==0.17.1
-    # via lerobot
+pettingzoo==1.24.3
+    # via gymnasium-robotics
 pexpect==4.9.0
     # via ipython
 pfzy==0.3.4
     # via inquirerpy
-pillow==12.0.0
+pillow==11.3.0
     # via
     #   diffusers
     #   imageio
-    #   lerobot
     #   matplotlib
     #   meshcat
     #   rerun-sdk
-    #   robosuite
     #   scikit-image
-    #   tensorboard
     #   torchvision
 pin==3.4.0
     # via placo
 placo==0.9.14
     # via lerobot
-platformdirs==4.5.0
+platformdirs==4.3.8
     # via
-    #   jupyter-core
     #   virtualenv
     #   wandb
 pluggy==1.6.0
     # via
     #   pytest
     #   pytest-cov
-pre-commit==4.3.0
+pre-commit==4.2.0
     # via lerobot
-prompt-toolkit==3.0.52
+prompt-toolkit==3.0.51
     # via
     #   inquirerpy
     #   ipython
-propcache==0.4.1
+propcache==0.3.2
     # via
     #   aiohttp
     #   yarl
@@ -545,17 +432,11 @@ protobuf==6.31.0
     #   dm-control
     #   grpcio-tools
     #   lerobot
-    #   reachy2-sdk
-    #   reachy2-sdk-api
-    #   tensorboard
-    #   tensorboardx
     #   wandb
-psutil==7.1.1
+psutil==7.0.0
     # via
     #   accelerate
     #   imageio
-    #   peft
-    #   robomimic
 ptyprocess==0.7.0
     # via pexpect
 pure-eval==0.2.3
@@ -564,13 +445,11 @@ pyarrow==21.0.0
     # via
     #   datasets
     #   rerun-sdk
-pycparser==2.23
+pycparser==2.22
     # via cffi
-pydantic==2.12.3
-    # via
-    #   fastapi
-    #   wandb
-pydantic-core==2.41.4
+pydantic==2.11.7
+    # via wandb
+pydantic-core==2.33.2
     # via pydantic
 pygame==2.6.1
     # via
@@ -585,22 +464,20 @@ pymunk==6.11.1
     # via
     #   gym-pusht
     #   lerobot
-pyngrok==7.4.1
+pyngrok==7.2.12
     # via meshcat
 pynput==1.8.1
     # via
     #   gym-hil
     #   lerobot
-pyopengl==3.1.10
+pyopengl==3.1.9
     # via
     #   dm-control
     #   mujoco
-pyparsing==3.2.5
+pyparsing==3.2.3
     # via
     #   dm-control
     #   matplotlib
-pyquaternion==0.9.9
-    # via reachy2-sdk
 pyrealsense2==2.56.5.9235
     # via lerobot
 pyserial==3.5
@@ -608,14 +485,12 @@ pyserial==3.5
     #   dynamixel-sdk
     #   feetech-servo-sdk
     #   lerobot
-pytest==8.4.2
+pytest==8.4.1
     # via
-    #   bddl
     #   lerobot
     #   pytest-cov
     #   pytest-timeout
-    #   teleop
-pytest-cov==7.0.0
+pytest-cov==6.2.1
     # via lerobot
 pytest-timeout==2.4.0
     # via lerobot
@@ -623,75 +498,48 @@ python-dateutil==2.9.0.post0
     # via
     #   matplotlib
     #   pandas
-python-dotenv==1.1.1
-    # via uvicorn
 python-xlib==0.33
     # via pynput
 pytz==2025.2
     # via pandas
-pyyaml==6.0.3
+pyyaml==6.0.2
     # via
     #   accelerate
     #   datasets
     #   draccus
-    #   hebi-py
     #   huggingface-hub
-    #   jupytext
-    #   omegaconf
-    #   peft
     #   pre-commit
     #   pyngrok
     #   pyyaml-include
-    #   timm
     #   transformers
-    #   uvicorn
     #   wandb
 pyyaml-include==1.4.1
     # via draccus
-pyzmq==27.1.0
+pyzmq==27.0.0
     # via
     #   lerobot
     #   meshcat
-reachy2-sdk==1.0.14
-    # via lerobot
-reachy2-sdk-api==1.0.21
-    # via reachy2-sdk
-referencing==0.37.0
-    # via
-    #   jsonschema
-    #   jsonschema-specifications
-regex==2025.10.23
+regex==2025.7.34
     # via
     #   diffusers
     #   transformers
-requests==2.32.5
+requests==2.32.4
     # via
     #   datasets
     #   diffusers
     #   dm-control
     #   huggingface-hub
-    #   teleop
     #   transformers
     #   wandb
-rerun-sdk==0.26.1
+rerun-sdk==0.22.1
     # via lerobot
 rhoban-cmeel-jsoncpp==1.9.4.9
     # via placo
-robomimic==0.2.0
-    # via libero
-robosuite==1.4.0
-    # via libero
-rpds-py==0.28.0
-    # via
-    #   jsonschema
-    #   referencing
-safetensors==0.6.2
+safetensors==0.5.3
     # via
     #   accelerate
     #   diffusers
     #   lerobot
-    #   peft
-    #   timm
     #   transformers
 scikit-image==0.25.2
     # via
@@ -700,12 +548,10 @@ scikit-image==0.25.2
 scipy==1.15.3
     # via
     #   dm-control
-    #   metaworld
-    #   robosuite
     #   scikit-image
-sentry-sdk==2.42.1
+sentry-sdk==2.34.1
     # via wandb
-shapely==2.1.2
+shapely==2.1.1
     # via gym-pusht
 six==1.17.0
     # via
@@ -714,109 +560,66 @@ six==1.17.0
     #   python-xlib
 smmap==5.0.2
     # via gitdb
-sniffio==1.3.1
-    # via anyio
 stack-data==0.6.3
     # via ipython
-starlette==0.48.0
-    # via fastapi
 sympy==1.14.0
     # via torch
-teleop==0.1.2
-    # via lerobot
-tensorboard==2.20.0
-    # via robomimic
-tensorboard-data-server==0.7.2
-    # via tensorboard
-tensorboardx==2.6.4
-    # via robomimic
 termcolor==3.1.0
-    # via
-    #   lerobot
-    #   robomimic
-thop==0.1.1.post2209072238
-    # via libero
+    # via lerobot
 tifffile==2025.5.10
     # via scikit-image
-timm==1.0.20
-    # via lerobot
-tokenizers==0.22.1
+tokenizers==0.21.4
     # via transformers
 toml==0.10.2
     # via draccus
-tomli==2.3.0
+tomli==2.2.1
     # via
     #   cmeel
     #   coverage
-    #   jupytext
     #   pytest
 torch==2.7.1
     # via
     #   accelerate
-    #   flash-attn
     #   lerobot
-    #   peft
-    #   robomimic
-    #   thop
-    #   timm
     #   torchvision
 torchcodec==0.5
     # via lerobot
 torchvision==0.22.1
-    # via
-    #   lerobot
-    #   robomimic
-    #   timm
-tornado==6.5.2
+    # via lerobot
+tornado==6.5.1
     # via meshcat
 tqdm==4.67.1
     # via
     #   datasets
     #   dm-control
     #   huggingface-hub
-    #   peft
-    #   robomimic
     #   transformers
 traitlets==5.14.3
     # via
     #   ipython
-    #   jupyter-core
     #   matplotlib-inline
-    #   nbformat
-transformers==4.57.1
-    # via
-    #   lerobot
-    #   libero
-    #   peft
-transforms3d==0.4.2
-    # via teleop
+transformers==4.51.3
+    # via lerobot
 triton==3.3.1
     # via torch
-typing-extensions==4.15.0
+typing-extensions==4.14.1
     # via
     #   aiosignal
-    #   anyio
-    #   etils
     #   exceptiongroup
-    #   fastapi
     #   gymnasium
     #   huggingface-hub
     #   ipython
     #   multidict
     #   pydantic
     #   pydantic-core
-    #   referencing
     #   rerun-sdk
-    #   starlette
     #   torch
     #   typing-inspect
     #   typing-inspection
-    #   uvicorn
-    #   virtualenv
     #   wandb
 typing-inspect==0.9.0
     # via draccus
-typing-inspection==0.4.2
+typing-inspection==0.4.1
     # via pydantic
 tzdata==2025.2
     # via pandas
@@ -826,36 +629,22 @@ urllib3==2.5.0
     # via
     #   requests
     #   sentry-sdk
-uvicorn[standard]==0.38.0
-    # via teleop
-uvloop==0.22.1
-    # via uvicorn
-virtualenv==20.35.3
+virtualenv==20.32.0
     # via pre-commit
-wandb==0.21.4
-    # via
-    #   lerobot
-    #   libero
-watchfiles==1.1.1
-    # via uvicorn
-wcwidth==0.2.14
+wandb==0.21.0
+    # via lerobot
+wcwidth==0.2.13
     # via prompt-toolkit
-websocket-client==1.9.0
-    # via teleop
-websockets==15.0.1
-    # via uvicorn
 werkzeug==3.1.3
-    # via tensorboard
-wrapt==2.0.0
+    # via flask
+wrapt==1.17.2
     # via dm-tree
-xxhash==3.6.0
+xxhash==3.5.0
     # via datasets
-yarl==1.22.0
+yarl==1.20.1
     # via aiohttp
 zipp==3.23.0
-    # via
-    #   etils
-    #   importlib-metadata
+    # via importlib-metadata
 
 # The following packages are considered to be unsafe in a requirements file:
 # setuptools

requirements.in

@@ -1,9 +1,9 @@
 # requirements.in
 
-# requirements-macos.txt was generated on macOS and is platform-specific (macOS 26.0.1 25A362 arm64).
-# Darwin MacBook-Pro.local 25.0.0 Darwin Kernel Version 25.0.0: Wed Sep 17 21:42:08 PDT 2025; root:xnu-12377.1.9~141/RELEASE_ARM64_T8132 arm64
+# requirements-macos.txt was generated on macOS and is platform-specific (macOS 15.5 24F74 arm64).
+# Darwin MacBook-Pro.local 24.5.0 Darwin Kernel Version 24.5.0: Tue Apr 22 19:54:43 PDT 2025; root:xnu-11417.121.6~2/RELEASE_ARM64_T8132 arm64
 
-# requirements-ubuntu.txt was generated on Linux and is platform-specific (Ubuntu 24.04.3 LTS x86_64).
-# Linux mlerobot-linux 6.14.0-33-generic #33~24.04.1-Ubuntu SMP PREEMPT_DYNAMIC Fri Sep 19 17:02:30 UTC 2 x86_64 x86_64 x86_64 GNU/Linux
+# requirements-ubuntu.txt was generated on Linux and is platform-specific (Ubuntu 24.04.2 LTS x86_64).
+# Linux mlerobot-linux 6.14.0-27-generic #27~24.04.1-Ubuntu SMP PREEMPT_DYNAMIC Tue Jul 22 17:38:49 UTC 2 x86_64 x86_64 x86_64 GNU/Linux
 
 -e .[all]

src/lerobot/__init__.py

@@ -57,6 +57,7 @@ available_tasks_per_env = {
         "AlohaTransferCube-v0",
     ],
     "pusht": ["PushT-v0"],
+    "xarm": ["XarmLift-v0"],
 }
 available_envs = list(available_tasks_per_env.keys())
 
@@ -74,6 +75,16 @@ 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 = [
@@ -184,6 +195,7 @@ available_motors = [
 available_policies_per_env = {
     "aloha": ["act"],
     "pusht": ["diffusion", "vqbet"],
+    "xarm": ["tdmpc"],
     "koch_real": ["act_koch_real"],
     "aloha_real": ["act_aloha_real"],
 }

src/lerobot/async_inference/helpers.py

@@ -16,7 +16,7 @@ import logging
 import logging.handlers
 import os
 import time
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from pathlib import Path
 
 import torch
@@ -268,7 +268,6 @@ class RemotePolicyConfig:
     lerobot_features: dict[str, PolicyFeature]
     actions_per_chunk: int
     device: str = "cpu"
-    rename_map: dict[str, str] = field(default_factory=dict)
 
 
 def _compare_observation_states(obs1_state: torch.Tensor, obs2_state: torch.Tensor, atol: float) -> bool:

src/lerobot/async_inference/policy_server.py

@@ -159,10 +159,7 @@ class PolicyServer(services_pb2_grpc.AsyncInferenceServicer):
         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,
-                "rename_observations_processor": {"rename_map": policy_specs.rename_map},
-            },
+            preprocessor_overrides={"device_processor": device_override},
             postprocessor_overrides={"device_processor": device_override},
         )
 

src/lerobot/cameras/camera.py

@@ -17,7 +17,7 @@
 import abc
 from typing import Any
 
-from numpy.typing import NDArray  # type: ignore  # TODO: add type stubs for numpy.typing
+import numpy as np
 
 from .configs import CameraConfig, ColorMode
 
@@ -89,7 +89,7 @@ class Camera(abc.ABC):
         pass
 
     @abc.abstractmethod
-    def read(self, color_mode: ColorMode | None = None) -> NDArray[Any]:
+    def read(self, color_mode: ColorMode | None = None) -> np.ndarray:
         """Capture and return a single frame from the camera.
 
         Args:
@@ -102,7 +102,7 @@ class Camera(abc.ABC):
         pass
 
     @abc.abstractmethod
-    def async_read(self, timeout_ms: float = ...) -> NDArray[Any]:
+    def async_read(self, timeout_ms: float = ...) -> np.ndarray:
         """Asynchronously capture and return a single frame from the camera.
 
         Args:

src/lerobot/cameras/configs.py

@@ -18,7 +18,7 @@ import abc
 from dataclasses import dataclass
 from enum import Enum
 
-import draccus  # type: ignore  # TODO: add type stubs for draccus
+import draccus
 
 
 class ColorMode(str, Enum):
@@ -34,11 +34,11 @@ class Cv2Rotation(int, Enum):
 
 
 @dataclass(kw_only=True)
-class CameraConfig(draccus.ChoiceRegistry, abc.ABC):  # type: ignore  # TODO: add type stubs for draccus
+class CameraConfig(draccus.ChoiceRegistry, abc.ABC):
     fps: int | None = None
     width: int | None = None
     height: int | None = None
 
     @property
     def type(self) -> str:
-        return str(self.get_choice_name(self.__class__))
+        return self.get_choice_name(self.__class__)

src/lerobot/cameras/opencv/__init__.py

@@ -14,5 +14,3 @@
 
 from .camera_opencv import OpenCVCamera
 from .configuration_opencv import OpenCVCameraConfig
-
-__all__ = ["OpenCVCamera", "OpenCVCameraConfig"]

src/lerobot/cameras/opencv/camera_opencv.py

@@ -25,12 +25,11 @@ from pathlib import Path
 from threading import Event, Lock, Thread
 from typing import Any
 
-from numpy.typing import NDArray  # type: ignore  # TODO: add type stubs for numpy.typing
-
 # Fix MSMF hardware transform compatibility for Windows before importing cv2
 if platform.system() == "Windows" and "OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS" not in os.environ:
     os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
-import cv2  # type: ignore  # TODO: add type stubs for OpenCV
+import cv2
+import numpy as np
 
 from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
@@ -122,7 +121,7 @@ class OpenCVCamera(Camera):
         self.thread: Thread | None = None
         self.stop_event: Event | None = None
         self.frame_lock: Lock = Lock()
-        self.latest_frame: NDArray[Any] | None = None
+        self.latest_frame: np.ndarray | None = None
         self.new_frame_event: Event = Event()
 
         self.rotation: int | None = get_cv2_rotation(config.rotation)
@@ -141,7 +140,7 @@ class OpenCVCamera(Camera):
         """Checks if the camera is currently connected and opened."""
         return isinstance(self.videocapture, cv2.VideoCapture) and self.videocapture.isOpened()
 
-    def connect(self, warmup: bool = True) -> None:
+    def connect(self, warmup: bool = True):
         """
         Connects to the OpenCV camera specified in the configuration.
 
@@ -181,14 +180,12 @@ class OpenCVCamera(Camera):
 
     def _configure_capture_settings(self) -> None:
         """
-        Applies the specified FOURCC, FPS, width, and height settings to the connected camera.
+        Applies the specified FPS, width, and height settings to the connected camera.
 
         This method attempts to set the camera properties via OpenCV. It checks if
         the camera successfully applied the settings and raises an error if not.
-        FOURCC is set first (if specified) as it can affect the available FPS and resolution options.
 
         Args:
-            fourcc: The desired FOURCC code (e.g., "MJPG", "YUYV"). If None, auto-detect.
             fps: The desired frames per second. If None, the setting is skipped.
             width: The desired capture width. If None, the setting is skipped.
             height: The desired capture height. If None, the setting is skipped.
@@ -202,11 +199,10 @@ class OpenCVCamera(Camera):
         if not self.is_connected:
             raise DeviceNotConnectedError(f"Cannot configure settings for {self} as it is not connected.")
 
-        # Set FOURCC first (if specified) as it can affect available FPS/resolution options
-        if self.config.fourcc is not None:
-            self._validate_fourcc()
-        if self.videocapture is None:
-            raise DeviceNotConnectedError(f"{self} videocapture is not initialized")
+        if self.fps is None:
+            self.fps = self.videocapture.get(cv2.CAP_PROP_FPS)
+        else:
+            self._validate_fps()
 
         default_width = int(round(self.videocapture.get(cv2.CAP_PROP_FRAME_WIDTH)))
         default_height = int(round(self.videocapture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
@@ -220,56 +216,18 @@ class OpenCVCamera(Camera):
         else:
             self._validate_width_and_height()
 
-        if self.fps is None:
-            self.fps = self.videocapture.get(cv2.CAP_PROP_FPS)
-        else:
-            self._validate_fps()
-
     def _validate_fps(self) -> None:
         """Validates and sets the camera's frames per second (FPS)."""
 
-        if self.videocapture is None:
-            raise DeviceNotConnectedError(f"{self} videocapture is not initialized")
-
-        if self.fps is None:
-            raise ValueError(f"{self} FPS is not set")
-
         success = self.videocapture.set(cv2.CAP_PROP_FPS, float(self.fps))
         actual_fps = self.videocapture.get(cv2.CAP_PROP_FPS)
         # Use math.isclose for robust float comparison
         if not success or not math.isclose(self.fps, actual_fps, rel_tol=1e-3):
             raise RuntimeError(f"{self} failed to set fps={self.fps} ({actual_fps=}).")
 
-    def _validate_fourcc(self) -> None:
-        """Validates and sets the camera's FOURCC code."""
-
-        fourcc_code = cv2.VideoWriter_fourcc(*self.config.fourcc)
-
-        if self.videocapture is None:
-            raise DeviceNotConnectedError(f"{self} videocapture is not initialized")
-
-        success = self.videocapture.set(cv2.CAP_PROP_FOURCC, fourcc_code)
-        actual_fourcc_code = self.videocapture.get(cv2.CAP_PROP_FOURCC)
-
-        # Convert actual FOURCC code back to string for comparison
-        actual_fourcc_code_int = int(actual_fourcc_code)
-        actual_fourcc = "".join([chr((actual_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)])
-
-        if not success or actual_fourcc != self.config.fourcc:
-            logger.warning(
-                f"{self} failed to set fourcc={self.config.fourcc} (actual={actual_fourcc}, success={success}). "
-                f"Continuing with default format."
-            )
-
     def _validate_width_and_height(self) -> None:
         """Validates and sets the camera's frame capture width and height."""
 
-        if self.videocapture is None:
-            raise DeviceNotConnectedError(f"{self} videocapture is not initialized")
-
-        if self.capture_width is None or self.capture_height is None:
-            raise ValueError(f"{self} capture_width or capture_height is not set")
-
         width_success = self.videocapture.set(cv2.CAP_PROP_FRAME_WIDTH, float(self.capture_width))
         height_success = self.videocapture.set(cv2.CAP_PROP_FRAME_HEIGHT, float(self.capture_height))
 
@@ -300,12 +258,11 @@ class OpenCVCamera(Camera):
         """
         found_cameras_info = []
 
-        targets_to_scan: list[str | int]
         if platform.system() == "Linux":
             possible_paths = sorted(Path("/dev").glob("video*"), key=lambda p: p.name)
             targets_to_scan = [str(p) for p in possible_paths]
         else:
-            targets_to_scan = [int(i) for i in range(MAX_OPENCV_INDEX)]
+            targets_to_scan = list(range(MAX_OPENCV_INDEX))
 
         for target in targets_to_scan:
             camera = cv2.VideoCapture(target)
@@ -314,12 +271,6 @@ class OpenCVCamera(Camera):
                 default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
                 default_fps = camera.get(cv2.CAP_PROP_FPS)
                 default_format = camera.get(cv2.CAP_PROP_FORMAT)
-
-                # Get FOURCC code and convert to string
-                default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC)
-                default_fourcc_code_int = int(default_fourcc_code)
-                default_fourcc = "".join([chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)])
-
                 camera_info = {
                     "name": f"OpenCV Camera @ {target}",
                     "type": "OpenCV",
@@ -327,7 +278,6 @@ class OpenCVCamera(Camera):
                     "backend_api": camera.getBackendName(),
                     "default_stream_profile": {
                         "format": default_format,
-                        "fourcc": default_fourcc,
                         "width": default_width,
                         "height": default_height,
                         "fps": default_fps,
@@ -339,7 +289,7 @@ class OpenCVCamera(Camera):
 
         return found_cameras_info
 
-    def read(self, color_mode: ColorMode | None = None) -> NDArray[Any]:
+    def read(self, color_mode: ColorMode | None = None) -> np.ndarray:
         """
         Reads a single frame synchronously from the camera.
 
@@ -367,9 +317,6 @@ class OpenCVCamera(Camera):
 
         start_time = time.perf_counter()
 
-        if self.videocapture is None:
-            raise DeviceNotConnectedError(f"{self} videocapture is not initialized")
-
         ret, frame = self.videocapture.read()
 
         if not ret or frame is None:
@@ -382,7 +329,7 @@ class OpenCVCamera(Camera):
 
         return processed_frame
 
-    def _postprocess_image(self, image: NDArray[Any], color_mode: ColorMode | None = None) -> NDArray[Any]:
+    def _postprocess_image(self, image: np.ndarray, color_mode: ColorMode | None = None) -> np.ndarray:
         """
         Applies color conversion, dimension validation, and rotation to a raw frame.
 
@@ -425,7 +372,7 @@ class OpenCVCamera(Camera):
 
         return processed_image
 
-    def _read_loop(self) -> None:
+    def _read_loop(self):
         """
         Internal loop run by the background thread for asynchronous reading.
 
@@ -436,9 +383,6 @@ class OpenCVCamera(Camera):
 
         Stops on DeviceNotConnectedError, logs other errors and continues.
         """
-        if self.stop_event is None:
-            raise RuntimeError(f"{self}: stop_event is not initialized before starting read loop.")
-
         while not self.stop_event.is_set():
             try:
                 color_image = self.read()
@@ -475,7 +419,7 @@ class OpenCVCamera(Camera):
         self.thread = None
         self.stop_event = None
 
-    def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
+    def async_read(self, timeout_ms: float = 200) -> np.ndarray:
         """
         Reads the latest available frame asynchronously.
 
@@ -518,7 +462,7 @@ class OpenCVCamera(Camera):
 
         return frame
 
-    def disconnect(self) -> None:
+    def disconnect(self):
         """
         Disconnects from the camera and cleans up resources.
 

src/lerobot/cameras/opencv/configuration_opencv.py

@@ -17,8 +17,6 @@ from pathlib import Path
 
 from ..configs import CameraConfig, ColorMode, Cv2Rotation
 
-__all__ = ["OpenCVCameraConfig", "ColorMode", "Cv2Rotation"]
-
 
 @CameraConfig.register_subclass("opencv")
 @dataclass
@@ -35,9 +33,8 @@ class OpenCVCameraConfig(CameraConfig):
     OpenCVCameraConfig(0, 30, 1280, 720)   # 1280x720 @ 30FPS
     OpenCVCameraConfig(/dev/video4, 60, 640, 480)   # 640x480 @ 60FPS
 
-    # Advanced configurations with FOURCC format
-    OpenCVCameraConfig(128422271347, 30, 640, 480, rotation=Cv2Rotation.ROTATE_90, fourcc="MJPG")     # With 90° rotation and MJPG format
-    OpenCVCameraConfig(0, 30, 1280, 720, fourcc="YUYV")     # With YUYV format
+    # Advanced configurations
+    OpenCVCameraConfig(128422271347, 30, 640, 480, rotation=Cv2Rotation.ROTATE_90)     # With 90° rotation
     ```
 
     Attributes:
@@ -49,21 +46,17 @@ class OpenCVCameraConfig(CameraConfig):
         color_mode: Color mode for image output (RGB or BGR). Defaults to RGB.
         rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
         warmup_s: Time reading frames before returning from connect (in seconds)
-        fourcc: FOURCC code for video format (e.g., "MJPG", "YUYV", "I420"). Defaults to None (auto-detect).
 
     Note:
         - Only 3-channel color output (RGB/BGR) is currently supported.
-        - FOURCC codes must be 4-character strings (e.g., "MJPG", "YUYV"). Some common FOUCC codes: https://learn.microsoft.com/en-us/windows/win32/medfound/video-fourccs#fourcc-constants
-        - Setting FOURCC can help achieve higher frame rates on some cameras.
     """
 
     index_or_path: int | Path
     color_mode: ColorMode = ColorMode.RGB
     rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
     warmup_s: int = 1
-    fourcc: str | None = None
 
-    def __post_init__(self) -> None:
+    def __post_init__(self):
         if self.color_mode not in (ColorMode.RGB, ColorMode.BGR):
             raise ValueError(
                 f"`color_mode` is expected to be {ColorMode.RGB.value} or {ColorMode.BGR.value}, but {self.color_mode} is provided."
@@ -78,8 +71,3 @@ class OpenCVCameraConfig(CameraConfig):
             raise ValueError(
                 f"`rotation` is expected to be in {(Cv2Rotation.NO_ROTATION, Cv2Rotation.ROTATE_90, Cv2Rotation.ROTATE_180, Cv2Rotation.ROTATE_270)}, but {self.rotation} is provided."
             )
-
-        if self.fourcc is not None and (not isinstance(self.fourcc, str) or len(self.fourcc) != 4):
-            raise ValueError(
-                f"`fourcc` must be a 4-character string (e.g., 'MJPG', 'YUYV'), but '{self.fourcc}' is provided."
-            )

src/lerobot/cameras/reachy2_camera/configuration_reachy2_camera.py

@@ -16,8 +16,6 @@ from dataclasses import dataclass
 
 from ..configs import CameraConfig, ColorMode
 
-__all__ = ["CameraConfig", "ColorMode", "Reachy2CameraConfig"]
-
 
 @CameraConfig.register_subclass("reachy2_camera")
 @dataclass
@@ -64,7 +62,7 @@ class Reachy2CameraConfig(CameraConfig):
     port: int = 50065
     # use_depth: bool = False
 
-    def __post_init__(self) -> None:
+    def __post_init__(self):
         if self.name not in ["teleop", "depth"]:
             raise ValueError(f"`name` is expected to be 'teleop' or 'depth', but {self.name} is provided.")
         if (self.name == "teleop" and self.image_type not in ["left", "right"]) or (

src/lerobot/cameras/reachy2_camera/reachy2_camera.py

@@ -23,17 +23,13 @@ import time
 from threading import Event, Lock, Thread
 from typing import Any
 
-from numpy.typing import NDArray  # type: ignore  # TODO: add type stubs for numpy.typing
-
 # Fix MSMF hardware transform compatibility for Windows before importing cv2
 if platform.system() == "Windows" and "OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS" not in os.environ:
     os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
-import cv2  # type: ignore  # TODO: add type stubs for OpenCV
-import numpy as np  # type: ignore  # TODO: add type stubs for numpy
-from reachy2_sdk.media.camera import CameraView  # type: ignore  # TODO: add type stubs for reachy2_sdk
-from reachy2_sdk.media.camera_manager import (  # type: ignore  # TODO: add type stubs for reachy2_sdk
-    CameraManager,
-)
+import cv2
+import numpy as np
+from reachy2_sdk.media.camera import CameraView
+from reachy2_sdk.media.camera_manager import CameraManager
 
 from lerobot.utils.errors import DeviceNotConnectedError
 
@@ -77,7 +73,7 @@ class Reachy2Camera(Camera):
         self.thread: Thread | None = None
         self.stop_event: Event | None = None
         self.frame_lock: Lock = Lock()
-        self.latest_frame: NDArray[Any] | None = None
+        self.latest_frame: np.ndarray | None = None
         self.new_frame_event: Event = Event()
 
     def __str__(self) -> str:
@@ -87,17 +83,13 @@ class Reachy2Camera(Camera):
     def is_connected(self) -> bool:
         """Checks if the camera is currently connected and opened."""
         if self.config.name == "teleop":
-            return bool(
-                self.cam_manager._grpc_connected and self.cam_manager.teleop if self.cam_manager else False
-            )
+            return self.cam_manager._grpc_connected and self.cam_manager.teleop if self.cam_manager else False
         elif self.config.name == "depth":
-            return bool(
-                self.cam_manager._grpc_connected and self.cam_manager.depth if self.cam_manager else False
-            )
+            return self.cam_manager._grpc_connected and self.cam_manager.depth if self.cam_manager else False
         else:
             raise ValueError(f"Invalid camera name '{self.config.name}'. Expected 'teleop' or 'depth'.")
 
-    def connect(self, warmup: bool = True) -> None:
+    def connect(self, warmup: bool = True):
         """
         Connects to the Reachy2 CameraManager as specified in the configuration.
         """
@@ -139,7 +131,7 @@ class Reachy2Camera(Camera):
         camera_manager.disconnect()
         return initialized_cameras
 
-    def read(self, color_mode: ColorMode | None = None) -> NDArray[Any]:
+    def read(self, color_mode: ColorMode | None = None) -> np.ndarray:
         """
         Reads a single frame synchronously from the camera.
 
@@ -160,7 +152,7 @@ class Reachy2Camera(Camera):
 
         start_time = time.perf_counter()
 
-        frame: NDArray[Any] = np.empty((0, 0, 3), dtype=np.uint8)
+        frame = None
 
         if self.cam_manager is None:
             raise DeviceNotConnectedError(f"{self} is not connected.")
@@ -187,7 +179,7 @@ class Reachy2Camera(Camera):
 
         return frame
 
-    def _read_loop(self) -> None:
+    def _read_loop(self):
         """
         Internal loop run by the background thread for asynchronous reading.
 
@@ -198,9 +190,6 @@ class Reachy2Camera(Camera):
 
         Stops on DeviceNotConnectedError, logs other errors and continues.
         """
-        if self.stop_event is None:
-            raise RuntimeError(f"{self}: stop_event is not initialized before starting read loop.")
-
         while not self.stop_event.is_set():
             try:
                 color_image = self.read()
@@ -237,7 +226,7 @@ class Reachy2Camera(Camera):
         self.thread = None
         self.stop_event = None
 
-    def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
+    def async_read(self, timeout_ms: float = 200) -> np.ndarray:
         """
         Reads the latest available frame asynchronously.
 
@@ -280,7 +269,7 @@ class Reachy2Camera(Camera):
 
         return frame
 
-    def disconnect(self) -> None:
+    def disconnect(self):
         """
         Stops the background read thread (if running).
 

src/lerobot/cameras/realsense/camera_realsense.py

@@ -21,12 +21,11 @@ import time
 from threading import Event, Lock, Thread
 from typing import Any
 
-import cv2  # type: ignore  # TODO: add type stubs for OpenCV
-import numpy as np  # type: ignore  # TODO: add type stubs for numpy
-from numpy.typing import NDArray  # type: ignore  # TODO: add type stubs for numpy.typing
+import cv2
+import numpy as np
 
 try:
-    import pyrealsense2 as rs  # type: ignore  # TODO: add type stubs for pyrealsense2
+    import pyrealsense2 as rs
 except Exception as e:
     logging.info(f"Could not import realsense: {e}")
 
@@ -133,7 +132,7 @@ class RealSenseCamera(Camera):
         self.thread: Thread | None = None
         self.stop_event: Event | None = None
         self.frame_lock: Lock = Lock()
-        self.latest_frame: NDArray[Any] | None = None
+        self.latest_frame: np.ndarray | None = None
         self.new_frame_event: Event = Event()
 
         self.rotation: int | None = get_cv2_rotation(config.rotation)
@@ -151,7 +150,7 @@ class RealSenseCamera(Camera):
         """Checks if the camera pipeline is started and streams are active."""
         return self.rs_pipeline is not None and self.rs_profile is not None
 
-    def connect(self, warmup: bool = True) -> None:
+    def connect(self, warmup: bool = True):
         """
         Connects to the RealSense camera specified in the configuration.
 
@@ -265,7 +264,7 @@ class RealSenseCamera(Camera):
         serial_number = str(found_devices[0]["serial_number"])
         return serial_number
 
-    def _configure_rs_pipeline_config(self, rs_config: Any) -> None:
+    def _configure_rs_pipeline_config(self, rs_config):
         """Creates and configures the RealSense pipeline configuration object."""
         rs.config.enable_device(rs_config, self.serial_number)
 
@@ -294,9 +293,6 @@ class RealSenseCamera(Camera):
         if not self.is_connected:
             raise DeviceNotConnectedError(f"Cannot validate settings for {self} as it is not connected.")
 
-        if self.rs_profile is None:
-            raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
-
         stream = self.rs_profile.get_stream(rs.stream.color).as_video_stream_profile()
 
         if self.fps is None:
@@ -312,7 +308,7 @@ class RealSenseCamera(Camera):
                 self.width, self.height = actual_width, actual_height
                 self.capture_width, self.capture_height = actual_width, actual_height
 
-    def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
+    def read_depth(self, timeout_ms: int = 200) -> np.ndarray:
         """
         Reads a single frame (depth) synchronously from the camera.
 
@@ -340,9 +336,6 @@ class RealSenseCamera(Camera):
 
         start_time = time.perf_counter()
 
-        if self.rs_pipeline is None:
-            raise RuntimeError(f"{self}: rs_pipeline must be initialized before use.")
-
         ret, frame = self.rs_pipeline.try_wait_for_frames(timeout_ms=timeout_ms)
 
         if not ret or frame is None:
@@ -358,7 +351,7 @@ class RealSenseCamera(Camera):
 
         return depth_map_processed
 
-    def read(self, color_mode: ColorMode | None = None, timeout_ms: int = 200) -> NDArray[Any]:
+    def read(self, color_mode: ColorMode | None = None, timeout_ms: int = 200) -> np.ndarray:
         """
         Reads a single frame (color) synchronously from the camera.
 
@@ -383,9 +376,6 @@ class RealSenseCamera(Camera):
 
         start_time = time.perf_counter()
 
-        if self.rs_pipeline is None:
-            raise RuntimeError(f"{self}: rs_pipeline must be initialized before use.")
-
         ret, frame = self.rs_pipeline.try_wait_for_frames(timeout_ms=timeout_ms)
 
         if not ret or frame is None:
@@ -402,8 +392,8 @@ class RealSenseCamera(Camera):
         return color_image_processed
 
     def _postprocess_image(
-        self, image: NDArray[Any], color_mode: ColorMode | None = None, depth_frame: bool = False
-    ) -> NDArray[Any]:
+        self, image: np.ndarray, color_mode: ColorMode | None = None, depth_frame: bool = False
+    ) -> np.ndarray:
         """
         Applies color conversion, dimension validation, and rotation to a raw color frame.
 
@@ -448,7 +438,7 @@ class RealSenseCamera(Camera):
 
         return processed_image
 
-    def _read_loop(self) -> None:
+    def _read_loop(self):
         """
         Internal loop run by the background thread for asynchronous reading.
 
@@ -459,9 +449,6 @@ class RealSenseCamera(Camera):
 
         Stops on DeviceNotConnectedError, logs other errors and continues.
         """
-        if self.stop_event is None:
-            raise RuntimeError(f"{self}: stop_event is not initialized before starting read loop.")
-
         while not self.stop_event.is_set():
             try:
                 color_image = self.read(timeout_ms=500)
@@ -487,7 +474,7 @@ class RealSenseCamera(Camera):
         self.thread.daemon = True
         self.thread.start()
 
-    def _stop_read_thread(self) -> None:
+    def _stop_read_thread(self):
         """Signals the background read thread to stop and waits for it to join."""
         if self.stop_event is not None:
             self.stop_event.set()
@@ -499,7 +486,7 @@ class RealSenseCamera(Camera):
         self.stop_event = None
 
     # NOTE(Steven): Missing implementation for depth for now
-    def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
+    def async_read(self, timeout_ms: float = 200) -> np.ndarray:
         """
         Reads the latest available frame data (color) asynchronously.
 
@@ -542,7 +529,7 @@ class RealSenseCamera(Camera):
 
         return frame
 
-    def disconnect(self) -> None:
+    def disconnect(self):
         """
         Disconnects from the camera, stops the pipeline, and cleans up resources.
 

src/lerobot/cameras/realsense/configuration_realsense.py

@@ -59,7 +59,7 @@ class RealSenseCameraConfig(CameraConfig):
     rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
     warmup_s: int = 1
 
-    def __post_init__(self) -> None:
+    def __post_init__(self):
         if self.color_mode not in (ColorMode.RGB, ColorMode.BGR):
             raise ValueError(
                 f"`color_mode` is expected to be {ColorMode.RGB.value} or {ColorMode.BGR.value}, but {self.color_mode} is provided."

src/lerobot/cameras/utils.py

@@ -53,14 +53,14 @@ def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> dict[s
 
 
 def get_cv2_rotation(rotation: Cv2Rotation) -> int | None:
-    import cv2  # type: ignore  # TODO: add type stubs for OpenCV
+    import cv2
 
     if rotation == Cv2Rotation.ROTATE_90:
-        return int(cv2.ROTATE_90_CLOCKWISE)
+        return cv2.ROTATE_90_CLOCKWISE
     elif rotation == Cv2Rotation.ROTATE_180:
-        return int(cv2.ROTATE_180)
+        return cv2.ROTATE_180
     elif rotation == Cv2Rotation.ROTATE_270:
-        return int(cv2.ROTATE_90_COUNTERCLOCKWISE)
+        return cv2.ROTATE_90_COUNTERCLOCKWISE
     else:
         return None
 
@@ -69,8 +69,8 @@ def get_cv2_backend() -> int:
     import cv2
 
     if platform.system() == "Windows":
-        return int(cv2.CAP_MSMF)  # Use MSMF for Windows instead of AVFOUNDATION
+        return cv2.CAP_MSMF  # Use MSMF for Windows instead of AVFOUNDATION
     # elif platform.system() == "Darwin":  # macOS
     #     return cv2.CAP_AVFOUNDATION
     else:  # Linux and others
-        return int(cv2.CAP_ANY)
+        return cv2.CAP_ANY

src/lerobot/configs/default.py

@@ -57,7 +57,7 @@ class EvalConfig:
     # `use_async_envs` specifies whether to use asynchronous environments (multiprocessing).
     use_async_envs: bool = False
 
-    def __post_init__(self) -> None:
+    def __post_init__(self):
         if self.batch_size > self.n_episodes:
             raise ValueError(
                 "The eval batch size is greater than the number of eval episodes "

src/lerobot/configs/eval.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 
 import datetime as dt
+import logging
 from dataclasses import dataclass, field
-from logging import getLogger
 from pathlib import Path
 
 from lerobot import envs, policies  # noqa: F401
@@ -22,8 +22,6 @@ from lerobot.configs import parser
 from lerobot.configs.default import EvalConfig
 from lerobot.configs.policies import PreTrainedConfig
 
-logger = getLogger(__name__)
-
 
 @dataclass
 class EvalPipelineConfig:
@@ -36,31 +34,25 @@ class EvalPipelineConfig:
     output_dir: Path | None = None
     job_name: str | None = None
     seed: int | None = 1000
-    # Rename map for the observation to override the image and state keys
-    rename_map: dict[str, str] = field(default_factory=dict)
 
-    def __post_init__(self) -> None:
+    def __post_init__(self):
         # HACK: We parse again the cli args here to get the pretrained path if there was one.
         policy_path = parser.get_path_arg("policy")
         if policy_path:
             cli_overrides = parser.get_cli_overrides("policy")
             self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
-            self.policy.pretrained_path = Path(policy_path)
+            self.policy.pretrained_path = policy_path
 
         else:
-            logger.warning(
+            logging.warning(
                 "No pretrained path was provided, evaluated policy will be built from scratch (random weights)."
             )
 
         if not self.job_name:
             if self.env is None:
-                self.job_name = f"{self.policy.type if self.policy is not None else 'scratch'}"
+                self.job_name = f"{self.policy.type}"
             else:
-                self.job_name = (
-                    f"{self.env.type}_{self.policy.type if self.policy is not None else 'scratch'}"
-                )
-
-            logger.warning(f"No job name provided, using '{self.job_name}' as job name.")
+                self.job_name = f"{self.env.type}_{self.policy.type}"
 
         if not self.output_dir:
             now = dt.datetime.now()

src/lerobot/configs/parser.py

@@ -16,19 +16,14 @@ import inspect
 import pkgutil
 import sys
 from argparse import ArgumentError
-from collections.abc import Callable, Iterable, Sequence
+from collections.abc import Sequence
 from functools import wraps
 from pathlib import Path
-from pkgutil import ModuleInfo
-from types import ModuleType
-from typing import Any, TypeVar, cast
 
 import draccus
 
 from lerobot.utils.utils import has_method
 
-F = TypeVar("F", bound=Callable[..., object])
-
 PATH_KEY = "path"
 PLUGIN_DISCOVERY_SUFFIX = "discover_packages_path"
 
@@ -65,7 +60,7 @@ def parse_arg(arg_name: str, args: Sequence[str] | None = None) -> str | None:
     return None
 
 
-def parse_plugin_args(plugin_arg_suffix: str, args: Sequence[str]) -> dict[str, str]:
+def parse_plugin_args(plugin_arg_suffix: str, args: Sequence[str]) -> dict:
     """Parse plugin-related arguments from command-line arguments.
 
     This function extracts arguments from command-line arguments that match a specified suffix pattern.
@@ -132,7 +127,7 @@ def load_plugin(plugin_path: str) -> None:
             f"Failed to load plugin '{plugin_path}'. Verify the path and installation: {str(e)}"
         ) from e
 
-    def iter_namespace(ns_pkg: ModuleType) -> Iterable[ModuleInfo]:
+    def iter_namespace(ns_pkg):
         return pkgutil.iter_modules(ns_pkg.__path__, ns_pkg.__name__ + ".")
 
     try:
@@ -153,8 +148,6 @@ def get_type_arg(field_name: str, args: Sequence[str] | None = None) -> str | No
 
 
 def filter_arg(field_to_filter: str, args: Sequence[str] | None = None) -> list[str]:
-    if args is None:
-        return []
     return [arg for arg in args if not arg.startswith(f"--{field_to_filter}=")]
 
 
@@ -178,8 +171,7 @@ def filter_path_args(fields_to_filter: str | list[str], args: Sequence[str] | No
     if isinstance(fields_to_filter, str):
         fields_to_filter = [fields_to_filter]
 
-    filtered_args = [] if args is None else list(args)
-
+    filtered_args = args
     for field in fields_to_filter:
         if get_path_arg(field, args):
             if get_type_arg(field, args):
@@ -192,7 +184,7 @@ def filter_path_args(fields_to_filter: str | list[str], args: Sequence[str] | No
     return filtered_args
 
 
-def wrap(config_path: Path | None = None) -> Callable[[F], F]:
+def wrap(config_path: Path | None = None):
     """
     HACK: Similar to draccus.wrap but does three additional things:
         - Will remove '.path' arguments from CLI in order to process them later on.
@@ -203,9 +195,9 @@ def wrap(config_path: Path | None = None) -> Callable[[F], F]:
             from the CLI '.type' arguments
     """
 
-    def wrapper_outer(fn: F) -> F:
+    def wrapper_outer(fn):
         @wraps(fn)
-        def wrapper_inner(*args: Any, **kwargs: Any) -> Any:
+        def wrapper_inner(*args, **kwargs):
             argspec = inspect.getfullargspec(fn)
             argtype = argspec.annotations[argspec.args[0]]
             if len(args) > 0 and type(args[0]) is argtype:
@@ -233,6 +225,6 @@ def wrap(config_path: Path | None = None) -> Callable[[F], F]:
             response = fn(cfg, *args, **kwargs)
             return response
 
-        return cast(F, wrapper_inner)
+        return wrapper_inner
 
-    return cast(Callable[[F], F], wrapper_outer)
+    return wrapper_outer

src/lerobot/configs/policies.py

@@ -14,12 +14,12 @@
 import abc
 import builtins
 import json
+import logging
 import os
 import tempfile
 from dataclasses import dataclass, field
-from logging import getLogger
 from pathlib import Path
-from typing import Any, TypeVar
+from typing import TypeVar
 
 import draccus
 from huggingface_hub import hf_hub_download
@@ -34,11 +34,10 @@ from lerobot.utils.hub import HubMixin
 from lerobot.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
 
 T = TypeVar("T", bound="PreTrainedConfig")
-logger = getLogger(__name__)
 
 
 @dataclass
-class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: ignore[misc,name-defined] #TODO: draccus issue
+class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
     """
     Base configuration class for policy models.
 
@@ -58,12 +57,12 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: igno
     input_features: dict[str, PolicyFeature] = field(default_factory=dict)
     output_features: dict[str, PolicyFeature] = field(default_factory=dict)
 
-    device: str | None = None  # e.g. "cuda", "cuda:0", "cpu", or "mps"
+    device: str | None = None  # cuda | cpu | mp
     # `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
     # automatic gradient scaling is used.
     use_amp: bool = False
 
-    push_to_hub: bool = True  # type: ignore[assignment] # TODO: use a different name to avoid override
+    push_to_hub: bool = True
     repo_id: str | None = None
 
     # Upload on private repository on the Hugging Face hub.
@@ -74,41 +73,38 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: igno
     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: Path | None = None
+    pretrained_path: str | None = None
 
-    def __post_init__(self) -> None:
+    def __post_init__(self):
         if not self.device or not is_torch_device_available(self.device):
             auto_device = auto_select_torch_device()
-            logger.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
+            logging.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
             self.device = auto_device.type
 
         # Automatically deactivate AMP if necessary
         if self.use_amp and not is_amp_available(self.device):
-            logger.warning(
+            logging.warning(
                 f"Automatic Mixed Precision (amp) is not available on device '{self.device}'. Deactivating AMP."
             )
             self.use_amp = False
 
     @property
     def type(self) -> str:
-        choice_name = self.get_choice_name(self.__class__)
-        if not isinstance(choice_name, str):
-            raise TypeError(f"Expected string from get_choice_name, got {type(choice_name)}")
-        return choice_name
+        return self.get_choice_name(self.__class__)
 
     @property
     @abc.abstractmethod
-    def observation_delta_indices(self) -> list | None:  # type: ignore[type-arg] #TODO: No implementation
+    def observation_delta_indices(self) -> list | None:
         raise NotImplementedError
 
     @property
     @abc.abstractmethod
-    def action_delta_indices(self) -> list | None:  # type: ignore[type-arg]    #TODO: No implementation
+    def action_delta_indices(self) -> list | None:
         raise NotImplementedError
 
     @property
     @abc.abstractmethod
-    def reward_delta_indices(self) -> list | None:  # type: ignore[type-arg]    #TODO: No implementation
+    def reward_delta_indices(self) -> list | None:
         raise NotImplementedError
 
     @abc.abstractmethod
@@ -158,13 +154,13 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: igno
         pretrained_name_or_path: str | Path,
         *,
         force_download: bool = False,
-        resume_download: bool | None = None,
-        proxies: dict[Any, Any] | None = None,
+        resume_download: bool = None,
+        proxies: dict | None = None,
         token: str | bool | None = None,
         cache_dir: str | Path | None = None,
         local_files_only: bool = False,
         revision: str | None = None,
-        **policy_kwargs: Any,
+        **policy_kwargs,
     ) -> T:
         model_id = str(pretrained_name_or_path)
         config_file: str | None = None
@@ -172,7 +168,7 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: igno
             if CONFIG_NAME in os.listdir(model_id):
                 config_file = os.path.join(model_id, CONFIG_NAME)
             else:
-                logger.error(f"{CONFIG_NAME} not found in {Path(model_id).resolve()}")
+                print(f"{CONFIG_NAME} not found in {Path(model_id).resolve()}")
         else:
             try:
                 config_file = hf_hub_download(
@@ -198,9 +194,6 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: igno
         with draccus.config_type("json"):
             orig_config = draccus.parse(cls, config_file, args=[])
 
-        if config_file is None:
-            raise FileNotFoundError(f"{CONFIG_NAME} not found in {model_id}")
-
         with open(config_file) as f:
             config = json.load(f)
 

src/lerobot/configs/train.py

@@ -16,7 +16,6 @@ import datetime as dt
 import os
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any
 
 import draccus
 from huggingface_hub import hf_hub_download
@@ -64,18 +63,18 @@ class TrainPipelineConfig(HubMixin):
     scheduler: LRSchedulerConfig | None = None
     eval: EvalConfig = field(default_factory=EvalConfig)
     wandb: WandBConfig = field(default_factory=WandBConfig)
-    checkpoint_path: Path | None = field(init=False, default=None)
-    # Rename map for the observation to override the image and state keys
-    rename_map: dict[str, str] = field(default_factory=dict)
 
-    def validate(self) -> None:
+    def __post_init__(self):
+        self.checkpoint_path = None
+
+    def validate(self):
         # HACK: We parse again the cli args here to get the pretrained paths if there was some.
         policy_path = parser.get_path_arg("policy")
         if policy_path:
             # Only load the policy config
             cli_overrides = parser.get_cli_overrides("policy")
             self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
-            self.policy.pretrained_path = Path(policy_path)
+            self.policy.pretrained_path = policy_path
         elif self.resume:
             # The entire train config is already loaded, we just need to get the checkpoint dir
             config_path = parser.parse_arg("config_path")
@@ -83,22 +82,14 @@ class TrainPipelineConfig(HubMixin):
                 raise ValueError(
                     f"A config_path is expected when resuming a run. Please specify path to {TRAIN_CONFIG_NAME}"
                 )
-
             if not Path(config_path).resolve().exists():
                 raise NotADirectoryError(
                     f"{config_path=} is expected to be a local path. "
                     "Resuming from the hub is not supported for now."
                 )
-
-            policy_dir = Path(config_path).parent
-            if self.policy is not None:
-                self.policy.pretrained_path = policy_dir
-            self.checkpoint_path = policy_dir.parent
-
-        if self.policy is None:
-            raise ValueError(
-                "Policy is not configured. Please specify a pretrained policy with `--policy.path`."
-            )
+            policy_path = Path(config_path).parent
+            self.policy.pretrained_path = policy_path
+            self.checkpoint_path = policy_path.parent
 
         if not self.job_name:
             if self.env is None:
@@ -135,8 +126,8 @@ class TrainPipelineConfig(HubMixin):
         """This enables the parser to load config from the policy using `--policy.path=local/dir`"""
         return ["policy"]
 
-    def to_dict(self) -> dict[str, Any]:
-        return draccus.encode(self)  # type: ignore[no-any-return]  # because of the third-party library draccus uses Any as the return type
+    def to_dict(self) -> dict:
+        return draccus.encode(self)
 
     def _save_pretrained(self, save_directory: Path) -> None:
         with open(save_directory / TRAIN_CONFIG_NAME, "w") as f, draccus.config_type("json"):
@@ -148,13 +139,13 @@ class TrainPipelineConfig(HubMixin):
         pretrained_name_or_path: str | Path,
         *,
         force_download: bool = False,
-        resume_download: bool | None = None,
-        proxies: dict[Any, Any] | None = None,
+        resume_download: bool = None,
+        proxies: dict | None = None,
         token: str | bool | None = None,
         cache_dir: str | Path | None = None,
         local_files_only: bool = False,
         revision: str | None = None,
-        **kwargs: Any,
+        **kwargs,
     ) -> "TrainPipelineConfig":
         model_id = str(pretrained_name_or_path)
         config_file: str | None = None
@@ -190,6 +181,4 @@ class TrainPipelineConfig(HubMixin):
 
 @dataclass(kw_only=True)
 class TrainRLServerPipelineConfig(TrainPipelineConfig):
-    # NOTE: In RL, we don't need an offline dataset
-    # TODO: Make `TrainPipelineConfig.dataset` optional
-    dataset: DatasetConfig | None = None  # type: ignore[assignment] # because the parent class has made it's type non-optional
+    dataset: DatasetConfig | None = None  # NOTE: In RL, we don't need an offline dataset

src/lerobot/configs/types.py

@@ -42,11 +42,4 @@ class NormalizationMode(str, Enum):
 @dataclass
 class PolicyFeature:
     type: FeatureType
-    shape: tuple[int, ...]
-
-
-class RTCAttentionSchedule(str, Enum):
-    ZEROS = "ZEROS"
-    ONES = "ONES"
-    LINEAR = "LINEAR"
-    EXP = "EXP"
+    shape: tuple

src/lerobot/datasets/dataset_tools.py

@@ -28,10 +28,8 @@ import shutil
 from collections.abc import Callable
 from pathlib import Path
 
-import datasets
 import numpy as np
 import pandas as pd
-import pyarrow.parquet as pq
 import torch
 from tqdm import tqdm
 
@@ -39,13 +37,13 @@ from lerobot.datasets.aggregate import aggregate_datasets
 from lerobot.datasets.compute_stats import aggregate_stats
 from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
 from lerobot.datasets.utils import (
-    DATA_DIR,
     DEFAULT_CHUNK_SIZE,
     DEFAULT_DATA_FILE_SIZE_IN_MB,
     DEFAULT_DATA_PATH,
     DEFAULT_EPISODES_PATH,
     get_parquet_file_size_in_mb,
     load_episodes,
+    to_parquet_with_hf_images,
     update_chunk_file_indices,
     write_info,
     write_stats,
@@ -270,79 +268,39 @@ def merge_datasets(
     return merged_dataset
 
 
-def modify_features(
+def add_feature(
     dataset: LeRobotDataset,
-    add_features: dict[str, tuple[np.ndarray | torch.Tensor | Callable, dict]] | None = None,
-    remove_features: str | list[str] | None = None,
+    feature_name: str,
+    feature_values: np.ndarray | torch.Tensor | Callable,
+    feature_info: dict,
     output_dir: str | Path | None = None,
     repo_id: str | None = None,
 ) -> LeRobotDataset:
-    """Modify a LeRobotDataset by adding and/or removing features in a single pass.
-
-    This is the most efficient way to modify features, as it only copies the dataset once
-    regardless of how many features are being added or removed.
+    """Add a new feature to a LeRobotDataset.
 
     Args:
         dataset: The source LeRobotDataset.
-        add_features: Optional dict mapping feature names to (feature_values, feature_info) tuples.
-        remove_features: Optional feature name(s) to remove. Can be a single string or list.
+        feature_name: Name of the new feature.
+        feature_values: Either:
+            - Array/tensor of shape (num_frames, ...) with values for each frame
+            - Callable that takes (frame_dict, episode_index, frame_index) and returns feature value
+        feature_info: Dictionary with feature metadata (dtype, shape, names).
         output_dir: Directory to save the new dataset. If None, uses default location.
         repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
-
-    Returns:
-        New dataset with features modified.
-
-    Example:
-        new_dataset = modify_features(
-            dataset,
-            add_features={
-                "reward": (reward_array, {"dtype": "float32", "shape": [1], "names": None}),
-            },
-            remove_features=["old_feature"],
-            output_dir="./output",
-        )
     """
-    if add_features is None and remove_features is None:
-        raise ValueError("Must specify at least one of add_features or remove_features")
-
-    remove_features_list: list[str] = []
-    if remove_features is not None:
-        remove_features_list = [remove_features] if isinstance(remove_features, str) else remove_features
-
-    if add_features:
-        required_keys = {"dtype", "shape"}
-        for feature_name, (_, feature_info) in add_features.items():
-            if feature_name in dataset.meta.features:
-                raise ValueError(f"Feature '{feature_name}' already exists in dataset")
-
-            if not required_keys.issubset(feature_info.keys()):
-                raise ValueError(f"feature_info for '{feature_name}' must contain keys: {required_keys}")
-
-    if remove_features_list:
-        for name in remove_features_list:
-            if name not in dataset.meta.features:
-                raise ValueError(f"Feature '{name}' not found in dataset")
-
-        required_features = {"timestamp", "frame_index", "episode_index", "index", "task_index"}
-        if any(name in required_features for name in remove_features_list):
-            raise ValueError(f"Cannot remove required features: {required_features}")
+    if feature_name in dataset.meta.features:
+        raise ValueError(f"Feature '{feature_name}' already exists in dataset")
 
     if repo_id is None:
         repo_id = f"{dataset.repo_id}_modified"
     output_dir = Path(output_dir) if output_dir is not None else HF_LEROBOT_HOME / repo_id
 
+    required_keys = {"dtype", "shape"}
+    if not required_keys.issubset(feature_info.keys()):
+        raise ValueError(f"feature_info must contain keys: {required_keys}")
+
     new_features = dataset.meta.features.copy()
-
-    if remove_features_list:
-        for name in remove_features_list:
-            new_features.pop(name, None)
-
-    if add_features:
-        for feature_name, (_, feature_info) in add_features.items():
-            new_features[feature_name] = feature_info
-
-    video_keys_to_remove = [name for name in remove_features_list if name in dataset.meta.video_keys]
-    remaining_video_keys = [k for k in dataset.meta.video_keys if k not in video_keys_to_remove]
+    new_features[feature_name] = feature_info
 
     new_meta = LeRobotDatasetMetadata.create(
         repo_id=repo_id,
@@ -350,18 +308,17 @@ def modify_features(
         features=new_features,
         robot_type=dataset.meta.robot_type,
         root=output_dir,
-        use_videos=len(remaining_video_keys) > 0,
+        use_videos=len(dataset.meta.video_keys) > 0,
     )
 
     _copy_data_with_feature_changes(
         dataset=dataset,
         new_meta=new_meta,
-        add_features=add_features,
-        remove_features=remove_features_list if remove_features_list else None,
+        add_features={feature_name: (feature_values, feature_info)},
     )
 
-    if new_meta.video_keys:
-        _copy_videos(dataset, new_meta, exclude_keys=video_keys_to_remove if video_keys_to_remove else None)
+    if dataset.meta.video_keys:
+        _copy_videos(dataset, new_meta)
 
     new_dataset = LeRobotDataset(
         repo_id=repo_id,
@@ -374,46 +331,6 @@ def modify_features(
     return new_dataset
 
 
-def add_features(
-    dataset: LeRobotDataset,
-    features: dict[str, tuple[np.ndarray | torch.Tensor | Callable, dict]],
-    output_dir: str | Path | None = None,
-    repo_id: str | None = None,
-) -> LeRobotDataset:
-    """Add multiple features to a LeRobotDataset in a single pass.
-
-    This is more efficient than calling add_feature() multiple times, as it only
-    copies the dataset once regardless of how many features are being added.
-
-    Args:
-        dataset: The source LeRobotDataset.
-        features: Dictionary mapping feature names to (feature_values, feature_info) tuples.
-        output_dir: Directory to save the new dataset. If None, uses default location.
-        repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
-
-    Returns:
-        New dataset with all features added.
-
-    Example:
-        features = {
-            "task_embedding": (task_emb_array, {"dtype": "float32", "shape": [384], "names": None}),
-            "cam1_embedding": (cam1_emb_array, {"dtype": "float32", "shape": [768], "names": None}),
-            "cam2_embedding": (cam2_emb_array, {"dtype": "float32", "shape": [768], "names": None}),
-        }
-        new_dataset = add_features(dataset, features, output_dir="./output", repo_id="my_dataset")
-    """
-    if not features:
-        raise ValueError("No features provided")
-
-    return modify_features(
-        dataset=dataset,
-        add_features=features,
-        remove_features=None,
-        output_dir=output_dir,
-        repo_id=repo_id,
-    )
-
-
 def remove_feature(
     dataset: LeRobotDataset,
     feature_names: str | list[str],
@@ -428,17 +345,56 @@ def remove_feature(
         output_dir: Directory to save the new dataset. If None, uses default location.
         repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
 
-    Returns:
-        New dataset with features removed.
     """
-    return modify_features(
-        dataset=dataset,
-        add_features=None,
-        remove_features=feature_names,
-        output_dir=output_dir,
+    if isinstance(feature_names, str):
+        feature_names = [feature_names]
+
+    for name in feature_names:
+        if name not in dataset.meta.features:
+            raise ValueError(f"Feature '{name}' not found in dataset")
+
+    required_features = {"timestamp", "frame_index", "episode_index", "index", "task_index"}
+    if any(name in required_features for name in feature_names):
+        raise ValueError(f"Cannot remove required features: {required_features}")
+
+    if repo_id is None:
+        repo_id = f"{dataset.repo_id}_modified"
+    output_dir = Path(output_dir) if output_dir is not None else HF_LEROBOT_HOME / repo_id
+
+    new_features = {k: v for k, v in dataset.meta.features.items() if k not in feature_names}
+
+    video_keys_to_remove = [name for name in feature_names if name in dataset.meta.video_keys]
+
+    remaining_video_keys = [k for k in dataset.meta.video_keys if k not in video_keys_to_remove]
+
+    new_meta = LeRobotDatasetMetadata.create(
         repo_id=repo_id,
+        fps=dataset.meta.fps,
+        features=new_features,
+        robot_type=dataset.meta.robot_type,
+        root=output_dir,
+        use_videos=len(remaining_video_keys) > 0,
     )
 
+    _copy_data_with_feature_changes(
+        dataset=dataset,
+        new_meta=new_meta,
+        remove_features=feature_names,
+    )
+
+    if new_meta.video_keys:
+        _copy_videos(dataset, new_meta, exclude_keys=video_keys_to_remove)
+
+    new_dataset = LeRobotDataset(
+        repo_id=repo_id,
+        root=output_dir,
+        image_transforms=dataset.image_transforms,
+        delta_timestamps=dataset.delta_timestamps,
+        tolerance_s=dataset.tolerance_s,
+    )
+
+    return new_dataset
+
 
 def _fractions_to_episode_indices(
     total_episodes: int,
@@ -545,7 +501,10 @@ def _copy_and_reindex_data(
         dst_path = dst_meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
         dst_path.parent.mkdir(parents=True, exist_ok=True)
 
-        _write_parquet(df, dst_path, dst_meta)
+        if len(dst_meta.image_keys) > 0:
+            to_parquet_with_hf_images(df, dst_path)
+        else:
+            df.to_parquet(dst_path, index=False)
 
         for ep_old_idx in episodes_to_keep:
             ep_new_idx = episode_mapping[ep_old_idx]
@@ -903,25 +862,6 @@ def _copy_and_reindex_episodes_metadata(
     write_stats(filtered_stats, dst_meta.root)
 
 
-def _write_parquet(df: pd.DataFrame, path: Path, meta: LeRobotDatasetMetadata) -> None:
-    """Write DataFrame to parquet
-
-    This ensures images are properly embedded and the file can be loaded correctly by HF datasets.
-    """
-    from lerobot.datasets.utils import embed_images, get_hf_features_from_features
-
-    hf_features = get_hf_features_from_features(meta.features)
-    ep_dataset = datasets.Dataset.from_dict(df.to_dict(orient="list"), features=hf_features, split="train")
-
-    if len(meta.image_keys) > 0:
-        ep_dataset = embed_images(ep_dataset)
-
-    table = ep_dataset.with_format("arrow")[:]
-    writer = pq.ParquetWriter(path, schema=table.schema, compression="snappy", use_dictionary=True)
-    writer.write_table(table)
-    writer.close()
-
-
 def _save_data_chunk(
     df: pd.DataFrame,
     meta: LeRobotDatasetMetadata,
@@ -937,7 +877,10 @@ def _save_data_chunk(
     path = meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
     path.parent.mkdir(parents=True, exist_ok=True)
 
-    _write_parquet(df, path, meta)
+    if len(meta.image_keys) > 0:
+        to_parquet_with_hf_images(df, path)
+    else:
+        df.to_parquet(path, index=False)
 
     episode_metadata = {}
     for ep_idx in df["episode_index"].unique():
@@ -963,29 +906,19 @@ def _copy_data_with_feature_changes(
     remove_features: list[str] | None = None,
 ) -> None:
     """Copy data while adding or removing features."""
-    data_dir = dataset.root / DATA_DIR
-    parquet_files = sorted(data_dir.glob("*/*.parquet"))
-
-    if not parquet_files:
-        raise ValueError(f"No parquet files found in {data_dir}")
+    file_paths = set()
+    for ep_idx in range(dataset.meta.total_episodes):
+        file_paths.add(dataset.meta.get_data_file_path(ep_idx))
 
     frame_idx = 0
 
-    for src_path in tqdm(parquet_files, desc="Processing data files"):
-        df = pd.read_parquet(src_path).reset_index(drop=True)
-
-        relative_path = src_path.relative_to(dataset.root)
-        chunk_dir = relative_path.parts[1]
-        file_name = relative_path.parts[2]
-
-        chunk_idx = int(chunk_dir.split("-")[1])
-        file_idx = int(file_name.split("-")[1].split(".")[0])
+    for src_path in tqdm(sorted(file_paths), desc="Processing data files"):
+        df = pd.read_parquet(dataset.root / src_path).reset_index(drop=True)
 
         if remove_features:
             df = df.drop(columns=remove_features, errors="ignore")
 
         if add_features:
-            end_idx = frame_idx + len(df)
             for feature_name, (values, _) in add_features.items():
                 if callable(values):
                     feature_values = []
@@ -998,18 +931,15 @@ def _copy_data_with_feature_changes(
                         feature_values.append(value)
                     df[feature_name] = feature_values
                 else:
+                    end_idx = frame_idx + len(df)
                     feature_slice = values[frame_idx:end_idx]
                     if len(feature_slice.shape) > 1 and feature_slice.shape[1] == 1:
                         df[feature_name] = feature_slice.flatten()
                     else:
                         df[feature_name] = feature_slice
-            frame_idx = end_idx
+                    frame_idx = end_idx
 
-        # Write using the same chunk/file structure as source
-        dst_path = new_meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
-        dst_path.parent.mkdir(parents=True, exist_ok=True)
-
-        _write_parquet(df, dst_path, new_meta)
+        _save_data_chunk(df, new_meta)
 
     _copy_episodes_metadata_and_stats(dataset, new_meta)
 

src/lerobot/datasets/lerobot_dataset.py

@@ -430,7 +430,9 @@ class LeRobotDatasetMetadata:
         video_keys = [video_key] if video_key is not None else self.video_keys
         for key in video_keys:
             if not self.features[key].get("info", None):
-                video_path = self.root / self.video_path.format(video_key=key, chunk_index=0, file_index=0)
+                video_path = self.root / self.video_path.format(
+                    video_key=video_key, chunk_index=0, file_index=0
+                )
                 self.info["features"][key]["info"] = get_video_info(video_path)
 
     def update_chunk_settings(
@@ -684,7 +686,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         self.episode_buffer = None
         self.writer = None
         self.latest_episode = None
-        self._current_file_start_frame = None  # Track the starting frame index of the current parquet file
 
         self.root.mkdir(exist_ok=True, parents=True)
 
@@ -707,8 +708,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
             if not self._check_cached_episodes_sufficient():
                 raise FileNotFoundError("Cached dataset doesn't contain all requested episodes")
         except (AssertionError, FileNotFoundError, NotADirectoryError):
-            if is_valid_version(self.revision):
-                self.revision = get_safe_version(self.repo_id, self.revision)
+            self.revision = get_safe_version(self.repo_id, self.revision)
             self.download(download_videos)
             self.hf_dataset = self.load_hf_dataset()
 
@@ -835,14 +835,14 @@ class LeRobotDataset(torch.utils.data.Dataset):
         return hf_dataset
 
     def _check_cached_episodes_sufficient(self) -> bool:
-        """Check if the cached dataset contains all requested episodes and their video files."""
+        """Check if the cached dataset contains all requested episodes."""
         if self.hf_dataset is None or len(self.hf_dataset) == 0:
             return False
 
         # Get available episode indices from cached dataset
         available_episodes = {
             ep_idx.item() if isinstance(ep_idx, torch.Tensor) else ep_idx
-            for ep_idx in self.hf_dataset.unique("episode_index")
+            for ep_idx in self.hf_dataset["episode_index"]
         }
 
         # Determine requested episodes
@@ -854,18 +854,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
             requested_episodes = set(self.episodes)
 
         # Check if all requested episodes are available in cached data
-        if not requested_episodes.issubset(available_episodes):
-            return False
-
-        # Check if all required video files exist
-        if len(self.meta.video_keys) > 0:
-            for ep_idx in requested_episodes:
-                for vid_key in self.meta.video_keys:
-                    video_path = self.root / self.meta.get_video_file_path(ep_idx, vid_key)
-                    if not video_path.exists():
-                        return False
-
-        return True
+        return requested_episodes.issubset(available_episodes)
 
     def create_hf_dataset(self) -> datasets.Dataset:
         features = get_hf_features_from_features(self.features)
@@ -940,26 +929,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
         return query_timestamps
 
     def _query_hf_dataset(self, query_indices: dict[str, list[int]]) -> dict:
-        """
-        Query dataset for indices across keys, skipping video keys.
-
-        Tries column-first [key][indices] for speed, falls back to row-first.
-
-        Args:
-            query_indices: Dict mapping keys to index lists to retrieve
-
-        Returns:
-            Dict with stacked tensors of queried data (video keys excluded)
-        """
-        result: dict = {}
-        for key, q_idx in query_indices.items():
-            if key in self.meta.video_keys:
-                continue
-            try:
-                result[key] = torch.stack(self.hf_dataset[key][q_idx])
-            except (KeyError, TypeError, IndexError):
-                result[key] = torch.stack(self.hf_dataset[q_idx][key])
-        return result
+        return {
+            key: torch.stack(self.hf_dataset[q_idx][key])
+            for key, q_idx in query_indices.items()
+            if key not in self.meta.video_keys
+        }
 
     def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict[str, torch.Tensor]:
         """Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
@@ -1257,7 +1231,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
             # Initialize indices and frame count for a new dataset made of the first episode data
             chunk_idx, file_idx = 0, 0
             global_frame_index = 0
-            self._current_file_start_frame = 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
@@ -1269,7 +1242,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
                 # When resuming, move to the next file
                 chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, self.meta.chunks_size)
-                self._current_file_start_frame = global_frame_index
         else:
             # Retrieve information from the latest parquet file
             latest_ep = self.latest_episode
@@ -1280,7 +1252,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
             latest_path = self.root / self.meta.data_path.format(chunk_index=chunk_idx, file_index=file_idx)
             latest_size_in_mb = get_file_size_in_mb(latest_path)
 
-            frames_in_current_file = global_frame_index - self._current_file_start_frame
+            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
             )
@@ -1294,7 +1266,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
                 chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, self.meta.chunks_size)
                 self._close_writer()
                 self._writer_closed_for_reading = False
-                self._current_file_start_frame = global_frame_index
 
         ep_dict["data/chunk_index"] = chunk_idx
         ep_dict["data/file_index"] = file_idx
@@ -1501,7 +1472,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec()
         obj.writer = None
         obj.latest_episode = None
-        obj._current_file_start_frame = None
         # Initialize tracking for incremental recording
         obj._lazy_loading = False
         obj._recorded_frames = 0

src/lerobot/datasets/transforms.py

@@ -206,11 +206,6 @@ class ImageTransformsConfig:
                 type="SharpnessJitter",
                 kwargs={"sharpness": (0.5, 1.5)},
             ),
-            "affine": ImageTransformConfig(
-                weight=1.0,
-                type="RandomAffine",
-                kwargs={"degrees": (-5.0, 5.0), "translate": (0.05, 0.05)},
-            ),
         }
     )
 
@@ -222,8 +217,6 @@ def make_transform_from_config(cfg: ImageTransformConfig):
         return v2.ColorJitter(**cfg.kwargs)
     elif cfg.type == "SharpnessJitter":
         return SharpnessJitter(**cfg.kwargs)
-    elif cfg.type == "RandomAffine":
-        return v2.RandomAffine(**cfg.kwargs)
     else:
         raise ValueError(f"Transform '{cfg.type}' is not valid.")
 

src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py

@@ -69,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,
@@ -98,7 +98,7 @@ OLD
 videos/chunk-000/CAMERA/episode_000000.mp4
 
 NEW
-videos/CAMERA/chunk-000/file_000.mp4
+videos/chunk-000/file_000.mp4
 -------------------------
 OLD
 episodes.jsonl
@@ -310,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_file_size_in_mb(ep_path)
+        ep_size_in_mb = get_video_size_in_mb(ep_path)
         ep_duration_in_s = get_video_duration_in_s(ep_path)
 
         # Check if adding this episode would exceed the limit

src/lerobot/datasets/video_utils.py

@@ -342,8 +342,8 @@ def encode_video_frames(
     # Define video output frame size (assuming all input frames are the same size)
     if len(input_list) == 0:
         raise FileNotFoundError(f"No images found in {imgs_dir}.")
-    with Image.open(input_list[0]) as dummy_image:
-        width, height = dummy_image.size
+    dummy_image = Image.open(input_list[0])
+    width, height = dummy_image.size
 
     # Define video codec options
     video_options = {}
@@ -373,12 +373,11 @@ def encode_video_frames(
 
         # Loop through input frames and encode them
         for input_data in input_list:
-            with Image.open(input_data) as input_image:
-                input_image = input_image.convert("RGB")
-                input_frame = av.VideoFrame.from_image(input_image)
-                packet = output_stream.encode(input_frame)
-                if packet:
-                    output.mux(packet)
+            input_image = Image.open(input_data).convert("RGB")
+            input_frame = av.VideoFrame.from_image(input_image)
+            packet = output_stream.encode(input_frame)
+            if packet:
+                output.mux(packet)
 
         # Flush the encoder
         packet = output_stream.encode()

src/lerobot/envs/__init__.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configs import AlohaEnv, EnvConfig, PushtEnv  # noqa: F401
+from .configs import AlohaEnv, EnvConfig, PushtEnv, XarmEnv  # noqa: F401

src/lerobot/envs/configs.py

@@ -37,16 +37,6 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC):
     def type(self) -> str:
         return self.get_choice_name(self.__class__)
 
-    @property
-    def package_name(self) -> str:
-        """Package name to import if environment not found in gym registry"""
-        return f"gym_{self.type}"
-
-    @property
-    def gym_id(self) -> str:
-        """ID string used in gym.make() to instantiate the environment"""
-        return f"{self.package_name}/{self.task}"
-
     @property
     @abc.abstractmethod
     def gym_kwargs(self) -> dict:
@@ -143,6 +133,45 @@ 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
@@ -277,45 +306,3 @@ 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,
-        }

src/lerobot/envs/factory.py

@@ -16,10 +16,8 @@
 import importlib
 
 import gymnasium as gym
-from gymnasium.envs.registration import registry as gym_registry
 
-from lerobot.envs.configs import AlohaEnv, EnvConfig, LiberoEnv, PushtEnv
-from lerobot.envs.utils import _call_make_env, _download_hub_file, _import_hub_module, _normalize_hub_result
+from lerobot.envs.configs import AlohaEnv, EnvConfig, LiberoEnv, PushtEnv, XarmEnv
 
 
 def make_env_config(env_type: str, **kwargs) -> EnvConfig:
@@ -27,6 +25,8 @@ 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:
@@ -34,24 +34,15 @@ def make_env_config(env_type: str, **kwargs) -> EnvConfig:
 
 
 def make_env(
-    cfg: EnvConfig | str,
-    n_envs: int = 1,
-    use_async_envs: bool = False,
-    hub_cache_dir: str | None = None,
-    trust_remote_code: bool = False,
+    cfg: EnvConfig, n_envs: int = 1, use_async_envs: bool = False
 ) -> dict[str, dict[int, gym.vector.VectorEnv]]:
-    """Makes a gym vector environment according to the config or Hub reference.
+    """Makes a gym vector environment according to the config.
 
     Args:
-        cfg (EnvConfig | str): Either an `EnvConfig` object describing the environment to build locally,
-            or a Hugging Face Hub repository identifier (e.g. `"username/repo"`). In the latter case,
-            the repo must include a Python file (usually `env.py`).
+        cfg (EnvConfig): the config of the environment to instantiate.
         n_envs (int, optional): The number of parallelized env to return. Defaults to 1.
         use_async_envs (bool, optional): Whether to return an AsyncVectorEnv or a SyncVectorEnv. Defaults to
             False.
-        hub_cache_dir (str | None): Optional cache path for downloaded hub files.
-        trust_remote_code (bool): **Explicit consent** to execute remote code from the Hub.
-            Default False — must be set to True to import/exec hub `env.py`.
 
     Raises:
         ValueError: if n_envs < 1
@@ -64,21 +55,6 @@ def make_env(
             - For single-task environments: a single suite entry (cfg.type) with task_id=0.
 
     """
-    # if user passed a hub id string (e.g., "username/repo", "username/repo@main:env.py")
-    # simplified: only support hub-provided `make_env`
-    if isinstance(cfg, str):
-        # _download_hub_file will raise the same RuntimeError if trust_remote_code is False
-        repo_id, file_path, local_file, revision = _download_hub_file(cfg, trust_remote_code, hub_cache_dir)
-
-        # import and surface clear import errors
-        module = _import_hub_module(local_file, repo_id)
-
-        # call the hub-provided make_env
-        raw_result = _call_make_env(module, n_envs=n_envs, use_async_envs=use_async_envs)
-
-        # normalize the return into {suite: {task_id: vec_env}}
-        return _normalize_hub_result(raw_result)
-
     if n_envs < 1:
         raise ValueError("`n_envs` must be at least 1")
 
@@ -98,38 +74,20 @@ 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")
+    package_name = f"gym_{cfg.type}"
+    try:
+        importlib.import_module(package_name)
+    except ModuleNotFoundError as e:
+        print(f"{package_name} is not installed. Please install it with `pip install 'lerobot[{cfg.type}]'`")
+        raise e
 
-        return create_metaworld_envs(
-            task=cfg.task,
-            n_envs=n_envs,
-            gym_kwargs=cfg.gym_kwargs,
-            env_cls=env_cls,
-        )
-
-    if cfg.gym_id not in gym_registry:
-        print(f"gym id '{cfg.gym_id}' not found, attempting to import '{cfg.package_name}'...")
-        try:
-            importlib.import_module(cfg.package_name)
-        except ModuleNotFoundError as e:
-            raise ModuleNotFoundError(
-                f"Package '{cfg.package_name}' required for env '{cfg.type}' not found. "
-                f"Please install it or check PYTHONPATH."
-            ) from e
-
-        if cfg.gym_id not in gym_registry:
-            raise gym.error.NameNotFound(
-                f"Environment '{cfg.gym_id}' not registered even after importing '{cfg.package_name}'."
-            )
+    gym_handle = f"{package_name}/{cfg.task}"
 
     def _make_one():
-        return gym.make(cfg.gym_id, disable_env_checker=cfg.disable_env_checker, **(cfg.gym_kwargs or {}))
+        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)], autoreset_mode=gym.vector.AutoresetMode.SAME_STEP)
+    vec = env_cls([_make_one for _ in range(n_envs)])
 
     # normalize to {suite: {task_id: vec_env}} for consistency
     suite_name = cfg.type  # e.g., "pusht", "aloha"

src/lerobot/envs/libero.py

@@ -260,23 +260,19 @@ class LiberoEnv(gym.Env):
 
         is_success = self._env.check_success()
         terminated = done or is_success
-        info.update(
-            {
-                "task": self.task,
-                "task_id": self.task_id,
-                "done": done,
-                "is_success": is_success,
-            }
-        )
+        info["is_success"] = is_success
+
         observation = self._format_raw_obs(raw_obs)
-        if terminated:
-            info["final_info"] = {
-                "task": self.task,
-                "task_id": self.task_id,
-                "done": bool(done),
-                "is_success": bool(is_success),
-            }
+        if done:
             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
 

src/lerobot/envs/metaworld.py

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

src/lerobot/envs/metaworld_config.json

@@ -1,121 +0,0 @@
-{
-  "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"
-  }
-}

src/lerobot/envs/utils.py

@@ -13,8 +13,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import importlib.util
-import os
 import warnings
 from collections.abc import Mapping, Sequence
 from functools import singledispatch
@@ -24,7 +22,6 @@ import einops
 import gymnasium as gym
 import numpy as np
 import torch
-from huggingface_hub import hf_hub_download, snapshot_download
 from torch import Tensor
 
 from lerobot.configs.types import FeatureType, PolicyFeature
@@ -198,132 +195,3 @@ def _(envs: Sequence) -> None:
 @close_envs.register
 def _(env: gym.Env) -> None:
     _close_single_env(env)
-
-
-# helper to safely load a python file as a module
-def _load_module_from_path(path: str, module_name: str | None = None):
-    module_name = module_name or f"hub_env_{os.path.basename(path).replace('.', '_')}"
-    spec = importlib.util.spec_from_file_location(module_name, path)
-    if spec is None:
-        raise ImportError(f"Could not load module spec for {module_name} from {path}")
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)  # type: ignore
-    return module
-
-
-# helper to parse hub string (supports "user/repo", "user/repo@rev", optional path)
-# examples:
-#   "user/repo" -> will look for env.py at repo root
-#   "user/repo@main:envs/my_env.py" -> explicit revision and path
-def _parse_hub_url(hub_uri: str):
-    # very small parser: [repo_id][@revision][:path]
-    # repo_id is required (user/repo or org/repo)
-    revision = None
-    file_path = "env.py"
-    if "@" in hub_uri:
-        repo_and_rev, *rest = hub_uri.split(":", 1)
-        repo_id, rev = repo_and_rev.split("@", 1)
-        revision = rev
-        if rest:
-            file_path = rest[0]
-    else:
-        repo_id, *rest = hub_uri.split(":", 1)
-        if rest:
-            file_path = rest[0]
-    return repo_id, revision, file_path
-
-
-def _download_hub_file(
-    cfg_str: str,
-    trust_remote_code: bool,
-    hub_cache_dir: str | None,
-) -> tuple[str, str, str, str]:
-    """
-    Parse `cfg_str` (hub URL), enforce `trust_remote_code`, and return
-    (repo_id, file_path, local_file, revision).
-    """
-    if not trust_remote_code:
-        raise RuntimeError(
-            f"Refusing to execute remote code from the Hub for '{cfg_str}'. "
-            "Executing hub env modules runs arbitrary Python code from third-party repositories. "
-            "If you trust this repo and understand the risks, call `make_env(..., trust_remote_code=True)` "
-            "and prefer pinning to a specific revision: 'user/repo@<commit-hash>:env.py'."
-        )
-
-    repo_id, revision, file_path = _parse_hub_url(cfg_str)
-
-    try:
-        local_file = hf_hub_download(
-            repo_id=repo_id, filename=file_path, revision=revision, cache_dir=hub_cache_dir
-        )
-    except Exception as e:
-        # fallback to snapshot download
-        snapshot_dir = snapshot_download(repo_id=repo_id, revision=revision, cache_dir=hub_cache_dir)
-        local_file = os.path.join(snapshot_dir, file_path)
-        if not os.path.exists(local_file):
-            raise FileNotFoundError(
-                f"Could not find {file_path} in repository {repo_id}@{revision or 'main'}"
-            ) from e
-
-    return repo_id, file_path, local_file, revision
-
-
-def _import_hub_module(local_file: str, repo_id: str) -> Any:
-    """
-    Import the downloaded file as a module and surface helpful import error messages.
-    """
-    module_name = f"hub_env_{repo_id.replace('/', '_')}"
-    try:
-        module = _load_module_from_path(local_file, module_name=module_name)
-    except ModuleNotFoundError as e:
-        missing = getattr(e, "name", None) or str(e)
-        raise ModuleNotFoundError(
-            f"Hub env '{repo_id}:{os.path.basename(local_file)}' failed to import because the dependency "
-            f"'{missing}' is not installed locally.\n\n"
-        ) from e
-    except ImportError as e:
-        raise ImportError(
-            f"Failed to load hub env module '{repo_id}:{os.path.basename(local_file)}'. Import error: {e}\n\n"
-        ) from e
-    return module
-
-
-def _call_make_env(module: Any, n_envs: int, use_async_envs: bool) -> Any:
-    """
-    Ensure module exposes make_env and call it.
-    """
-    if not hasattr(module, "make_env"):
-        raise AttributeError(
-            f"The hub module {getattr(module, '__name__', 'hub_module')} must expose `make_env(n_envs=int, use_async_envs=bool)`."
-        )
-    entry_fn = module.make_env
-    return entry_fn(n_envs=n_envs, use_async_envs=use_async_envs)
-
-
-def _normalize_hub_result(result: Any) -> dict[str, dict[int, gym.vector.VectorEnv]]:
-    """
-    Normalize possible return types from hub `make_env` into the mapping:
-      { suite_name: { task_id: vector_env } }
-    Accepts:
-      - dict (assumed already correct)
-      - gym.vector.VectorEnv
-      - gym.Env (will be wrapped into SyncVectorEnv)
-    """
-    if isinstance(result, dict):
-        return result
-
-    # VectorEnv: use its spec.id if available
-    if isinstance(result, gym.vector.VectorEnv):
-        suite_name = getattr(result, "spec", None) and getattr(result.spec, "id", None) or "hub_env"
-        return {suite_name: {0: result}}
-
-    # Single Env: wrap into SyncVectorEnv
-    if isinstance(result, gym.Env):
-        vec = gym.vector.SyncVectorEnv([lambda: result])
-        suite_name = getattr(result, "spec", None) and getattr(result.spec, "id", None) or "hub_env"
-        return {suite_name: {0: vec}}
-
-    raise ValueError(
-        "Hub `make_env` must return either a mapping {suite: {task_id: vec_env}}, "
-        "a gym.vector.VectorEnv, or a single gym.Env."
-    )

src/lerobot/model/kinematics.py

@@ -22,18 +22,18 @@ class RobotKinematics:
         self,
         urdf_path: str,
         target_frame_name: str = "gripper_frame_link",
-        joint_names: list[str] | None = None,
+        joint_names: list[str] = None,
     ):
         """
         Initialize placo-based kinematics solver.
 
         Args:
-            urdf_path (str): Path to the robot URDF file
-            target_frame_name (str): Name of the end-effector frame in the URDF
-            joint_names (list[str] | None): List of joint names to use for the kinematics solver
+            urdf_path: Path to the robot URDF file
+            target_frame_name: Name of the end-effector frame in the URDF
+            joint_names: List of joint names to use for the kinematics solver
         """
         try:
-            import placo  # type: ignore[import-not-found] # C++ library with Python bindings, no type stubs available. TODO: Create stub file or request upstream typing support.
+            import placo
         except ImportError as e:
             raise ImportError(
                 "placo is required for RobotKinematics. "
@@ -52,7 +52,7 @@ class RobotKinematics:
         # Initialize frame task for IK
         self.tip_frame = self.solver.add_frame_task(self.target_frame_name, np.eye(4))
 
-    def forward_kinematics(self, joint_pos_deg: np.ndarray) -> np.ndarray:
+    def forward_kinematics(self, joint_pos_deg):
         """
         Compute forward kinematics for given joint configuration given the target frame name in the constructor.
 
@@ -77,12 +77,8 @@ class RobotKinematics:
         return self.robot.get_T_world_frame(self.target_frame_name)
 
     def inverse_kinematics(
-        self,
-        current_joint_pos: np.ndarray,
-        desired_ee_pose: np.ndarray,
-        position_weight: float = 1.0,
-        orientation_weight: float = 0.01,
-    ) -> np.ndarray:
+        self, current_joint_pos, desired_ee_pose, position_weight=1.0, orientation_weight=0.01
+    ):
         """
         Compute inverse kinematics using placo solver.
 

src/lerobot/motors/dynamixel/dynamixel.py

@@ -60,7 +60,7 @@ class OperatingMode(Enum):
 
     # This mode controls position. This mode is identical to the Multi-turn Position Control from existing
     # DYNAMIXEL. 512 turns are supported(-256[rev] ~ 256[rev]). This mode is ideal for multi-turn wrists or
-    # conveyor systems or a system that requires an additional reduction gear. Note that Max Position
+    # conveyer systems or a system that requires an additional reduction gear. Note that Max Position
     # Limit(48), Min Position Limit(52) are not used on Extended Position Control Mode.
     EXTENDED_POSITION = 4
 

src/lerobot/motors/feetech/tables.py

@@ -206,12 +206,8 @@ MODEL_BAUDRATE_TABLE = {
 # Sign-Magnitude encoding bits
 STS_SMS_SERIES_ENCODINGS_TABLE = {
     "Homing_Offset": 11,
-    "Goal_Position": 15,
     "Goal_Velocity": 15,
-    "Goal_Speed": 15,
-    "Present_Position": 15,
     "Present_Velocity": 15,
-    "Present_Speed": 15,
 }
 
 MODEL_ENCODING_TABLE = {

src/lerobot/optim/schedulers.py

@@ -14,7 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import abc
-import logging
 import math
 from dataclasses import asdict, dataclass
 from pathlib import Path
@@ -80,11 +79,7 @@ class VQBeTSchedulerConfig(LRSchedulerConfig):
 @LRSchedulerConfig.register_subclass("cosine_decay_with_warmup")
 @dataclass
 class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig):
-    """Used by Physical Intelligence to train Pi0.
-
-    Automatically scales warmup and decay steps if num_training_steps < num_decay_steps.
-    This ensures the learning rate schedule completes properly even with shorter training runs.
-    """
+    """Used by Physical Intelligence to train Pi0"""
 
     num_warmup_steps: int
     num_decay_steps: int
@@ -92,39 +87,23 @@ class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig):
     decay_lr: float
 
     def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR:
-        # Auto-scale scheduler parameters if training steps are shorter than configured decay steps
-        actual_warmup_steps = self.num_warmup_steps
-        actual_decay_steps = self.num_decay_steps
-
-        if num_training_steps < self.num_decay_steps:
-            # Calculate scaling factor to fit the schedule into the available training steps
-            scale_factor = num_training_steps / self.num_decay_steps
-            actual_warmup_steps = int(self.num_warmup_steps * scale_factor)
-            actual_decay_steps = num_training_steps
-
-            logging.info(
-                f"Auto-scaling LR scheduler: "
-                f"num_training_steps ({num_training_steps}) < num_decay_steps ({self.num_decay_steps}). "
-                f"Scaling warmup: {self.num_warmup_steps} → {actual_warmup_steps}, "
-                f"decay: {self.num_decay_steps} → {actual_decay_steps} "
-                f"(scale factor: {scale_factor:.3f})"
-            )
+        del num_training_steps
 
         def lr_lambda(current_step):
             def linear_warmup_schedule(current_step):
                 if current_step <= 0:
-                    return 1 / (actual_warmup_steps + 1)
-                frac = 1 - current_step / actual_warmup_steps
-                return (1 / (actual_warmup_steps + 1) - 1) * frac + 1
+                    return 1 / (self.num_warmup_steps + 1)
+                frac = 1 - current_step / self.num_warmup_steps
+                return (1 / (self.num_warmup_steps + 1) - 1) * frac + 1
 
             def cosine_decay_schedule(current_step):
-                step = min(current_step, actual_decay_steps)
-                cosine_decay = 0.5 * (1 + math.cos(math.pi * step / actual_decay_steps))
+                step = min(current_step, self.num_decay_steps)
+                cosine_decay = 0.5 * (1 + math.cos(math.pi * step / self.num_decay_steps))
                 alpha = self.decay_lr / self.peak_lr
                 decayed = (1 - alpha) * cosine_decay + alpha
                 return decayed
 
-            if current_step < actual_warmup_steps:
+            if current_step < self.num_warmup_steps:
                 return linear_warmup_schedule(current_step)
 
             return cosine_decay_schedule(current_step)

src/lerobot/policies/__init__.py

@@ -14,7 +14,6 @@
 
 from .act.configuration_act import ACTConfig as ACTConfig
 from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
-from .groot.configuration_groot import GrootConfig as GrootConfig
 from .pi0.configuration_pi0 import PI0Config as PI0Config
 from .pi05.configuration_pi05 import PI05Config as PI05Config
 from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
@@ -30,5 +29,4 @@ __all__ = [
     "SmolVLAConfig",
     "TDMPCConfig",
     "VQBeTConfig",
-    "GrootConfig",
 ]

src/lerobot/policies/act/modeling_act.py

@@ -626,8 +626,8 @@ class ACTDecoderLayer(nn.Module):
             x: (Decoder Sequence, Batch, Channel) tensor of input tokens.
             encoder_out: (Encoder Sequence, B, C) output features from the last layer of the encoder we are
                 cross-attending with.
-            encoder_pos_embed: (ES, 1, C) positional embedding for keys (from the encoder).
-            decoder_pos_embed: (DS, 1, C) positional embedding for the queries (from the decoder).
+            decoder_pos_embed: (ES, 1, C) positional embedding for keys (from the encoder).
+            encoder_pos_embed: (DS, 1, C) Positional_embedding for the queries (from the decoder).
         Returns:
             (DS, B, C) tensor of decoder output features.
         """

src/lerobot/policies/diffusion/configuration_diffusion.py

@@ -45,7 +45,7 @@ class DiffusionConfig(PreTrainedConfig):
     Args:
         n_obs_steps: Number of environment steps worth of observations to pass to the policy (takes the
             current step and additional steps going back).
-        horizon: Diffusion model action prediction size as detailed in `DiffusionPolicy.select_action`.
+        chunk_size: Diffusion model action prediction size as detailed in `DiffusionPolicy.select_action`.
         n_action_steps: The number of action steps to run in the environment for one invocation of the policy.
             See `DiffusionPolicy.select_action` for more details.
         input_shapes: A dictionary defining the shapes of the input data for the policy. The key represents
@@ -105,7 +105,7 @@ class DiffusionConfig(PreTrainedConfig):
 
     # Inputs / output structure.
     n_obs_steps: int = 2
-    horizon: int = 16
+    chunk_size: int = 16
     n_action_steps: int = 8
 
     normalization_mapping: dict[str, NormalizationMode] = field(
@@ -118,7 +118,7 @@ class DiffusionConfig(PreTrainedConfig):
 
     # The original implementation doesn't sample frames for the last 7 steps,
     # which avoids excessive padding and leads to improved training results.
-    drop_n_last_frames: int = 7  # horizon - n_action_steps - n_obs_steps + 1
+    drop_n_last_frames: int = 7  # chunk_size - n_action_steps - n_obs_steps + 1
 
     # Architecture / modeling.
     # Vision backbone.
@@ -180,13 +180,13 @@ class DiffusionConfig(PreTrainedConfig):
                 f"Got {self.noise_scheduler_type}."
             )
 
-        # Check that the horizon size and U-Net downsampling is compatible.
+        # Check that the chunk size and U-Net downsampling is compatible.
         # U-Net downsamples by 2 with each stage.
         downsampling_factor = 2 ** len(self.down_dims)
-        if self.horizon % downsampling_factor != 0:
+        if self.chunk_size % downsampling_factor != 0:
             raise ValueError(
-                "The horizon should be an integer multiple of the downsampling factor (which is determined "
-                f"by `len(down_dims)`). Got {self.horizon=} and {self.down_dims=}"
+                "The chunk_size should be an integer multiple of the downsampling factor (which is determined "
+                f"by `len(down_dims)`). Got {self.chunk_size=} and {self.down_dims=}"
             )
 
     def get_optimizer_preset(self) -> AdamConfig:
@@ -231,7 +231,7 @@ class DiffusionConfig(PreTrainedConfig):
 
     @property
     def action_delta_indices(self) -> list:
-        return list(range(1 - self.n_obs_steps, 1 - self.n_obs_steps + self.horizon))
+        return list(range(1 - self.n_obs_steps, 1 - self.n_obs_steps + self.chunk_size))
 
     @property
     def reward_delta_indices(self) -> None:

src/lerobot/policies/diffusion/modeling_diffusion.py

@@ -99,25 +99,25 @@ class DiffusionPolicy(PreTrainedPolicy):
         return actions
 
     @torch.no_grad()
-    def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
+    def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs) -> Tensor:
         """Select a single action given environment observations.
 
         This method handles caching a history of observations and an action trajectory generated by the
         underlying diffusion model. Here's how it works:
           - `n_obs_steps` steps worth of observations are cached (for the first steps, the observation is
             copied `n_obs_steps` times to fill the cache).
-          - The diffusion model generates `horizon` steps worth of actions.
+          - The diffusion model generates `chunk_size` steps worth of actions.
           - `n_action_steps` worth of actions are actually kept for execution, starting from the current step.
         Schematically this looks like:
             ----------------------------------------------------------------------------------------------
-            (legend: o = n_obs_steps, h = horizon, a = n_action_steps)
+            (legend: o = n_obs_steps, c = chunk_size, a = n_action_steps)
             |timestep            | n-o+1 | n-o+2 | ..... | n     | ..... | n+a-1 | n+a   | ..... | n-o+h |
             |observation is used | YES   | YES   | YES   | YES   | NO    | NO    | NO    | NO    | NO    |
             |action is generated | YES   | YES   | YES   | YES   | YES   | YES   | YES   | YES   | YES   |
             |action is used      | NO    | NO    | NO    | YES   | YES   | YES   | NO    | NO    | NO    |
             ----------------------------------------------------------------------------------------------
-        Note that this means we require: `n_action_steps <= horizon - n_obs_steps + 1`. Also, note that
-        "horizon" may not the best name to describe what the variable actually means, because this period is
+        Note that this means we require: `n_action_steps <= chunk_size - n_obs_steps + 1`. Also, note that
+        this period is
         actually measured from the first observation which (if `n_obs_steps` > 1) happened in the past.
         """
         # NOTE: for offline evaluation, we have action in the batch, so we need to pop it out
@@ -213,7 +213,7 @@ class DiffusionModel(nn.Module):
             noise
             if noise is not None
             else torch.randn(
-                size=(batch_size, self.config.horizon, self.config.action_feature.shape[0]),
+                size=(batch_size, self.config.chunk_size, self.config.action_feature.shape[0]),
                 dtype=dtype,
                 device=device,
                 generator=generator,
@@ -309,16 +309,16 @@ class DiffusionModel(nn.Module):
                 AND/OR
             "observation.environment_state": (B, n_obs_steps, environment_dim)
 
-            "action": (B, horizon, action_dim)
-            "action_is_pad": (B, horizon)
+            "action": (B, chunk_size, action_dim)
+            "action_is_pad": (B, chunk_size)
         }
         """
         # Input validation.
         assert set(batch).issuperset({OBS_STATE, ACTION, "action_is_pad"})
         assert OBS_IMAGES in batch or OBS_ENV_STATE in batch
         n_obs_steps = batch[OBS_STATE].shape[1]
-        horizon = batch[ACTION].shape[1]
-        assert horizon == self.config.horizon
+        chunk_size = batch[ACTION].shape[1]
+        assert chunk_size == self.config.chunk_size
         assert n_obs_steps == self.config.n_obs_steps
 
         # Encode image features and concatenate them all together along with the state vector.

src/lerobot/policies/dsrl/configuration_dsrl.py

@@ -0,0 +1,244 @@
+# !/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.types import NormalizationMode
+from lerobot.optim.optimizers import MultiAdamConfig
+from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE
+
+
+def is_image_feature(key: str) -> bool:
+    """Check if a feature key represents an image feature.
+
+    Args:
+        key: The feature key to check
+
+    Returns:
+        True if the key represents an image feature, False otherwise
+    """
+    return key.startswith(OBS_IMAGE)
+
+
+@dataclass
+class ConcurrencyConfig:
+    """Configuration for the concurrency of the actor and learner.
+    Possible values are:
+    - "threads": Use threads for the actor and learner.
+    - "processes": Use processes for the actor and learner.
+    """
+
+    actor: str = "threads"
+    learner: str = "threads"
+
+
+@dataclass
+class ActorLearnerConfig:
+    learner_host: str = "127.0.0.1"
+    learner_port: int = 50051
+    policy_parameters_push_frequency: int = 4
+    queue_get_timeout: float = 2
+
+
+@dataclass
+class CriticNetworkConfig:
+    hidden_dims: list[int] = field(default_factory=lambda: [256, 256])
+    activate_final: bool = True
+    final_activation: str | None = None
+
+
+@dataclass
+class ActorNetworkConfig:
+    hidden_dims: list[int] = field(default_factory=lambda: [256, 256])
+    activate_final: bool = True
+    use_layer_norm: bool = True
+
+
+@dataclass
+class NoiseActorConfig:
+    """Configuration for the noise actor in DSRL.
+    The noise actor outputs noise that gets fed to the diffusion policy.
+    """
+
+    use_tanh_squash: bool = False  # Whether to bound the noise output
+    std_min: float = 1e-5
+    std_max: float = 2.0
+    init_final: float = 0.05
+
+
+@PreTrainedConfig.register_subclass("dsrl")
+@dataclass
+class DSRLConfig(PreTrainedConfig):
+    """Diffusion Steering via Reinforcement Learning (DSRL) configuration."""
+
+    # Mapping of feature types to normalization modes
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.MEAN_STD,
+            "STATE": NormalizationMode.MIN_MAX,
+            "ENV": NormalizationMode.MIN_MAX,
+            "ACTION": NormalizationMode.MIN_MAX,
+        }
+    )
+
+    # Statistics for normalizing different types of inputs
+    dataset_stats: dict[str, dict[str, list[float]]] | None = field(
+        default_factory=lambda: {
+            OBS_IMAGE: {
+                "mean": [0.485, 0.456, 0.406],
+                "std": [0.229, 0.224, 0.225],
+            },
+            OBS_STATE: {
+                "min": [0.0, 0.0],
+                "max": [1.0, 1.0],
+            },
+            ACTION: {
+                "min": [0.0, 0.0, 0.0],
+                "max": [1.0, 1.0, 1.0],
+            },
+        }
+    )
+
+    # Architecture specifics
+    # Device to run the model on (e.g., "cuda", "cpu")
+    device: str = "cpu"
+    # Device to store the model on
+    storage_device: str = "cpu"
+    # Name of the vision encoder model (Set to "helper2424/resnet10" for hil serl resnet10)
+    vision_encoder_name: str | None = None
+    # Whether to freeze the vision encoder during training
+    freeze_vision_encoder: bool = True
+    # Hidden dimension size for the image encoder
+    image_encoder_hidden_dim: int = 32
+    # Whether to use a shared encoder for actor and critic
+    shared_encoder: bool = True
+    # Number of discrete actions, eg for gripper actions
+    num_discrete_actions: int | None = None
+    # Dimension of the image embedding pooling
+    image_embedding_pooling_dim: int = 8
+
+    # Name of the action policy
+    action_policy_name: str = "pi0"
+    action_policy_weights: str | None = "lerobot/pi0_base"
+
+    # Training parameter
+    # Number of steps for online training
+    online_steps: int = 1000000
+    # Number of steps for offline training
+    offline_steps: int = 100000
+    # Capacity of the online replay buffer
+    online_buffer_capacity: int = 100000
+    # Capacity of the offline replay buffer
+    offline_buffer_capacity: int = 100000
+    # Whether to use asynchronous prefetching for the buffers
+    async_prefetch: bool = False
+    # Number of steps before learning starts
+    online_step_before_learning: int = 100
+    # Frequency of policy updates
+    policy_update_freq: int = 1
+
+    # SAC algorithm parameters
+    discount: float = 0.99
+    # Initial temperature value
+    temperature_init: float = 1.0
+    # Number of critics in the ensemble
+    num_critics: int = 2
+    # Number of subsampled critics for training
+    num_subsample_critics: int | None = None
+    # Learning rate for the critic network
+    critic_lr: float = 3e-4
+    # Learning rate for the actor network
+    actor_lr: float = 3e-4
+    # Learning rate for the temperature parameter
+    temperature_lr: float = 3e-4
+    # Weight for the critic target update
+    critic_target_update_weight: float = 0.005
+    # Update-to-data ratio for the UTD algorithm (If you want enable utd_ratio, you need to set it to >1)
+    utd_ratio: int = 1
+    # Hidden dimension size for the state encoder
+    state_encoder_hidden_dim: int = 256
+    # Dimension of the latent space
+    latent_dim: int = 256
+    # Target entropy for the SAC algorithm
+    target_entropy: float | None = None
+    # Whether to use backup entropy for the SAC algorithm
+    use_backup_entropy: bool = True
+    # Gradient clipping norm for the SAC algorithm
+    grad_clip_norm: float = 40.0
+
+    # Network configuration
+    # Configuration for the critic network architecture
+    critic_network_kwargs: CriticNetworkConfig = field(default_factory=CriticNetworkConfig)
+    # Configuration for the noise critic network architecture
+    noise_critic_network_kwargs: CriticNetworkConfig = field(default_factory=CriticNetworkConfig)
+    # Configuration for the noise actor network architecture
+    noise_actor_network_kwargs: ActorNetworkConfig = field(default_factory=ActorNetworkConfig)
+    # Configuration for the noise actor specific parameters
+    noise_actor_kwargs: NoiseActorConfig = field(default_factory=NoiseActorConfig)
+    # Configuration for actor-learner architecture
+    actor_learner_config: ActorLearnerConfig = field(default_factory=ActorLearnerConfig)
+    # Configuration for concurrency settings (you can use threads or processes for the actor and learner)
+    concurrency: ConcurrencyConfig = field(default_factory=ConcurrencyConfig)
+
+    # Optimizations
+    use_torch_compile: bool = True
+
+    def __post_init__(self):
+        super().__post_init__()
+
+    def get_optimizer_preset(self) -> MultiAdamConfig:
+        return MultiAdamConfig(
+            weight_decay=0.0,
+            optimizer_groups={
+                "critic_action": {"lr": self.critic_lr},
+                "critic_noise": {"lr": self.critic_lr},
+                "noise_actor": {"lr": self.actor_lr},
+                "temperature": {"lr": self.temperature_lr},
+            },
+        )
+
+    def get_scheduler_preset(self) -> None:
+        return None
+
+    def validate_features(self) -> None:
+        has_image = any(is_image_feature(key) for key in self.input_features)
+        has_state = OBS_STATE in self.input_features
+
+        if not (has_state or has_image):
+            raise ValueError(
+                "You must provide either 'observation.state' or an image observation (key starting with 'observation.image') in the input features"
+            )
+
+        if ACTION not in self.output_features:
+            raise ValueError("You must provide 'action' in the output features")
+
+    @property
+    def image_features(self) -> list[str]:
+        return [key for key in self.input_features if is_image_feature(key)]
+
+    @property
+    def observation_delta_indices(self) -> list:
+        return None
+
+    @property
+    def action_delta_indices(self) -> list:
+        return None
+
+    @property
+    def reward_delta_indices(self) -> None:
+        return None

src/lerobot/policies/dsrl/processor_dsrl.py

@@ -0,0 +1,89 @@
+# !/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.
+"""
+Processor for DSRL policy.
+
+DSRL uses a similar processing pipeline as SAC since it operates on
+state-action transitions. The main difference is that internally it
+also works with noise, but that's handled within the policy itself.
+"""
+
+from typing import Any
+
+import torch
+
+from lerobot.policies.dsrl.configuration_dsrl import DSRLConfig
+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_dsrl_pre_post_processors(
+    config: DSRLConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict, dict],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """Create preprocessor and postprocessor pipelines for DSRL policy.
+
+    Args:
+        config: DSRL policy configuration
+        dataset_stats: Optional dataset statistics for normalization
+
+    Returns:
+        Tuple of (preprocessor, postprocessor) pipelines
+    """
+    input_steps = [
+        RenameObservationsProcessorStep(rename_map={}),
+        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,
+        ),
+    )

src/lerobot/policies/factory.py

@@ -30,15 +30,15 @@ from lerobot.envs.configs import EnvConfig
 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.groot.configuration_groot import GrootConfig
+from lerobot.policies.dsrl.configuration_dsrl import DSRLConfig
 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
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
-from lerobot.policies.utils import validate_visual_features_consistency
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.processor import PolicyAction, PolicyProcessorPipeline
 from lerobot.processor.converters import (
@@ -59,7 +59,7 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
 
     Args:
         name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
-              "vqbet", "pi0", "pi05", "sac", "reward_classifier", "smolvla".
+              "vqbet", "pi0", "pi0fast", "sac", "reward_classifier", "smolvla", "dsrl".
 
     Returns:
         The policy class corresponding to the given name.
@@ -83,6 +83,10 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
         from lerobot.policies.vqbet.modeling_vqbet import VQBeTPolicy
 
         return VQBeTPolicy
+    elif name == "pi0fast":
+        from lerobot.policies.pi0fast.modeling_pi0fast import PI0FASTPolicy
+
+        return PI0FASTPolicy
     elif name == "pi0":
         from lerobot.policies.pi0.modeling_pi0 import PI0Policy
 
@@ -103,10 +107,10 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
         from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
 
         return SmolVLAPolicy
-    elif name == "groot":
-        from lerobot.policies.groot.modeling_groot import GrootPolicy
+    elif name == "dsrl":
+        from lerobot.policies.dsrl.modeling_dsrl import DSRLPolicy
 
-        return GrootPolicy
+        return DSRLPolicy
     else:
         raise NotImplementedError(f"Policy with name {name} is not implemented.")
 
@@ -120,8 +124,8 @@ 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", "pi05", "sac", "smolvla",
-                     "reward_classifier".
+                     "diffusion", "act", "vqbet", "pi0", "pi0fast", "sac", "smolvla",
+                     "reward_classifier", "dsrl".
         **kwargs: Keyword arguments to be passed to the configuration class constructor.
 
     Returns:
@@ -138,6 +142,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return ACTConfig(**kwargs)
     elif policy_type == "vqbet":
         return VQBeTConfig(**kwargs)
+    elif policy_type == "pi0fast":
+        return PI0FASTConfig(**kwargs)
     elif policy_type == "pi0":
         return PI0Config(**kwargs)
     elif policy_type == "pi05":
@@ -148,8 +154,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return SmolVLAConfig(**kwargs)
     elif policy_type == "reward_classifier":
         return RewardClassifierConfig(**kwargs)
-    elif policy_type == "groot":
-        return GrootConfig(**kwargs)
+    elif policy_type == "dsrl":
+        return DSRLConfig(**kwargs)
     else:
         raise ValueError(f"Policy type '{policy_type}' is not available.")
 
@@ -207,27 +213,6 @@ def make_pre_post_processors(
             policy configuration type.
     """
     if pretrained_path:
-        # TODO(Steven): Temporary patch, implement correctly the processors for Gr00t
-        if isinstance(policy_cfg, GrootConfig):
-            # GROOT handles normalization in groot_pack_inputs_v3 step
-            # Need to override both stats AND normalize_min_max since saved config might be empty
-            preprocessor_overrides = {}
-            postprocessor_overrides = {}
-            preprocessor_overrides["groot_pack_inputs_v3"] = {
-                "stats": kwargs.get("dataset_stats"),
-                "normalize_min_max": True,
-            }
-
-            # Also ensure postprocessing slices to env action dim and unnormalizes with dataset stats
-            env_action_dim = policy_cfg.output_features["action"].shape[0]
-            postprocessor_overrides["groot_action_unpack_unnormalize_v1"] = {
-                "stats": kwargs.get("dataset_stats"),
-                "normalize_min_max": True,
-                "env_action_dim": env_action_dim,
-            }
-            kwargs["preprocessor_overrides"] = preprocessor_overrides
-            kwargs["postprocessor_overrides"] = postprocessor_overrides
-
         return (
             PolicyProcessorPipeline.from_pretrained(
                 pretrained_model_name_or_path=pretrained_path,
@@ -282,6 +267,14 @@ 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
+
+        processors = make_pi0fast_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
     elif isinstance(policy_cfg, PI0Config):
         from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors
 
@@ -321,11 +314,10 @@ def make_pre_post_processors(
             config=policy_cfg,
             dataset_stats=kwargs.get("dataset_stats"),
         )
+    elif isinstance(policy_cfg, DSRLConfig):
+        from lerobot.policies.dsrl.processor_dsrl import make_dsrl_pre_post_processors
 
-    elif isinstance(policy_cfg, GrootConfig):
-        from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
-
-        processors = make_groot_pre_post_processors(
+        processors = make_dsrl_pre_post_processors(
             config=policy_cfg,
             dataset_stats=kwargs.get("dataset_stats"),
         )
@@ -340,7 +332,6 @@ def make_policy(
     cfg: PreTrainedConfig,
     ds_meta: LeRobotDatasetMetadata | None = None,
     env_cfg: EnvConfig | None = None,
-    rename_map: dict[str, str] | None = None,
 ) -> PreTrainedPolicy:
     """
     Instantiate a policy model.
@@ -357,8 +348,6 @@ def make_policy(
                  statistics for normalization layers.
         env_cfg: Environment configuration used to infer feature shapes and types.
                  One of `ds_meta` or `env_cfg` must be provided.
-        rename_map: Optional mapping of dataset or environment feature keys to match
-                 expected policy feature names (e.g., `"left"` → `"camera1"`).
 
     Returns:
         An instantiated and device-placed policy model.
@@ -400,10 +389,8 @@ def make_policy(
             raise ValueError("env_cfg cannot be None when ds_meta is not provided")
         features = env_to_policy_features(env_cfg)
 
-    if not cfg.output_features:
-        cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
-    if not cfg.input_features:
-        cfg.input_features = {key: ft for key, ft in features.items() if key not in cfg.output_features}
+    cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
+    cfg.input_features = {key: ft for key, ft in features.items() if key not in cfg.output_features}
     kwargs["config"] = cfg
 
     if cfg.pretrained_path:
@@ -420,8 +407,4 @@ def make_policy(
 
     # policy = torch.compile(policy, mode="reduce-overhead")
 
-    if not rename_map:
-        validate_visual_features_consistency(cfg, features)
-        # TODO: (jadechoghari) - add a check_state(cfg, features) and check_action(cfg, features)
-
     return policy

src/lerobot/policies/groot/README.md

@@ -1 +0,0 @@
-../../../../docs/source/policy_groot_README.md

src/lerobot/policies/groot/action_head/action_encoder.py

@@ -1,54 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# 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 as nn
-
-
-def swish(x):
-    return x * torch.sigmoid(x)
-
-
-class SinusoidalPositionalEncoding(nn.Module):
-    """
-    Produces a sinusoidal encoding of shape (B, T, w)
-    given timesteps of shape (B, T).
-    """
-
-    def __init__(self, embedding_dim):
-        super().__init__()
-        self.embedding_dim = embedding_dim
-
-    def forward(self, timesteps):
-        # timesteps: shape (B, T)
-        # We'll compute sin/cos frequencies across dim T
-        timesteps = timesteps.float()  # ensure float
-
-        b, t = timesteps.shape
-        device = timesteps.device
-
-        half_dim = self.embedding_dim // 2
-        # typical log space frequencies for sinusoidal encoding
-        exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * (
-            torch.log(torch.tensor(10000.0)) / half_dim
-        )
-        # Expand timesteps to (B, T, 1) then multiply
-        freqs = timesteps.unsqueeze(-1) * exponent.exp()  # (B, T, half_dim)
-
-        sin = torch.sin(freqs)
-        cos = torch.cos(freqs)
-        enc = torch.cat([sin, cos], dim=-1)  # (B, T, w)
-
-        return enc