fix(policies): vla jepa prepare model input to take index 0 and not index -1

feat(robot): mirror new configs from SO10X to its Bi manual counterpart (#4238 )
fix(datasets): stop frame errors being treated as shard exhaustion in StreamingLeRobotDataset (#4237 )
2026-07-30 21:19:40 +00:00 · 2026-07-30 15:19:48 +00:00 · 2026-07-30 17:11:58 +02:00 · 2026-07-30 16:19:37 +02:00 · 2026-07-30 15:53:34 +02:00 · 2026-07-30 15:27:51 +02:00
234 changed files with 4783 additions and 2029 deletions
@@ -0,0 +1,11 @@
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+    cooldown:
+      default-days: 7
+    groups:
+      actions:
+        patterns: ["*"]
@@ -51,6 +51,7 @@ pre-commit run --all-files                           # Lint + format (ruff, typo
 ## Notes

 - **Mypy is gradual**: strict only for `lerobot.envs`, `lerobot.configs`, `lerobot.optim`, `lerobot.model`, `lerobot.cameras`, `lerobot.motors`, `lerobot.transport`. Add type annotations when modifying these modules.
- **Optional dependencies**: many policies, envs, and robots are behind extras (e.g., `lerobot[aloha]`). New imports for optional packages must be guarded or lazy. See `pyproject.toml [project.optional-dependencies]`.
+- **Imports**: prefer top-level imports; relative (`from .sibling import X`) across sibling files within a module, absolute (`from lerobot.module import X`) across modules.
+- **Optional dependencies**: many policies, envs, and robots are behind extras (e.g., `lerobot[aloha]`, see `pyproject.toml`). Guard optional imports with `TYPE_CHECKING or _foo_available` at module top + a `require_package(...)` check at use time. Reuse the `_foo_available` flags in `utils/import_utils.py`; don't call `is_package_available`.
 - **Video decoding**: datasets can store observations as video files. `LeRobotDataset` handles frame extraction, but tests need ffmpeg installed.
 - **Prioritize use of `uv run`** to execute Python commands (not raw `python` or `pip`).
@@ -61,15 +61,19 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
 **4.1 Install**

 ```bash
-pip install 'lerobot[feetech]'              # SO-100/SO-101 motor stack
-# pip install 'lerobot[all]'                # everything
-# pip install 'lerobot[aloha,pusht]'        # specific features
-# pip install 'lerobot[smolvla]'            # add SmolVLA deps
-git lfs install && git lfs pull
-hf auth login                               # required to push datasets/policies
-```
+# uv (recommended — see AGENTS.md and CLAUDE.md)
+uv sync --locked --extra feetech          # SO-100/SO-101 motor stack
+# uv sync --locked --extra all            # everything
+# uv sync --locked --extra smolvla        # add SmolVLA deps

-Contributors can alternatively use `uv sync --locked --extra feetech` (see `AGENTS.md`).
+# pip (alternative, e.g. when not working from source)
+# pip install 'lerobot[feetech]'
+# pip install 'lerobot[all]'
+# pip install 'lerobot[smolvla]'
+
+git lfs install && git lfs pull
+hf auth login                             # required to push datasets/policies
+```

 **4.2 Find USB ports** — run once per arm, unplug when prompted.

@@ -83,7 +83,7 @@ episode_index=0
 print(f"{dataset[episode_index]['action'].shape=}\n")
 ```

-Learn more about it in the [LeRobotDataset Documentation](https://huggingface.co/docs/lerobot/lerobot-dataset-v3)
+Learn more about it in the [LeRobotDataset Documentation](https://huggingface.co/docs/lerobot/lerobot-dataset-v3).

 ## SoTA Models

@@ -109,7 +109,7 @@ lerobot-train \
 | **World Models**           | [VLA-JEPA](./docs/source/vla_jepa.mdx), [LingBot-VA](./docs/source/lingbot_va.mdx), [FastWAM](./docs/source/fastwam.mdx)                                                                                                                                                                                                                                                                   |
 | **Reward Models**          | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx)                                                                                                                                                                                                                                                                         |

-Similarly to the hardware, you can easily implement your own policy & leverage LeRobot's data collection, training, and visualization tools, and share your model to the HF Hub
+Similarly to the hardware, you can easily implement your own policy & leverage LeRobot's data collection, training, and visualization tools, and share your model to the HF Hub.

 For detailed policy setup guides, see the [Policy Documentation](https://huggingface.co/docs/lerobot/bring_your_own_policies). For GPU/RAM requirements and expected training time per policy, see the [Compute Hardware Guide](https://huggingface.co/docs/lerobot/hardware_guide).

@@ -126,7 +126,7 @@ lerobot-eval \
  --eval.n_episodes=10
 ```

-Learn how to implement your own simulation environment or benchmark and distribute it from the HF Hub by following the [EnvHub Documentation](https://huggingface.co/docs/lerobot/envhub)
+Learn how to implement your own simulation environment or benchmark and distribute it from the HF Hub by following the [EnvHub Documentation](https://huggingface.co/docs/lerobot/envhub).

 ## Resources

@@ -68,17 +68,16 @@ ENV HOME=/home/user_lerobot \
 # issues with MuJoCo and OpenGL drivers.
 RUN uv venv --python python${PYTHON_VERSION}

-# Install Python dependencies for caching
+# Install third-party dependencies separately for layer caching
 COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
-COPY --chown=user_lerobot:user_lerobot src/ src/
-
-RUN uv sync --locked --extra all --no-cache
+RUN uv sync --locked --extra all --no-install-project --no-cache

 RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas

-# Copy the rest of the application source code
+# Copy the application source code and install the local project
 # Make sure to have the git-LFS files for testing
 COPY --chown=user_lerobot:user_lerobot . .
+RUN uv sync --locked --extra all --no-cache

 # Set the default command
 CMD ["/bin/bash"]
@@ -60,15 +60,14 @@ ENV HOME=/home/user_lerobot \
 # run other Python projects in the same container without dependency conflicts.
 RUN uv venv

-# Install Python dependencies for caching
+# Install third-party dependencies separately for layer caching
 COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
-COPY --chown=user_lerobot:user_lerobot src/ src/
+RUN uv sync --locked --extra all --no-install-project --no-cache

-RUN uv sync --locked --extra all --no-cache
-
-# Copy the rest of the application code
+# Copy the application code and install the local project
 # Make sure to have the git-LFS files for testing
 COPY --chown=user_lerobot:user_lerobot . .
+RUN uv sync --locked --extra all --no-cache

 # Set the default command
 CMD ["/bin/bash"]
@@ -89,8 +89,8 @@ subtask.

 The resulting spans are then stitched into a gap-free, full-episode
 cover, so **every frame has exactly one active subtask**. See
-[`run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py)
-for the production settings (single camera, timestamped contact sheets,
+[Running on Hugging Face Jobs](#running-on-hugging-face-jobs) for the
+production settings (single camera, timestamped contact sheets,
 auto-windowed subtask generation).

 ### Tools
@@ -110,28 +110,67 @@ not-yet-implemented.

 ## Running on Hugging Face Jobs

-Annotation runs on [Hugging Face Jobs](https://huggingface.co/docs/hub/en/jobs).
-The repo ships a launcher script you copy and tweak for your dataset:
+Annotating a real dataset needs a GPU big enough to serve the VLM, so
+`lerobot-annotate` can dispatch itself to
+[Hugging Face Jobs](https://huggingface.co/docs/hub/en/jobs) — same as
+`lerobot-train`. Add `--job.target=<flavor>` to the exact command you'd
+run locally and it runs on that hardware instead:

 ```bash
-HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
+hf auth login   # once
+
+uv run lerobot-annotate \
+    --repo_id=user/my_dataset \
+    --new_repo_id=user/my_dataset_annotated \
+    --push_to_hub=true \
+    --vlm.model_id=Qwen/Qwen3.6-27B \
+    --vlm.num_gpus=1 \
+    --vlm.serve_command="vllm serve Qwen/Qwen3.6-27B --tensor-parallel-size 1 \
+        --max-model-len 32768 --gpu-memory-utilization 0.8 \
+        --uvicorn-log-level warning --port {port}" \
+    --vlm.serve_ready_timeout_s=1800 \
+    --vlm.chat_template_kwargs='{"enable_thinking": false}' \
+    --job.target=h200
 ```

-[`run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py)
-starts a single-GPU `h200` job (bump it to `h200x4` for big datasets)
-that:
+That submits a single-GPU `h200` job that:

-1. installs `lerobot` (from `main`) plus the annotation extras,
-2. boots one vLLM server per GPU (using the `vllm/vllm-openai` image) and
-   drives it over the OpenAI-compatible API,
-3. runs the `plan` / `interjections` / `vqa` modules across the dataset
-   with `lerobot-annotate`,
+1. starts from the `vllm/vllm-openai` image and installs `lerobot` on top,
+2. boots one vLLM server per GPU and drives it over the OpenAI-compatible API,
+3. runs the `plan` / `interjections` / `vqa` modules across the dataset,
 4. with `--push_to_hub=true`, uploads the result to `--new_repo_id` (or
   back to `--repo_id` in place if you leave that unset).

-To use a different dataset, model, or hub repo, edit the `CMD` block in
-the script. Every flag there maps directly to a `lerobot-annotate` flag
-(run `lerobot-annotate --help` for the full list).
+The command streams the job's logs; `Ctrl-C` detaches without cancelling
+it. List the available flavors and their pricing with `hf jobs hardware`.
+
+<Tip warning={true}>
+
+Qwen3.6 ships with thinking enabled, which eats the token budget the
+annotator needs for its JSON answer — `--vlm.chat_template_kwargs='{"enable_thinking": false}'`
+turns it off. Without `--push_to_hub=true` the annotated dataset is
+discarded when the pod exits.
+
+</Tip>
+
+### Job options
+
+| Flag                | Default                   | What it does                                                                    |
+| ------------------- | ------------------------- | ------------------------------------------------------------------------------- |
+| `--job.target`      | `local`                   | HF Jobs flavor to run on (e.g. `h200`, `h200x4`). Omitted/`local` runs here.    |
+| `--job.image`       | `vllm/vllm-openai:latest` | Runtime image for the pod.                                                      |
+| `--job.timeout`     | `2h`                      | Wall-clock cap. Raise it for large datasets.                                    |
+| `--job.detach`      | `false`                   | Submit and exit instead of streaming logs.                                      |
+| `--job.lerobot_ref` | `main`                    | Git ref of lerobot installed on the pod — point it at a branch to test changes. |
+| `--job.tags`        | `[]`                      | Extra tags on the job and on any dataset it pushes (`lerobot` is always added). |
+
+For a bigger dataset, scale to `h200x4` and raise
+`--vlm.parallel_servers` / `--vlm.num_gpus` to match, and give the job
+more headroom with e.g. `--job.timeout=8h`.
+
+Remote runs need `--repo_id` (the pod pulls the dataset from the Hub;
+`--root` names a directory only your machine has). A dataset that exists
+only in your local cache is pushed to a **private** repo first.

 ## Key options

@@ -58,7 +58,7 @@ final_action = postprocessor(action)

 ## Hardware API redesign

-PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is a overview of what changed and how you can continue to work with datasets created before this pull request.
+PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is an overview of what changed and how you can continue to work with datasets created before this pull request.

 ### What changed?

@@ -129,8 +129,8 @@ python examples/backward_compatibility/replay.py \

 Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied.

-To find these transformations, we recommend to first try and and replay an episode of the dataset your policy was trained on using the section above.
-Then, add these same transformations on your inference script (shown here in the `record.py` script):
+To find these transformations, we recommend first replaying an episode of the dataset your policy was trained on using the section above.
+Then, add these same transformations to your inference script (shown here in the `record.py` script):

 ```diff
 action_values = predict_action(
@@ -165,6 +165,8 @@ Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constant

 LeRobot uses `PolicyProcessorPipeline`s to normalize inputs and de-normalize outputs around your policy. For a concrete reference, see [`processor_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/processor_act.py) or [`processor_diffusion.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/diffusion/processor_diffusion.py).

+Pay close attention here: processors are the most common reproducibility pain point. A mismatch in normalization mode (`IDENTITY` vs `MEAN_STD` vs `MIN_MAX` vs `QUANTILES`/`QUANTILE10`) or in which features get normalized will train and eval without erroring, yet silently wreck results. Make sure the modes match how the checkpoint was trained, that the required stats exist (e.g. `QUANTILES` needs `q01`/`q99`), and that the pre- and post-processors stay consistent.
+
 ```python
 # processor_my_policy.py
 from typing import Any
@@ -304,7 +306,9 @@ Mirror an existing policy that's structurally similar to yours; the diff is smal

 ### Heavy / optional dependencies

-Most policies need a heavy backbone (transformers, diffusers, a specific VLM SDK). The convention is **two-step gating**: a `TYPE_CHECKING`-guarded import at module top, and a `require_package` runtime check in the constructor. [`modeling_diffusion.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/diffusion/modeling_diffusion.py) is the canonical reference:
+Most policies need a heavy backbone (transformers, diffusers, a specific VLM SDK). Wherever one exists, prefer loading it e.g from `transformers` or `diffusers` rather than re-implementing the architecture in-tree.
+
+The convention is **two-step gating**: a `TYPE_CHECKING`-guarded import at module top, and a `require_package` runtime check in the constructor. [`modeling_diffusion.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/diffusion/modeling_diffusion.py) is the canonical reference:

 ```python
 from typing import TYPE_CHECKING
@@ -374,6 +378,7 @@ The general expectations are in [`CONTRIBUTING.md`](https://github.com/huggingfa
 - [ ] Optional deps live behind a `[project.optional-dependencies]` extra and the `TYPE_CHECKING + require_package` guard.
 - [ ] `tests/policies/` updated; backward-compat artifact committed & policy-specific tests.
 - [ ] `src/lerobot/policies/<name>/README.md` symlinked into `docs/source/policy_<name>_README.md`; user-facing `docs/source/<name>.mdx` written and added to `_toctree.yml`.
+- [ ] `lerobot-train --policy.type my_policy ...` runs end-to-end for at least a few steps + save a checkpoint that can be loaded and run by `lerobot-eval` or `lerobot-rollout`.
 - [ ] `templates/lerobot_modelcard_template.md` has a description entry and a `policy_docs` link for your policy.
 - [ ] The models table in the root `README.md` lists your policy in the right category, linking to your doc page.
 - [ ] At least one reproducible benchmark eval in the policy MDX with a published checkpoint (sim benchmark, or real-robot dataset + checkpoint).
@@ -136,6 +136,10 @@ config = RealSenseCameraConfig(
    height=480,
    color_mode=ColorMode.RGB,
    use_depth=True,
+    # Optional fixed color controls. Omit them to leave the current sensor settings unchanged.
+    exposure=120,
+    gain=64,
+    white_balance=4600,
    rotation=Cv2Rotation.NO_ROTATION
 )

@@ -154,6 +158,15 @@ finally:
 ```
 <!-- prettier-ignore-end -->

+Manual color controls disable the corresponding automatic exposure or white-balance mode. Their
+supported ranges vary by camera model; an invalid value raises an error at connection time that
+includes the range reported by the sensor. Requesting an unsupported control also raises an error.
+Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
+require `use_rgb=True`.
+
+On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
+exposure or gain also affects the depth stream.
+
 </hfoption>
 </hfoptions>

@@ -88,20 +88,6 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)

 The same policy can work with different environment processors, and the same environment processor can work with different policies:

-````python
-# Use SmolVLA policy with LIBERO environment
-# Use SmolVLA policy with LIBERO environment
-libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
-    env_cfg=libero_cfg,
-    policy_cfg=smolvla_cfg,
-)
-smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
-# Or use ACT policy with the same LIBERO environment
-libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
-    env_cfg=libero_cfg,
-    policy_cfg=act_cfg,
-)
-act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
 ```python
 # Use SmolVLA policy with LIBERO environment
 libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
@@ -116,6 +102,7 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
    policy_cfg=act_cfg,
 )
 act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
+```

 ### 3. **Easier Experimentation**

@@ -145,7 +132,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
        state = torch.cat([eef_pos, eef_axisangle, eef_vel,
                          gripper_pos, gripper_vel], dim=-1)  # 14D
        return state
-````
+```

 ### 4. **Cleaner Environment Code**

@@ -40,10 +40,10 @@ This tutorial guides you through updating the firmware of Feetech motors using t
 For each motor you want to update:

 1. **Select the motor** from the list by clicking on it
-2. **Click on Upgrade tab**:
-3. **Click on Online button**:
-   - If an potential firmware update is found, it will be displayed in the box
-4. **Click on Upgrade button**:
+2. **Click the Upgrade tab**:
+3. **Click the Online button**:
+   - If a potential firmware update is found, it will be displayed in the box
+4. **Click the Upgrade button**:
   - The update progress will be displayed

 ## Step 6: Verify Update
@@ -59,6 +59,7 @@ The `lerobot-rollout --strategy.type=dagger` mode requires **teleoperators with

 - `bi_openarm_mini` - Bimanual OpenArm Mini
 - `so_leader` - SO100 / SO101 leader arm
+- `bi_so_leader` - Bimanual SO100 / SO101 leader arms

 > [!IMPORTANT]
 > The provided commands default to `bi_openarm_follower` + `bi_openarm_mini`.
@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.

 ### Record

-This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
+This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data).

 ```bash
 lerobot-record \
@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
 - [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) – for controlling Dynamixel servos

 Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
-For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py)
+For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py)

 Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):

@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
 pip install -e ".[lekiwi]"
 ```

-Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
+Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
 Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.

 # Step-by-Step Assembly Instructions
@@ -1,3 +1,11 @@
+# OMX
+
+<img
+  src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/omx_mainimage.png"
+  alt="OMX"
+  width=600
+/>
+
 ## Order and Assemble the parts

 First, assemble the OMX hardware following the official assembly guide.
@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs

 ## Reproducing π₀Fast results

-We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
+We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).

 The finetuned model can be found here:

@@ -22,7 +22,7 @@ With processors, you choose the learning features you want to use for your polic
 ## Three pipelines

 We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
-Each of these pipelines handle different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
+Each of these pipelines handles different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.

 1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets)
 2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
@@ -74,15 +74,15 @@ In the phone to SO-100 follower examples we use the following adapters:
 - `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
 - `transition_to_robot_action`: transforms the pipeline transition to a robot action dict.
 - `observation_to_transition`: transforms the robot observation dict to a pipeline transition.
- `transition_to_observation`: transforms the pipeline transition to a observation dict.
+- `transition_to_observation`: transforms the pipeline transition to an observation dict.

-Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
+Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.

 ## Dataset feature contracts

 Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`.

-Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
+Below is an example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:

 ```python
    def transform_features(
@@ -82,6 +82,8 @@ By default the env samples objects only from the `lightwheel` registry (what `--

 All eval snippets below mirror the CI command (see `.github/workflows/benchmark_tests.yml`). The `--rename_map` argument maps RoboCasa's native camera keys (`robot0_agentview_left` / `robot0_eye_in_hand` / `robot0_agentview_right`) onto the three-camera (`camera1` / `camera2` / `camera3`) input layout the released `smolvla_robocasa` policy was trained on.

+By default, each task uses the rollout horizon registered by RoboCasa. Set `--env.episode_length=<steps>` to apply the same explicit horizon to every selected task.
+
 ### Single-task evaluation (recommended for quick iteration)

 ```bash
@@ -57,7 +57,7 @@ policy_cfg.rtc_config = RTCConfig(
 policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")

 # Now use predict_action_chunk with RTC parameters
-inference_delay = 4  # How many steps of inference latency, this values should be calculated based on the inference latency of the policy
+inference_delay = 4  # How many steps of inference latency, this value should be calculated based on the inference latency of the policy

 # Initialize the action queue
 action_queue = ActionQueue(policy_cfg.rtc_config)
@@ -100,7 +100,7 @@ Typical values: 8-12 steps
 RTCConfig(execution_horizon=10)
 ```

-**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is a optimal value.
+**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is an optimal value.

 **`prefix_attention_schedule`**: How to weight consistency across the overlap region.

@@ -93,7 +93,7 @@ lerobot-train --help

 ## Evaluate the finetuned model and run it in real-time

-Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
+Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset).
 Once you are logged in, you can run inference in your setup by doing:

 ```bash
@@ -338,7 +338,7 @@ It is advisable to install one 3-pin cable in the motor after placing them befor
 <hfoption id="Leader">

 - Mount the leader holder onto the wrist and secure it with 4 M3x6mm screws.
- Attach the handle to motor 5 using 1 M2x6mm screw.
+- Attach the handle to the leader holder using 1 M2x6mm screw.
 - Insert the gripper motor, secure it with 2 M2x6mm screws on each side, attach a motor horn using a M3x6mm horn screw.
 - Attach the follower trigger with 4 M3x6mm screws.

@@ -50,11 +50,11 @@ lerobot-edit-dataset \
 Divide a dataset into multiple subsets.

 ```bash
-# Split by fractions (e.g. 80% train, 20% test, 20% val)
+# Split by fractions (e.g. 60% train, 20% val, 20% test)
 lerobot-edit-dataset \
    --repo_id lerobot/pusht \
    --operation.type split \
-    --operation.splits '{"train": 0.8, "test": 0.2, "val": 0.2}'
+    --operation.splits '{"train": 0.6, "val": 0.2, "test": 0.2}'

 # Split by specific episode indices
 lerobot-edit-dataset \
@@ -252,6 +252,10 @@ lerobot-dataset-viz \
    --episode-index 0
 ```

+For a private or gated dataset, authenticate first with `hf auth login`, or set the
+`HF_TOKEN` environment variable. The Hub client then discovers the credential
+automatically; no token argument is needed.
+
 **From a local folder:**
 Add the `--root` option and set `--mode local`. For example, to search in `./my_local_data_dir/lerobot/pusht`:

@@ -1,80 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 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.
-"""Launch ``lerobot-annotate`` on a Hugging Face job (vllm + Qwen3.6-27B VLM).
-
-Spawns one single-GPU ``h200`` job that:
-
-  1. installs ``lerobot`` from ``main`` plus the annotation extras,
-  2. boots one vllm server with Qwen3.6-27B (dense VLM),
-  3. runs the plan / interjections / vqa modules across the dataset
-     in free-form mode (each episode generates its own subtasks +
-     memory),
-  4. uploads the annotated dataset to ``--new_repo_id`` (when set)
-     or back to ``--repo_id``.
-
-Usage:
-
-    HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
-
-Adjust ``CMD`` (dataset, model, hub repo) and ``flavor`` below for your
-run. For larger datasets, scale to ``h200x4`` and raise
-``--vlm.parallel_servers`` / ``--vlm.num_gpus`` to match.
-"""
-
-import os
-
-from huggingface_hub import get_token, run_job
-
-token = os.environ.get("HF_TOKEN") or get_token()
-if not token:
-    raise RuntimeError("No HF token. Run `huggingface-cli login` or `export HF_TOKEN=hf_...`")
-
-CMD = (
-    "apt-get update -qq && apt-get install -y -qq git ffmpeg && "
-    "pip install --no-deps "
-    "'lerobot @ git+https://github.com/huggingface/lerobot.git@main' && "
-    # Pins mirror pyproject.toml — unpinned installs pull av 18 / datasets 5 /
-    # draccus 0.11, which break lerobot at import time.
-    "pip install --upgrade-strategy only-if-needed "
-    "'datasets>=4.7.0,<5.0.0' 'pyarrow>=21.0.0,<30.0.0' 'av>=15.0.0,<16.0.0' 'draccus==0.10.0' "
-    "'pandas>=2.0.0,<3.0.0' jsonlines gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
-    "openai && "
-    "export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
-    "export VLLM_VIDEO_BACKEND=pyav && "
-    "lerobot-annotate "
-    "--repo_id=pepijn223/robocasa_pretrain_human300_v4 "
-    "--new_repo_id=pepijn223/robocasa_pretrain_human300_v4_annotated "
-    "--push_to_hub=true "
-    "--vlm.backend=openai "
-    "--vlm.model_id=Qwen/Qwen3.6-27B "
-    "--vlm.num_gpus=1 "
-    '--vlm.serve_command="vllm serve Qwen/Qwen3.6-27B '
-    "--tensor-parallel-size 1 --max-model-len 32768 "
-    '--gpu-memory-utilization 0.8 --uvicorn-log-level warning --port {port}" '
-    "--vlm.serve_ready_timeout_s=1800 "
-    # Qwen3.6 ships with thinking on; annotation wants plain JSON answers.
-    "--vlm.chat_template_kwargs='{\"enable_thinking\": false}'"
-)
-
-job = run_job(
-    image="vllm/vllm-openai:latest",
-    command=["bash", "-c", CMD],
-    flavor="h200",
-    secrets={"HF_TOKEN": token},
-    timeout="2h",
-)
-print(f"Job URL: {job.url}")
-print(f"Job ID:  {job.id}")
@@ -44,6 +44,7 @@ from typing import Protocol

 import numpy as np

+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
    RobotProcessorPipeline,
@@ -56,7 +57,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
    EEBoundsAndSafety,
    InverseKinematicsEEToJoints,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, HF_LEROBOT_HOME, TELEOPERATORS
 from lerobot.utils.robot_utils import precise_sleep

@@ -38,7 +38,7 @@ from typing import TYPE_CHECKING

 import numpy as np

-from lerobot.types import RobotAction
+from lerobot.lerobot_types import RobotAction

 from .base import _GRIPPER_MOTOR_SCALE, IsaacTeleopTeleoperator, _isaacteleop_available
 from .config_isaac_teleop import SO101LeaderArmConfig
@@ -32,7 +32,7 @@ from typing import TYPE_CHECKING, Any

 import numpy as np

-from lerobot.types import RobotAction
+from lerobot.lerobot_types import RobotAction

 from .base import IsaacTeleopTeleoperator, _isaacteleop_available
 from .config_isaac_teleop import XRControllerConfig
@@ -26,8 +26,8 @@ from __future__ import annotations
 from dataclasses import dataclass

 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import RobotAction
 from lerobot.processor import ProcessorStepRegistry, RobotActionProcessorStep
-from lerobot.types import RobotAction
 from lerobot.utils.rotation import Rotation

 from .base import _GRIPPER_MOTOR_SCALE
@@ -21,6 +21,7 @@ from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.common.control_utils import predict_action
 from lerobot.configs import FeatureType, PolicyFeature
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.policies import make_pre_post_processors
 from lerobot.policies.act import ACTPolicy
@@ -38,7 +39,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
    ForwardKinematicsJointsToEE,
    InverseKinematicsEEToJoints,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts
 from lerobot.utils.keyboard_input import init_keyboard_listener
@@ -16,6 +16,7 @@

 from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
    RobotProcessorPipeline,
@@ -36,7 +37,6 @@ from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.phone import Phone, PhoneConfig
 from lerobot.teleoperators.phone.config_phone import PhoneOS
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.feature_utils import combine_feature_dicts
 from lerobot.utils.keyboard_input import init_keyboard_listener
 from lerobot.utils.utils import log_say
@@ -17,6 +17,7 @@
 import time

 from lerobot.datasets import LeRobotDataset
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
    RobotProcessorPipeline,
@@ -27,7 +28,6 @@ from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
 from lerobot.robots.so_follower.robot_kinematic_processor import (
    InverseKinematicsEEToJoints,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say
@@ -27,6 +27,7 @@ Highlight, or DAgger via ``lerobot-rollout --strategy.type=...``.

 from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.configs import PreTrainedConfig
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
    RobotProcessorPipeline,
@@ -43,7 +44,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
 from lerobot.rollout.inference import SyncInferenceConfig
 from lerobot.rollout.strategies import BaseStrategy
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.process import ProcessSignalHandler
 from lerobot.utils.utils import init_logging

@@ -15,6 +15,7 @@

 import time

+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
    RobotProcessorPipeline,
@@ -31,7 +32,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 from lerobot.teleoperators.phone import Phone, PhoneConfig
 from lerobot.teleoperators.phone.config_phone import PhoneOS
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.visualization_utils import init_rerun, log_rerun_data

@@ -21,6 +21,7 @@ from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.common.control_utils import predict_action
 from lerobot.configs import FeatureType, PolicyFeature
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.policies import make_pre_post_processors
 from lerobot.policies.act import ACTPolicy
@@ -38,7 +39,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
    ForwardKinematicsJointsToEE,
    InverseKinematicsEEToJoints,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts
 from lerobot.utils.keyboard_input import init_keyboard_listener
@@ -17,6 +17,7 @@

 from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
    RobotProcessorPipeline,
@@ -33,7 +34,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 )
 from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.feature_utils import combine_feature_dicts
 from lerobot.utils.keyboard_input import init_keyboard_listener
 from lerobot.utils.utils import log_say
@@ -18,6 +18,7 @@
 import time

 from lerobot.datasets import LeRobotDataset
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
    RobotProcessorPipeline,
@@ -28,7 +29,6 @@ from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
 from lerobot.robots.so_follower.robot_kinematic_processor import (
    InverseKinematicsEEToJoints,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say
@@ -25,6 +25,7 @@ forward/inverse kinematics.

 from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.configs import PreTrainedConfig
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
    RobotProcessorPipeline,
@@ -41,7 +42,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
 from lerobot.rollout.inference import SyncInferenceConfig
 from lerobot.rollout.strategies import BaseStrategy
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.process import ProcessSignalHandler
 from lerobot.utils.utils import init_logging

@@ -16,6 +16,7 @@

 import time

+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
    RobotProcessorPipeline,
@@ -30,7 +31,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
    InverseKinematicsEEToJoints,
 )
 from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.visualization_utils import init_rerun, log_rerun_data

@@ -67,7 +67,7 @@ dependencies = [
    "einops>=0.8.0,<0.9.0",

    # Config & Hub
-    "draccus==0.10.0", # TODO: Relax version constraint
+    "draccus>=0.11.6,<0.12.0",
    "huggingface-hub>=1.0.0,<2.0.0",
    "requests>=2.32.0,<3.0.0",

@@ -155,7 +155,7 @@ accelerate-dep = ["accelerate>=1.14.0,<2.0.0"]
 can-dep = ["python-can>=4.2.0,<5.0.0"]
 peft-dep = ["peft>=0.18.0,<1.0.0"]
 scipy-dep = ["scipy>=1.14.0,<2.0.0"]
-diffusers-dep = ["diffusers>=0.27.2,<0.36.0"]
+diffusers-dep = ["diffusers>=0.38.0,<0.40.0"]
 qwen-vl-utils-dep = ["qwen-vl-utils>=0.0.11,<0.1.0"]
 matplotlib-dep = ["matplotlib>=3.10.3,<4.0.0", "contourpy>=1.3.0,<2.0.0"] # NOTE: Explicitly listing contourpy helps the resolver converge faster.
 pyserial-dep = ["pyserial>=3.5,<4.0"]
@@ -494,6 +494,19 @@ ignore_errors = true
 module = "lerobot.envs.*"
 ignore_errors = false

+[[tool.mypy.overrides]]
+module = "lerobot.annotations.*"
+ignore_errors = false
+disallow_untyped_defs = true
+disallow_incomplete_defs = true
+check_untyped_defs = true
+
+[[tool.mypy.overrides]]
+module = "lerobot.transforms.*"
+ignore_errors = false
+disallow_untyped_defs = true
+disallow_incomplete_defs = true
+check_untyped_defs = true

 # [[tool.mypy.overrides]]
 # module = "lerobot.utils.*"
@@ -20,6 +20,29 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any

+from lerobot.configs.default import JobConfig
+
+# The annotation pipeline boots its own vLLM server, so the pod starts from the
+# official vLLM runtime rather than the prebuilt `lerobot-gpu` training image;
+# `lerobot` is pip-installed on top (see `lerobot.jobs.annotate`).
+DEFAULT_ANNOTATE_JOB_IMAGE = "vllm/vllm-openai:latest"
+
+
+@dataclass
+class AnnotationJobConfig(JobConfig):
+    """`JobConfig` with the annotation runtime's defaults.
+
+    Adds `lerobot_ref` because the vLLM image ships no lerobot: the pod installs
+    it from git, and the ref decides which code actually annotates. Point it at a
+    branch/tag/SHA to try unmerged changes remotely.
+    """
+
+    image: str = DEFAULT_ANNOTATE_JOB_IMAGE
+    # Annotation is a bounded pass over a dataset; a tighter cap than training's
+    # "2d" keeps a wedged vLLM server from burning a day of GPU time.
+    timeout: str | None = "2h"
+    lerobot_ref: str = "main"
+

@dataclass
 class PlanConfig:
@@ -207,6 +230,11 @@ class AnnotationPipelineConfig:
    vlm: VlmConfig = field(default_factory=VlmConfig)
    executor: ExecutorConfig = field(default_factory=ExecutorConfig)

+    # Where the annotation runs: omitted / "local" annotates on this machine, any
+    # other value is an HF Jobs flavor (e.g. "h200") and submits the run there.
+    # List flavors + pricing with `hf jobs hardware`.
+    job: AnnotationJobConfig = field(default_factory=AnnotationJobConfig)
+
    skip_validation: bool = False
    only_episodes: tuple[int, ...] | None = None

@@ -30,8 +30,8 @@ Phase 3 is why the ``plan`` module must be re-entered after the
 timestamps.

 Distributed execution is provided by Hugging Face Jobs (see
-``examples/annotations/run_hf_job.py``); the runner inside the job
-invokes ``lerobot-annotate`` which uses this in-process executor.
+``lerobot.jobs.annotate``, reached via ``--job.target=<flavor>``); the pod
+inside the job invokes ``lerobot-annotate`` which uses this in-process executor.
 Episode-level concurrency is controlled by
 ``ExecutorConfig.episode_parallelism``.
 """
@@ -194,12 +194,13 @@ def make_vlm_client(config: VlmConfig) -> VlmClient:
    """Build the shared VLM client.

    Only the ``openai`` backend is supported for now. The shipped workflow
-    is Hugging Face Jobs (``examples/annotations/run_hf_job.py``): it boots
-    a vLLM server inside the ``vllm/vllm-openai`` image and the pipeline
-    talks to it over the OpenAI-compatible API (``--vlm.backend=openai``,
-    optionally auto-spawning the server via ``auto_serve`` /
-    ``serve_command``). The former in-process ``vllm`` / ``transformers``
-    backends were removed to keep the support surface to the HF Jobs path.
+    is Hugging Face Jobs (``lerobot-annotate --job.target=<flavor>``): it
+    boots a vLLM server inside the ``vllm/vllm-openai`` image and the
+    pipeline talks to it over the OpenAI-compatible API
+    (``--vlm.backend=openai``, optionally auto-spawning the server via
+    ``auto_serve`` / ``serve_command``). The former in-process ``vllm`` /
+    ``transformers`` backends were removed to keep the support surface to
+    the HF Jobs path.

    For ``stub``, construct :class:`StubVlmClient` directly with a responder
    callable; it is rejected here to make accidental misuse obvious.
@@ -213,8 +214,8 @@ def make_vlm_client(config: VlmConfig) -> VlmClient:
    if config.backend in {"vllm", "transformers"}:
        raise ValueError(
            f"backend={config.backend!r} (in-process local model) is not supported for now — "
-            "only backend='openai' (the Hugging Face Jobs flow) is. Run the pipeline via "
-            "examples/annotations/run_hf_job.py, which serves the model with vLLM in the "
+            "only backend='openai' (the Hugging Face Jobs flow) is. Run the pipeline with "
+            "`lerobot-annotate --job.target=<flavor>`, which serves the model with vLLM in the "
            "vllm/vllm-openai image and talks to it over the OpenAI-compatible API."
        )
    raise ValueError(f"Unknown VLM backend: {config.backend!r}")
@@ -38,6 +38,7 @@ import draccus
 import grpc
 import torch

+from lerobot.lerobot_types import PolicyAction
 from lerobot.policies import get_policy_class, make_pre_post_processors
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.transport import (
@@ -45,7 +46,6 @@ from lerobot.transport import (
    services_pb2_grpc,  # type: ignore
 )
 from lerobot.transport.utils import receive_bytes_in_chunks
-from lerobot.types import PolicyAction

 from .configs import PolicyServerConfig
 from .constants import SUPPORTED_POLICIES
@@ -120,14 +120,22 @@ class OpenCVCamera(Camera):
        self.rotation: int | None = get_cv2_rotation(config.rotation)
        self.backend: int = config.backend

-        if self.height and self.width:
-            self.capture_width, self.capture_height = self.width, self.height
-            if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
-                self.capture_width, self.capture_height = self.height, self.width
+        self.capture_width: int | None = None
+        self.capture_height: int | None = None
+        self._reset_connection_settings()

    def __str__(self) -> str:
        return f"{self.__class__.__name__}({self.index_or_path})"

+    def _reset_connection_settings(self) -> None:
+        """Restore settings that may have been auto-detected during a failed connection."""
+        self.fps = self.config.fps
+        self.width = self.config.width
+        self.height = self.config.height
+        self.capture_width, self.capture_height = self.width, self.height
+        if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
+            self.capture_width, self.capture_height = self.height, self.width
+
    @property
    def is_connected(self) -> bool:
        """Checks if the camera is currently connected and opened."""
@@ -164,17 +172,25 @@ class OpenCVCamera(Camera):
                f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
            )

-        self._configure_capture_settings()
-        self._start_read_thread()
+        try:
+            self._configure_capture_settings()
+            self._start_read_thread()

-        if warmup and self.warmup_s > 0:
-            start_time = time.time()
-            while time.time() - start_time < self.warmup_s:
-                self.async_read(timeout_ms=self.warmup_s * 1000)
-                time.sleep(0.1)
-            with self.frame_lock:
-                if self.latest_frame is None:
-                    raise ConnectionError(f"{self} failed to capture frames during warmup.")
+            if warmup and self.warmup_s > 0:
+                start_time = time.time()
+                while time.time() - start_time < self.warmup_s:
+                    self.async_read(timeout_ms=self.warmup_s * 1000)
+                    time.sleep(0.1)
+                with self.frame_lock:
+                    if self.latest_frame is None:
+                        raise ConnectionError(f"{self} failed to capture frames during warmup.")
+        except BaseException:
+            try:
+                self._cleanup_resources()
+            except Exception:
+                logger.exception(f"Failed to fully clean up {self} after connect() failed.")
+            self._reset_connection_settings()
+            raise

        logger.info(f"{self} connected.")

@@ -312,32 +328,36 @@ class OpenCVCamera(Camera):

        for target in targets_to_scan:
            camera = cv2.VideoCapture(target)
-            if camera.isOpened():
-                default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
-                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)
+            try:
+                if camera.isOpened():
+                    default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
+                    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)])
+                    # 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",
-                    "id": target,
-                    "backend_api": camera.getBackendName(),
-                    "default_stream_profile": {
-                        "format": default_format,
-                        "fourcc": default_fourcc,
-                        "width": default_width,
-                        "height": default_height,
-                        "fps": default_fps,
-                    },
-                }
+                    camera_info = {
+                        "name": f"OpenCV Camera @ {target}",
+                        "type": "OpenCV",
+                        "id": target,
+                        "backend_api": camera.getBackendName(),
+                        "default_stream_profile": {
+                            "format": default_format,
+                            "fourcc": default_fourcc,
+                            "width": default_width,
+                            "height": default_height,
+                            "fps": default_fps,
+                        },
+                    }

-                found_cameras_info.append(camera_info)
+                    found_cameras_info.append(camera_info)
+            finally:
                camera.release()

        return found_cameras_info
@@ -496,6 +516,26 @@ class OpenCVCamera(Camera):
            self.latest_timestamp = None
            self.new_frame_event.clear()

+    def _cleanup_resources(self) -> None:
+        """Stop background reads and release the capture, including after partial setup."""
+        read_thread = self.thread
+        videocapture = self.videocapture
+
+        try:
+            self._stop_read_thread()
+        finally:
+            self.videocapture = None
+            try:
+                if videocapture is not None:
+                    videocapture.release()
+            finally:
+                # Releasing the device may unblock a hardware read that outlived
+                # the first bounded join in _stop_read_thread().
+                if read_thread is not None and read_thread.is_alive():
+                    read_thread.join(timeout=2.0)
+                    if read_thread.is_alive():  # pragma: no cover
+                        logger.warning(f"{self} read thread remained alive after releasing the capture.")
+
    @check_if_not_connected
    def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
        """
@@ -586,16 +626,6 @@ class OpenCVCamera(Camera):
        if not self.is_connected and self.thread is None:
            raise DeviceNotConnectedError(f"{self} not connected.")

-        if self.thread is not None:
-            self._stop_read_thread()
-
-        if self.videocapture is not None:
-            self.videocapture.release()
-            self.videocapture = None
-
-        with self.frame_lock:
-            self.latest_frame = None
-            self.latest_timestamp = None
-            self.new_frame_event.clear()
+        self._cleanup_resources()

        logger.info(f"{self} disconnected.")
@@ -173,7 +173,8 @@ class Reachy2Camera(Camera):
            raise ValueError(
                f"Invalid color mode '{self.color_mode}'. Expected {ColorMode.RGB} or {ColorMode.BGR}."
            )
-        if self.color_mode == ColorMode.RGB:
+        is_depth_frame = self.config.name == "depth" and self.config.image_type == "depth"
+        if not is_depth_frame and self.color_mode == ColorMode.RGB:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        self.latest_frame = frame
@@ -121,6 +121,9 @@ class RealSenseCamera(Camera):

        self.config = config

+        self.width: int | None = config.width
+        self.height: int | None = config.height
+
        if config.serial_number_or_name.isdigit():
            self.serial_number = config.serial_number_or_name
        else:
@@ -131,6 +134,9 @@ class RealSenseCamera(Camera):
        self.use_rgb = config.use_rgb
        self.use_depth = config.use_depth
        self.warmup_s = config.warmup_s
+        self.exposure: int | None = config.exposure
+        self.gain: int | None = config.gain
+        self.white_balance: int | None = config.white_balance

        self.rs_pipeline: rs.pipeline | None = None
        self.rs_profile: rs.pipeline_profile | None = None
@@ -145,14 +151,23 @@ class RealSenseCamera(Camera):

        self.rotation: int | None = get_cv2_rotation(config.rotation)

-        if self.height and self.width:
-            self.capture_width, self.capture_height = self.width, self.height
-            if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
-                self.capture_width, self.capture_height = self.height, self.width
+        self.capture_width: int | None = None
+        self.capture_height: int | None = None
+        self._reset_connection_settings()

    def __str__(self) -> str:
        return f"{self.__class__.__name__}({self.serial_number})"

+    def _reset_connection_settings(self) -> None:
+        """Restore settings that may have been auto-detected during a failed connection."""
+        self.fps = self.config.fps
+        self.width = self.config.width
+        self.height = self.config.height
+        self.warmup_s = self.config.warmup_s
+        self.capture_width, self.capture_height = self.width, self.height
+        if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
+            self.capture_width, self.capture_height = self.height, self.width
+
    @property
    def is_connected(self) -> bool:
        """Checks if the camera pipeline is started and streams are active."""
@@ -172,7 +187,8 @@ class RealSenseCamera(Camera):

        Raises:
            DeviceAlreadyConnectedError: If the camera is already connected.
-            ValueError: If the configuration is invalid (e.g., missing serial/name, name not unique).
+            ValueError: If the configuration is invalid, a requested sensor option is unsupported,
+                or a requested sensor value is invalid.
            ConnectionError: If the camera is found but fails to start the pipeline or no RealSense devices are detected at all.
            RuntimeError: If the pipeline starts but fails to apply requested settings.
        """
@@ -190,22 +206,31 @@ class RealSenseCamera(Camera):
                f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
            ) from e

-        self._configure_capture_settings()
-        self._start_read_thread()
+        try:
+            self._configure_capture_settings()
+            self._configure_sensor_options()
+            self._start_read_thread()

-        # NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise.
-        self.warmup_s = max(self.warmup_s, 1)
+            # NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise.
+            self.warmup_s = max(self.warmup_s, 1)

-        warmup_read = self.async_read if self.use_rgb else self.async_read_depth
-        start_time = time.time()
-        while time.time() - start_time < self.warmup_s:
-            warmup_read(timeout_ms=self.warmup_s * 1000)
-            time.sleep(0.1)
-        with self.frame_lock:
-            if (self.use_rgb and self.latest_color_frame is None) or (
-                self.use_depth and self.latest_depth_frame is None
-            ):
-                raise ConnectionError(f"{self} failed to capture frames during warmup.")
+            warmup_read = self.async_read if self.use_rgb else self.async_read_depth
+            start_time = time.time()
+            while time.time() - start_time < self.warmup_s:
+                warmup_read(timeout_ms=self.warmup_s * 1000)
+                time.sleep(0.1)
+            with self.frame_lock:
+                if (self.use_rgb and self.latest_color_frame is None) or (
+                    self.use_depth and self.latest_depth_frame is None
+                ):
+                    raise ConnectionError(f"{self} failed to capture frames during warmup.")
+        except BaseException:
+            try:
+                self._cleanup_resources()
+            except Exception:
+                logger.exception(f"Failed to fully clean up {self} after connect() failed.")
+            self._reset_connection_settings()
+            raise

        logger.info(f"{self} connected.")

@@ -339,6 +364,111 @@ class RealSenseCamera(Camera):
        self.new_frame_event.clear()
        return self._async_read(timeout_ms=10000, read_depth=read_depth)

+    def _get_color_sensor(self) -> "rs.sensor":
+        """Returns the sensor that controls the color stream.
+
+        Most RealSense cameras expose "RGB Camera" for color. The D405 has no
+        separate RGB module — its color stream comes from "Stereo Module".
+        We try RGB Camera first, then fall back to Stereo Module.
+        """
+        if self.rs_profile is None:
+            raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
+
+        device = self.rs_profile.get_device()
+        sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
+
+        for name in ("RGB Camera", "Stereo Module"):
+            if name in sensors:
+                return sensors[name]
+
+        available = list(sensors.keys())
+        raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
+
+    def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
+        """Sets a sensor option, re-raising range errors with actionable diagnostics."""
+        try:
+            sensor.set_option(option, value)
+        except Exception as e:
+            range_info = ""
+            try:
+                option_range = sensor.get_option_range(option)
+                range_info = (
+                    f" (supported range: min={option_range.min}, max={option_range.max}, "
+                    f"step={option_range.step}, default={option_range.default})"
+                )
+            except Exception:
+                range_info = " (option range unavailable)"
+            raise ValueError(
+                f"{self}: failed to set {label} to {value}{range_info}. Original error: {e}"
+            ) from e
+
+    def _configure_sensor_options(self) -> None:
+        """Applies manual sensor options (exposure, gain, white balance) to the color sensor.
+
+        When exposure or gain is set, auto-exposure is disabled first. When white_balance
+        is set, auto white balance is disabled first. An omitted option is left unchanged,
+        and configuration is skipped entirely if all options are omitted.
+
+        Raises:
+            ValueError: If the sensor does not support a requested option or a requested
+                value is invalid. Invalid-value errors include the option name, requested
+                value, and supported range when available.
+        """
+        if self.exposure is None and self.gain is None and self.white_balance is None:
+            return
+
+        color_sensor = self._get_color_sensor()
+
+        requested_options = (
+            (rs.option.exposure, self.exposure, "exposure"),
+            (rs.option.gain, self.gain, "gain"),
+            (rs.option.white_balance, self.white_balance, "white balance"),
+        )
+        unsupported_options = [
+            label
+            for option, value, label in requested_options
+            if value is not None and not color_sensor.supports(option)
+        ]
+        if unsupported_options:
+            raise ValueError(
+                f"{self}: color sensor does not support requested manual options: {unsupported_options}."
+            )
+
+        manual_exposure_requested = self.exposure is not None or self.gain is not None
+        if manual_exposure_requested:
+            if color_sensor.supports(rs.option.enable_auto_exposure):
+                self._set_sensor_option(color_sensor, rs.option.enable_auto_exposure, 0, "auto-exposure")
+                logger.info(f"{self} auto-exposure disabled.")
+            else:
+                logger.warning(
+                    f"{self} sensor does not support disabling auto-exposure; "
+                    "applying manual exposure/gain directly."
+                )
+
+        if self.exposure is not None:
+            self._set_sensor_option(color_sensor, rs.option.exposure, self.exposure, "exposure")
+            logger.info(f"{self} exposure set to {self.exposure}.")
+
+        if self.gain is not None:
+            self._set_sensor_option(color_sensor, rs.option.gain, self.gain, "gain")
+            logger.info(f"{self} gain set to {self.gain}.")
+
+        if self.white_balance is not None:
+            if color_sensor.supports(rs.option.enable_auto_white_balance):
+                self._set_sensor_option(
+                    color_sensor, rs.option.enable_auto_white_balance, 0, "auto white balance"
+                )
+                logger.info(f"{self} auto white balance disabled.")
+            else:
+                logger.warning(
+                    f"{self} sensor does not support disabling auto white balance; "
+                    "applying manual white balance directly."
+                )
+            self._set_sensor_option(
+                color_sensor, rs.option.white_balance, self.white_balance, "white balance"
+            )
+            logger.info(f"{self} white balance set to {self.white_balance}.")
+
    @check_if_not_connected
    def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
        """
@@ -453,7 +583,7 @@ class RealSenseCamera(Camera):
            )

        processed_image = image
-        if self.color_mode == ColorMode.BGR:
+        if not depth_frame and self.color_mode == ColorMode.BGR:
            processed_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

        if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE, cv2.ROTATE_180]:
@@ -541,6 +671,27 @@ class RealSenseCamera(Camera):
            self.latest_timestamp = None
            self.new_frame_event.clear()

+    def _cleanup_resources(self) -> None:
+        """Stop background reads and stop the pipeline, including after partial setup."""
+        read_thread = self.thread
+        rs_pipeline = self.rs_pipeline
+
+        try:
+            self._stop_read_thread()
+        finally:
+            self.rs_pipeline = None
+            self.rs_profile = None
+            try:
+                if rs_pipeline is not None:
+                    rs_pipeline.stop()
+            finally:
+                # Stopping the pipeline may unblock a hardware read that outlived
+                # the first bounded join in _stop_read_thread().
+                if read_thread is not None and read_thread.is_alive():
+                    read_thread.join(timeout=2.0)
+                    if read_thread.is_alive():  # pragma: no cover
+                        logger.warning(f"{self} read thread remained alive after stopping the pipeline.")
+
    def _async_read(self, timeout_ms: float, read_depth: bool = False) -> NDArray[Any]:
        """Shared helper for :meth:`async_read`/:meth:`async_read_depth`: return the latest buffered frame."""
        if self.thread is None or not self.thread.is_alive():
@@ -684,18 +835,5 @@ class RealSenseCamera(Camera):
                f"Attempted to disconnect {self}, but it appears already disconnected."
            )

-        if self.thread is not None:
-            self._stop_read_thread()
-
-        if self.rs_pipeline is not None:
-            self.rs_pipeline.stop()
-            self.rs_pipeline = None
-            self.rs_profile = None
-
-        with self.frame_lock:
-            self.latest_color_frame = None
-            self.latest_depth_frame = None
-            self.latest_timestamp = None
-            self.new_frame_event.clear()
-
+        self._cleanup_resources()
        logger.info(f"{self} disconnected.")
@@ -46,6 +46,17 @@ class RealSenseCameraConfig(CameraConfig):
        use_depth: Whether to enable depth stream. Defaults to False.
        rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
        warmup_s: Time reading frames before returning from connect (in seconds)
+        exposure: Manual exposure value for the color sensor. When set, auto-exposure is
+            disabled and this fixed value is used. Valid ranges are camera-model specific
+            and reported if the value is rejected. Defaults to None (leave unchanged).
+        gain: Manual gain value for the color sensor. When set, auto-exposure is disabled
+            and this fixed gain is used, which also freezes exposure at its current value
+            when no exposure is configured. Valid ranges are camera-model specific and
+            reported if the value is rejected. Defaults to None (leave unchanged).
+        white_balance: Manual white balance value for the color sensor. When set, auto
+            white balance is disabled and this fixed value is used. Valid ranges are
+            camera-model specific and reported if the value is rejected. Defaults to None
+            (leave unchanged).

    Note:
        - Either name or serial_number must be specified.
@@ -61,6 +72,9 @@ class RealSenseCameraConfig(CameraConfig):
    use_depth: bool = False
    rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
    warmup_s: int = 1
+    exposure: int | None = None
+    gain: int | None = None
+    white_balance: int | None = None

    def __post_init__(self) -> None:
        self.color_mode = ColorMode(self.color_mode)
@@ -69,6 +83,18 @@ class RealSenseCameraConfig(CameraConfig):
        if not self.use_rgb and not self.use_depth:
            raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.")

+        manual_color_options = {
+            "exposure": self.exposure,
+            "gain": self.gain,
+            "white_balance": self.white_balance,
+        }
+        configured_color_options = [name for name, value in manual_color_options.items() if value is not None]
+        if configured_color_options and not self.use_rgb:
+            raise ValueError(
+                "Manual color sensor options require `use_rgb=True`. "
+                f"Configured options: {configured_color_options}."
+            )
+
        values = (self.fps, self.width, self.height)
        if any(v is not None for v in values) and any(v is None for v in values):
            raise ValueError(
@@ -35,9 +35,9 @@ else:

 if TYPE_CHECKING:
    from lerobot.datasets import LeRobotDataset
+from lerobot.lerobot_types import PolicyAction
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.robots import Robot
-from lerobot.types import PolicyAction


 def predict_action(
@@ -71,13 +71,19 @@ class DatasetRecordConfig:
    # Number of threads per encoder instance. None = auto (codec default).
    # Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
    encoder_threads: int | None = None
+    # Skip appending the date-time tag to repo_id, keeping the user-provided name as-is
+    # (e.g. self-managed versioned names intended for a later `lerobot-edit-dataset merge`).
+    no_stamp: bool = False

    def stamp_repo_id(self) -> None:
        """Append a date-time tag to ``repo_id`` so each recording session gets a unique name.

        Must be called explicitly at dataset *creation* time — not on resume,
        where the existing ``repo_id`` (already stamped) must be preserved.
+        No-op when ``no_stamp`` is set, preserving a user-managed ``repo_id``.
        """
+        if self.no_stamp:
+            return
        if self.repo_id:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            self.repo_id = f"{self.repo_id}_{timestamp}"
@@ -163,8 +163,10 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: igno
        return None

    def _save_pretrained(self, save_directory: Path) -> None:
-        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
-            draccus.dump(self, f, indent=4)
+        # Encode against the base class so draccus includes the choice "type" key,
+        # which `from_pretrained` needs to resolve the concrete subclass.
+        with open(save_directory / CONFIG_NAME, "w") as f:
+            json.dump(draccus.encode(self, PreTrainedConfig), f, indent=4)

    @classmethod
    def from_pretrained(
@@ -103,8 +103,10 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
        pass

    def _save_pretrained(self, save_directory: Path) -> None:
-        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
-            draccus.dump(self, f, indent=4)
+        # Encode against the base class so draccus includes the choice "type" key,
+        # which `from_pretrained` needs to resolve the concrete subclass.
+        with open(save_directory / CONFIG_NAME, "w") as f:
+            json.dump(draccus.encode(self, RewardModelConfig), f, indent=4)

    @classmethod
    def from_pretrained(
@@ -14,6 +14,7 @@
 import builtins
 import datetime as dt
 import json
+import multiprocessing
 import os
 import tempfile
 from dataclasses import dataclass, field
@@ -101,6 +102,12 @@ class TrainPipelineConfig(HubMixin):
    batch_size: int = 8
    prefetch_factor: int = 4
    persistent_workers: bool = True
+    # DataLoader worker start method. "spawn" is safer than "fork" with
+    # non-fork-safe libs (PyAV / torchcodec / ffmpeg), but adds some
+    # worker-startup time per run since workers re-import modules instead
+    # of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
+    # when appropriate, or set it to `null` to use Python's platform default.
+    dataloader_multiprocessing_context: str | None = "spawn"
    steps: int = 100_000
    # Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
    env_eval_freq: int = 20_000
@@ -187,7 +194,11 @@ class TrainPipelineConfig(HubMixin):
            )

        if Path(config_path).resolve().exists():
-            policy_dir = Path(config_path).parent
+            # `config_path` may point at the checkpoint's train_config.json or at its
+            # pretrained_model/ directory (both documented above) — resolve either to
+            # the pretrained_model/ directory.
+            config_path_obj = Path(config_path)
+            policy_dir = config_path_obj.parent if config_path_obj.is_file() else config_path_obj
            self.checkpoint_path = policy_dir.parent
        elif self.job.is_remote:
            return
@@ -212,6 +223,17 @@ class TrainPipelineConfig(HubMixin):
            self.reward_model.pretrained_path = str(policy_dir)

    def validate(self) -> None:
+        available_contexts = multiprocessing.get_all_start_methods()
+        if (
+            self.dataloader_multiprocessing_context is not None
+            and self.dataloader_multiprocessing_context not in available_contexts
+        ):
+            raise ValueError(
+                "`dataloader_multiprocessing_context` must be None or one of "
+                f"{available_contexts} on this platform, got "
+                f"{self.dataloader_multiprocessing_context!r}."
+            )
+
        self._resolve_pretrained_from_cli()

        if self.policy is None and self.reward_model is None:
@@ -19,6 +19,7 @@ import copy
 import logging
 import shutil
 from pathlib import Path
+from typing import Any, NotRequired, TypedDict

 import datasets
 import pandas as pd
@@ -49,8 +50,32 @@ from .utils import (
 )
 from .video_utils import concatenate_video_files, get_video_duration_in_s

+logger = logging.getLogger(__name__)

-def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> dict[str, dict]:
+type FeatureDict = dict[str, dict[str, Any]]
+type ChunkFile = tuple[int, int]
+
+
+class IndexState(TypedDict):
+    chunk: int
+    file: int
+    src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
+
+
+class VideoIndex(TypedDict):
+    chunk: int
+    file: int
+    latest_duration: float
+    episode_duration: float
+    src_to_offset: NotRequired[dict[ChunkFile, float]]
+    src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
+    dst_file_durations: NotRequired[dict[ChunkFile, float]]
+
+
+type VideoIndexState = dict[str, VideoIndex]
+
+
+def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> FeatureDict:
    """Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged.

    Args:
@@ -59,14 +84,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
    Returns:
        dict: A dictionary of merged video feature info.
    """
-    merged_info = copy.deepcopy(all_metadata[0].features)
+    merged_info: FeatureDict = copy.deepcopy(all_metadata[0].features)
    video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"]

    for vk in video_keys:
        video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
        base_video_info = video_infos[0]

-        merged_encoder_info: dict = {}
+        merged_encoder_info: dict[str, Any] = {}
        fallback_keys: list[str] = []
        for info_key in VIDEO_ENCODER_INFO_KEYS:
            values = [info.get(info_key, None) for info in video_infos]
@@ -80,7 +105,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
                merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None

        if fallback_keys:
-            logging.warning(
+            logger.warning(
                f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
                f"Setting these keys to null: {fallback_keys}.",
            )
@@ -92,7 +117,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
    return merged_info


-def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
+def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[int, str | None, FeatureDict]:
    """Validates that all dataset metadata have consistent properties.

    Ensures all datasets have the same fps, robot_type, and features to guarantee
@@ -129,7 +154,9 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
    return fps, robot_type, features


-def update_data_df(df, src_meta, dst_meta):
+def update_data_df(
+    df: pd.DataFrame, src_meta: LeRobotDatasetMetadata, dst_meta: LeRobotDatasetMetadata
+) -> pd.DataFrame:
    """Updates a data DataFrame with new indices and task mappings for aggregation.

    Adjusts episode indices, frame indices, and task indices to account for
@@ -154,12 +181,12 @@ def update_data_df(df, src_meta, dst_meta):


 def update_meta_data(
-    df,
-    dst_meta,
-    meta_idx,
-    data_idx,
-    videos_idx,
-):
+    df: pd.DataFrame,
+    dst_meta: LeRobotDatasetMetadata,
+    meta_idx: IndexState,
+    data_idx: IndexState,
+    videos_idx: VideoIndexState,
+) -> pd.DataFrame:
    """Updates metadata DataFrame with new chunk, file, and timestamp indices.

    Adjusts all indices and timestamps to account for previously aggregated
@@ -289,7 +316,7 @@ def aggregate_datasets(
    chunk_size: int | None = None,
    concatenate_videos: bool = True,
    concatenate_data: bool = True,
-):
+) -> None:
    """Aggregates multiple LeRobot datasets into a single unified dataset.

    This is the main function that orchestrates the aggregation process by:
@@ -309,7 +336,7 @@ def aggregate_datasets(
        concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
        concatenate_data: When False, keep one parquet per source file instead of packing into shards.
    """
-    logging.info("Start aggregate_datasets")
+    logger.info("Start aggregate_datasets")

    if data_files_size_in_mb is None:
        data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
@@ -341,15 +368,15 @@ def aggregate_datasets(
        video_files_size_in_mb=video_files_size_in_mb,
    )

-    logging.info("Find all tasks")
+    logger.info("Find all tasks")
    unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
    dst_meta.tasks = pd.DataFrame(
        {"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
    )

-    meta_idx = {"chunk": 0, "file": 0}
-    data_idx = {"chunk": 0, "file": 0}
-    videos_idx = {
+    meta_idx: IndexState = {"chunk": 0, "file": 0}
+    data_idx: IndexState = {"chunk": 0, "file": 0}
+    videos_idx: VideoIndexState = {
        key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
    }

@@ -373,12 +400,17 @@ def aggregate_datasets(
        dst_meta.info.total_frames += src_meta.total_frames

    finalize_aggregation(dst_meta, all_metadata)
-    logging.info("Aggregation complete.")
+    logger.info("Aggregation complete.")


 def aggregate_videos(
-    src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True
-):
+    src_meta: LeRobotDatasetMetadata,
+    dst_meta: LeRobotDatasetMetadata,
+    videos_idx: VideoIndexState,
+    video_files_size_in_mb: float,
+    chunk_size: int,
+    concatenate_videos: bool = True,
+) -> VideoIndexState:
    """Aggregates video chunks from a source dataset into the destination dataset.

    Handles video file concatenation and rotation based on file size limits.
@@ -406,15 +438,16 @@ def aggregate_videos(
            videos_idx[key]["dst_file_durations"] = {}

    for key, video_idx in videos_idx.items():
-        unique_chunk_file_pairs = {
-            (chunk, file)
-            for chunk, file in zip(
-                src_meta.episodes[f"videos/{key}/chunk_index"],
-                src_meta.episodes[f"videos/{key}/file_index"],
-                strict=False,
-            )
-        }
-        unique_chunk_file_pairs = sorted(unique_chunk_file_pairs)
+        unique_chunk_file_pairs: list[ChunkFile] = sorted(
+            {
+                (chunk, file)
+                for chunk, file in zip(
+                    src_meta.episodes[f"videos/{key}/chunk_index"],
+                    src_meta.episodes[f"videos/{key}/file_index"],
+                    strict=False,
+                )
+            }
+        )

        chunk_idx = video_idx["chunk"]
        file_idx = video_idx["file"]
@@ -489,7 +522,14 @@ def aggregate_videos(
    return videos_idx


-def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data=True):
+def aggregate_data(
+    src_meta: LeRobotDatasetMetadata,
+    dst_meta: LeRobotDatasetMetadata,
+    data_idx: IndexState,
+    data_files_size_in_mb: float,
+    chunk_size: int,
+    concatenate_data: bool = True,
+) -> IndexState:
    """Aggregates data chunks from a source dataset into the destination dataset.

    Reads source data files, updates indices to match the aggregated dataset,
@@ -510,14 +550,16 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
    Returns:
        dict: Updated data_idx with current chunk and file indices.
    """
-    unique_chunk_file_ids = {
-        (c, f)
-        for c, f in zip(
-            src_meta.episodes["data/chunk_index"], src_meta.episodes["data/file_index"], strict=False
-        )
-    }
-
-    unique_chunk_file_ids = sorted(unique_chunk_file_ids)
+    unique_chunk_file_ids: list[ChunkFile] = sorted(
+        {
+            (c, f)
+            for c, f in zip(
+                src_meta.episodes["data/chunk_index"],
+                src_meta.episodes["data/file_index"],
+                strict=False,
+            )
+        }
+    )
    contains_images = len(dst_meta.image_keys) > 0

    # retrieve features schema for proper image typing in parquet
@@ -525,7 +567,7 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si

    # Track source to destination file mapping for metadata update
    # This is critical for handling datasets that are already results of a merge
-    src_to_dst: dict[tuple[int, int], tuple[int, int]] = {}
+    src_to_dst: dict[ChunkFile, ChunkFile] = {}

    for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
        src_path = src_meta.root / DEFAULT_DATA_PATH.format(
@@ -564,7 +606,13 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
    return data_idx


-def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
+def aggregate_metadata(
+    src_meta: LeRobotDatasetMetadata,
+    dst_meta: LeRobotDatasetMetadata,
+    meta_idx: IndexState,
+    data_idx: IndexState,
+    videos_idx: VideoIndexState,
+) -> IndexState:
    """Aggregates metadata from a source dataset into the destination dataset.

    Reads source metadata files, updates all indices and timestamps,
@@ -580,16 +628,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
    Returns:
        dict: Updated meta_idx with current chunk and file indices.
    """
-    chunk_file_ids = {
-        (c, f)
-        for c, f in zip(
-            src_meta.episodes["meta/episodes/chunk_index"],
-            src_meta.episodes["meta/episodes/file_index"],
-            strict=False,
-        )
-    }
-
-    chunk_file_ids = sorted(chunk_file_ids)
+    chunk_file_ids: list[ChunkFile] = sorted(
+        {
+            (c, f)
+            for c, f in zip(
+                src_meta.episodes["meta/episodes/chunk_index"],
+                src_meta.episodes["meta/episodes/file_index"],
+                strict=False,
+            )
+        }
+    )
    for chunk_idx, file_idx in chunk_file_ids:
        src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
        df = pd.read_parquet(src_path)
@@ -622,16 +670,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
 def append_or_create_parquet_file(
    df: pd.DataFrame,
    src_path: Path,
-    idx: dict[str, int],
+    idx: IndexState,
    max_mb: float,
    chunk_size: int,
    default_path: str,
    contains_images: bool = False,
-    aggr_root: Path = None,
+    aggr_root: Path | None = None,
    hf_features: datasets.Features | None = None,
    concatenate: bool = True,
    one_row_group_per_episode: bool = False,
-) -> tuple[dict[str, int], tuple[int, int]]:
+) -> tuple[IndexState, ChunkFile]:
    """Appends data to an existing parquet file or creates a new one based on size constraints.

    Manages file rotation when size limits are exceeded to prevent individual files
@@ -654,7 +702,13 @@ def append_or_create_parquet_file(
    Returns:
        tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
               and (dst_chunk, dst_file) is the actual destination file the data was written to.
+
+    Raises:
+        ValueError: If aggr_root is not provided.
    """
+    if aggr_root is None:
+        raise ValueError("aggr_root must be provided.")
+
    dst_chunk, dst_file = idx["chunk"], idx["file"]
    dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)

@@ -698,7 +752,9 @@ def append_or_create_parquet_file(
    return idx, (dst_chunk, dst_file)


-def finalize_aggregation(aggr_meta, all_metadata):
+def finalize_aggregation(
+    aggr_meta: LeRobotDatasetMetadata, all_metadata: list[LeRobotDatasetMetadata]
+) -> None:
    """Finalizes the dataset aggregation by writing summary files and statistics.

    Writes the tasks file, info file with total counts and splits, and
@@ -708,16 +764,16 @@ def finalize_aggregation(aggr_meta, all_metadata):
        aggr_meta: Aggregated dataset metadata.
        all_metadata: List of all source dataset metadata objects.
    """
-    logging.info("write tasks")
+    logger.info("write tasks")
    write_tasks(aggr_meta.tasks, aggr_meta.root)

-    logging.info("write info")
+    logger.info("write info")
    aggr_meta.info.total_tasks = len(aggr_meta.tasks)
    aggr_meta.info.total_episodes = sum(m.total_episodes for m in all_metadata)
    aggr_meta.info.total_frames = sum(m.total_frames for m in all_metadata)
    aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"}
    write_info(aggr_meta.info, aggr_meta.root)

-    logging.info("write stats")
+    logger.info("write stats")
    aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
    write_stats(aggr_meta.stats, aggr_meta.root)
@@ -73,6 +73,8 @@ class LeRobotDatasetMetadata:
        revision: str | None = None,
        force_cache_sync: bool = False,
        metadata_buffer_size: int = 10,
+        *,
+        token: str | bool | None = None,
    ):
        """Load or download metadata for an existing LeRobot dataset.

@@ -94,6 +96,10 @@ class LeRobotDatasetMetadata:
                even when local files exist.
            metadata_buffer_size: Number of episode metadata records to buffer
                in memory before flushing to parquet.
+            token: Authentication token used for Hub requests. Pass a string
+                token, ``True`` to require the locally stored token, ``False``
+                to disable authentication, or ``None`` to use the Hugging Face
+                Hub default.
        """
        self.repo_id = repo_id
        self.revision = revision if revision else CODEBASE_VERSION
@@ -113,9 +119,12 @@ class LeRobotDatasetMetadata:
            self._load_metadata()
        except (FileNotFoundError, NotADirectoryError):
            if is_valid_version(self.revision):
-                self.revision = get_safe_version(self.repo_id, self.revision)
+                if token is None:
+                    self.revision = get_safe_version(self.repo_id, self.revision)
+                else:
+                    self.revision = get_safe_version(self.repo_id, self.revision, token=token)

-            self._pull_from_repo(allow_patterns="meta/")
+            self._pull_from_repo(allow_patterns="meta/", token=token)
            self._load_metadata()

    def _flush_metadata_buffer(self) -> None:
@@ -179,8 +188,8 @@ class LeRobotDatasetMetadata:
    def _load_metadata(self):
        self.info = load_info(self.root)
        check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
-        self.tasks = load_tasks(self.root)
-        self.episodes = load_episodes(self.root)
+        self.tasks = load_tasks(self.root) if self.total_tasks > 0 else None
+        self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None
        self.stats = load_stats(self.root)

    def ensure_readable(self) -> None:
@@ -220,7 +229,10 @@ class LeRobotDatasetMetadata:
        self,
        allow_patterns: list[str] | str | None = None,
        ignore_patterns: list[str] | str | None = None,
+        *,
+        token: str | bool | None = None,
    ) -> None:
+        token_kwargs = {} if token is None else {"token": token}
        if self._requested_root is None:
            self.root = Path(
                snapshot_download(
@@ -230,6 +242,7 @@ class LeRobotDatasetMetadata:
                    cache_dir=HF_LEROBOT_HUB_CACHE,
                    allow_patterns=allow_patterns,
                    ignore_patterns=ignore_patterns,
+                    **token_kwargs,
                )
            )
            return
@@ -242,6 +255,7 @@ class LeRobotDatasetMetadata:
            local_dir=self._requested_root,
            allow_patterns=allow_patterns,
            ignore_patterns=ignore_patterns,
+            **token_kwargs,
        )
        self.root = self._requested_root

@@ -172,6 +172,23 @@ class DatasetWriter:
    def _get_image_file_dir(self, episode_index: int, image_key: str) -> Path:
        return self._get_image_file_path(episode_index, image_key, frame_index=0).parent

+    def _get_episode_buffer_index(self) -> int:
+        episode_index = self.episode_buffer["episode_index"]
+        # episode_index is `int` when freshly created, but becomes `np.ndarray` after
+        # save_episode() mutates the buffer. Handle both types here.
+        if isinstance(episode_index, np.ndarray):
+            episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
+        return int(episode_index)
+
+    def _delete_camera_frame_dirs(self, camera_keys: list[str]) -> None:
+        if self.image_writer is not None:
+            self._wait_image_writer()
+        episode_index = self._get_episode_buffer_index()
+        for camera_key in camera_keys:
+            img_dir = self._get_image_file_dir(episode_index, camera_key)
+            if img_dir.is_dir():
+                shutil.rmtree(img_dir)
+
    def _save_image(
        self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1
    ) -> None:
@@ -369,7 +386,9 @@ class DatasetWriter:
                self._episodes_since_last_encoding = 0

        if episode_data is None:
-            self.clear_episode_buffer(delete_images=len(self._meta.image_keys) > 0)
+            if len(self._meta.image_keys) > 0:
+                self._delete_camera_frame_dirs(self._meta.image_keys)
+            self.episode_buffer = self._create_episode_buffer()

    def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None) -> None:
        """Batch save videos for multiple episodes."""
@@ -561,10 +580,10 @@ class DatasetWriter:
        return metadata

    def clear_episode_buffer(self, delete_images: bool = True) -> None:
-        """Discard the current episode buffer and optionally delete temp images.
+        """Discard the current episode buffer and optionally delete temp camera frames.

        Args:
-            delete_images: If ``True``, remove temporary image directories
+            delete_images: If ``True``, remove temporary camera frame directories
                written for the current episode.
        """
        # Cancel streaming encoder if active
@@ -572,17 +591,7 @@ class DatasetWriter:
            self._streaming_encoder.cancel_episode()

        if delete_images:
-            if self.image_writer is not None:
-                self._wait_image_writer()
-            episode_index = self.episode_buffer["episode_index"]
-            # episode_index is `int` when freshly created, but becomes `np.ndarray` after
-            # save_episode() mutates the buffer. Handle both types here.
-            if isinstance(episode_index, np.ndarray):
-                episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
-            for cam_key in self._meta.image_keys:
-                img_dir = self._get_image_file_dir(episode_index, cam_key)
-                if img_dir.is_dir():
-                    shutil.rmtree(img_dir)
+            self._delete_camera_frame_dirs(self._meta.camera_keys)

        self.episode_buffer = self._create_episode_buffer()

@@ -65,6 +65,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
        encoder_threads: int | None = None,
        streaming_encoding: bool = False,
        encoder_queue_maxsize: int = 30,
+        *,
+        token: str | bool | None = None,
    ):
        """
        2 modes are available for instantiating this class, depending on 2 different use cases:
@@ -197,6 +199,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
                instead of writing PNG images first. This makes save_episode() near-instant. Defaults to False.
            encoder_queue_maxsize (int, optional): Maximum number of frames to buffer per camera when using
                streaming encoding. Defaults to 30 (~1s at 30fps).
+            token: Authentication token used while downloading this dataset
+                from the Hub. Pass a string token, ``True`` to require the
+                locally stored token, ``False`` to disable authentication, or
+                ``None`` to use the Hugging Face Hub default. The token is not
+                retained on the dataset instance after initialization.

        Note:
            Write-mode parameters (``streaming_encoding``, ``batch_encoding_size``) passed to
@@ -220,7 +227,11 @@ class LeRobotDataset(torch.utils.data.Dataset):

        # Load metadata (sets self.root once from the resolved metadata root)
        self.meta = LeRobotDatasetMetadata(
-            self.repo_id, self._requested_root, self.revision, force_cache_sync=force_cache_sync
+            self.repo_id,
+            self._requested_root,
+            self.revision,
+            force_cache_sync=force_cache_sync,
+            token=token,
        )
        self.root = self.meta.root
        self.revision = self.meta.revision
@@ -260,8 +271,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
        # Load actual data
        if force_cache_sync or not self.reader.try_load():
            if is_valid_version(self.revision):
-                self.revision = get_safe_version(self.repo_id, self.revision)
-            self._download(download_videos)
+                if token is None:
+                    self.revision = get_safe_version(self.repo_id, self.revision)
+                else:
+                    self.revision = get_safe_version(self.repo_id, self.revision, token=token)
+            self._download(download_videos, token=token)
            self.reader.load_and_activate()

        # Detect write-mode params for backward compatibility
@@ -478,18 +492,19 @@ class LeRobotDataset(torch.utils.data.Dataset):
        """Return the number of frames in the selected episodes."""
        return self.num_frames

-    def __getitem__(self, idx) -> dict:
-        """Return a single frame by index, with all transforms applied.
+    def __getitem__(self, idx: int | slice) -> dict | list[dict]:
+        """Return one frame or a slice of frames, with all transforms applied.

        Loads the frame from the underlying HF dataset, expands delta-timestamp
        windows, decodes video frames, and applies image transforms. Delegates
-        the core logic to :meth:`DatasetReader.get_item`.
+        the core logic to :class:`DatasetReader`.

        Args:
-            idx: Index into the (possibly episode-filtered) dataset.
+            idx: Integer index or slice into the possibly episode-filtered dataset.

        Returns:
-            Dict mapping feature names to their tensor values for this frame.
+            A frame dictionary for an integer index, or a list of frame
+            dictionaries for a slice.

        Raises:
            RuntimeError: If the dataset is currently being recorded and
@@ -499,6 +514,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
            raise RuntimeError(
                "Cannot read from a dataset that is being recorded. Call finalize() first, then access items."
            )
+        if isinstance(idx, slice):
+            return [self[item_idx] for item_idx in range(*idx.indices(len(self)))]
+
        reader = self._ensure_reader()
        if reader.hf_dataset is None:
            # One-shot load after finalize()
@@ -622,10 +640,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
                hub_api.delete_tag(self.repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
            hub_api.create_tag(self.repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")

-    def _download(self, download_videos: bool = True) -> None:
+    def _download(self, download_videos: bool = True, *, token: str | bool | None = None) -> None:
        """Downloads the dataset from the given 'repo_id' at the provided version."""
        ignore_patterns = None if download_videos else "videos/"
        files = None
+        token_kwargs = {} if token is None else {"token": token}
        if self.episodes is not None:
            # Reader is guaranteed to exist here (created in __init__ before _download)
            files = self.reader.get_episodes_file_paths()
@@ -639,6 +658,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
                    cache_dir=HF_LEROBOT_HUB_CACHE,
                    allow_patterns=files,
                    ignore_patterns=ignore_patterns,
+                    **token_kwargs,
                )
            )
        else:
@@ -650,6 +670,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
                local_dir=self._requested_root,
                allow_patterns=files,
                ignore_patterns=ignore_patterns,
+                **token_kwargs,
            )
            self.meta.root = self._requested_root

@@ -789,6 +810,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
        image_writer_threads: int = 0,
        streaming_encoding: bool = False,
        encoder_queue_maxsize: int = 30,
+        *,
+        token: str | bool | None = None,
    ) -> "LeRobotDataset":
        """Resume recording on an existing dataset.

@@ -822,6 +845,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
            streaming_encoding: If ``True``, encode video in real-time during
                capture.
            encoder_queue_maxsize: Max buffered frames per camera for streaming.
+            token: Authentication token used if metadata must be downloaded
+                from the Hub. The token is not retained on the dataset instance.

        Returns:
            A :class:`LeRobotDataset` in write mode, ready to append episodes.
@@ -850,7 +875,11 @@ class LeRobotDataset(torch.utils.data.Dataset):

        # Load metadata (revision-safe when root is not provided)
        obj.meta = LeRobotDatasetMetadata(
-            obj.repo_id, obj._requested_root, obj.revision, force_cache_sync=force_cache_sync
+            obj.repo_id,
+            obj._requested_root,
+            obj.revision,
+            force_cache_sync=force_cache_sync,
+            token=token,
        )

        obj._encoder_threads = encoder_threads
@@ -48,6 +48,8 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
        tolerances_s: dict | None = None,
        download_videos: bool = True,
        video_backend: str | None = None,
+        *,
+        token: str | bool | None = None,
    ):
        super().__init__()
        self.repo_ids = repo_ids
@@ -65,6 +67,7 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
                tolerance_s=self.tolerances_s[repo_id],
                download_videos=download_videos,
                video_backend=video_backend,
+                token=token,
            )
            for repo_id in repo_ids
        ]
@@ -17,8 +17,8 @@ from collections.abc import Sequence
 from typing import Any

 from lerobot.configs import PipelineFeatureType
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.processor import DataProcessorPipeline
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR
 from lerobot.utils.feature_utils import hw_to_dataset_features

@@ -58,6 +58,10 @@ class LookAheadError(Exception):
    pass


+class _ShardExhaustedError(Exception):
+    """Raised when a streaming dataset shard has no more items."""
+
+
 class Backtrackable[T]:
    """
    Wrap any iterator/iterable so you can step back up to `history` items
@@ -178,7 +182,7 @@ class Backtrackable[T]:
        """
        Check if we can go back `steps` items without raising an IndexError.
        """
-        return steps <= len(self._back_buf) + self._cursor
+        return steps < len(self._back_buf) + self._cursor

    def can_peek_ahead(self, steps: int = 1) -> bool:
        """
@@ -256,6 +260,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
        shuffle: bool = True,
        return_uint8: bool = False,
        depth_output_unit: str = DEFAULT_DEPTH_UNIT,
+        *,
+        token: str | bool | None = None,
    ):
        """Initialize a StreamingLeRobotDataset.

@@ -278,6 +284,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
            shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True.
            depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm").
                Defaults to "mm".
+            token: Authentication token used while streaming this dataset from
+                the Hub. Pass a string token, ``True`` to require the locally
+                stored token, ``False`` to disable authentication, or ``None``
+                to use the Hugging Face Hub default. The token is not retained
+                on the dataset instance after initialization.
        """
        super().__init__()
        self.repo_id = repo_id
@@ -306,7 +317,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):

        # Load metadata
        self.meta = LeRobotDatasetMetadata(
-            self.repo_id, self._requested_root, self.revision, force_cache_sync=force_cache_sync
+            self.repo_id,
+            self._requested_root,
+            self.revision,
+            force_cache_sync=force_cache_sync,
+            token=token,
        )
        self.root = self.meta.root
        self.revision = self.meta.revision
@@ -334,12 +349,14 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
            self.delta_timestamps = delta_timestamps
            self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps)

+        token_kwargs = {} if token is None or self.streaming_from_local else {"token": token}
        self.hf_dataset: datasets.IterableDataset = load_dataset(
            self.repo_id if not self.streaming_from_local else str(self.root),
            split="train",
            streaming=self.streaming,
            data_files="data/*/*.parquet",
            revision=self.revision,
+            **token_kwargs,
        )

        self.num_shards = min(self.hf_dataset.num_shards, max_num_shards)
@@ -409,10 +426,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
                    else:
                        frames_buffer.append(frame)
                    break  # random shard sampled, switch shard
-            except (
-                RuntimeError,
-                StopIteration,
-            ):  # NOTE: StopIteration inside a generator throws a RuntimeError since python 3.7
+            except _ShardExhaustedError:
                del idx_to_backtrack_dataset[shard_key]  # Remove exhausted shard, onto another shard

        # Once shards are all exhausted, shuffle the buffer and yield the remaining frames
@@ -490,7 +504,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):

    def make_frame(self, dataset_iterator: Backtrackable) -> Generator:
        """Makes a frame starting from a dataset iterator"""
-        item = next(dataset_iterator)
+        try:
+            item = next(dataset_iterator)
+        except StopIteration as e:
+            # Translate exhaustion here, before PEP 479 turns it into an indistinguishable RuntimeError.
+            raise _ShardExhaustedError from e
        item = item_to_torch(item)

        updates = []  # list of "updates" to apply to the item retrieved from hf_dataset (w/o camera features)
@@ -325,16 +325,19 @@ def check_version_compatibility(
        logging.warning(FUTURE_MESSAGE.format(repo_id=repo_id, version=v_check))


-def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
+def get_repo_versions(repo_id: str, *, token: str | bool | None = None) -> list[packaging.version.Version]:
    """Return available valid versions (branches and tags) on a given Hub repo.

    Args:
        repo_id (str): The repository ID on the Hugging Face Hub.
+        token: Authentication token used for Hub requests. Pass a string token,
+            ``True`` to require the locally stored token, ``False`` to disable
+            authentication, or ``None`` to use the Hugging Face Hub default.

    Returns:
        list[packaging.version.Version]: A list of valid versions found.
    """
-    api = HfApi()
+    api = HfApi() if token is None else HfApi(token=token)
    repo_refs = api.list_repo_refs(repo_id, repo_type="dataset")
    repo_refs = [b.name for b in repo_refs.branches + repo_refs.tags]
    repo_versions = []
@@ -345,7 +348,12 @@ def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
    return repo_versions


-def get_safe_version(repo_id: str, version: str | packaging.version.Version) -> str:
+def get_safe_version(
+    repo_id: str,
+    version: str | packaging.version.Version,
+    *,
+    token: str | bool | None = None,
+) -> str:
    """Return the specified version if available on repo, or the latest compatible one.

    If the exact version is not found, it looks for the latest version with the
@@ -354,6 +362,7 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) ->
    Args:
        repo_id (str): The repository ID on the Hugging Face Hub.
        version (str | packaging.version.Version): The target version.
+        token: Authentication token forwarded to the Hub version lookup.

    Returns:
        str: The safe version string (e.g., "v1.2.3") to use as a revision.
@@ -366,7 +375,7 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) ->
    target_version = (
        packaging.version.parse(version) if not isinstance(version, packaging.version.Version) else version
    )
-    hub_versions = get_repo_versions(repo_id)
+    hub_versions = get_repo_versions(repo_id) if token is None else get_repo_versions(repo_id, token=token)

    if not hub_versions:
        raise RevisionNotFoundError(
@@ -322,7 +322,7 @@ class HILSerlRobotEnvConfig(EnvConfig):
 class LiberoEnv(EnvConfig):
    task: str = "libero_10"  # can also choose libero_spatial, libero_object, etc.
    task_ids: list[int] | None = None
-    fps: int = 30
+    fps: int = 20  # Must match robosuite's default control_freq (20 Hz)
    episode_length: int | None = None
    obs_type: str = "pixels_agent_pos"
    render_mode: str = "rgb_array"
@@ -354,6 +354,9 @@ class LiberoEnv(EnvConfig):
    control_mode: str = "relative"  # or "absolute"

    def __post_init__(self):
+        if self.fps <= 0:
+            raise ValueError(f"fps must be positive, got {self.fps}")
+
        if self.obs_type == "pixels":
            self.features[LIBERO_KEY_PIXELS_AGENTVIEW] = PolicyFeature(
                type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
@@ -412,6 +415,7 @@ class LiberoEnv(EnvConfig):
            "render_mode": self.render_mode,
            "observation_height": self.observation_height,
            "observation_width": self.observation_width,
+            "control_freq": self.fps,
        }
        if self.task_ids is not None:
            kwargs["task_ids"] = self.task_ids
@@ -503,7 +507,7 @@ class MetaworldEnv(EnvConfig):
 class RoboCasaEnv(EnvConfig):
    task: str = "CloseFridge"
    fps: int = 20
-    episode_length: int = 1000
+    episode_length: int | None = None
    obs_type: str = "pixels_agent_pos"
    render_mode: str = "rgb_array"
    camera_name: str = "robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right"
@@ -30,7 +30,7 @@ from gymnasium import spaces
 from libero.libero import benchmark, get_libero_path
 from libero.libero.envs import OffScreenRenderEnv

-from lerobot.types import RobotObservation
+from lerobot.lerobot_types import RobotObservation

 from .utils import _LazyAsyncVectorEnv, parse_camera_names

@@ -125,10 +125,13 @@ class LiberoEnv(gym.Env):
        n_envs: int = 1,
        camera_name_mapping: dict[str, str] | None = None,
        num_steps_wait: int = 10,
+        control_freq: int = 20,
        control_mode: str = "relative",
        is_libero_plus: bool = False,
    ):
        super().__init__()
+        if control_freq <= 0:
+            raise ValueError(f"control_freq must be positive, got {control_freq}")
        self.task_id = task_id
        self.is_libero_plus = is_libero_plus
        self.obs_type = obs_type
@@ -154,6 +157,7 @@ class LiberoEnv(gym.Env):
            }
        self.camera_name_mapping = camera_name_mapping
        self.num_steps_wait = num_steps_wait
+        self.control_freq = control_freq
        self.episode_index = episode_index
        self.episode_length = episode_length
        # Load once and keep
@@ -260,6 +264,7 @@ class LiberoEnv(gym.Env):
            bddl_file_name=self._task_bddl_file,
            camera_heights=self.observation_height,
            camera_widths=self.observation_width,
+            control_freq=self.control_freq,
        )
        env.reset()
        self._env = env
@@ -379,7 +384,12 @@ class LiberoEnv(gym.Env):

    def close(self):
        if self._env is not None:
-            self._env.close()
+            try:
+                self._env.close()
+            finally:
+                # LIBERO deletes its inner env on close, so this wrapper must
+                # be recreated before the next reset.
+                self._env = None


 def _make_env_fns(
@@ -25,7 +25,7 @@ import metaworld.policies as policies
 import numpy as np
 from gymnasium import spaces

-from lerobot.types import RobotObservation
+from lerobot.lerobot_types import RobotObservation

 from .utils import _LazyAsyncVectorEnv

@@ -155,6 +155,7 @@ class MetaworldEnv(gym.Env):
            env.model.cam_pos[2] = [0.75, 0.075, 0.7]
        env.reset()
        env._freeze_rand_vec = False  # otherwise no randomization
+        env.seeded_rand_vec = True  # use seeded RNG so reset(seed=X) controls object positions
        self._env = env

    def render(self) -> np.ndarray:
@@ -220,6 +221,8 @@ class MetaworldEnv(gym.Env):
        self._ensure_env()
        super().reset(seed=seed)

+        if seed is not None:
+            self._env.seed(seed)
        raw_obs, info = self._env.reset(seed=seed)

        observation = self._format_raw_obs(raw_obs)
@@ -25,7 +25,7 @@ import gymnasium as gym
 import numpy as np
 from gymnasium import spaces

-from lerobot.types import RobotObservation
+from lerobot.lerobot_types import RobotObservation

 from .utils import _LazyAsyncVectorEnv, parse_camera_names

@@ -98,6 +98,19 @@ def _resolve_tasks(task: str) -> tuple[list[str], str | None]:
    return names, None


+def _get_task_horizon(task: str) -> int:
+    """Return the rollout horizon registered by RoboCasa for a task."""
+    from robocasa.utils.dataset_registry_utils import get_task_horizon
+
+    try:
+        return int(get_task_horizon(task))
+    except ValueError as exc:
+        raise ValueError(
+            f"No RoboCasa horizon is registered for task '{task}'. "
+            "Set `--env.episode_length=<steps>` explicitly."
+        ) from exc
+
+
 def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
    """Split a flat (12,) action vector into a RoboCasa action dict.

@@ -154,7 +167,7 @@ class RoboCasaEnv(gym.Env):

        self.camera_name = parse_camera_names(camera_name)

-        self._max_episode_steps = episode_length if episode_length is not None else 1000
+        self._max_episode_steps = episode_length if episode_length is not None else _get_task_horizon(task)

        # Deferred — created on first reset() inside the worker subprocess
        # to avoid inheriting stale GPU/EGL contexts across fork().
@@ -28,7 +28,7 @@ import numpy as np
 import torch
 from gymnasium import spaces

-from lerobot.types import RobotObservation
+from lerobot.lerobot_types import RobotObservation
 from lerobot.utils.import_utils import _scipy_available

 from .utils import _LazyAsyncVectorEnv
@@ -384,7 +384,9 @@ class RoboTwinEnv(gym.Env):

        self._env: Any | None = None  # deferred — created on first reset() inside worker
        self._step_count: int = 0
-        self._black_frame = np.zeros((self.observation_height, self.observation_width, 3), dtype=np.uint8)
+        self._black_frame: np.ndarray = np.zeros(
+            (self.observation_height, self.observation_width, 3), dtype=np.uint8
+        )

        image_spaces = {
            cam: spaces.Box(
@@ -37,7 +37,7 @@ import numpy as np
 from gymnasium import spaces
 from scipy.spatial.transform import Rotation

-from lerobot.types import RobotObservation
+from lerobot.lerobot_types import RobotObservation

 from .utils import _LazyAsyncVectorEnv

@@ -373,7 +373,7 @@ class VLABenchEnv(gym.Env):

        if action.shape[0] != 7:
            # Unknown layout — fall back to zero-pad so the sim doesn't crash.
-            padded = np.zeros(ctrl_dim, dtype=np.float64)
+            padded: np.ndarray = np.zeros(ctrl_dim, dtype=np.float64)
            padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
            return padded

@@ -18,6 +18,7 @@ from lerobot.utils.import_utils import require_package
 # guard the optional dependency here so importing this package fails loudly if it's missing.
 require_package("datasets", extra="dataset")

+from .annotate import submit_annotate_to_hf
 from .hf import submit_to_hf

-__all__ = ["submit_to_hf"]
+__all__ = ["submit_annotate_to_hf", "submit_to_hf"]
@@ -0,0 +1,176 @@
+# Copyright 2026 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.
+"""Run ``lerobot-annotate`` on HF Jobs (HuggingFace GPUs).
+
+Same shape as the training submitter in ``hf.py``, with one difference: the
+annotation pipeline serves its own VLM, so the pod starts from the official
+``vllm/vllm-openai`` image (which has no lerobot) instead of the prebuilt
+``lerobot-gpu`` image, and installs lerobot on top before running.
+
+Because there is no config repo to stage, the pod replays the user's own CLI
+flags — everything except the client-only ``--job.*`` and the host-local
+``--root``, which is replaced by ``--repo_id`` so the pod pulls the dataset
+from the Hub.
+"""
+
+from __future__ import annotations
+
+import shlex
+import sys
+from dataclasses import is_dataclass
+from typing import TYPE_CHECKING
+
+from huggingface_hub import HfApi, get_token, run_job
+
+from .dataset import ensure_dataset_available
+
+# Package-internal reuse of the training submitter's job plumbing: following a
+# submitted job and forwarding argv are identical for annotation runs.
+from .hf import _pod_forwarded_args, follow_job, resolve_job_tags
+
+if TYPE_CHECKING:
+    from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
+
+LEROBOT_GIT_URL = "https://github.com/huggingface/lerobot.git"
+
+# Mirrors the pins in pyproject.toml. The vLLM image resolves dependencies on its
+# own otherwise, and pulls av 18 / datasets 5 / draccus 0.11 — each of which breaks
+# lerobot at import time. `--upgrade-strategy only-if-needed` keeps vLLM's own
+# (torch, transformers, ...) pins intact.
+_RUNTIME_REQUIREMENTS = (
+    "'datasets>=4.7.0,<5.0.0' 'pyarrow>=21.0.0,<30.0.0' 'av>=15.0.0,<16.0.0' 'draccus==0.10.0' "
+    "'pandas>=2.0.0,<3.0.0' jsonlines gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
+    "openai"
+)
+
+# Flags the submitter resolves itself instead of forwarding verbatim: `--root`
+# names a directory only this machine has, `--repo_id` is re-emitted from the
+# config, and the config-file args name local files (rejected up front by
+# `submit_annotate_to_hf`). `--job.*` is dropped separately, by prefix; bare
+# `--job` is not, hence its entry here — it is the one arg that could smuggle a
+# remote `target` onto the pod and have the job recursively submit itself.
+_SUBMITTER_OWNED_ARGS = ("--root", "--repo_id", "--config_path", "--job")
+
+
+def _local_config_file_args(cfg: AnnotationPipelineConfig) -> list[str]:
+    """The CLI args that name a config file on the client's disk.
+
+    draccus exposes ``--config_path`` for the whole config plus a ``--<field>``
+    for every nested dataclass (``--vlm``, ``--plan``, ``--job``, ...). The pod has
+    none of those files, so a remote run has to reject them rather than silently
+    drop the settings they carry.
+    """
+    return ["--config_path", *(f"--{name}" for name in vars(cfg) if is_dataclass(getattr(cfg, name)))]
+
+
+def build_pod_setup(lerobot_ref: str) -> str:
+    """Shell prelude that turns the vLLM image into a ``lerobot-annotate`` runtime."""
+    spec = f"lerobot @ git+{LEROBOT_GIT_URL}@{lerobot_ref}"
+    return (
+        # git to install from the repo, ffmpeg to decode the dataset's videos.
+        "apt-get update -qq && apt-get install -y -qq git ffmpeg && "
+        f"pip install --no-deps {shlex.quote(spec)} && "
+        f"pip install --upgrade-strategy only-if-needed {_RUNTIME_REQUIREMENTS} && "
+        # vLLM's cudagraph memory estimate over-reserves and starves the KV cache;
+        # PyAV is the video backend the server can decode our frames with.
+        "export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
+        "export VLLM_VIDEO_BACKEND=pyav"
+    )
+
+
+def build_pod_command(repo_id: str, lerobot_ref: str, argv: list[str]) -> list[str]:
+    """Build the ``bash -c`` command the pod runs: setup prelude, then annotation.
+
+    ``argv`` is the user's CLI (``sys.argv[1:]``) minus the flags in
+    ``_SUBMITTER_OWNED_ARGS``; ``--repo_id`` is re-added from the config so the pod
+    always annotates the dataset we just made sure is reachable on the Hub.
+    ``--job.target=local`` stops the pod from re-dispatching to itself.
+    """
+    forwarded = _pod_forwarded_args(argv, drop_names=_SUBMITTER_OWNED_ARGS, drop_prefixes=("--job.",))
+    annotate = shlex.join(["lerobot-annotate", f"--repo_id={repo_id}", *forwarded, "--job.target=local"])
+    return ["bash", "-c", f"{build_pod_setup(lerobot_ref)} && {annotate}"]
+
+
+def submit_annotate_to_hf(cfg: AnnotationPipelineConfig) -> None:
+    """Submit an annotation run to HF Jobs infrastructure.
+
+    Resolves credentials, makes sure the source dataset is reachable from the pod,
+    submits the job, then tails its logs until the job reaches a terminal stage —
+    or returns immediately with ``--job.detach``. Ctrl-C detaches without
+    cancelling the remote job.
+    """
+    token = get_token()
+    if not token:
+        raise RuntimeError("Not logged in to Hugging Face. Run `hf auth login` first.")
+
+    if cfg.repo_id is None:
+        raise ValueError(
+            "Remote annotation requires --repo_id: the pod downloads the dataset from the Hub, "
+            "and --root only names a directory on this machine."
+        )
+
+    argv = sys.argv[1:]
+    passed = {tok.split("=", 1)[0] for tok in argv}
+    used_config_files = sorted(passed.intersection(_local_config_file_args(cfg)))
+    if used_config_files:
+        raise ValueError(
+            f"{', '.join(used_config_files)} cannot be used with a remote --job.target: the pod "
+            "cannot read config files from this machine. Pass the settings as CLI flags instead."
+        )
+
+    if not cfg.push_to_hub:
+        # The pod's filesystem is discarded when the job ends, so without a push the
+        # run produces nothing. Warn rather than fail: a smoke test over
+        # --only_episodes that only inspects the logs is a legitimate use.
+        print(
+            "WARNING: --push_to_hub is off. The annotated dataset lives only on the pod and is "
+            "discarded when the job ends. Pass --push_to_hub=true to keep the result."
+        )
+
+    api = HfApi(token=token)
+    tags = resolve_job_tags(cfg.job.tags)
+    ensure_dataset_available(cfg.repo_id, api=api, tags=tags)
+
+    command = build_pod_command(cfg.repo_id, cfg.job.lerobot_ref, argv)
+
+    print(f"Submitting job to HF Jobs (flavor={cfg.job.target}, image={cfg.job.image}) ...")
+    job_info = run_job(
+        image=cfg.job.image,
+        command=command,
+        flavor=cfg.job.target,
+        secrets={"HF_TOKEN": token},
+        timeout=cfg.job.timeout,
+        # HF Jobs labels are key/value; expose each tag as a queryable label.
+        labels=dict.fromkeys(tags, "true"),
+    )
+    job_id = job_info.id
+    job_url = getattr(job_info, "url", None)
+    print(f"Job submitted: {job_id}")
+    if job_url:
+        print(f"  Job page:     {job_url}")
+    target_repo_id = cfg.new_repo_id or cfg.repo_id
+    if cfg.push_to_hub:
+        print(f"  Dataset repo: https://huggingface.co/datasets/{target_repo_id}")
+    print(f"  Monitor:      hf jobs logs {job_id}")
+    print(f"  Cancel:       hf jobs cancel {job_id}")
+
+    # No success marker: `lerobot-annotate` keeps working after the upload log line
+    # (dataset card, version tag), so completion has to be stage-based.
+    if not follow_job(job_id, detach=cfg.job.detach):
+        return
+
+    if cfg.push_to_hub:
+        print(f"\nAnnotation complete — dataset pushed to https://huggingface.co/datasets/{target_repo_id}")
+    else:
+        print("\nAnnotation complete. Note: --push_to_hub was off, so the result stayed on the pod.")
@@ -223,6 +223,74 @@ def _poll_until_done(
    return None


+def follow_job(job_id: str, *, detach: bool = False, success_marker: str | None = None) -> bool:
+    """Watch a submitted job to the end, streaming its logs to stdout.
+
+    Returns True when the job finished successfully and False when we stopped watching
+    without a verdict — `detach`, or the user pressing Ctrl-C, which detaches rather than
+    cancelling the remote job. Raises RuntimeError when the job reaches a terminal stage
+    other than COMPLETED.
+
+    `success_marker` finishes as soon as that string appears in the logs instead of waiting
+    out the platform's post-run finalization (~30s). Callers that have a log line meaning
+    "the artifact is on the Hub" should pass it; without one, completion is stage-based.
+    """
+    if detach:
+        return False
+
+    done = threading.Event()
+    detached = threading.Event()
+    marker_seen = threading.Event()
+    stage_holder: dict[str, str | None] = {}
+
+    def _poll() -> None:
+        stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder)
+
+    poll_thread = threading.Thread(target=_poll, daemon=True)
+    poll_thread.start()
+    log_thread = threading.Thread(
+        target=_tail_logs, args=(job_id, done, success_marker, marker_seen), daemon=True
+    )
+    log_thread.start()
+
+    def _detach(sig, frame):
+        detached.set()
+        done.set()
+        print("\nDetached. Job is still running.")
+        print(f"  Monitor: hf jobs logs {job_id}")
+        print(f"  Cancel:  hf jobs cancel {job_id}")
+
+    # signal.signal only works on the main thread; when called from a worker thread
+    # (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels
+    # handler rather than crashing with ValueError.
+    install_sigint = threading.current_thread() is threading.main_thread()
+    original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None
+    if install_sigint:
+        signal.signal(signal.SIGINT, _detach)
+    try:
+        # Timeout-based join so SIGINT is delivered to the main thread promptly.
+        while poll_thread.is_alive():
+            poll_thread.join(timeout=0.5)
+        log_thread.join(timeout=5)
+    finally:
+        if install_sigint:
+            signal.signal(signal.SIGINT, original_sigint)
+
+    if detached.is_set():
+        return False
+    if marker_seen.is_set():
+        return True
+
+    stage = stage_holder.get("stage")
+    if stage != "COMPLETED":
+        message = stage_holder.get("message")
+        detail = f" ({message})" if message else ""
+        raise RuntimeError(
+            f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}"
+        )
+    return True
+
+
 def _pod_forwarded_args(
    argv: list[str], drop_names: tuple[str, ...] = (), drop_prefixes: tuple[str, ...] = ()
 ) -> list[str]:
@@ -362,64 +430,11 @@ def submit_to_hf(cfg: TrainPipelineConfig) -> None:
    print(f"  Monitor:    hf jobs logs {job_id}")
    print(f"  Cancel:     hf jobs cancel {job_id}")

-    if cfg.job.detach:
-        return
-
-    done = threading.Event()
-    detached = threading.Event()
-    pushed_ok = threading.Event()
-    stage_holder: dict[str, str | None] = {}
-
-    def _poll() -> None:
-        stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder)
-
-    poll_thread = threading.Thread(target=_poll, daemon=True)
-    poll_thread.start()
    # Finish as soon as the model is pushed, rather than waiting out the platform's
    # post-run finalization before the job stage flips to COMPLETED. This matches the
    # exact log line emitted by PreTrainedPolicy.push_model_to_hub — the two must stay
    # in sync. If it ever stops matching we just fall back to stage-based completion
    # (~30s slower), so the contract is an optimization, not a correctness requirement.
    success_marker = f"Model pushed to https://huggingface.co/{repo_id}"
-    log_thread = threading.Thread(
-        target=_tail_logs, args=(job_id, done, success_marker, pushed_ok), daemon=True
-    )
-    log_thread.start()
-
-    def _detach(sig, frame):
-        detached.set()
-        done.set()
-        print("\nDetached. Job is still running.")
-        print(f"  Monitor: hf jobs logs {job_id}")
-        print(f"  Cancel:  hf jobs cancel {job_id}")
-
-    # signal.signal only works on the main thread; when called from a worker thread
-    # (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels
-    # handler rather than crashing with ValueError.
-    install_sigint = threading.current_thread() is threading.main_thread()
-    original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None
-    if install_sigint:
-        signal.signal(signal.SIGINT, _detach)
-    try:
-        # Timeout-based join so SIGINT is delivered to the main thread promptly.
-        while poll_thread.is_alive():
-            poll_thread.join(timeout=0.5)
-        log_thread.join(timeout=5)
-    finally:
-        if install_sigint:
-            signal.signal(signal.SIGINT, original_sigint)
-
-    if detached.is_set():
-        return
-
-    if pushed_ok.is_set():
+    if follow_job(job_id, detach=cfg.job.detach, success_marker=success_marker):
        print(f"\nTraining complete — model pushed to https://huggingface.co/{repo_id}")
-        return
-
-    stage = stage_holder.get("stage")
-    if stage != "COMPLETED":
-        message = stage_holder.get("message")
-        detail = f" ({message})" if message else ""
-        raise RuntimeError(
-            f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}"
-        )
@@ -20,7 +20,6 @@ import logging
 import time
 from contextlib import contextmanager
 from copy import deepcopy
-from functools import cached_property
 from typing import TYPE_CHECKING, Any, TypedDict

 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
@@ -854,7 +853,7 @@ class DamiaoMotorsBus(MotorsBusBase):
        else:
            raise ValueError(f"Motor {motor_obj} doesn't have a valid recv_id (None).")

-    @cached_property
+    @property
    def is_calibrated(self) -> bool:
        """Check if motors are calibrated."""
        return bool(self.calibration)
@@ -122,6 +122,9 @@ MODEL_ENCODING_TABLE = {
    "xm430-w350": X_SERIES_ENCODINGS_TABLE,
    "xm540-w270": X_SERIES_ENCODINGS_TABLE,
    "xc430-w150": X_SERIES_ENCODINGS_TABLE,
+    "xh540-w150": X_SERIES_ENCODINGS_TABLE,
+    "xc330-t288": X_SERIES_ENCODINGS_TABLE,
+    "xc330-t181": X_SERIES_ENCODINGS_TABLE,
 }

 # {model: model_resolution}
@@ -134,6 +137,9 @@ MODEL_RESOLUTION = {
    "xm430-w350": 4096,
    "xm540-w270": 4096,
    "xc430-w150": 4096,
+    "xh540-w150": 4096,
+    "xc330-t288": 4096,
+    "xc330-t181": 4096,
 }

 # {model: model_number}
@@ -145,6 +151,9 @@ MODEL_NUMBER_TABLE = {
    "xm430-w350": 1020,
    "xm540-w270": 1120,
    "xc430-w150": 1070,
+    "xh540-w150": 1110,
+    "xc330-t288": 1220,
+    "xc330-t181": 1210,
 }

 # {model: available_operating_modes}
@@ -156,6 +165,9 @@ MODEL_OPERATING_MODES = {
    "xm430-w350": [0, 1, 3, 4, 5, 16],
    "xm540-w270": [0, 1, 3, 4, 5, 16],
    "xc430-w150": [1, 3, 4, 16],
+    "xh540-w150": [0, 1, 3, 4, 5, 16],
+    "xc330-t288": [0, 1, 3, 4, 5, 16],
+    "xc330-t181": [0, 1, 3, 4, 5, 16],
 }

 MODEL_CONTROL_TABLE = {
@@ -166,6 +178,9 @@ MODEL_CONTROL_TABLE = {
    "xm430-w350": X_SERIES_CONTROL_TABLE,
    "xm540-w270": X_SERIES_CONTROL_TABLE,
    "xc430-w150": X_SERIES_CONTROL_TABLE,
+    "xh540-w150": X_SERIES_CONTROL_TABLE,
+    "xc330-t288": X_SERIES_CONTROL_TABLE,
+    "xc330-t181": X_SERIES_CONTROL_TABLE,
 }

 MODEL_BAUDRATE_TABLE = {
@@ -176,6 +191,9 @@ MODEL_BAUDRATE_TABLE = {
    "xm430-w350": X_SERIES_BAUDRATE_TABLE,
    "xm540-w270": X_SERIES_BAUDRATE_TABLE,
    "xc430-w150": X_SERIES_BAUDRATE_TABLE,
+    "xh540-w150": X_SERIES_BAUDRATE_TABLE,
+    "xc330-t288": X_SERIES_BAUDRATE_TABLE,
+    "xc330-t181": X_SERIES_BAUDRATE_TABLE,
 }

 AVAILABLE_BAUDRATES = [
@@ -23,6 +23,7 @@ from __future__ import annotations

 import abc
 import logging
+import time
 from collections.abc import Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
@@ -818,13 +819,13 @@ class SerialMotorsBus(MotorsBusBase):
        """
        motor_names = self._get_motors_list(motors)

-        start_positions = self.sync_read("Present_Position", motor_names, normalize=False)
+        start_positions = self.sync_read("Present_Position", motor_names, normalize=False, num_retry=5)
        mins = start_positions.copy()
        maxes = start_positions.copy()

        user_pressed_enter = False
        while not user_pressed_enter:
-            positions = self.sync_read("Present_Position", motor_names, normalize=False)
+            positions = self.sync_read("Present_Position", motor_names, normalize=False, num_retry=5)
            mins = {motor: min(positions[motor], min_) for motor, min_ in mins.items()}
            maxes = {motor: max(positions[motor], max_) for motor, max_ in maxes.items()}

@@ -837,9 +838,12 @@ class SerialMotorsBus(MotorsBusBase):
            if enter_pressed():
                user_pressed_enter = True

-            if display_values and not user_pressed_enter:
-                # Move cursor up to overwrite the previous output
-                move_cursor_up(len(motor_names) + 3)
+            if not user_pressed_enter:
+                if display_values:
+                    # Move cursor up to overwrite the previous output
+                    move_cursor_up(len(motor_names) + 3)
+                # Throttle reads even when the live table is disabled.
+                time.sleep(0.02)

        same_min_max = [motor for motor in motor_names if mins[motor] == maxes[motor]]
        if same_min_max:
@@ -79,6 +79,8 @@ class DiffusionConfig(PreTrainedConfig):
        use_film_scale_modulation: FiLM (https://huggingface.co/papers/1709.07871) is used for the Unet conditioning.
            Bias modulation is used be default, while this parameter indicates whether to also use scale
            modulation.
+        gradient_checkpointing: Whether to checkpoint the Unet residual blocks during training. This reduces
+            activation memory at the cost of recomputing those blocks during the backward pass.
        noise_scheduler_type: Name of the noise scheduler to use. Supported options: ["DDPM", "DDIM"].
        num_train_timesteps: Number of diffusion steps for the forward diffusion schedule.
        beta_schedule: Name of the diffusion beta schedule as per DDPMScheduler from Hugging Face diffusers.
@@ -132,6 +134,7 @@ class DiffusionConfig(PreTrainedConfig):
    n_groups: int = 8
    diffusion_step_embed_dim: int = 128
    use_film_scale_modulation: bool = True
+    gradient_checkpointing: bool = False
    # Noise scheduler.
    noise_scheduler_type: str = "DDPM"
    num_train_timesteps: int = 100
@@ -31,6 +31,7 @@ import torch
 import torch.nn.functional as F  # noqa: N812
 import torchvision
 from torch import Tensor, nn
+from torch.utils.checkpoint import checkpoint

 from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
 from lerobot.utils.import_utils import _diffusers_available, require_package
@@ -727,22 +728,35 @@ class DiffusionConditionalUnet1d(nn.Module):
        else:
            global_feature = timesteps_embed

+        use_gc = self.config.gradient_checkpointing and self.training
+
        # Run encoder, keeping track of skip features to pass to the decoder.
        encoder_skip_features: list[Tensor] = []
        for resnet, resnet2, downsample in self.down_modules:
-            x = resnet(x, global_feature)
-            x = resnet2(x, global_feature)
+            if use_gc:
+                x = checkpoint(resnet, x, global_feature, use_reentrant=False)
+                x = checkpoint(resnet2, x, global_feature, use_reentrant=False)
+            else:
+                x = resnet(x, global_feature)
+                x = resnet2(x, global_feature)
            encoder_skip_features.append(x)
            x = downsample(x)

        for mid_module in self.mid_modules:
-            x = mid_module(x, global_feature)
+            if use_gc:
+                x = checkpoint(mid_module, x, global_feature, use_reentrant=False)
+            else:
+                x = mid_module(x, global_feature)

        # Run decoder, using the skip features from the encoder.
        for resnet, resnet2, upsample in self.up_modules:
            x = torch.cat((x, encoder_skip_features.pop()), dim=1)
-            x = resnet(x, global_feature)
-            x = resnet2(x, global_feature)
+            if use_gc:
+                x = checkpoint(resnet, x, global_feature, use_reentrant=False)
+                x = checkpoint(resnet2, x, global_feature, use_reentrant=False)
+            else:
+                x = resnet(x, global_feature)
+                x = resnet2(x, global_feature)
            x = upsample(x)

        x = self.final_conv(x)
@@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, Any
 import torch

 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import TransitionKey
 from lerobot.processor import (
    ComplementaryDataProcessorStep,
    PolicyAction,
@@ -31,7 +32,6 @@ from lerobot.processor import (
    make_default_policy_processor_steps,
    make_policy_processor_pipelines,
 )
-from lerobot.types import TransitionKey
 from lerobot.utils.constants import OBS_STATE
 from lerobot.utils.import_utils import _transformers_available, require_package

@@ -42,6 +42,9 @@ class Evo1Policy(PreTrainedPolicy):
    config_class = Evo1Config
    name = "evo1"

+    def supports_rtc(self) -> bool:
+        return True
+
    def __init__(self, config: Evo1Config, *, vlm_hub_kwargs: dict | None = None, **kwargs):
        super().__init__(config)
        config.validate_features()
@@ -21,6 +21,7 @@ from typing import Any
 import torch

 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
    AddBatchDimensionProcessorStep,
    DeviceProcessorStep,
@@ -40,7 +41,6 @@ from lerobot.processor.converters import (
    policy_action_to_transition,
    transition_to_policy_action,
 )
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
    ACTION,
    DONE,
@@ -302,6 +302,33 @@ def _pad_evo1_stats(
    return padded_stats


+def _refresh_evo1_normalization_steps(
+    config: Evo1Config,
+    preprocessor: PolicyProcessorPipeline,
+    postprocessor: PolicyProcessorPipeline,
+) -> None:
+    """Re-pad checkpoint-loaded (un)normalizer stats/features to EVO1's fixed widths.
+
+    Loading a checkpoint injects the raw dataset stats (unpadded to max_state_dim/max_action_dim)
+    into the (un)normalizer via the generic override path in make_pre_post_processors. Those stats
+    and their declared features must be re-padded/reshaped to EVO1's fixed widths, otherwise
+    normalization fails against the padded state/action tensors (e.g. state padded to 24 vs. 8-dim
+    LIBERO stats). Padding is a no-op when stats are already at the target width.
+    """
+    normalization_features = _evo1_normalization_features(config)
+    action_features = _evo1_action_features(config)
+    for step in preprocessor.steps:
+        if isinstance(step, NormalizerProcessorStep):
+            step.features = normalization_features
+            step.stats = _pad_evo1_stats(config, step.stats)
+            step.to(device=step.device, dtype=step.dtype)
+    for step in postprocessor.steps:
+        if isinstance(step, UnnormalizerProcessorStep):
+            step.features = action_features
+            step.stats = _pad_evo1_stats(config, step.stats)
+            step.to(device=step.device, dtype=step.dtype)
+
+
 def reconcile_evo1_processors(
    config: Evo1Config,
    preprocessor: PolicyProcessorPipeline,
@@ -309,16 +336,19 @@ def reconcile_evo1_processors(
 ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
    """Reconcile checkpoint-loaded pipelines with the current EVO1 config.

-    Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
-    (converters are plain functions and are never serialized), and eval-time CLI overrides of the
-    action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This
-    restores the converter and rebuilds the action step from the current config so those overrides
-    take effect.
+    Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
+    (converters are plain functions and are never serialized), eval-time CLI overrides of the
+    action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the
+    (un)normalizer stats/features when the generic override path injects raw, unpadded dataset
+    stats. This restores the converter, re-pads the normalization stats to EVO1's fixed widths, and
+    rebuilds the action step from the current config so those overrides take effect.
    """
    # Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
    # non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
    preprocessor.to_transition = evo1_batch_to_transition

+    _refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
+
    action_step = Evo1ActionProcessorStep(
        action_dim=_evo1_action_dim(config),
        binarize_gripper=config.binarize_gripper,
@@ -28,6 +28,7 @@ if TYPE_CHECKING:

 from lerobot.configs import FeatureType, PreTrainedConfig
 from lerobot.envs import EnvConfig, env_to_policy_features
+from lerobot.lerobot_types import PolicyAction
 from lerobot.processor import (
    AbsoluteActionsProcessorStep,
    PolicyProcessorPipeline,
@@ -37,19 +38,25 @@ from lerobot.processor import (
    transition_to_batch,
    transition_to_policy_action,
 )
-from lerobot.types import PolicyAction
 from lerobot.utils.constants import (
    ACTION,
    POLICY_POSTPROCESSOR_DEFAULT_NAME,
    POLICY_PREPROCESSOR_DEFAULT_NAME,
 )
 from lerobot.utils.feature_utils import dataset_to_policy_features
+from lerobot.utils.import_utils import _peft_available, require_package

 from .evo1.configuration_evo1 import Evo1Config
 from .groot.configuration_groot import GrootConfig
 from .pretrained import PreTrainedPolicy
 from .utils import validate_visual_features_consistency

+if TYPE_CHECKING or _peft_available:
+    from peft import PeftConfig, PeftModel
+else:
+    PeftConfig = None
+    PeftModel = None
+

 def _reconnect_relative_absolute_steps(
    preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
@@ -177,6 +184,7 @@ def make_pre_post_processors(
            return make_groot_pre_post_processors_from_pretrained(
                config=policy_cfg,
                pretrained_path=pretrained_path,
+                revision=pretrained_revision,
                dataset_stats=kwargs.get("dataset_stats"),
                dataset_meta=kwargs.get("dataset_meta"),
                preprocessor_overrides=kwargs.get("preprocessor_overrides"),
@@ -333,12 +341,15 @@ def make_policy(
        # Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
        # of the adapter and the adapter's config contains the path to the base policy. So we need the
        # adapter config first, then load the correct policy and then apply PEFT.
-        from peft import PeftConfig, PeftModel
+        require_package("peft", extra="peft")

        logging.info("Loading policy's PEFT adapter.")

        peft_pretrained_path = str(cfg.pretrained_path)
-        peft_config = PeftConfig.from_pretrained(peft_pretrained_path)
+        peft_config = PeftConfig.from_pretrained(
+            peft_pretrained_path,
+            revision=cfg.pretrained_revision,
+        )

        kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path
        if not kwargs["pretrained_name_or_path"]:
@@ -349,9 +360,14 @@ def make_policy(
                "the adapter was trained."
            )

+        kwargs["revision"] = peft_config.revision
        policy = policy_cls.from_pretrained(**kwargs)
        policy = PeftModel.from_pretrained(
-            policy, peft_pretrained_path, config=peft_config, is_trainable=True
+            policy,
+            peft_pretrained_path,
+            config=peft_config,
+            revision=cfg.pretrained_revision,
+            is_trainable=True,
        )

    else:
@@ -37,13 +37,19 @@ def is_image_feature(key: str) -> bool:
@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.
+
+    ``multiprocessing_context`` selects the process-wide start method when
+    processes are used. Set it to ``None`` to preserve Python's default or a
+    method already selected by the embedding application.
    """

    actor: str = "threads"
    learner: str = "threads"
+    multiprocessing_context: str | None = "spawn"


@dataclass
@@ -68,6 +68,9 @@ class GrootPolicy(PreTrainedPolicy):
    name = "groot"
    config_class = GrootConfig

+    def supports_rtc(self) -> bool:
+        return True
+
    def __init__(self, config: GrootConfig, **kwargs):
        """Initialize Groot policy wrapper."""
        require_package("transformers", extra="groot")
@@ -50,6 +50,7 @@ if TYPE_CHECKING or _datasets_available:
 else:
    LeRobotDataset = None

+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
    AbsoluteActionsProcessorStep,
    AddBatchDimensionProcessorStep,
@@ -66,7 +67,6 @@ from lerobot.processor import (
    transition_to_batch,
    transition_to_policy_action,
 )
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
    ACTION,
    OBS_IMAGE,
@@ -475,6 +475,7 @@ def make_groot_pre_post_processors_from_pretrained(
    config: GrootConfig,
    pretrained_path: str,
    *,
+    revision: str | None = None,
    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
    dataset_meta: Any | None = None,
    preprocessor_overrides: dict[str, Any] | None = None,
@@ -511,6 +512,7 @@ def make_groot_pre_post_processors_from_pretrained(

    preprocessor, postprocessor = _load_groot_processor_pipelines(
        pretrained_path,
+        revision=revision,
        preprocessor_overrides=preprocessor_overrides,
        postprocessor_overrides=postprocessor_overrides,
        preprocessor_config_filename=preprocessor_config_filename,
@@ -526,6 +528,7 @@ def make_groot_pre_post_processors_from_pretrained(
 def _load_groot_processor_pipelines(
    pretrained_path: str,
    *,
+    revision: str | None,
    preprocessor_overrides: dict[str, Any],
    postprocessor_overrides: dict[str, Any],
    preprocessor_config_filename: str,
@@ -540,6 +543,7 @@ def _load_groot_processor_pipelines(
    preprocessor = PolicyProcessorPipeline.from_pretrained(
        pretrained_model_name_or_path=pretrained_path,
        config_filename=preprocessor_config_filename,
+        revision=revision,
        overrides=preprocessor_overrides,
        to_transition=batch_to_transition,
        to_output=transition_to_batch,
@@ -547,6 +551,7 @@ def _load_groot_processor_pipelines(
    postprocessor = PolicyProcessorPipeline.from_pretrained(
        pretrained_model_name_or_path=pretrained_path,
        config_filename=postprocessor_config_filename,
+        revision=revision,
        overrides=postprocessor_overrides,
        to_transition=policy_action_to_transition,
        to_output=transition_to_policy_action,
@@ -43,11 +43,22 @@ from torch.distributions import Beta

 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.utils.constants import ACTION
-from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
+from lerobot.utils.import_utils import (
+    _peft_available,
+    _scipy_available,
+    _transformers_available,
+    require_package,
+)

 from ..rtc.modeling_rtc import RTCProcessor
 from .configuration_molmoact2 import MolmoAct2Config

+if TYPE_CHECKING or _peft_available:
+    from peft import LoraConfig, get_peft_model
+else:
+    LoraConfig = None
+    get_peft_model = None
+
 logger = logging.getLogger(__name__)


@@ -509,6 +520,9 @@ class MolmoAct2Policy(PreTrainedPolicy):
    config_class = MolmoAct2Config
    name = "molmoact2"

+    def supports_rtc(self) -> bool:
+        return self.config.inference_action_mode == "continuous"
+
    def __init__(
        self,
        config: MolmoAct2Config,
@@ -1731,13 +1745,11 @@ class MolmoAct2Policy(PreTrainedPolicy):

    def _build_inner_lora_config(self):
        require_package("peft", extra="molmoact2")
-        from peft import LoraConfig

        return LoraConfig(**self._get_inner_peft_targets())

    def _apply_lora_adapters(self) -> None:
        require_package("peft", extra="molmoact2")
-        from peft import get_peft_model

        peft_config = self._build_inner_lora_config()
        self._validate_peft_config(peft_config)
@@ -36,6 +36,7 @@ import torch
 from torch import Tensor

 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
    AddBatchDimensionProcessorStep,
    DeviceProcessorStep,
@@ -49,7 +50,6 @@ from lerobot.processor import (
    policy_action_to_transition,
    transition_to_policy_action,
 )
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
    ACTION,
    OBS_IMAGES,
@@ -749,6 +749,9 @@ class PI0Policy(PreTrainedPolicy):
    config_class = PI0Config
    name = "pi0"

+    def supports_rtc(self) -> bool:
+        return True
+
    def __init__(
        self,
        config: PI0Config,
@@ -524,8 +524,6 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`

    def embed_suffix(self, noisy_actions, timestep):
        """Embed noisy_actions, timestep to prepare for Expert Gemma processing."""
-        embs = []
-        pad_masks = []
        att_masks = []

        # Embed timestep using sine-cosine positional encoding
@@ -551,23 +549,17 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
            return F.silu(x)

        time_emb = self._apply_checkpoint(time_mlp_func, time_emb)
-        action_time_emb = action_emb
        adarms_cond = time_emb

-        embs.append(action_time_emb)
-        bsize, action_time_dim = action_time_emb.shape[:2]
-        action_time_mask = torch.ones(bsize, action_time_dim, dtype=torch.bool, device=timestep.device)
-        pad_masks.append(action_time_mask)
+        bsize, action_time_dim = action_emb.shape[:2]
+        pad_masks = torch.ones(bsize, action_time_dim, dtype=torch.bool, device=timestep.device)

        # Set attention masks so that image, language and state inputs do not attend to action tokens
        att_masks += [1] + ([0] * (self.config.chunk_size - 1))
-
-        embs = torch.cat(embs, dim=1)
-        pad_masks = torch.cat(pad_masks, dim=1)
-        att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
+        att_masks = torch.tensor(att_masks, dtype=action_emb.dtype, device=action_emb.device)
        att_masks = att_masks[None, :].expand(bsize, len(att_masks))

-        return embs, pad_masks, att_masks, adarms_cond
+        return action_emb, pad_masks, att_masks, adarms_cond

    def forward(self, images, img_masks, tokens, masks, actions, noise, time) -> Tensor:
        """Do a full training forward pass and compute the loss."""
@@ -722,6 +714,9 @@ class PI05Policy(PreTrainedPolicy):
    config_class = PI05Config
    name = "pi05"

+    def supports_rtc(self) -> bool:
+        return True
+
    def __init__(
        self,
        config: PI05Config,
@@ -22,6 +22,7 @@ import numpy as np
 import torch

 from lerobot.configs import PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
    AbsoluteActionsProcessorStep,
    PolicyAction,
@@ -33,7 +34,6 @@ from lerobot.processor import (
    make_default_policy_processor_steps,
    make_policy_processor_pipelines,
 )
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import OBS_STATE

 from .configuration_pi05 import PI05Config
@@ -22,6 +22,7 @@ import numpy as np
 import torch

 from lerobot.configs import PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
    AbsoluteActionsProcessorStep,
    ActionTokenizerProcessorStep,
@@ -34,7 +35,6 @@ from lerobot.processor import (
    make_default_policy_processor_steps,
    make_policy_processor_pipelines,
 )
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import OBS_STATE

 from .configuration_pi0_fast import PI0FastConfig
@@ -34,14 +34,22 @@ from lerobot.configs import PreTrainedConfig
 from lerobot.configs.train import TrainPipelineConfig
 from lerobot.utils.device_utils import resolve_safetensors_device
 from lerobot.utils.hub import HubMixin
+from lerobot.utils.import_utils import _peft_available, require_package

 from .utils import log_model_loading_keys

-T = TypeVar("T", bound="PreTrainedPolicy")
+if TYPE_CHECKING or _peft_available:
+    from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType, get_peft_model
+else:
+    PEFT_TYPE_TO_CONFIG_MAPPING = None
+    PeftType = None
+    get_peft_model = None

 if TYPE_CHECKING:
    from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata

+T = TypeVar("T", bound="PreTrainedPolicy")
+

 def _build_card_context(
    cfg: TrainPipelineConfig | None,
@@ -241,6 +249,10 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
        """
        raise NotImplementedError

+    def supports_rtc(self) -> bool:
+        """Whether this policy implements Real-Time Chunking inference semantics."""
+        return False
+
    # TODO(aliberts, rcadene): split into 'forward' and 'compute_loss'?
    @abc.abstractmethod
    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
@@ -384,7 +396,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
            peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.)
                These are merged with policy defaults to build the final config.
        """
-        from peft import get_peft_model
+        require_package("peft", extra="peft")

        # If user provided a complete config, use it directly (with overrides)
        if peft_config is not None:
@@ -455,7 +467,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
        Returns:
            Preprocessed dict with renamed keys and init_type mapped to method-specific key.
        """
-        from peft import PeftType
+        require_package("peft", extra="peft")

        cli_overrides = cli_overrides.copy()

@@ -480,7 +492,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):

    def _build_peft_config(self, cli_overrides: dict):
        """Build a PEFT config from policy defaults and CLI overrides."""
-        from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
+        require_package("peft", extra="peft")

        # Determine PEFT method type (default to LORA)
        method_type_str = cli_overrides.get("method_type") or "lora"
@@ -507,7 +519,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):

    def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
        """Apply CLI overrides to an existing PEFT config."""
-        from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
+        require_package("peft", extra="peft")

        # Get method type from existing config or CLI override
        method_type_str = cli_overrides.get("method_type")
@@ -236,16 +236,11 @@ class ActionQueue:
        if action_index_before_inference is not None:
            indexes_diff = max(0, self.last_index - action_index_before_inference)
            if indexes_diff != real_delay:
-                # take the min of both to avoid discarding actions that were not
-                # actually consumed during inference, which would cause a jump in the queue
-                resolved = min(real_delay, indexes_diff)
                logger.warning(
-                    "Indexes diff != real delay (indexes_diff=%d, real_delay=%d); "
-                    "clamping discard to %d to avoid a queue-splice jump.",
+                    "Indexes diff is not equal to real delay. indexes_diff=%d, real_delay=%d",
                    indexes_diff,
                    real_delay,
-                    resolved,
                )
-                return resolved
+                return real_delay

        return effective_delay
@@ -145,6 +145,9 @@ class SmolVLAPolicy(PreTrainedPolicy):
    config_class = SmolVLAConfig
    name = "smolvla"

+    def supports_rtc(self) -> bool:
+        return True
+
    def __init__(
        self,
        config: SmolVLAConfig,
@@ -168,14 +168,23 @@ class SmolVLMWithExpertModel(nn.Module):
                last_layers.append(self.num_vlm_layers - 2)
            frozen_layers = [
                "lm_head",
-                "text_model.model.norm.weight",
+                "text_model.norm.weight",
            ]
            for layer in last_layers:
-                frozen_layers.append(f"text_model.model.layers.{layer}.")
+                frozen_layers.append(f"text_model.layers.{layer}.")

+            unmatched_patterns = set(frozen_layers)
            for name, params in self.vlm.named_parameters():
-                if any(k in name for k in frozen_layers):
+                matched_patterns = [k for k in frozen_layers if k in name]
+                if matched_patterns:
                    params.requires_grad = False
+                    unmatched_patterns.difference_update(matched_patterns)
+            if unmatched_patterns:
+                raise RuntimeError(
+                    "Some frozen layer patterns matched no VLM parameters, so the corresponding layers "
+                    "would silently remain trainable (parameter naming may have changed in transformers): "
+                    f"{sorted(unmatched_patterns)}"
+                )
        # To avoid unused params issue with distributed training
        for name, params in self.lm_expert.named_parameters():
            if "lm_head" in name:
@@ -22,7 +22,7 @@ import torch
 from torch import nn

 from lerobot.configs import FeatureType, PolicyFeature, PreTrainedConfig
-from lerobot.types import PolicyAction, RobotAction, RobotObservation
+from lerobot.lerobot_types import PolicyAction, RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.feature_utils import build_dataset_frame

@@ -15,11 +15,13 @@
 from __future__ import annotations

 from dataclasses import dataclass, field
+from typing import Any

 from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import NormalizationMode
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
+from lerobot.utils.constants import OBS_STATE


@PreTrainedConfig.register_subclass("vla_jepa")
@@ -122,6 +124,13 @@ class VLAJEPAConfig(PreTrainedConfig):
        if self.robot_state_feature is not None:
            self.state_dim = self.robot_state_feature.shape[0]

+    def set_dataset_feature_metadata(self, dataset_features: dict[str, Any]) -> None:
+        """Add `observation.state` to `input_features` if missing, so it gets normalized."""
+        if OBS_STATE in self.input_features or OBS_STATE not in dataset_features:
+            return
+        shape = tuple(dataset_features[OBS_STATE]["shape"])
+        self.input_features[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=shape)
+
    def get_optimizer_preset(self) -> AdamWConfig:
        return AdamWConfig(
            lr=self.optimizer_lr,
@@ -399,7 +399,8 @@ class VLAJEPAPolicy(PreTrainedPolicy):
        state = batch.get(OBS_STATE)
        if state is not None:
            if state.ndim > 2:
-                state = state[:, -1, :]
+                # deltas are forward-looking here, so index 0 is the current observation, not -1.
+                state = state[:, 0, :]
            inputs["state"] = (state.unsqueeze(1) if state.ndim == 2 else state).float()  # [B, 1, dim]

        return inputs
@@ -150,7 +150,7 @@ class XVLAModel(nn.Module):
        # Freeze or unfreeze policy transformer
        if not self.config.train_policy_transformer:
            for name, param in self.transformer.named_parameters():
-                if "soft_prompts" not in name:
+                if "soft_prompt" not in name:
                    param.requires_grad = False

        # Freeze or unfreeze soft prompts
@@ -21,6 +21,7 @@ import numpy as np
 import torch

 from lerobot.configs import PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
    ObservationProcessorStep,
    PolicyAction,
@@ -31,7 +32,6 @@ from lerobot.processor import (
    make_default_policy_processor_steps,
    make_policy_processor_pipelines,
 )
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
    IMAGENET_STATS,
    OBS_IMAGES,
@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

-from lerobot.types import (
+from lerobot.lerobot_types import (
    EnvAction,
    EnvTransition,
    PolicyAction,
@@ -25,7 +25,7 @@ from dataclasses import dataclass, field
 from torch import Tensor

 from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.types import EnvTransition, PolicyAction
+from lerobot.lerobot_types import EnvTransition, PolicyAction
 from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE

 from .pipeline import (
--- a/Show More
+++ b/Show More